neo-go/pkg/core/mpt/trie_test.go

package mpt

import (
	"testing"

	"github.com/nspcc-dev/neo-go/internal/random"
	"github.com/nspcc-dev/neo-go/pkg/core/storage"
	"github.com/stretchr/testify/require"
)

func newTestStore() *storage.MemCachedStore {
	return storage.NewMemCachedStore(storage.NewMemoryStore())
}

func newTestTrie(t *testing.T) *Trie {
	b := NewBranchNode()

	l1 := NewLeafNode([]byte{0xAB, 0xCD})
	b.Children[0] = NewExtensionNode([]byte{0x01}, l1)

	l2 := NewLeafNode([]byte{0x22, 0x22})
	b.Children[9] = NewExtensionNode([]byte{0x09}, l2)

	v := NewLeafNode([]byte("hello"))
	h := NewHashNode(v.Hash())
	b.Children[10] = NewExtensionNode([]byte{0x0e}, h)

	e := NewExtensionNode(toNibbles([]byte{0xAC}), b)
	tr := NewTrie(e, false, newTestStore())

	tr.putToStore(e)
	tr.putToStore(b)
	tr.putToStore(l1)
	tr.putToStore(l2)
	tr.putToStore(v)
	tr.putToStore(b.Children[0])
	tr.putToStore(b.Children[9])
	tr.putToStore(b.Children[10])

	return tr
}

func testTrieRefcount(t *testing.T, key1, key2 []byte) {
	tr := NewTrie(nil, true, storage.NewMemCachedStore(storage.NewMemoryStore()))
	require.NoError(t, tr.Put(key1, []byte{1}))
	tr.Flush()
	require.NoError(t, tr.Put(key2, []byte{1}))
	tr.Flush()
	tr.testHas(t, key1, []byte{1})
	tr.testHas(t, key2, []byte{1})

	// remove first, keep second
	require.NoError(t, tr.Delete(key1))
	tr.Flush()
	tr.testHas(t, key1, nil)
	tr.testHas(t, key2, []byte{1})

	// no-op
	require.NoError(t, tr.Put(key1, []byte{1}))
	require.NoError(t, tr.Delete(key1))
	tr.Flush()
	tr.testHas(t, key1, nil)
	tr.testHas(t, key2, []byte{1})

	// error on delete, refcount should not be updated
	require.Error(t, tr.Delete(key1))
	tr.Flush()
	tr.testHas(t, key1, nil)
	tr.testHas(t, key2, []byte{1})
}

func TestTrie_Refcount(t *testing.T) {
	t.Run("Leaf", func(t *testing.T) {
		testTrieRefcount(t, []byte{0x11}, []byte{0x12})
	})
	t.Run("Extension", func(t *testing.T) {
		testTrieRefcount(t, []byte{0x10, 11}, []byte{0x11, 12})
	})
}

func TestTrie_PutIntoBranchNode(t *testing.T) {
	b := NewBranchNode()
	l := NewLeafNode([]byte{0x8})
	b.Children[0x7] = NewHashNode(l.Hash())
	b.Children[0x8] = NewHashNode(random.Uint256())
	tr := NewTrie(b, false, newTestStore())

	// next
	require.NoError(t, tr.Put([]byte{}, []byte{0x12, 0x34}))
	tr.testHas(t, []byte{}, []byte{0x12, 0x34})

	// empty hash node child
	require.NoError(t, tr.Put([]byte{0x66}, []byte{0x56}))
	tr.testHas(t, []byte{0x66}, []byte{0x56})
	require.True(t, isValid(tr.root))

	// missing hash
	require.Error(t, tr.Put([]byte{0x70}, []byte{0x42}))
	require.True(t, isValid(tr.root))

	// hash is in store
	tr.putToStore(l)
	require.NoError(t, tr.Put([]byte{0x70}, []byte{0x42}))
	require.True(t, isValid(tr.root))
}

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())

	// missing hash
	require.Error(t, tr.Put(key, []byte{0x42}))

	tr.putToStore(l)
	require.NoError(t, tr.Put(key, []byte{0x42}))
	tr.testHas(t, key, []byte{0x42})
	require.True(t, isValid(tr.root))
}

