Merge pull request #1976 from dolthub/zachmu/cleanup

Cleanup from transaction migration

go/cmd/dolt/commands/add.go

@@ -75,26 +75,41 @@ func (cmd AddCmd) Exec(ctx context.Context, commandStr string, args []string, dE
 
 	roots, err := dEnv.Roots(ctx)
 	if err != nil {
-		cli.PrintErrln(err.Error())
-		return 1
+		return handleStageError(err)
 	}
 
 	if apr.NArg() == 0 && !allFlag {
 		cli.Println("Nothing specified, nothing added.\n Maybe you wanted to say 'dolt add .'?")
 	} else if allFlag || apr.NArg() == 1 && apr.Arg(0) == "." {
-		err = actions.StageAllTables(ctx, roots, dEnv.DbData())
+		roots, err = actions.StageAllTables(ctx, roots, dEnv.Docs)
+		if err != nil {
+			return handleStageError(err)
+		}
 	} else {
-		err = actions.StageTables(ctx, roots, dEnv.DbData(), apr.Args())
+		tables, docs, err := actions.GetTablesOrDocs(dEnv.DocsReadWriter(), apr.Args())
+		if err != nil {
+			return handleStageError(err)
+		}
+
+		roots, err = actions.StageTables(ctx, roots, docs, tables)
+		if err != nil {
+			return handleStageError(err)
+		}
 	}
 
+	err = dEnv.UpdateRoots(ctx, roots)
 	if err != nil {
-		cli.PrintErrln(toAddVErr(err).Verbose())
-		return 1
+		return handleStageError(err)
 	}
 
 	return 0
 }
 
