diff --git a/internal/trie/node/branch.go b/internal/trie/node/branch.go index c8cea365bad..0bbf0554440 100644 --- a/internal/trie/node/branch.go +++ b/internal/trie/node/branch.go @@ -30,6 +30,12 @@ type Branch struct { // which is updated to match the trie Generation once they are // inserted, moved or iterated over. Generation uint64 + + // Statistics + + // Descendants is the number of descendant nodes for + // this particular node. + Descendants uint32 } // NewBranch creates a new branch using the arguments given. @@ -62,6 +68,7 @@ func (b *Branch) StringNode() (stringNode *gotree.Node) { stringNode.Appendf("Dirty: %t", b.Dirty) stringNode.Appendf("Key: " + bytesToString(b.Key)) stringNode.Appendf("Value: " + bytesToString(b.Value)) + stringNode.Appendf("Descendants: %d", b.Descendants) stringNode.Appendf("Calculated encoding: " + bytesToString(b.Encoding)) stringNode.Appendf("Calculated digest: " + bytesToString(b.HashDigest)) diff --git a/internal/trie/node/branch_test.go b/internal/trie/node/branch_test.go index 2e890b5a0c6..8e3d56e6bd7 100644 --- a/internal/trie/node/branch_test.go +++ b/internal/trie/node/branch_test.go @@ -82,14 +82,16 @@ func Test_Branch_String(t *testing.T) { ├── Dirty: false ├── Key: nil ├── Value: nil +├── Descendants: 0 ├── Calculated encoding: nil └── Calculated digest: nil`, }, "branch with value smaller than 1024": { branch: &Branch{ - Key: []byte{1, 2}, - Value: []byte{3, 4}, - Dirty: true, + Key: []byte{1, 2}, + Value: []byte{3, 4}, + Dirty: true, + Descendants: 3, Children: [16]Node{ nil, nil, nil, &Leaf{}, @@ -105,6 +107,7 @@ func Test_Branch_String(t *testing.T) { ├── Dirty: true ├── Key: 0x0102 ├── Value: 0x0304 +├── Descendants: 3 ├── Calculated encoding: nil ├── Calculated digest: nil ├── Child 3 @@ -121,6 +124,7 @@ func Test_Branch_String(t *testing.T) { | ├── Dirty: false | ├── Key: nil | ├── Value: nil +| ├── Descendants: 0 | ├── Calculated encoding: nil | └── Calculated digest: nil └── Child 11 @@ -134,9 +138,10 @@ func Test_Branch_String(t *testing.T) { }, "branch with value higher than 1024": { branch: &Branch{ - Key: []byte{1, 2}, - Value: make([]byte, 1025), - Dirty: true, + Key: []byte{1, 2}, + Value: make([]byte, 1025), + Dirty: true, + Descendants: 3, Children: [16]Node{ nil, nil, nil, &Leaf{}, @@ -152,6 +157,7 @@ func Test_Branch_String(t *testing.T) { ├── Dirty: true ├── Key: 0x0102 ├── Value: 0x0000000000000000...0000000000000000 +├── Descendants: 3 ├── Calculated encoding: nil ├── Calculated digest: nil ├── Child 3 @@ -168,6 +174,7 @@ func Test_Branch_String(t *testing.T) { | ├── Dirty: false | ├── Key: nil | ├── Value: nil +| ├── Descendants: 0 | ├── Calculated encoding: nil | └── Calculated digest: nil └── Child 11 diff --git a/internal/trie/node/copy.go b/internal/trie/node/copy.go index c69161e5f89..9a26ff6b757 100644 --- a/internal/trie/node/copy.go +++ b/internal/trie/node/copy.go @@ -8,8 +8,9 @@ package node // children as well. func (b *Branch) Copy(copyChildren bool) Node { cpy := &Branch{ - Dirty: b.Dirty, - Generation: b.Generation, + Dirty: b.Dirty, + Generation: b.Generation, + Descendants: b.GetDescendants(), } if copyChildren { diff --git a/internal/trie/node/decode.go b/internal/trie/node/decode.go index c67e84e49af..c08a63dc13f 100644 --- a/internal/trie/node/decode.go +++ b/internal/trie/node/decode.go @@ -98,6 +98,8 @@ func decodeBranch(reader io.Reader, header byte) (branch *Branch, err error) { if (childrenBitmap[i/8]>>(i%8))&1 != 1 { continue } + branch.AddDescendants(1) + var hash []byte err := sd.Decode(&hash) if err != nil { diff --git a/internal/trie/node/decode_test.go b/internal/trie/node/decode_test.go index c6f683aece3..207208b627f 100644 --- a/internal/trie/node/decode_test.go +++ b/internal/trie/node/decode_test.go @@ -176,7 +176,8 @@ func Test_decodeBranch(t *testing.T) { HashDigest: []byte{1, 2, 3, 4, 5}, }, }, - Dirty: true, + Dirty: true, + Descendants: 1, }, }, "value decoding error for node type 3": { @@ -211,7 +212,8 @@ func Test_decodeBranch(t *testing.T) { HashDigest: []byte{1, 2, 3, 4, 5}, }, }, - Dirty: true, + Dirty: true, + Descendants: 1, }, }, } diff --git a/internal/trie/node/encode_decode_test.go b/internal/trie/node/encode_decode_test.go index f4bb1686971..fe174a67aa9 100644 --- a/internal/trie/node/encode_decode_test.go +++ b/internal/trie/node/encode_decode_test.go @@ -60,7 +60,8 @@ func Test_Branch_Encode_Decode(t *testing.T) { HashDigest: []byte{0x41, 0x9, 0x4, 0xa}, }, }, - Dirty: true, + Dirty: true, + Descendants: 1, }, }, } diff --git a/internal/trie/node/stats.go b/internal/trie/node/stats.go new file mode 100644 index 00000000000..61ad8045ccc --- /dev/null +++ b/internal/trie/node/stats.go @@ -0,0 +1,19 @@ +// Copyright 2022 ChainSafe Systems (ON) +// SPDX-License-Identifier: LGPL-3.0-only + +package node + +// GetDescendants returns the number of descendants in the branch. +func (b *Branch) GetDescendants() (descendants uint32) { + return b.Descendants +} + +// AddDescendants adds descendant nodes count to the node stats. +func (b *Branch) AddDescendants(n uint32) { + b.Descendants += n +} + +// SubDescendants subtracts descendant nodes count from the node stats. +func (b *Branch) SubDescendants(n uint32) { + b.Descendants -= n +} diff --git a/internal/trie/node/stats_test.go b/internal/trie/node/stats_test.go new file mode 100644 index 00000000000..d96563d0dc5 --- /dev/null +++ b/internal/trie/node/stats_test.go @@ -0,0 +1,60 @@ +// Copyright 2022 ChainSafe Systems (ON) +// SPDX-License-Identifier: LGPL-3.0-only + +package node + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func Test_Branch_GetDescendants(t *testing.T) { + t.Parallel() + + const descendants uint32 = 10 + branch := &Branch{ + Descendants: descendants, + } + result := branch.GetDescendants() + + assert.Equal(t, descendants, result) +} + +func Test_Branch_AddDescendants(t *testing.T) { + t.Parallel() + + const ( + initialDescendants uint32 = 10 + addDescendants uint32 = 2 + finalDescendants uint32 = 12 + ) + branch := &Branch{ + Descendants: initialDescendants, + } + branch.AddDescendants(addDescendants) + expected := &Branch{ + Descendants: finalDescendants, + } + + assert.Equal(t, expected, branch) +} + +func Test_Branch_SubDescendants(t *testing.T) { + t.Parallel() + + const ( + initialDescendants uint32 = 10 + subDescendants uint32 = 2 + finalDescendants uint32 = 8 + ) + branch := &Branch{ + Descendants: initialDescendants, + } + branch.SubDescendants(subDescendants) + expected := &Branch{ + Descendants: finalDescendants, + } + + assert.Equal(t, expected, branch) +} diff --git a/lib/trie/print_test.go b/lib/trie/print_test.go index 8cb47ba56ce..cf2343569b9 100644 --- a/lib/trie/print_test.go +++ b/lib/trie/print_test.go @@ -39,8 +39,9 @@ func Test_Trie_String(t *testing.T) { "branch root": { trie: Trie{ root: &node.Branch{ - Key: nil, - Value: []byte{1, 2}, + Key: nil, + Value: []byte{1, 2}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ Key: []byte{1, 2, 3}, @@ -61,6 +62,7 @@ func Test_Trie_String(t *testing.T) { ├── Dirty: false ├── Key: nil ├── Value: 0x0102 +├── Descendants: 2 ├── Calculated encoding: nil ├── Calculated digest: nil ├── Child 0 diff --git a/lib/trie/trie.go b/lib/trie/trie.go index 0b81cca93cb..97a09b9071b 100644 --- a/lib/trie/trie.go +++ b/lib/trie/trie.go @@ -326,19 +326,20 @@ func (t *Trie) Put(keyLE, value []byte) { } func (t *Trie) put(key, value []byte) { - t.root = t.insert(t.root, key, value) + t.root, _ = t.insert(t.root, key, value) } // 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) (newParent Node) { +func (t *Trie) insert(parent Node, key, value []byte) (newParent Node, nodesCreated uint32) { if parent == nil { + const nodesCreated = 1 return &node.Leaf{ Key: key, Value: value, Generation: t.generation, Dirty: true, - } + }, nodesCreated } // TODO ensure all values have dirty set to true @@ -353,16 +354,17 @@ func (t *Trie) insert(parent Node, key, value []byte) (newParent Node) { } } -func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, - value []byte) (newParent Node) { +func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, value []byte) ( + newParent Node, nodesCreated uint32) { if bytes.Equal(parentLeaf.Key, key) { + nodesCreated = 0 if bytes.Equal(value, parentLeaf.Value) { - return parentLeaf + return parentLeaf, nodesCreated } parentLeaf = t.prepLeafForMutation(parentLeaf) parentLeaf.Value = value - return parentLeaf + return parentLeaf, nodesCreated } commonPrefixLength := lenCommonPrefix(key, parentLeaf.Key) @@ -385,9 +387,11 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, childIndex := parentLeafKey[commonPrefixLength] parentLeaf.Key = parentLeaf.Key[commonPrefixLength+1:] newBranchParent.Children[childIndex] = parentLeaf + newBranchParent.AddDescendants(1) + nodesCreated++ } - return newBranchParent + return newBranchParent, nodesCreated } if len(parentLeaf.Key) == commonPrefixLength { @@ -399,6 +403,8 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, childIndex := parentLeafKey[commonPrefixLength] parentLeaf.Key = parentLeaf.Key[commonPrefixLength+1:] newBranchParent.Children[childIndex] = parentLeaf + newBranchParent.AddDescendants(1) + nodesCreated++ } childIndex := key[commonPrefixLength] newBranchParent.Children[childIndex] = &node.Leaf{ @@ -407,16 +413,19 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, Generation: t.generation, Dirty: true, } + newBranchParent.AddDescendants(1) + nodesCreated++ - return newBranchParent + return newBranchParent, nodesCreated } -func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (newParent Node) { +func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) ( + newParent Node, nodesCreated uint32) { parentBranch = t.prepBranchForMutation(parentBranch) if bytes.Equal(key, parentBranch.Key) { parentBranch.Value = value - return parentBranch + return parentBranch, 0 } if bytes.HasPrefix(key, parentBranch.Key) { @@ -433,16 +442,19 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (new Generation: t.generation, Dirty: true, } + nodesCreated = 1 } else { - child = t.insert(child, remainingKey, value) + child, nodesCreated = t.insert(child, remainingKey, value) } parentBranch.Children[childIndex] = child - return parentBranch + parentBranch.AddDescendants(nodesCreated) + return parentBranch, nodesCreated } // we need to branch out at the point where the keys diverge // update partial keys, new branch has key up to matching length + nodesCreated = 1 commonPrefixLength := lenCommonPrefix(key, parentBranch.Key) newParentBranch := &node.Branch{ Key: key[:commonPrefixLength], @@ -455,16 +467,20 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (new parentBranch.Key = remainingOldParentKey newParentBranch.Children[oldParentIndex] = parentBranch + newParentBranch.AddDescendants(1 + parentBranch.GetDescendants()) if len(key) <= commonPrefixLength { newParentBranch.Value = value } else { childIndex := key[commonPrefixLength] remainingKey := key[commonPrefixLength+1:] - newParentBranch.Children[childIndex] = t.insert(nil, remainingKey, value) + var additionalNodesCreated uint32 + newParentBranch.Children[childIndex], additionalNodesCreated = t.insert(nil, remainingKey, value) + nodesCreated += additionalNodesCreated + newParentBranch.AddDescendants(additionalNodesCreated) } - return newParentBranch + return newParentBranch, nodesCreated } // LoadFromMap loads the given data mapping of key to value into the trie. @@ -652,7 +668,7 @@ func (t *Trie) ClearPrefixLimit(prefixLE []byte, limit uint32) (deleted uint32, prefix := codec.KeyLEToNibbles(prefixLE) prefix = bytes.TrimSuffix(prefix, []byte{0}) - t.root, deleted, allDeleted = t.clearPrefixLimit(t.root, prefix, limit) + t.root, deleted, _, allDeleted = t.clearPrefixLimit(t.root, prefix, limit) return deleted, allDeleted } @@ -660,9 +676,9 @@ func (t *Trie) ClearPrefixLimit(prefixLE []byte, limit uint32) (deleted uint32, // It returns the updated node newParent, the number of deleted values valuesDeleted and the // allDeleted boolean indicating if there is no key left with the prefix. func (t *Trie) clearPrefixLimit(parent Node, prefix []byte, limit uint32) ( - newParent Node, valuesDeleted uint32, allDeleted bool) { + newParent Node, valuesDeleted, nodesRemoved uint32, allDeleted bool) { if parent == nil { - return nil, 0, true + return nil, 0, 0, true } if parent.Type() == node.LeafType { @@ -671,10 +687,10 @@ func (t *Trie) clearPrefixLimit(parent Node, prefix []byte, limit uint32) ( // TODO check this is the same behaviour as in substrate const allDeleted = true if bytes.HasPrefix(leaf.Key, prefix) { - valuesDeleted = 1 - return nil, valuesDeleted, allDeleted + valuesDeleted, nodesRemoved = 1, 1 + return nil, valuesDeleted, nodesRemoved, allDeleted } - return parent, 0, allDeleted + return parent, 0, 0, allDeleted } branch := parent.(*node.Branch) @@ -682,14 +698,14 @@ func (t *Trie) clearPrefixLimit(parent Node, prefix []byte, limit uint32) ( } func (t *Trie) clearPrefixLimitBranch(branch *node.Branch, prefix []byte, limit uint32) ( - newParent Node, valuesDeleted uint32, allDeleted bool) { + newParent Node, valuesDeleted uint32, nodesRemoved uint32, allDeleted bool) { newParent = branch if bytes.HasPrefix(branch.Key, prefix) { nilPrefix := ([]byte)(nil) - newParent, valuesDeleted = t.deleteNodesLimit(branch, nilPrefix, limit) + newParent, valuesDeleted, nodesRemoved = t.deleteNodesLimit(branch, nilPrefix, limit) allDeleted = newParent == nil - return newParent, valuesDeleted, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted } if len(prefix) == len(branch.Key)+1 && @@ -701,67 +717,80 @@ func (t *Trie) clearPrefixLimitBranch(branch *node.Branch, prefix []byte, limit noPrefixForNode := len(prefix) <= len(branch.Key) || lenCommonPrefix(branch.Key, prefix) < len(branch.Key) if noPrefixForNode { - valuesDeleted = 0 + valuesDeleted, nodesRemoved = 0, 0 allDeleted = true - return newParent, valuesDeleted, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted } childIndex := prefix[len(branch.Key)] childPrefix := prefix[len(branch.Key)+1:] child := branch.Children[childIndex] - child, valuesDeleted, allDeleted = t.clearPrefixLimit(child, childPrefix, limit) + child, valuesDeleted, nodesRemoved, allDeleted = t.clearPrefixLimit(child, childPrefix, limit) if valuesDeleted == 0 { - return branch, valuesDeleted, allDeleted + return branch, valuesDeleted, nodesRemoved, allDeleted } branch = t.prepBranchForMutation(branch) branch.Children[childIndex] = child - newParent = handleDeletion(branch, prefix) - return newParent, valuesDeleted, allDeleted + branch.SubDescendants(nodesRemoved) + newParent, branchChildMerged := handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } + + return newParent, valuesDeleted, nodesRemoved, allDeleted } func (t *Trie) clearPrefixLimitChild(branch *node.Branch, prefix []byte, limit uint32) ( - newParent Node, valuesDeleted uint32, allDeleted bool) { + newParent Node, valuesDeleted, nodesRemoved uint32, allDeleted bool) { newParent = branch childIndex := prefix[len(branch.Key)] child := branch.Children[childIndex] if child == nil { + const valuesDeleted, nodesRemoved = 0, 0 // TODO ensure this is the same behaviour as in substrate allDeleted = true - return newParent, 0, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted } nilPrefix := ([]byte)(nil) - child, valuesDeleted = t.deleteNodesLimit(child, nilPrefix, limit) + child, valuesDeleted, nodesRemoved = t.deleteNodesLimit(child, nilPrefix, limit) if valuesDeleted == 0 { allDeleted = branch.Children[childIndex] == nil - return branch, valuesDeleted, allDeleted + return branch, valuesDeleted, nodesRemoved, allDeleted } branch = t.prepBranchForMutation(branch) branch.Children[childIndex] = child + branch.SubDescendants(nodesRemoved) + + newParent, branchChildMerged := handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } - newParent = handleDeletion(branch, prefix) allDeleted = branch.Children[childIndex] == nil - return newParent, valuesDeleted, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted } func (t *Trie) deleteNodesLimit(parent Node, prefix []byte, limit uint32) ( - newParent Node, valuesDeleted uint32) { + newParent Node, valuesDeleted, nodesRemoved uint32) { if limit == 0 { - return parent, 0 + valuesDeleted, nodesRemoved = 0, 0 + return parent, valuesDeleted, nodesRemoved } if parent == nil { - return nil, 0 + valuesDeleted, nodesRemoved = 0, 0 + return nil, valuesDeleted, nodesRemoved } if parent.Type() == node.LeafType { - valuesDeleted = 1 - return nil, valuesDeleted + valuesDeleted, nodesRemoved = 1, 1 + return nil, valuesDeleted, nodesRemoved } branch := parent.(*node.Branch) @@ -770,36 +799,46 @@ func (t *Trie) deleteNodesLimit(parent Node, prefix []byte, limit uint32) ( nilChildren := node.ChildrenCapacity - branch.NumChildren() - var newDeleted uint32 + var newDeleted, newNodesRemoved uint32 + var branchChildMerged bool for i, child := range branch.Children { if child == nil { continue } branch = t.prepBranchForMutation(branch) - branch.Children[i], newDeleted = t.deleteNodesLimit(child, fullKey, limit) + branch.Children[i], newDeleted, newNodesRemoved = t.deleteNodesLimit(child, fullKey, limit) if branch.Children[i] == nil { nilChildren++ } limit -= newDeleted valuesDeleted += newDeleted + nodesRemoved += newNodesRemoved + branch.SubDescendants(newNodesRemoved) + + branch.SetDirty(true) + + newParent, branchChildMerged = handleDeletion(branch, fullKey) + if branchChildMerged { + nodesRemoved++ + } - newParent = handleDeletion(branch, fullKey) if nilChildren == node.ChildrenCapacity && branch.Value == nil { - return nil, valuesDeleted + return nil, valuesDeleted, nodesRemoved } if limit == 0 { - return newParent, valuesDeleted + return newParent, valuesDeleted, nodesRemoved } } + nodesRemoved++ if branch.Value != nil { valuesDeleted++ } - return nil, valuesDeleted + return nil, valuesDeleted, nodesRemoved } // ClearPrefix deletes all nodes in the trie for which the key contains the @@ -817,17 +856,23 @@ func (t *Trie) ClearPrefix(prefixLE []byte) { } func (t *Trie) clearPrefix(parent Node, prefix []byte) ( - newParent Node, updated bool) { + newParent Node, nodesRemoved uint32) { if parent == nil { - return nil, false + const nodesRemoved = 0 + return nil, nodesRemoved } if bytes.HasPrefix(parent.GetKey(), prefix) { - return nil, true + nodesRemoved = 1 + if parent.Type() != node.LeafType { // branch + nodesRemoved += parent.(*node.Branch).GetDescendants() + } + return nil, nodesRemoved } if parent.Type() == node.LeafType { - return parent, false + const nodesRemoved = 0 + return parent, nodesRemoved } branch := parent.(*node.Branch) @@ -839,34 +884,46 @@ func (t *Trie) clearPrefix(parent Node, prefix []byte) ( child := branch.Children[childIndex] if child == nil { - return parent, false + const nodesRemoved = 0 + return parent, nodesRemoved } + nodesRemoved = 1 branch = t.prepBranchForMutation(branch) branch.Children[childIndex] = nil - newParent = handleDeletion(branch, prefix) - return newParent, true + var branchChildMerged bool + newParent, branchChildMerged = handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } + return newParent, nodesRemoved } noPrefixForNode := len(prefix) <= len(branch.Key) || lenCommonPrefix(branch.Key, prefix) < len(branch.Key) if noPrefixForNode { - return parent, false + const nodesRemoved = 0 + return parent, nodesRemoved } childIndex := prefix[len(branch.Key)] childPrefix := prefix[len(branch.Key)+1:] child := branch.Children[childIndex] - child, updated = t.clearPrefix(child, childPrefix) - if !updated { - return parent, false + child, nodesRemoved = t.clearPrefix(child, childPrefix) + if nodesRemoved == 0 { + return parent, nodesRemoved } branch = t.prepBranchForMutation(branch) + branch.SubDescendants(nodesRemoved) branch.Children[childIndex] = child - newParent = handleDeletion(branch, prefix) - return newParent, true + newParent, branchChildMerged := handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } + + return newParent, nodesRemoved } // Delete removes the node of the trie with the key @@ -874,19 +931,23 @@ func (t *Trie) clearPrefix(parent Node, prefix []byte) ( // If no node is found at this key, nothing is deleted. func (t *Trie) Delete(keyLE []byte) { key := codec.KeyLEToNibbles(keyLE) - t.root, _ = t.delete(t.root, key) + t.root, _, _ = t.delete(t.root, key) } -func (t *Trie) delete(parent Node, key []byte) (newParent Node, deleted bool) { +func (t *Trie) delete(parent Node, key []byte) ( + newParent Node, deleted bool, nodesRemoved uint32) { if parent == nil { - return nil, false + const nodesRemoved = 0 + return nil, false, nodesRemoved } if parent.Type() == node.LeafType { if deleteLeaf(parent, key) == nil { - return nil, true + const nodesRemoved = 1 + return nil, true, nodesRemoved } - return parent, false + const nodesRemoved = 0 + return parent, false, nodesRemoved } branch := parent.(*node.Branch) @@ -900,11 +961,18 @@ func deleteLeaf(parent Node, key []byte) (newParent Node) { return parent } -func (t *Trie) deleteBranch(branch *node.Branch, key []byte) (newParent Node, deleted bool) { +func (t *Trie) deleteBranch(branch *node.Branch, key []byte) ( + newParent Node, deleted bool, nodesRemoved uint32) { if len(key) == 0 || bytes.Equal(branch.Key, key) { branch = t.prepBranchForMutation(branch) branch.Value = nil - return handleDeletion(branch, key), true + deleted = true + var branchChildMerged bool + newParent, branchChildMerged = handleDeletion(branch, key) + if branchChildMerged { + nodesRemoved = 1 + } + return newParent, deleted, nodesRemoved } commonPrefixLength := lenCommonPrefix(branch.Key, key) @@ -912,23 +980,31 @@ func (t *Trie) deleteBranch(branch *node.Branch, key []byte) (newParent Node, de childKey := key[commonPrefixLength+1:] child := branch.Children[childIndex] - newChild, deleted := t.delete(child, childKey) + newChild, deleted, nodesRemoved := t.delete(child, childKey) if !deleted { - return branch, false + const nodesRemoved = 0 + return branch, false, nodesRemoved } branch = t.prepBranchForMutation(branch) + branch.SubDescendants(nodesRemoved) branch.Children[childIndex] = newChild - newParent = handleDeletion(branch, key) - return newParent, true + + newParent, branchChildMerged := handleDeletion(branch, key) + if branchChildMerged { + nodesRemoved++ + } + + return newParent, true, nodesRemoved } // handleDeletion is called when a value is deleted from a branch to handle // the eventual mutation of the branch depending on its children. // If the branch has no value and a single child, it will be combined with this child. +// 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.Branch, key []byte) (newNode Node) { - // TODO try to remove key argument just use branch.Key instead? +func handleDeletion(branch *node.Branch, key []byte) (newNode Node, branchChildMerged bool) { childrenCount := 0 firstChildIndex := -1 for i, child := range branch.Children { @@ -943,16 +1019,19 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { switch { default: - return branch + const branchChildMerged = false + return branch, branchChildMerged case childrenCount == 0 && branch.Value != nil: + const branchChildMerged = false commonPrefixLength := lenCommonPrefix(branch.Key, key) return &node.Leaf{ Key: key[:commonPrefixLength], Value: branch.Value, Dirty: true, Generation: branch.Generation, - } + }, branchChildMerged case childrenCount == 1 && branch.Value == nil: + const branchChildMerged = true childIndex := firstChildIndex child := branch.Children[firstChildIndex] @@ -964,7 +1043,7 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { Value: child.GetValue(), Dirty: true, Generation: branch.Generation, - } + }, branchChildMerged } childBranch := child.(*node.Branch) @@ -974,6 +1053,8 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { Value: childBranch.Value, Generation: branch.Generation, Dirty: true, + // this is the descendants of the original branch minus one + Descendants: childBranch.GetDescendants(), } // Adopt the grand-children @@ -985,7 +1066,7 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { } } - return newBranch + return newBranch, branchChildMerged } } diff --git a/lib/trie/trie_endtoend_test.go b/lib/trie/trie_endtoend_test.go index 7baf467c7a2..bf982b04467 100644 --- a/lib/trie/trie_endtoend_test.go +++ b/lib/trie/trie_endtoend_test.go @@ -1111,3 +1111,70 @@ func Test_encodeRoot_fuzz(t *testing.T) { require.NotEmpty(t, buffer.Bytes()) } } + +func countNodesRecursively(root Node) (nodesCount uint32) { + if root == nil { + return 0 + } else if root.Type() == node.LeafType { + return 1 + } + branch := root.(*node.Branch) + for _, child := range branch.Children { + nodesCount += countNodesRecursively(child) + } + + return 1 + nodesCount +} + +func countNodesFromStats(root Node) (nodesCount uint32) { + if root == nil { + return 0 + } else if root.Type() == node.LeafType { + return 1 + } + return 1 + root.(*node.Branch).GetDescendants() +} + +func testDescendants(t *testing.T, root Node) { + t.Helper() + expectedCount := countNodesRecursively(root) + statsCount := countNodesFromStats(root) + require.Equal(t, int(expectedCount), int(statsCount)) +} + +func Test_Trie_Descendants_Fuzz(t *testing.T) { + generator := newGenerator() + const kvSize = 5000 + kv := generateKeyValues(t, generator, kvSize) + + trie := NewEmptyTrie() + + keys := make([][]byte, 0, len(kv)) + for key := range kv { + keys = append(keys, []byte(key)) + } + sort.Slice(keys, func(i, j int) bool { + return bytes.Compare(keys[i], keys[j]) < 0 + }) + + for _, key := range keys { + trie.Put(key, kv[string(key)]) + } + + testDescendants(t, trie.root) + + require.Greater(t, kvSize, 3) + + trie.ClearPrefix(keys[0]) + + testDescendants(t, trie.root) + + trie.ClearPrefixLimit(keys[1], 100) + + testDescendants(t, trie.root) + + trie.Delete(keys[2]) + trie.Delete(keys[3]) + + testDescendants(t, trie.root) +} diff --git a/lib/trie/trie_test.go b/lib/trie/trie_test.go index 7ee8858c512..1ec6d13e116 100644 --- a/lib/trie/trie_test.go +++ b/lib/trie/trie_test.go @@ -461,8 +461,9 @@ func Test_Trie_Hash(t *testing.T) { "branch root": { trie: Trie{ root: &node.Branch{ - Key: []byte{1, 2, 3}, - Value: []byte("branch"), + Key: []byte{1, 2, 3}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{9}}, }, @@ -475,8 +476,9 @@ func Test_Trie_Hash(t *testing.T) { 0xbe, 0x27, 0xab, 0x13, 0xcb, 0xf0, 0xfd, 0xd7}, expectedTrie: Trie{ root: &node.Branch{ - Key: []byte{1, 2, 3}, - Value: []byte("branch"), + Key: []byte{1, 2, 3}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ Key: []byte{9}, @@ -539,8 +541,9 @@ func Test_Trie_Entries(t *testing.T) { t.Parallel() root := &node.Branch{ - Key: []byte{0xa}, - Value: []byte("root"), + Key: []byte{0xa}, + Value: []byte("root"), + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ // index 0 Key: []byte{2, 0xb}, @@ -571,13 +574,15 @@ func Test_Trie_Entries(t *testing.T) { t.Parallel() root := &node.Branch{ - Key: []byte{0xa, 0xb}, - Value: []byte("root"), + Key: []byte{0xa, 0xb}, + Value: []byte("root"), + Descendants: 5, Children: [16]node.Node{ nil, nil, nil, &node.Branch{ // branch with value at child index 3 - Key: []byte{0xb}, - Value: []byte("branch 1"), + Key: []byte{0xb}, + Value: []byte("branch 1"), + Descendants: 1, Children: [16]node.Node{ nil, nil, nil, &node.Leaf{ // leaf at child index 3 @@ -593,8 +598,9 @@ func Test_Trie_Entries(t *testing.T) { }, nil, &node.Branch{ // branch without value at child index 9 - Key: []byte{0xe}, - Value: []byte("branch 2"), + Key: []byte{0xe}, + Value: []byte("branch 2"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // leaf at child index 0 Key: []byte{0xf}, @@ -741,8 +747,9 @@ func Test_nextKey(t *testing.T) { "key smaller than root branch full key": { trie: Trie{ root: &node.Branch{ - Key: []byte{2}, - Value: []byte("branch"), + Key: []byte{2}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ Key: []byte{1}, @@ -756,8 +763,9 @@ func Test_nextKey(t *testing.T) { "key equal to root branch full key": { trie: Trie{ root: &node.Branch{ - Key: []byte{2}, - Value: []byte("branch"), + Key: []byte{2}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ Key: []byte{1}, @@ -770,8 +778,9 @@ func Test_nextKey(t *testing.T) { "key smaller than leaf full key": { trie: Trie{ root: &node.Branch{ - Key: []byte{1}, - Value: []byte("branch"), + Key: []byte{1}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ nil, nil, &node.Leaf{ @@ -787,8 +796,9 @@ func Test_nextKey(t *testing.T) { "key equal to leaf full key": { trie: Trie{ root: &node.Branch{ - Key: []byte{1}, - Value: []byte("branch"), + Key: []byte{1}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ nil, nil, &node.Leaf{ @@ -803,8 +813,9 @@ func Test_nextKey(t *testing.T) { "key greater than leaf full key": { trie: Trie{ root: &node.Branch{ - Key: []byte{1}, - Value: []byte("branch"), + Key: []byte{1}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ nil, nil, &node.Leaf{ @@ -819,14 +830,16 @@ func Test_nextKey(t *testing.T) { "next key branch with value": { trie: Trie{ root: &node.Branch{ - Key: []byte{1}, - Value: []byte("top branch"), + Key: []byte{1}, + Value: []byte("top branch"), + Descendants: 2, Children: [16]node.Node{ nil, nil, &node.Branch{ // full key [1, 2, 3] - Key: []byte{3}, - Value: []byte("branch 1"), + Key: []byte{3}, + Value: []byte("branch 1"), + Descendants: 1, Children: [16]node.Node{ nil, nil, nil, nil, &node.Leaf{ @@ -845,12 +858,14 @@ func Test_nextKey(t *testing.T) { "next key go through branch without value": { trie: Trie{ root: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ nil, nil, &node.Branch{ // full key [1, 2, 3] - Key: []byte{3}, + Key: []byte{3}, + Descendants: 1, Children: [16]node.Node{ nil, nil, nil, nil, &node.Leaf{ @@ -869,13 +884,15 @@ func Test_nextKey(t *testing.T) { "next key leaf from bottom branch": { trie: Trie{ root: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ nil, nil, &node.Branch{ // full key [1, 2, 3] - Key: []byte{3}, - Value: []byte("bottom branch"), + Key: []byte{3}, + Value: []byte("bottom branch"), + Descendants: 1, Children: [16]node.Node{ nil, nil, nil, nil, &node.Leaf{ @@ -894,13 +911,15 @@ func Test_nextKey(t *testing.T) { "next key greater than branch": { trie: Trie{ root: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ nil, nil, &node.Branch{ // full key [1, 2, 3] - Key: []byte{3}, - Value: []byte("bottom branch"), + Key: []byte{3}, + Value: []byte("bottom branch"), + Descendants: 1, Children: [16]node.Node{ nil, nil, nil, nil, &node.Leaf{ @@ -919,8 +938,9 @@ func Test_nextKey(t *testing.T) { "key smaller length and greater than root branch full key": { trie: Trie{ root: &node.Branch{ - Key: []byte{2, 0}, - Value: []byte("branch"), + Key: []byte{2, 0}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -976,9 +996,10 @@ func Test_Trie_Put(t *testing.T) { expectedTrie: Trie{ generation: 1, root: &node.Branch{ - Key: []byte{1, 2}, - Generation: 1, - Dirty: true, + Key: []byte{1, 2}, + Generation: 1, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ Key: []byte{5}, @@ -1075,9 +1096,10 @@ func Test_Trie_put(t *testing.T) { expectedTrie: Trie{ generation: 1, root: &node.Branch{ - Key: []byte{1}, - Generation: 1, - Dirty: true, + Key: []byte{1}, + Generation: 1, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ Key: []byte{5}, @@ -1114,11 +1136,12 @@ func Test_Trie_insert(t *testing.T) { t.Parallel() testCases := map[string]struct { - trie Trie - parent Node - key []byte - value []byte - newNode Node + trie Trie + parent Node + key []byte + value []byte + newNode Node + nodesCreated uint32 }{ "nil parent": { trie: Trie{ @@ -1132,14 +1155,16 @@ func Test_Trie_insert(t *testing.T) { Generation: 1, Dirty: true, }, + nodesCreated: 1, }, "branch parent": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte("branch"), + Key: []byte{1}, + Value: []byte("branch"), + Descendants: 1, Children: [16]node.Node{ nil, &node.Leaf{Key: []byte{2}}, @@ -1148,10 +1173,11 @@ func Test_Trie_insert(t *testing.T) { key: []byte{1, 0}, value: []byte("leaf"), newNode: &node.Branch{ - Key: []byte{1}, - Value: []byte("branch"), - Generation: 1, - Dirty: true, + Key: []byte{1}, + Value: []byte("branch"), + Generation: 1, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ Key: []byte{}, @@ -1162,6 +1188,7 @@ func Test_Trie_insert(t *testing.T) { &node.Leaf{Key: []byte{2}}, }, }, + nodesCreated: 1, }, "override leaf parent": { trie: Trie{ @@ -1180,7 +1207,7 @@ func Test_Trie_insert(t *testing.T) { Dirty: true, }, }, - "write same leaf value as child to parent leaf": { + "write same leaf value to leaf parent": { trie: Trie{ generation: 1, }, @@ -1206,10 +1233,11 @@ func Test_Trie_insert(t *testing.T) { key: []byte{1, 0}, value: []byte("leaf"), newNode: &node.Branch{ - Key: []byte{1}, - Value: []byte("original leaf"), - Dirty: true, - Generation: 1, + Key: []byte{1}, + Value: []byte("original leaf"), + Dirty: true, + Generation: 1, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ Key: []byte{}, @@ -1219,6 +1247,7 @@ func Test_Trie_insert(t *testing.T) { }, }, }, + nodesCreated: 1, }, "write leaf as divergent child next to parent leaf": { trie: Trie{ @@ -1231,9 +1260,10 @@ func Test_Trie_insert(t *testing.T) { key: []byte{2, 3}, value: []byte("leaf"), newNode: &node.Branch{ - Key: []byte{}, - Dirty: true, - Generation: 1, + Key: []byte{}, + Dirty: true, + Generation: 1, + Descendants: 2, Children: [16]node.Node{ nil, &node.Leaf{ @@ -1250,8 +1280,9 @@ func Test_Trie_insert(t *testing.T) { }, }, }, + nodesCreated: 2, }, - "write leaf into nil leaf": { + "write leaf into nil value leaf": { trie: Trie{ generation: 1, }, @@ -1277,10 +1308,11 @@ func Test_Trie_insert(t *testing.T) { key: []byte{1}, value: []byte("leaf"), newNode: &node.Branch{ - Key: []byte{1}, - Value: []byte("leaf"), - Dirty: true, - Generation: 1, + Key: []byte{1}, + Value: []byte("leaf"), + Dirty: true, + Generation: 1, + Descendants: 1, Children: [16]node.Node{ nil, nil, &node.Leaf{ @@ -1290,6 +1322,7 @@ func Test_Trie_insert(t *testing.T) { }, }, }, + nodesCreated: 1, }, } @@ -1301,9 +1334,10 @@ func Test_Trie_insert(t *testing.T) { trie := testCase.trie expectedTrie := *trie.DeepCopy() - newNode := trie.insert(testCase.parent, testCase.key, testCase.value) + newNode, nodesCreated := trie.insert(testCase.parent, testCase.key, testCase.value) assert.Equal(t, testCase.newNode, newNode) + assert.Equal(t, testCase.nodesCreated, nodesCreated) assert.Equal(t, expectedTrie, trie) }) } @@ -1313,15 +1347,17 @@ func Test_Trie_insertInBranch(t *testing.T) { t.Parallel() testCases := map[string]struct { - parent *node.Branch - key []byte - value []byte - newNode Node + parent *node.Branch + key []byte + value []byte + newNode Node + nodesCreated uint32 }{ "update with branch": { parent: &node.Branch{ - Key: []byte{2}, - Value: []byte("old"), + Key: []byte{2}, + Value: []byte("old"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1329,9 +1365,10 @@ func Test_Trie_insertInBranch(t *testing.T) { key: []byte{2}, value: []byte("new"), newNode: &node.Branch{ - Key: []byte{2}, - Value: []byte("new"), - Dirty: true, + Key: []byte{2}, + Value: []byte("new"), + Dirty: true, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1339,8 +1376,9 @@ func Test_Trie_insertInBranch(t *testing.T) { }, "update with leaf": { parent: &node.Branch{ - Key: []byte{2}, - Value: []byte("old"), + Key: []byte{2}, + Value: []byte("old"), + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1348,9 +1386,10 @@ func Test_Trie_insertInBranch(t *testing.T) { key: []byte{2}, value: []byte("new"), newNode: &node.Branch{ - Key: []byte{2}, - Value: []byte("new"), - Dirty: true, + Key: []byte{2}, + Value: []byte("new"), + Dirty: true, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1358,8 +1397,9 @@ func Test_Trie_insertInBranch(t *testing.T) { }, "add leaf as direct child": { parent: &node.Branch{ - Key: []byte{2}, - Value: []byte{5}, + Key: []byte{2}, + Value: []byte{5}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1367,9 +1407,10 @@ func Test_Trie_insertInBranch(t *testing.T) { key: []byte{2, 3, 4, 5}, value: []byte{6}, newNode: &node.Branch{ - Key: []byte{2}, - Value: []byte{5}, - Dirty: true, + Key: []byte{2}, + Value: []byte{5}, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, nil, nil, @@ -1380,15 +1421,18 @@ func Test_Trie_insertInBranch(t *testing.T) { }, }, }, + nodesCreated: 1, }, "add leaf as nested child": { parent: &node.Branch{ - Key: []byte{2}, - Value: []byte{5}, + Key: []byte{2}, + Value: []byte{5}, + Descendants: 2, Children: [16]node.Node{ nil, nil, nil, &node.Branch{ - Key: []byte{4}, + Key: []byte{4}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1398,14 +1442,16 @@ func Test_Trie_insertInBranch(t *testing.T) { key: []byte{2, 3, 4, 5, 6}, value: []byte{6}, newNode: &node.Branch{ - Key: []byte{2}, - Value: []byte{5}, - Dirty: true, + Key: []byte{2}, + Value: []byte{5}, + Dirty: true, + Descendants: 3, Children: [16]node.Node{ nil, nil, nil, &node.Branch{ - Key: []byte{4}, - Dirty: true, + Key: []byte{4}, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, nil, nil, nil, nil, @@ -1418,11 +1464,13 @@ func Test_Trie_insertInBranch(t *testing.T) { }, }, }, + nodesCreated: 1, }, "split branch for longer key": { parent: &node.Branch{ - Key: []byte{2, 3}, - Value: []byte{5}, + Key: []byte{2, 3}, + Value: []byte{5}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1430,14 +1478,16 @@ func Test_Trie_insertInBranch(t *testing.T) { key: []byte{2, 4, 5, 6}, value: []byte{6}, newNode: &node.Branch{ - Key: []byte{2}, - Dirty: true, + Key: []byte{2}, + Dirty: true, + Descendants: 3, Children: [16]node.Node{ nil, nil, nil, &node.Branch{ - Key: []byte{}, - Value: []byte{5}, - Dirty: true, + Key: []byte{}, + Value: []byte{5}, + Dirty: true, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1449,11 +1499,13 @@ func Test_Trie_insertInBranch(t *testing.T) { }, }, }, + nodesCreated: 2, }, "split root branch": { parent: &node.Branch{ - Key: []byte{2, 3}, - Value: []byte{5}, + Key: []byte{2, 3}, + Value: []byte{5}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1461,14 +1513,16 @@ func Test_Trie_insertInBranch(t *testing.T) { key: []byte{3}, value: []byte{6}, newNode: &node.Branch{ - Key: []byte{}, - Dirty: true, + Key: []byte{}, + Dirty: true, + Descendants: 3, Children: [16]node.Node{ nil, nil, &node.Branch{ - Key: []byte{3}, - Value: []byte{5}, - Dirty: true, + Key: []byte{3}, + Value: []byte{5}, + Dirty: true, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1480,11 +1534,13 @@ func Test_Trie_insertInBranch(t *testing.T) { }, }, }, + nodesCreated: 2, }, "update with leaf at empty key": { parent: &node.Branch{ - Key: []byte{2}, - Value: []byte{5}, + Key: []byte{2}, + Value: []byte{5}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1492,21 +1548,24 @@ func Test_Trie_insertInBranch(t *testing.T) { key: []byte{}, value: []byte{6}, newNode: &node.Branch{ - Key: []byte{}, - Value: []byte{6}, - Dirty: true, + Key: []byte{}, + Value: []byte{6}, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ nil, nil, &node.Branch{ - Key: []byte{}, - Value: []byte{5}, - Dirty: true, + Key: []byte{}, + Value: []byte{5}, + Dirty: true, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, }, }, }, + nodesCreated: 1, }, } @@ -1517,9 +1576,10 @@ func Test_Trie_insertInBranch(t *testing.T) { trie := new(Trie) - newNode := trie.insertInBranch(testCase.parent, testCase.key, testCase.value) + newNode, nodesCreated := trie.insertInBranch(testCase.parent, testCase.key, testCase.value) assert.Equal(t, testCase.newNode, newNode) + assert.Equal(t, testCase.nodesCreated, nodesCreated) assert.Equal(t, new(Trie), trie) // check no mutation }) } @@ -1561,9 +1621,10 @@ func Test_Trie_LoadFromMap(t *testing.T) { }, expectedTrie: Trie{ root: &node.Branch{ - Key: []byte{00, 01}, - Value: []byte{6}, - Dirty: true, + Key: []byte{00, 01}, + Value: []byte{6}, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ nil, nil, &node.Leaf{ @@ -1583,9 +1644,10 @@ func Test_Trie_LoadFromMap(t *testing.T) { "override trie": { trie: Trie{ root: &node.Branch{ - Key: []byte{00, 01}, - Value: []byte{106}, - Dirty: true, + Key: []byte{00, 01}, + Value: []byte{106}, + Dirty: true, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ Value: []byte{9}, @@ -1606,9 +1668,10 @@ func Test_Trie_LoadFromMap(t *testing.T) { }, expectedTrie: Trie{ root: &node.Branch{ - Key: []byte{00, 01}, - Value: []byte{6}, - Dirty: true, + Key: []byte{00, 01}, + Value: []byte{6}, + Dirty: true, + Descendants: 3, Children: [16]node.Node{ &node.Leaf{ Value: []byte{9}, @@ -1658,10 +1721,12 @@ func Test_Trie_GetKeysWithPrefix(t *testing.T) { "some trie": { trie: Trie{ root: &node.Branch{ - Key: []byte{0, 1}, + Key: []byte{0, 1}, + Descendants: 4, Children: [16]node.Node{ &node.Branch{ // full key 0, 1, 0, 3 - Key: []byte{3}, + Key: []byte{3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ // full key 0, 1, 0, 0, 4 Key: []byte{4}, @@ -1714,7 +1779,8 @@ func Test_getKeysWithPrefix(t *testing.T) { }, "common prefix for parent branch and search key": { parent: &node.Branch{ - Key: []byte{1, 2, 3}, + Key: []byte{1, 2, 3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{4}}, &node.Leaf{Key: []byte{5}}, @@ -1729,7 +1795,8 @@ func Test_getKeysWithPrefix(t *testing.T) { }, "parent branch and empty key": { parent: &node.Branch{ - Key: []byte{1, 2, 3}, + Key: []byte{1, 2, 3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{4}}, &node.Leaf{Key: []byte{5}}, @@ -1744,7 +1811,8 @@ func Test_getKeysWithPrefix(t *testing.T) { }, "search key smaller than branch key with no full common prefix": { parent: &node.Branch{ - Key: []byte{1, 2, 3}, + Key: []byte{1, 2, 3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{4}}, &node.Leaf{Key: []byte{5}}, @@ -1756,7 +1824,8 @@ func Test_getKeysWithPrefix(t *testing.T) { }, "common prefix smaller tan search key": { parent: &node.Branch{ - Key: []byte{1, 2}, + Key: []byte{1, 2}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{4}}, &node.Leaf{Key: []byte{5}}, @@ -1768,7 +1837,8 @@ func Test_getKeysWithPrefix(t *testing.T) { }, "recursive call": { parent: &node.Branch{ - Key: []byte{1, 2, 3}, + Key: []byte{1, 2, 3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{4}}, &node.Leaf{Key: []byte{5}}, @@ -1867,7 +1937,8 @@ func Test_addAllKeys(t *testing.T) { }, "parent branch without value": { parent: &node.Branch{ - Key: []byte{1, 2, 3}, + Key: []byte{1, 2, 3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{4}}, &node.Leaf{Key: []byte{5}}, @@ -1881,8 +1952,9 @@ func Test_addAllKeys(t *testing.T) { }, "parent branch with empty value": { parent: &node.Branch{ - Key: []byte{1, 2, 3}, - Value: []byte{}, + Key: []byte{1, 2, 3}, + Value: []byte{}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{4}}, &node.Leaf{Key: []byte{5}}, @@ -1921,12 +1993,14 @@ func Test_Trie_Get(t *testing.T) { "some trie": { trie: Trie{ root: &node.Branch{ - Key: []byte{0, 1}, - Value: []byte{1, 3}, + Key: []byte{0, 1}, + Value: []byte{1, 3}, + Descendants: 3, Children: [16]node.Node{ &node.Branch{ // full key 0, 1, 0, 3 - Key: []byte{3}, - Value: []byte{1, 2}, + Key: []byte{3}, + Value: []byte{1, 2}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1983,8 +2057,9 @@ func Test_retrieve(t *testing.T) { }, "branch key match": { parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{2}, + Key: []byte{1}, + Value: []byte{2}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -1994,8 +2069,9 @@ func Test_retrieve(t *testing.T) { }, "branch key with empty search key": { parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{2}, + Key: []byte{1}, + Value: []byte{2}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2004,8 +2080,9 @@ func Test_retrieve(t *testing.T) { }, "branch key mismatch with shorter search key": { parent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{2}, + Key: []byte{1, 2}, + Value: []byte{2}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2014,13 +2091,15 @@ func Test_retrieve(t *testing.T) { }, "bottom leaf in branch": { parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 2, Children: [16]node.Node{ nil, nil, &node.Branch{ // full key 1, 2, 3 - Key: []byte{3}, - Value: []byte{2}, + Key: []byte{3}, + Value: []byte{2}, + Descendants: 1, Children: [16]node.Node{ nil, nil, nil, nil, &node.Leaf{ // full key 1, 2, 3, 4, 5 @@ -2071,8 +2150,9 @@ func Test_Trie_ClearPrefixLimit(t *testing.T) { "clear prefix limit": { trie: Trie{ root: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ nil, nil, nil, &node.Leaf{ @@ -2114,6 +2194,7 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { limit uint32 newParent Node valuesDeleted uint32 + nodesRemoved uint32 allDeleted bool }{ "limit is zero": { @@ -2130,6 +2211,7 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1}, limit: 1, valuesDeleted: 1, + nodesRemoved: 1, allDeleted: true, }, "leaf parent with key equal prefix": { @@ -2139,6 +2221,7 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1}, limit: 1, valuesDeleted: 1, + nodesRemoved: 1, allDeleted: true, }, "leaf parent with key no common prefix": { @@ -2171,7 +2254,8 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { }, "branch without value parent with common prefix": { parent: &node.Branch{ - Key: []byte{1, 2}, + Key: []byte{1, 2}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2180,11 +2264,13 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1}, limit: 3, valuesDeleted: 2, + nodesRemoved: 3, allDeleted: true, }, "branch without value with key equal prefix": { parent: &node.Branch{ - Key: []byte{1, 2}, + Key: []byte{1, 2}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2193,6 +2279,7 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 2}, limit: 3, valuesDeleted: 2, + nodesRemoved: 3, allDeleted: true, }, "branch without value with no common prefix": { @@ -2200,7 +2287,8 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1, 2}, + Key: []byte{1, 2}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2209,7 +2297,8 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 3}, limit: 1, newParent: &node.Branch{ - Key: []byte{1, 2}, + Key: []byte{1, 2}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2222,7 +2311,8 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2231,7 +2321,8 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 2, 3}, limit: 1, newParent: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2244,7 +2335,8 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2253,7 +2345,8 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 2}, limit: 1, newParent: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2263,8 +2356,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { }, "branch with value with common prefix": { parent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2272,12 +2366,14 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1}, limit: 2, valuesDeleted: 2, + nodesRemoved: 2, allDeleted: true, }, "branch with value with key equal prefix": { parent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2285,6 +2381,7 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 2}, limit: 2, valuesDeleted: 2, + nodesRemoved: 2, allDeleted: true, }, "branch with value with no common prefix": { @@ -2292,8 +2389,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2301,8 +2399,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 3}, limit: 1, newParent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2314,8 +2413,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2323,8 +2423,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 2, 3}, limit: 1, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2336,8 +2437,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2345,8 +2447,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 2}, limit: 1, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2358,8 +2461,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{3}}, &node.Leaf{Key: []byte{4}}, @@ -2368,21 +2472,24 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1}, limit: 1, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, - Dirty: true, - Generation: 1, + Key: []byte{1}, + Value: []byte{1}, + Dirty: true, + Generation: 1, + Descendants: 1, Children: [16]node.Node{ nil, &node.Leaf{Key: []byte{4}}, }, }, valuesDeleted: 1, + nodesRemoved: 1, }, "delete only child of branch": { parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{3}}, }, @@ -2395,6 +2502,7 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { Dirty: true, }, valuesDeleted: 1, + nodesRemoved: 1, allDeleted: true, }, "fully delete children of branch with value": { @@ -2402,8 +2510,9 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{3}}, &node.Leaf{Key: []byte{4}}, @@ -2418,10 +2527,12 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { Generation: 1, }, valuesDeleted: 2, + nodesRemoved: 2, }, "fully delete children of branch without value": { parent: &node.Branch{ - Key: []byte{1}, + Key: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{3}}, &node.Leaf{Key: []byte{4}}, @@ -2430,6 +2541,7 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1}, limit: 2, valuesDeleted: 2, + nodesRemoved: 3, allDeleted: true, }, "partially delete child of branch": { @@ -2437,12 +2549,14 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 3, Children: [16]node.Node{ &node.Branch{ // full key 1, 0, 3 - Key: []byte{3}, - Value: []byte{1}, + Key: []byte{3}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // full key 1, 0, 3, 0, 5 Key: []byte{5}, @@ -2457,10 +2571,11 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { prefix: []byte{1, 0}, limit: 1, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, - Dirty: true, - Generation: 1, + Key: []byte{1}, + Value: []byte{1}, + Dirty: true, + Generation: 1, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ // full key 1, 0, 3 Key: []byte{3}, @@ -2475,18 +2590,21 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { }, }, valuesDeleted: 1, + nodesRemoved: 1, }, "update child of branch": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Branch{ // full key 1, 0, 2 - Key: []byte{2}, - Value: []byte{1}, + Key: []byte{2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, }, @@ -2502,6 +2620,53 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { Generation: 1, }, valuesDeleted: 2, + nodesRemoved: 2, + allDeleted: true, + }, + "delete one of two children of branch without value": { + trie: Trie{ + generation: 1, + }, + parent: &node.Branch{ + Key: []byte{1}, + Descendants: 2, + Children: [16]node.Node{ + &node.Leaf{Key: []byte{3}}, + &node.Leaf{Key: []byte{4}}, + }, + }, + prefix: []byte{1, 0, 3}, + limit: 3, + newParent: &node.Leaf{ + Key: []byte{1, 1, 4}, + Dirty: true, + Generation: 1, + }, + valuesDeleted: 1, + nodesRemoved: 2, + allDeleted: true, + }, + "delete one of two children of branch": { + trie: Trie{ + generation: 1, + }, + parent: &node.Branch{ + Key: []byte{1}, + Descendants: 2, + Children: [16]node.Node{ + &node.Leaf{Key: []byte{3}}, + &node.Leaf{Key: []byte{4}}, + }, + }, + prefix: []byte{1, 0}, + limit: 3, + newParent: &node.Leaf{ + Key: []byte{1, 1, 4}, + Dirty: true, + Generation: 1, + }, + valuesDeleted: 1, + nodesRemoved: 2, allDeleted: true, }, "delete child of branch with limit reached": { @@ -2534,11 +2699,12 @@ func Test_Trie_clearPrefixLimit(t *testing.T) { trie := testCase.trie expectedTrie := *trie.DeepCopy() - newParent, valuesDeleted, allDeleted := trie.clearPrefixLimit(testCase.parent, - testCase.prefix, testCase.limit) + newParent, valuesDeleted, nodesRemoved, allDeleted := + trie.clearPrefixLimit(testCase.parent, testCase.prefix, testCase.limit) assert.Equal(t, testCase.newParent, newParent) assert.Equal(t, testCase.valuesDeleted, valuesDeleted) + assert.Equal(t, testCase.nodesRemoved, nodesRemoved) assert.Equal(t, testCase.allDeleted, allDeleted) assert.Equal(t, expectedTrie, trie) }) @@ -2555,6 +2721,7 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { limit uint32 newNode Node valuesDeleted uint32 + nodesRemoved uint32 }{ "zero limit": { trie: Trie{ @@ -2574,9 +2741,11 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { parent: &node.Leaf{}, limit: 2, valuesDeleted: 1, + nodesRemoved: 1, }, "delete branch without value": { parent: &node.Branch{ + Descendants: 2, Children: [16]node.Node{ &node.Leaf{}, &node.Leaf{}, @@ -2584,21 +2753,25 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { }, limit: 3, valuesDeleted: 2, + nodesRemoved: 3, }, "delete branch with value": { parent: &node.Branch{ - Key: []byte{3}, - Value: []byte{1}, + Key: []byte{3}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, }, limit: 3, valuesDeleted: 2, + nodesRemoved: 2, }, "delete branch and all children": { parent: &node.Branch{ - Key: []byte{3}, + Key: []byte{3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2606,14 +2779,16 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { }, limit: 10, valuesDeleted: 2, + nodesRemoved: 3, }, "delete branch one child only": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{3}, - Value: []byte{1, 2, 3}, + Key: []byte{3}, + Value: []byte{1, 2, 3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2621,24 +2796,27 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { }, limit: 1, newNode: &node.Branch{ - Key: []byte{3}, - Value: []byte{1, 2, 3}, - Dirty: true, - Generation: 1, + Key: []byte{3}, + Value: []byte{1, 2, 3}, + Dirty: true, + Generation: 1, + Descendants: 1, Children: [16]node.Node{ nil, &node.Leaf{Key: []byte{2}}, }, }, valuesDeleted: 1, + nodesRemoved: 1, }, "delete branch children only": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{3}, - Value: []byte{1, 2, 3}, + Key: []byte{3}, + Value: []byte{1, 2, 3}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{1}}, &node.Leaf{Key: []byte{2}}, @@ -2652,13 +2830,15 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { Generation: 1, }, valuesDeleted: 2, + nodesRemoved: 2, }, "delete branch all children except one": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{3}, + Key: []byte{3}, + Descendants: 3, Children: [16]node.Node{ nil, &node.Leaf{Key: []byte{1}}, @@ -2676,6 +2856,7 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { Dirty: true, }, valuesDeleted: 2, + nodesRemoved: 3, }, } @@ -2687,10 +2868,12 @@ func Test_Trie_deleteNodesLimit(t *testing.T) { trie := testCase.trie expectedTrie := *trie.DeepCopy() - newNode, valuesDeleted := trie.deleteNodesLimit(testCase.parent, testCase.prefix, testCase.limit) + newNode, valuesDeleted, nodesRemoved := + trie.deleteNodesLimit(testCase.parent, testCase.prefix, testCase.limit) assert.Equal(t, testCase.newNode, newNode) assert.Equal(t, testCase.valuesDeleted, valuesDeleted) + assert.Equal(t, testCase.nodesRemoved, nodesRemoved) assert.Equal(t, expectedTrie, trie) }) } @@ -2721,7 +2904,8 @@ func Test_Trie_ClearPrefix(t *testing.T) { "clear prefix": { trie: Trie{ root: &node.Branch{ - Key: []byte{1, 2}, + Key: []byte{1, 2}, + Descendants: 3, Children: [16]node.Node{ &node.Leaf{ // full key in nibbles 1, 2, 0, 5 Key: []byte{5}, @@ -2772,26 +2956,46 @@ func Test_Trie_clearPrefix(t *testing.T) { t.Parallel() testCases := map[string]struct { - trie Trie - parent Node - prefix []byte - newParent Node - updated bool + trie Trie + parent Node + prefix []byte + newParent Node + nodesRemoved uint32 }{ + "delete one of two children of branch": { + trie: Trie{ + generation: 1, + }, + parent: &node.Branch{ + Key: []byte{1}, + Descendants: 2, + Children: [16]node.Node{ + &node.Leaf{Key: []byte{3}}, + &node.Leaf{Key: []byte{4}}, + }, + }, + prefix: []byte{1, 0}, + newParent: &node.Leaf{ + Key: []byte{1, 1, 4}, + Dirty: true, + Generation: 1, + }, + nodesRemoved: 2, + }, "nil parent": {}, "leaf parent with common prefix": { parent: &node.Leaf{ Key: []byte{1, 2}, }, - prefix: []byte{1}, - updated: true, + prefix: []byte{1}, + nodesRemoved: 1, }, "leaf parent with key equal prefix": { parent: &node.Leaf{ Key: []byte{1}, }, - prefix: []byte{1}, - updated: true, + prefix: []byte{1}, + nodesRemoved: 1, }, "leaf parent with key no common prefix": { trie: Trie{ @@ -2819,41 +3023,45 @@ func Test_Trie_clearPrefix(t *testing.T) { }, "branch parent with common prefix": { parent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, }, - prefix: []byte{1}, - updated: true, + prefix: []byte{1}, + nodesRemoved: 2, }, "branch with key equal prefix": { parent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, }, - prefix: []byte{1, 2}, - updated: true, + prefix: []byte{1, 2}, + nodesRemoved: 2, }, "branch with no common prefix": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, }, prefix: []byte{1, 3}, newParent: &node.Branch{ - Key: []byte{1, 2}, - Value: []byte{1}, + Key: []byte{1, 2}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, @@ -2864,16 +3072,18 @@ func Test_Trie_clearPrefix(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, }, prefix: []byte{1, 2, 3}, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, @@ -2884,16 +3094,18 @@ func Test_Trie_clearPrefix(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, }, prefix: []byte{1, 2}, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, @@ -2904,8 +3116,9 @@ func Test_Trie_clearPrefix(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{Key: []byte{3}}, &node.Leaf{Key: []byte{4}}, @@ -2913,24 +3126,26 @@ func Test_Trie_clearPrefix(t *testing.T) { }, prefix: []byte{1, 0, 3}, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, - Dirty: true, - Generation: 1, + Key: []byte{1}, + Value: []byte{1}, + Dirty: true, + Generation: 1, + Descendants: 1, Children: [16]node.Node{ nil, &node.Leaf{Key: []byte{4}}, }, }, - updated: true, + nodesRemoved: 1, }, "fully delete child of branch": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{Key: []byte{3}}, }, @@ -2942,19 +3157,21 @@ func Test_Trie_clearPrefix(t *testing.T) { Dirty: true, Generation: 1, }, - updated: true, + nodesRemoved: 1, }, "partially delete child of branch": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 2, Children: [16]node.Node{ &node.Branch{ // full key 1, 0, 3 - Key: []byte{3}, - Value: []byte{1}, + Key: []byte{3}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // full key 1, 0, 3, 0, 5 Key: []byte{5}, @@ -2965,10 +3182,11 @@ func Test_Trie_clearPrefix(t *testing.T) { }, prefix: []byte{1, 0, 3, 0}, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, - Dirty: true, - Generation: 1, + Key: []byte{1}, + Value: []byte{1}, + Dirty: true, + Generation: 1, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // full key 1, 0, 3 Key: []byte{3}, @@ -2978,7 +3196,27 @@ func Test_Trie_clearPrefix(t *testing.T) { }, }, }, - updated: true, + nodesRemoved: 1, + }, + "delete one of two children of branch without value": { + trie: Trie{ + generation: 1, + }, + parent: &node.Branch{ + Key: []byte{1}, + Descendants: 2, + Children: [16]node.Node{ + &node.Leaf{Key: []byte{3}}, // full key 1, 0, 3 + &node.Leaf{Key: []byte{4}}, // full key 1, 1, 4 + }, + }, + prefix: []byte{1, 0, 3}, + newParent: &node.Leaf{ + Key: []byte{1, 1, 4}, + Dirty: true, + Generation: 1, + }, + nodesRemoved: 2, }, } @@ -2990,11 +3228,11 @@ func Test_Trie_clearPrefix(t *testing.T) { trie := testCase.trie expectedTrie := *trie.DeepCopy() - newParent, updated := trie.clearPrefix(testCase.parent, - testCase.prefix) + newParent, nodesRemoved := + trie.clearPrefix(testCase.parent, testCase.prefix) assert.Equal(t, testCase.newParent, newParent) - assert.Equal(t, testCase.updated, updated) + assert.Equal(t, testCase.nodesRemoved, nodesRemoved) assert.Equal(t, expectedTrie, trie) }) } @@ -3025,15 +3263,17 @@ func Test_Trie_Delete(t *testing.T) { trie: Trie{ generation: 1, root: &node.Branch{ - Key: []byte{1, 2}, + Key: []byte{1, 2}, + Descendants: 3, Children: [16]node.Node{ &node.Leaf{ Key: []byte{5}, Value: []byte{97}, }, &node.Branch{ // full key in nibbles 1, 2, 1, 6 - Key: []byte{6}, - Value: []byte{98}, + Key: []byte{6}, + Value: []byte{98}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // full key in nibbles 1, 2, 1, 6, 0, 7 Key: []byte{7}, @@ -3048,9 +3288,10 @@ func Test_Trie_Delete(t *testing.T) { expectedTrie: Trie{ generation: 1, root: &node.Branch{ - Key: []byte{1, 2}, - Dirty: true, - Generation: 1, + Key: []byte{1, 2}, + Dirty: true, + Generation: 1, + Descendants: 2, Children: [16]node.Node{ &node.Leaf{ Key: []byte{5}, @@ -3092,11 +3333,12 @@ func Test_Trie_delete(t *testing.T) { t.Parallel() testCases := map[string]struct { - trie Trie - parent Node - key []byte - newParent Node - updated bool + trie Trie + parent Node + key []byte + newParent Node + updated bool + nodesRemoved uint32 }{ "nil parent": { key: []byte{1}, @@ -3105,21 +3347,24 @@ func Test_Trie_delete(t *testing.T) { parent: &node.Leaf{ Key: []byte{1}, }, - updated: true, + updated: true, + nodesRemoved: 1, }, "leaf parent and empty key": { parent: &node.Leaf{ Key: []byte{1}, }, - key: []byte{}, - updated: true, + key: []byte{}, + updated: true, + nodesRemoved: 1, }, "leaf parent matches key": { parent: &node.Leaf{ Key: []byte{1}, }, - key: []byte{1}, - updated: true, + key: []byte{1}, + updated: true, + nodesRemoved: 1, }, "leaf parent mismatches key": { trie: Trie{ @@ -3138,8 +3383,9 @@ func Test_Trie_delete(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ Key: []byte{2}, @@ -3151,15 +3397,17 @@ func Test_Trie_delete(t *testing.T) { Dirty: true, Generation: 1, }, - updated: true, + updated: true, + nodesRemoved: 1, }, "branch parent and empty key": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ Key: []byte{2}, @@ -3172,15 +3420,17 @@ func Test_Trie_delete(t *testing.T) { Dirty: true, Generation: 1, }, - updated: true, + updated: true, + nodesRemoved: 1, }, "branch parent matches key": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ Key: []byte{2}, @@ -3193,15 +3443,17 @@ func Test_Trie_delete(t *testing.T) { Dirty: true, Generation: 1, }, - updated: true, + updated: true, + nodesRemoved: 1, }, "branch parent child matches key": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // full key 1, 0, 2 Key: []byte{2}, @@ -3215,23 +3467,26 @@ func Test_Trie_delete(t *testing.T) { Dirty: true, Generation: 1, }, - updated: true, + updated: true, + nodesRemoved: 1, }, "branch parent mismatches key": { trie: Trie{ generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, }, key: []byte{2}, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{}, }, @@ -3242,8 +3497,9 @@ func Test_Trie_delete(t *testing.T) { generation: 1, }, parent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // full key 1, 0, 2 Key: []byte{2}, @@ -3252,8 +3508,9 @@ func Test_Trie_delete(t *testing.T) { }, key: []byte{1, 0, 3}, newParent: &node.Branch{ - Key: []byte{1}, - Value: []byte{1}, + Key: []byte{1}, + Value: []byte{1}, + Descendants: 1, Children: [16]node.Node{ &node.Leaf{ // full key 1, 0, 2 Key: []byte{2}, @@ -3261,6 +3518,60 @@ func Test_Trie_delete(t *testing.T) { }, }, }, + "delete branch child and merge branch and left child": { + trie: Trie{ + generation: 1, + }, + parent: &node.Branch{ + Key: []byte{1}, + Descendants: 1, + Children: [16]node.Node{ + &node.Leaf{ // full key 1, 0, 2 + Key: []byte{2}, + Value: []byte{1}, + }, + &node.Leaf{ // full key 1, 1, 2 + Key: []byte{2}, + Value: []byte{2}, + }, + }, + }, + key: []byte{1, 0, 2}, + newParent: &node.Leaf{ + Key: []byte{1, 1, 2}, + Value: []byte{2}, + Dirty: true, + Generation: 1, + }, + updated: true, + nodesRemoved: 2, + }, + "delete branch and keep two children": { + trie: Trie{ + generation: 1, + }, + parent: &node.Branch{ + Key: []byte{1}, + Value: []byte{1}, + Descendants: 2, + Children: [16]node.Node{ + &node.Leaf{Key: []byte{2}}, + &node.Leaf{Key: []byte{2}}, + }, + }, + key: []byte{1}, + newParent: &node.Branch{ + Key: []byte{1}, + Generation: 1, + Dirty: true, + Descendants: 2, + Children: [16]node.Node{ + &node.Leaf{Key: []byte{2}}, + &node.Leaf{Key: []byte{2}}, + }, + }, + updated: true, + }, } for name, testCase := range testCases { @@ -3276,10 +3587,11 @@ func Test_Trie_delete(t *testing.T) { } expectedTrie := *testCase.trie.DeepCopy() - newParent, updated := testCase.trie.delete(testCase.parent, testCase.key) + newParent, updated, nodesRemoved := testCase.trie.delete(testCase.parent, testCase.key) 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, expectedKey, testCase.key) }) @@ -3290,9 +3602,10 @@ func Test_handleDeletion(t *testing.T) { t.Parallel() testCases := map[string]struct { - branch *node.Branch - deletedKey []byte - newNode Node + branch *node.Branch + deletedKey []byte + newNode Node + branchChildMerged bool }{ "branch with value and without children": { branch: &node.Branch{ @@ -3350,6 +3663,7 @@ func Test_handleDeletion(t *testing.T) { Generation: 1, Dirty: true, }, + branchChildMerged: true, }, "branch without value and a single branch child": { branch: &node.Branch{ @@ -3379,6 +3693,7 @@ func Test_handleDeletion(t *testing.T) { &node.Leaf{Key: []byte{8}}, }, }, + branchChildMerged: true, }, } @@ -3394,9 +3709,10 @@ func Test_handleDeletion(t *testing.T) { copy(expectedKey, testCase.deletedKey) } - newNode := handleDeletion(testCase.branch, testCase.deletedKey) + newNode, branchChildMerged := handleDeletion(testCase.branch, testCase.deletedKey) assert.Equal(t, testCase.newNode, newNode) + assert.Equal(t, testCase.branchChildMerged, branchChildMerged) assert.Equal(t, expectedKey, testCase.deletedKey) }) }