func TestTrie_PutIntoHashNode(t *testing.T) {
	b := NewBranchNode()
	l := NewLeafNode(random.Bytes(5))
	e := NewExtensionNode([]byte{0x02}, l)
	b.Children[1] = NewHashNode(e.Hash())
	b.Children[9] = NewHashNode(random.Uint256())
	tr := NewTrie(b, false, newTestStore())

	tr.putToStore(e)

	t.Run("MissingLeafHash", func(t *testing.T) {
		_, err := tr.Get([]byte{0x12})
		require.Error(t, err)
	})

	tr.putToStore(l)

	val := random.Bytes(3)
	require.NoError(t, tr.Put([]byte{0x12, 0x34}, val))
	tr.testHas(t, []byte{0x12, 0x34}, val)
	tr.testHas(t, []byte{0x12}, l.value)
	require.True(t, isValid(tr.root))
}

func TestTrie_Put(t *testing.T) {
	trExp := newTestTrie(t)

	trAct := NewTrie(nil, false, newTestStore())
	require.NoError(t, trAct.Put([]byte{0xAC, 0x01}, []byte{0xAB, 0xCD}))
	require.NoError(t, trAct.Put([]byte{0xAC, 0x99}, []byte{0x22, 0x22}))
	require.NoError(t, trAct.Put([]byte{0xAC, 0xAE}, []byte("hello")))

	// Note: the exact tries differ because of ("acae":"hello") node is stored as Hash node in test trie.
	require.Equal(t, trExp.root.Hash(), trAct.root.Hash())
	require.True(t, isValid(trAct.root))
}

func TestTrie_PutInvalid(t *testing.T) {
	tr := NewTrie(nil, false, newTestStore())
	key, value := []byte("key"), []byte("value")

	// big key
	require.Error(t, tr.Put(make([]byte, MaxKeyLength+1), value))

	// big value
	require.Error(t, tr.Put(key, make([]byte, MaxValueLength+1)))

	// this is ok though
	require.NoError(t, tr.Put(key, value))
	tr.testHas(t, key, value)
}

func TestTrie_BigPut(t *testing.T) {
	tr := NewTrie(nil, false, newTestStore())
	items := []struct{ k, v string }{
		{"item with long key", "value1"},
		{"item with matching prefix", "value2"},
		{"another prefix", "value3"},
		{"another prefix 2", "value4"},
		{"another ", "value5"},
	}

	for i := range items {
		require.NoError(t, tr.Put([]byte(items[i].k), []byte(items[i].v)))
	}

	for i := range items {
		tr.testHas(t, []byte(items[i].k), []byte(items[i].v))
	}

	t.Run("Rewrite", func(t *testing.T) {
		k, v := []byte(items[0].k), []byte{0x01, 0x23}
		require.NoError(t, tr.Put(k, v))
		tr.testHas(t, k, v)
	})

	t.Run("Remove", func(t *testing.T) {
		k := []byte(items[1].k)
		require.NoError(t, tr.Put(k, []byte{}))
		tr.testHas(t, k, nil)
	})
}

func (tr *Trie) putToStore(n Node) {
	if n.Type() == HashT {
		panic("can't put hash node in trie")
	}
	if tr.refcountEnabled {
		tr.refcount[n.Hash()] = &cachedNode{
			bytes:    n.Bytes(),
			refcount: 1,
		}
		tr.updateRefCount(n.Hash())
	} else {
		_ = tr.Store.Put(makeStorageKey(n.Hash().BytesBE()), n.Bytes())
	}
}

func (tr *Trie) testHas(t *testing.T, key, value []byte) {
	v, err := tr.Get(key)
	if value == nil {
		require.Error(t, err)
		return
	}
	require.NoError(t, err)
	require.Equal(t, value, v)
}

