mssmt: add new batched ms-smt leaf insertion with recursive merge update

mssmt/compacted_tree.go

@@ -2,7 +2,9 @@
 
 import (
 	"context"
+	"bytes"
 	"fmt"
+	"sort"
 )
 
 // CompactedTree represents a compacted Merkle-Sum Sparse Merkle Tree (MS-SMT).
@@ -15,6 +17,276 @@
 	store TreeStore
 }
 
+// batched_insert handles the insertion of multiple entries in one go.
+func (t *CompactedTree) batched_insert(tx TreeStoreUpdateTx, entries []BatchedInsertionEntry, height int, root *BranchNode) (*BranchNode, error) {
+	// Base-case: If we've reached the bottom, simply return the current branch.
+	if height >= lastBitIndex {
+		return root, nil
+	}
+
+	// Guard against empty batch.
+	if len(entries) == 0 {
+		return root, nil
+	}
+
+	// Partition entries into two groups based on bit at current height.
+	var leftEntries, rightEntries []BatchedInsertionEntry
+	for _, entry := range entries {
+		if bitIndex(uint8(height), &entry.Key) == 0 {
+			leftEntries = append(leftEntries, entry)
+		} else {
+			rightEntries = append(rightEntries, entry)
+		}
+	}
+
+	// Get the current children from the node.
+	leftChild, rightChild, err := tx.GetChildren(height, root.NodeHash())
+	if err != nil {
+		return nil, err
+	}
+
+	// Process left subtree:
+	var newLeft Node
+	if len(leftEntries) > 0 {
+		// Check if the current left child is not empty.
+		if leftChild != EmptyTree[height+1] {
+			// If the existing child is a compacted leaf, we must handle potential collisions.
+			if cl, ok := leftChild.(*CompactedLeafNode); ok {
+				if len(leftEntries) == 1 {
+					entry := leftEntries[0]
+					if entry.Key == cl.Key() {
+						// Replacement: update the compacted leaf.
+						newLeaf := NewCompactedLeafNode(height+1, &entry.Key, entry.Leaf)
+						if err := tx.DeleteCompactedLeaf(cl.NodeHash()); err != nil {
+							return nil, err
+						}
+						if err := tx.InsertCompactedLeaf(newLeaf); err != nil {
+							return nil, err
+						}
+						newLeft = newLeaf
+					} else {
+						// Collision – keys differ: call merge to combine the new entry with the existing compacted leaf.
+						newLeft, err = t.merge(tx, height+1, entry.Key, entry.Leaf, cl.Key(), cl.LeafNode)
+						if err != nil {
+							return nil, err
+						}
+					}
+				} else {
+					// Multiple batch entries – check if they all match the existing key.
+					allMatch := true
+					for _, entry := range leftEntries {
+						if entry.Key != cl.Key() {
+							allMatch = false
+							break
+						}
+					}
+					if allMatch {
+						// All entries match; take the last one as replacement.
+						lastEntry := leftEntries[len(leftEntries)-1]
+						newLeaf := NewCompactedLeafNode(height+1, &lastEntry.Key, lastEntry.Leaf)
+						if err := tx.DeleteCompactedLeaf(cl.NodeHash()); err != nil {
+							return nil, err
+						}
+						if err := tx.InsertCompactedLeaf(newLeaf); err != nil {
+							return nil, err
+						}
+						newLeft = newLeaf
+					} else {
+						// At least one entry has a different key – merge using the first differing entry.
+						var mergeEntry *BatchedInsertionEntry
+						for _, entry := range leftEntries {
+							if entry.Key != cl.Key() {
+								mergeEntry = &entry
+								break
+							}
+						}
+						if mergeEntry == nil {
+							return nil, fmt.Errorf("unexpected nil merge entry")
+						}
+						newLeft, err = t.merge(tx, height+1, mergeEntry.Key, mergeEntry.Leaf, cl.Key(), cl.LeafNode)
+						if err != nil {
+							return nil, err
+						}
+					}
+				}
+			} else {
+				// leftChild is not a compacted leaf, so it must be a branch; recurse normally.
+				baseLeft := leftChild.(*BranchNode)
+				newLeft, err = t.batched_insert(tx, leftEntries, height+1, baseLeft)
+				if err != nil {
+					return nil, err
+				}
+			}
+		} else {
+			// The left child is empty.
+			if len(leftEntries) == 1 {
+				entry := leftEntries[0]
+				newLeft = NewCompactedLeafNode(height+1, &entry.Key, entry.Leaf)
+				if err := tx.InsertCompactedLeaf(newLeft.(*CompactedLeafNode)); err != nil {
+					return nil, err
+				}
+			} else {
+				baseLeft := EmptyTree[height+1].(*BranchNode)
+				newLeft, err = t.batched_insert(tx, leftEntries, height+1, baseLeft)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+		// Use newLeft as the computed left child.
+		leftChild = newLeft
+	}
+
+	// Process right subtree:
+	var newRight Node
+	if len(rightEntries) > 0 {
+		// Check if the current right child is not empty.
+		if rightChild != EmptyTree[height+1] {
+			// If the existing child is a compacted leaf, we must handle potential collisions.
+			if cr, ok := rightChild.(*CompactedLeafNode); ok {
+				if len(rightEntries) == 1 {
+					entry := rightEntries[0]
+					if entry.Key == cr.Key() {
+						// Replacement: update the compacted leaf.
+						newLeaf := NewCompactedLeafNode(height+1, &entry.Key, entry.Leaf)
+						if err := tx.DeleteCompactedLeaf(cr.NodeHash()); err != nil {
+							return nil, err
+						}
+						if err := tx.InsertCompactedLeaf(newLeaf); err != nil {
+							return nil, err
+						}
+						newRight = newLeaf
+					} else {
+						// Collision – keys differ: call merge to combine the new entry with the existing compacted leaf.
+						newRight, err = t.merge(tx, height+1, entry.Key, entry.Leaf, cr.Key(), cr.LeafNode)
+						if err != nil {
+							return nil, err
+						}
+					}
+				} else {
+					// Multiple batch entries – check if they all match the existing key.
+					allMatch := true
+					for _, entry := range rightEntries {
+						if entry.Key != cr.Key() {
+							allMatch = false
+							break
+						}
+					}
+					if allMatch {
+						// All entries match; take the last one as replacement.
+						lastEntry := rightEntries[len(rightEntries)-1]
+						newLeaf := NewCompactedLeafNode(height+1, &lastEntry.Key, lastEntry.Leaf)
+						if err := tx.DeleteCompactedLeaf(cr.NodeHash()); err != nil {
+							return nil, err
+						}
+						if err := tx.InsertCompactedLeaf(newLeaf); err != nil {
+							return nil, err
+						}
+						newRight = newLeaf
+					} else {
+						// At least one entry has a different key – merge using the first differing entry.
+						var mergeEntry *BatchedInsertionEntry
+						for _, entry := range rightEntries {
+							if entry.Key != cr.Key() {
+								mergeEntry = &entry
+								break
+							}
+						}
+						if mergeEntry == nil {
+							return nil, fmt.Errorf("unexpected nil merge entry")
+						}
+						newRight, err = t.merge(tx, height+1, mergeEntry.Key, mergeEntry.Leaf, cr.Key(), cr.LeafNode)
+						if err != nil {
+							return nil, err
+						}
+					}
+				}
+			} else {
+				// rightChild is not a compacted leaf, so it must be a branch; recurse normally.
+				baseRight := rightChild.(*BranchNode)
+				newRight, err = t.batched_insert(tx, rightEntries, height+1, baseRight)
+				if err != nil {
+					return nil, err
+				}
+			}
+		} else {
+			// The right child is empty.
+			if len(rightEntries) == 1 {
+				entry := rightEntries[0]
+				newRight = NewCompactedLeafNode(height+1, &entry.Key, entry.Leaf)
+				if err := tx.InsertCompactedLeaf(newRight.(*CompactedLeafNode)); err != nil {
+					return nil, err
+				}
+			} else {
+				baseRight := EmptyTree[height+1].(*BranchNode)
+				newRight, err = t.batched_insert(tx, rightEntries, height+1, baseRight)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+		// Use newRight as the computed right child.
+		rightChild = newRight
+	}
+
+	// Create the updated branch from the new left and right children.
+	var updatedBranch *BranchNode
+	updatedBranch = NewBranch(leftChild, rightChild)
+
+	// Delete the old branch and insert the new one.
+	if root != EmptyTree[height] {
+		if err := tx.DeleteBranch(root.NodeHash()); err != nil {
+			return nil, err
+		}
+	}
+	if !IsEqualNode(updatedBranch, EmptyTree[height]) {
+		if err := tx.InsertBranch(updatedBranch); err != nil {
+			return nil, err
+		}
+	}
+
+	return updatedBranch, nil
+}
+
+// BatchedInsert inserts multiple leaf nodes at the given keys within the MS-SMT.
+func (t *CompactedTree) BatchedInsert(ctx context.Context, entries []BatchedInsertionEntry) (Tree, error) {
+	sort.Slice(entries, func(i, j int) bool {
+		return bytes.Compare(entries[i].Key[:], entries[j].Key[:]) < 0
+	})
+
+	err := t.store.Update(ctx, func(tx TreeStoreUpdateTx) error {
+		currentRoot, err := tx.RootNode()
+		if err != nil {
+			return err
+		}
+		branchRoot := currentRoot.(*BranchNode)
+
+		// (Optional) Loop over entries and check for sum overflow.
+		for _, entry := range entries {
+			if err := CheckSumOverflowUint64(branchRoot.NodeSum(), entry.Leaf.NodeSum()); err != nil {
+				return fmt.Errorf("batched insert key %v sum overflow: %w", entry.Key, err)
+			}
+		}
+
+		// Call the new batched_insert method.
+		newRoot, err := t.batched_insert(tx, entries, 0, branchRoot)
+		if err != nil {
+			return err
+		}
+		return tx.UpdateRoot(newRoot)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return t, nil
+}
+
+// BatchedInsertionEntry represents one leaf insertion.
+type BatchedInsertionEntry struct {
+	Key  [hashSize]byte
+	Leaf *LeafNode
+}
+
 var _ Tree = (*CompactedTree)(nil)
 
 // NewCompactedTree initializes an empty MS-SMT backed by `store`.
@@ -68,19 +340,9 @@
 
 		switch node := next.(type) {
 		case *CompactedLeafNode:
-			// Our next node is a compacted leaf. We just need to
-			// expand it so we can continue our walk down the tree.
-			next = node.Extract(i)
-
-			// Sibling might be a compacted leaf too, in which case
-			// we need to extract it as well.
-			if compSibling, ok := sibling.(*CompactedLeafNode); ok {
-				sibling = compSibling.Extract(i)
-			}
+			// Our next node is a compacted leaf. We simply return the underlying leaf.
+			return node.LeafNode, nil
 
-			// Now that all required branches are reconstructed we
-			// can continue the search for the leaf matching the
-			// passed key.
 			for j := i; j <= lastBitIndex; j++ {
 				if iter != nil {
 					err := iter(j, next, sibling, current)
@@ -91,9 +353,6 @@
 				current = next
 
 				if j < lastBitIndex {
-					// Since we have all the branches we
-					// need extracted already we can just
-					// continue walking down.
 					branch := current.(*BranchNode)
 					next, sibling = stepOrder(
 						j+1, key, branch.Left,

mssmt/node.go

@@ -121,6 +121,9 @@ type CompactedLeafNode struct {
 	// compactedNodeHash holds the topmost (omitted) node's node hash in the
 	// subtree.
 	compactedNodeHash NodeHash
+
+	// Height is the level at which this compacted leaf was created.
+	Height int
 }
 
 // NewCompactedLeafNode creates a new compacted leaf at the passed height with
@@ -144,6 +147,7 @@ func NewCompactedLeafNode(height int, key *[32]byte,
 		compactedNodeHash: nodeHash,
 	}
 
+	node.Height = height
 	return node
 }
 
@@ -157,24 +161,35 @@ func (c *CompactedLeafNode) Key() [32]byte {
 	return c.key
 }
 
-// Extract extracts the subtree represented by this compacted leaf and returns
-// the topmost node in the tree.
-func (c *CompactedLeafNode) Extract(height int) Node {
-	var current Node = c.LeafNode
-
-	// Walk up and recreate the missing branches.
-	for j := MaxTreeLevels; j > height+1; j-- {
-		var left, right Node
-		if bitIndex(uint8(j-1), &c.key) == 0 {
-			left, right = current, EmptyTree[j]
-		} else {
-			left, right = EmptyTree[j], current
-		}
-
-		current = NewBranch(left, right)
-	}
-
-	return current
+func (c *CompactedLeafNode) Extract(requestedHeight int) Node {
+    target := requestedHeight
+    if target >= c.Height {
+        target = c.Height - 1
+    }
+
+    var current Node = c.LeafNode
+
+    for j := c.Height - 1; j >= target+1; j-- {
+        if bitIndex(uint8(j), &c.key) == 0 {
+            current = NewBranch(current, EmptyTree[j+1])
+        } else {
+            current = NewBranch(EmptyTree[j+1], current)
+        }
+    }
+
+    sibling := EmptyTree[target+1]
+    if target+1 == MaxTreeLevels {
+        sibling = NewBranch(sibling, sibling)
+    }
+
+    var result Node
+    if bitIndex(uint8(target), &c.key) == 0 {
+        result = NewBranch(current, sibling)
+    } else {
+        result = NewBranch(sibling, current)
+    }
+
+    return result
 }
 
 // Copy returns a deep copy of the compacted leaf node.