Skip to content

Commit

Permalink
backwards iteration
Browse files Browse the repository at this point in the history 
Added `DescendRange`, `DescendLessOrEqual`, `DescendGreaterThan`,
and `Descend` functions. These are modeled after the existing
`Ascend..` methods.
  • Loading branch information
@tidwall
tidwall committed May 30, 2016
1 parent 7d79101 commit 1826edc
Show file tree
Hide file tree
Showing 2 changed files with 205 additions and 38 deletions.
143 changes: 105 additions & 38 deletions btree.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@
// widely used ordered tree implementation in the Go ecosystem currently.
// Its functions, therefore, exactly mirror those of
// llrb.LLRB where possible. Unlike gollrb, though, we currently don't
// support storing multiple equivalent values or backwards iteration.
// support storing multiple equivalent values.
package btree

import (
Expand Down Expand Up @@ -419,34 +419,77 @@ func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove)
return n.remove(item, minItems, typ)
}

type direction int

const (
descend = direction(-1)
ascend = direction(+1)
)

// iterate provides a simple method for iterating over elements in the tree.
// It could probably use some work to be extra-efficient (it calls from() a
// little more than it should), but it works pretty well for now.
//
// It requires that 'from' and 'to' both return true for values we should hit
// with the iterator. It should also be the case that 'from' returns true for
// values less than or equal to values 'to' returns true for, and 'to'
// returns true for values greater than or equal to those that 'from'
// does.
func (n *node) iterate(from, to func(Item) bool, iter ItemIterator) bool {
for i, item := range n.items {
if !from(item) {
continue
// When ascending, the 'start' should be less than 'stop' and when descending,
// the 'start' should be greater than 'stop'. Setting 'includeStart' to true
// will force the iterator to include the first item when it equals 'start',
// thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
// "greaterThan" or "lessThan" queries.
func (n *node) iterate(dir direction, start, stop Item, includeStart bool, hit bool, iter ItemIterator) (bool, bool) {
var ok bool
switch dir {
case ascend:
for i := 0; i < len(n.items); i++ {
if start != nil && n.items[i].Less(start) {
continue
}
if len(n.children) > 0 {
if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter); !ok {
return hit, false
}
}
if !includeStart && !hit && start != nil && !start.Less(n.items[i]) {
hit = true
continue
}
hit = true
if stop != nil && !n.items[i].Less(stop) {
return hit, false
}
if !iter(n.items[i]) {
return hit, false
}
}
if len(n.children) > 0 && !n.children[i].iterate(from, to, iter) {
return false
if len(n.children) > 0 {
if hit, ok = n.children[len(n.children)-1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
return hit, false
}
}
if !to(item) {
return false
case descend:
for i := len(n.items) - 1; i >= 0; i-- {
if start != nil && !n.items[i].Less(start) {
if !includeStart || hit || start.Less(n.items[i]) {
continue
}
}
if len(n.children) > 0 {
if hit, ok = n.children[i+1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
return hit, false
}
}
if stop != nil && !stop.Less(n.items[i]) {
return hit, false // continue
}
hit = true
if !iter(n.items[i]) {
return hit, false
}
}
if !iter(item) {
return false
if len(n.children) > 0 {
if hit, ok = n.children[0].iterate(dir, start, stop, includeStart, hit, iter); !ok {
return hit, false
}
}
}
if len(n.children) > 0 {
return n.children[len(n.children)-1].iterate(from, to, iter)
}
return true
return hit, true
}

// Used for testing/debugging purposes.
Expand Down Expand Up @@ -569,10 +612,7 @@ func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator
if t.root == nil {
return
}
t.root.iterate(
func(a Item) bool { return !a.Less(greaterOrEqual) },
func(a Item) bool { return a.Less(lessThan) },
iterator)
t.root.iterate(ascend, greaterOrEqual, lessThan, true, false, iterator)
}

// AscendLessThan calls the iterator for every value in the tree within the range
Expand All @@ -581,10 +621,7 @@ func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
if t.root == nil {
return
}
t.root.iterate(
func(a Item) bool { return true },
func(a Item) bool { return a.Less(pivot) },
iterator)
t.root.iterate(ascend, nil, pivot, false, false, iterator)
}

// AscendGreaterOrEqual calls the iterator for every value in the tree within
Expand All @@ -593,10 +630,7 @@ func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
if t.root == nil {
return
}
t.root.iterate(
func(a Item) bool { return !a.Less(pivot) },
func(a Item) bool { return true },
iterator)
t.root.iterate(ascend, pivot, nil, true, false, iterator)
}

// Ascend calls the iterator for every value in the tree within the range
Expand All @@ -605,10 +639,43 @@ func (t *BTree) Ascend(iterator ItemIterator) {
if t.root == nil {
return
}
t.root.iterate(
func(a Item) bool { return true },
func(a Item) bool { return true },
iterator)
t.root.iterate(ascend, nil, nil, false, false, iterator)
}

// DescendRange calls the iterator for every value in the tree within the range
// [lessOrEqual, greaterThan), until iterator returns false.
func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
if t.root == nil {
return
}
t.root.iterate(descend, lessOrEqual, greaterThan, true, false, iterator)
}

// DescendLessOrEqual calls the iterator for every value in the tree within the range
// [pivot, first], until iterator returns false.
func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
if t.root == nil {
return
}
t.root.iterate(descend, pivot, nil, true, false, iterator)
}

