Optimize MTrie Checkpoint (regCount & regSize): -9GB alloc/op, -110 milllion allocs/op, -4GB file size

ledger/complete/ledger.go

@@ -178,7 +178,6 @@ func (l *Ledger) Set(update *ledger.Update) (newState ledger.State, trieUpdate *
 	}
 
 	l.metrics.UpdateCount()
-	l.metrics.UpdateValuesNumber(uint64(len(trieUpdate.Paths)))
 
 	walChan := make(chan error)
 
@@ -345,7 +344,7 @@ func (l *Ledger) ExportCheckpointAt(
 
 	// no need to prune the data since it has already been prunned through migrations
 	applyPruning := false
-	newTrie, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, applyPruning)
+	newTrie, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, applyPruning)
 	if err != nil {
 		return ledger.State(hash.DummyHash), fmt.Errorf("constructing updated trie failed: %w", err)
 	}

ledger/complete/mtrie/flattener/encoding.go

@@ -22,19 +22,19 @@ const (
 const (
 	encNodeTypeSize      = 1
 	encHeightSize        = 2
-	encMaxDepthSize      = 2
 	encRegCountSize      = 8
+	encRegSizeSize       = 8
 	encHashSize          = hash.HashLen
 	encPathSize          = ledger.PathLen
 	encNodeIndexSize     = 8
 	encPayloadLengthSize = 4
+
+	encodedTrieSize = encNodeIndexSize + encRegCountSize + encRegSizeSize + encHashSize
 )
 
 // encodeLeafNode encodes leaf node in the following format:
 // - node type (1 byte)
 // - height (2 bytes)
-// - max depth (2 bytes)
-// - reg count (8 bytes)
 // - hash (32 bytes)
 // - path (32 bytes)
 // - payload (4 bytes + n bytes)
@@ -52,8 +52,6 @@ func encodeLeafNode(n *node.Node, scratch []byte) []byte {
 
 	encodedNodeSize := encNodeTypeSize +
 		encHeightSize +
-		encMaxDepthSize +
-		encRegCountSize +
 		encHashSize +
 		encPathSize +
 		encPayloadLengthSize +
@@ -78,14 +76,6 @@ func encodeLeafNode(n *node.Node, scratch []byte) []byte {
 	binary.BigEndian.PutUint16(buf[pos:], uint16(n.Height()))
 	pos += encHeightSize
 
-	// Encode max depth (2 bytes Big Endian)
-	binary.BigEndian.PutUint16(buf[pos:], n.MaxDepth())
-	pos += encMaxDepthSize
-
-	// Encode reg count (8 bytes Big Endian)
-	binary.BigEndian.PutUint64(buf[pos:], n.RegCount())
-	pos += encRegCountSize
-
 	// Encode hash (32 bytes hashValue)
 	hash := n.Hash()
 	copy(buf[pos:], hash[:])
@@ -110,8 +100,6 @@ func encodeLeafNode(n *node.Node, scratch []byte) []byte {
 // encodeInterimNode encodes interim node in the following format:
 // - node type (1 byte)
 // - height (2 bytes)
-// - max depth (2 bytes)
-// - reg count (8 bytes)
 // - hash (32 bytes)
 // - lchild index (8 bytes)
 // - rchild index (8 bytes)
@@ -126,8 +114,6 @@ func encodeInterimNode(n *node.Node, lchildIndex uint64, rchildIndex uint64, scr
 
 	const encodedNodeSize = encNodeTypeSize +
 		encHeightSize +
-		encMaxDepthSize +
-		encRegCountSize +
 		encHashSize +
 		encNodeIndexSize +
 		encNodeIndexSize
@@ -151,14 +137,6 @@ func encodeInterimNode(n *node.Node, lchildIndex uint64, rchildIndex uint64, scr
 	binary.BigEndian.PutUint16(buf[pos:], uint16(n.Height()))
 	pos += encHeightSize
 
-	// Encode max depth (2 bytes Big Endian)
-	binary.BigEndian.PutUint16(buf[pos:], n.MaxDepth())
-	pos += encMaxDepthSize
-
-	// Encode reg count (8 bytes Big Endian)
-	binary.BigEndian.PutUint64(buf[pos:], n.RegCount())
-	pos += encRegCountSize
-
 	// Encode hash (32 bytes hashValue)
 	h := n.Hash()
 	copy(buf[pos:], h[:])
@@ -204,11 +182,7 @@ func ReadNode(reader io.Reader, scratch []byte, getNode func(nodeIndex uint64) (
 	}
 
 	// fixLengthSize is the size of shared data of leaf node and interim node
-	const fixLengthSize = encNodeTypeSize +
-		encHeightSize +
-		encMaxDepthSize +
-		encRegCountSize +
-		encHashSize
+	const fixLengthSize = encNodeTypeSize + encHeightSize + encHashSize
 
 	_, err := io.ReadFull(reader, scratch[:fixLengthSize])
 	if err != nil {
@@ -229,14 +203,6 @@ func ReadNode(reader io.Reader, scratch []byte, getNode func(nodeIndex uint64) (
 	height := binary.BigEndian.Uint16(scratch[pos:])
 	pos += encHeightSize
 
-	// Decode max depth (2 bytes)
-	maxDepth := binary.BigEndian.Uint16(scratch[pos:])
-	pos += encMaxDepthSize
-
-	// Decode reg count (8 bytes)
-	regCount := binary.BigEndian.Uint64(scratch[pos:])
-	pos += encRegCountSize
-
 	// Decode and create hash.Hash (32 bytes)
 	nodeHash, err := hash.ToHash(scratch[pos : pos+encHashSize])
 	if err != nil {
@@ -264,7 +230,7 @@ func ReadNode(reader io.Reader, scratch []byte, getNode func(nodeIndex uint64) (
 			return nil, fmt.Errorf("failed to read and decode payload of serialized node: %w", err)
 		}
 
-		node := node.NewNode(int(height), nil, nil, path, payload, nodeHash, maxDepth, regCount)
+		node := node.NewNode(int(height), nil, nil, path, payload, nodeHash)
 		return node, nil
 	}
 
@@ -297,58 +263,55 @@ func ReadNode(reader io.Reader, scratch []byte, getNode func(nodeIndex uint64) (
 		return nil, fmt.Errorf("failed to find right child node of serialized node: %w", err)
 	}
 
-	n := node.NewNode(int(height), lchild, rchild, ledger.DummyPath, nil, nodeHash, maxDepth, regCount)
+	n := node.NewNode(int(height), lchild, rchild, ledger.DummyPath, nil, nodeHash)
 	return n, nil
 }
 
 // EncodeTrie encodes trie in the following format:
 // - root node index (8 byte)
+// - allocated reg count (8 byte)
+// - allocated reg size (8 byte)
 // - root node hash (32 bytes)
 // Scratch buffer is used to avoid allocs.
 // WARNING: The returned buffer is likely to share the same underlying array as
 // the scratch buffer. Caller is responsible for copying or using returned buffer
 // before scratch buffer is used again.
-func EncodeTrie(rootNode *node.Node, rootIndex uint64, scratch []byte) []byte {
-
-	const encodedTrieSize = encNodeIndexSize + encHashSize
-
-	// Get root hash
-	var rootHash ledger.RootHash
-	if rootNode == nil {
-		rootHash = trie.EmptyTrieRootHash()
-	} else {
-		rootHash = ledger.RootHash(rootNode.Hash())
-	}
-
+func EncodeTrie(trie *trie.MTrie, rootIndex uint64, scratch []byte) []byte {
+	buf := scratch
 	if len(scratch) < encodedTrieSize {
-		scratch = make([]byte, encodedTrieSize)
+		buf = make([]byte, encodedTrieSize)
 	}
 
 	pos := 0
 
 	// Encode root node index (8 bytes Big Endian)
-	binary.BigEndian.PutUint64(scratch, rootIndex)
+	binary.BigEndian.PutUint64(buf, rootIndex)
 	pos += encNodeIndexSize
 
+	// Encode trie reg count (8 bytes Big Endian)
+	binary.BigEndian.PutUint64(buf[pos:], trie.AllocatedRegCount())
+	pos += encRegCountSize
+
+	// Encode trie reg size (8 bytes Big Endian)
+	binary.BigEndian.PutUint64(buf[pos:], trie.AllocatedRegSize())
+	pos += encRegSizeSize
+
 	// Encode hash (32-bytes hashValue)
-	copy(scratch[pos:], rootHash[:])
+	rootHash := trie.RootHash()
+	copy(buf[pos:], rootHash[:])
 	pos += encHashSize
 
-	return scratch[:pos]
+	return buf[:pos]
 }
 
 // ReadTrie reconstructs a trie from data read from reader.
 func ReadTrie(reader io.Reader, scratch []byte, getNode func(nodeIndex uint64) (*node.Node, error)) (*trie.MTrie, error) {
 
-	// encodedTrieSize is a failsafe and is only used when len(scratch) is much smaller
-	// than expected (4096 by default).
-	const encodedTrieSize = encNodeIndexSize + encHashSize
-
 	if len(scratch) < encodedTrieSize {
 		scratch = make([]byte, encodedTrieSize)
 	}
 
-	// Read encoded trie (8 + 32 bytes)
+	// Read encoded trie
 	_, err := io.ReadFull(reader, scratch[:encodedTrieSize])
 	if err != nil {
 		return nil, fmt.Errorf("failed to read serialized trie: %w", err)
@@ -360,6 +323,14 @@ func ReadTrie(reader io.Reader, scratch []byte, getNode func(nodeIndex uint64) (
 	rootIndex := binary.BigEndian.Uint64(scratch)
 	pos += encNodeIndexSize
 
+	// Decode trie reg count (8 bytes)
+	regCount := binary.BigEndian.Uint64(scratch[pos:])
+	pos += encRegCountSize
+
+	// Decode trie reg size (8 bytes)
+	regSize := binary.BigEndian.Uint64(scratch[pos:])
+	pos += encRegSizeSize
+
 	// Decode root node hash
 	readRootHash, err := hash.ToHash(scratch[pos : pos+encHashSize])
 	if err != nil {
@@ -371,7 +342,7 @@ func ReadTrie(reader io.Reader, scratch []byte, getNode func(nodeIndex uint64) (
 		return nil, fmt.Errorf("failed to find root node of serialized trie: %w", err)
 	}
 
-	mtrie, err := trie.NewMTrie(rootNode)
+	mtrie, err := trie.NewMTrie(rootNode, regCount, regSize)
 	if err != nil {
 		return nil, fmt.Errorf("failed to restore serialized trie: %w", err)
 	}