Refactor to move feature version checking logic into CheckoutWouldStompWorkingSetChanges

go/libraries/doltcore/env/actions/checkout.go

@@ -471,7 +471,30 @@ func writeTableHashes(ctx context.Context, head *doltdb.RootValue, tblHashes map
 // This means that if both working sets are present (ie there are changes on both source and dest branches),
 // we check if the changes are identical before allowing a clobbering checkout.
 // Working set errors are ignored by this function, because they are properly handled elsewhere.
-func CheckoutWouldStompWorkingSetChanges(sourceRoots, destRoots doltdb.Roots) bool {
+func CheckoutWouldStompWorkingSetChanges(ctx context.Context, sourceRoots, destRoots doltdb.Roots) (bool, error) {
+
+	wouldStomp := doRootsHaveIncompatibleChanges(sourceRoots, destRoots)
+
+	if !wouldStomp {
+		return false, nil
+	}
+
+	// In some cases, a working set differs from its head only by the feature version.
+	// If this is the case, moving the working set is safe.
+	modifiedSourceRoots, err := clearFeatureVersion(ctx, sourceRoots)
+	if err != nil {
+		return true, err
+	}
+
+	modifiedDestRoots, err := clearFeatureVersion(ctx, destRoots)
+	if err != nil {
+		return true, err
+	}
+
+	return doRootsHaveIncompatibleChanges(modifiedSourceRoots, modifiedDestRoots), nil
+}
+
+func doRootsHaveIncompatibleChanges(sourceRoots, destRoots doltdb.Roots) bool {
 	sourceHasChanges, sourceWorkingHash, sourceStagedHash, err := RootHasUncommittedChanges(sourceRoots)
 	if err != nil {
 		return false
@@ -487,6 +510,32 @@ func CheckoutWouldStompWorkingSetChanges(sourceRoots, destRoots doltdb.Roots) bo
 	return sourceHasChanges && destHasChanges && (sourceWorkingHash != destWorkingHash || sourceStagedHash != destStagedHash)
 }
 
+// clearFeatureVersion creates a new version of the provided roots where all three roots have the same
+// feature version. By hashing these new roots, we can easily determine whether the roots differ only by
+// their feature version.
+func clearFeatureVersion(ctx context.Context, roots doltdb.Roots) (doltdb.Roots, error) {
+	currentBranchFeatureVersion, _, err := roots.Head.GetFeatureVersion(ctx)
+	if err != nil {
+		return doltdb.Roots{}, err
+	}
+
+	modifiedWorking, err := roots.Working.SetFeatureVersion(currentBranchFeatureVersion)
+	if err != nil {
+		return doltdb.Roots{}, err
+	}
+
+	modifiedStaged, err := roots.Staged.SetFeatureVersion(currentBranchFeatureVersion)
+	if err != nil {
+		return doltdb.Roots{}, err
+	}
+
+	return doltdb.Roots{
+		Head:    roots.Head,
+		Working: modifiedWorking,
+		Staged:  modifiedStaged,
+	}, nil
+}
+
 // RootHasUncommittedChanges returns whether the roots given have uncommitted changes, and the hashes of the working and staged roots
 func RootHasUncommittedChanges(roots doltdb.Roots) (hasChanges bool, workingHash hash.Hash, stagedHash hash.Hash, err error) {
 	headHash, err := roots.Head.HashOf()

go/libraries/doltcore/sqle/dprocedures/dolt_checkout_helpers.go

@@ -15,7 +15,6 @@
 package dprocedures
 
 import (
-	"context"
 	"errors"
 	"fmt"
 
@@ -79,20 +78,12 @@ func MoveWorkingSetToBranch(ctx *sql.Context, brName string, force bool, isNewBr
 		if err != nil {
 			return err
 		}
-		if !isNewBranch && actions.CheckoutWouldStompWorkingSetChanges(currentRoots, newBranchRoots) {
-			// In some cases, a working set differs from its head only by the feature version.
-			// If this is the case, moving the working set is safe.
-			modifiedCurrentRoots, err := clearFeatureVersion(ctx, currentRoots)
+		if !isNewBranch {
+			wouldStomp, err := actions.CheckoutWouldStompWorkingSetChanges(ctx, currentRoots, newBranchRoots)
 			if err != nil {
 				return err
 			}
-
-			modifiedNewRoots, err := clearFeatureVersion(ctx, newBranchRoots)
-			if err != nil {
-				return err
-			}
-
-			if !isNewBranch && actions.CheckoutWouldStompWorkingSetChanges(modifiedCurrentRoots, modifiedNewRoots) {
+			if wouldStomp {
 				return actions.ErrWorkingSetsOnBothBranches
 			}
 		}
@@ -154,32 +145,6 @@ func MoveWorkingSetToBranch(ctx *sql.Context, brName string, force bool, isNewBr
 	return nil
 }
 
-// clearFeatureVersion creates a new version of the provided roots where all three roots have the same
-// feature version. By hashing these new roots, we can easily determine whether the roots differ only by
-// their feature version.
-func clearFeatureVersion(ctx context.Context, roots doltdb.Roots) (doltdb.Roots, error) {
-	currentBranchFeatureVersion, _, err := roots.Head.GetFeatureVersion(ctx)
-	if err != nil {
-		return doltdb.Roots{}, err
-	}
-
-	modifiedWorking, err := roots.Working.SetFeatureVersion(currentBranchFeatureVersion)
-	if err != nil {
-		return doltdb.Roots{}, err
-	}
-
-	modifiedStaged, err := roots.Staged.SetFeatureVersion(currentBranchFeatureVersion)
-	if err != nil {
-		return doltdb.Roots{}, err
-	}
-
-	return doltdb.Roots{
-		Head:    roots.Head,
-		Working: modifiedWorking,
-		Staged:  modifiedStaged,
-	}, nil
-}
-
 // transferWorkingChanges computes new roots for `branchRef` by applying the changes from the staged and working sets
 // of `initialRoots` onto the branch head specified by `branchHead`. This is a DESTRUCTIVE ACTION used during command
 // line checkout, to move the working set changes onto a new branch.