mpt: add the notion of MPT mode

It directly affects the storage format, so it's important. ModeGC is not used
at the moment, but defined for future extensibility.

pkg/core/blockchain.go

@@ -320,7 +320,7 @@ func (bc *Blockchain) init() error {
 		if err != nil {
 			return err
 		}
-		if err := bc.stateRoot.Init(0, bc.config.KeepOnlyLatestState); err != nil {
+		if err := bc.stateRoot.Init(0); err != nil {
 			return fmt.Errorf("can't init MPT: %w", err)
 		}
 		return bc.storeBlock(genesisBlock, nil)
@@ -426,7 +426,7 @@ func (bc *Blockchain) init() error {
 	}
 	bc.blockHeight = bHeight
 	bc.persistedHeight = bHeight
-	if err = bc.stateRoot.Init(bHeight, bc.config.KeepOnlyLatestState); err != nil {
+	if err = bc.stateRoot.Init(bHeight); err != nil {
 		return fmt.Errorf("can't init MPT at height %d: %w", bHeight, err)
 	}
 
@@ -599,7 +599,7 @@ func (bc *Blockchain) jumpToStateInternal(p uint32, stage stateJumpStage) error
 	if err = bc.stateRoot.JumpToState(&state.MPTRoot{
 		Index: p,
 		Root:  block.PrevStateRoot,
-	}, bc.config.KeepOnlyLatestState); err != nil {
+	}); err != nil {
 		return fmt.Errorf("can't perform MPT jump to height %d: %w", p, err)
 	}
 

pkg/core/mpt/batch_test.go

@@ -57,7 +57,7 @@ func testIncompletePut(t *testing.T, ps pairs, n int, tr1, tr2 *Trie) {
 
 	t.Run("test restore", func(t *testing.T) {
 		tr2.Flush()
-		tr3 := NewTrie(NewHashNode(tr2.StateRoot()), false, storage.NewMemCachedStore(tr2.Store))
+		tr3 := NewTrie(NewHashNode(tr2.StateRoot()), ModeAll, storage.NewMemCachedStore(tr2.Store))
 		for _, p := range ps[:n] {
 			val, err := tr3.Get(p[0])
 			if p[1] == nil {
@@ -76,8 +76,8 @@ func testPut(t *testing.T, ps pairs, tr1, tr2 *Trie) {
 
 func TestTrie_PutBatchLeaf(t *testing.T) {
 	prepareLeaf := func(t *testing.T) (*Trie, *Trie) {
-		tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-		tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 		require.NoError(t, tr1.Put([]byte{0}, []byte("value")))
 		require.NoError(t, tr2.Put([]byte{0}, []byte("value")))
 		return tr1, tr2
@@ -118,8 +118,8 @@ func TestTrie_PutBatchLeaf(t *testing.T) {
 
 func TestTrie_PutBatchExtension(t *testing.T) {
 	prepareExtension := func(t *testing.T) (*Trie, *Trie) {
-		tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-		tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 		require.NoError(t, tr1.Put([]byte{1, 2}, []byte("value1")))
 		require.NoError(t, tr2.Put([]byte{1, 2}, []byte("value1")))
 		return tr1, tr2
@@ -170,8 +170,8 @@ func TestTrie_PutBatchExtension(t *testing.T) {
 
 func TestTrie_PutBatchBranch(t *testing.T) {
 	prepareBranch := func(t *testing.T) (*Trie, *Trie) {
-		tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-		tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 		require.NoError(t, tr1.Put([]byte{0x00, 2}, []byte("value1")))
 		require.NoError(t, tr2.Put([]byte{0x00, 2}, []byte("value1")))
 		require.NoError(t, tr1.Put([]byte{0x10, 3}, []byte("value2")))
@@ -201,8 +201,8 @@ func TestTrie_PutBatchBranch(t *testing.T) {
 			require.IsType(t, (*ExtensionNode)(nil), tr1.root)
 		})
 		t.Run("non-empty child is last node", func(t *testing.T) {
-			tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+			tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-			tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+			tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 			require.NoError(t, tr1.Put([]byte{0x00, 2}, []byte("value1")))
 			require.NoError(t, tr2.Put([]byte{0x00, 2}, []byte("value1")))
 			require.NoError(t, tr1.Put([]byte{0x00}, []byte("value2")))
@@ -248,8 +248,8 @@ func TestTrie_PutBatchBranch(t *testing.T) {
 
 func TestTrie_PutBatchHash(t *testing.T) {
 	prepareHash := func(t *testing.T) (*Trie, *Trie) {
-		tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-		tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 		require.NoError(t, tr1.Put([]byte{0x10}, []byte("value1")))
 		require.NoError(t, tr2.Put([]byte{0x10}, []byte("value1")))
 		require.NoError(t, tr1.Put([]byte{0x20}, []byte("value2")))
@@ -283,8 +283,8 @@ func TestTrie_PutBatchHash(t *testing.T) {
 
 func TestTrie_PutBatchEmpty(t *testing.T) {
 	t.Run("good", func(t *testing.T) {
-		tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-		tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 		var ps = pairs{
 			{[]byte{0}, []byte("value0")},
 			{[]byte{1}, []byte("value1")},
@@ -299,15 +299,15 @@ func TestTrie_PutBatchEmpty(t *testing.T) {
 			{[]byte{2}, nil},
 			{[]byte{3}, []byte("replace3")},
 		}
-		tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-		tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 		testIncompletePut(t, ps, 4, tr1, tr2)
 	})
 }
 
 // For the sake of coverage.
 func TestTrie_InvalidNodeType(t *testing.T) {
-	tr := NewTrie(EmptyNode{}, false, newTestStore())
+	tr := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 	var b Batch
 	b.Add([]byte{1}, []byte("value"))
 	tr.root = Node(nil)
@@ -315,8 +315,8 @@ func TestTrie_InvalidNodeType(t *testing.T) {
 }
 
 func TestTrie_PutBatch(t *testing.T) {
-	tr1 := NewTrie(EmptyNode{}, false, newTestStore())
+	tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
-	tr2 := NewTrie(EmptyNode{}, false, newTestStore())
+	tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 	var ps = pairs{
 		{[]byte{1}, []byte{1}},
 		{[]byte{2}, []byte{3}},

pkg/core/mpt/billet.go

@@ -32,19 +32,19 @@ type Billet struct {
 	Store             *storage.MemCachedStore
 
 	root Node
-	refcountEnabled bool
+	mode TrieMode
 }
 
 // NewBillet returns new billet for MPT trie restoring. It accepts a MemCachedStore
 // to decouple storage errors from logic errors so that all storage errors are
 // processed during `store.Persist()` at the caller. This also has the benefit,
 // that every `Put` can be considered an atomic operation.
-func NewBillet(rootHash util.Uint256, enableRefCount bool, prefix storage.KeyPrefix, store *storage.MemCachedStore) *Billet {
+func NewBillet(rootHash util.Uint256, mode TrieMode, prefix storage.KeyPrefix, store *storage.MemCachedStore) *Billet {
 	return &Billet{
 		TempStoragePrefix: prefix,
 		Store:             store,
 		root:              NewHashNode(rootHash),
-		refcountEnabled:   enableRefCount,
+		mode:              mode,
 	}
 }
 
@@ -178,7 +178,7 @@ func (b *Billet) putIntoHash(curr *HashNode, path []byte, val Node) (Node, error
 
 func (b *Billet) incrementRefAndStore(h util.Uint256, bs []byte) {
 	key := makeStorageKey(h)
-	if b.refcountEnabled {
+	if b.mode.RC() {
 		var (
 			err  error
 			data []byte
@@ -337,7 +337,7 @@ func (b *Billet) GetFromStore(h util.Uint256) (Node, error) {
 		return nil, r.Err
 	}
 
-	if b.refcountEnabled {
+	if b.mode.RC() {
 		data = data[:len(data)-4]
 	}
 	n.Node.(flushedNode).setCache(data, h)

pkg/core/mpt/billet_test.go

@@ -32,7 +32,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 			b.Children[5] = NewExtensionNode([]byte{0x01}, NewLeafNode([]byte{0xAB, 0xDE}))
 			path := toNibbles([]byte{0xAC})
 			e := NewExtensionNode(path, NewHashNode(b.Hash()))
-			tr := NewBillet(e.Hash(), true, storage.STTempStorage, newTestStore())
+			tr := NewBillet(e.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 			tr.root = e
 
 			// OK
@@ -61,7 +61,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 			l := NewLeafNode([]byte{0xAB, 0xCD})
 			path := toNibbles([]byte{0xAC})
 			e := NewExtensionNode(path, NewHashNode(l.Hash()))
-			tr := NewBillet(e.Hash(), true, storage.STTempStorage, newTestStore())
+			tr := NewBillet(e.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 			tr.root = e
 
 			// OK
@@ -87,7 +87,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 			h := NewHashNode(util.Uint256{1, 2, 3})
 			path := toNibbles([]byte{0xAC})
 			e := NewExtensionNode(path, h)
-			tr := NewBillet(e.Hash(), true, storage.STTempStorage, newTestStore())
+			tr := NewBillet(e.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 			tr.root = e
 
 			// no-op
@@ -99,7 +99,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 	t.Run("parent is Leaf", func(t *testing.T) {
 		l := NewLeafNode([]byte{0xAB, 0xCD})
 		path := []byte{}
-		tr := NewBillet(l.Hash(), true, storage.STTempStorage, newTestStore())
+		tr := NewBillet(l.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 		tr.root = l
 
 		// Already restored => panic expected
@@ -121,7 +121,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 			b := NewBranchNode()
 			b.Children[5] = NewHashNode(l1.Hash())
 			b.Children[lastChild] = NewHashNode(l2.Hash())
-			tr := NewBillet(b.Hash(), true, storage.STTempStorage, newTestStore())
+			tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 			tr.root = b
 
 			// OK
@@ -152,7 +152,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 			b := NewBranchNode()
 			b.Children[5] = NewHashNode(l1.Hash())
 			b.Children[lastChild] = NewHashNode(l2.Hash())
-			tr := NewBillet(b.Hash(), true, storage.STTempStorage, newTestStore())
+			tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 			tr.root = b
 
 			// OK
@@ -179,7 +179,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 			// two same hashnodes => leaf's refcount expected to be 2 in the end.
 			b.Children[3] = NewHashNode(l.Hash())
 			b.Children[4] = NewHashNode(l.Hash())
-			tr := NewBillet(b.Hash(), true, storage.STTempStorage, newTestStore())
+			tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 			tr.root = b
 
 			// OK
@@ -202,7 +202,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 		b := NewBranchNode()
 		b.Children[3] = NewHashNode(l.Hash())
 		b.Children[4] = NewHashNode(l.Hash())
-		tr := NewBillet(b.Hash(), true, storage.STTempStorage, newTestStore())
+		tr := NewBillet(b.Hash(), ModeLatest, storage.STTempStorage, newTestStore())
 
 		// Should fail, because if it's a hash node with non-empty path, then the node
 		// has already been collapsed.

pkg/core/mpt/compat_test.go

@@ -23,7 +23,7 @@ func prepareMPTCompat() *Trie {
 	b.Children[16] = v2
 	b.Children[15] = NewHashNode(e4.Hash())
 
-	tr := NewTrie(r, true, newTestStore())
+	tr := NewTrie(r, ModeLatest, newTestStore())
 	tr.putToStore(r)
 	tr.putToStore(b)
 	tr.putToStore(e1)
@@ -132,7 +132,7 @@ func TestCompatibility(t *testing.T) {
 		b.Children[0] = e1
 		b.Children[15] = NewHashNode(e4.Hash())
 
-		tr := NewTrie(NewHashNode(r.Hash()), false, newTestStore())
+		tr := NewTrie(NewHashNode(r.Hash()), ModeAll, newTestStore())
 		tr.putToStore(r)
 		tr.putToStore(b)
 		tr.putToStore(e1)
@@ -152,7 +152,7 @@ func TestCompatibility(t *testing.T) {
 		tr.testHas(t, []byte{0xac, 0x02}, []byte{0xab, 0xcd})
 		tr.Flush()
 
-		tr2 := NewTrie(NewHashNode(tr.root.Hash()), false, tr.Store)
+		tr2 := NewTrie(NewHashNode(tr.root.Hash()), ModeAll, tr.Store)
 		tr2.testHas(t, []byte{0xac, 0x02}, []byte{0xab, 0xcd})
 	})
 
@@ -189,7 +189,7 @@ func TestCompatibility(t *testing.T) {
 		b.Children[16] = v2
 		b.Children[15] = NewHashNode(e4.Hash())
 
-		tr := NewTrie(NewHashNode(r.Hash()), true, mainTrie.Store)
+		tr := NewTrie(NewHashNode(r.Hash()), ModeLatest, mainTrie.Store)
 		require.Equal(t, r.Hash(), tr.root.Hash())
 
 		// Tail bytes contain reference counter thus check for prefix.
@@ -352,7 +352,7 @@ func TestCompatibility(t *testing.T) {
 }
 
 func copyTrie(t *Trie) *Trie {
-	return NewTrie(NewHashNode(t.root.Hash()), t.refcountEnabled, t.Store)
+	return NewTrie(NewHashNode(t.root.Hash()), t.mode, t.Store)
 }
 
 func checkBatchSize(t *testing.T, tr *Trie, n int) {
@@ -372,7 +372,7 @@ func testGetProof(t *testing.T, tr *Trie, key []byte, size int) [][]byte {
 }
 
 func newFilledTrie(t *testing.T, args ...[]byte) *Trie {
-	tr := NewTrie(nil, true, newTestStore())
+	tr := NewTrie(nil, ModeLatest, newTestStore())
 	for i := 0; i < len(args); i += 2 {
 		require.NoError(t, tr.Put(args[i], args[i+1]))
 	}
@@ -381,7 +381,7 @@ func newFilledTrie(t *testing.T, args ...[]byte) *Trie {
 
 func TestCompatibility_Find(t *testing.T) {
 	check := func(t *testing.T, from []byte, expectedResLen int) {
-		tr := NewTrie(nil, false, newTestStore())
+		tr := NewTrie(nil, ModeAll, newTestStore())
 		require.NoError(t, tr.Put([]byte("aa"), []byte("02")))
 		require.NoError(t, tr.Put([]byte("aa10"), []byte("03")))
 		require.NoError(t, tr.Put([]byte("aa50"), []byte("04")))
@@ -407,7 +407,7 @@ func TestCompatibility_Find(t *testing.T) {
 		check(t, []byte{}, 2) // without `from` key
 	})
 	t.Run("TestFindStatesIssue652", func(t *testing.T) {
-		tr := NewTrie(nil, false, newTestStore())
+		tr := NewTrie(nil, ModeAll, newTestStore())
 		// root is an extension node with key=abc; next=branch
 		require.NoError(t, tr.Put([]byte("abc1"), []byte("01")))
 		require.NoError(t, tr.Put([]byte("abc3"), []byte("02")))

pkg/core/mpt/node_test.go

@@ -93,7 +93,7 @@ func TestNode_Serializable(t *testing.T) {
 
 // https://github.com/neo-project/neo/blob/neox-2.x/neo.UnitTests/UT_MPTTrie.cs#L198
 func TestJSONSharp(t *testing.T) {
-	tr := NewTrie(nil, false, newTestStore())
+	tr := NewTrie(nil, ModeAll, newTestStore())
 	require.NoError(t, tr.Put([]byte{0xac, 0x11}, []byte{0xac, 0x11}))
 	require.NoError(t, tr.Put([]byte{0xac, 0x22}, []byte{0xac, 0x22}))
 	require.NoError(t, tr.Put([]byte{0xac}, []byte{0xac}))

pkg/core/mpt/proof.go

@@ -66,7 +66,7 @@ func (t *Trie) getProof(curr Node, path []byte, proofs *[][]byte) (Node, error)
 // It also returns value for the key.
 func VerifyProof(rh util.Uint256, key []byte, proofs [][]byte) ([]byte, bool) {
 	path := toNibbles(key)
-	tr := NewTrie(NewHashNode(rh), false, storage.NewMemCachedStore(storage.NewMemoryStore()))
+	tr := NewTrie(NewHashNode(rh), ModeAll, storage.NewMemCachedStore(storage.NewMemoryStore()))
 	for i := range proofs {
 		h := hash.DoubleSha256(proofs[i])
 		// no errors in Put to memory store

pkg/core/mpt/proof_test.go

@@ -15,7 +15,7 @@ func newProofTrie(t *testing.T, missingHashNode bool) *Trie {
 	b.Children[4] = NewHashNode(e.Hash())
 	b.Children[5] = e2
 
-	tr := NewTrie(b, false, newTestStore())
+	tr := NewTrie(b, ModeAll, newTestStore())
 	require.NoError(t, tr.Put([]byte{0x12, 0x31}, []byte("value1")))
 	require.NoError(t, tr.Put([]byte{0x12, 0x32}, []byte("value2")))
 	tr.putToStore(l)

pkg/core/mpt/trie.go

@@ -12,12 +12,28 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/util/slice"
 )
 
+// TrieMode is the storage mode of trie, it affects the DB scheme.
+type TrieMode byte
+
+// TrieMode is the storage mode of trie.
+const (
+	// ModeAll is used to store everything.
+	ModeAll TrieMode = 0
+	// ModeLatest is used to only store the latest root.
+	ModeLatest TrieMode = 0x01
+	// ModeGCFlag is a flag for GC.
+	ModeGCFlag TrieMode = 0x02
+	// ModeGC is used to store a set of roots with GC possible, it combines
+	// GCFlag and Latest (because it needs RC, but it has GC enabled).
+	ModeGC TrieMode = 0x03
+)
+
 // Trie is an MPT trie storing all key-value pairs.
 type Trie struct {
 	Store *storage.MemCachedStore
 
 	root     Node
-	refcountEnabled bool
+	mode     TrieMode
 	refcount map[util.Uint256]*cachedNode
 }
 
@@ -30,10 +46,20 @@ type cachedNode struct {
 // ErrNotFound is returned when requested trie item is missing.
 var ErrNotFound = errors.New("item not found")
 
+// RC returns true when reference counting is enabled.
+func (m TrieMode) RC() bool {
+	return m&ModeLatest != 0
+}
+
+// GC returns true when garbage collection is enabled.
+func (m TrieMode) GC() bool {
+	return m&ModeGCFlag != 0
+}
+
 // NewTrie returns new MPT trie. It accepts a MemCachedStore to decouple storage errors from logic errors
 // so that all storage errors are processed during `store.Persist()` at the caller.
 // This also has the benefit, that every `Put` can be considered an atomic operation.
-func NewTrie(root Node, enableRefCount bool, store *storage.MemCachedStore) *Trie {
+func NewTrie(root Node, mode TrieMode, store *storage.MemCachedStore) *Trie {
 	if root == nil {
 		root = EmptyNode{}
 	}
@@ -42,7 +68,7 @@ func NewTrie(root Node, enableRefCount bool, store *storage.MemCachedStore) *Tri
 		Store: store,
 		root:  root,
 
-		refcountEnabled: enableRefCount,
+		mode:     mode,
 		refcount: make(map[util.Uint256]*cachedNode),
 	}
 }
@@ -386,7 +412,7 @@ func (t *Trie) Flush() {
 			if node.bytes == nil {
 				panic("item not in trie")
 			}
-			if t.refcountEnabled {
+			if t.mode.RC() {
 				node.initial = t.updateRefCount(h)
 				if node.initial == 0 {
 					delete(t.refcount, h)
@@ -403,7 +429,7 @@ func (t *Trie) Flush() {
 
 // updateRefCount should be called only when refcounting is enabled.
 func (t *Trie) updateRefCount(h util.Uint256) int32 {
-	if !t.refcountEnabled {
+	if !t.mode.RC() {
 		panic("`updateRefCount` is called, but GC is disabled")
 	}
 	var data []byte
@@ -478,7 +504,7 @@ func (t *Trie) getFromStore(h util.Uint256) (Node, error) {
 		return nil, r.Err
 	}
 
-	if t.refcountEnabled {
+	if t.mode.RC() {
 		data = data[:len(data)-4]
 		node := t.refcount[h]
 		if node != nil {
@@ -566,7 +592,7 @@ func (t *Trie) Find(prefix, from []byte, max int) ([]storage.KeyValue, error) {
 		res   []storage.KeyValue
 		count int
 	)
-	b := NewBillet(t.root.Hash(), false, 0, t.Store)
+	b := NewBillet(t.root.Hash(), t.mode, 0, t.Store)
 	process := func(pathToNode []byte, node Node, _ []byte) bool {
 		if leaf, ok := node.(*LeafNode); ok {
 			if from == nil || !bytes.Equal(pathToNode, from) { // (*Billet).traverse includes `from` path into result if so. Need to filter out manually.

pkg/core/mpt/trie_test.go

@@ -29,7 +29,7 @@ func newTestTrie(t *testing.T) *Trie {
 	b.Children[10] = NewExtensionNode([]byte{0x0e}, h)
 
 	e := NewExtensionNode(toNibbles([]byte{0xAC}), b)
-	tr := NewTrie(e, false, newTestStore())
+	tr := NewTrie(e, ModeAll, newTestStore())
 
 	tr.putToStore(e)
 	tr.putToStore(b)
@@ -46,7 +46,7 @@ func newTestTrie(t *testing.T) *Trie {
 }
 
 func testTrieRefcount(t *testing.T, key1, key2 []byte) {
-	tr := NewTrie(nil, true, storage.NewMemCachedStore(storage.NewMemoryStore()))
+	tr := NewTrie(nil, ModeLatest, storage.NewMemCachedStore(storage.NewMemoryStore()))
 	require.NoError(t, tr.Put(key1, []byte{1}))
 	tr.Flush()
 	require.NoError(t, tr.Put(key2, []byte{1}))
@@ -89,7 +89,7 @@ func TestTrie_PutIntoBranchNode(t *testing.T) {
 		l := NewLeafNode([]byte{0x8})
 		b.Children[0x7] = NewHashNode(l.Hash())
 		b.Children[0x8] = NewHashNode(random.Uint256())
-		tr := NewTrie(b, false, newTestStore())
+		tr := NewTrie(b, ModeAll, newTestStore())
 
 		// empty hash node child
 		require.NoError(t, tr.Put([]byte{0x66}, value))
@@ -119,7 +119,7 @@ func TestTrie_PutIntoExtensionNode(t *testing.T) {
 		l := NewLeafNode([]byte{0x11})
 		key := []byte{0x12}
 		e := NewExtensionNode(toNibbles(key), NewHashNode(l.Hash()))
-		tr := NewTrie(e, false, newTestStore())
+		tr := NewTrie(e, ModeAll, newTestStore())
 
 		// missing hash
 		require.Error(t, tr.Put(key, value))
@@ -145,7 +145,7 @@ func TestTrie_PutIntoHashNode(t *testing.T) {
 		e := NewExtensionNode([]byte{0x02}, l)
 		b.Children[1] = NewHashNode(e.Hash())
 		b.Children[9] = NewHashNode(random.Uint256())
-		tr := NewTrie(b, false, newTestStore())
+		tr := NewTrie(b, ModeAll, newTestStore())
 
 		tr.putToStore(e)
 
@@ -174,7 +174,7 @@ func TestTrie_PutIntoHashNode(t *testing.T) {
 func TestTrie_Put(t *testing.T) {
 	trExp := newTestTrie(t)
 
-	trAct := NewTrie(nil, false, newTestStore())
+	trAct := NewTrie(nil, ModeAll, newTestStore())
 	require.NoError(t, trAct.Put([]byte{0xAC, 0x01}, []byte{0xAB, 0xCD}))
 	require.NoError(t, trAct.Put([]byte{0xAC, 0x13}, []byte{}))
 	require.NoError(t, trAct.Put([]byte{0xAC, 0x99}, []byte{0x22, 0x22}))
@@ -186,7 +186,7 @@ func TestTrie_Put(t *testing.T) {
 }
 
 func TestTrie_PutInvalid(t *testing.T) {
-	tr := NewTrie(nil, false, newTestStore())
+	tr := NewTrie(nil, ModeAll, newTestStore())
 	key, value := []byte("key"), []byte("value")
 
 	// empty key
@@ -204,7 +204,7 @@ func TestTrie_PutInvalid(t *testing.T) {
 }
 
 func TestTrie_BigPut(t *testing.T) {
-	tr := NewTrie(nil, false, newTestStore())
+	tr := NewTrie(nil, ModeAll, newTestStore())
 	items := []struct{ k, v string }{
 		{"item with long key", "value1"},
 		{"item with matching prefix", "value2"},
@@ -244,7 +244,7 @@ func (tr *Trie) putToStore(n Node) {
 	if n.Type() == HashT {
 		panic("can't put hash node in trie")
 	}
-	if tr.refcountEnabled {
+	if tr.mode.RC() {
 		tr.refcount[n.Hash()] = &cachedNode{
 			bytes:    n.Bytes(),
 			refcount: 1,
@@ -298,7 +298,7 @@ func TestTrie_Get(t *testing.T) {
 	})
 	t.Run("UnfoldRoot", func(t *testing.T) {
 		tr := newTestTrie(t)
-		single := NewTrie(NewHashNode(tr.root.Hash()), false, tr.Store)
+		single := NewTrie(NewHashNode(tr.root.Hash()), ModeAll, tr.Store)
 		single.testHas(t, []byte{0xAC}, nil)
 		single.testHas(t, []byte{0xAC, 0x01}, []byte{0xAB, 0xCD})
 		single.testHas(t, []byte{0xAC, 0x99}, []byte{0x22, 0x22})
@@ -313,13 +313,13 @@ func TestTrie_Flush(t *testing.T) {
 		"key2": []byte("value2"),
 	}
 
-	tr := NewTrie(nil, false, newTestStore())
+	tr := NewTrie(nil, ModeAll, newTestStore())
 	for k, v := range pairs {
 		require.NoError(t, tr.Put([]byte(k), v))
 	}
 
 	tr.Flush()
-	tr = NewTrie(NewHashNode(tr.StateRoot()), false, tr.Store)
+	tr = NewTrie(NewHashNode(tr.StateRoot()), ModeAll, tr.Store)
 	for k, v := range pairs {
 		actual, err := tr.Get([]byte(k))
 		require.NoError(t, err)
@@ -337,10 +337,14 @@ func TestTrie_Delete(t *testing.T) {
 }
 
 func testTrieDelete(t *testing.T, enableGC bool) {
+	var mode TrieMode
+	if enableGC {
+		mode = ModeLatest
+	}
 	t.Run("Hash", func(t *testing.T) {
 		t.Run("FromStore", func(t *testing.T) {
 			l := NewLeafNode([]byte{0x12})
-			tr := NewTrie(NewHashNode(l.Hash()), enableGC, newTestStore())
+			tr := NewTrie(NewHashNode(l.Hash()), mode, newTestStore())
 			t.Run("NotInStore", func(t *testing.T) {
 				require.Error(t, tr.Delete([]byte{}))
 			})
@@ -352,7 +356,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 		})
 
 		t.Run("Empty", func(t *testing.T) {
-			tr := NewTrie(nil, enableGC, newTestStore())
+			tr := NewTrie(nil, mode, newTestStore())
 			require.NoError(t, tr.Delete([]byte{}))
 		})
 	})
@@ -360,7 +364,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 	t.Run("Leaf", func(t *testing.T) {
 		check := func(t *testing.T, value []byte) {
 			l := NewLeafNode(value)
-			tr := NewTrie(l, enableGC, newTestStore())
+			tr := NewTrie(l, mode, newTestStore())
 			t.Run("NonExistentKey", func(t *testing.T) {
 				require.NoError(t, tr.Delete([]byte{0x12}))
 				tr.testHas(t, []byte{}, value)
@@ -381,7 +385,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 			check := func(t *testing.T, value []byte) {
 				l := NewLeafNode(value)
 				e := NewExtensionNode([]byte{0x0A, 0x0B}, l)
-				tr := NewTrie(e, enableGC, newTestStore())
+				tr := NewTrie(e, mode, newTestStore())
 
 				t.Run("NonExistentKey", func(t *testing.T) {
 					require.NoError(t, tr.Delete([]byte{}))
@@ -405,7 +409,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 				b.Children[0] = NewExtensionNode([]byte{0x01}, NewLeafNode(value))
 				b.Children[6] = NewExtensionNode([]byte{0x07}, NewLeafNode([]byte{0x56, 0x78}))
 				e := NewExtensionNode([]byte{0x01, 0x02}, b)
-				tr := NewTrie(e, enableGC, newTestStore())
+				tr := NewTrie(e, mode, newTestStore())
 
 				h := e.Hash()
 				require.NoError(t, tr.Delete([]byte{0x12, 0x01}))
@@ -432,7 +436,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 				b.Children[lastChild] = NewLeafNode([]byte{0x12})
 				b.Children[0] = NewExtensionNode([]byte{0x01}, NewLeafNode([]byte{0x34}))
 				b.Children[1] = NewExtensionNode([]byte{0x06}, NewLeafNode(value))
-				tr := NewTrie(b, enableGC, newTestStore())
+				tr := NewTrie(b, mode, newTestStore())
 				require.NoError(t, tr.Delete([]byte{0x16}))
 				tr.testHas(t, []byte{}, []byte{0x12})
 				tr.testHas(t, []byte{0x01}, []byte{0x34})
@@ -454,7 +458,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 						l := NewLeafNode([]byte{0x34})
 						e := NewExtensionNode([]byte{0x06}, l)
 						b.Children[5] = NewHashNode(e.Hash())
-						tr := NewTrie(b, enableGC, newTestStore())
+						tr := NewTrie(b, mode, newTestStore())
 						tr.putToStore(l)
 						tr.putToStore(e)
 						require.NoError(t, tr.Delete([]byte{}))
@@ -478,7 +482,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 							b.Children[3] = NewExtensionNode([]byte{4}, NewLeafNode(value))
 							b.Children[lastChild] = NewHashNode(h)
 
-							tr := NewTrie(NewExtensionNode([]byte{1, 2}, b), enableGC, newTestStore())
+							tr := NewTrie(NewExtensionNode([]byte{1, 2}, b), mode, newTestStore())
 							tr.putToStore(ch)
 
 							require.NoError(t, tr.Delete([]byte{0x12, 0x34}))
@@ -505,7 +509,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 						b := NewBranchNode()
 						b.Children[lastChild] = NewLeafNode(value)
 						b.Children[5] = c
-						tr := NewTrie(b, enableGC, newTestStore())
+						tr := NewTrie(b, mode, newTestStore())
 
 						require.NoError(t, tr.Delete([]byte{}))
 						tr.testHas(t, []byte{}, nil)
@@ -529,7 +533,7 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 					l := NewLeafNode(value)
 					e := NewExtensionNode([]byte{0x06}, l)
 					b.Children[5] = NewHashNode(e.Hash())
-					tr := NewTrie(b, enableGC, newTestStore())
+					tr := NewTrie(b, mode, newTestStore())
 					tr.putToStore(l)
 					tr.putToStore(e)
 					require.NoError(t, tr.Delete([]byte{0x56}))
@@ -579,7 +583,7 @@ func TestTrie_Collapse(t *testing.T) {
 		b.Children[0] = e
 		hb := b.Hash()
 
-		tr := NewTrie(b, false, newTestStore())
+		tr := NewTrie(b, ModeAll, newTestStore())
 		tr.Collapse(1)
 
 		newb, ok := tr.root.(*BranchNode)
@@ -593,7 +597,7 @@ func TestTrie_Collapse(t *testing.T) {
 		hl := l.Hash()
 		e := NewExtensionNode([]byte{0x01}, l)
 		h := e.Hash()
-		tr := NewTrie(e, false, newTestStore())
+		tr := NewTrie(e, ModeAll, newTestStore())
 		tr.Collapse(1)
 
 		newe, ok := tr.root.(*ExtensionNode)
@@ -604,19 +608,19 @@ func TestTrie_Collapse(t *testing.T) {
 	})
 	t.Run("Leaf", func(t *testing.T) {
 		l := NewLeafNode([]byte("value"))
-		tr := NewTrie(l, false, newTestStore())
+		tr := NewTrie(l, ModeAll, newTestStore())
 		tr.Collapse(10)
 		require.Equal(t, NewLeafNode([]byte("value")), tr.root)
 	})
 	t.Run("Empty Leaf", func(t *testing.T) {
 		l := NewLeafNode([]byte{})
-		tr := NewTrie(l, false, newTestStore())
+		tr := NewTrie(l, ModeAll, newTestStore())
 		tr.Collapse(10)
 		require.Equal(t, NewLeafNode([]byte{}), tr.root)
 	})
 	t.Run("Hash", func(t *testing.T) {
 		t.Run("EmptyNode", func(t *testing.T) {
-			tr := NewTrie(EmptyNode{}, false, newTestStore())
+			tr := NewTrie(EmptyNode{}, ModeAll, newTestStore())
 			require.NotPanics(t, func() { tr.Collapse(1) })
 			_, ok := tr.root.(EmptyNode)
 			require.True(t, ok)
@@ -624,7 +628,7 @@ func TestTrie_Collapse(t *testing.T) {
 
 		h := random.Uint256()
 		hn := NewHashNode(h)
-		tr := NewTrie(hn, false, newTestStore())
+		tr := NewTrie(hn, ModeAll, newTestStore())
 		tr.Collapse(10)
 
 		newRoot, ok := tr.root.(*HashNode)

pkg/core/stateroot/module.go

@@ -28,6 +28,7 @@ type (
 		Store    *storage.MemCachedStore
 		network  netmode.Magic
 		srInHead bool
+		mode     mpt.TrieMode
 		mpt      *mpt.Trie
 		verifier VerifierFunc
 		log      *zap.Logger
@@ -52,9 +53,14 @@ type (
 
 // NewModule returns new instance of stateroot module.
 func NewModule(cfg config.ProtocolConfiguration, verif VerifierFunc, log *zap.Logger, s *storage.MemCachedStore) *Module {
+	var mode mpt.TrieMode
+	if cfg.KeepOnlyLatestState {
+		mode |= mpt.ModeLatest
+	}
 	return &Module{
 		network:  cfg.Magic,
 		srInHead: cfg.StateRootInHeader,
+		mode:     mode,
 		verifier: verif,
 		log:      log,
 		Store:    s,
@@ -63,7 +69,7 @@ func NewModule(cfg config.ProtocolConfiguration, verif VerifierFunc, log *zap.Lo
 
 // GetState returns value at the specified key fom the MPT with the specified root.
 func (s *Module) GetState(root util.Uint256, key []byte) ([]byte, error) {
-	tr := mpt.NewTrie(mpt.NewHashNode(root), false, storage.NewMemCachedStore(s.Store))
+	tr := mpt.NewTrie(mpt.NewHashNode(root), s.mode, storage.NewMemCachedStore(s.Store))
 	return tr.Get(key)
 }
 
@@ -73,13 +79,13 @@ func (s *Module) GetState(root util.Uint256, key []byte) ([]byte, error) {
 // item with key equals to prefix is included into result; if empty `start` specified,
 // then item with key equals to prefix is not included into result.
 func (s *Module) FindStates(root util.Uint256, prefix, start []byte, max int) ([]storage.KeyValue, error) {
-	tr := mpt.NewTrie(mpt.NewHashNode(root), false, storage.NewMemCachedStore(s.Store))
+	tr := mpt.NewTrie(mpt.NewHashNode(root), s.mode, storage.NewMemCachedStore(s.Store))
 	return tr.Find(prefix, start, max)
 }
 
 // GetStateProof returns proof of having key in the MPT with the specified root.
 func (s *Module) GetStateProof(root util.Uint256, key []byte) ([][]byte, error) {
-	tr := mpt.NewTrie(mpt.NewHashNode(root), false, storage.NewMemCachedStore(s.Store))
+	tr := mpt.NewTrie(mpt.NewHashNode(root), s.mode, storage.NewMemCachedStore(s.Store))
 	return tr.GetProof(key)
 }
 
@@ -104,14 +110,14 @@ func (s *Module) CurrentValidatedHeight() uint32 {
 }
 
 // Init initializes state root module at the given height.
-func (s *Module) Init(height uint32, enableRefCount bool) error {
+func (s *Module) Init(height uint32) error {
 	data, err := s.Store.Get([]byte{byte(storage.DataMPT), prefixValidated})
 	if err == nil {
 		s.validatedHeight.Store(binary.LittleEndian.Uint32(data))
 	}
 
 	if height == 0 {
-		s.mpt = mpt.NewTrie(nil, enableRefCount, s.Store)
+		s.mpt = mpt.NewTrie(nil, s.mode, s.Store)
 		s.currentLocal.Store(util.Uint256{})
 		return nil
 	}
@@ -121,7 +127,7 @@ func (s *Module) Init(height uint32, enableRefCount bool) error {
 	}
 	s.currentLocal.Store(r.Root)
 	s.localHeight.Store(r.Index)
-	s.mpt = mpt.NewTrie(mpt.NewHashNode(r.Root), enableRefCount, s.Store)
+	s.mpt = mpt.NewTrie(mpt.NewHashNode(r.Root), s.mode, s.Store)
 	return nil
 }
 
@@ -157,7 +163,7 @@ func (s *Module) CleanStorage() error {
 }
 
 // JumpToState performs jump to the state specified by given stateroot index.
-func (s *Module) JumpToState(sr *state.MPTRoot, enableRefCount bool) error {
+func (s *Module) JumpToState(sr *state.MPTRoot) error {
 	if err := s.addLocalStateRoot(s.Store, sr); err != nil {
 		return fmt.Errorf("failed to store local state root: %w", err)
 	}
@@ -171,7 +177,7 @@ func (s *Module) JumpToState(sr *state.MPTRoot, enableRefCount bool) error {
 
 	s.currentLocal.Store(sr.Root)
 	s.localHeight.Store(sr.Index)
-	s.mpt = mpt.NewTrie(mpt.NewHashNode(sr.Root), enableRefCount, s.Store)
+	s.mpt = mpt.NewTrie(mpt.NewHashNode(sr.Root), s.mode, s.Store)
 	return nil
 }
 

pkg/core/statesync/module.go

@@ -221,7 +221,11 @@ func (s *Module) defineSyncStage() error {
 		if err != nil {
 			return fmt.Errorf("failed to get header to initialize MPT billet: %w", err)
 		}
-		s.billet = mpt.NewBillet(header.PrevStateRoot, s.bc.GetConfig().KeepOnlyLatestState,
+		var mode mpt.TrieMode
+		if s.bc.GetConfig().KeepOnlyLatestState {
+			mode |= mpt.ModeLatest
+		}
+		s.billet = mpt.NewBillet(header.PrevStateRoot, mode,
 			TemporaryPrefix(s.dao.Version.StoragePrefix), s.dao.Store)
 		s.log.Info("MPT billet initialized",
 			zap.Uint32("height", s.syncPoint),
@@ -494,7 +498,11 @@ func (s *Module) Traverse(root util.Uint256, process func(node mpt.Node, nodeByt
 	s.lock.RLock()
 	defer s.lock.RUnlock()
 
-	b := mpt.NewBillet(root, s.bc.GetConfig().KeepOnlyLatestState, 0, storage.NewMemCachedStore(s.dao.Store))
+	var mode mpt.TrieMode
+	if s.bc.GetConfig().KeepOnlyLatestState {
+		mode |= mpt.ModeLatest
+	}
+	b := mpt.NewBillet(root, mode, 0, storage.NewMemCachedStore(s.dao.Store))
 	return b.Traverse(func(pathToNode []byte, node mpt.Node, nodeBytes []byte) bool {
 		return process(node, nodeBytes)
 	}, false)

pkg/core/statesync/module_test.go

@@ -15,7 +15,7 @@ import (
 
 func TestModule_PR2019_discussion_r689629704(t *testing.T) {
 	expectedStorage := storage.NewMemCachedStore(storage.NewMemoryStore())
-	tr := mpt.NewTrie(nil, true, expectedStorage)
+	tr := mpt.NewTrie(nil, mpt.ModeLatest, expectedStorage)
 	require.NoError(t, tr.Put([]byte{0x03}, []byte("leaf1")))
 	require.NoError(t, tr.Put([]byte{0x01, 0xab, 0x02}, []byte("leaf2")))
 	require.NoError(t, tr.Put([]byte{0x01, 0xab, 0x04}, []byte("leaf3")))
@@ -57,7 +57,7 @@ func TestModule_PR2019_discussion_r689629704(t *testing.T) {
 		dao:          dao.NewSimple(actualStorage, true, false),
 		mptpool:      NewPool(),
 	}
-	stateSync.billet = mpt.NewBillet(sr, true,
+	stateSync.billet = mpt.NewBillet(sr, mpt.ModeLatest,
 		TemporaryPrefix(stateSync.dao.Version.StoragePrefix), actualStorage)
 	stateSync.mptpool.Add(sr, []byte{})