From 103c45850afd58b8e4b7f4bf2fc419fec6885db7 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 28 May 2020 11:53:19 +0300
Subject: [PATCH] mpt: implement (*Trie).Collapse()

Because trie size is rather big, it can't be stored in memory.
Thus some form of caching should also be implemented. To avoid
marshaling/unmarshaling of items which are close to root and are used
very frequenly we can save them across the persists.
This commit implements pruning items at the specified depth,
replacing them by hash nodes.
---
 pkg/core/mpt/trie.go      | 32 +++++++++++++++++
 pkg/core/mpt/trie_test.go | 73 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/pkg/core/mpt/trie.go b/pkg/core/mpt/trie.go
index f9589fde3..c5093d614 100644
--- a/pkg/core/mpt/trie.go
+++ b/pkg/core/mpt/trie.go
@@ -355,3 +355,35 @@ func (t *Trie) getFromStore(h util.Uint256) (Node, error) {
 	}
 	return n.Node, nil
 }
+
+// Collapse compresses all nodes at depth n to the hash nodes.
+// Note: this function does not perform any kind of storage flushing so
+// `Flush()` should be called explicitly before invoking function.
+func (t *Trie) Collapse(depth int) {
+	if depth < 0 {
+		panic("negative depth")
+	}
+	t.root = collapse(depth, t.root)
+}
+
+func collapse(depth int, node Node) Node {
+	if _, ok := node.(*HashNode); ok {
+		return node
+	} else if depth == 0 {
+		return NewHashNode(node.Hash())
+	}
+
+	switch n := node.(type) {
+	case *BranchNode:
+		for i := range n.Children {
+			n.Children[i] = collapse(depth-1, n.Children[i])
+		}
+	case *ExtensionNode:
+		n.next = collapse(depth-1, n.next)
+	case *LeafNode:
+	case *HashNode:
+	default:
+		panic("invalid MPT node type")
+	}
+	return node
+}
diff --git a/pkg/core/mpt/trie_test.go b/pkg/core/mpt/trie_test.go
index 470e0c8e5..d06e08168 100644
--- a/pkg/core/mpt/trie_test.go
+++ b/pkg/core/mpt/trie_test.go
@@ -371,3 +371,76 @@ func TestTrie_PanicInvalidRoot(t *testing.T) {
 	require.Panics(t, func() { _, _ = tr.Get([]byte{1}) })
 	require.Panics(t, func() { _ = tr.Delete([]byte{1}) })
 }
+
+func TestTrie_Collapse(t *testing.T) {
+	t.Run("PanicNegative", func(t *testing.T) {
+		tr := newTestTrie(t)
+		require.Panics(t, func() { tr.Collapse(-1) })
+	})
+	t.Run("Depth=0", func(t *testing.T) {
+		tr := newTestTrie(t)
+		h := tr.root.Hash()
+
+		_, ok := tr.root.(*HashNode)
+		require.False(t, ok)
+
+		tr.Collapse(0)
+		_, ok = tr.root.(*HashNode)
+		require.True(t, ok)
+		require.Equal(t, h, tr.root.Hash())
+	})
+	t.Run("Branch,Depth=1", func(t *testing.T) {
+		b := NewBranchNode()
+		e := NewExtensionNode([]byte{0x01}, NewLeafNode([]byte("value1")))
+		he := e.Hash()
+		b.Children[0] = e
+		hb := b.Hash()
+
+		tr := NewTrie(b, newTestStore())
+		tr.Collapse(1)
+
+		newb, ok := tr.root.(*BranchNode)
+		require.True(t, ok)
+		require.Equal(t, hb, newb.Hash())
+		require.IsType(t, (*HashNode)(nil), b.Children[0])
+		require.Equal(t, he, b.Children[0].Hash())
+	})
+	t.Run("Extension,Depth=1", func(t *testing.T) {
+		l := NewLeafNode([]byte("value"))
+		hl := l.Hash()
+		e := NewExtensionNode([]byte{0x01}, l)
+		h := e.Hash()
+		tr := NewTrie(e, newTestStore())
+		tr.Collapse(1)
+
+		newe, ok := tr.root.(*ExtensionNode)
+		require.True(t, ok)
+		require.Equal(t, h, newe.Hash())
+		require.IsType(t, (*HashNode)(nil), newe.next)
+		require.Equal(t, hl, newe.next.Hash())
+	})
+	t.Run("Leaf", func(t *testing.T) {
+		l := NewLeafNode([]byte("value"))
+		tr := NewTrie(l, newTestStore())
+		tr.Collapse(10)
+		require.Equal(t, NewLeafNode([]byte("value")), tr.root)
+	})
+	t.Run("Hash", func(t *testing.T) {
+		t.Run("Empty", func(t *testing.T) {
+			tr := NewTrie(new(HashNode), newTestStore())
+			require.NotPanics(t, func() { tr.Collapse(1) })
+			hn, ok := tr.root.(*HashNode)
+			require.True(t, ok)
+			require.True(t, hn.IsEmpty())
+		})
+
+		h := random.Uint256()
+		hn := NewHashNode(h)
+		tr := NewTrie(hn, newTestStore())
+		tr.Collapse(10)
+
+		newRoot, ok := tr.root.(*HashNode)
+		require.True(t, ok)
+		require.Equal(t, NewHashNode(h), newRoot)
+	})
+}