// DescendGreaterThan calls the iterator for every value in the tree within
// the range [last, pivot), until iterator returns false.
func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
if t.root == nil {
return
}
t.root.iterate(descend, nil, pivot, false, false, iterator)
}

// Descend calls the iterator for every value in the tree within the range
// [last, first], until iterator returns false.
func (t *BTree) Descend(iterator ItemIterator) {
if t.root == nil {
return
}
t.root.iterate(descend, nil, nil, false, false, iterator)
}

// Get looks for the key item in the tree, returning it. It returns nil if
Expand Down
100 changes: 100 additions & 0 deletions btree_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,23 @@ func all(t *BTree) (out []Item) {
return
}

// rangerev returns a reversed ordered list of Int items in the range [0, n).
func rangrev(n int) (out []Item) {
for i := n - 1; i >= 0; i-- {
out = append(out, Int(i))
}
return
}

// allrev extracts all items from a tree in reverse order as a slice.
func allrev(t *BTree) (out []Item) {
t.Descend(func(a Item) bool {
out = append(out, a)
return true
})
return
}

var btreeDegree = flag.Int("degree", 32, "B-Tree degree")

func TestBTree(t *testing.T) {
Expand Down Expand Up @@ -87,6 +104,13 @@ func TestBTree(t *testing.T) {
if !reflect.DeepEqual(got, want) {
t.Fatalf("mismatch:\n got: %v\nwant: %v", got, want)
}

gotrev := allrev(tr)
wantrev := rangrev(treeSize)
if !reflect.DeepEqual(gotrev, wantrev) {
t.Fatalf("mismatch:\n got: %v\nwant: %v", got, want)
}

for _, item := range perm(treeSize) {
if x := tr.Delete(item); x == nil {
t.Fatalf("didn't find %v", item)
Expand Down Expand Up @@ -188,6 +212,31 @@ func TestAscendRange(t *testing.T) {
}
}

func TestDescendRange(t *testing.T) {
tr := New(2)
for _, v := range perm(100) {
tr.ReplaceOrInsert(v)
}
var got []Item
tr.DescendRange(Int(60), Int(40), func(a Item) bool {
got = append(got, a)
return true
})
if want := rangrev(100)[39:59]; !reflect.DeepEqual(got, want) {
t.Fatalf("descendrange:\n got: %v\nwant: %v", got, want)
}
got = got[:0]
tr.DescendRange(Int(60), Int(40), func(a Item) bool {
if a.(Int) < 50 {
return false
}
got = append(got, a)
return true
})
if want := rangrev(100)[39:50]; !reflect.DeepEqual(got, want) {
t.Fatalf("descendrange:\n got: %v\nwant: %v", got, want)
}
}
func TestAscendLessThan(t *testing.T) {
tr := New(*btreeDegree)
for _, v := range perm(100) {
Expand All @@ -214,6 +263,31 @@ func TestAscendLessThan(t *testing.T) {
}
}

func TestDescendLessOrEqual(t *testing.T) {
tr := New(*btreeDegree)
for _, v := range perm(100) {
tr.ReplaceOrInsert(v)
}
var got []Item
tr.DescendLessOrEqual(Int(40), func(a Item) bool {
got = append(got, a)
return true
})
if want := rangrev(100)[59:]; !reflect.DeepEqual(got, want) {
t.Fatalf("descendlessorequal:\n got: %v\nwant: %v", got, want)
}
got = got[:0]
tr.DescendLessOrEqual(Int(60), func(a Item) bool {
if a.(Int) < 50 {
return false
}
got = append(got, a)
return true
})
if want := rangrev(100)[39:50]; !reflect.DeepEqual(got, want) {
t.Fatalf("descendlessorequal:\n got: %v\nwant: %v", got, want)
}
}
func TestAscendGreaterOrEqual(t *testing.T) {
tr := New(*btreeDegree)
for _, v := range perm(100) {
Expand All @@ -240,6 +314,32 @@ func TestAscendGreaterOrEqual(t *testing.T) {
}
}

func TestDescendGreaterThan(t *testing.T) {
tr := New(*btreeDegree)
for _, v := range perm(100) {
tr.ReplaceOrInsert(v)
}
var got []Item
tr.DescendGreaterThan(Int(40), func(a Item) bool {
got = append(got, a)
return true
})
if want := rangrev(100)[:59]; !reflect.DeepEqual(got, want) {
t.Fatalf("descendgreaterthan:\n got: %v\nwant: %v", got, want)
}
got = got[:0]
tr.DescendGreaterThan(Int(40), func(a Item) bool {
if a.(Int) < 50 {
return false
}
got = append(got, a)
return true
})
if want := rangrev(100)[:50]; !reflect.DeepEqual(got, want) {
t.Fatalf("descendgreaterthan:\n got: %v\nwant: %v", got, want)
}
}

const benchmarkTreeSize = 10000

func BenchmarkInsert(b *testing.B) {
Expand Down

0 comments on commit 1826edc

Please sign in to comment.