From 622410bde016a0203276b4b24fc08e5f57063c70 Mon Sep 17 00:00:00 2001 From: Cassandra Heart Date: Mon, 10 Feb 2025 18:19:00 -0600 Subject: [PATCH] inline metadata calculation --- node/crypto/proof_tree.go | 112 +++++++++++++++++++------------------- 1 file changed, 57 insertions(+), 55 deletions(-) diff --git a/node/crypto/proof_tree.go b/node/crypto/proof_tree.go index 6ad3fce..b98bd0f 100644 --- a/node/crypto/proof_tree.go +++ b/node/crypto/proof_tree.go @@ -223,39 +223,6 @@ func getNibblesUntilDiverge(key1, key2 []byte, startDepth int) ([]int, int) { } } -func recalcMetadata(node VectorCommitmentNode) ( - leafCount int, - longestBranch int, - size *big.Int, -) { - switch n := node.(type) { - case *VectorCommitmentLeafNode: - // A leaf counts as one, and its depth (from itself) is zero. - return 1, 0, n.Size - case *VectorCommitmentBranchNode: - totalLeaves := 0 - maxChildDepth := 0 - size := new(big.Int) - for _, child := range n.Children { - if child != nil { - cLeaves, cDepth, cSize := recalcMetadata(child) - totalLeaves += cLeaves - size.Add(size, cSize) - if cDepth > maxChildDepth { - maxChildDepth = cDepth - } - } - } - // Store the aggregated values in the branch node. - n.LeafCount = totalLeaves - // The branch’s longest branch is one more than its deepest child. - n.LongestBranch = maxChildDepth + 1 - n.Size = size - return totalLeaves, n.LongestBranch, n.Size - } - return 0, 0, new(big.Int) -} - // Insert adds or updates a key-value pair in the tree func (t *VectorCommitmentTree) Insert( key, value, hashTarget []byte, @@ -264,7 +231,6 @@ func (t *VectorCommitmentTree) Insert( if len(key) == 0 { return errors.New("empty key not allowed") } - var insert func(node VectorCommitmentNode, depth int) VectorCommitmentNode insert = func(node VectorCommitmentNode, depth int) VectorCommitmentNode { if node == nil { @@ -291,7 +257,10 @@ func (t *VectorCommitmentTree) Insert( // Create single branch node with shared prefix branch := &VectorCommitmentBranchNode{ - Prefix: sharedNibbles, + Prefix: sharedNibbles, + LeafCount: 2, + LongestBranch: 1, + Size: new(big.Int).Add(n.Size, size), } // Add both leaves at their final positions @@ -315,7 +284,10 @@ func (t *VectorCommitmentTree) Insert( if actualNibble != expectedNibble { // Create new branch with shared prefix subset newBranch := &VectorCommitmentBranchNode{ - Prefix: n.Prefix[:i], + Prefix: n.Prefix[:i], + LeafCount: n.LeafCount + 1, + LongestBranch: n.LongestBranch + 1, + Size: new(big.Int).Add(n.Size, size), } // Position old branch and new leaf newBranch.Children[expectedNibble] = n @@ -326,26 +298,45 @@ func (t *VectorCommitmentTree) Insert( HashTarget: hashTarget, Size: size, } - recalcMetadata(newBranch) return newBranch } } // Key matches prefix, continue with final nibble finalNibble := getNextNibble(key, depth+len(n.Prefix)*BranchBits) - n.Children[finalNibble] = insert( + inserted := insert( n.Children[finalNibble], depth+len(n.Prefix)*BranchBits+BranchBits, ) + n.Children[finalNibble] = inserted n.Commitment = nil - recalcMetadata(n) + n.LeafCount += 1 + switch i := inserted.(type) { + case *VectorCommitmentBranchNode: + if n.LongestBranch <= i.LongestBranch { + n.LongestBranch = i.LongestBranch + 1 + } + case *VectorCommitmentLeafNode: + n.LongestBranch = 1 + } + n.Size = n.Size.Add(n.Size, size) return n } else { // Simple branch without prefix nibble := getNextNibble(key, depth) - n.Children[nibble] = insert(n.Children[nibble], depth+BranchBits) + inserted := insert(n.Children[nibble], depth+BranchBits) + n.Children[nibble] = inserted n.Commitment = nil - recalcMetadata(n) + n.LeafCount += 1 + switch i := inserted.(type) { + case *VectorCommitmentBranchNode: + if n.LongestBranch <= i.LongestBranch { + n.LongestBranch = i.LongestBranch + 1 + } + case *VectorCommitmentLeafNode: + n.LongestBranch = 1 + } + n.Size = n.Size.Add(n.Size, size) return n } } @@ -490,30 +481,31 @@ func (t *VectorCommitmentTree) Delete(key []byte) error { return errors.New("empty key not allowed") } - var remove func(node VectorCommitmentNode, depth int) VectorCommitmentNode - remove = func(node VectorCommitmentNode, depth int) VectorCommitmentNode { + var remove func(node VectorCommitmentNode, depth int) (*big.Int, VectorCommitmentNode) + remove = func(node VectorCommitmentNode, depth int) (*big.Int, VectorCommitmentNode) { if node == nil { - return nil + return big.NewInt(0), nil } switch n := node.(type) { case *VectorCommitmentLeafNode: if bytes.Equal(n.Key, key) { - return nil + return n.Size, nil } - return n + return big.NewInt(0), n case *VectorCommitmentBranchNode: for i, expectedNibble := range n.Prefix { currentNibble := getNextNibble(key, depth+i*BranchBits) if currentNibble != expectedNibble { - return n + return big.NewInt(0), n } } finalNibble := getNextNibble(key, depth+len(n.Prefix)*BranchBits) - n.Children[finalNibble] = + var size *big.Int + size, n.Children[finalNibble] = remove(n.Children[finalNibble], depth+len(n.Prefix)*BranchBits+BranchBits) n.Commitment = nil @@ -521,11 +513,22 @@ func (t *VectorCommitmentTree) Delete(key []byte) error { childCount := 0 var lastChild VectorCommitmentNode var lastChildIndex int + longestBranch := 1 + leaves := 0 for i, child := range n.Children { if child != nil { childCount++ lastChild = child lastChildIndex = i + switch c := child.(type) { + case *VectorCommitmentBranchNode: + leaves += c.LeafCount + if longestBranch < c.LongestBranch+1 { + longestBranch = c.LongestBranch + 1 + } + case *VectorCommitmentLeafNode: + leaves += 1 + } } } @@ -549,20 +552,19 @@ func (t *VectorCommitmentTree) Delete(key []byte) error { retNode = lastChild } default: + n.LongestBranch = longestBranch + n.LeafCount = leaves + n.Size = n.Size.Sub(n.Size, size) retNode = n } - if branch, ok := retNode.(*VectorCommitmentBranchNode); ok { - recalcMetadata(branch) - } - - return retNode + return size, retNode default: - return node + return big.NewInt(0), node } } - t.Root = remove(t.Root, 0) + _, t.Root = remove(t.Root, 0) return nil }