[Execution node] Trie pruning (backport to master)

ledger/complete/ledger.go

@@ -322,7 +322,9 @@ func (l *Ledger) ExportCheckpointAt(
 
 	emptyTrie := trie.NewEmptyMTrie()
 
-	newTrie, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads)
+	// no need to prune the data since it has already been prunned through migrations
+	applyPruning := false
+	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/README.md

@@ -246,7 +246,7 @@ This results in the following `Update` algorithm. When applying the updates `(pa
 
 
 ```golang
-FUNCTION Update(height Int, node Node, paths []Path, payloads []Payload, compactLeaf Node) Node {
+FUNCTION Update(height Int, node Node, paths []Path, payloads []Payload, compactLeaf Node, prune bool) Node {
  if len(paths) == 0 {
   // If a compactLeaf from a higher height is carried over, then we are necessarily in case 2.a 
   // (node == nil and only one register to create)
@@ -330,7 +330,15 @@ FUNCTION Update(height Int, node Node, paths []Path, payloads []Payload, compact
  if lChild == newlChild && rChild == newrChild {
   return node
  }
- return NewInterimNode(height, newlChild, newrChild)
-}
-```
 
+nodeToBeReturned := NewInterimNode(height, newlChild, newrChild)
+ // if pruning is enabled, check if we could compactify the nodes after the update
+ // a common example of this is when we update a register payload to nil from a non-nil value
+ // therefore at least one of the children might be a default node (any node that has hashvalue equal to the default hashValue for the given height)
+ if prune { 
+    return nodeToBeReturned.Compactify()
+ }
+
+ return nodeToBeReturned
+}
+```

ledger/complete/mtrie/flattener/forest.go

@@ -173,7 +173,13 @@ func RebuildNodes(storableNodes []*StorableNode) ([]*node.Node, error) {
 			if err != nil {
 				return nil, fmt.Errorf("failed to decode a hash of a storableNode %w", err)
 			}
-			node := node.NewNode(int(snode.Height), nodes[snode.LIndex], nodes[snode.RIndex], path, payload, nodeHash, snode.MaxDepth, snode.RegCount)
+			// make a copy of payload
+			var pl *ledger.Payload
+			if payload != nil {
+				pl = payload.DeepCopy()
+			}
+
+			node := node.NewNode(int(snode.Height), nodes[snode.LIndex], nodes[snode.RIndex], path, pl, nodeHash, snode.MaxDepth, snode.RegCount)
 			nodes = append(nodes, node)
 			continue
 		}

ledger/complete/mtrie/flattener/iterator_test.go

@@ -39,7 +39,7 @@ func TestPopulatedTrie(t *testing.T) {
 	paths := []ledger.Path{p1, p2}
 	payloads := []ledger.Payload{*v1, *v2}
 
-	testTrie, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads)
+	testTrie, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
 	require.NoError(t, err)
 
 	for itr := flattener.NewNodeIterator(testTrie); itr.Next(); {

ledger/complete/mtrie/flattener/trie_test.go

@@ -36,7 +36,7 @@ func TestTrieStoreAndLoad(t *testing.T) {
 	paths := []ledger.Path{p1, p2, p3, p4, p5}
 	payloads := []ledger.Payload{*v1, *v2, *v3, *v4, *v5}
 
-	newTrie, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads)
+	newTrie, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
 	require.NoError(t, err)
 
 	flattedTrie, err := flattener.FlattenTrie(newTrie)

ledger/complete/mtrie/forest.go

@@ -163,7 +163,9 @@ func (f *Forest) Update(u *ledger.TrieUpdate) (ledger.RootHash, error) {
 	// TODO rename metrics names
 	f.metrics.UpdateValuesSize(uint64(totalPayloadSize))
 
-	newTrie, err := trie.NewTrieWithUpdatedRegisters(parentTrie, deduplicatedPaths, deduplicatedPayloads)
+	// apply pruning on update
+	applyPruning := true
+	newTrie, err := trie.NewTrieWithUpdatedRegisters(parentTrie, deduplicatedPaths, deduplicatedPayloads, applyPruning)
 	if err != nil {
 		return emptyHash, fmt.Errorf("constructing updated trie failed: %w", err)
 	}
@@ -216,7 +218,12 @@ func (f *Forest) Proofs(r *ledger.TrieRead) (*ledger.TrieBatchProof, error) {
 
 	// if we have to insert empty values
 	if len(notFoundPaths) > 0 {
-		newTrie, err := trie.NewTrieWithUpdatedRegisters(stateTrie, notFoundPaths, notFoundPayloads)
+		// for proofs, we have to set the pruning to false,
+		// currently batch proofs are only consists of inclusion proofs
+		// so for non-inclusion proofs we expand the trie with nil value and use an inclusion proof
+		// instead. if pruning is enabled it would break this trick and return the exact trie.
+		applyPruning := false
+		newTrie, err := trie.NewTrieWithUpdatedRegisters(stateTrie, notFoundPaths, notFoundPayloads, applyPruning)
 		if err != nil {
 			return nil, err
 		}

ledger/complete/mtrie/forest_test.go

@@ -31,7 +31,7 @@ func TestTrieOperations(t *testing.T) {
 	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
 	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
 
-	updatedTrie, err := trie.NewTrieWithUpdatedRegisters(nt, []ledger.Path{p1}, []ledger.Payload{*v1})
+	updatedTrie, err := trie.NewTrieWithUpdatedRegisters(nt, []ledger.Path{p1}, []ledger.Payload{*v1}, true)
 	require.NoError(t, err)
 
 	// Add trie

ledger/complete/mtrie/node/node.go

@@ -57,18 +57,13 @@ func NewNode(height int,
 	maxDepth uint16,
 	regCount uint64) *Node {
 
-	var pl *ledger.Payload
-	if payload != nil {
-		pl = payload.DeepCopy()
-	}
-
 	n := &Node{
 		lChild:    lchild,
 		rChild:    rchild,
 		height:    height,
 		path:      path,
 		hashValue: hashValue,
-		payload:   pl,
+		payload:   payload,
 		maxDepth:  maxDepth,
 		regCount:  regCount,
 	}
@@ -78,6 +73,7 @@ func NewNode(height int,
 // NewLeaf creates a compact leaf Node.
 // UNCHECKED requirement: height must be non-negative
 // UNCHECKED requirement: payload is non nil
+// UNCHECKED requirement: payload should be deep copied if received from external sources
 func NewLeaf(path ledger.Path,
 	payload *ledger.Payload,
 	height int) *Node {
@@ -87,7 +83,7 @@ func NewLeaf(path ledger.Path,
 		rChild:   nil,
 		height:   height,
 		path:     path,
-		payload:  payload.DeepCopy(),
+		payload:  payload,
 		maxDepth: 0,
 		regCount: uint64(1),
 	}
@@ -122,12 +118,99 @@ func NewInterimNode(height int, lchild, rchild *Node) *Node {
 	return n
 }
 
+// CopyAndPromoteLeafNode makes a copy of a node and moves it one level higher to replace
+// the parent node, this method should only be called for leaf nodes
+// depending on where this node is located to the parent the
+// hash value would be different, if isRight is set to true the original place
+// of the node n was right child of its parent so the hash value would be adjusted accordingly
+func (n *Node) copyAndPromoteLeafNode(isRight bool) *Node {
+	// note path is an arrays (not slice) so it would be coppied
+	newNode := &Node{
+		height:   n.height + 1,
+		path:     n.path,
+		payload:  n.payload,
+		maxDepth: n.maxDepth,
+		regCount: n.regCount,
+	}
+	if isRight {
+		newNode.hashValue = hash.HashInterNode(ledger.GetDefaultHashForHeight(n.height), n.hashValue)
+	} else {
+		newNode.hashValue = hash.HashInterNode(n.hashValue, ledger.GetDefaultHashForHeight(n.height))
+	}
+	return newNode
+}
+
 // computeAndStoreHash computes the node's hash value and
 // stores the result in the nodes internal `hashValue` field
 func (n *Node) computeAndStoreHash() {
 	n.hashValue = n.computeHash()
 }
 
+// Compactify checks if the subtree represented by an interim-node can be simplified to its most concise representation by looking only at its direct children. The
+// compactified representation of a default node is `nil`. For a node that only has a
+// _single_ child that is itself a leaf, this method returns a new, fully compactified leaf.
+// Returns:
+//  * n: if the node cannot be compactified, we return the original node `n`
+//  * cn: if the node can be compactified, where cn is a newly created compactified leaf
+func (n *Node) Compactify() *Node {
+
+	// if is a default node return nil instaed
+	if n.isDefaultNode() {
+		return nil
+	}
+
+	// if is a non default leaf, return it as is (no need for deep copy)
+	if n.IsLeaf() {
+		return n
+	}
+
+	// if leaf return it as is
+	// if non leaf bubble up
+	lChildEmpty := true
+	rChildEmpty := true
+	if n.lChild != nil {
+		lChildEmpty = n.lChild.isDefaultNode()
+	}
+	if n.rChild != nil {
+		rChildEmpty = n.rChild.isDefaultNode()
+	}
+	if rChildEmpty && lChildEmpty {
+		// if both children are empty this is the same as as a default leafs
+		return nil
+	}
+
+	// if childNode is non empty
+	if rChildEmpty {
+		if !lChildEmpty && n.lChild.IsLeaf() {
+			return n.lChild.copyAndPromoteLeafNode(false)
+		}
+		// this would replace empty nodes with nil
+		n.rChild = nil
+	}
+	if lChildEmpty {
+		if !rChildEmpty && n.rChild.IsLeaf() {
+			return n.rChild.copyAndPromoteLeafNode(true)
+		}
+		// this would replace empty nodes with nil
+		n.lChild = nil
+	}
+
+	// else no change needed
+	return n
+}
+
+// isDefaultNode returns true if either the node is nil
+// or the node's hash value is equal to the default hash value
+// for that height, in other words, it
+// does not contains any non-empty value in its sub-trie
+func (n *Node) isDefaultNode() bool {
+	if n == nil {
+		return true
+	}
+	return n.hashValue == ledger.GetDefaultHashForHeight(n.height)
+
+}
+
 // computeHash returns the hashValue of the node
 func (n *Node) computeHash() hash.Hash {
 	// check for leaf node

ledger/complete/mtrie/node/node_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/ledger/common/hash"
 	"github.com/onflow/flow-go/ledger/common/utils"
 	"github.com/onflow/flow-go/ledger/complete/mtrie/node"
@@ -135,6 +136,125 @@ func Test_VerifyCachedHash(t *testing.T) {
 	require.True(t, n5.VerifyCachedHash())
 }
 
+func Test_Compactify(t *testing.T) {
+	// Paths are not acurate in this case which causes the compact value be wrong
+	path0 := utils.PathByUint16(0)             // 0000...
+	path1 := utils.PathByUint16(1<<14 + 1<<13) // 01100...
+	path2 := utils.PathByUint16(1 << 15)       // 1000...
+	payload1 := utils.LightPayload(2, 2)
+	payload2 := utils.LightPayload(2, 4)
+	emptyPayload := &ledger.Payload{}
+
+	t.Run("non-leaf non-empty on right and leaf empty on left", func(t *testing.T) {
+		//          n5
+		//       /     \
+		//      n3      n4(-)
+		//   /    \
+		//  n1(p1) n2(p2)
+		//
+		// n4 would be replaced with nil, but n5 won't be prunned
+		n1 := node.NewLeaf(path0, payload1, 254)
+		n2 := node.NewLeaf(path1, payload2, 254)
+		n3 := node.NewInterimNode(255, n1, n2)
+		n4 := node.NewLeaf(path2, emptyPayload, 255)
+		n5 := node.NewInterimNode(256, n3, n4)
+
+		nn5 := n5.Compactify()
+		require.Equal(t, n5.MaxDepth(), nn5.MaxDepth())
+		require.True(t, nn5.VerifyCachedHash())
+		require.True(t, nn5.VerifyCachedHash())
+		require.Nil(t, nn5.RightChild())
+		require.Equal(t, nn5, n5)
+	})
+
+	t.Run("lowest level right leaf be empty", func(t *testing.T) {
+		//          n5
+		//       /     \
+		//      n3      n4(p2)
+		//   /    \
+		//  n1(p1) n2(-)
+		//
+		// n2 represents an unallocated/empty register
+		// pruning n3 should result in
+		//
+		//          nn5
+		//       /     \
+		//     nn3(p1)  n4(p2)
+		// and nn5 pruning should result in no change
+		n1 := node.NewLeaf(path0, payload1, 254)
+		n2 := node.NewLeaf(path1, emptyPayload, 254)
+		n3 := node.NewInterimNode(255, n1, n2)
+		n4 := node.NewLeaf(path2, payload2, 255)
+		n5 := node.NewInterimNode(256, n3, n4)
+
+		nn3 := n3.Compactify()
+		require.True(t, nn3.VerifyCachedHash())
+		require.Equal(t, n3.Hash(), nn3.Hash())
+		require.Equal(t, payload1, nn3.Payload())
+		require.True(t, nn3.IsLeaf())
+
+		nn5 := n5.Compactify()
+		require.Equal(t, nn5, n5)
+	})
+
+	t.Run("lowest level left leaf be empty", func(t *testing.T) {
+		//          n5
+		//       /     \
+		//      n3      n4(p2)
+		//   /    \
+		//  n1(-) n2(p1)
+		//
+		// n1 represents an unallocated/empty register
+		// pruning should result in
+		//          nn5
+		//       /     \
+		//     nn3(p1)  n4(p2)
+		n1 := node.NewLeaf(path0, emptyPayload, 254)
+		n2 := node.NewLeaf(path1, payload1, 254)
+		n3 := node.NewInterimNode(255, n1, n2)
+		n4 := node.NewLeaf(path2, payload2, 255)
+		n5 := node.NewInterimNode(256, n3, n4)
+		require.True(t, n2.VerifyCachedHash())
+
+		nn3 := n3.Compactify()
+		require.True(t, nn3.VerifyCachedHash())
+		require.Equal(t, nn3.Hash(), n3.Hash())
+		require.Equal(t, nn3.Payload(), payload1)
+		require.True(t, nn3.IsLeaf())
+
+		nn5 := n5.Compactify()
+		require.Equal(t, nn5, n5)
+	})
+
+	t.Run("lowest level left and right leaves be empty", func(t *testing.T) {
+		//          n5
+		//       /     \
+		//      n3      n4(p1)
+		//   /    \
+		//  n1(-) n2(-)
+		//
+		// n1 and n2 represent unallocated/empty registers
+		// pruning should result in
+		//          nn5 (p1)
+		n1 := node.NewLeaf(path0, emptyPayload, 254)
+		n2 := node.NewLeaf(path1, emptyPayload, 254)
+		n3 := node.NewInterimNode(255, n1, n2)
+		n4 := node.NewLeaf(path2, payload1, 255)
+		n5 := node.NewInterimNode(256, n3, n4)
+		require.True(t, n2.VerifyCachedHash())
+
+		nn3 := n3.Compactify()
+		require.Nil(t, nn3)
+
+		nn5 := n5.Compactify()
+		require.True(t, nn5.VerifyCachedHash())
+		require.Equal(t, nn5.Hash(), n5.Hash())
+		require.Equal(t, nn5.Payload(), payload1)
+		require.True(t, nn5.IsLeaf())
+	})
+
+}
+
 func hashToString(hash hash.Hash) string {
 	return hex.EncodeToString(hash[:])
 }