From 6ca22027d5991fd8a5119030b7b64b5ccf02bde1 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.

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 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)
+	})
+}