trie/bintrie: fix grouped InternalNode serialization path mismatch (#569)

trie/bintrie/binary_node.go

@@ -78,7 +78,7 @@ type BinaryNode interface {
 // It traverses up to `remainingDepth` levels, storing hashes of bottom-layer children.
 // position tracks the current index (0 to 2^groupDepth - 1) for bitmap placement.
 // hashes collects the hashes of present children, bitmap tracks which positions are present.
-func serializeSubtree(node BinaryNode, remainingDepth int, position int, bitmap []byte, hashes *[]common.Hash) {
+func serializeSubtree(node BinaryNode, remainingDepth int, position int, absoluteDepth int, bitmap []byte, hashes *[]common.Hash) {
 	if remainingDepth == 0 {
 		// Bottom layer: store hash if not empty
 		switch node.(type) {
@@ -98,18 +98,30 @@ func serializeSubtree(node BinaryNode, remainingDepth int, position int, bitmap
 		// Recurse into left (bit 0) and right (bit 1) children
 		leftPos := position * 2
 		rightPos := position*2 + 1
-		serializeSubtree(n.left, remainingDepth-1, leftPos, bitmap, hashes)
-		serializeSubtree(n.right, remainingDepth-1, rightPos, bitmap, hashes)
+		serializeSubtree(n.left, remainingDepth-1, leftPos, absoluteDepth+1, bitmap, hashes)
+		serializeSubtree(n.right, remainingDepth-1, rightPos, absoluteDepth+1, bitmap, hashes)
 	case Empty:
 		// Empty subtree: all positions in this subtree are empty (bits already 0)
 		return
 	default:
-		// StemNode or HashedNode before reaching bottom: store hash at current position
-		// This creates a variable-depth group where this branch terminates early.
-		// We need to mark this single position and all its would-be descendants as "this hash".
-		// For simplicity, we store the hash at the first leaf position of this subtree.
-		firstLeafPos := position << remainingDepth
-		bitmap[firstLeafPos/8] |= 1 << (7 - (firstLeafPos % 8))
+		// StemNode or HashedNode encountered before reaching the group's bottom
+		// layer. Compute the leaf bitmap position where this node's hash will
+		// be stored.
+		leafPos := position
+		switch sn := node.(type) {
+		case *StemNode:
+			// Extend position using the stem's key bits so that
+			// GetValuesAtStem traversal (which follows key bits) finds the hash.
+			for d := 0; d < remainingDepth; d++ {
+				bit := sn.Stem[(absoluteDepth+d)/8] >> (7 - ((absoluteDepth + d) % 8)) & 1
+				leafPos = leafPos*2 + int(bit)
+			}
+		default:
+			// HashedNode or unknown: extend all-left (no key bits available).
+			// This matches the all-zero path that resolveNode would follow.
+			leafPos = position << remainingDepth
+		}
+		bitmap[leafPos/8] |= 1 << (7 - (leafPos % 8))
 		*hashes = append(*hashes, node.Hash())
 	}
 }
@@ -127,12 +139,16 @@ func SerializeNode(node BinaryNode, groupDepth int) []byte {
 		bitmap := make([]byte, bitmapSize)
 		var hashes []common.Hash
 
-		serializeSubtree(n, groupDepth, 0, bitmap, &hashes)
+		serializeSubtree(n, groupDepth, 0, n.depth, bitmap, &hashes)
 
 		// Build serialized output
 		serializedLen := NodeTypeBytes + 1 + bitmapSize + len(hashes)*HashSize
 		serialized := make([]byte, serializedLen)
 		serialized[0] = nodeTypeInternal
+		// Store the group depth so deserialization knows the bitmap size.
+		// The bottom layer of the internal subtree may be sparse (e.g. a
+		// StemNode terminates a branch early), making the depth necessary
+		// to correctly interpret the variable-length bitmap that follows.
 		serialized[1] = byte(groupDepth)
 		copy(serialized[2:2+bitmapSize], bitmap)
 

trie/bintrie/group_debug_test.go

@@ -344,9 +344,10 @@ func buildDeepTreeUnique(depth, maxDepth, position int) BinaryNode {
 		return HashedNode(h)
 	}
 	return &InternalNode{
-		depth: depth,
-		left:  buildDeepTreeUnique(depth+1, maxDepth, position*2),
-		right: buildDeepTreeUnique(depth+1, maxDepth, position*2+1),
+		depth:         depth,
+		left:          buildDeepTreeUnique(depth+1, maxDepth, position*2),
+		right:         buildDeepTreeUnique(depth+1, maxDepth, position*2+1),
+		mustRecompute: true,
 	}
 }
 

trie/bintrie/internal_node.go

@@ -238,7 +238,12 @@ func (bt *InternalNode) collectChildGroups(path []byte, flushfn NodeFlushFn, gro
 		return nil
 	}
 
-	// Continue traversing within the group
+	// Continue traversing within the group.
+	// When a non-InternalNode (StemNode, HashedNode) appears mid-group, its
+	// hash gets projected to a leaf bitmap position by serializeSubtree. The
+	// storage path must be extended to match that projected position so that
+	// lookups after deserialization find the node at the correct path.
+	childDepth := bt.depth + 1
 	if bt.left != nil {
 		switch n := bt.left.(type) {
 		case *InternalNode:
@@ -246,8 +251,8 @@ func (bt *InternalNode) collectChildGroups(path []byte, flushfn NodeFlushFn, gro
 				return err
 			}
 		default:
-			// StemNode, HashedNode, or Empty - they handle their own collection
-			if err := bt.left.CollectNodes(appendBit(path, 0), flushfn, groupDepth); err != nil {
+			extPath := extendPathToGroupLeaf(appendBit(path, 0), bt.left, remainingLevels, childDepth)
+			if err := bt.left.CollectNodes(extPath, flushfn, groupDepth); err != nil {
 				return err
 			}
 		}
@@ -259,8 +264,8 @@ func (bt *InternalNode) collectChildGroups(path []byte, flushfn NodeFlushFn, gro
 				return err
 			}
 		default:
-			// StemNode, HashedNode, or Empty - they handle their own collection
-			if err := bt.right.CollectNodes(appendBit(path, 1), flushfn, groupDepth); err != nil {
+			extPath := extendPathToGroupLeaf(appendBit(path, 1), bt.right, remainingLevels, childDepth)
+			if err := bt.right.CollectNodes(extPath, flushfn, groupDepth); err != nil {
 				return err
 			}
 		}
@@ -268,6 +273,29 @@ func (bt *InternalNode) collectChildGroups(path []byte, flushfn NodeFlushFn, gro
 	return nil
 }
 
+// extendPathToGroupLeaf extends a storage path to the group's leaf boundary,
+// matching the projection done by serializeSubtree. For StemNodes, the path
+// is extended using the stem's key bits (same as serializeSubtree). For other
+// node types, the path is extended with all-zero (left) bits.
+func extendPathToGroupLeaf(path []byte, node BinaryNode, remainingLevels int, absoluteDepth int) []byte {
+	if remainingLevels <= 0 {
+		return path
+	}
+	if sn, ok := node.(*StemNode); ok {
+		for d := 0; d < remainingLevels; d++ {
+			bit := sn.Stem[(absoluteDepth+d)/8] >> (7 - ((absoluteDepth + d) % 8)) & 1
+			path = appendBit(path, bit)
+		}
+	} else {
+		// HashedNode or other: all-left extension (matches serializeSubtree's
+		// position << remainingDepth behavior).
+		for d := 0; d < remainingLevels; d++ {
+			path = appendBit(path, 0)
+		}
+	}
+	return path
+}
+
 // appendBit appends a bit to a path, returning a new slice
 func appendBit(path []byte, bit byte) []byte {
 	var p [256]byte

trie/bintrie/internal_node_test.go

@@ -391,11 +391,14 @@ func TestInternalNodeCollectNodes(t *testing.T) {
 		t.Errorf("Expected 3 collected nodes, got %d", len(collectedNodes))
 	}
 
-	// Check paths
+	// Check paths — with the groupDepth fix, mid-group StemNodes get their
+	// storage path extended to the group leaf boundary using stem key bits.
+	// Left stem (all-zero stem): extends with 7 zero bits from depth 1 to 8.
+	// Right stem (stem[0]=0x80, bit 1=1): extends with 1,0,0,0,0,0,0 from depth 1 to 8.
 	expectedPaths := [][]byte{
-		{1, 0}, // left child
-		{1, 1}, // right child
-		{1},    // internal node itself
+		{1, 0, 0, 0, 0, 0, 0, 0, 0}, // left child: path + 0 + 7 zero extension bits
+		{1, 1, 0, 0, 0, 0, 0, 0, 0}, // right child: path + 1 + stem-bit extension (0,0,0,0,0,0,0)
+		{1},                           // internal node itself (at group boundary)
 	}
 
 	for i, expectedPath := range expectedPaths {