Remove cached node bytes from merkle nodes

x/merkledb/codec.go

@@ -29,18 +29,17 @@ const (
 	minDBNodeLen         = minMaybeByteSliceLen + minVarIntLen
 	minChildLen          = minVarIntLen + minKeyLen + ids.IDLen + boolLen
 
-	estimatedKeyLen           = 64
-	estimatedValueLen         = 64
-	estimatedCompressedKeyLen = 8
-	// Child index, child compressed key, child ID, child has value
-	estimatedNodeChildLen = minVarIntLen + estimatedCompressedKeyLen + ids.IDLen + boolLen
+	estimatedKeyLen   = 64
+	estimatedValueLen = 64
 	// Child index, child ID
 	hashValuesChildLen = minVarIntLen + ids.IDLen
 )
 
 var (
 	_ encoderDecoder = (*codecImpl)(nil)
 
+	log128 = math.Log(128)
+
 	trueBytes  = []byte{trueByte}
 	falseBytes = []byte{falseByte}
 
@@ -62,6 +61,7 @@ type encoderDecoder interface {
 type encoder interface {
 	// Assumes [n] is non-nil.
 	encodeDBNode(n *dbNode) []byte
+	encodedDBNodeSize(n *dbNode) int
 
 	// Returns the bytes that will be hashed to generate [n]'s ID.
 	// Assumes [n] is non-nil.
@@ -91,16 +91,42 @@ type codecImpl struct {
 	varIntPool sync.Pool
 }
 
-func (c *codecImpl) encodeDBNode(n *dbNode) []byte {
-	var (
-		numChildren = len(n.children)
-		// Estimate size of [n] to prevent memory allocations
-		estimatedLen = estimatedValueLen + minVarIntLen + estimatedNodeChildLen*numChildren
-		buf          = bytes.NewBuffer(make([]byte, 0, estimatedLen))
-	)
+func (c *codecImpl) encodedDBNodeSize(n *dbNode) int {
+	// total the number of children pointers + bool indicating if it has a value + the value + the child entries for n.children
+	total := c.uintSize(uint64(len(n.children))) + boolLen
+	if n.value.HasValue() {
+		total += c.uintSize(uint64(len(n.value.Value()))) + len(n.value.Value())
+	}
+	// for each non-nil entry, we add the additional size of the child entry
+	for index, entry := range n.children {
+		total += c.childSize(index, entry)
+	}
+	return total
+}
 
+func (c *codecImpl) childSize(index byte, childEntry *child) int {
+	return c.uintSize(uint64(index)) + ids.IDLen + c.keySize(childEntry.compressedKey) + boolLen
+}
+
+// based on the current implementation of codecImpl.encodeUint which uses binary.PutUvarint
+func (*codecImpl) uintSize(value uint64) int {
+	// catch the call to log(0) and return early
+	if value == 0 {
+		return 1
+	}
+	// binary.PutUvarint repeatedly divides by 128 until the value is under 128,
+	// so count the number of times that occurs
+	return 1 + int(math.Log(float64(value))/log128)
+}
+
+func (c *codecImpl) keySize(p Key) int {
+	return c.uintSize(uint64(p.length)) + bytesNeeded(p.length)
+}
+
+func (c *codecImpl) encodeDBNode(n *dbNode) []byte {
+	buf := bytes.NewBuffer(make([]byte, 0, c.encodedDBNodeSize(n)))
 	c.encodeMaybeByteSlice(buf, n.value)
-	c.encodeUint(buf, uint64(numChildren))
+	c.encodeUint(buf, uint64(len(n.children)))
 	// Note we insert children in order of increasing index
 	// for determinism.
 	keys := maps.Keys(n.children)

x/merkledb/codec_test.go

@@ -165,7 +165,7 @@ func FuzzCodecDBNodeDeterministic(f *testing.F) {
 				}
 
 				nodeBytes := codec.encodeDBNode(&node)
-
+				require.Len(nodeBytes, codec.encodedDBNodeSize(&node))
 				var gotNode dbNode
 				require.NoError(codec.decodeDBNode(nodeBytes, &gotNode))
 				require.Equal(node, gotNode)

x/merkledb/db.go

@@ -1335,6 +1335,5 @@ func cacheEntrySize(key Key, n *node) int {
 	if n == nil {
 		return len(key.Bytes())
 	}
-	// nodes cache their bytes representation so the total memory consumed is roughly twice that
-	return len(key.Bytes()) + 2*len(n.bytes())
+	return len(key.Bytes()) + codec.encodedDBNodeSize(&n.dbNode)
 }

x/merkledb/node.go

@@ -29,7 +29,6 @@ type child struct {
 type node struct {
 	dbNode
 	key         Key
-	nodeBytes   []byte
 	valueDigest maybe.Maybe[[]byte]
 }
 
@@ -50,9 +49,8 @@ func parseNode(key Key, nodeBytes []byte) (*node, error) {
 		return nil, err
 	}
 	result := &node{
-		dbNode:    n,
-		key:       key,
-		nodeBytes: nodeBytes,
+		dbNode: n,
+		key:    key,
 	}
 
 	result.setValueDigest()
@@ -66,17 +64,7 @@ func (n *node) hasValue() bool {
 
 // Returns the byte representation of this node.
 func (n *node) bytes() []byte {
-	if n.nodeBytes == nil {
-		n.nodeBytes = codec.encodeDBNode(&n.dbNode)
-	}
-
-	return n.nodeBytes
-}
-
-// clear the cached values that will need to be recalculated whenever the node changes
-// for example, node ID and byte representation
-func (n *node) onNodeChanged() {
-	n.nodeBytes = nil
+	return codec.encodeDBNode(&n.dbNode)
 }
 
 // Returns and caches the ID of this node.
@@ -88,7 +76,6 @@ func (n *node) calculateID(metrics merkleMetrics) ids.ID {
 
 // Set [n]'s value to [val].
 func (n *node) setValue(val maybe.Maybe[[]byte]) {
-	n.onNodeChanged()
 	n.value = val
 	n.setValueDigest()
 }
@@ -116,13 +103,11 @@ func (n *node) addChild(childNode *node, tokenSize int) {
 
 // Adds a child to [n] without a reference to the child node.
 func (n *node) setChildEntry(index byte, childEntry *child) {
-	n.onNodeChanged()
 	n.children[index] = childEntry
 }
 
 // Removes [child] from [n]'s children.
 func (n *node) removeChild(child *node, tokenSize int) {
-	n.onNodeChanged()
 	delete(n.children, child.key.Token(n.key.length, tokenSize))
 }
 
@@ -138,7 +123,6 @@ func (n *node) clone() *node {
 			children: make(map[byte]*child, len(n.children)),
 		},
 		valueDigest: n.valueDigest,
-		nodeBytes:   n.nodeBytes,
 	}
 	for key, existing := range n.children {
 		result.children[key] = &child{

x/merkledb/trieview.go

@@ -284,7 +284,6 @@ func (t *trieView) calculateNodeIDsHelper(n *node) ids.ID {
 			// This child wasn't changed.
 			continue
 		}
-		n.onNodeChanged()
 		childEntry.hasValue = childNodeChange.after.hasValue()
 
 		// Try updating the child and its descendants in a goroutine.