Merge pull request #4971 from dolthub/pavel/reset-branch-after-checkout

Reset source branch after checkout

go/cmd/dolt/commands/checkout.go

@@ -120,7 +120,12 @@ func (cmd CheckoutCmd) Exec(ctx context.Context, commandStr string, args []strin
 		if err != nil {
 			return HandleVErrAndExitCode(errhand.BuildDError(err.Error()).Build(), usagePrt)
 		}
-		verr := actions.ResetHard(ctx, dEnv, "HEAD", roots)
+		headRef := dEnv.RepoStateReader().CWBHeadRef()
+		ws, err := dEnv.WorkingSet(ctx)
+		if err != nil {
+			HandleVErrAndExitCode(errhand.BuildDError(err.Error()).Build(), usagePrt)
+		}
+		verr := actions.ResetHard(ctx, dEnv, "HEAD", roots, headRef, ws)
 		return handleResetError(verr, usagePrt)
 	}
 

go/cmd/dolt/commands/reset.go

@@ -106,7 +106,13 @@ func (cmd ResetCmd) Exec(ctx context.Context, commandStr string, args []string,
 			arg = apr.Arg(0)
 		}
 
-		err = actions.ResetHard(ctx, dEnv, arg, roots)
+		headRef := dEnv.RepoStateReader().CWBHeadRef()
+		ws, err := dEnv.WorkingSet(ctx)
+		if err != nil {
+			return HandleVErrAndExitCode(errhand.VerboseErrorFromError(err), usage)
+		}
+
+		err = actions.ResetHard(ctx, dEnv, arg, roots, headRef, ws)
 	} else {
 		// Check whether the input argument is a ref.
 		if apr.NArg() == 1 {

go/libraries/doltcore/env/actions/branch.go

@@ -306,6 +306,7 @@ func checkoutBranchNoDocs(ctx context.Context, roots doltdb.Roots, branchRoot *d
 
 func CheckoutBranch(ctx context.Context, dEnv *env.DoltEnv, brName string, force bool) error {
 	branchRef := ref.NewBranchRef(brName)
+	branchHeadRef := dEnv.RepoStateReader().CWBHeadRef()
 
 	db := dEnv.DoltDB
 	hasRef, err := db.HasRef(ctx, branchRef)
@@ -325,8 +326,14 @@ func CheckoutBranch(ctx context.Context, dEnv *env.DoltEnv, brName string, force
 		return err
 	}
 
+	currentWs, err := dEnv.WorkingSet(ctx)
+	if err != nil {
+		// working set does not exist, skip check
+		return nil
+	}
+
 	if !force {
-		err = checkWorkingSetCompatibility(ctx, dEnv, branchRef)
+		err = checkWorkingSetCompatibility(ctx, dEnv, branchRef, currentWs)
 		if err != nil {
 			return err
 		}
@@ -342,15 +349,20 @@ func CheckoutBranch(ctx context.Context, dEnv *env.DoltEnv, brName string, force
 		return err
 	}
 
+	err = checkoutBranchNoDocs(ctx, roots, branchRoot, dEnv.RepoStateWriter(), branchRef, force)
+	if err != nil {
+		return err
+	}
+
 	if shouldResetWorkingSet {
-		// reset the working set to the branch head, leaving the branch unchanged
-		err = ResetHard(ctx, dEnv, "", roots)
+		// reset the source branch's working set to the branch head, leaving the source branch unchanged
+		err = ResetHard(ctx, dEnv, "", roots, branchHeadRef, currentWs)
 		if err != nil {
 			return err
 		}
 	}
 
-	return checkoutBranchNoDocs(ctx, roots, branchRoot, dEnv.RepoStateWriter(), branchRef, force)
+	return nil
 }
 
 // BranchRoot returns the root value at the branch with the name given
@@ -483,13 +495,7 @@ func overwriteRoot(ctx context.Context, head *doltdb.RootValue, tblHashes map[st
 // 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 checkWorkingSetCompatibility(ctx context.Context, dEnv *env.DoltEnv, branchRef ref.BranchRef) error {
-	currentWs, err := dEnv.WorkingSet(ctx)
-	if err != nil {
-		// working set does not exist, skip check
-		return nil
-	}
-
+func checkWorkingSetCompatibility(ctx context.Context, dEnv *env.DoltEnv, branchRef ref.BranchRef, currentWs *doltdb.WorkingSet) error {
 	db := dEnv.DoltDB
 	destWsRef, err := ref.WorkingSetRefForHead(branchRef)
 	if err != nil {

go/libraries/doltcore/env/actions/reset.go

@@ -18,10 +18,10 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
-
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
+	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
+	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
 	"github.com/dolthub/dolt/go/libraries/utils/argparser"
 )
 
@@ -136,7 +136,17 @@ func ResetHardTables(ctx context.Context, dbData env.DbData, cSpecStr string, ro
 	return resetHardTables(ctx, dbData, cSpecStr, roots)
 }
 
-func ResetHard(ctx context.Context, dEnv *env.DoltEnv, cSpecStr string, roots doltdb.Roots) error {
+// ResetHard resets the working, staged, and head to the ones in the provided roots and head ref.
+// The reset can be performed on a non-current branch and working set.
+// Returns an error if the reset fails.
+func ResetHard(
+	ctx context.Context,
+	dEnv *env.DoltEnv,
+	cSpecStr string,
+	roots doltdb.Roots,
+	headRef ref.DoltRef,
+	ws *doltdb.WorkingSet,
+) error {
 	dbData := dEnv.DbData()
 
 	newHead, roots, err := resetHardTables(ctx, dbData, cSpecStr, roots)
@@ -144,18 +154,13 @@ func ResetHard(ctx context.Context, dEnv *env.DoltEnv, cSpecStr string, roots do
 		return err
 	}
 
-	ws, err := dEnv.WorkingSet(ctx)
-	if err != nil {
-		return err
-	}
-
 	err = dEnv.UpdateWorkingSet(ctx, ws.WithWorkingRoot(roots.Working).WithStagedRoot(roots.Staged).ClearMerge())
 	if err != nil {
 		return err
 	}
 
 	if newHead != nil {
-		err = dEnv.DoltDB.SetHeadToCommit(ctx, dEnv.RepoStateReader().CWBHeadRef(), newHead)
+		err = dEnv.DoltDB.SetHeadToCommit(ctx, headRef, newHead)
 		if err != nil {
 			return err
 		}