diff --git a/btree.go b/btree.go index eb74b1d..fc118f0 100644 --- a/btree.go +++ b/btree.go @@ -298,6 +298,20 @@ func (n *node) maybeSplitChild(i, maxItems int) bool { return true } +// reset will free all nodes within the given cow ctx. +func (n *node) reset(c *copyOnWriteContext) { + for i := len(n.items) - 1; i >= 0; i-- { + if len(n.children) > 0 { + n.children[i+1].reset(c) + } + } + if len(n.children) > 0 { + n.children[0].reset(c) + } else { + c.freeNode(n) + } +} + // insert inserts an item into the subtree rooted at this node, making sure // no nodes in the subtree exceed maxItems items. Should an equivalent item be // be found/replaced by insert, it will be returned. @@ -618,6 +632,20 @@ func (t *BTree) Clone() (t2 *BTree) { return &out } +// Reset removes all items from the btree. +func (t *BTree) Reset() { + if t.root == nil || len(t.root.items) == 0 { + return + } + t.root = t.root.mutableFor(t.cow) + t.root.reset(t.cow) + for _, n := range t.root.children { + t.cow.freeNode(n) + } + t.cow.freeNode(t.root) + t.root, t.length = nil, 0 +} + // maxItems returns the max number of items to allow per node. func (t *BTree) maxItems() int { return t.degree*2 - 1 diff --git a/btree_test.go b/btree_test.go index 5da9d8b..81fcb89 100644 --- a/btree_test.go +++ b/btree_test.go @@ -431,6 +431,48 @@ func BenchmarkDelete(b *testing.B) { } } +func BenchmarkDeleteAll(b *testing.B) { + b.StopTimer() + insertP := perm(benchmarkTreeSize) + b.StartTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + tr := New(*btreeDegree) + for _, v := range insertP { + tr.ReplaceOrInsert(v) + } + + dels := make([]Item, 0, tr.Len()) + tr.Ascend(ItemIterator(func(b Item) bool { + dels = append(dels, b) + return true + })) + for _, del := range dels { + tr.Delete(del) + } + if tr.Len() > 0 { + b.Fatal(`len must be zero`) + } + } +} + +func BenchmarkReset(b *testing.B) { + b.StopTimer() + insertP := perm(benchmarkTreeSize) + b.StartTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + tr := New(*btreeDegree) + for _, v := range insertP { + tr.ReplaceOrInsert(v) + } + tr.Reset() + if tr.Len() > 0 { + b.Fatal(`len must be zero`) + } + } +} + func BenchmarkGet(b *testing.B) { b.StopTimer() insertP := perm(benchmarkTreeSize)