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 }