github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/core/mpt/batch_test.go

package mpt

import (
	"encoding/hex"
	"fmt"
	"testing"

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

func TestBatchAdd(t *testing.T) {
	b := MapToMPTBatch(map[string][]byte{
		"a\x01":     {2},
		"a\x02\x10": {3},
		"a\x00\x01": {5},
		"a\x02\x00": {6},
	})

	expected := []keyValue{
		{[]byte{0, 0, 0, 1}, []byte{5}},
		{[]byte{0, 1}, []byte{2}},
		{[]byte{0, 2, 0, 0}, []byte{6}},
		{[]byte{0, 2, 1, 0}, []byte{3}},
	}
	require.Equal(t, expected, b.kv)
}

type pairs = [][2][]byte

func testIncompletePut(t *testing.T, ps pairs, n int, tr1, tr2 *Trie) {
	var m = make(map[string][]byte)
	for i, p := range ps {
		if i < n {
			if p[1] == nil {
				require.NoError(t, tr1.Delete(p[0]), "item %d", i)
			} else {
				require.NoError(t, tr1.Put(p[0], p[1]), "item %d", i)
			}
		} else if i == n {
			if p[1] == nil {
				require.Error(t, tr1.Delete(p[0]), "item %d", i)
			} else {
				require.Error(t, tr1.Put(p[0], p[1]), "item %d", i)
			}
		}
		m["a"+string(p[0])] = p[1]
	}

	b := MapToMPTBatch(m)
	num, err := tr2.PutBatch(b)
	if n == len(ps) {
		require.NoError(t, err)
	} else {
		require.Error(t, err)
	}
	require.Equal(t, n, num)
	require.Equal(t, tr1.StateRoot(), tr2.StateRoot())

	t.Run("test restore", func(t *testing.T) {
		tr2.Flush(0)
		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 {
				require.Error(t, err)
				continue
			}
			require.NoError(t, err, "key: %s", hex.EncodeToString(p[0]))
			require.Equal(t, p[1], val)
		}
	})
}

func testPut(t *testing.T, ps pairs, tr1, tr2 *Trie) {
	testIncompletePut(t, ps, len(ps), tr1, tr2)
}

