mssmt: add InsertMany method to full and compacted trees

This commit introduces the InsertMany method to both the FullTree and
CompactedTree implementations of the MS-SMT. This method allows for the
insertion of multiple leaf nodes in a single database transaction,
improving efficiency when adding multiple leaves at once.

The InsertMany method is added to the Tree interface and implemented in
both FullTree and CompactedTree. The implementation includes sum
overflow checks before each insertion and updates the root within the
transaction for consistency.

A new test case, TestInsertMany, is added to verify the functionality of the
InsertMany method in both FullTree and CompactedTree. The test inserts a
random set of leaves using InsertMany and verifies the resulting root and
retrieved leaves.

The Copy method in both FullTree and CompactedTree is updated to use
InsertMany for efficiency when copying leaves to the target tree.

Author: Roasbeef

mssmt/compacted_tree.go

@@ -543,15 +543,73 @@ func (t *CompactedTree) Copy(ctx context.Context, targetTree Tree) error {
 		return err
 	}
 
-	// Insert all found leaves into the target tree.
-	for key, leaf := range leaves {
-		// Use the target tree's Insert method.
-		_, err := targetTree.Insert(ctx, key, leaf)
+	// Insert all found leaves into the target tree using InsertMany for
+	// efficiency.
+	_, err = targetTree.InsertMany(ctx, leaves)
+	if err != nil {
+		return fmt.Errorf("error inserting leaves into "+
+			"target tree: %w", err)
+	}
+
+	return nil
+}
+
+// InsertMany inserts multiple leaf nodes provided in the leaves map within a
+// single database transaction.
+func (t *CompactedTree) InsertMany(ctx context.Context,
+	leaves map[[hashSize]byte]*LeafNode) (Tree, error) {
+
+	if len(leaves) == 0 {
+		return t, nil
+	}
+
+	dbErr := t.store.Update(ctx, func(tx TreeStoreUpdateTx) error {
+		currentRoot, err := tx.RootNode()
 		if err != nil {
-			return fmt.Errorf("error inserting leaf with key %x "+
-				"into target tree: %w", key, err)
+			return err
 		}
+		rootBranch := currentRoot.(*BranchNode)
+
+		for key, leaf := range leaves {
+			// Check for potential sum overflow before each
+			// insertion.
+			sumRoot := rootBranch.NodeSum()
+			sumLeaf := leaf.NodeSum()
+			err = CheckSumOverflowUint64(sumRoot, sumLeaf)
+			if err != nil {
+				return fmt.Errorf("compact tree leaf insert "+
+					"sum overflow, root: %d, leaf: %d; %w",
+					sumRoot, sumLeaf, err)
+			}
+
+			// Insert the leaf using the internal helper.
+			newRoot, err := t.insert(
+				tx, &key, 0, rootBranch, leaf,
+			)
+			if err != nil {
+				return fmt.Errorf("error inserting leaf "+
+					"with key %x: %w", key, err)
+			}
+			rootBranch = newRoot
+
+			// Update the root within the transaction for
+			// consistency, even though the insert logic passes the
+			// root explicitly.
+			err = tx.UpdateRoot(rootBranch)
+			if err != nil {
+				return fmt.Errorf("error updating root "+
+					"during InsertMany: %w", err)
+			}
+		}
+
+		// The root is already updated by the last iteration of the
+		// loop. No final update needed here, but returning nil error
+		// signals success.
+		return nil
+	})
+	if dbErr != nil {
+		return nil, dbErr
 	}
 
-	return nil
+	return t, nil
 }

mssmt/interface.go

