diff --git a/dot/state/offline_pruner.go b/dot/state/offline_pruner.go index 72d400260b9..4c907db3ed9 100644 --- a/dot/state/offline_pruner.go +++ b/dot/state/offline_pruner.go @@ -121,7 +121,8 @@ func (p *OfflinePruner) SetBloomFilter() (err error) { return err } - trie.PopulateMerkleValues(tr.RootNode(), merkleValues) + const excludeInlined = false + trie.PopulateMerkleValues(tr.RootNode(), merkleValues, excludeInlined) // get parent header of current block header, err = p.blockState.GetHeader(header.ParentHash) diff --git a/dot/state/test_helpers.go b/dot/state/test_helpers.go index cd02d9b358f..3bdde6364d9 100644 --- a/dot/state/test_helpers.go +++ b/dot/state/test_helpers.go @@ -244,7 +244,8 @@ func generateBlockWithRandomTrie(t *testing.T, serv *Service, rand := time.Now().UnixNano() key := []byte("testKey" + fmt.Sprint(rand)) value := []byte("testValue" + fmt.Sprint(rand)) - trieState.Set(key, value) + err = trieState.Set(key, value) + require.NoError(t, err) trieStateRoot, err := trieState.Root() require.NoError(t, err) diff --git a/lib/runtime/interface.go b/lib/runtime/interface.go index 6fc9f3a01a5..fd81c51c3a2 100644 --- a/lib/runtime/interface.go +++ b/lib/runtime/interface.go @@ -50,22 +50,24 @@ type Instance interface { // Storage interface type Storage interface { - Set(key []byte, value []byte) + Set(key []byte, value []byte) (err error) Get(key []byte) []byte Root() (common.Hash, error) SetChild(keyToChild []byte, child *trie.Trie) error SetChildStorage(keyToChild, key, value []byte) error GetChildStorage(keyToChild, key []byte) ([]byte, error) - Delete(key []byte) - DeleteChild(keyToChild []byte) - DeleteChildLimit(keyToChild []byte, limit *[]byte) (uint32, bool, error) + Delete(key []byte) (err error) + DeleteChild(keyToChild []byte) (err error) + DeleteChildLimit(keyToChild []byte, limit *[]byte) ( + deleted uint32, allDeleted bool, err error) ClearChildStorage(keyToChild, key []byte) error NextKey([]byte) []byte ClearPrefixInChild(keyToChild, prefix []byte) error GetChildNextKey(keyToChild, key []byte) ([]byte, error) GetChild(keyToChild []byte) (*trie.Trie, error) - ClearPrefix(prefix []byte) - ClearPrefixLimit(prefix []byte, limit uint32) (uint32, bool) + ClearPrefix(prefix []byte) (err error) + ClearPrefixLimit(prefix []byte, limit uint32) ( + deleted uint32, allDeleted bool, err error) BeginStorageTransaction() CommitStorageTransaction() RollbackStorageTransaction() diff --git a/lib/runtime/storage/trie.go b/lib/runtime/storage/trie.go index e7dc4ca72c3..7955a8507d9 100644 --- a/lib/runtime/storage/trie.go +++ b/lib/runtime/storage/trie.go @@ -5,6 +5,7 @@ package storage import ( "encoding/binary" + "fmt" "sort" "sync" @@ -68,10 +69,10 @@ func (s *TrieState) RollbackStorageTransaction() { } // Set sets a key-value pair in the trie -func (s *TrieState) Set(key, value []byte) { +func (s *TrieState) Set(key, value []byte) (err error) { s.lock.Lock() defer s.lock.Unlock() - s.t.Put(key, value) + return s.t.Put(key, value) } // Get gets a value from the trie @@ -97,15 +98,20 @@ func (s *TrieState) Has(key []byte) bool { } // Delete deletes a key from the trie -func (s *TrieState) Delete(key []byte) { +func (s *TrieState) Delete(key []byte) (err error) { val := s.t.Get(key) if val == nil { - return + return nil } s.lock.Lock() defer s.lock.Unlock() - s.t.Delete(key) + err = s.t.Delete(key) + if err != nil { + return err + } + + return nil } // NextKey returns the next key in the trie in lexicographical order. If it does not exist, it returns nil. @@ -116,19 +122,19 @@ func (s *TrieState) NextKey(key []byte) []byte { } // ClearPrefix deletes all key-value pairs from the trie where the key starts with the given prefix -func (s *TrieState) ClearPrefix(prefix []byte) { +func (s *TrieState) ClearPrefix(prefix []byte) (err error) { s.lock.Lock() defer s.lock.Unlock() - s.t.ClearPrefix(prefix) + return s.t.ClearPrefix(prefix) } // ClearPrefixLimit deletes key-value pairs from the trie where the key starts with the given prefix till limit reached -func (s *TrieState) ClearPrefixLimit(prefix []byte, limit uint32) (uint32, bool) { +func (s *TrieState) ClearPrefixLimit(prefix []byte, limit uint32) ( + deleted uint32, allDeleted bool, err error) { s.lock.Lock() defer s.lock.Unlock() - num, del := s.t.ClearPrefixLimit(prefix, limit) - return num, del + return s.t.ClearPrefixLimit(prefix, limit) } // TrieEntries returns every key-value pair in the trie @@ -167,24 +173,29 @@ func (s *TrieState) GetChildStorage(keyToChild, key []byte) ([]byte, error) { } // DeleteChild deletes a child trie from the main trie -func (s *TrieState) DeleteChild(key []byte) { +func (s *TrieState) DeleteChild(key []byte) (err error) { s.lock.Lock() defer s.lock.Unlock() - s.t.DeleteChild(key) + return s.t.DeleteChild(key) } -// DeleteChildLimit deletes up to limit of database entries by lexicographic order, return number -// deleted, true if all delete otherwise false -func (s *TrieState) DeleteChildLimit(key []byte, limit *[]byte) (uint32, bool, error) { +// DeleteChildLimit deletes up to limit of database entries by lexicographic order. +func (s *TrieState) DeleteChildLimit(key []byte, limit *[]byte) ( + deleted uint32, allDeleted bool, err error) { s.lock.Lock() defer s.lock.Unlock() tr, err := s.t.GetChild(key) if err != nil { return 0, false, err } + qtyEntries := uint32(len(tr.Entries())) if limit == nil { - s.t.DeleteChild(key) + err = s.t.DeleteChild(key) + if err != nil { + return 0, false, err + } + return qtyEntries, true, nil } limitUint := binary.LittleEndian.Uint32(*limit) @@ -194,20 +205,24 @@ func (s *TrieState) DeleteChildLimit(key []byte, limit *[]byte) (uint32, bool, e keys = append(keys, k) } sort.Strings(keys) - deleted := uint32(0) for _, k := range keys { - tr.Delete([]byte(k)) + // TODO have a transactional/atomic way to delete multiple keys in trie. + // If one deletion fails, the child trie and its parent trie are then in + // a bad intermediary state. Take also care of the caching of deleted Merkle + // values within the tries, which is used for online pruning. + err = tr.Delete([]byte(k)) + if err != nil { + return deleted, allDeleted, fmt.Errorf("deleting from child trie at key 0x%x: %w", key, err) + } + deleted++ if deleted == limitUint { break } } - if deleted == qtyEntries { - return deleted, true, nil - } - - return deleted, false, nil + allDeleted = deleted == qtyEntries + return deleted, allDeleted, nil } // ClearChildStorage removes the child storage entry from the trie @@ -230,7 +245,11 @@ func (s *TrieState) ClearPrefixInChild(keyToChild, prefix []byte) error { return nil } - child.ClearPrefix(prefix) + err = child.ClearPrefix(prefix) + if err != nil { + return fmt.Errorf("clearing prefix in child trie at key 0x%x: %w", keyToChild, err) + } + return nil } diff --git a/lib/runtime/wasmer/helpers.go b/lib/runtime/wasmer/helpers.go index 6065f2d83ce..81629a37f53 100644 --- a/lib/runtime/wasmer/helpers.go +++ b/lib/runtime/wasmer/helpers.go @@ -159,7 +159,7 @@ func toWasmMemoryFixedSizeOptional(context wasmer.InstanceContext, data []byte) return toWasmMemory(context, encodedOptionalFixedSize) } -func storageAppend(storage runtime.Storage, key, valueToAppend []byte) error { +func storageAppend(storage runtime.Storage, key, valueToAppend []byte) (err error) { // this function assumes the item in storage is a SCALE encoded array of items // the valueToAppend is a new item, so it appends the item and increases the length prefix by 1 currentValue := storage.Get(key) @@ -210,7 +210,16 @@ func storageAppend(storage runtime.Storage, key, valueToAppend []byte) error { } } - logger.Debugf("resulting value: 0x%x", value) - storage.Set(key, value) + err = storage.Set(key, value) + if err != nil { + return fmt.Errorf("settings key and value in storage: %w", err) + } + return nil } + +func panicOnError(err error) { + if err != nil { + panic(err) + } +} diff --git a/lib/runtime/wasmer/helpers_test.go b/lib/runtime/wasmer/helpers_test.go index 279dc499666..f89af3a58a9 100644 --- a/lib/runtime/wasmer/helpers_test.go +++ b/lib/runtime/wasmer/helpers_test.go @@ -5,6 +5,7 @@ package wasmer import ( "encoding/json" + "errors" "os" "path/filepath" "testing" @@ -66,3 +67,13 @@ func Test_pointerSize(t *testing.T) { }) } } + +func Test_panicOnError(t *testing.T) { + t.Parallel() + + err := (error)(nil) + assert.NotPanics(t, func() { panicOnError(err) }) + + err = errors.New("test error") + assert.PanicsWithValue(t, err, func() { panicOnError(err) }) +} diff --git a/lib/runtime/wasmer/imports.go b/lib/runtime/wasmer/imports.go index 5c88fa5a4bf..b342b543195 100644 --- a/lib/runtime/wasmer/imports.go +++ b/lib/runtime/wasmer/imports.go @@ -823,7 +823,12 @@ func ext_trie_blake2_256_root_version_1(context unsafe.Pointer, dataSpan C.int64 } for _, kv := range kvs { - t.Put(kv.Key, kv.Value) + err := t.Put(kv.Key, kv.Value) + if err != nil { + logger.Errorf("failed putting key 0x%x and value 0x%x into trie: %s", + kv.Key, kv.Value, err) + return 0 + } } // allocate memory for value and copy value to memory @@ -871,7 +876,12 @@ func ext_trie_blake2_256_ordered_root_version_1(context unsafe.Pointer, dataSpan "put key=0x%x and value=0x%x", key, value) - t.Put(key, value) + err = t.Put(key, value) + if err != nil { + logger.Errorf("failed putting key 0x%x and value 0x%x into trie: %s", + key, value, err) + return 0 + } } // allocate memory for value and copy value to memory @@ -1187,7 +1197,8 @@ func ext_default_child_storage_storage_kill_version_1(context unsafe.Pointer, ch storage := ctx.Storage childStorageKey := asMemorySlice(instanceContext, childStorageKeySpan) - storage.DeleteChild(childStorageKey) + err := storage.DeleteChild(childStorageKey) + panicOnError(err) } //export ext_default_child_storage_storage_kill_version_2 @@ -1850,7 +1861,8 @@ func ext_storage_clear_version_1(context unsafe.Pointer, keySpan C.int64_t) { key := asMemorySlice(instanceContext, keySpan) logger.Debugf("key: 0x%x", key) - storage.Delete(key) + err := storage.Delete(key) + panicOnError(err) } //export ext_storage_clear_prefix_version_1 @@ -1863,7 +1875,8 @@ func ext_storage_clear_prefix_version_1(context unsafe.Pointer, prefixSpan C.int prefix := asMemorySlice(instanceContext, prefixSpan) logger.Debugf("prefix: 0x%x", prefix) - storage.ClearPrefix(prefix) + err := storage.ClearPrefix(prefix) + panicOnError(err) } //export ext_storage_clear_prefix_version_2 @@ -1893,7 +1906,13 @@ func ext_storage_clear_prefix_version_2(context unsafe.Pointer, prefixSpan, lim } limitUint := binary.LittleEndian.Uint32(limit) - numRemoved, all := storage.ClearPrefixLimit(prefix, limitUint) + numRemoved, all, err := storage.ClearPrefixLimit(prefix, limitUint) + if err != nil { + logger.Errorf("failed to clear prefix limit: %s", err) + ret, _ := toWasmMemory(instanceContext, nil) + return C.int64_t(ret) + } + encBytes, err := toKillStorageResultEnum(all, numRemoved) if err != nil { logger.Errorf("failed to allocate memory: %s", err) @@ -2056,7 +2075,8 @@ func ext_storage_set_version_1(context unsafe.Pointer, keySpan, valueSpan C.int6 logger.Debugf( "key 0x%x has value 0x%x", key, value) - storage.Set(key, cp) + err := storage.Set(key, cp) + panicOnError(err) } //export ext_storage_start_transaction_version_1 diff --git a/lib/trie/child_storage.go b/lib/trie/child_storage.go index fc80da533f3..6eb869b2a9f 100644 --- a/lib/trie/child_storage.go +++ b/lib/trie/child_storage.go @@ -28,7 +28,11 @@ func (t *Trie) PutChild(keyToChild []byte, child *Trie) error { copy(key, ChildStorageKeyPrefix) copy(key[len(ChildStorageKeyPrefix):], keyToChild) - t.Put(key, childHash.ToBytes()) + err = t.Put(key, childHash.ToBytes()) + if err != nil { + return fmt.Errorf("putting child trie root hash %s in main trie: %w", childHash, err) + } + t.childTries[childHash] = child return nil } @@ -59,7 +63,11 @@ func (t *Trie) PutIntoChild(keyToChild, key, value []byte) error { return err } - child.Put(key, value) + err = child.Put(key, value) + if err != nil { + return fmt.Errorf("putting into child trie at key 0x%x: %w", keyToChild, err) + } + childHash, err := child.Hash() if err != nil { return err @@ -88,12 +96,16 @@ func (t *Trie) GetFromChild(keyToChild, key []byte) ([]byte, error) { } // DeleteChild deletes the child storage trie -func (t *Trie) DeleteChild(keyToChild []byte) { +func (t *Trie) DeleteChild(keyToChild []byte) (err error) { key := make([]byte, len(ChildStorageKeyPrefix)+len(keyToChild)) copy(key, ChildStorageKeyPrefix) copy(key[len(ChildStorageKeyPrefix):], keyToChild) - t.Delete(key) + err = t.Delete(key) + if err != nil { + return fmt.Errorf("deleting child trie at key 0x%x: %w", key, err) + } + return nil } // ClearFromChild removes the child storage entry @@ -105,6 +117,11 @@ func (t *Trie) ClearFromChild(keyToChild, key []byte) error { if child == nil { return fmt.Errorf("%w at key 0x%x%x", ErrChildTrieDoesNotExist, ChildStorageKeyPrefix, keyToChild) } - child.Delete(key) + + err = child.Delete(key) + if err != nil { + return fmt.Errorf("deleting from child trie at key 0x%x: %w", keyToChild, err) + } + return nil } diff --git a/lib/trie/database.go b/lib/trie/database.go index e5c59f8897e..9e1cac57fe0 100644 --- a/lib/trie/database.go +++ b/lib/trie/database.go @@ -189,7 +189,7 @@ func (t *Trie) loadNode(db Database, n *Node) error { // all its descendant nodes as keys to the map merkleValues. // It is assumed the node and its descendant nodes have their Merkle value already // computed. -func PopulateMerkleValues(n *Node, merkleValues map[string]struct{}) { +func PopulateMerkleValues(n *Node, merkleValues map[string]struct{}, excludeInlined bool) { if n == nil { return } @@ -198,7 +198,10 @@ func PopulateMerkleValues(n *Node, merkleValues map[string]struct{}) { panic(fmt.Sprintf("node with key 0x%x has no Merkle value computed", n.Key)) } - merkleValues[string(n.MerkleValue)] = struct{}{} + isInlined := len(n.MerkleValue) < 32 + if !excludeInlined || !isInlined { + merkleValues[string(n.MerkleValue)] = struct{}{} + } if n.Kind() == node.Leaf { return @@ -206,7 +209,7 @@ func PopulateMerkleValues(n *Node, merkleValues map[string]struct{}) { branch := n for _, child := range branch.Children { - PopulateMerkleValues(child, merkleValues) + PopulateMerkleValues(child, merkleValues, excludeInlined) } } diff --git a/lib/trie/database_test.go b/lib/trie/database_test.go index 63a2883d1a9..1aefa1c63c4 100644 --- a/lib/trie/database_test.go +++ b/lib/trie/database_test.go @@ -162,9 +162,10 @@ func Test_PopulateMerkleValues(t *testing.T) { t.Parallel() testCases := map[string]struct { - node *Node - merkleValues map[string]struct{} - panicValue interface{} + node *Node + excludeInlined bool + merkleValues map[string]struct{} + panicValue interface{} }{ "nil node": { merkleValues: map[string]struct{}{}, @@ -191,6 +192,16 @@ func Test_PopulateMerkleValues(t *testing.T) { "b": {}, }, }, + "exclude inlined": { + node: &Node{ + MerkleValue: []byte("a"), + Children: padRightChildren([]*Node{ + {MerkleValue: []byte("b")}, + }), + }, + excludeInlined: true, + merkleValues: map[string]struct{}{}, + }, "nested branch node": { node: &Node{ MerkleValue: []byte("a"), @@ -222,12 +233,12 @@ func Test_PopulateMerkleValues(t *testing.T) { if testCase.panicValue != nil { assert.PanicsWithValue(t, testCase.panicValue, func() { - PopulateMerkleValues(testCase.node, merkleValues) + PopulateMerkleValues(testCase.node, merkleValues, testCase.excludeInlined) }) return } - PopulateMerkleValues(testCase.node, merkleValues) + PopulateMerkleValues(testCase.node, merkleValues, testCase.excludeInlined) assert.Equal(t, testCase.merkleValues, merkleValues) }) diff --git a/lib/trie/trie.go b/lib/trie/trie.go index 735d32c2487..686d111b708 100644 --- a/lib/trie/trie.go +++ b/lib/trie/trie.go @@ -82,50 +82,79 @@ func (t *Trie) handleTrackedDeltas(success bool, pendingDeletedMerkleValues map[ func (t *Trie) prepLeafForMutation(currentLeaf *Node, copySettings node.CopySettings, - pendingDeletedMerkleValues map[string]struct{}) (newLeaf *Node) { + pendingDeletedMerkleValues map[string]struct{}) ( + newLeaf *Node, err error) { if currentLeaf.Generation == t.generation { // no need to deep copy and update generation // of current leaf. newLeaf = currentLeaf } else { - newLeaf = updateGeneration(currentLeaf, t.generation, pendingDeletedMerkleValues, copySettings) + newLeaf, err = t.updateGeneration(currentLeaf, pendingDeletedMerkleValues, copySettings) + if err != nil { + return nil, fmt.Errorf("updating generation: %w", err) + } } newLeaf.SetDirty() - return newLeaf + return newLeaf, nil } func (t *Trie) prepBranchForMutation(currentBranch *Node, copySettings node.CopySettings, - pendingDeletedMerkleValues map[string]struct{}) (newBranch *Node) { + pendingDeletedMerkleValues map[string]struct{}) ( + newBranch *Node, err error) { if currentBranch.Generation == t.generation { // no need to deep copy and update generation // of current branch. newBranch = currentBranch } else { - newBranch = updateGeneration(currentBranch, t.generation, pendingDeletedMerkleValues, copySettings) + newBranch, err = t.updateGeneration(currentBranch, pendingDeletedMerkleValues, copySettings) + if err != nil { + return nil, fmt.Errorf("updating generation: %w", err) + } } newBranch.SetDirty() - return newBranch + return newBranch, nil } // updateGeneration is called when the currentNode is from // an older trie generation (snapshot) so we deep copy the // node and update the generation on the newer copy. -func updateGeneration(currentNode *Node, trieGeneration uint64, - deletedMerkleValues map[string]struct{}, copySettings node.CopySettings) ( - newNode *Node) { +func (t *Trie) updateGeneration(currentNode *Node, + pendingDeletedMerkleValues map[string]struct{}, + copySettings node.CopySettings) (newNode *Node, err error) { + isRoot := currentNode == t.root + err = registerDeletedMerkleValue(currentNode, isRoot, pendingDeletedMerkleValues) + if err != nil { + return nil, fmt.Errorf("registering deleted merkle value: %w", err) + } + newNode = currentNode.Copy(copySettings) - newNode.Generation = trieGeneration + newNode.Generation = t.generation + + return newNode, nil +} + +func registerDeletedMerkleValue(node *Node, isRoot bool, + pendingDeletedMerkleValues map[string]struct{}) (err error) { + err = ensureMerkleValueIsCalculated(node, isRoot) + if err != nil { + return fmt.Errorf("ensuring Merkle value is calculated: %w", err) + } + + if len(node.MerkleValue) < 32 { + // Merkle values which are less than 32 bytes are inlined + // in the parent branch and are not stored on disk, so there + // is no need to track their deletion for the online pruning. + return nil + } - // The hash of the node from a previous snapshotted trie - // is usually already computed. - deletedMerkleValue := currentNode.MerkleValue - if len(deletedMerkleValue) > 0 { - deletedMerkleValueString := string(deletedMerkleValue) - deletedMerkleValues[deletedMerkleValueString] = struct{}{} + if !node.Dirty { + // Only register deleted nodes that were not previously modified + // since the last trie snapshot. + pendingDeletedMerkleValues[string(node.MerkleValue)] = struct{}{} } - return newNode + return nil } // DeepCopy deep copies the trie and returns @@ -335,25 +364,31 @@ func findNextKeyChild(children []*Node, startIndex byte, // Put inserts a value into the trie at the // key specified in little Endian format. -func (t *Trie) Put(keyLE, value []byte) { +func (t *Trie) Put(keyLE, value []byte) (err error) { pendingDeletedMerkleValues := make(map[string]struct{}) defer func() { const success = true t.handleTrackedDeltas(success, pendingDeletedMerkleValues) }() - t.insertKeyLE(keyLE, value, pendingDeletedMerkleValues) + return t.insertKeyLE(keyLE, value, pendingDeletedMerkleValues) } -func (t *Trie) insertKeyLE(keyLE, value []byte, deletedMerkleValues map[string]struct{}) { +func (t *Trie) insertKeyLE(keyLE, value []byte, + deletedMerkleValues map[string]struct{}) (err error) { nibblesKey := codec.KeyLEToNibbles(keyLE) - t.root, _, _ = t.insert(t.root, nibblesKey, value, deletedMerkleValues) + root, _, _, err := t.insert(t.root, nibblesKey, value, deletedMerkleValues) + if err != nil { + return err + } + t.root = root + return nil } // insert inserts a value in the trie at the key specified. // It may create one or more new nodes or update an existing node. func (t *Trie) insert(parent *Node, key, value []byte, deletedMerkleValues map[string]struct{}) (newParent *Node, - mutated bool, nodesCreated uint32) { + mutated bool, nodesCreated uint32, err error) { if parent == nil { mutated = true nodesCreated = 1 @@ -362,33 +397,51 @@ func (t *Trie) insert(parent *Node, key, value []byte, SubValue: value, Generation: t.generation, Dirty: true, - }, mutated, nodesCreated + }, mutated, nodesCreated, nil } // TODO ensure all values have dirty set to true if parent.Kind() == node.Branch { - return t.insertInBranch(parent, key, value, deletedMerkleValues) + newParent, mutated, nodesCreated, err = t.insertInBranch( + parent, key, value, deletedMerkleValues) + if err != nil { + // `insertInBranch` may call `insert` so do not wrap the + // error since this may be a deep recursive call. + return nil, false, 0, err + } + return newParent, mutated, nodesCreated, nil } - return t.insertInLeaf(parent, key, value, deletedMerkleValues) + + newParent, mutated, nodesCreated, err = t.insertInLeaf( + parent, key, value, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("inserting in leaf: %w", err) + } + + return newParent, mutated, nodesCreated, nil } func (t *Trie) insertInLeaf(parentLeaf *Node, key, value []byte, deletedMerkleValues map[string]struct{}) ( - newParent *Node, mutated bool, nodesCreated uint32) { + newParent *Node, mutated bool, nodesCreated uint32, err error) { if bytes.Equal(parentLeaf.Key, key) { nodesCreated = 0 if parentLeaf.SubValueEqual(value) { mutated = false - return parentLeaf, mutated, nodesCreated + return parentLeaf, mutated, nodesCreated, nil } copySettings := node.DefaultCopySettings copySettings.CopyValue = false - parentLeaf = t.prepLeafForMutation(parentLeaf, copySettings, deletedMerkleValues) + parentLeaf, err = t.prepLeafForMutation(parentLeaf, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing leaf for mutation: %w", err) + } + parentLeaf.SubValue = value mutated = true - return parentLeaf, mutated, nodesCreated + return parentLeaf, mutated, nodesCreated, nil } commonPrefixLength := lenCommonPrefix(key, parentLeaf.Key) @@ -413,7 +466,10 @@ func (t *Trie) insertInLeaf(parentLeaf *Node, key, value []byte, childIndex := parentLeafKey[commonPrefixLength] newParentLeafKey := parentLeaf.Key[commonPrefixLength+1:] if !bytes.Equal(parentLeaf.Key, newParentLeafKey) { - parentLeaf = t.prepLeafForMutation(parentLeaf, copySettings, deletedMerkleValues) + parentLeaf, err = t.prepLeafForMutation(parentLeaf, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing leaf for mutation: %w", err) + } parentLeaf.Key = newParentLeafKey } newBranchParent.Children[childIndex] = parentLeaf @@ -421,7 +477,7 @@ func (t *Trie) insertInLeaf(parentLeaf *Node, key, value []byte, nodesCreated++ } - return newBranchParent, mutated, nodesCreated + return newBranchParent, mutated, nodesCreated, nil } if len(parentLeaf.Key) == commonPrefixLength { @@ -433,7 +489,10 @@ func (t *Trie) insertInLeaf(parentLeaf *Node, key, value []byte, childIndex := parentLeafKey[commonPrefixLength] newParentLeafKey := parentLeaf.Key[commonPrefixLength+1:] if !bytes.Equal(parentLeaf.Key, newParentLeafKey) { - parentLeaf = t.prepLeafForMutation(parentLeaf, copySettings, deletedMerkleValues) + parentLeaf, err = t.prepLeafForMutation(parentLeaf, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing leaf for mutation: %w", err) + } parentLeaf.Key = newParentLeafKey } newBranchParent.Children[childIndex] = parentLeaf @@ -450,23 +509,26 @@ func (t *Trie) insertInLeaf(parentLeaf *Node, key, value []byte, newBranchParent.Descendants++ nodesCreated++ - return newBranchParent, mutated, nodesCreated + return newBranchParent, mutated, nodesCreated, nil } func (t *Trie) insertInBranch(parentBranch *Node, key, value []byte, deletedMerkleValues map[string]struct{}) ( - newParent *Node, mutated bool, nodesCreated uint32) { + newParent *Node, mutated bool, nodesCreated uint32, err error) { copySettings := node.DefaultCopySettings if bytes.Equal(key, parentBranch.Key) { if parentBranch.SubValueEqual(value) { mutated = false - return parentBranch, mutated, 0 + return parentBranch, mutated, 0, nil + } + parentBranch, err = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing branch for mutation: %w", err) } - parentBranch = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) parentBranch.SubValue = value mutated = true - return parentBranch, mutated, 0 + return parentBranch, mutated, 0, nil } if bytes.HasPrefix(key, parentBranch.Key) { @@ -484,22 +546,32 @@ func (t *Trie) insertInBranch(parentBranch *Node, key, value []byte, Dirty: true, } nodesCreated = 1 - parentBranch = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) + parentBranch, err = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing branch for mutation: %w", err) + } parentBranch.Children[childIndex] = child parentBranch.Descendants += nodesCreated mutated = true - return parentBranch, mutated, nodesCreated + return parentBranch, mutated, nodesCreated, nil + } + + child, mutated, nodesCreated, err = t.insert(child, remainingKey, value, deletedMerkleValues) + if err != nil { + // do not wrap error since `insert` may call `insertInBranch` recursively + return nil, false, 0, err + } else if !mutated { + return parentBranch, mutated, 0, nil } - child, mutated, nodesCreated = t.insert(child, remainingKey, value, deletedMerkleValues) - if !mutated { - return parentBranch, mutated, 0 + parentBranch, err = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing branch for mutation: %w", err) } - parentBranch = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) parentBranch.Children[childIndex] = child parentBranch.Descendants += nodesCreated - return parentBranch, mutated, nodesCreated + return parentBranch, mutated, nodesCreated, nil } // we need to branch out at the point where the keys diverge @@ -518,7 +590,11 @@ func (t *Trie) insertInBranch(parentBranch *Node, key, value []byte, remainingOldParentKey := parentBranch.Key[commonPrefixLength+1:] // Note: parentBranch.Key != remainingOldParentKey - parentBranch = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) + parentBranch, err = t.prepBranchForMutation(parentBranch, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing branch for mutation: %w", err) + } + parentBranch.Key = remainingOldParentKey newParentBranch.Children[oldParentIndex] = parentBranch newParentBranch.Descendants += 1 + parentBranch.Descendants @@ -529,13 +605,18 @@ func (t *Trie) insertInBranch(parentBranch *Node, key, value []byte, childIndex := key[commonPrefixLength] remainingKey := key[commonPrefixLength+1:] var additionalNodesCreated uint32 - newParentBranch.Children[childIndex], _, additionalNodesCreated = t.insert( + newParentBranch.Children[childIndex], _, additionalNodesCreated, err = t.insert( nil, remainingKey, value, deletedMerkleValues) + if err != nil { + // do not wrap error since `insert` may call `insertInBranch` recursively + return nil, false, 0, err + } + nodesCreated += additionalNodesCreated newParentBranch.Descendants += additionalNodesCreated } - return newParentBranch, mutated, nodesCreated + return newParentBranch, mutated, nodesCreated, nil } // LoadFromMap loads the given data mapping of key to value into a new empty trie. @@ -560,7 +641,10 @@ func LoadFromMap(data map[string]string) (trie Trie, err error) { return Trie{}, fmt.Errorf("cannot convert value hex to bytes: %w", err) } - trie.insertKeyLE(keyLEBytes, valueBytes, pendingDeletedMerkleValues) + err = trie.insertKeyLE(keyLEBytes, valueBytes, pendingDeletedMerkleValues) + if err != nil { + return Trie{}, fmt.Errorf("inserting key value pair in trie: %w", err) + } } return trie, nil @@ -714,7 +798,8 @@ func retrieveFromBranch(branch *Node, key []byte) (value []byte) { // Endian format for up to `limit` keys. It returns the number of deleted // keys and a boolean indicating if all keys with the prefix were deleted // within the limit. -func (t *Trie) ClearPrefixLimit(prefixLE []byte, limit uint32) (deleted uint32, allDeleted bool) { +func (t *Trie) ClearPrefixLimit(prefixLE []byte, limit uint32) ( + deleted uint32, allDeleted bool, err error) { pendingDeletedMerkleValues := make(map[string]struct{}) defer func() { const success = true @@ -722,15 +807,21 @@ func (t *Trie) ClearPrefixLimit(prefixLE []byte, limit uint32) (deleted uint32, }() if limit == 0 { - return 0, false + return 0, false, nil } prefix := codec.KeyLEToNibbles(prefixLE) prefix = bytes.TrimSuffix(prefix, []byte{0}) - t.root, deleted, _, allDeleted = t.clearPrefixLimitAtNode( + t.root, deleted, _, allDeleted, err = t.clearPrefixLimitAtNode( t.root, prefix, limit, pendingDeletedMerkleValues) - return deleted, allDeleted + if err != nil { + // Note: no need to wrap the error really since the private function has + // the same name as the exported function `ClearPrefixLimit`. + return 0, false, err + } + + return deleted, allDeleted, nil } // clearPrefixLimitAtNode deletes the keys having the prefix until the value deletion limit is reached. @@ -738,9 +829,9 @@ func (t *Trie) ClearPrefixLimit(prefixLE []byte, limit uint32) (deleted uint32, // allDeleted boolean indicating if there is no key left with the prefix. func (t *Trie) clearPrefixLimitAtNode(parent *Node, prefix []byte, limit uint32, deletedMerkleValues map[string]struct{}) ( - newParent *Node, valuesDeleted, nodesRemoved uint32, allDeleted bool) { + newParent *Node, valuesDeleted, nodesRemoved uint32, allDeleted bool, err error) { if parent == nil { - return nil, 0, 0, true + return nil, 0, 0, true, nil } if parent.Kind() == node.Leaf { @@ -748,25 +839,37 @@ func (t *Trie) clearPrefixLimitAtNode(parent *Node, prefix []byte, // TODO check this is the same behaviour as in substrate const allDeleted = true if bytes.HasPrefix(parent.Key, prefix) { + isRoot := parent == t.root + err = registerDeletedMerkleValue(parent, isRoot, deletedMerkleValues) + if err != nil { + return nil, 0, 0, false, + fmt.Errorf("registering deleted Merkle value: %w", err) + } + valuesDeleted, nodesRemoved = 1, 1 - return nil, valuesDeleted, nodesRemoved, allDeleted + return nil, valuesDeleted, nodesRemoved, allDeleted, nil } - return parent, 0, 0, allDeleted + return parent, 0, 0, allDeleted, nil } + // Note: `clearPrefixLimitBranch` may call `clearPrefixLimitAtNode` so do not wrap + // the error since that could be a deep recursive call. return t.clearPrefixLimitBranch(parent, prefix, limit, deletedMerkleValues) } func (t *Trie) clearPrefixLimitBranch(branch *Node, prefix []byte, limit uint32, deletedMerkleValues map[string]struct{}) ( - newParent *Node, valuesDeleted, nodesRemoved uint32, allDeleted bool) { + newParent *Node, valuesDeleted, nodesRemoved uint32, allDeleted bool, err error) { newParent = branch if bytes.HasPrefix(branch.Key, prefix) { - newParent, valuesDeleted, nodesRemoved = t.deleteNodesLimit( + newParent, valuesDeleted, nodesRemoved, err = t.deleteNodesLimit( branch, limit, deletedMerkleValues) + if err != nil { + return nil, 0, 0, false, fmt.Errorf("deleting nodes: %w", err) + } allDeleted = newParent == nil - return newParent, valuesDeleted, nodesRemoved, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted, nil } if len(prefix) == len(branch.Key)+1 && @@ -780,34 +883,44 @@ func (t *Trie) clearPrefixLimitBranch(branch *Node, prefix []byte, limit uint32, if noPrefixForNode { valuesDeleted, nodesRemoved = 0, 0 allDeleted = true - return newParent, valuesDeleted, nodesRemoved, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted, nil } childIndex := prefix[len(branch.Key)] childPrefix := prefix[len(branch.Key)+1:] child := branch.Children[childIndex] - child, valuesDeleted, nodesRemoved, allDeleted = t.clearPrefixLimitAtNode( + child, valuesDeleted, nodesRemoved, allDeleted, err = t.clearPrefixLimitAtNode( child, childPrefix, limit, deletedMerkleValues) - if valuesDeleted == 0 { - return branch, valuesDeleted, nodesRemoved, allDeleted + if err != nil { + return nil, 0, 0, false, fmt.Errorf("clearing prefix limit at node: %w", err) + } else if valuesDeleted == 0 { + return branch, valuesDeleted, nodesRemoved, allDeleted, nil } copySettings := node.DefaultCopySettings - branch = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + branch, err = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + if err != nil { + return nil, 0, 0, false, fmt.Errorf("preparing branch for mutation: %w", err) + } + branch.Children[childIndex] = child branch.Descendants -= nodesRemoved - newParent, branchChildMerged := handleDeletion(branch, prefix) + newParent, branchChildMerged, err := handleDeletion(branch, prefix, deletedMerkleValues) + if err != nil { + return nil, 0, 0, false, fmt.Errorf("handling deletion: %w", err) + } + if branchChildMerged { nodesRemoved++ } - return newParent, valuesDeleted, nodesRemoved, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted, nil } func (t *Trie) clearPrefixLimitChild(branch *Node, prefix []byte, limit uint32, deletedMerkleValues map[string]struct{}) ( - newParent *Node, valuesDeleted, nodesRemoved uint32, allDeleted bool) { + newParent *Node, valuesDeleted, nodesRemoved uint32, allDeleted bool, err error) { newParent = branch childIndex := prefix[len(branch.Key)] @@ -817,51 +930,80 @@ func (t *Trie) clearPrefixLimitChild(branch *Node, prefix []byte, limit uint32, const valuesDeleted, nodesRemoved = 0, 0 // TODO ensure this is the same behaviour as in substrate allDeleted = true - return newParent, valuesDeleted, nodesRemoved, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted, nil } - child, valuesDeleted, nodesRemoved = t.deleteNodesLimit( + child, valuesDeleted, nodesRemoved, err = t.deleteNodesLimit( child, limit, deletedMerkleValues) + if err != nil { + // Note: do not wrap error since this is recursive. + return nil, 0, 0, false, err + } + if valuesDeleted == 0 { allDeleted = branch.Children[childIndex] == nil - return branch, valuesDeleted, nodesRemoved, allDeleted + return branch, valuesDeleted, nodesRemoved, allDeleted, nil } copySettings := node.DefaultCopySettings - branch = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + branch, err = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + if err != nil { + return nil, 0, 0, false, fmt.Errorf("preparing branch for mutation: %w", err) + } + branch.Children[childIndex] = child branch.Descendants -= nodesRemoved - newParent, branchChildMerged := handleDeletion(branch, prefix) + newParent, branchChildMerged, err := handleDeletion(branch, prefix, deletedMerkleValues) + if err != nil { + return nil, 0, 0, false, fmt.Errorf("handling deletion: %w", err) + } + if branchChildMerged { nodesRemoved++ } allDeleted = branch.Children[childIndex] == nil - return newParent, valuesDeleted, nodesRemoved, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted, nil } func (t *Trie) deleteNodesLimit(parent *Node, limit uint32, deletedMerkleValues map[string]struct{}) ( - newParent *Node, valuesDeleted, nodesRemoved uint32) { + newParent *Node, valuesDeleted, nodesRemoved uint32, err error) { if limit == 0 { valuesDeleted, nodesRemoved = 0, 0 - return parent, valuesDeleted, nodesRemoved + return parent, valuesDeleted, nodesRemoved, nil } if parent == nil { valuesDeleted, nodesRemoved = 0, 0 - return nil, valuesDeleted, nodesRemoved + return nil, valuesDeleted, nodesRemoved, nil } if parent.Kind() == node.Leaf { + isRoot := parent == t.root + err = registerDeletedMerkleValue(parent, isRoot, deletedMerkleValues) + if err != nil { + return nil, 0, 0, fmt.Errorf("registering deleted merkle value: %w", err) + } valuesDeleted, nodesRemoved = 1, 1 - return nil, valuesDeleted, nodesRemoved + return nil, valuesDeleted, nodesRemoved, nil } branch := parent nilChildren := node.ChildrenCapacity - branch.NumChildren() + if nilChildren == node.ChildrenCapacity { + panic("got branch with all nil children") + } + + // Note: there is at least one non-nil child and the limit isn't zero, + // therefore it is safe to prepare the branch for mutation. + copySettings := node.DefaultCopySettings + branch, err = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + if err != nil { + return nil, 0, 0, fmt.Errorf("preparing branch for mutation: %w", err) + } var newDeleted, newNodesRemoved uint32 var branchChildMerged bool @@ -870,14 +1012,13 @@ func (t *Trie) deleteNodesLimit(parent *Node, limit uint32, continue } - copySettings := node.DefaultCopySettings - branch = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) - - branch.Children[i], newDeleted, newNodesRemoved = t.deleteNodesLimit( + branch.Children[i], newDeleted, newNodesRemoved, err = t.deleteNodesLimit( child, limit, deletedMerkleValues) - // Note: newDeleted can never be zero here since the limit isn't zero - // and the child is not nil. Therefore it is safe to prepare the branch - // for mutation right before this call. + if err != nil { + // `deleteNodesLimit` is recursive, so do not wrap error. + return nil, 0, 0, err + } + if branch.Children[i] == nil { nilChildren++ } @@ -886,20 +1027,22 @@ func (t *Trie) deleteNodesLimit(parent *Node, limit uint32, nodesRemoved += newNodesRemoved branch.Descendants -= newNodesRemoved - branch.SetDirty() + newParent, branchChildMerged, err = handleDeletion(branch, branch.Key, deletedMerkleValues) + if err != nil { + return nil, 0, 0, fmt.Errorf("handling deletion: %w", err) + } - newParent, branchChildMerged = handleDeletion(branch, branch.Key) if branchChildMerged { nodesRemoved++ } if nilChildren == node.ChildrenCapacity && branch.SubValue == nil { - return nil, valuesDeleted, nodesRemoved + return nil, valuesDeleted, nodesRemoved, nil } if limit == 0 { - return newParent, valuesDeleted, nodesRemoved + return newParent, valuesDeleted, nodesRemoved, nil } } @@ -908,12 +1051,12 @@ func (t *Trie) deleteNodesLimit(parent *Node, limit uint32, valuesDeleted++ } - return nil, valuesDeleted, nodesRemoved + return nil, valuesDeleted, nodesRemoved, nil } // ClearPrefix deletes all nodes in the trie for which the key contains the // prefix given in little Endian format. -func (t *Trie) ClearPrefix(prefixLE []byte) { +func (t *Trie) ClearPrefix(prefixLE []byte) (err error) { pendingDeletedMerkleValues := make(map[string]struct{}) defer func() { const success = true @@ -921,6 +1064,14 @@ func (t *Trie) ClearPrefix(prefixLE []byte) { }() if len(prefixLE) == 0 { + const isRoot = true + err = ensureMerkleValueIsCalculated(t.root, isRoot) + if err != nil { + return fmt.Errorf("ensuring Merkle values are calculated: %w", err) + } + + const excludeInlined = true + PopulateMerkleValues(t.root, pendingDeletedMerkleValues, excludeInlined) t.root = nil return } @@ -928,25 +1079,39 @@ func (t *Trie) ClearPrefix(prefixLE []byte) { prefix := codec.KeyLEToNibbles(prefixLE) prefix = bytes.TrimSuffix(prefix, []byte{0}) - t.root, _ = t.clearPrefixAtNode(t.root, prefix, pendingDeletedMerkleValues) + t.root, _, err = t.clearPrefixAtNode(t.root, prefix, pendingDeletedMerkleValues) + if err != nil { + return fmt.Errorf("clearing prefix at root node: %w", err) + } + + return nil } func (t *Trie) clearPrefixAtNode(parent *Node, prefix []byte, deletedMerkleValues map[string]struct{}) ( - newParent *Node, nodesRemoved uint32) { + newParent *Node, nodesRemoved uint32, err error) { if parent == nil { const nodesRemoved = 0 - return nil, nodesRemoved + return nil, nodesRemoved, nil } if bytes.HasPrefix(parent.Key, prefix) { + isRoot := parent == t.root + err = ensureMerkleValueIsCalculated(parent, isRoot) + if err != nil { + nodesRemoved = 0 + return parent, nodesRemoved, fmt.Errorf("ensuring Merkle values are calculated: %w", err) + } + + const excludeInlined = true + PopulateMerkleValues(parent, deletedMerkleValues, excludeInlined) nodesRemoved = 1 + parent.Descendants - return nil, nodesRemoved + return nil, nodesRemoved, nil } if parent.Kind() == node.Leaf { const nodesRemoved = 0 - return parent, nodesRemoved + return parent, nodesRemoved, nil } branch := parent @@ -958,54 +1123,80 @@ func (t *Trie) clearPrefixAtNode(parent *Node, prefix []byte, if child == nil { const nodesRemoved = 0 - return parent, nodesRemoved + return parent, nodesRemoved, nil } nodesRemoved = 1 + child.Descendants copySettings := node.DefaultCopySettings - branch = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + branch, err = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + if err != nil { + return nil, 0, fmt.Errorf("preparing branch for mutation: %w", err) + } + + const isRoot = false // child so it cannot be the root + err = registerDeletedMerkleValue(child, isRoot, deletedMerkleValues) + if err != nil { + return nil, 0, fmt.Errorf("registering deleted merkle value for child: %w", err) + } + branch.Children[childIndex] = nil branch.Descendants -= nodesRemoved var branchChildMerged bool - newParent, branchChildMerged = handleDeletion(branch, prefix) + newParent, branchChildMerged, err = handleDeletion(branch, prefix, deletedMerkleValues) + if err != nil { + return nil, 0, fmt.Errorf("handling deletion: %w", err) + } + if branchChildMerged { nodesRemoved++ } - return newParent, nodesRemoved + return newParent, nodesRemoved, nil } noPrefixForNode := len(prefix) <= len(branch.Key) || lenCommonPrefix(branch.Key, prefix) < len(branch.Key) if noPrefixForNode { const nodesRemoved = 0 - return parent, nodesRemoved + return parent, nodesRemoved, nil } childIndex := prefix[len(branch.Key)] childPrefix := prefix[len(branch.Key)+1:] child := branch.Children[childIndex] - child, nodesRemoved = t.clearPrefixAtNode(child, childPrefix, deletedMerkleValues) - if nodesRemoved == 0 { - return parent, nodesRemoved + child, nodesRemoved, err = t.clearPrefixAtNode(child, childPrefix, deletedMerkleValues) + if err != nil { + nodesRemoved = 0 + // Note: do not wrap error since this is recursive + return parent, nodesRemoved, err + } else if nodesRemoved == 0 { + return parent, nodesRemoved, nil } copySettings := node.DefaultCopySettings - branch = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + branch, err = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + if err != nil { + return nil, 0, fmt.Errorf("preparing branch for mutation: %w", err) + } + branch.Descendants -= nodesRemoved branch.Children[childIndex] = child - newParent, branchChildMerged := handleDeletion(branch, prefix) + newParent, branchChildMerged, err := handleDeletion(branch, prefix, deletedMerkleValues) + if err != nil { + return nil, 0, fmt.Errorf("handling deletion: %w", err) + } + if branchChildMerged { nodesRemoved++ } - return newParent, nodesRemoved + return newParent, nodesRemoved, nil } // Delete removes the node of the trie with the key // matching the key given in little Endian format. // If no node is found at this key, nothing is deleted. -func (t *Trie) Delete(keyLE []byte) { +func (t *Trie) Delete(keyLE []byte) (err error) { pendingDeletedMerkleValues := make(map[string]struct{}) defer func() { const success = true @@ -1013,80 +1204,129 @@ func (t *Trie) Delete(keyLE []byte) { }() key := codec.KeyLEToNibbles(keyLE) - t.root, _, _ = t.deleteAtNode(t.root, key, pendingDeletedMerkleValues) + root, _, _, err := t.deleteAtNode(t.root, key, pendingDeletedMerkleValues) + if err != nil { + return err + } + t.root = root + return nil } func (t *Trie) deleteAtNode(parent *Node, key []byte, deletedMerkleValues map[string]struct{}) ( - newParent *Node, deleted bool, nodesRemoved uint32) { + newParent *Node, deleted bool, nodesRemoved uint32, err error) { if parent == nil { const nodesRemoved = 0 - return nil, false, nodesRemoved + return nil, false, nodesRemoved, nil } if parent.Kind() == node.Leaf { - if deleteLeaf(parent, key) == nil { + newParent, err = t.deleteLeaf(parent, key, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("deleting leaf: %w", err) + } + + if newParent == nil { const nodesRemoved = 1 - return nil, true, nodesRemoved + return nil, true, nodesRemoved, nil } const nodesRemoved = 0 - return parent, false, nodesRemoved + return parent, false, nodesRemoved, nil } - return t.deleteBranch(parent, key, deletedMerkleValues) + + newParent, deleted, nodesRemoved, err = t.deleteBranch(parent, key, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("deleting branch: %w", err) + } + + return newParent, deleted, nodesRemoved, nil } -func deleteLeaf(parent *Node, key []byte) (newParent *Node) { - if len(key) == 0 || bytes.Equal(key, parent.Key) { - return nil +func (t *Trie) deleteLeaf(parent *Node, key []byte, + deletedMerkleValues map[string]struct{}) ( + newParent *Node, err error) { + if len(key) > 0 && !bytes.Equal(key, parent.Key) { + return parent, nil + } + + newParent = nil + + isRoot := parent == t.root + err = registerDeletedMerkleValue(parent, isRoot, deletedMerkleValues) + if err != nil { + return nil, fmt.Errorf("registering deleted merkle value: %w", err) } - return parent + + return newParent, nil } func (t *Trie) deleteBranch(branch *Node, key []byte, deletedMerkleValues map[string]struct{}) ( - newParent *Node, deleted bool, nodesRemoved uint32) { + newParent *Node, deleted bool, nodesRemoved uint32, err error) { if len(key) == 0 || bytes.Equal(branch.Key, key) { copySettings := node.DefaultCopySettings copySettings.CopyValue = false - branch = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + branch, err = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing branch for mutation: %w", err) + } + // we need to set to nil if the branch has the same generation // as the current trie. branch.SubValue = nil deleted = true var branchChildMerged bool - newParent, branchChildMerged = handleDeletion(branch, key) + newParent, branchChildMerged, err = handleDeletion(branch, key, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("handling deletion: %w", err) + } + if branchChildMerged { nodesRemoved = 1 } - return newParent, deleted, nodesRemoved + return newParent, deleted, nodesRemoved, nil } commonPrefixLength := lenCommonPrefix(branch.Key, key) keyDoesNotExist := commonPrefixLength == len(key) if keyDoesNotExist { - return branch, false, 0 + return branch, false, 0, nil } childIndex := key[commonPrefixLength] childKey := key[commonPrefixLength+1:] child := branch.Children[childIndex] - newChild, deleted, nodesRemoved := t.deleteAtNode(child, childKey, deletedMerkleValues) + newChild, deleted, nodesRemoved, err := t.deleteAtNode(child, childKey, deletedMerkleValues) + if err != nil { + // deleteAtNode may call deleteBranch so don't wrap the error + // since this may be a recursive call. + return nil, false, 0, err + } + if !deleted { const nodesRemoved = 0 - return branch, false, nodesRemoved + return branch, false, nodesRemoved, nil } copySettings := node.DefaultCopySettings - branch = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + branch, err = t.prepBranchForMutation(branch, copySettings, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("preparing branch for mutation: %w", err) + } + branch.Descendants -= nodesRemoved branch.Children[childIndex] = newChild - newParent, branchChildMerged := handleDeletion(branch, key) + newParent, branchChildMerged, err := handleDeletion(branch, key, deletedMerkleValues) + if err != nil { + return nil, false, 0, fmt.Errorf("handling deletion: %w", err) + } + if branchChildMerged { nodesRemoved++ } - return newParent, true, nodesRemoved + return newParent, true, nodesRemoved, nil } // handleDeletion is called when a value is deleted from a branch to handle @@ -1095,7 +1335,9 @@ func (t *Trie) deleteBranch(branch *Node, key []byte, // In this first case, branchChildMerged is returned as true to keep track of the removal // of one node in callers. // If the branch has a value and no child, it will be changed into a leaf. -func handleDeletion(branch *Node, key []byte) (newNode *Node, branchChildMerged bool) { +func handleDeletion(branch *Node, key []byte, + deletedMerkleValues map[string]struct{}) ( + newNode *Node, branchChildMerged bool, err error) { childrenCount := 0 firstChildIndex := -1 for i, child := range branch.Children { @@ -1111,8 +1353,11 @@ func handleDeletion(branch *Node, key []byte) (newNode *Node, branchChildMerged switch { default: const branchChildMerged = false - return branch, branchChildMerged + return branch, branchChildMerged, nil case childrenCount == 0 && branch.SubValue != nil: + // The branch passed to handleDeletion is always a modified branch + // so the original branch Merkle value is already tracked in the deleted + // Merkle values map. const branchChildMerged = false commonPrefixLength := lenCommonPrefix(branch.Key, key) return &Node{ @@ -1120,11 +1365,19 @@ func handleDeletion(branch *Node, key []byte) (newNode *Node, branchChildMerged SubValue: branch.SubValue, Dirty: true, Generation: branch.Generation, - }, branchChildMerged + }, branchChildMerged, nil case childrenCount == 1 && branch.SubValue == nil: + // The branch passed to handleDeletion is always a modified branch + // so the original branch Merkle value is already tracked in the deleted + // Merkle values map. const branchChildMerged = true childIndex := firstChildIndex child := branch.Children[firstChildIndex] + const isRoot = false // child so it cannot be the root node + err = registerDeletedMerkleValue(child, isRoot, deletedMerkleValues) + if err != nil { + return nil, false, fmt.Errorf("registering deleted merkle value: %w", err) + } if child.Kind() == node.Leaf { newLeafKey := concatenateSlices(branch.Key, intToByteSlice(childIndex), child.Key) @@ -1133,7 +1386,7 @@ func handleDeletion(branch *Node, key []byte) (newNode *Node, branchChildMerged SubValue: child.SubValue, Dirty: true, Generation: branch.Generation, - }, branchChildMerged + }, branchChildMerged, nil } childBranch := child @@ -1157,8 +1410,32 @@ func handleDeletion(branch *Node, key []byte) (newNode *Node, branchChildMerged } } - return newBranch, branchChildMerged + return newBranch, branchChildMerged, nil + } +} + +// ensureMerkleValueIsCalculated is used before calling PopulateMerkleValues +// to ensure the parent node and all its descendant nodes have their Merkle +// value computed and ready to be used. This has a close to zero performance +// impact if the parent node Merkle value is already computed. +func ensureMerkleValueIsCalculated(parent *Node, isRoot bool) (err error) { + if parent == nil { + return nil } + + if isRoot { + _, err = parent.CalculateRootMerkleValue() + if err != nil { + return fmt.Errorf("calculating Merkle value of root node: %w", err) + } + } else { + _, err = parent.CalculateMerkleValue() + if err != nil { + return fmt.Errorf("calculating Merkle value of node: %w", err) + } + } + + return nil } // lenCommonPrefix returns the length of the diff --git a/lib/trie/trie_endtoend_test.go b/lib/trie/trie_endtoend_test.go index 90261d665a3..6d532331c9c 100644 --- a/lib/trie/trie_endtoend_test.go +++ b/lib/trie/trie_endtoend_test.go @@ -314,8 +314,15 @@ func TestTrieDiff(t *testing.T) { for _, test := range tests { newTrie.Put(test.key, test.value) } + deletedMerkleValues := newTrie.deletedMerkleValues - require.Len(t, deletedMerkleValues, 3) + expectedDeletedMerkleValues := map[string]struct{}{ + // root branch Merkle value which was modified (by its descendants). + // Other nodes result in an encoding of less than 32B so they are not + // tracked since they are inlined in the branch. + "\xa9v\xfaUme$<\x03\x80\x89\xd4\x15\r\xb1\x9a䶊`\xe5M\xeah\x9c\xab\xbf\xbb\xc0\xfcrH": {}, + } + assert.Equal(t, expectedDeletedMerkleValues, deletedMerkleValues) err = newTrie.WriteDirty(storageDB) require.NoError(t, err) @@ -837,7 +844,8 @@ func TestTrie_ClearPrefixLimit(t *testing.T) { trieClearPrefix.Put(test.key, test.value) } - num, allDeleted := trieClearPrefix.ClearPrefixLimit(prefix, uint32(lim)) + num, allDeleted, err := trieClearPrefix.ClearPrefixLimit(prefix, uint32(lim)) + require.NoError(t, err) deleteCount := uint32(0) isAllDeleted := true @@ -961,7 +969,8 @@ func TestTrie_ClearPrefixLimitSnapshot(t *testing.T) { require.Equal(t, tHash, dcTrieHash) require.Equal(t, dcTrieHash, ssTrieHash) - num, allDeleted := ssTrie.ClearPrefixLimit(prefix, uint32(lim)) + num, allDeleted, err := ssTrie.ClearPrefixLimit(prefix, uint32(lim)) + require.NoError(t, err) deleteCount := uint32(0) isAllDeleted := true diff --git a/lib/trie/trie_test.go b/lib/trie/trie_test.go index 0e9ef8d8adb..ef0c3c124a6 100644 --- a/lib/trie/trie_test.go +++ b/lib/trie/trie_test.go @@ -97,15 +97,20 @@ func Test_Trie_updateGeneration(t *testing.T) { t.Parallel() testCases := map[string]struct { - trieGeneration uint64 - node *Node - copySettings node.CopySettings - newNode *Node - copied bool - expectedDeletedMerkleValues map[string]struct{} + trie Trie + node *Node + pendingDeletedMerkleValues map[string]struct{} + copySettings node.CopySettings + newNode *Node + errSentinel error + errMessage string + copied bool + expectedPendingDeletedMerkleValues map[string]struct{} }{ - "trie generation higher and empty hash": { - trieGeneration: 2, + "update without registering deleted merkle value": { + trie: Trie{ + generation: 2, + }, node: &Node{ Generation: 1, Key: []byte{1}, @@ -115,24 +120,35 @@ func Test_Trie_updateGeneration(t *testing.T) { Generation: 2, Key: []byte{1}, }, - copied: true, - expectedDeletedMerkleValues: map[string]struct{}{}, + copied: true, }, - "trie generation higher and hash": { - trieGeneration: 2, + "update and register deleted Merkle value": { + trie: Trie{ + generation: 2, + }, + pendingDeletedMerkleValues: map[string]struct{}{}, node: &Node{ - Generation: 1, - Key: []byte{1}, - MerkleValue: []byte{1, 2, 3}, + Generation: 1, + Key: []byte{1}, + SubValue: []byte{ + 1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24, + 25, 26, 27, 28, 29, 30, 31, 32}, }, copySettings: node.DefaultCopySettings, newNode: &Node{ Generation: 2, Key: []byte{1}, + SubValue: []byte{ + 1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24, + 25, 26, 27, 28, 29, 30, 31, 32}, }, copied: true, - expectedDeletedMerkleValues: map[string]struct{}{ - string([]byte{1, 2, 3}): {}, + expectedPendingDeletedMerkleValues: map[string]struct{}{ + "\x98\xfc\xd6k\xa3\x12\u009e\xf1\x93\x05/\xd0\xc1Ln8\xb1X\xbd\\\x025\x06E\x94\xca\xcc\x1a\xb5\x96]": {}, }, }, } @@ -142,13 +158,18 @@ func Test_Trie_updateGeneration(t *testing.T) { t.Run(name, func(t *testing.T) { t.Parallel() - deletedMerkleValues := make(map[string]struct{}) + expectedTrie := *testCase.trie.DeepCopy() - newNode := updateGeneration(testCase.node, testCase.trieGeneration, - deletedMerkleValues, testCase.copySettings) + newNode, err := testCase.trie.updateGeneration( + testCase.node, testCase.pendingDeletedMerkleValues, testCase.copySettings) + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } assert.Equal(t, testCase.newNode, newNode) - assert.Equal(t, testCase.expectedDeletedMerkleValues, deletedMerkleValues) + assert.Equal(t, expectedTrie, testCase.trie) + assert.Equal(t, testCase.expectedPendingDeletedMerkleValues, testCase.pendingDeletedMerkleValues) // Check for deep copy if newNode != nil && testCase.copied { @@ -163,6 +184,66 @@ func Test_Trie_updateGeneration(t *testing.T) { } } +func Test_registerDeletedMerkleValue(t *testing.T) { + t.Parallel() + + someSmallNode := &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + } + + testCases := map[string]struct { + node *Node + isRoot bool + pendingDeletedMerkleValues map[string]struct{} + expectedPendingDeletedMerkleValues map[string]struct{} + }{ + "dirty node not registered": { + node: &Node{Dirty: true}, + }, + "clean root node registered": { + node: someSmallNode, + isRoot: true, + pendingDeletedMerkleValues: map[string]struct{}{}, + expectedPendingDeletedMerkleValues: map[string]struct{}{ + "`Qm\v\xb6\xe1\xbb\xfb\x12\x93\xf1\xb2v\xea\x95\x05\xe9\xf4\xa4\xe7ُb\r\x05\x11^\v\x85'J\xe1": {}, + }, + }, + "clean node with inlined Merkle value not registered": { + node: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + }, + }, + "clean node with hash Merkle value registered": { + node: &Node{ + Key: []byte{1}, + SubValue: []byte{ + 1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24, + 25, 26, 27, 28, 29, 30, 31, 32}, + }, + pendingDeletedMerkleValues: map[string]struct{}{}, + expectedPendingDeletedMerkleValues: map[string]struct{}{ + "\x98\xfc\xd6k\xa3\x12\u009e\xf1\x93\x05/\xd0\xc1Ln8\xb1X\xbd\\\x025\x06E\x94\xca\xcc\x1a\xb5\x96]": {}, + }, + }, + } + + for name, testCase := range testCases { + testCase := testCase + t.Run(name, func(t *testing.T) { + t.Parallel() + + err := registerDeletedMerkleValue(testCase.node, testCase.isRoot, testCase.pendingDeletedMerkleValues) + + require.NoError(t, err) + assert.Equal(t, testCase.expectedPendingDeletedMerkleValues, testCase.pendingDeletedMerkleValues) + }) + } +} + func getPointer(x interface{}) (pointer uintptr, ok bool) { func() { defer func() { @@ -1007,7 +1088,8 @@ func Test_Trie_Put(t *testing.T) { }{ "trie with key and value": { trie: Trie{ - generation: 1, + generation: 1, + deletedMerkleValues: map[string]struct{}{}, root: &Node{ Key: []byte{1, 2, 0, 5}, SubValue: []byte{1}, @@ -1017,6 +1099,9 @@ func Test_Trie_Put(t *testing.T) { value: []byte{2}, expectedTrie: Trie{ generation: 1, + deletedMerkleValues: map[string]struct{}{ + "\xa1\x95\b\x9c>\x8f\x8b[6\x97\x87\x00\xad\x95J\xed\x99\xe0\x84\x13\xcf\xc1\xe2\xb4\xc0\n]\x06J\xbef\xa9": {}, + }, root: &Node{ Key: []byte{1, 2}, Generation: 1, @@ -1110,7 +1195,12 @@ func Test_Trie_insert(t *testing.T) { Generation: 1, Dirty: true, }, - {Key: []byte{2}, SubValue: []byte{1}}, + { + Key: []byte{2}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x02, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x02, 0x04, 0x01}, + }, }), }, mutated: true, @@ -1268,10 +1358,11 @@ func Test_Trie_insert(t *testing.T) { trie := testCase.trie expectedTrie := *trie.DeepCopy() - newNode, mutated, nodesCreated := trie.insert( + newNode, mutated, nodesCreated, err := trie.insert( testCase.parent, testCase.key, testCase.value, testCase.deletedMerkleValues) + require.NoError(t, err) assert.Equal(t, testCase.newNode, newNode) assert.Equal(t, testCase.mutated, mutated) assert.Equal(t, testCase.nodesCreated, nodesCreated) @@ -1284,13 +1375,16 @@ func Test_Trie_insertInBranch(t *testing.T) { t.Parallel() testCases := map[string]struct { - parent *Node - key []byte - value []byte - deletedMerkleValues map[string]struct{} - newNode *Node - mutated bool - nodesCreated uint32 + parent *Node + key []byte + value []byte + deletedMerkleValues map[string]struct{} + newNode *Node + mutated bool + nodesCreated uint32 + errSentinel error + errMessage string + expectedDeletedMerkleValues map[string]struct{} }{ "insert existing value to branch": { parent: &Node{ @@ -1562,14 +1656,19 @@ func Test_Trie_insertInBranch(t *testing.T) { trie := new(Trie) - newNode, mutated, nodesCreated := trie.insertInBranch( + newNode, mutated, nodesCreated, err := trie.insertInBranch( testCase.parent, testCase.key, testCase.value, testCase.deletedMerkleValues) + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } assert.Equal(t, testCase.newNode, newNode) assert.Equal(t, testCase.mutated, mutated) assert.Equal(t, testCase.nodesCreated, nodesCreated) assert.Equal(t, new(Trie), trie) // check no mutation + assert.Equal(t, testCase.expectedDeletedMerkleValues, testCase.deletedMerkleValues) }) } } @@ -2121,6 +2220,8 @@ func Test_Trie_ClearPrefixLimit(t *testing.T) { limit uint32 deleted uint32 allDeleted bool + errSentinel error + errMessage string expectedTrie Trie }{ "limit is zero": {}, @@ -2153,8 +2254,12 @@ func Test_Trie_ClearPrefixLimit(t *testing.T) { trie := testCase.trie - deleted, allDeleted := trie.ClearPrefixLimit(testCase.prefix, testCase.limit) + deleted, allDeleted, err := trie.ClearPrefixLimit(testCase.prefix, testCase.limit) + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } assert.Equal(t, testCase.deleted, deleted) assert.Equal(t, testCase.allDeleted, allDeleted) assert.Equal(t, testCase.expectedTrie, trie) @@ -2175,6 +2280,8 @@ func Test_Trie_clearPrefixLimitAtNode(t *testing.T) { valuesDeleted uint32 nodesRemoved uint32 allDeleted bool + errSentinel error + errMessage string }{ "limit is zero": { allDeleted: true, @@ -2464,7 +2571,12 @@ func Test_Trie_clearPrefixLimitAtNode(t *testing.T) { Descendants: 1, Children: padRightChildren([]*Node{ nil, - {Key: []byte{4}, SubValue: []byte{1}}, + { + Key: []byte{4}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x04, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x04, 0x04, 0x01}, + }, }), }, valuesDeleted: 1, @@ -2529,6 +2641,7 @@ func Test_Trie_clearPrefixLimitAtNode(t *testing.T) { nodesRemoved: 3, allDeleted: true, }, + "partially delete child of branch": { trie: Trie{ generation: 1, @@ -2574,6 +2687,8 @@ func Test_Trie_clearPrefixLimitAtNode(t *testing.T) { Key: []byte{6}, SubValue: []byte{1}, // Not modified so same generation as before + Encoding: []byte{0x41, 0x06, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x06, 0x04, 0x01}, }, }), }, @@ -2691,10 +2806,14 @@ func Test_Trie_clearPrefixLimitAtNode(t *testing.T) { trie := testCase.trie expectedTrie := *trie.DeepCopy() - newParent, valuesDeleted, nodesRemoved, allDeleted := + newParent, valuesDeleted, nodesRemoved, allDeleted, err := trie.clearPrefixLimitAtNode(testCase.parent, testCase.prefix, testCase.limit, testCase.deletedMerkleValues) + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } assert.Equal(t, testCase.newParent, newParent) assert.Equal(t, testCase.valuesDeleted, valuesDeleted) assert.Equal(t, testCase.nodesRemoved, nodesRemoved) @@ -2715,6 +2834,8 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { newNode *Node valuesDeleted uint32 nodesRemoved uint32 + errSentinel error + errMessage string }{ "zero limit": { trie: Trie{ @@ -2800,7 +2921,12 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { Descendants: 1, Children: padRightChildren([]*Node{ nil, - {Key: []byte{2}, SubValue: []byte{1}}, + { + Key: []byte{2}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x02, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x02, 0x04, 0x01}, + }, }), }, valuesDeleted: 1, @@ -2865,10 +2991,14 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { trie := testCase.trie expectedTrie := *trie.DeepCopy() - newNode, valuesDeleted, nodesRemoved := + newNode, valuesDeleted, nodesRemoved, err := trie.deleteNodesLimit(testCase.parent, testCase.limit, testCase.deletedMerkleValues) + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } assert.Equal(t, testCase.newNode, newNode) assert.Equal(t, testCase.valuesDeleted, valuesDeleted) assert.Equal(t, testCase.nodesRemoved, nodesRemoved) @@ -2887,14 +3017,26 @@ func Test_Trie_ClearPrefix(t *testing.T) { }{ "nil prefix": { trie: Trie{ - root: &Node{SubValue: []byte{1}}, + root: &Node{SubValue: []byte{1}}, + deletedMerkleValues: map[string]struct{}{}, + }, + expectedTrie: Trie{ + deletedMerkleValues: map[string]struct{}{ + "\xf9jt\x15\"\xbc\xc1O\n\xea/p`DR$\x1dY\xb5\xf2ݫ\x9aiH\xfd\xb3\xfe\xf5\xf9\x86C": {}, + }, }, }, "empty prefix": { trie: Trie{ - root: &Node{SubValue: []byte{1}}, + root: &Node{SubValue: []byte{1}}, + deletedMerkleValues: map[string]struct{}{}, }, prefix: []byte{}, + expectedTrie: Trie{ + deletedMerkleValues: map[string]struct{}{ + "\xf9jt\x15\"\xbc\xc1O\n\xea/p`DR$\x1dY\xb5\xf2ݫ\x9aiH\xfd\xb3\xfe\xf5\xf9\x86C": {}, + }, + }, }, "empty trie": { prefix: []byte{0x12}, @@ -2963,6 +3105,7 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { deletedMerkleValues map[string]struct{} newParent *Node nodesRemoved uint32 + expectedTrie Trie }{ "delete one of two children of branch": { trie: Trie{ @@ -2984,6 +3127,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { Generation: 1, }, nodesRemoved: 2, + expectedTrie: Trie{ + generation: 1, + }, }, "nil parent": {}, "leaf parent with common prefix": { @@ -3015,6 +3161,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { Key: []byte{1, 2}, SubValue: []byte{1}, }, + expectedTrie: Trie{ + generation: 1, + }, }, "leaf parent with key smaller than prefix": { trie: Trie{ @@ -3029,6 +3178,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { Key: []byte{1}, SubValue: []byte{1}, }, + expectedTrie: Trie{ + generation: 1, + }, }, "branch parent with common prefix": { parent: &Node{ @@ -3075,6 +3227,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { {}, }), }, + expectedTrie: Trie{ + generation: 1, + }, }, "branch with key smaller than prefix by more than one": { trie: Trie{ @@ -3097,6 +3252,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { {}, }), }, + expectedTrie: Trie{ + generation: 1, + }, }, "branch with key smaller than prefix by one": { trie: Trie{ @@ -3119,6 +3277,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { {}, }), }, + expectedTrie: Trie{ + generation: 1, + }, }, "delete one child of branch": { trie: Trie{ @@ -3142,10 +3303,18 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { Descendants: 1, Children: padRightChildren([]*Node{ nil, - {Key: []byte{4}, SubValue: []byte{1}}, + { + Key: []byte{4}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x04, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x04, 0x04, 0x01}, + }, }), }, nodesRemoved: 1, + expectedTrie: Trie{ + generation: 1, + }, }, "fully delete child of branch": { trie: Trie{ @@ -3167,6 +3336,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { Generation: 1, }, nodesRemoved: 1, + expectedTrie: Trie{ + generation: 1, + }, }, "partially delete child of branch": { trie: Trie{ @@ -3207,6 +3379,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { }), }, nodesRemoved: 1, + expectedTrie: Trie{ + generation: 1, + }, }, "delete one of two children of branch without value": { trie: Trie{ @@ -3228,6 +3403,9 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { Generation: 1, }, nodesRemoved: 2, + expectedTrie: Trie{ + generation: 1, + }, }, } @@ -3237,14 +3415,14 @@ func Test_Trie_clearPrefixAtNode(t *testing.T) { t.Parallel() trie := testCase.trie - expectedTrie := *trie.DeepCopy() - newParent, nodesRemoved := trie.clearPrefixAtNode( + newParent, nodesRemoved, err := trie.clearPrefixAtNode( testCase.parent, testCase.prefix, testCase.deletedMerkleValues) + require.NoError(t, err) assert.Equal(t, testCase.newParent, newParent) assert.Equal(t, testCase.nodesRemoved, nodesRemoved) - assert.Equal(t, expectedTrie, trie) + assert.Equal(t, testCase.expectedTrie, trie) }) } } @@ -3259,12 +3437,24 @@ func Test_Trie_Delete(t *testing.T) { }{ "nil key": { trie: Trie{ - root: &Node{SubValue: []byte{1}}, + root: &Node{SubValue: []byte{1}}, + deletedMerkleValues: map[string]struct{}{}, + }, + expectedTrie: Trie{ + deletedMerkleValues: map[string]struct{}{ + "\xf9jt\x15\"\xbc\xc1O\n\xea/p`DR$\x1dY\xb5\xf2ݫ\x9aiH\xfd\xb3\xfe\xf5\xf9\x86C": {}, + }, }, }, "empty key": { trie: Trie{ - root: &Node{SubValue: []byte{1}}, + root: &Node{SubValue: []byte{1}}, + deletedMerkleValues: map[string]struct{}{}, + }, + expectedTrie: Trie{ + deletedMerkleValues: map[string]struct{}{ + "\xf9jt\x15\"\xbc\xc1O\n\xea/p`DR$\x1dY\xb5\xf2ݫ\x9aiH\xfd\xb3\xfe\xf5\xf9\x86C": {}, + }, }, }, "empty trie": { @@ -3294,6 +3484,7 @@ func Test_Trie_Delete(t *testing.T) { }, }), }, + deletedMerkleValues: map[string]struct{}{}, }, key: []byte{0x12, 0x16}, expectedTrie: Trie{ @@ -3305,8 +3496,10 @@ func Test_Trie_Delete(t *testing.T) { Descendants: 2, Children: padRightChildren([]*Node{ { - Key: []byte{5}, - SubValue: []byte{97}, + Key: []byte{5}, + SubValue: []byte{97}, + Encoding: []byte{0x41, 0x05, 0x04, 0x61}, + MerkleValue: []byte{0x41, 0x05, 0x04, 0x61}, }, { // full key in nibbles 1, 2, 1, 6 Key: []byte{6, 0, 7}, @@ -3316,6 +3509,9 @@ func Test_Trie_Delete(t *testing.T) { }, }), }, + deletedMerkleValues: map[string]struct{}{ + "=\x1b=r~\xe4\x04T\x9a]%1\xaa\xb9\xff\xf0\xee\xddŋ\xc3\v\xfe/\xe8+\x1a\f\xfe~v\xd5": {}, + }, }, }, } @@ -3351,6 +3547,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { newParent *Node updated bool nodesRemoved uint32 + errSentinel error + errMessage string + expectedTrie Trie }{ "nil parent": { key: []byte{1}, @@ -3394,6 +3593,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { Key: []byte{1}, SubValue: []byte{1}, }, + expectedTrie: Trie{ + generation: 1, + }, }, "branch parent and nil key": { trie: Trie{ @@ -3418,6 +3620,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { }, updated: true, nodesRemoved: 1, + expectedTrie: Trie{ + generation: 1, + }, }, "branch parent and empty key": { trie: Trie{ @@ -3440,6 +3645,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { }, updated: true, nodesRemoved: 1, + expectedTrie: Trie{ + generation: 1, + }, }, "branch parent matches key": { trie: Trie{ @@ -3462,6 +3670,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { }, updated: true, nodesRemoved: 1, + expectedTrie: Trie{ + generation: 1, + }, }, "branch parent child matches key": { trie: Trie{ @@ -3487,6 +3698,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { }, updated: true, nodesRemoved: 1, + expectedTrie: Trie{ + generation: 1, + }, }, "branch parent mismatches key": { trie: Trie{ @@ -3509,6 +3723,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { {}, }), }, + expectedTrie: Trie{ + generation: 1, + }, }, "branch parent child mismatches key": { trie: Trie{ @@ -3537,6 +3754,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { }, }), }, + expectedTrie: Trie{ + generation: 1, + }, }, "delete branch child and merge branch and left child": { trie: Trie{ @@ -3565,6 +3785,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { }, updated: true, nodesRemoved: 2, + expectedTrie: Trie{ + generation: 1, + }, }, "delete branch and keep two children": { trie: Trie{ @@ -3586,11 +3809,24 @@ func Test_Trie_deleteAtNode(t *testing.T) { Dirty: true, Descendants: 2, Children: padRightChildren([]*Node{ - {Key: []byte{2}, SubValue: []byte{1}}, - {Key: []byte{2}, SubValue: []byte{1}}, + { + Key: []byte{2}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x02, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x02, 0x04, 0x01}, + }, + { + Key: []byte{2}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x02, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x02, 0x04, 0x01}, + }, }), }, updated: true, + expectedTrie: Trie{ + generation: 1, + }, }, "handle nonexistent key (no op)": { trie: Trie{ @@ -3625,6 +3861,9 @@ func Test_Trie_deleteAtNode(t *testing.T) { }, }), }, + expectedTrie: Trie{ + generation: 1, + }, }, } @@ -3639,15 +3878,18 @@ func Test_Trie_deleteAtNode(t *testing.T) { expectedKey = make([]byte, len(testCase.key)) copy(expectedKey, testCase.key) } - expectedTrie := *testCase.trie.DeepCopy() - newParent, updated, nodesRemoved := testCase.trie.deleteAtNode( + newParent, updated, nodesRemoved, err := testCase.trie.deleteAtNode( testCase.parent, testCase.key, testCase.deletedMerkleValues) + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } assert.Equal(t, testCase.newParent, newParent) assert.Equal(t, testCase.updated, updated) assert.Equal(t, testCase.nodesRemoved, nodesRemoved) - assert.Equal(t, expectedTrie, testCase.trie) + assert.Equal(t, testCase.expectedTrie, testCase.trie) assert.Equal(t, expectedKey, testCase.key) }) } @@ -3657,10 +3899,14 @@ func Test_handleDeletion(t *testing.T) { t.Parallel() testCases := map[string]struct { - branch *Node - deletedKey []byte - newNode *Node - branchChildMerged bool + branch *Node + deletedKey []byte + deletedMerkleValues map[string]struct{} + newNode *Node + branchChildMerged bool + errSentinel error + errMessage string + expectedDeletedMerkleValues map[string]struct{} }{ "branch with value and without children": { branch: &Node{ @@ -3743,9 +3989,19 @@ func Test_handleDeletion(t *testing.T) { Generation: 1, Dirty: true, Children: padRightChildren([]*Node{ - {Key: []byte{7}, SubValue: []byte{1}}, + { + Key: []byte{7}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x07, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x07, 0x04, 0x01}, + }, nil, - {Key: []byte{8}, SubValue: []byte{1}}, + { + Key: []byte{8}, + SubValue: []byte{1}, + Encoding: []byte{0x41, 0x08, 0x04, 0x01}, + MerkleValue: []byte{0x41, 0x08, 0x04, 0x01}, + }, }), }, branchChildMerged: true, @@ -3764,11 +4020,135 @@ func Test_handleDeletion(t *testing.T) { copy(expectedKey, testCase.deletedKey) } - newNode, branchChildMerged := handleDeletion(testCase.branch, testCase.deletedKey) + newNode, branchChildMerged, err := handleDeletion( + testCase.branch, testCase.deletedKey, testCase.deletedMerkleValues) + + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } assert.Equal(t, testCase.newNode, newNode) assert.Equal(t, testCase.branchChildMerged, branchChildMerged) assert.Equal(t, expectedKey, testCase.deletedKey) + assert.Equal(t, testCase.expectedDeletedMerkleValues, testCase.deletedMerkleValues) + }) + } +} + +func Test_ensureMerkleValueIsCalculated(t *testing.T) { + t.Parallel() + + testCases := map[string]struct { + parent *Node + isRoot bool + errSentinel error + errMessage string + expectedNode *Node + }{ + "nil parent": {}, + "root node without Merkle value": { + parent: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + }, + isRoot: true, + expectedNode: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + Encoding: []byte{0x41, 0x01, 0x04, 0x02}, + MerkleValue: []byte{ + 0x60, 0x51, 0x6d, 0xb, 0xb6, 0xe1, 0xbb, 0xfb, + 0x12, 0x93, 0xf1, 0xb2, 0x76, 0xea, 0x95, 0x5, + 0xe9, 0xf4, 0xa4, 0xe7, 0xd9, 0x8f, 0x62, 0xd, + 0x5, 0x11, 0x5e, 0xb, 0x85, 0x27, 0x4a, 0xe1}, + }, + }, + "root node with inlined Merkle value": { + parent: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + MerkleValue: []byte{3}, + }, + isRoot: true, + expectedNode: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + Encoding: []byte{0x41, 0x01, 0x04, 0x02}, + MerkleValue: []byte{ + 0x60, 0x51, 0x6d, 0xb, 0xb6, 0xe1, 0xbb, 0xfb, + 0x12, 0x93, 0xf1, 0xb2, 0x76, 0xea, 0x95, 0x5, + 0xe9, 0xf4, 0xa4, 0xe7, 0xd9, 0x8f, 0x62, 0xd, + 0x5, 0x11, 0x5e, 0xb, 0x85, 0x27, 0x4a, 0xe1}, + }, + }, + "root node with hash Merkle value": { + parent: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + MerkleValue: []byte{ + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8}, + }, + isRoot: true, + expectedNode: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + MerkleValue: []byte{ + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8}, + }, + }, + "non root node without Merkle value": { + parent: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + }, + expectedNode: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + Encoding: []byte{0x41, 0x01, 0x04, 0x02}, + MerkleValue: []byte{0x41, 0x1, 0x4, 0x2}, + }, + }, + "non root node with Merkle value": { + parent: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + MerkleValue: []byte{3}, + }, + expectedNode: &Node{ + Key: []byte{1}, + SubValue: []byte{2}, + MerkleValue: []byte{3}, + }, + }, + } + + for name, testCase := range testCases { + testCase := testCase + t.Run(name, func(t *testing.T) { + t.Parallel() + + err := ensureMerkleValueIsCalculated(testCase.parent, testCase.isRoot) + + // Ensure all Merkle values are set in the trie rooted + // at the parent node. + assert.NotPanics(t, func() { + m := map[string]struct{}{} + const excludeInlined = false + PopulateMerkleValues(testCase.parent, m, excludeInlined) + }) + + assert.ErrorIs(t, err, testCase.errSentinel) + if testCase.errSentinel != nil { + assert.EqualError(t, err, testCase.errMessage) + } + assert.Equal(t, testCase.expectedNode, testCase.parent) }) } } @@ -3915,4 +4295,12 @@ func Benchmark_concatSlices(b *testing.B) { concatenated[0] = 1 } }) + + // 16453 ns/op 204800 B/op 1 allocs/op + b.Run("bytes.Join", func(b *testing.B) { + for i := 0; i < b.N; i++ { + concatenated := bytes.Join([][]byte{slice1, slice2}, nil) + concatenated[0] = 1 + } + }) }