func TestTrie_PutBatchLeaf(t *testing.T) {
	prepareLeaf := func(t *testing.T) (*Trie, *Trie) {
		tr1 := NewTrie(EmptyNode{}, ModeAll, 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
	}

	t.Run("remove", func(t *testing.T) {
		tr1, tr2 := prepareLeaf(t)
		var ps = pairs{{[]byte{0}, nil}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("empty value", func(t *testing.T) {
		tr1, tr2 := prepareLeaf(t)
		var ps = pairs{{[]byte{0}, []byte{}}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("replace", func(t *testing.T) {
		tr1, tr2 := prepareLeaf(t)
		var ps = pairs{{[]byte{0}, []byte("replace")}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("remove and replace", func(t *testing.T) {
		tr1, tr2 := prepareLeaf(t)
		var ps = pairs{
			{[]byte{0}, nil},
			{[]byte{0, 2}, []byte("replace2")},
		}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("empty value and replace", func(t *testing.T) {
		tr1, tr2 := prepareLeaf(t)
		var ps = pairs{
			{[]byte{0}, []byte{}},
			{[]byte{0, 2}, []byte("replace2")},
		}
		testPut(t, ps, tr1, tr2)
	})
}

func TestTrie_PutBatchExtension(t *testing.T) {
	prepareExtension := func(t *testing.T) (*Trie, *Trie) {
		tr1 := NewTrie(EmptyNode{}, ModeAll, 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
	}

	t.Run("split, key len > 1", func(t *testing.T) {
		tr1, tr2 := prepareExtension(t)
		var ps = pairs{{[]byte{2, 3}, []byte("value2")}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("split, key len = 1", func(t *testing.T) {
		tr1, tr2 := prepareExtension(t)
		var ps = pairs{{[]byte{1, 3}, []byte("value2")}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("add to next", func(t *testing.T) {
		tr1, tr2 := prepareExtension(t)
		var ps = pairs{{[]byte{1, 2, 3}, []byte("value2")}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("add to next with leaf", func(t *testing.T) {
		tr1, tr2 := prepareExtension(t)
		var ps = pairs{
			{[]byte{0}, []byte("value3")},
			{[]byte{1, 2, 3}, []byte("value2")},
		}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("remove value", func(t *testing.T) {
		tr1, tr2 := prepareExtension(t)
		var ps = pairs{{[]byte{1, 2}, nil}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("empty value", func(t *testing.T) {
		tr1, tr2 := prepareExtension(t)
		var ps = pairs{{[]byte{1, 2}, []byte{}}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("add to next, merge extension", func(t *testing.T) {
		tr1, tr2 := prepareExtension(t)
		var ps = pairs{
			{[]byte{1, 2}, nil},
			{[]byte{1, 2, 3}, []byte("value2")},
		}
		testPut(t, ps, tr1, tr2)
	})
}

func TestTrie_PutBatchBranch(t *testing.T) {
	prepareBranch := func(t *testing.T) (*Trie, *Trie) {
		tr1 := NewTrie(EmptyNode{}, ModeAll, 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")))
		require.NoError(t, tr2.Put([]byte{0x10, 3}, []byte("value2")))
		return tr1, tr2
	}

	t.Run("simple add", func(t *testing.T) {
		tr1, tr2 := prepareBranch(t)
		var ps = pairs{{[]byte{0x20, 4}, []byte("value3")}}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("remove 1, transform to extension", func(t *testing.T) {
		tr1, tr2 := prepareBranch(t)
		var ps = pairs{{[]byte{0x00, 2}, nil}}
		testPut(t, ps, tr1, tr2)

		t.Run("non-empty child is hash node", func(t *testing.T) {
			tr1, tr2 := prepareBranch(t)
			tr1.Flush(0)
			tr1.Collapse(1)
			tr2.Flush(0)
			tr2.Collapse(1)

			var ps = pairs{{[]byte{0x00, 2}, nil}}
			testPut(t, ps, tr1, tr2)
			require.IsType(t, (*ExtensionNode)(nil), tr1.root)
		})
		t.Run("non-empty child is last node", func(t *testing.T) {
			tr1 := NewTrie(EmptyNode{}, ModeAll, 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")))
			require.NoError(t, tr2.Put([]byte{0x00}, []byte("value2")))

			tr1.Flush(0)
			tr1.Collapse(1)
			tr2.Flush(0)
			tr2.Collapse(1)

			var ps = pairs{{[]byte{0x00, 2}, nil}}
			testPut(t, ps, tr1, tr2)
		})
	})
	t.Run("incomplete put, transform to extension", func(t *testing.T) {
		tr1, tr2 := prepareBranch(t)
		var ps = pairs{
			{[]byte{0x00, 2}, nil},
			{[]byte{0x20, 2}, nil},
			{[]byte{0x30, 3}, []byte("won't be put")},
		}
		testIncompletePut(t, ps, 3, tr1, tr2)
	})
	t.Run("incomplete put, transform to empty", func(t *testing.T) {
		tr1, tr2 := prepareBranch(t)
		var ps = pairs{
			{[]byte{0x00, 2}, nil},
			{[]byte{0x10, 3}, nil},
			{[]byte{0x20, 2}, nil},
			{[]byte{0x30, 3}, []byte("won't be put")},
		}
		testIncompletePut(t, ps, 4, tr1, tr2)
	})
	t.Run("remove 2, become empty", func(t *testing.T) {
		tr1, tr2 := prepareBranch(t)
		var ps = pairs{
			{[]byte{0x00, 2}, nil},
			{[]byte{0x10, 3}, nil},
		}
		testPut(t, ps, tr1, tr2)
	})
}

func TestTrie_PutBatchHash(t *testing.T) {
	prepareHash := func(t *testing.T) (*Trie, *Trie) {
		tr1 := NewTrie(EmptyNode{}, ModeAll, 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")))
		require.NoError(t, tr2.Put([]byte{0x20}, []byte("value2")))
		tr1.Flush(0)
		tr2.Flush(0)
		return tr1, tr2
	}

	t.Run("good", func(t *testing.T) {
		tr1, tr2 := prepareHash(t)
		var ps = pairs{{[]byte{2}, []byte("value2")}}
		tr1.Collapse(0)
		tr1.Collapse(0)
		testPut(t, ps, tr1, tr2)
	})
	t.Run("incomplete, second hash not found", func(t *testing.T) {
		tr1, tr2 := prepareHash(t)
		var ps = pairs{
			{[]byte{0x10}, []byte("replace1")},
			{[]byte{0x20}, []byte("replace2")},
		}
		tr1.Collapse(1)
		tr2.Collapse(1)
		key := makeStorageKey(tr1.root.(*BranchNode).Children[2].Hash())
		tr1.Store.Delete(key)
		tr2.Store.Delete(key)
		testIncompletePut(t, ps, 1, tr1, tr2)
	})
}

func TestTrie_PutBatchEmpty(t *testing.T) {
	t.Run("good", func(t *testing.T) {
		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
		var ps = pairs{
			{[]byte{0}, []byte("value0")},
			{[]byte{1}, []byte("value1")},
			{[]byte{3}, []byte("value3")},
		}
		testPut(t, ps, tr1, tr2)
	})
	t.Run("incomplete", func(t *testing.T) {
		var ps = pairs{
			{[]byte{0}, []byte("replace0")},
			{[]byte{1}, []byte("replace1")},
			{[]byte{2}, nil},
			{[]byte{3}, []byte("replace3")},
		}
		tr1 := NewTrie(EmptyNode{}, ModeAll, 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{}, ModeAll, newTestStore())
	var b = Batch{kv: []keyValue{{
		key:   []byte{0, 1},
		value: []byte("value"),
	}}}
	tr.root = Node(nil)
	require.Panics(t, func() { _, _ = tr.PutBatch(b) })
}

func TestTrie_PutBatch(t *testing.T) {
	tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
	tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
	var ps = pairs{
		{[]byte{1}, []byte{1}},
		{[]byte{2}, []byte{3}},
		{[]byte{4}, []byte{5}},
	}
	testPut(t, ps, tr1, tr2)

	ps = pairs{[2][]byte{{4}, {6}}}
	testPut(t, ps, tr1, tr2)

	ps = pairs{[2][]byte{{4}, nil}}
	testPut(t, ps, tr1, tr2)

	testPut(t, pairs{}, tr1, tr2)
}

var _ = printNode

// This function is unused, but is helpful for debugging
// as it provides more readable Trie representation compared to
// `spew.Dump()`.
func printNode(prefix string, n Node) {
	switch tn := n.(type) {
	case EmptyNode:
		fmt.Printf("%s empty\n", prefix)
		return
	case *HashNode:
		fmt.Printf("%s %s\n", prefix, tn.Hash().StringLE())
	case *BranchNode:
		for i, c := range tn.Children {
			if isEmpty(c) {
				continue
			}
			fmt.Printf("%s [%2d] ->\n", prefix, i)
			printNode(prefix+" ", c)
		}
	case *ExtensionNode:
		fmt.Printf("%s extension-> %s\n", prefix, hex.EncodeToString(tn.key))
		printNode(prefix+" ", tn.next)
	case *LeafNode:
		fmt.Printf("%s leaf-> %s\n", prefix, hex.EncodeToString(tn.value))
	}
}