@@ -31,6 +31,11 @@ type Tree interface {
 	// leaf.
 	MerkleProof(ctx context.Context, key [hashSize]byte) (*Proof, error)
 
+	// InsertMany inserts multiple leaf nodes provided in the leaves map
+	// within a single database transaction.
+	InsertMany(ctx context.Context, leaves map[[hashSize]byte]*LeafNode) (
+		Tree, error)
+
 	// Copy copies all the key-value pairs from the source tree into the
 	// target tree. The target tree is assumed to be empty.
 	Copy(ctx context.Context, targetTree Tree) error

mssmt/tree.go

@@ -423,20 +423,71 @@ func (t *FullTree) Copy(ctx context.Context, targetTree Tree) error {
 		return err
 	}
 
-	// Insert all found leaves into the target tree. We assume the target
-	// tree handles batching or individual inserts efficiently.
-	for key, leaf := range leaves {
-		// Use the target tree's Insert method. We ignore the returned
-		// tree as we are modifying the targetTree in place via its
-		// store.
-		_, err := targetTree.Insert(ctx, key, leaf)
+	// Insert all found leaves into the target tree using InsertMany for
+	// efficiency.
+	_, err = targetTree.InsertMany(ctx, leaves)
+	if err != nil {
+		return fmt.Errorf("error inserting leaves into target tree: %w", err)
+	}
+
+	return nil
+}
+
+// InsertMany inserts multiple leaf nodes provided in the leaves map within a
+// single database transaction.
+func (t *FullTree) InsertMany(ctx context.Context,
+	leaves map[[hashSize]byte]*LeafNode) (Tree, error) {
+
+	if len(leaves) == 0 {
+		return t, nil
+	}
+
+	err := t.store.Update(ctx, func(tx TreeStoreUpdateTx) error {
+		currentRoot, err := tx.RootNode()
 		if err != nil {
-			return fmt.Errorf("error inserting leaf with key %x "+
-				"into target tree: %w", key, err)
+			return err
 		}
+		rootBranch := currentRoot.(*BranchNode)
+
+		for key, leaf := range leaves {
+			// Check for potential sum overflow before each
+			// insertion.
+			sumRoot := rootBranch.NodeSum()
+			sumLeaf := leaf.NodeSum()
+			err = CheckSumOverflowUint64(sumRoot, sumLeaf)
+			if err != nil {
+				return fmt.Errorf("full tree leaf insert sum "+
+					"overflow, root: %d, leaf: %d; %w",
+					sumRoot, sumLeaf, err)
+			}
+
+			// Insert the leaf using the internal helper.
+			newRoot, err := t.insert(tx, &key, leaf)
+			if err != nil {
+				return fmt.Errorf("error inserting leaf "+
+					"with key %x: %w", key, err)
+			}
+			rootBranch = newRoot
+
+			// Update the root within the transaction so subsequent
+			// inserts in this batch read the correct state.
+			err = tx.UpdateRoot(rootBranch)
+			if err != nil {
+				return fmt.Errorf("error updating root "+
+					"during InsertMany: %w", err)
+			}
+		}
+
+		// The root is already updated by the last iteration of the
+		// loop. No final update needed here, but returning nil error
+		// signals success.
+		return nil
+	})
+	if err != nil {
+		return nil, err
 	}
 
-	return nil
+	return t, nil
 }
 
 // VerifyMerkleProof determines whether a merkle proof for the leaf found at the

mssmt/tree_test.go

@@ -918,6 +918,68 @@ func TestTreeCopy(t *testing.T) {
 	}
 }
 
+// TestInsertMany tests inserting multiple leaves using the InsertMany method.
+func TestInsertMany(t *testing.T) {
+	t.Parallel()
+
+	leavesToInsert := randTree(50)
+	leavesMap := make(map[[hashSize]byte]*mssmt.LeafNode)
+	for _, item := range leavesToInsert {
+		leavesMap[item.key] = item.leaf
+	}
+
+	// Calculate expected root after individual insertions for comparison.
+	tempStore := mssmt.NewDefaultStore()
+	tempTree := mssmt.NewFullTree(tempStore)
+	ctx := context.Background()
+	for key, leaf := range leavesMap {
+		_, err := tempTree.Insert(ctx, key, leaf)
+		require.NoError(t, err)
+	}
+	expectedRoot, err := tempTree.Root(ctx)
+	require.NoError(t, err)
+
+	runTest := func(t *testing.T, name string,
+		makeTree func(mssmt.TreeStore) mssmt.Tree,
+		makeStore makeTestTreeStoreFunc) {
+
+		t.Run(name, func(t *testing.T) {
+			store, err := makeStore()
+			require.NoError(t, err)
+			tree := makeTree(store)
+
+			// Test inserting an empty map (should be a no-op).
+			_, err = tree.InsertMany(ctx, make(map[[hashSize]byte]*mssmt.LeafNode))
+			require.NoError(t, err)
+			initialRoot, err := tree.Root(ctx)
+			require.NoError(t, err)
+			require.True(t, mssmt.IsEqualNode(mssmt.EmptyTree[0], initialRoot))
+
+			// Insert the leaves using InsertMany.
+			_, err = tree.InsertMany(ctx, leavesMap)
+			require.NoError(t, err)
+
+			// Verify the root.
+			finalRoot, err := tree.Root(ctx)
+			require.NoError(t, err)
+			require.True(t, mssmt.IsEqualNode(expectedRoot, finalRoot))
+
+			// Verify each leaf can be retrieved.
+			for key, expectedLeaf := range leavesMap {
+				retrievedLeaf, err := tree.Get(ctx, key)
+				require.NoError(t, err)
+				require.Equal(t, expectedLeaf, retrievedLeaf)
+			}
+		})
+	}
+
+	for storeName, makeStore := range genTestStores(t) {
+		t.Run(storeName, func(t *testing.T) {
+			runTest(t, "full SMT", makeFullTree, makeStore)
+			runTest(t, "smol SMT", makeSmolTree, makeStore)
+		})
+	}
+}
 
 // runBIPTestVector runs the tests in a single BIP test vector file.
 func runBIPTestVector(t *testing.T, testVectors *mssmt.TestVectors) {