508 lines
14 KiB
Go
508 lines
14 KiB
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)
|
|
})
|
|
}
|