From 3e571fb67a009a521608ac3772bdb8a87fda5d90 Mon Sep 17 00:00:00 2001
From: Andy Arthur <andy@dolthub.com>
Date: Mon, 25 Jul 2022 13:15:04 -0700
Subject: [PATCH] copy blob tree along with map data

---
 go/libraries/doltcore/migrate/transform.go    | 39 ++++++++++++++-----
 .../sqle/json/noms_json_value_test.go         |  6 ++-
 go/store/cmd/noms/noms_show_test.go           |  3 +-
 go/store/nbs/frag/main.go                     |  3 +-
 go/store/types/ref.go                         |  5 +--
 5 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/go/libraries/doltcore/migrate/transform.go b/go/libraries/doltcore/migrate/transform.go
index 9652405d94..b83eeb5e34 100644
--- a/go/libraries/doltcore/migrate/transform.go
+++ b/go/libraries/doltcore/migrate/transform.go
@@ -186,13 +186,34 @@ func MigrateIndexSet(ctx context.Context, sch schema.Schema, oldSet durable.Inde
 
 func MigrateNomsMap(ctx context.Context, idx durable.Index, old, new types.ValueReadWriter) error {
 	m := durable.NomsMapFromIndex(idx)
-	_, _, err := types.VisitMapLevelOrder(m, func(h hash.Hash) (int64, error) {
-		v, err := old.ReadValue(ctx, h)
-		if err != nil {
-			return 0, err
-		}
-		_, err = new.WriteValue(ctx, v)
-		return 0, err
-	})
-	return err
+	return copyTreeFromValue(ctx, m, old, new)
+}
+
+// copyTreeFromValue recursively copies |v| and all its children from |old| to |new|.
+func copyTreeFromValue(ctx context.Context, v types.Value, old, new types.ValueReadWriter) error {
+	if _, err := new.WriteValue(ctx, v); err != nil {
+		return err
+	}
+	return types.WalkAddrs(v, old.Format(), func(h hash.Hash, isleaf bool) error {
+		if err := copyValue(ctx, h, old, new); err != nil {
+			return err
+		}
+		if isleaf {
+			return nil
+		}
+		val, err := old.ReadValue(ctx, h)
+		if err != nil {
+			return err
+		}
+		return copyTreeFromValue(ctx, val, old, new)
+	})
+}
+
+func copyValue(ctx context.Context, addr hash.Hash, old, new types.ValueReadWriter) (err error) {
+	var v types.Value
+	if v, err = old.ReadValue(ctx, addr); err != nil {
+		return err
+	}
+	_, err = new.WriteValue(ctx, v)
+	return
 }
diff --git a/go/libraries/doltcore/sqle/json/noms_json_value_test.go b/go/libraries/doltcore/sqle/json/noms_json_value_test.go
index fe8e624dac..8a7f63db7e 100644
--- a/go/libraries/doltcore/sqle/json/noms_json_value_test.go
+++ b/go/libraries/doltcore/sqle/json/noms_json_value_test.go
@@ -184,16 +184,18 @@ func TestJSONStructuralSharing(t *testing.T) {
 		val := MustNomsJSONWithVRW(vrw, sb.String())
 
 		json_refs := make(hash.HashSet)
-		err := types.WalkAddrs(types.JSON(val), vrw.Format(), func(h hash.Hash, _ bool) {
+		err := types.WalkAddrs(types.JSON(val), vrw.Format(), func(h hash.Hash, _ bool) error {
 			json_refs.Insert(h)
+			return nil
 		})
 		require.NoError(t, err)
 
 		tup, err := types.NewTuple(types.Format_Default, types.Int(12), types.JSON(val))
 		require.NoError(t, err)
 		tuple_refs := make(hash.HashSet)
-		types.WalkAddrs(tup, vrw.Format(), func(h hash.Hash, _ bool) {
+		types.WalkAddrs(tup, vrw.Format(), func(h hash.Hash, _ bool) error {
 			tuple_refs.Insert(h)
+			return nil
 		})
 
 		assert.Greater(t, len(json_refs), 0)
diff --git a/go/store/cmd/noms/noms_show_test.go b/go/store/cmd/noms/noms_show_test.go
index 89ab65ed66..ad3a114181 100644
--- a/go/store/cmd/noms/noms_show_test.go
+++ b/go/store/cmd/noms/noms_show_test.go
@@ -152,8 +152,9 @@ func (s *nomsShowTestSuite) TestNomsShowRaw() {
 	s.NoError(err)
 
 	numChildChunks := 0
-	err = types.WalkAddrs(l, vrw.Format(), func(_ hash.Hash, _ bool) {
+	err = types.WalkAddrs(l, vrw.Format(), func(_ hash.Hash, _ bool) error {
 		numChildChunks++
+		return nil
 	})
 	s.NoError(err)
 	s.True(numChildChunks > 0)
diff --git a/go/store/nbs/frag/main.go b/go/store/nbs/frag/main.go
index 143edc6067..860d9e36ba 100644
--- a/go/store/nbs/frag/main.go
+++ b/go/store/nbs/frag/main.go
@@ -130,11 +130,12 @@ func main() {
 		orderedChildren := hash.HashSlice{}
 		nextLevel := hash.HashSlice{}
 		for _, h := range current {
-			_ = types.WalkAddrs(currentValues[h], types.Format_Default, func(h hash.Hash, isleaf bool) {
+			_ = types.WalkAddrs(currentValues[h], types.Format_Default, func(h hash.Hash, isleaf bool) error {
 				orderedChildren = append(orderedChildren, h)
 				if !visited[h] && !isleaf {
 					nextLevel = append(nextLevel, h)
 				}
+				return nil
 			})
 		}
 
diff --git a/go/store/types/ref.go b/go/store/types/ref.go
index 175ee218b2..50e48b142b 100644
--- a/go/store/types/ref.go
+++ b/go/store/types/ref.go
@@ -242,9 +242,8 @@ func WalkAddrsForNBF(nbf *NomsBinFormat) func(chunks.Chunk, func(h hash.Hash, is
 	}
 }
 
-func WalkAddrs(v Value, nbf *NomsBinFormat, cb func(h hash.Hash, isleaf bool)) error {
+func WalkAddrs(v Value, nbf *NomsBinFormat, cb func(h hash.Hash, isleaf bool) error) error {
 	return v.walkRefs(nbf, func(r Ref) error {
-		cb(r.TargetHash(), r.Height() == 1)
-		return nil
+		return cb(r.TargetHash(), r.Height() == 1)
 	})
 }