diff --git a/internal/trie/node/branch.go b/internal/trie/node/branch.go index f7e4d28799b..5f600b860bb 100644 --- a/internal/trie/node/branch.go +++ b/internal/trie/node/branch.go @@ -33,6 +33,7 @@ type Branch struct { // inserted, moved or iterated over. Generation uint64 sync.RWMutex + Stats Stats } // NewBranch creates a new branch using the arguments given. diff --git a/internal/trie/node/stats.go b/internal/trie/node/stats.go new file mode 100644 index 00000000000..a438663abe9 --- /dev/null +++ b/internal/trie/node/stats.go @@ -0,0 +1,23 @@ +package node + +// Stats contains statistical information for a node. +type Stats struct { + // Descendants is the number of descendant nodes for + // this particular node. + Descendants uint32 +} + +// GetDescendants returns the number of descendants in the branch. +func (b *Branch) GetDescendants() (descendants uint32) { + return b.Stats.Descendants +} + +// AddDescendants adds descendant nodes count to the node stats. +func (b *Branch) AddDescendants(n uint32) { + b.Stats.Descendants += n +} + +// SubDescendants subtracts descendant nodes count from the node stats. +func (b *Branch) SubDescendants(n uint32) { + b.Stats.Descendants -= n +} diff --git a/internal/trie/node/stats_test.go b/internal/trie/node/stats_test.go new file mode 100644 index 00000000000..2b4023a62f5 --- /dev/null +++ b/internal/trie/node/stats_test.go @@ -0,0 +1 @@ +package node diff --git a/lib/trie/trie.go b/lib/trie/trie.go index 6e960e2cba7..7f8f28881d5 100644 --- a/lib/trie/trie.go +++ b/lib/trie/trie.go @@ -306,19 +306,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 { + nodesCreated = 1 return &node.Leaf{ Key: key, Value: value, Generation: t.generation, Dirty: true, - } + }, nodesCreated } // TODO ensure all values have dirty set to true @@ -337,15 +338,15 @@ 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) { if !bytes.Equal(value, parentLeaf.Value) { parentLeaf.Value = value parentLeaf.Generation = t.generation parentLeaf.SetDirty(true) } - return parentLeaf + return parentLeaf, 0 } commonPrefixLength := lenCommonPrefix(key, parentLeaf.Key) @@ -368,9 +369,11 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, parentLeaf.Key = parentLeaf.Key[commonPrefixLength+1:] parentLeaf.SetDirty(true) newBranchParent.Children[childIndex] = parentLeaf + newBranchParent.AddDescendants(1) + nodesCreated++ } - return newBranchParent + return newBranchParent, nodesCreated } if len(parentLeaf.Key) == commonPrefixLength { @@ -382,6 +385,8 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key, parentLeaf.Key = parentLeaf.Key[commonPrefixLength+1:] parentLeaf.SetDirty(true) newBranchParent.Children[childIndex] = parentLeaf + newBranchParent.AddDescendants(1) + nodesCreated++ } childIndex := key[commonPrefixLength] newBranchParent.Children[childIndex] = &node.Leaf{ @@ -390,16 +395,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) { if bytes.Equal(key, parentBranch.Key) { parentBranch.SetDirty(true) parentBranch.Generation = t.generation parentBranch.Value = value - return parentBranch + return parentBranch, 0 } if bytes.HasPrefix(key, parentBranch.Key) { @@ -416,19 +424,23 @@ 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) child.SetDirty(true) } parentBranch.Children[childIndex] = child + parentBranch.AddDescendants(nodesCreated) parentBranch.SetDirty(true) parentBranch.Generation = t.generation - return parentBranch + + 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], @@ -443,16 +455,20 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (new parentBranch.Key = remainingOldParentKey parentBranch.Generation = t.generation 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. @@ -640,16 +656,16 @@ 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 } // clearPrefixLimit deletes the keys having the prefix till limit reached and returns updated trie root node, // true if any node in the trie got updated, and next bool returns true if there is no keys left with 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 } newParent = parent @@ -663,17 +679,18 @@ 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 } // not modified so return the leaf of the original // trie generation. The copied leaf newParent will be // garbage collected. - return parent, 0, allDeleted + return parent, 0, 0, allDeleted } branch := newParent.(*node.Branch) - newParent, valuesDeleted, allDeleted = t.clearPrefixLimitBranch(branch, prefix, limit) + newParent, valuesDeleted, nodesRemoved, allDeleted = + t.clearPrefixLimitBranch(branch, prefix, limit) if valuesDeleted == 0 { // not modified so return the node of the original // trie generation. The copied newParent will be @@ -681,18 +698,18 @@ func (t *Trie) clearPrefixLimit(parent Node, prefix []byte, limit uint32) ( newParent = parent } - return newParent, valuesDeleted, allDeleted + return newParent, valuesDeleted, nodesRemoved, allDeleted } 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 && @@ -704,9 +721,9 @@ 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)] @@ -714,46 +731,56 @@ func (t *Trie) clearPrefixLimitBranch(branch *node.Branch, prefix []byte, limit child := branch.Children[childIndex] newParent = branch // mostly just a reminder for the reader - branch.Children[childIndex], valuesDeleted, allDeleted = t.clearPrefixLimit(child, childPrefix, limit) + branch.Children[childIndex], valuesDeleted, nodesRemoved, allDeleted = t.clearPrefixLimit(child, childPrefix, limit) if valuesDeleted > 0 { branch.SetDirty(true) - newParent = handleDeletion(branch, prefix) + var branchChildMerged bool + newParent, branchChildMerged = handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } } - return newParent, valuesDeleted, allDeleted + 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 { + 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) - branch.Children[childIndex], valuesDeleted = t.deleteNodesLimit(child, nilPrefix, limit) + branch.Children[childIndex], valuesDeleted, nodesRemoved = t.deleteNodesLimit(child, nilPrefix, limit) branch.SetDirty(true) - newParent = handleDeletion(branch, prefix) + newParent, branchChildMerged := handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } 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 } newParent = parent @@ -762,8 +789,8 @@ func (t *Trie) deleteNodesLimit(parent Node, prefix []byte, limit uint32) ( } if newParent.Type() == node.LeafType { - valuesDeleted = 1 - return nil, valuesDeleted + valuesDeleted, nodesRemoved = 1, 1 + return nil, valuesDeleted, nodesRemoved } branch := newParent.(*node.Branch) @@ -772,36 +799,45 @@ 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.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 = handleDeletion(branch, fullKey) + + newParent, branchChildMerged = handleDeletion(branch, fullKey) + if branchChildMerged { + nodesRemoved++ + } + 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 @@ -815,13 +851,13 @@ func (t *Trie) ClearPrefix(prefixLE []byte) { prefix := codec.KeyLEToNibbles(prefixLE) prefix = bytes.TrimSuffix(prefix, []byte{0}) - t.root, _ = t.clearPrefix(t.root, prefix) + t.root, _, _ = t.clearPrefix(t.root, prefix) } func (t *Trie) clearPrefix(parent Node, prefix []byte) ( - newParent Node, updated bool) { + newParent Node, updated bool, nodesRemoved uint32) { if parent == nil { - return nil, false + return nil, false, 0 } newParent = parent @@ -830,14 +866,18 @@ func (t *Trie) clearPrefix(parent Node, prefix []byte) ( } if bytes.HasPrefix(newParent.GetKey(), prefix) { - return nil, true + nodesRemoved = 1 // leaf + if newParent.Type() != node.LeafType { // branch + nodesRemoved = newParent.(*node.Branch).GetDescendants() + } + return nil, true, nodesRemoved } if newParent.Type() == node.LeafType { // not modified so return the leaf of the original // trie generation. The copied newParent will be // garbage collected. - return parent, false + return parent, false, 0 } branch := newParent.(*node.Branch) @@ -853,13 +893,17 @@ func (t *Trie) clearPrefix(parent Node, prefix []byte) ( // node is not modified so return the branch of the original // trie generation. The copied newParent will be // garbage collected. - return parent, false + const nodesRemoved = 0 + return parent, false, nodesRemoved } branch.Children[childIndex] = nil branch.SetDirty(true) - newParent = handleDeletion(branch, prefix) - return newParent, true + newParent, branchChildMerged := handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } + return newParent, true, nodesRemoved } noPrefixForNode := len(prefix) <= len(branch.Key) || @@ -868,24 +912,30 @@ func (t *Trie) clearPrefix(parent Node, prefix []byte) ( // not modified so return the branch of the original // trie generation. The copied newParent will be // garbage collected. - return parent, false + const nodesRemoved = 0 + return parent, false, nodesRemoved } childIndex := prefix[len(branch.Key)] childPrefix := prefix[len(branch.Key)+1:] child := branch.Children[childIndex] - branch.Children[childIndex], updated = t.clearPrefix(child, childPrefix) + branch.Children[childIndex], updated, nodesRemoved = t.clearPrefix(child, childPrefix) if !updated { // branch not modified so return the branch of the original // trie generation. The copied newParent will be // garbage collected. - return parent, false + const nodesRemoved = 0 + return parent, false, nodesRemoved } branch.SetDirty(true) - newParent = handleDeletion(branch, prefix) - return newParent, true + newParent, branchChildMerged := handleDeletion(branch, prefix) + if branchChildMerged { + nodesRemoved++ + } + + return newParent, true, nodesRemoved } // Delete removes the node of the trie with the key @@ -893,12 +943,14 @@ 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 } newParent = parent @@ -909,24 +961,29 @@ func (t *Trie) delete(parent Node, key []byte) (newParent Node, deleted bool) { if newParent.Type() == node.LeafType { newParent = deleteLeaf(newParent, key) if newParent == nil { - return nil, true + const nodesRemoved = 1 + return nil, true, nodesRemoved } // The leaf was not deleted so return the original // parent without its generation updated. // The copied newParent will be garbage collected. - return parent, false + const nodesRemoved = 0 + return parent, false, nodesRemoved } branch := newParent.(*node.Branch) - newParent, deleted = t.deleteBranch(branch, key) + newParent, deleted, nodesRemoved = t.deleteBranch(branch, key) + // Note deleted is true and nodesRemoved is 0 if the branch + // has more than one child since no node is deleted. if !deleted { // Nothing was deleted so return the original // parent without its generation updated. // The copied newParent will be garbage collected. - return parent, false + const nodesRemoved = 0 + return parent, false, nodesRemoved } - return newParent, true + return newParent, true, nodesRemoved } func deleteLeaf(parent Node, key []byte) (newParent Node) { @@ -936,11 +993,17 @@ 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.Value = nil branch.SetDirty(true) - return handleDeletion(branch, key), true + deleted = true + newParent, branchChildMerged := handleDeletion(branch, key) + if branchChildMerged { + nodesRemoved = 1 + } + return newParent, deleted, nodesRemoved } commonPrefixLength := lenCommonPrefix(branch.Key, key) @@ -948,22 +1011,30 @@ 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, 0 } branch.Children[childIndex] = newChild branch.SetDirty(true) - newParent = handleDeletion(branch, key) - return newParent, true + + newParent, branchChildMerged := handleDeletion(branch, key) + if branchChildMerged { + nodesRemoved = 1 + } + + 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) { +func handleDeletion(branch *node.Branch, key []byte) (newNode Node, branchChildMerged bool) { // TODO try to remove key argument just use branch.Key instead? childrenCount := 0 firstChildIndex := -1 @@ -979,16 +1050,18 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { switch { default: - return branch + return branch, false case childrenCount == 0 && branch.Value != nil: + 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: + branchChildMerged = true childIndex := firstChildIndex child := branch.Children[firstChildIndex] @@ -1000,7 +1073,7 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { Value: child.GetValue(), Dirty: true, Generation: branch.Generation, - } + }, branchChildMerged } childBranch := child.(*node.Branch) @@ -1010,6 +1083,10 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { Value: childBranch.Value, Generation: branch.Generation, Dirty: true, + Stats: node.Stats{ + // this is the descendants of the original branch minus one + Descendants: childBranch.GetDescendants(), + }, } // Adopt the grand-children @@ -1021,7 +1098,7 @@ func handleDeletion(branch *node.Branch, key []byte) (newNode Node) { } } - return newBranch + return newBranch, branchChildMerged } }