diff --git a/trees/btree/btree_test.go b/trees/btree/btree_test.go index dde69ef..71e86d2 100644 --- a/trees/btree/btree_test.go +++ b/trees/btree/btree_test.go @@ -384,6 +384,345 @@ func TestBTreeLeftAndRight(t *testing.T) { } } +func TestBTreeIteratorNextOnEmpty(t *testing.T) { + tree := NewWithIntComparator(3) + it := tree.Iterator() + for it.Next() { + t.Errorf("Shouldn't iterate on empty tree") + } +} + +func TestBTreeIteratorPrevOnEmpty(t *testing.T) { + tree := NewWithIntComparator(3) + it := tree.Iterator() + for it.Prev() { + t.Errorf("Shouldn't iterate on empty tree") + } +} + +func TestBTreeIterator1Next(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") + tree.Put(2, "b") + tree.Put(1, "a") //overwrite + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + key := it.Key() + switch key { + case count: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIterator1Prev(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") + tree.Put(2, "b") + tree.Put(1, "a") //overwrite + it := tree.Iterator() + for it.Next() { + } + countDown := tree.size + for it.Prev() { + key := it.Key() + switch key { + case countDown: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + countDown-- + } + if actualValue, expectedValue := countDown, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIterator2Next(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + key := it.Key() + switch key { + case count: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIterator2Prev(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + for it.Next() { + } + countDown := tree.size + for it.Prev() { + key := it.Key() + switch key { + case countDown: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + countDown-- + } + if actualValue, expectedValue := countDown, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIterator3Next(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(1, "a") + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + key := it.Key() + switch key { + case count: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIterator3Prev(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(1, "a") + it := tree.Iterator() + for it.Next() { + } + countDown := tree.size + for it.Prev() { + key := it.Key() + switch key { + case countDown: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + countDown-- + } + if actualValue, expectedValue := countDown, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIterator4Next(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(13, 5) + tree.Put(8, 3) + tree.Put(17, 7) + tree.Put(1, 1) + tree.Put(11, 4) + tree.Put(15, 6) + tree.Put(25, 9) + tree.Put(6, 2) + tree.Put(22, 8) + tree.Put(27, 10) + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + value := it.Value() + switch value { + case count: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIterator4Prev(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(13, 5) + tree.Put(8, 3) + tree.Put(17, 7) + tree.Put(1, 1) + tree.Put(11, 4) + tree.Put(15, 6) + tree.Put(25, 9) + tree.Put(6, 2) + tree.Put(22, 8) + tree.Put(27, 10) + it := tree.Iterator() + count := tree.Size() + for it.Next() { + } + for it.Prev() { + value := it.Value() + switch value { + case count: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + count-- + } + if actualValue, expectedValue := count, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestBTreeIteratorBegin(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.Begin() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + for it.Next() { + } + + it.Begin() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.Next() + if key, value := it.Key(), it.Value(); key != 1 || value != "a" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 1, "a") + } +} + +func TestBTreeIteratorEnd(t *testing.T) { + tree := NewWithIntComparator(3) + it := tree.Iterator() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.End() + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it.End() + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.Prev() + if key, value := it.Key(), it.Value(); key != 3 || value != "c" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 3, "c") + } +} + +func TestBTreeIteratorFirst(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + if actualValue, expectedValue := it.First(), true; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if key, value := it.Key(), it.Value(); key != 1 || value != "a" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 1, "a") + } +} + +func TestBTreeIteratorLast(t *testing.T) { + tree := NewWithIntComparator(3) + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + if actualValue, expectedValue := it.Last(), true; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if key, value := it.Key(), it.Value(); key != 3 || value != "c" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 3, "c") + } +} + func TestBTree_search(t *testing.T) { { tree := NewWithIntComparator(3) diff --git a/trees/btree/iterator.go b/trees/btree/iterator.go new file mode 100644 index 0000000..840db68 --- /dev/null +++ b/trees/btree/iterator.go @@ -0,0 +1,193 @@ +// Copyright (c) 2015, Emir Pasic. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package btree + +import "github.com/emirpasic/gods/containers" + +func assertIteratorImplementation() { + var _ containers.ReverseIteratorWithKey = (*Iterator)(nil) +} + +// Iterator holding the iterator's state +type Iterator struct { + tree *Tree + node *Node + entry *Entry + position position +} + +type position byte + +const ( + begin, between, end position = 0, 1, 2 +) + +// Iterator returns a stateful iterator whose elements are key/value pairs. +func (tree *Tree) Iterator() Iterator { + return Iterator{tree: tree, node: nil, position: begin} +} + +// Next moves the iterator to the next element and returns true if there was a next element in the container. +// If Next() returns true, then next element's key and value can be retrieved by Key() and Value(). +// If Next() was called for the first time, then it will point the iterator to the first element if it exists. +// Modifies the state of the iterator. +func (iterator *Iterator) Next() bool { + // If already at end, go to end + if iterator.position == end { + goto end + } + // If at beginning, get the left-most entry in the tree + if iterator.position == begin { + left := iterator.tree.Left() + if left == nil { + goto end + } + iterator.node = left + iterator.entry = left.Entries[0] + goto between + } + { + // Find current entry position in current node + e, _ := iterator.tree.search(iterator.node, iterator.entry.Key) + // Try to go down to the child right of the current entry + if e+1 < len(iterator.node.Children) { + iterator.node = iterator.node.Children[e+1] + // Try to go down to the child left of the current node + for len(iterator.node.Children) > 0 { + iterator.node = iterator.node.Children[0] + } + // Return the left-most entry + iterator.entry = iterator.node.Entries[0] + goto between + } + // Above assures that we have reached a leaf node, so return the next entry in current node (if any) + if e+1 < len(iterator.node.Entries) { + iterator.entry = iterator.node.Entries[e+1] + goto between + } + } + // Reached leaf node and there are no entries to the right of the current entry, so go up to the parent + for iterator.node.Parent != nil { + iterator.node = iterator.node.Parent + // Find next entry position in current node (note: search returns the first equal or bigger than entry) + e, _ := iterator.tree.search(iterator.node, iterator.entry.Key) + // Check that there is a next entry position in current node + if e < len(iterator.node.Entries) { + iterator.entry = iterator.node.Entries[e] + goto between + } + } + +end: + iterator.End() + return false + +between: + iterator.position = between + return true +} + +// Prev moves the iterator to the previous element and returns true if there was a previous element in the container. +// If Prev() returns true, then previous element's key and value can be retrieved by Key() and Value(). +// Modifies the state of the iterator. +func (iterator *Iterator) Prev() bool { + // If already at beginning, go to begin + if iterator.position == begin { + goto begin + } + // If at end, get the right-most entry in the tree + if iterator.position == end { + right := iterator.tree.Right() + if right == nil { + goto begin + } + iterator.node = right + iterator.entry = right.Entries[len(right.Entries)-1] + goto between + } + { + // Find current entry position in current node + e, _ := iterator.tree.search(iterator.node, iterator.entry.Key) + // Try to go down to the child left of the current entry + if e < len(iterator.node.Children) { + iterator.node = iterator.node.Children[e] + // Try to go down to the child right of the current node + for len(iterator.node.Children) > 0 { + iterator.node = iterator.node.Children[len(iterator.node.Children)-1] + } + // Return the right-most entry + iterator.entry = iterator.node.Entries[len(iterator.node.Entries)-1] + goto between + } + // Above assures that we have reached a leaf node, so return the previous entry in current node (if any) + if e-1 >= 0 { + iterator.entry = iterator.node.Entries[e-1] + goto between + } + } + // Reached leaf node and there are no entries to the left of the current entry, so go up to the parent + for iterator.node.Parent != nil { + iterator.node = iterator.node.Parent + // Find previous entry position in current node (note: search returns the first equal or bigger than entry) + e, _ := iterator.tree.search(iterator.node, iterator.entry.Key) + // Check that there is a previous entry position in current node + if e-1 >= 0 { + iterator.entry = iterator.node.Entries[e-1] + goto between + } + } + +begin: + iterator.Begin() + return false + +between: + iterator.position = between + return true +} + +// Value returns the current element's value. +// Does not modify the state of the iterator. +func (iterator *Iterator) Value() interface{} { + return iterator.entry.Value +} + +// Key returns the current element's key. +// Does not modify the state of the iterator. +func (iterator *Iterator) Key() interface{} { + return iterator.entry.Key +} + +// Begin resets the iterator to its initial state (one-before-first) +// Call Next() to fetch the first element if any. +func (iterator *Iterator) Begin() { + iterator.node = nil + iterator.position = begin + iterator.entry = nil +} + +// End moves the iterator past the last element (one-past-the-end). +// Call Prev() to fetch the last element if any. +func (iterator *Iterator) End() { + iterator.node = nil + iterator.position = end + iterator.entry = nil +} + +// First moves the iterator to the first element and returns true if there was a first element in the container. +// If First() returns true, then first element's key and value can be retrieved by Key() and Value(). +// Modifies the state of the iterator +func (iterator *Iterator) First() bool { + iterator.Begin() + return iterator.Next() +} + +// Last moves the iterator to the last element and returns true if there was a last element in the container. +// If Last() returns true, then last element's key and value can be retrieved by Key() and Value(). +// Modifies the state of the iterator. +func (iterator *Iterator) Last() bool { + iterator.End() + return iterator.Prev() +}