// isValid checks for 3 invariants:
// - BranchNode contains > 1 children
// - ExtensionNode do not contain another extension node
// - ExtensionNode do not have nil key
// It is used only during testing to catch possible bugs.
func isValid(curr Node) bool {
	switch n := curr.(type) {
	case *BranchNode:
		var count int
		for i := range n.Children {
			if !isValid(n.Children[i]) {
				return false
			}
			hn, ok := n.Children[i].(*HashNode)
			if !ok || !hn.IsEmpty() {
				count++
			}
		}
		return count > 1
	case *ExtensionNode:
		_, ok := n.next.(*ExtensionNode)
		return len(n.key) != 0 && !ok
	default:
		return true
	}
}

func TestTrie_Get(t *testing.T) {
	t.Run("HashNode", func(t *testing.T) {
		tr := newTestTrie(t)
		tr.testHas(t, []byte{0xAC, 0xAE}, []byte("hello"))
	})
	t.Run("UnfoldRoot", func(t *testing.T) {
		tr := newTestTrie(t)
		single := NewTrie(NewHashNode(tr.root.Hash()), false, 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})
		single.testHas(t, []byte{0xAC, 0xAE}, []byte("hello"))
	})
}

func TestTrie_Flush(t *testing.T) {
	pairs := map[string][]byte{
		"":     []byte("value0"),
		"key1": []byte("value1"),
		"key2": []byte("value2"),
	}

	tr := NewTrie(nil, false, newTestStore())
	for k, v := range pairs {
		require.NoError(t, tr.Put([]byte(k), v))
	}

	tr.Flush()
	tr = NewTrie(NewHashNode(tr.StateRoot()), false, tr.Store)
	for k, v := range pairs {
		actual, err := tr.Get([]byte(k))
		require.NoError(t, err)
		require.Equal(t, v, actual)
	}
}

func TestTrie_Delete(t *testing.T) {
	t.Run("No GC", func(t *testing.T) {
		testTrieDelete(t, false)
	})
	t.Run("With GC", func(t *testing.T) {
		testTrieDelete(t, true)
	})
}