+func handleStageError(err error) int {
+	cli.PrintErrln(toAddVErr(err).Verbose())
+	return 1
+}
+
 func toAddVErr(err error) errhand.VerboseError {
 	switch {
 	case doltdb.IsRootValUnreachable(err):

go/cmd/dolt/commands/commit.go

@@ -82,11 +82,10 @@ func (cmd CommitCmd) Exec(ctx context.Context, commandStr string, args []string,
 	}
 
 	if allFlag {
-		err = actions.StageAllTables(ctx, roots, dEnv.DbData())
-	}
-
-	if err != nil {
-		return handleCommitErr(ctx, dEnv, err, help)
+		roots, err = actions.StageAllTables(ctx, roots, dEnv.Docs)
+		if err != nil {
+			return handleCommitErr(ctx, dEnv, err, help)
+		}
 	}
 
 	var name, email string
@@ -116,14 +115,17 @@ func (cmd CommitCmd) Exec(ctx context.Context, commandStr string, args []string,
 		}
 	}
 
-	// TODO: refactor above stage funcs to modify roots in memory instead of writing to disk
-	dbData := dEnv.DbData()
-	roots, err = dEnv.Roots(context.Background())
+	ws, err := dEnv.WorkingSet(ctx)
 	if err != nil {
-		return HandleVErrAndExitCode(errhand.BuildDError("Couldn't get working root").AddCause(err).Build(), usage)
+		return HandleVErrAndExitCode(errhand.BuildDError("Couldn't get working set").AddCause(err).Build(), usage)
 	}
 
-	_, err = actions.CommitStaged(ctx, roots, dbData, actions.CommitStagedProps{
+	var mergeParentCommits []*doltdb.Commit
+	if ws.MergeActive() {
+		mergeParentCommits = []*doltdb.Commit{ws.MergeState().Commit()}
+	}
+
+	_, err = actions.CommitStaged(ctx, roots, ws.MergeActive(), mergeParentCommits, dEnv.DbData(), actions.CommitStagedProps{
 		Message:    msg,
 		Date:       t,
 		AllowEmpty: apr.Contains(cli.AllowEmptyFlag),
@@ -132,12 +134,17 @@ func (cmd CommitCmd) Exec(ctx context.Context, commandStr string, args []string,
 		Email:      email,
 	})
 
-	if err == nil {
-		// if the commit was successful, print it out using the log command
-		return LogCmd{}.Exec(ctx, "log", []string{"-n=1"}, dEnv)
+	if err != nil {
+		return handleCommitErr(ctx, dEnv, err, usage)
 	}
 
-	return handleCommitErr(ctx, dEnv, err, usage)
+	err = dEnv.ClearMerge(ctx)
+	if err != nil {
+		return HandleVErrAndExitCode(errhand.BuildDError("Couldn't update working set").AddCause(err).Build(), usage)
+	}
+
+	// if the commit was successful, print it out using the log command
+	return LogCmd{}.Exec(ctx, "log", []string{"-n=1"}, dEnv)
 }
 
 func handleCommitErr(ctx context.Context, dEnv *env.DoltEnv, err error, usage cli.UsagePrinter) int {

go/cmd/dolt/commands/merge.go

@@ -213,7 +213,7 @@ func mergeCommitSpec(ctx context.Context, apr *argparser.ArgParseResults, dEnv *
 		return errhand.VerboseErrorFromError(err)
 	}
 
-	tblNames, workingDiffs, err := env.MergeWouldStompChanges(ctx, roots.Working, cm2, dEnv.DbData())
+	tblNames, workingDiffs, err := merge.MergeWouldStompChanges(ctx, roots, cm2)
 
 	if err != nil {
 		return errhand.BuildDError("error: failed to determine mergability.").AddCause(err).Build()
@@ -295,7 +295,17 @@ func execNoFFMerge(ctx context.Context, apr *argparser.ArgParseResults, dEnv *en
 		return errhand.VerboseErrorFromError(err)
 	}
 
-	_, err = actions.CommitStaged(ctx, roots, dEnv.DbData(), actions.CommitStagedProps{
+	ws, err := dEnv.WorkingSet(ctx)
+	if err != nil {
+		return errhand.VerboseErrorFromError(err)
+	}
+
+	var mergeParentCommits []*doltdb.Commit
+	if ws.MergeActive() {
+		mergeParentCommits = []*doltdb.Commit{ws.MergeState().Commit()}
+	}
+
+	_, err = actions.CommitStaged(ctx, roots, ws.MergeActive(), mergeParentCommits, dEnv.DbData(), actions.CommitStagedProps{
 		Message:    msg,
 		Date:       t,
 		AllowEmpty: apr.Contains(cli.AllowEmptyFlag),
@@ -308,6 +318,11 @@ func execNoFFMerge(ctx context.Context, apr *argparser.ArgParseResults, dEnv *en
 		return errhand.BuildDError("error: committing").AddCause(err).Build()
 	}
 
+	err = dEnv.ClearMerge(ctx)
+	if err != nil {
+		return errhand.VerboseErrorFromError(err)
+	}
+
 	return nil
 }
 

go/cmd/dolt/commands/reset.go

@@ -117,8 +117,12 @@ func (cmd ResetCmd) Exec(ctx context.Context, commandStr string, args []string,
 
 		tables := apr.Args()
 
-		roots.Staged, err = actions.ResetSoft(ctx, dEnv.DbData(), tables, roots)
+		roots, err = actions.ResetSoft(ctx, dEnv.DbData(), tables, roots)
+		if err != nil {
+			return handleResetError(err, usage)
+		}
 
+		err = dEnv.UpdateRoots(ctx, roots)
 		if err != nil {
 			return handleResetError(err, usage)
 		}

go/libraries/doltcore/doltdb/workingset.go

@@ -307,16 +307,10 @@ func workingSetMetaFromNomsSt(st types.Struct) (WorkingSetMeta, error) {
 	}, nil
 }
 
-// RootValue returns the root value stored by this workingset
-// TODO: replace references with calls to WorkingRoot
-func (ws *WorkingSet) RootValue() *RootValue {
-	return ws.workingRoot
-}
-
 // HashOf returns the hash of the workingset struct, which is not the same as the hash of the root value stored in the
 // working set. This value is used for optimistic locking when updating a working set for a head ref.
 func (ws *WorkingSet) HashOf() (hash.Hash, error) {
-	if ws.st == nil {
+	if ws == nil || ws.st == nil {
 		return hash.Hash{}, nil
 	}
 	return ws.st.Hash(ws.format)

go/libraries/doltcore/doltdocs/docs_table.go

@@ -17,7 +17,6 @@ package doltdocs
 import (
 	"context"
 	"errors"
-	"fmt"
 	"strconv"
 
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
@@ -225,15 +224,6 @@ func getDocPKFromRow(r row.Row) (string, error) {
 	}
 }
 
-// getFileFromDoc returns the file obj associated with the doc
-func getFileFromDoc(docName string) (string, error) {
-	if doc, ok := IsSupportedDoc(docName); ok {
-		return doc.File, nil
-	}
-
-	return "", fmt.Errorf("Doc name not provided %s", docName)
-}
-
 // GetAllDocs takes a root value and returns all the docs available in the root.
 func GetAllDocs(ctx context.Context, root *doltdb.RootValue) (Docs, bool, error) {
 	if root == nil {

go/libraries/doltcore/dtestutils/environment.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
+	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema/encoding"
 	"github.com/dolthub/dolt/go/libraries/doltcore/table"
 	"github.com/dolthub/dolt/go/libraries/doltcore/table/editor"
@@ -94,8 +95,46 @@ func CreateEnvWithSeedData(t *testing.T) *env.DoltEnv {
 	require.NoError(t, err)
 	indexes, err := tbl.GetIndexData(ctx)
 	require.NoError(t, err)
-	err = dEnv.PutTableToWorking(ctx, sch, rows, indexes, TableName, nil)
+	err = putTableToWorking(ctx, dEnv, sch, rows, indexes, TableName, nil)
 	require.NoError(t, err)
 
 	return dEnv
 }
+
+func putTableToWorking(ctx context.Context, dEnv *env.DoltEnv, sch schema.Schema, rows types.Map, indexData types.Map, tableName string, autoVal types.Value) error {
+	root, err := dEnv.WorkingRoot(ctx)
+	if err != nil {
+		return doltdb.ErrNomsIO
+	}
+
+	vrw := dEnv.DoltDB.ValueReadWriter()
+	schVal, err := encoding.MarshalSchemaAsNomsValue(ctx, vrw, sch)
+	if err != nil {
+		return env.ErrMarshallingSchema
+	}
+
+	tbl, err := doltdb.NewTable(ctx, vrw, schVal, rows, indexData, autoVal)
+	if err != nil {
+		return err
+	}
+
+	newRoot, err := root.PutTable(ctx, tableName, tbl)
+	if err != nil {
+		return err
+	}
+
+	rootHash, err := root.HashOf()
+	if err != nil {
+		return err
+	}
+
+	newRootHash, err := newRoot.HashOf()
+	if err != nil {
+		return err
+	}
+	if rootHash == newRootHash {
+		return nil
+	}
+
+	return dEnv.UpdateWorkingRoot(ctx, newRoot)
+}

go/libraries/doltcore/dtestutils/schema.go

@@ -125,7 +125,7 @@ func CreateTestTable(t *testing.T, dEnv *env.DoltEnv, tableName string, sch sche
 	require.NoError(t, err)
 	indexes, err := tbl.GetIndexData(ctx)
 	require.NoError(t, err)
-	err = dEnv.PutTableToWorking(ctx, sch, rows, indexes, tableName, nil)
+	err = putTableToWorking(ctx, dEnv, sch, rows, indexes, tableName, nil)
 	require.NoError(t, err)
 }
 

go/libraries/doltcore/dtestutils/testcommands/command.go

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Package testcommands TODO: kill off this package, replace with the non-test commands directly
 package testcommands
 
 import (
@@ -48,7 +49,10 @@ func (a StageAll) Exec(t *testing.T, dEnv *env.DoltEnv) error {
 	roots, err := dEnv.Roots(context.Background())
 	require.NoError(t, err)
 
-	return actions.StageAllTables(context.Background(), roots, dEnv.DbData())
+	roots, err = actions.StageAllTables(context.Background(), roots, dEnv.Docs)
+	require.NoError(t, err)
+
+	return dEnv.UpdateRoots(context.Background(), roots)
 }
 
 type CommitStaged struct {
@@ -71,7 +75,17 @@ func (c CommitStaged) Exec(t *testing.T, dEnv *env.DoltEnv) error {
 
 	dbData := dEnv.DbData()
 
-	_, err = actions.CommitStaged(context.Background(), roots, dbData, actions.CommitStagedProps{
+	ws, err := dEnv.WorkingSet(context.Background())
+	if err != nil {
+		return errhand.VerboseErrorFromError(err)
+	}
+
+	var mergeParentCommits []*doltdb.Commit
+	if ws.MergeActive() {
+		mergeParentCommits = []*doltdb.Commit{ws.MergeState().Commit()}
+	}
+
+	_, err = actions.CommitStaged(context.Background(), roots, ws.MergeActive(), mergeParentCommits, dbData, actions.CommitStagedProps{
 		Message:    c.Message,
 		Date:       time.Now(),
 		AllowEmpty: false,
@@ -80,7 +94,11 @@ func (c CommitStaged) Exec(t *testing.T, dEnv *env.DoltEnv) error {
 		Email:      email,
 	})
 
-	return err
+	if err != nil {
+		return err
+	}
+
+	return dEnv.ClearMerge(context.Background())
 }
 
 type CommitAll struct {
@@ -95,7 +113,7 @@ func (c CommitAll) Exec(t *testing.T, dEnv *env.DoltEnv) error {
 	roots, err := dEnv.Roots(context.Background())
 	require.NoError(t, err)
 
-	err = actions.StageAllTables(context.Background(), roots, dEnv.DbData())
+	roots, err = actions.StageAllTables(context.Background(), roots, dEnv.Docs)
 	require.NoError(t, err)
 
 	name, email, err := actions.GetNameAndEmail(dEnv.Config)
@@ -104,12 +122,17 @@ func (c CommitAll) Exec(t *testing.T, dEnv *env.DoltEnv) error {
 		return err
 	}
 
-	dbData := dEnv.DbData()
-	// TODO: refactor StageAllTables to just modify roots in memory, not write to disk
-	roots, err = dEnv.Roots(context.Background())
-	require.NoError(t, err)
+	ws, err := dEnv.WorkingSet(context.Background())
+	if err != nil {
+		return errhand.VerboseErrorFromError(err)
+	}
 
-	_, err = actions.CommitStaged(context.Background(), roots, dbData, actions.CommitStagedProps{
+	var mergeParentCommits []*doltdb.Commit
+	if ws.MergeActive() {
+		mergeParentCommits = []*doltdb.Commit{ws.MergeState().Commit()}
+	}
+
+	_, err = actions.CommitStaged(context.Background(), roots, ws.MergeActive(), mergeParentCommits, dEnv.DbData(), actions.CommitStagedProps{
 		Message:    c.Message,
 		Date:       time.Now(),
 		AllowEmpty: false,
@@ -118,7 +141,11 @@ func (c CommitAll) Exec(t *testing.T, dEnv *env.DoltEnv) error {
 		Email:      email,
 	})
 
-	return err
+	if err != nil {
+		return err
+	}
+
+	return dEnv.ClearMerge(context.Background())
 }
 
 type ResetHard struct{}
@@ -237,10 +264,10 @@ func (m Merge) Exec(t *testing.T, dEnv *env.DoltEnv) error {
 	assert.NoError(t, err)
 	assert.NotEqual(t, h1, h2)
 
-	workingRoot, err := dEnv.WorkingRoot(context.Background())
+	roots, err := dEnv.Roots(context.Background())
 	require.NoError(t, err)
 
-	tblNames, _, err := env.MergeWouldStompChanges(context.Background(), workingRoot, cm2, dEnv.DbData())
+	tblNames, _, err := merge.MergeWouldStompChanges(context.Background(), roots, cm2)
 	if err != nil {
 		return err
 	}

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

@@ -43,7 +43,7 @@ func checkoutTables(ctx context.Context, dbData env.DbData, roots doltdb.Roots,
 		return doltdb.Roots{}, err
 	}
 
-	err = env.UpdateWorkingRoot(ctx, dbData.Rsw, roots.Working)
+	err = dbData.Rsw.UpdateWorkingRoot(ctx, roots.Working)
 	if err != nil {
 		return doltdb.Roots{}, err
 	}
@@ -112,7 +112,7 @@ func checkoutDocs(ctx context.Context, dbData env.DbData, roots doltdb.Roots, do
 		}
 	}
 
-	err := env.UpdateWorkingRoot(ctx, dbData.Rsw, roots.Working)
+	err := dbData.Rsw.UpdateWorkingRoot(ctx, roots.Working)
 	if err != nil {
 		return err
 	}

go/libraries/doltcore/env/actions/commit.go

@@ -57,7 +57,7 @@ func GetNameAndEmail(cfg config.ReadableConfig) (string, string, error) {
 }
 
 // CommitStaged adds a new commit to HEAD with the given props. Returns the new commit's hash as a string and an error.
-func CommitStaged(ctx context.Context, roots doltdb.Roots, dbData env.DbData, props CommitStagedProps) (*doltdb.Commit, error) {
+func CommitStaged(ctx context.Context, roots doltdb.Roots, mergeActive bool, mergeParents []*doltdb.Commit, dbData env.DbData, props CommitStagedProps) (*doltdb.Commit, error) {
 	ddb := dbData.Ddb
 	rsr := dbData.Rsr
 	rsw := dbData.Rsw
@@ -81,11 +81,6 @@ func CommitStaged(ctx context.Context, roots doltdb.Roots, dbData env.DbData, pr
 		stagedTblNames = append(stagedTblNames, n)
 	}
 
-	mergeActive, err := rsr.IsMergeActive(ctx)
-	if err != nil {
-		return nil, err
-	}
-
 	if len(staged) == 0 && !mergeActive && !props.AllowEmpty {
 		_, notStagedDocs, err := diff.GetDocDiffs(ctx, roots, drw)
 		if err != nil {
@@ -94,8 +89,7 @@ func CommitStaged(ctx context.Context, roots doltdb.Roots, dbData env.DbData, pr
 		return nil, NothingStaged{notStaged, notStagedDocs}
 	}
 
-	var mergeParentCommits []*doltdb.Commit
-	if mergeActive && !props.Force {
+	if !props.Force {
 		inConflict, err := roots.Working.TablesInConflict(ctx)
 		if err != nil {
 			return nil, err
@@ -110,13 +104,6 @@ func CommitStaged(ctx context.Context, roots doltdb.Roots, dbData env.DbData, pr
 		if len(violatesConstraints) > 0 {
 			return nil, NewTblHasConstraintViolations(violatesConstraints)
 		}
-
-		commit, err := rsr.GetMergeCommit(ctx)
-		if err != nil {
-			return nil, err
-		}
-
-		mergeParentCommits = []*doltdb.Commit{commit}
 	}
 
 	stagedRoot, err := roots.Staged.UpdateSuperSchemasFromOther(ctx, stagedTblNames, roots.Staged)
@@ -132,7 +119,7 @@ func CommitStaged(ctx context.Context, roots doltdb.Roots, dbData env.DbData, pr
 	}
 
 	// TODO: combine into a single update
-	err = env.UpdateStagedRoot(ctx, rsw, stagedRoot)
+	err = rsw.UpdateStagedRoot(ctx, stagedRoot)
 	if err != nil {
 		return nil, err
 	}
@@ -142,7 +129,7 @@ func CommitStaged(ctx context.Context, roots doltdb.Roots, dbData env.DbData, pr
 		return nil, err
 	}
 
-	err = env.UpdateWorkingRoot(ctx, rsw, workingRoot)
+	err = rsw.UpdateWorkingRoot(ctx, workingRoot)
 	if err != nil {
 		return nil, err
 	}
@@ -162,12 +149,7 @@ func CommitStaged(ctx context.Context, roots doltdb.Roots, dbData env.DbData, pr
 	// logrus.Errorf("staged root is %s", stagedRoot.DebugString(ctx, true))
 
 	// DoltDB resolves the current working branch head ref to provide a parent commit.
-	c, err := ddb.CommitWithParentCommits(ctx, h, rsr.CWBHeadRef(), mergeParentCommits, meta)
-	if err != nil {
-		return nil, err
-	}
-
-	err = rsw.ClearMerge(ctx)
+	c, err := ddb.CommitWithParentCommits(ctx, h, rsr.CWBHeadRef(), mergeParents, meta)
 	if err != nil {
 		return nil, err
 	}

go/libraries/doltcore/env/actions/docs.go

@@ -36,16 +36,6 @@ func SaveTrackedDocsFromWorking(ctx context.Context, dEnv *env.DoltEnv) error {
 	return SaveTrackedDocs(ctx, dEnv.DocsReadWriter(), workingRoot, workingRoot, localDocs)
 }
 
-// SaveDocsFromWorking saves docs from the working root to the filesystem, and could overwrite untracked docs.
-func SaveDocsFromWorking(ctx context.Context, dEnv *env.DoltEnv) error {
-	workingRoot, err := dEnv.WorkingRoot(ctx)
-	if err != nil {
-		return err
-	}
-
-	return SaveDocsFromRoot(ctx, workingRoot, dEnv)
-}
-
 // SaveDocsFromRoot saves docs from the root given to the filesystem, and could overwrite untracked docs.
 func SaveDocsFromRoot(ctx context.Context, root *doltdb.RootValue, dEnv *env.DoltEnv) error {
 	localDocs := dEnv.Docs

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

@@ -147,16 +147,17 @@ func ResetSoftTables(ctx context.Context, dbData env.DbData, apr *argparser.ArgP
 	return roots, nil
 }
 
-func ResetSoft(ctx context.Context, dbData env.DbData, tables []string, roots doltdb.Roots) (*doltdb.RootValue, error) {
+// ResetSoft resets the staged value from HEAD for the tables given and returns the updated roots.
+func ResetSoft(ctx context.Context, dbData env.DbData, tables []string, roots doltdb.Roots) (doltdb.Roots, error) {
 	tables, err := getUnionedTables(ctx, tables, roots.Staged, roots.Head)
 
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
 	tables, docs, err := GetTablesOrDocs(dbData.Drw, tables)
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
 	if len(docs) > 0 {
@@ -166,21 +167,20 @@ func ResetSoft(ctx context.Context, dbData env.DbData, tables []string, roots do
 	err = ValidateTables(context.TODO(), tables, roots.Staged, roots.Head)
 
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
-	roots.Staged, err = resetDocs(ctx, dbData, roots, docs)
+	roots, err = resetDocs(ctx, roots, docs)
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
-	roots.Staged, err = resetStaged(ctx, roots, dbData.Rsw, tables)
-
+	roots, err = resetStaged(ctx, roots, tables)
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
-	return roots.Staged, nil
+	return roots, nil
 }
 
 // ResetSoftToRef matches the `git reset --soft <REF>` pattern. It resets both staged and head to the previous ref
@@ -201,8 +201,8 @@ func ResetSoftToRef(ctx context.Context, dbData env.DbData, cSpecStr string) err
 		return err
 	}
 
-	// Changed the stage to old the root. Leave the working as is.
-	err = env.UpdateStagedRoot(ctx, dbData.Rsw, foundRoot)
+	// Changed the staged to the old root. Leave the working as is.
+	err = dbData.Rsw.UpdateStagedRoot(ctx, foundRoot)
 	if err != nil {
 		return err
 	}
@@ -229,38 +229,33 @@ func getUnionedTables(ctx context.Context, tables []string, stagedRoot, headRoot
 }
 
 // resetDocs resets the working and staged docs with docs from head.
-func resetDocs(ctx context.Context, dbData env.DbData, roots doltdb.Roots, docs doltdocs.Docs) (newStgRoot *doltdb.RootValue, err error) {
-	docs, err = doltdocs.GetDocsFromRoot(ctx, roots.Head, doltdocs.GetDocNamesFromDocs(docs)...)
+func resetDocs(ctx context.Context, roots doltdb.Roots, docs doltdocs.Docs) (doltdb.Roots, error) {
+	docs, err := doltdocs.GetDocsFromRoot(ctx, roots.Head, doltdocs.GetDocNamesFromDocs(docs)...)
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
 	roots.Working, err = doltdocs.UpdateRootWithDocs(ctx, roots.Working, docs)
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
-	err = env.UpdateWorkingRoot(ctx, dbData.Rsw, roots.Working)
+	roots.Staged, err = doltdocs.UpdateRootWithDocs(ctx, roots.Staged, docs)
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
-	return doltdocs.UpdateRootWithDocs(ctx, roots.Staged, docs)
+	return roots, nil
 }
 
-// TODO: this should just work in memory, not write to disk
-func resetStaged(ctx context.Context, roots doltdb.Roots, rsw env.RepoStateWriter, tbls []string) (*doltdb.RootValue, error) {
+func resetStaged(ctx context.Context, roots doltdb.Roots, tbls []string) (doltdb.Roots, error) {
 	newStaged, err := MoveTablesBetweenRoots(ctx, tbls, roots.Head, roots.Staged)
 	if err != nil {
-		return nil, err
+		return doltdb.Roots{}, err
 	}
 
-	err = rsw.UpdateStagedRoot(ctx, newStaged)
-	if err != nil {
-		return nil, err
-	}
-
-	return newStaged, nil
+	roots.Staged = newStaged
+	return roots, nil
 }
 
 // ValidateIsRef validates whether the input parameter is a valid cString

go/libraries/doltcore/env/actions/staged.go

@@ -16,69 +16,44 @@ package actions
 
 import (
 	"context"
-	"errors"
 
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdocs"
-	"github.com/dolthub/dolt/go/libraries/doltcore/env"
 )
 
-var ErrTablesInConflict = errors.New("table is in conflict")
-
-func StageTables(ctx context.Context, roots doltdb.Roots, dbData env.DbData, tbls []string) error {
-	rsw := dbData.Rsw
-	drw := dbData.Drw
-
-	tables, docs, err := GetTablesOrDocs(drw, tbls)
-	if err != nil {
-		return err
-	}
-
+func StageTables(ctx context.Context, roots doltdb.Roots, docs doltdocs.Docs, tbls []string) (doltdb.Roots, error) {
 	if len(docs) > 0 {
+		var err error
 		roots.Working, err = doltdocs.UpdateRootWithDocs(ctx, roots.Working, docs)
 		if err != nil {
-			return err
+			return doltdb.Roots{}, err
 		}
 	}
 
-	err = stageTables(ctx, roots, rsw, tables)
-	if err != nil {
-		env.ResetWorkingDocsToStagedDocs(ctx, roots, rsw)
-		return err
-	}
-	return nil
+	return stageTables(ctx, roots, tbls)
 }
 
 func StageTablesNoDocs(ctx context.Context, roots doltdb.Roots, tbls []string) (doltdb.Roots, error) {
-	return stageTablesNoEnvUpdate(ctx, roots, tbls)
+	return stageTables(ctx, roots, tbls)
 }
 
-func StageAllTables(ctx context.Context, roots doltdb.Roots, dbData env.DbData) error {
-	rsw := dbData.Rsw
-	drw := dbData.Drw
+func StageAllTables(ctx context.Context, roots doltdb.Roots, docs doltdocs.Docs) (doltdb.Roots, error) {
+	var err error
 
-	docs, err := drw.GetDocsOnDisk()
-	if err != nil {
-		return err
-	}
-
-	roots.Working, err = doltdocs.UpdateRootWithDocs(ctx, roots.Working, docs)
-	if err != nil {
-		return err
+	// To stage all docs for removal, use an empty slice instead of nil
+	if docs != nil {
+		roots.Working, err = doltdocs.UpdateRootWithDocs(ctx, roots.Working, docs)
+		if err != nil {
+			return doltdb.Roots{}, err
+		}
 	}
 
 	tbls, err := doltdb.UnionTableNames(ctx, roots.Staged, roots.Working)
 	if err != nil {
-		return err
+		return doltdb.Roots{}, err
 	}
 
-	err = stageTables(ctx, roots, rsw, tbls)
-	if err != nil {
-		env.ResetWorkingDocsToStagedDocs(ctx, roots, rsw)
-		return err
-	}
-
-	return nil
+	return stageTables(ctx, roots, tbls)
 }
 
 func StageAllTablesNoDocs(ctx context.Context, roots doltdb.Roots) (doltdb.Roots, error) {
@@ -87,10 +62,10 @@ func StageAllTablesNoDocs(ctx context.Context, roots doltdb.Roots) (doltdb.Roots
 		return doltdb.Roots{}, err
 	}
 
-	return stageTablesNoEnvUpdate(ctx, roots, tbls)
+	return stageTables(ctx, roots, tbls)
 }
 
-func stageTablesNoEnvUpdate(
+func stageTables(
 	ctx context.Context,
 	roots doltdb.Roots,
 	tbls []string,
@@ -114,27 +89,6 @@ func stageTablesNoEnvUpdate(
 	return roots, nil
 }
 
-func stageTables(
-	ctx context.Context,
-	roots doltdb.Roots,
-	rsw env.RepoStateWriter,
-	tbls []string,
-) error {
-	var err error
-	roots, err = stageTablesNoEnvUpdate(ctx, roots, tbls)
-	if err != nil {
-		return err
-	}
-
-	// TODO: combine to single operation
-	err = env.UpdateWorkingRoot(ctx, rsw, roots.Working)
-	if err != nil {
-		return err
-	}
-
-	return env.UpdateStagedRoot(ctx, rsw, roots.Staged)
-}
-
 // clearEmptyConflicts clears any 0-row conflicts from the tables named, and returns a new root.
 func clearEmptyConflicts(ctx context.Context, tbls []string, working *doltdb.RootValue) (*doltdb.RootValue, error) {
 	for _, tblName := range tbls {

go/libraries/doltcore/env/dolt_docs.go

@@ -1,57 +0,0 @@
-// Copyright 2019 Dolthub, Inc.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package env
-
-import (
-	"context"
-
-	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
-)
-
-// ResetWorkingDocsToStagedDocs resets the `dolt_docs` table on the working root to match the staged root.
-// If the `dolt_docs` table does not exist on the staged root, it will be removed from the working root.
-func ResetWorkingDocsToStagedDocs(
-	ctx context.Context,
-	roots doltdb.Roots,
-	rsw RepoStateWriter,
-) error {
-	stgDocTbl, stgDocsFound, err := roots.Staged.GetTable(ctx, doltdb.DocTableName)
-	if err != nil {
-		return err
-	}
-
-	_, wrkDocsFound, err := roots.Working.GetTable(ctx, doltdb.DocTableName)
-	if err != nil {
-		return err
-	}
-
-	if wrkDocsFound && !stgDocsFound {
-		newWrkRoot, err := roots.Working.RemoveTables(ctx, doltdb.DocTableName)
-		if err != nil {
-			return err
-		}
-		return UpdateWorkingRoot(ctx, rsw, newWrkRoot)
-	}
-
-	if stgDocsFound {
-		newWrkRoot, err := roots.Working.PutTable(ctx, doltdb.DocTableName, stgDocTbl)
-		if err != nil {
-			return err
-		}
-		return UpdateWorkingRoot(ctx, rsw, newWrkRoot)
-	}
-
-	return nil
-}

go/libraries/doltcore/env/environment.go

@@ -35,8 +35,6 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdocs"
 	"github.com/dolthub/dolt/go/libraries/doltcore/grpcendpoint"
 	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
-	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
-	"github.com/dolthub/dolt/go/libraries/doltcore/schema/encoding"
 	"github.com/dolthub/dolt/go/libraries/utils/filesys"
 	"github.com/dolthub/dolt/go/store/hash"
 	"github.com/dolthub/dolt/go/store/types"
@@ -422,6 +420,7 @@ type RootsProvider interface {
 	GetRoots(ctx context.Context) (doltdb.Roots, error)
 }
 
+// Roots returns the roots for this environment
 func (dEnv *DoltEnv) Roots(ctx context.Context) (doltdb.Roots, error) {
 	ws, err := dEnv.WorkingSet(ctx)
 	if err != nil {
@@ -440,6 +439,23 @@ func (dEnv *DoltEnv) Roots(ctx context.Context) (doltdb.Roots, error) {
 	}, nil
 }
 
+// UpdateRoots updates the working and staged roots for this environment
+func (dEnv *DoltEnv) UpdateRoots(ctx context.Context, roots doltdb.Roots) error {
+	ws, err := dEnv.WorkingSet(ctx)
+	if err == doltdb.ErrWorkingSetNotFound {
+		// first time updating roots
+		wsRef, err := ref.WorkingSetRefForHead(dEnv.RepoState.CWBHeadRef())
+		if err != nil {
+			return err
+		}
+		ws = doltdb.EmptyWorkingSet(wsRef)
+	} else if err != nil {
+		return err
+	}
+
+	return dEnv.UpdateWorkingSet(ctx, ws.WithWorkingRoot(roots.Working).WithStagedRoot(roots.Staged))
+}
+
 // WorkingRoot returns the working root for the current working branch
 func (dEnv *DoltEnv) WorkingRoot(ctx context.Context) (*doltdb.RootValue, error) {
 	workingSet, err := dEnv.WorkingSet(ctx)
@@ -447,7 +463,7 @@ func (dEnv *DoltEnv) WorkingRoot(ctx context.Context) (*doltdb.RootValue, error)
 		return nil, err
 	}
 
-	return workingSet.RootValue(), nil
+	return workingSet.WorkingRoot(), nil
 }
 
 func (dEnv *DoltEnv) WorkingSet(ctx context.Context) (*doltdb.WorkingSet, error) {
@@ -510,14 +526,6 @@ type repoStateReader struct {
 	dEnv *DoltEnv
 }
 
-func (r *repoStateReader) StagedRoot(ctx context.Context) (*doltdb.RootValue, error) {
-	return r.dEnv.StagedRoot(ctx)
-}
-
-func (r *repoStateReader) WorkingRoot(ctx context.Context) (*doltdb.RootValue, error) {
-	return r.dEnv.WorkingRoot(ctx)
-}
-
 func (r *repoStateReader) CWBHeadRef() ref.DoltRef {
 	return r.dEnv.RepoState.CWBHeadRef()
 }
@@ -526,34 +534,6 @@ func (r *repoStateReader) CWBHeadSpec() *doltdb.CommitSpec {
 	return r.dEnv.RepoState.CWBHeadSpec()
 }
 
-func (r *repoStateReader) CWBHeadHash(ctx context.Context) (hash.Hash, error) {
-	ref := r.CWBHeadRef()
-	cm, err := r.dEnv.DoltDB.ResolveCommitRef(ctx, ref)
-
-	if err != nil {
-		return hash.Hash{}, err
-	}
-
-	return cm.HashOf()
-}
-
-func (r *repoStateReader) StagedHash() hash.Hash {
-	return hash.Parse(r.dEnv.RepoState.staged)
-}
-
-func (r *repoStateReader) IsMergeActive(ctx context.Context) (bool, error) {
-	return r.dEnv.IsMergeActive(ctx)
-}
-
-func (r *repoStateReader) GetMergeCommit(ctx context.Context) (*doltdb.Commit, error) {
-	ws, err := r.dEnv.WorkingSet(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	return ws.MergeState().Commit(), nil
-}
-
 func (dEnv *DoltEnv) RepoStateReader() RepoStateReader {
 	return &repoStateReader{dEnv}
 }
@@ -573,19 +553,6 @@ func (r *repoStateWriter) SetCWBHeadRef(ctx context.Context, marshalableRef ref.
 	return nil
 }
 
-// TODO: kill merge methods
-func (r *repoStateWriter) AbortMerge(ctx context.Context) error {
-	return r.DoltEnv.AbortMerge(ctx)
-}
-
-func (r *repoStateWriter) ClearMerge(ctx context.Context) error {
-	return r.DoltEnv.ClearMerge(ctx)
-}
-
-func (r *repoStateWriter) StartMerge(ctx context.Context, commit *doltdb.Commit) error {
-	return r.DoltEnv.StartMerge(ctx, commit)
-}
-
 func (dEnv *DoltEnv) RepoStateWriter() RepoStateWriter {
 	return &repoStateWriter{dEnv}
 }
@@ -730,45 +697,6 @@ func (dEnv *DoltEnv) StartMerge(ctx context.Context, commit *doltdb.Commit) erro
 	return dEnv.DoltDB.UpdateWorkingSet(ctx, ws.Ref(), ws.StartMerge(commit), h, dEnv.workingSetMeta())
 }
 
-// todo: move this out of env to actions
-func (dEnv *DoltEnv) PutTableToWorking(ctx context.Context, sch schema.Schema, rows types.Map, indexData types.Map, tableName string, autoVal types.Value) error {
-	root, err := dEnv.WorkingRoot(ctx)
-	if err != nil {
-		return doltdb.ErrNomsIO
-	}
-
-	vrw := dEnv.DoltDB.ValueReadWriter()
-	schVal, err := encoding.MarshalSchemaAsNomsValue(ctx, vrw, sch)
-	if err != nil {
-		return ErrMarshallingSchema
-	}
-
-	tbl, err := doltdb.NewTable(ctx, vrw, schVal, rows, indexData, autoVal)
-	if err != nil {
-		return err
-	}
-
-	newRoot, err := root.PutTable(ctx, tableName, tbl)
-	if err != nil {
-		return err
-	}
-
-	rootHash, err := root.HashOf()
-	if err != nil {
-		return err
-	}
-
-	newRootHash, err := newRoot.HashOf()
-	if err != nil {
-		return err
-	}
-	if rootHash == newRootHash {
-		return nil
-	}
-
-	return dEnv.UpdateWorkingRoot(ctx, newRoot)
-}
-
 func (dEnv *DoltEnv) IsMergeActive(ctx context.Context) (bool, error) {
 	ws, err := dEnv.WorkingSet(ctx)
 	if err != nil {
@@ -1016,3 +944,60 @@ func (dEnv *DoltEnv) GetUserHomeDir() (string, error) {
 func (dEnv *DoltEnv) TempTableFilesDir() string {
 	return mustAbs(dEnv, dEnv.GetDoltDir(), tempTablesDir)
 }
+
+// GetGCKeepers returns the hashes of all the objects in the environment provided that should be perserved during GC.
+// TODO: this should be unnecessary since we now store the working set in a noms dataset, remove it
+func GetGCKeepers(ctx context.Context, env *DoltEnv) ([]hash.Hash, error) {
+	workingRoot, err := env.WorkingRoot(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	workingHash, err := workingRoot.HashOf()
+	if err != nil {
+		return nil, err
+	}
+
+	stagedRoot, err := env.StagedRoot(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	stagedHash, err := stagedRoot.HashOf()
+	if err != nil {
+		return nil, err
+	}
+
+	keepers := []hash.Hash{
+		workingHash,
+		stagedHash,
+	}
+
+	mergeActive, err := env.IsMergeActive(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if mergeActive {
+		ws, err := env.WorkingSet(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		cm := ws.MergeState().Commit()
+		ch, err := cm.HashOf()
+		if err != nil {
+			return nil, err
+		}
+
+		pmw := ws.MergeState().PreMergeWorkingRoot()
+		pmwh, err := pmw.HashOf()
+		if err != nil {
+			return nil, err
+		}
+
+		keepers = append(keepers, ch, pmwh)
+	}
+
+	return keepers, nil
+}

go/libraries/doltcore/env/repo_state.go

@@ -28,10 +28,6 @@ import (
 type RepoStateReader interface {
 	CWBHeadRef() ref.DoltRef
 	CWBHeadSpec() *doltdb.CommitSpec
-	// TODO: get rid of this
-	IsMergeActive(ctx context.Context) (bool, error)
-	// TODO: get rid of this
-	GetMergeCommit(ctx context.Context) (*doltdb.Commit, error)
 }
 
 type RepoStateWriter interface {
@@ -40,12 +36,6 @@ type RepoStateWriter interface {
 	// TODO: get rid of this
 	UpdateWorkingRoot(ctx context.Context, newRoot *doltdb.RootValue) error
 	SetCWBHeadRef(context.Context, ref.MarshalableRef) error
-	// TODO: get rid of this
-	AbortMerge(ctx context.Context) error
-	// TODO: get rid of this
-	ClearMerge(ctx context.Context) error
-	// TODO: get rid of this
-	StartMerge(ctx context.Context, commit *doltdb.Commit) error
 }
 
 type DocsReadWriter interface {
@@ -210,159 +200,3 @@ func (rs *RepoState) CWBHeadSpec() *doltdb.CommitSpec {
 func (rs *RepoState) AddRemote(r Remote) {
 	rs.Remotes[r.Name] = r
 }
-
-// Updates the working root.
-func UpdateWorkingRoot(ctx context.Context, rsw RepoStateWriter, newRoot *doltdb.RootValue) error {
-	// logrus.Infof("Updating working root with value %s", newRoot.DebugString(ctx, true))
-
-	err := rsw.UpdateWorkingRoot(ctx, newRoot)
-	if err != nil {
-		return ErrStateUpdate
-	}
-
-	return nil
-}
-
-// Returns the head root.
-func HeadRoot(ctx context.Context, ddb *doltdb.DoltDB, rsr RepoStateReader) (*doltdb.RootValue, error) {
-	commit, err := ddb.ResolveCommitRef(ctx, rsr.CWBHeadRef())
-
-	if err != nil {
-		return nil, err
-	}
-
-	return commit.GetRootValue()
-}
-
-// Updates the staged root.
-// TODO: remove this
-func UpdateStagedRoot(ctx context.Context, rsw RepoStateWriter, newRoot *doltdb.RootValue) error {
-	err := rsw.UpdateStagedRoot(ctx, newRoot)
-	if err != nil {
-		return ErrStateUpdate
-	}
-
-	return nil
-}
-
-// TODO: this needs to be a function in the merge package, not repo state
-func MergeWouldStompChanges(ctx context.Context, workingRoot *doltdb.RootValue, mergeCommit *doltdb.Commit, dbData DbData) ([]string, map[string]hash.Hash, error) {
-	headRoot, err := HeadRoot(ctx, dbData.Ddb, dbData.Rsr)
-	if err != nil {
-		return nil, nil, err
-	}
-
-	mergeRoot, err := mergeCommit.GetRootValue()
-	if err != nil {
-		return nil, nil, err
-	}
-
-	headTableHashes, err := headRoot.MapTableHashes(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-
-	workingTableHashes, err := workingRoot.MapTableHashes(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-
-	mergeTableHashes, err := mergeRoot.MapTableHashes(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-
-	headWorkingDiffs := diffTableHashes(headTableHashes, workingTableHashes)
-	mergedHeadDiffs := diffTableHashes(headTableHashes, mergeTableHashes)
-
-	stompedTables := make([]string, 0, len(headWorkingDiffs))
-	for tName, _ := range headWorkingDiffs {
-		if _, ok := mergedHeadDiffs[tName]; ok {
-			// even if the working changes match the merge changes, don't allow (matches git behavior).
-			stompedTables = append(stompedTables, tName)
-		}
-	}
-
-	return stompedTables, headWorkingDiffs, nil
-}
-
-// GetGCKeepers queries |rsr| to find a list of values that need to be temporarily saved during GC.
-// TODO: move this out of repo_state.go
-func GetGCKeepers(ctx context.Context, env *DoltEnv) ([]hash.Hash, error) {
-	workingRoot, err := env.WorkingRoot(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	workingHash, err := workingRoot.HashOf()
-	if err != nil {
-		return nil, err
-	}
-
-	stagedRoot, err := env.StagedRoot(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	stagedHash, err := stagedRoot.HashOf()
-	if err != nil {
-		return nil, err
-	}
-
-	keepers := []hash.Hash{
-		workingHash,
-		stagedHash,
-	}
-
-	mergeActive, err := env.IsMergeActive(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	if mergeActive {
-		ws, err := env.WorkingSet(ctx)
-		if err != nil {
-			return nil, err
-		}
-
-		cm := ws.MergeState().Commit()
-		ch, err := cm.HashOf()
-		if err != nil {
-			return nil, err
-		}
-
-		pmw := ws.MergeState().PreMergeWorkingRoot()
-		pmwh, err := pmw.HashOf()
-		if err != nil {
-			return nil, err
-		}
-
-		keepers = append(keepers, ch, pmwh)
-	}
-
-	return keepers, nil
-}
-
-func diffTableHashes(headTableHashes, otherTableHashes map[string]hash.Hash) map[string]hash.Hash {
-	diffs := make(map[string]hash.Hash)
-	for tName, hh := range headTableHashes {
-		if h, ok := otherTableHashes[tName]; ok {
-			if h != hh {
-				// modification
-				diffs[tName] = h
-			}
-		} else {
-			// deletion
-			diffs[tName] = hash.Hash{}
-		}
-	}
-
-	for tName, h := range otherTableHashes {
-		if _, ok := headTableHashes[tName]; !ok {
-			// addition
-			diffs[tName] = h
-		}
-	}
-
-	return diffs
-}

go/libraries/doltcore/merge/merge.go

@@ -1041,3 +1041,62 @@ func GetDocsInConflict(ctx context.Context, workingRoot *doltdb.RootValue, drw e
 
 	return diff.NewDocDiffs(ctx, workingRoot, nil, docs)
 }
+
+func MergeWouldStompChanges(ctx context.Context, roots doltdb.Roots, mergeCommit *doltdb.Commit) ([]string, map[string]hash.Hash, error) {
+	mergeRoot, err := mergeCommit.GetRootValue()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	headTableHashes, err := roots.Head.MapTableHashes(ctx)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	workingTableHashes, err := roots.Working.MapTableHashes(ctx)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	mergeTableHashes, err := mergeRoot.MapTableHashes(ctx)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	headWorkingDiffs := diffTableHashes(headTableHashes, workingTableHashes)
+	mergedHeadDiffs := diffTableHashes(headTableHashes, mergeTableHashes)
+
+	stompedTables := make([]string, 0, len(headWorkingDiffs))
+	for tName, _ := range headWorkingDiffs {
+		if _, ok := mergedHeadDiffs[tName]; ok {
+			// even if the working changes match the merge changes, don't allow (matches git behavior).
+			stompedTables = append(stompedTables, tName)
+		}
+	}
+
+	return stompedTables, headWorkingDiffs, nil
+}
+
+func diffTableHashes(headTableHashes, otherTableHashes map[string]hash.Hash) map[string]hash.Hash {
+	diffs := make(map[string]hash.Hash)
+	for tName, hh := range headTableHashes {
+		if h, ok := otherTableHashes[tName]; ok {
+			if h != hh {
+				// modification
+				diffs[tName] = h
+			}
+		} else {
+			// deletion
+			diffs[tName] = hash.Hash{}
+		}
+	}
+
+	for tName, h := range otherTableHashes {
+		if _, ok := headTableHashes[tName]; !ok {
+			// addition
+			diffs[tName] = h
+		}
+	}
+
+	return diffs
+}

go/libraries/doltcore/rebase/rebase.go

@@ -184,12 +184,12 @@ func rebaseRefs(ctx context.Context, dbData env.DbData, replay ReplayCommitFn, n
 	}
 
 	// TODO: this should be a single update to repo state, not two
-	err = env.UpdateStagedRoot(ctx, rsw, r)
+	err = rsw.UpdateStagedRoot(ctx, r)
 	if err != nil {
 		return err
 	}
 
-	return env.UpdateWorkingRoot(ctx, rsw, r)
+	return rsw.UpdateWorkingRoot(ctx, r)
 }
 
 func rebase(ctx context.Context, ddb *doltdb.DoltDB, replay ReplayCommitFn, nerf NeedsRebaseFn, origins ...*doltdb.Commit) ([]*doltdb.Commit, error) {

go/libraries/doltcore/sqle/database.go

@@ -30,7 +30,6 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env/actions/commitwalk"
-	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
 	"github.com/dolthub/dolt/go/libraries/doltcore/row"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema/alterschema"
@@ -81,28 +80,6 @@ func (db Database) StartTransaction(ctx *sql.Context) (sql.Transaction, error) {
 	return dsession.StartTransaction(ctx, db.Name())
 }
 
-func (db Database) setHeadHash(ctx *sql.Context, ref ref.WorkingSetRef) error {
-	// TODO: use the session HEAD ref here instead of the repo state one
-	// headRef, err := ref.ToHeadRef()
-	// if err != nil {
-	// 	return err
-	// }
-
-	headCommit, err := db.ddb.Resolve(ctx, db.rsr.CWBHeadSpec(), db.rsr.CWBHeadRef())
-	if err != nil {
-		return err
-	}
-	headHash, err := headCommit.HashOf()
-	if err != nil {
-		return err
-	}
-	if doltSession, ok := ctx.Session.(*dsess.Session); ok {
-		return doltSession.SetSessionVarDirectly(ctx, dsess.HeadKey(db.name), headHash.String())
-	} else {
-		return ctx.SetSessionVariable(ctx, dsess.HeadKey(db.name), headHash.String())
-	}
-}
-
 func (db Database) CommitTransaction(ctx *sql.Context, tx sql.Transaction) error {
 	dsession := dsess.DSessFromSess(ctx.Session)
 	return dsession.CommitTransaction(ctx, db.name, tx)

go/libraries/doltcore/sqle/dfunctions/dolt_commit.go

@@ -21,6 +21,7 @@ import (
 	"github.com/dolthub/vitess/go/vt/proto/query"
 
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
+	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env/actions"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
 )
@@ -115,8 +116,18 @@ func (d DoltCommitFunc) Eval(ctx *sql.Context, row sql.Row) (interface{}, error)
 	// Unsetting the transaction here ensures that it won't be re-committed when this statement concludes
 	ctx.SetTransaction(nil)
 
+	var mergeParentCommits []*doltdb.Commit
+	ws, err := dSess.WorkingSet(ctx, dbName)
+	if err != nil {
+		return nil, err
+	}
+
+	if ws.MergeActive() {
+		mergeParentCommits = []*doltdb.Commit{ws.MergeState().Commit()}
+	}
+
 	// Now do a Dolt commit
-	commit, err := dSess.CommitToDolt(ctx, roots, dbName, actions.CommitStagedProps{
+	commit, err := dSess.CommitToDolt(ctx, roots, mergeParentCommits, dbName, actions.CommitStagedProps{
 		Message:    msg,
 		Date:       t,
 		AllowEmpty: apr.Contains(cli.AllowEmptyFlag),

go/libraries/doltcore/sqle/dfunctions/dolt_merge.go

@@ -325,9 +325,14 @@ func executeNoFFMerge(
 	// The roots need refreshing after the above
 	roots, _ := dSess.GetRoots(ctx, dbName)
 
+	var mergeParentCommits []*doltdb.Commit
+	if ws.MergeActive() {
+		mergeParentCommits = []*doltdb.Commit{ws.MergeState().Commit()}
+	}
+
 	// TODO: this does several session state updates, and it really needs to just do one
 	//  We also need to commit any pending transaction before we do this.
-	_, err = actions.CommitStaged(ctx, roots, dbData, actions.CommitStagedProps{
+	_, err = actions.CommitStaged(ctx, roots, ws.MergeActive(), mergeParentCommits, dbData, actions.CommitStagedProps{
 		Message:    msg,
 		Date:       t,
 		AllowEmpty: apr.Contains(cli.AllowEmptyFlag),
@@ -339,7 +344,7 @@ func executeNoFFMerge(
 		return nil, err
 	}
 
-	return ws, dSess.SetWorkingSet(ctx, dbName, ws, nil)
+	return ws, dSess.SetWorkingSet(ctx, dbName, ws.ClearMerge(), nil)
 }
 
 // TODO: this copied from commands/merge.go because the latter isn't reusable. Fix that.

go/libraries/doltcore/sqle/dsess/session.go

@@ -389,6 +389,7 @@ func (sess *Session) CommitTransaction(ctx *sql.Context, dbName string, tx sql.T
 func (sess *Session) CommitToDolt(
 	ctx *sql.Context,
 	roots doltdb.Roots,
+	mergeParentCommits []*doltdb.Commit,
 	dbName string,
 	props actions.CommitStagedProps,
 ) (*doltdb.Commit, error) {
@@ -398,6 +399,11 @@ func (sess *Session) CommitToDolt(
 	}
 	dbData := sessionState.dbData
 
+	ws, err := sess.WorkingSet(ctx, dbName)
+	if err != nil {
+		return nil, err
+	}
+
 	// TODO: this does several session state updates, and it really needs to just do one
 	//  It's also not atomic with the above commit. We need a way to set both new HEAD and update the working
 	//  set together, atomically. We can't easily do this in noms right now, because the the data set is the unit of
@@ -405,7 +411,7 @@ func (sess *Session) CommitToDolt(
 	//  just no API which allows one to update more than one dataset in the same atomic transaction. We need to write
 	//  one.
 	//  Meanwhile, this is all kinds of thread-unsafe
-	commit, err := actions.CommitStaged(ctx, roots, dbData, props)
+	commit, err := actions.CommitStaged(ctx, roots, ws.MergeActive(), mergeParentCommits, dbData, props)
 	if err != nil {
 		return nil, err
 	}
@@ -416,7 +422,7 @@ func (sess *Session) CommitToDolt(
 	// repo state writer, so we're never persisting the new working set to disk like in a command line context.
 	// TODO: fix this mess
 
-	ws, err := sess.WorkingSet(ctx, dbName)
+	ws, err = sess.WorkingSet(ctx, dbName)
 	if err != nil {
 		return nil, err
 	}
@@ -463,7 +469,12 @@ func (sess *Session) CreateDoltCommit(ctx *sql.Context, dbName string) error {
 		return err
 	}
 
-	_, err = sess.CommitToDolt(ctx, roots, dbName, actions.CommitStagedProps{
+	var mergeParentCommits []*doltdb.Commit
+	if sessionState.WorkingSet.MergeActive() {
+		mergeParentCommits = []*doltdb.Commit{sessionState.WorkingSet.MergeState().Commit()}
+	}
+
+	_, err = sess.CommitToDolt(ctx, roots, mergeParentCommits, dbName, actions.CommitStagedProps{
 		Message:    fmt.Sprintf("Transaction commit at %s", ctx.QueryTime().UTC().Format("2006-01-02T15:04:05Z")),
 		Date:       ctx.QueryTime(),
 		AllowEmpty: false,

go/libraries/doltcore/sqle/dsess/transactions.go

@@ -124,7 +124,7 @@ func (tx *DoltTransaction) Commit(ctx *sql.Context, workingSet *doltdb.WorkingSe
 			return nil, err
 		}
 
-		existingWorkingRoot := ws.RootValue()
+		existingWorkingRoot := ws.WorkingRoot()
 
 		hash, err := ws.HashOf()
 		if err != nil {