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

     1  package mpt
     2  
     3  import (
     4  	"encoding/hex"
     5  	"fmt"
     6  	"testing"
     7  
     8  	"github.com/nspcc-dev/neo-go/pkg/core/storage"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestBatchAdd(t *testing.T) {
    13  	b := MapToMPTBatch(map[string][]byte{
    14  		"a\x01":     {2},
    15  		"a\x02\x10": {3},
    16  		"a\x00\x01": {5},
    17  		"a\x02\x00": {6},
    18  	})
    19  
    20  	expected := []keyValue{
    21  		{[]byte{0, 0, 0, 1}, []byte{5}},
    22  		{[]byte{0, 1}, []byte{2}},
    23  		{[]byte{0, 2, 0, 0}, []byte{6}},
    24  		{[]byte{0, 2, 1, 0}, []byte{3}},
    25  	}
    26  	require.Equal(t, expected, b.kv)
    27  }
    28  
    29  type pairs = [][2][]byte
    30  
    31  func testIncompletePut(t *testing.T, ps pairs, n int, tr1, tr2 *Trie) {
    32  	var m = make(map[string][]byte)
    33  	for i, p := range ps {
    34  		if i < n {
    35  			if p[1] == nil {
    36  				require.NoError(t, tr1.Delete(p[0]), "item %d", i)
    37  			} else {
    38  				require.NoError(t, tr1.Put(p[0], p[1]), "item %d", i)
    39  			}
    40  		} else if i == n {
    41  			if p[1] == nil {
    42  				require.Error(t, tr1.Delete(p[0]), "item %d", i)
    43  			} else {
    44  				require.Error(t, tr1.Put(p[0], p[1]), "item %d", i)
    45  			}
    46  		}
    47  		m["a"+string(p[0])] = p[1]
    48  	}
    49  
    50  	b := MapToMPTBatch(m)
    51  	num, err := tr2.PutBatch(b)
    52  	if n == len(ps) {
    53  		require.NoError(t, err)
    54  	} else {
    55  		require.Error(t, err)
    56  	}
    57  	require.Equal(t, n, num)
    58  	require.Equal(t, tr1.StateRoot(), tr2.StateRoot())
    59  
    60  	t.Run("test restore", func(t *testing.T) {
    61  		tr2.Flush(0)
    62  		tr3 := NewTrie(NewHashNode(tr2.StateRoot()), ModeAll, storage.NewMemCachedStore(tr2.Store))
    63  		for _, p := range ps[:n] {
    64  			val, err := tr3.Get(p[0])
    65  			if p[1] == nil {
    66  				require.Error(t, err)
    67  				continue
    68  			}
    69  			require.NoError(t, err, "key: %s", hex.EncodeToString(p[0]))
    70  			require.Equal(t, p[1], val)
    71  		}
    72  	})
    73  }
    74  
    75  func testPut(t *testing.T, ps pairs, tr1, tr2 *Trie) {
    76  	testIncompletePut(t, ps, len(ps), tr1, tr2)
    77  }
    78  
    79  func TestTrie_PutBatchLeaf(t *testing.T) {
    80  	prepareLeaf := func(t *testing.T) (*Trie, *Trie) {
    81  		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
    82  		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
    83  		require.NoError(t, tr1.Put([]byte{0}, []byte("value")))
    84  		require.NoError(t, tr2.Put([]byte{0}, []byte("value")))
    85  		return tr1, tr2
    86  	}
    87  
    88  	t.Run("remove", func(t *testing.T) {
    89  		tr1, tr2 := prepareLeaf(t)
    90  		var ps = pairs{{[]byte{0}, nil}}
    91  		testPut(t, ps, tr1, tr2)
    92  	})
    93  	t.Run("empty value", func(t *testing.T) {
    94  		tr1, tr2 := prepareLeaf(t)
    95  		var ps = pairs{{[]byte{0}, []byte{}}}
    96  		testPut(t, ps, tr1, tr2)
    97  	})
    98  	t.Run("replace", func(t *testing.T) {
    99  		tr1, tr2 := prepareLeaf(t)
   100  		var ps = pairs{{[]byte{0}, []byte("replace")}}
   101  		testPut(t, ps, tr1, tr2)
   102  	})
   103  	t.Run("remove and replace", func(t *testing.T) {
   104  		tr1, tr2 := prepareLeaf(t)
   105  		var ps = pairs{
   106  			{[]byte{0}, nil},
   107  			{[]byte{0, 2}, []byte("replace2")},
   108  		}
   109  		testPut(t, ps, tr1, tr2)
   110  	})
   111  	t.Run("empty value and replace", func(t *testing.T) {
   112  		tr1, tr2 := prepareLeaf(t)
   113  		var ps = pairs{
   114  			{[]byte{0}, []byte{}},
   115  			{[]byte{0, 2}, []byte("replace2")},
   116  		}
   117  		testPut(t, ps, tr1, tr2)
   118  	})
   119  }
   120  
   121  func TestTrie_PutBatchExtension(t *testing.T) {
   122  	prepareExtension := func(t *testing.T) (*Trie, *Trie) {
   123  		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   124  		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   125  		require.NoError(t, tr1.Put([]byte{1, 2}, []byte("value1")))
   126  		require.NoError(t, tr2.Put([]byte{1, 2}, []byte("value1")))
   127  		return tr1, tr2
   128  	}
   129  
   130  	t.Run("split, key len > 1", func(t *testing.T) {
   131  		tr1, tr2 := prepareExtension(t)
   132  		var ps = pairs{{[]byte{2, 3}, []byte("value2")}}
   133  		testPut(t, ps, tr1, tr2)
   134  	})
   135  	t.Run("split, key len = 1", func(t *testing.T) {
   136  		tr1, tr2 := prepareExtension(t)
   137  		var ps = pairs{{[]byte{1, 3}, []byte("value2")}}
   138  		testPut(t, ps, tr1, tr2)
   139  	})
   140  	t.Run("add to next", func(t *testing.T) {
   141  		tr1, tr2 := prepareExtension(t)
   142  		var ps = pairs{{[]byte{1, 2, 3}, []byte("value2")}}
   143  		testPut(t, ps, tr1, tr2)
   144  	})
   145  	t.Run("add to next with leaf", func(t *testing.T) {
   146  		tr1, tr2 := prepareExtension(t)
   147  		var ps = pairs{
   148  			{[]byte{0}, []byte("value3")},
   149  			{[]byte{1, 2, 3}, []byte("value2")},
   150  		}
   151  		testPut(t, ps, tr1, tr2)
   152  	})
   153  	t.Run("remove value", func(t *testing.T) {
   154  		tr1, tr2 := prepareExtension(t)
   155  		var ps = pairs{{[]byte{1, 2}, nil}}
   156  		testPut(t, ps, tr1, tr2)
   157  	})
   158  	t.Run("empty value", func(t *testing.T) {
   159  		tr1, tr2 := prepareExtension(t)
   160  		var ps = pairs{{[]byte{1, 2}, []byte{}}}
   161  		testPut(t, ps, tr1, tr2)
   162  	})
   163  	t.Run("add to next, merge extension", func(t *testing.T) {
   164  		tr1, tr2 := prepareExtension(t)
   165  		var ps = pairs{
   166  			{[]byte{1, 2}, nil},
   167  			{[]byte{1, 2, 3}, []byte("value2")},
   168  		}
   169  		testPut(t, ps, tr1, tr2)
   170  	})
   171  }
   172  
   173  func TestTrie_PutBatchBranch(t *testing.T) {
   174  	prepareBranch := func(t *testing.T) (*Trie, *Trie) {
   175  		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   176  		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   177  		require.NoError(t, tr1.Put([]byte{0x00, 2}, []byte("value1")))
   178  		require.NoError(t, tr2.Put([]byte{0x00, 2}, []byte("value1")))
   179  		require.NoError(t, tr1.Put([]byte{0x10, 3}, []byte("value2")))
   180  		require.NoError(t, tr2.Put([]byte{0x10, 3}, []byte("value2")))
   181  		return tr1, tr2
   182  	}
   183  
   184  	t.Run("simple add", func(t *testing.T) {
   185  		tr1, tr2 := prepareBranch(t)
   186  		var ps = pairs{{[]byte{0x20, 4}, []byte("value3")}}
   187  		testPut(t, ps, tr1, tr2)
   188  	})
   189  	t.Run("remove 1, transform to extension", func(t *testing.T) {
   190  		tr1, tr2 := prepareBranch(t)
   191  		var ps = pairs{{[]byte{0x00, 2}, nil}}
   192  		testPut(t, ps, tr1, tr2)
   193  
   194  		t.Run("non-empty child is hash node", func(t *testing.T) {
   195  			tr1, tr2 := prepareBranch(t)
   196  			tr1.Flush(0)
   197  			tr1.Collapse(1)
   198  			tr2.Flush(0)
   199  			tr2.Collapse(1)
   200  
   201  			var ps = pairs{{[]byte{0x00, 2}, nil}}
   202  			testPut(t, ps, tr1, tr2)
   203  			require.IsType(t, (*ExtensionNode)(nil), tr1.root)
   204  		})
   205  		t.Run("non-empty child is last node", func(t *testing.T) {
   206  			tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   207  			tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   208  			require.NoError(t, tr1.Put([]byte{0x00, 2}, []byte("value1")))
   209  			require.NoError(t, tr2.Put([]byte{0x00, 2}, []byte("value1")))
   210  			require.NoError(t, tr1.Put([]byte{0x00}, []byte("value2")))
   211  			require.NoError(t, tr2.Put([]byte{0x00}, []byte("value2")))
   212  
   213  			tr1.Flush(0)
   214  			tr1.Collapse(1)
   215  			tr2.Flush(0)
   216  			tr2.Collapse(1)
   217  
   218  			var ps = pairs{{[]byte{0x00, 2}, nil}}
   219  			testPut(t, ps, tr1, tr2)
   220  		})
   221  	})
   222  	t.Run("incomplete put, transform to extension", func(t *testing.T) {
   223  		tr1, tr2 := prepareBranch(t)
   224  		var ps = pairs{
   225  			{[]byte{0x00, 2}, nil},
   226  			{[]byte{0x20, 2}, nil},
   227  			{[]byte{0x30, 3}, []byte("won't be put")},
   228  		}
   229  		testIncompletePut(t, ps, 3, tr1, tr2)
   230  	})
   231  	t.Run("incomplete put, transform to empty", func(t *testing.T) {
   232  		tr1, tr2 := prepareBranch(t)
   233  		var ps = pairs{
   234  			{[]byte{0x00, 2}, nil},
   235  			{[]byte{0x10, 3}, nil},
   236  			{[]byte{0x20, 2}, nil},
   237  			{[]byte{0x30, 3}, []byte("won't be put")},
   238  		}
   239  		testIncompletePut(t, ps, 4, tr1, tr2)
   240  	})
   241  	t.Run("remove 2, become empty", func(t *testing.T) {
   242  		tr1, tr2 := prepareBranch(t)
   243  		var ps = pairs{
   244  			{[]byte{0x00, 2}, nil},
   245  			{[]byte{0x10, 3}, nil},
   246  		}
   247  		testPut(t, ps, tr1, tr2)
   248  	})
   249  }
   250  
   251  func TestTrie_PutBatchHash(t *testing.T) {
   252  	prepareHash := func(t *testing.T) (*Trie, *Trie) {
   253  		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   254  		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   255  		require.NoError(t, tr1.Put([]byte{0x10}, []byte("value1")))
   256  		require.NoError(t, tr2.Put([]byte{0x10}, []byte("value1")))
   257  		require.NoError(t, tr1.Put([]byte{0x20}, []byte("value2")))
   258  		require.NoError(t, tr2.Put([]byte{0x20}, []byte("value2")))
   259  		tr1.Flush(0)
   260  		tr2.Flush(0)
   261  		return tr1, tr2
   262  	}
   263  
   264  	t.Run("good", func(t *testing.T) {
   265  		tr1, tr2 := prepareHash(t)
   266  		var ps = pairs{{[]byte{2}, []byte("value2")}}
   267  		tr1.Collapse(0)
   268  		tr1.Collapse(0)
   269  		testPut(t, ps, tr1, tr2)
   270  	})
   271  	t.Run("incomplete, second hash not found", func(t *testing.T) {
   272  		tr1, tr2 := prepareHash(t)
   273  		var ps = pairs{
   274  			{[]byte{0x10}, []byte("replace1")},
   275  			{[]byte{0x20}, []byte("replace2")},
   276  		}
   277  		tr1.Collapse(1)
   278  		tr2.Collapse(1)
   279  		key := makeStorageKey(tr1.root.(*BranchNode).Children[2].Hash())
   280  		tr1.Store.Delete(key)
   281  		tr2.Store.Delete(key)
   282  		testIncompletePut(t, ps, 1, tr1, tr2)
   283  	})
   284  }
   285  
   286  func TestTrie_PutBatchEmpty(t *testing.T) {
   287  	t.Run("good", func(t *testing.T) {
   288  		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   289  		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   290  		var ps = pairs{
   291  			{[]byte{0}, []byte("value0")},
   292  			{[]byte{1}, []byte("value1")},
   293  			{[]byte{3}, []byte("value3")},
   294  		}
   295  		testPut(t, ps, tr1, tr2)
   296  	})
   297  	t.Run("incomplete", func(t *testing.T) {
   298  		var ps = pairs{
   299  			{[]byte{0}, []byte("replace0")},
   300  			{[]byte{1}, []byte("replace1")},
   301  			{[]byte{2}, nil},
   302  			{[]byte{3}, []byte("replace3")},
   303  		}
   304  		tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   305  		tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   306  		testIncompletePut(t, ps, 4, tr1, tr2)
   307  	})
   308  }
   309  
   310  // For the sake of coverage.
   311  func TestTrie_InvalidNodeType(t *testing.T) {
   312  	tr := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   313  	var b = Batch{kv: []keyValue{{
   314  		key:   []byte{0, 1},
   315  		value: []byte("value"),
   316  	}}}
   317  	tr.root = Node(nil)
   318  	require.Panics(t, func() { _, _ = tr.PutBatch(b) })
   319  }
   320  
   321  func TestTrie_PutBatch(t *testing.T) {
   322  	tr1 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   323  	tr2 := NewTrie(EmptyNode{}, ModeAll, newTestStore())
   324  	var ps = pairs{
   325  		{[]byte{1}, []byte{1}},
   326  		{[]byte{2}, []byte{3}},
   327  		{[]byte{4}, []byte{5}},
   328  	}
   329  	testPut(t, ps, tr1, tr2)
   330  
   331  	ps = pairs{[2][]byte{{4}, {6}}}
   332  	testPut(t, ps, tr1, tr2)
   333  
   334  	ps = pairs{[2][]byte{{4}, nil}}
   335  	testPut(t, ps, tr1, tr2)
   336  
   337  	testPut(t, pairs{}, tr1, tr2)
   338  }
   339  
   340  var _ = printNode
   341  
   342  // This function is unused, but is helpful for debugging
   343  // as it provides more readable Trie representation compared to
   344  // `spew.Dump()`.
   345  func printNode(prefix string, n Node) {
   346  	switch tn := n.(type) {
   347  	case EmptyNode:
   348  		fmt.Printf("%s empty\n", prefix)
   349  		return
   350  	case *HashNode:
   351  		fmt.Printf("%s %s\n", prefix, tn.Hash().StringLE())
   352  	case *BranchNode:
   353  		for i, c := range tn.Children {
   354  			if isEmpty(c) {
   355  				continue
   356  			}
   357  			fmt.Printf("%s [%2d] ->\n", prefix, i)
   358  			printNode(prefix+" ", c)
   359  		}
   360  	case *ExtensionNode:
   361  		fmt.Printf("%s extension-> %s\n", prefix, hex.EncodeToString(tn.key))
   362  		printNode(prefix+" ", tn.next)
   363  	case *LeafNode:
   364  		fmt.Printf("%s leaf-> %s\n", prefix, hex.EncodeToString(tn.value))
   365  	}
   366  }