func testTrieDelete(t *testing.T, enableGC bool) {
	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())
			t.Run("NotInStore", func(t *testing.T) {
				require.Error(t, tr.Delete([]byte{}))
			})

			tr.putToStore(l)
			tr.testHas(t, []byte{}, []byte{0x12})
			require.NoError(t, tr.Delete([]byte{}))
			tr.testHas(t, []byte{}, nil)
		})

		t.Run("Empty", func(t *testing.T) {
			tr := NewTrie(nil, enableGC, newTestStore())
			require.Error(t, tr.Delete([]byte{}))
		})
	})

	t.Run("Leaf", func(t *testing.T) {
		l := NewLeafNode([]byte{0x12, 0x34})
		tr := NewTrie(l, enableGC, newTestStore())
		t.Run("NonExistentKey", func(t *testing.T) {
			require.Error(t, tr.Delete([]byte{0x12}))
			tr.testHas(t, []byte{}, []byte{0x12, 0x34})
		})
		require.NoError(t, tr.Delete([]byte{}))
		tr.testHas(t, []byte{}, nil)
	})

	t.Run("Extension", func(t *testing.T) {
		t.Run("SingleKey", func(t *testing.T) {
			l := NewLeafNode([]byte{0x12, 0x34})
			e := NewExtensionNode([]byte{0x0A, 0x0B}, l)
			tr := NewTrie(e, enableGC, newTestStore())

			t.Run("NonExistentKey", func(t *testing.T) {
				require.Error(t, tr.Delete([]byte{}))
				tr.testHas(t, []byte{0xAB}, []byte{0x12, 0x34})
			})

			require.NoError(t, tr.Delete([]byte{0xAB}))
			require.True(t, tr.root.(*HashNode).IsEmpty())
		})

		t.Run("MultipleKeys", func(t *testing.T) {
			b := NewBranchNode()
			b.Children[0] = NewExtensionNode([]byte{0x01}, NewLeafNode([]byte{0x12, 0x34}))
			b.Children[6] = NewExtensionNode([]byte{0x07}, NewLeafNode([]byte{0x56, 0x78}))
			e := NewExtensionNode([]byte{0x01, 0x02}, b)
			tr := NewTrie(e, enableGC, newTestStore())

			h := e.Hash()
			require.NoError(t, tr.Delete([]byte{0x12, 0x01}))
			tr.testHas(t, []byte{0x12, 0x01}, nil)
			tr.testHas(t, []byte{0x12, 0x67}, []byte{0x56, 0x78})

			require.NotEqual(t, h, tr.root.Hash())
			require.Equal(t, toNibbles([]byte{0x12, 0x67}), e.key)
			require.IsType(t, (*LeafNode)(nil), e.next)
		})
	})

	t.Run("Branch", func(t *testing.T) {
		t.Run("3 Children", func(t *testing.T) {
			b := NewBranchNode()
			b.Children[lastChild] = NewLeafNode([]byte{0x12})
			b.Children[0] = NewExtensionNode([]byte{0x01}, NewLeafNode([]byte{0x34}))
			b.Children[1] = NewExtensionNode([]byte{0x06}, NewLeafNode([]byte{0x56}))
			tr := NewTrie(b, enableGC, newTestStore())
			require.NoError(t, tr.Delete([]byte{0x16}))
			tr.testHas(t, []byte{}, []byte{0x12})
			tr.testHas(t, []byte{0x01}, []byte{0x34})
			tr.testHas(t, []byte{0x16}, nil)
		})
		t.Run("2 Children", func(t *testing.T) {
			newt := func(t *testing.T) *Trie {
				b := NewBranchNode()
				b.Children[lastChild] = NewLeafNode([]byte{0x12})
				l := NewLeafNode([]byte{0x34})
				e := NewExtensionNode([]byte{0x06}, l)
				b.Children[5] = NewHashNode(e.Hash())
				tr := NewTrie(b, enableGC, newTestStore())
				tr.putToStore(l)
				tr.putToStore(e)
				return tr
			}

			t.Run("DeleteLast", func(t *testing.T) {
				t.Run("MergeExtension", func(t *testing.T) {
					tr := newt(t)
					require.NoError(t, tr.Delete([]byte{}))
					tr.testHas(t, []byte{}, nil)
					tr.testHas(t, []byte{0x56}, []byte{0x34})
					require.IsType(t, (*ExtensionNode)(nil), tr.root)
				})

				t.Run("LeaveLeaf", func(t *testing.T) {
					c := NewBranchNode()
					c.Children[5] = NewLeafNode([]byte{0x05})
					c.Children[6] = NewLeafNode([]byte{0x06})

					b := NewBranchNode()
					b.Children[lastChild] = NewLeafNode([]byte{0x12})
					b.Children[5] = c
					tr := NewTrie(b, enableGC, newTestStore())

					require.NoError(t, tr.Delete([]byte{}))
					tr.testHas(t, []byte{}, nil)
					tr.testHas(t, []byte{0x55}, []byte{0x05})
					tr.testHas(t, []byte{0x56}, []byte{0x06})
					require.IsType(t, (*ExtensionNode)(nil), tr.root)
				})
			})

			t.Run("DeleteMiddle", func(t *testing.T) {
				tr := newt(t)
				require.NoError(t, tr.Delete([]byte{0x56}))
				tr.testHas(t, []byte{}, []byte{0x12})
				tr.testHas(t, []byte{0x56}, nil)
				require.IsType(t, (*LeafNode)(nil), tr.root)
			})
		})
	})
}

func TestTrie_PanicInvalidRoot(t *testing.T) {
	tr := &Trie{Store: newTestStore()}
	require.Panics(t, func() { _ = tr.Put([]byte{1}, []byte{2}) })
	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, false, 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, false, 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, false, 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), false, 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, false, newTestStore())
		tr.Collapse(10)

		newRoot, ok := tr.root.(*HashNode)
		require.True(t, ok)
		require.Equal(t, NewHashNode(h), newRoot)
	})
}