- test iterator end on reverse-iterable data structures

- fix red-black tree

lists/arraylist/arraylist_test.go

@@ -393,6 +393,31 @@ func TestListIteratorBegin(t *testing.T) {
 	}
 }
 
+func TestListIteratorEnd(t *testing.T) {
+	list := New()
+	it := list.Iterator()
+
+	if index := it.Index(); index != -1 {
+		t.Errorf("Got %v expected %v", index, -1)
+	}
+
+	it.End()
+	if index := it.Index(); index != 0 {
+		t.Errorf("Got %v expected %v", index, 0)
+	}
+
+	list.Add("a", "b", "c")
+	it.End()
+	if index := it.Index(); index != list.Size() {
+		t.Errorf("Got %v expected %v", index, list.Size())
+	}
+
+	it.Prev()
+	if index, value := it.Index(), it.Value(); index != list.Size()-1 || value != "c" {
+		t.Errorf("Got %v,%v expected %v,%v", index, value, list.Size()-1, "c")
+	}
+}
+
 func TestListIteratorFirst(t *testing.T) {
 	list := New()
 	it := list.Iterator()

lists/doublylinkedlist/doublylinkedlist_test.go

@@ -393,6 +393,31 @@ func TestListIteratorBegin(t *testing.T) {
 	}
 }
 
+func TestListIteratorEnd(t *testing.T) {
+	list := New()
+	it := list.Iterator()
+
+	if index := it.Index(); index != -1 {
+		t.Errorf("Got %v expected %v", index, -1)
+	}
+
+	it.End()
+	if index := it.Index(); index != 0 {
+		t.Errorf("Got %v expected %v", index, 0)
+	}
+
+	list.Add("a", "b", "c")
+	it.End()
+	if index := it.Index(); index != list.Size() {
+		t.Errorf("Got %v expected %v", index, list.Size())
+	}
+
+	it.Prev()
+	if index, value := it.Index(), it.Value(); index != list.Size()-1 || value != "c" {
+		t.Errorf("Got %v,%v expected %v,%v", index, value, list.Size()-1, "c")
+	}
+}
+
 func TestListIteratorFirst(t *testing.T) {
 	list := New()
 	it := list.Iterator()

maps/treemap/treemap_test.go

@@ -377,7 +377,6 @@ func TestMapIteratorPrev(t *testing.T) {
 	}
 	countDown := m.Size()
 	for it.Prev() {
-		countDown--
 		key := it.Key()
 		value := it.Value()
 		switch key {
@@ -399,9 +398,9 @@ func TestMapIteratorPrev(t *testing.T) {
 		if actualValue, expectedValue := value, countDown; actualValue != expectedValue {
 			t.Errorf("Got %v expected %v", actualValue, expectedValue)
 		}
+		countDown--
 	}
-	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := countDown, 0; actualValue != expectedValue {
-	if actualValue, expectedValue := countDown, 1; actualValue != expectedValue {
 		t.Errorf("Got %v expected %v", actualValue, expectedValue)
 	}
 }
@@ -422,6 +421,19 @@ func TestMapIteratorBegin(t *testing.T) {
 	}
 }
 
+func TestMapTreeIteratorEnd(t *testing.T) {
+	m := NewWithIntComparator()
+	it := m.Iterator()
+	m.Put(3, "c")
+	m.Put(1, "a")
+	m.Put(2, "b")
+	it.End()
+	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 TestMapIteratorFirst(t *testing.T) {
 	m := NewWithIntComparator()
 	m.Put(3, "c")

sets/treeset/treeset_test.go

@@ -292,6 +292,31 @@ func TestSetIteratorBegin(t *testing.T) {
 	}
 }
 
+func TestSetIteratorEnd(t *testing.T) {
+	set := NewWithStringComparator()
+	it := set.Iterator()
+
+	if index := it.Index(); index != -1 {
+		t.Errorf("Got %v expected %v", index, -1)
+	}
+
+	it.End()
+	if index := it.Index(); index != 0 {
+		t.Errorf("Got %v expected %v", index, 0)
+	}
+
+	set.Add("a", "b", "c")
+	it.End()
+	if index := it.Index(); index != set.Size() {
+		t.Errorf("Got %v expected %v", index, set.Size())
+	}
+
+	it.Prev()
+	if index, value := it.Index(), it.Value(); index != set.Size()-1 || value != "c" {
+		t.Errorf("Got %v,%v expected %v,%v", index, value, set.Size()-1, "c")
+	}
+}
+
 func TestSetIteratorFirst(t *testing.T) {
 	set := NewWithStringComparator()
 	set.Add("a", "b", "c")

stacks/arraystack/arraystack_test.go

@@ -192,6 +192,33 @@ func TestStackIteratorBegin(t *testing.T) {
 	}
 }
 
+func TestStackIteratorEnd(t *testing.T) {
+	stack := New()
+	it := stack.Iterator()
+
+	if index := it.Index(); index != -1 {
+		t.Errorf("Got %v expected %v", index, -1)
+	}
+
+	it.End()
+	if index := it.Index(); index != 0 {
+		t.Errorf("Got %v expected %v", index, 0)
+	}
+
+	stack.Push("a")
+	stack.Push("b")
+	stack.Push("c")
+	it.End()
+	if index := it.Index(); index != stack.Size() {
+		t.Errorf("Got %v expected %v", index, stack.Size())
+	}
+
+	it.Prev()
+	if index, value := it.Index(), it.Value(); index != stack.Size()-1 || value != "a" {
+		t.Errorf("Got %v,%v expected %v,%v", index, value, stack.Size()-1, "a")
+	}
+}
+
 func TestStackIteratorFirst(t *testing.T) {
 	stack := New()
 	it := stack.Iterator()

trees/binaryheap/binaryheap_test.go

@@ -207,6 +207,33 @@ func TestBinaryHeapIteratorBegin(t *testing.T) {
 	}
 }
 
+func TestListIteratorEnd(t *testing.T) {
+	heap := NewWithIntComparator()
+	it := heap.Iterator()
+
+	if index := it.Index(); index != -1 {
+		t.Errorf("Got %v expected %v", index, -1)
+	}
+
+	it.End()
+	if index := it.Index(); index != 0 {
+		t.Errorf("Got %v expected %v", index, 0)
+	}
+
+	heap.Push(3) // [3]
+	heap.Push(2) // [2,3]
+	heap.Push(1) // [1,3,2](2 swapped with 1, hence last)
+	it.End()
+	if index := it.Index(); index != heap.Size() {
+		t.Errorf("Got %v expected %v", index, heap.Size())
+	}
+
+	it.Prev()
+	if index, value := it.Index(), it.Value(); index != heap.Size()-1 || value != 2 {
+		t.Errorf("Got %v,%v expected %v,%v", index, value, heap.Size()-1, 2)
+	}
+}
+
 func TestStackIteratorFirst(t *testing.T) {
 	heap := NewWithIntComparator()
 	it := heap.Iterator()

trees/redblacktree/redblacktree.go

@@ -280,13 +280,20 @@ func (tree *Tree) Clear() {
 
 // Iterator holding the iterator's state
 type Iterator struct {
-	tree *Tree
+	tree     *Tree
-	node *Node
+	node     *Node
+	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}
+	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.
@@ -294,55 +301,84 @@ func (tree *Tree) Iterator() Iterator {
 // 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 iterator.node == nil {
+	if iterator.position == end {
-		iterator.node = iterator.tree.Left()
+		goto end
-		return iterator.node != nil
+	}
+	if iterator.position == begin {
+		left := iterator.tree.Left()
+		if left == nil {
+			goto end
+		}
+		iterator.node = left
+		goto between
 	}
 	if iterator.node.Right != nil {
 		iterator.node = iterator.node.Right
 		for iterator.node.Left != nil {
 			iterator.node = iterator.node.Left
 		}
-		return true
+		goto between
 	}
 	if iterator.node.Parent != nil {
 		node := iterator.node
 		for iterator.node.Parent != nil {
 			iterator.node = iterator.node.Parent
 			if iterator.tree.Comparator(node.Key, iterator.node.Key) <= 0 {
-				return true
+				goto between
 			}
 		}
-		iterator.node = node // fix: if parent didn't satisfy the comparator criteria
 	}
+
+end:
+	iterator.node = nil
+	iterator.position = 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 iterator.node == nil {
+	if iterator.position == begin {
-		return false
+		goto begin
+	}
+	if iterator.position == end {
+		right := iterator.tree.Right()
+		if right == nil {
+			goto begin
+		}
+		iterator.node = right
+		goto between
 	}
 	if iterator.node.Left != nil {
 		iterator.node = iterator.node.Left
 		for iterator.node.Right != nil {
 			iterator.node = iterator.node.Right
 		}
-		return true
+		goto between
 	}
 	if iterator.node.Parent != nil {
 		node := iterator.node
 		for iterator.node.Parent != nil {
 			iterator.node = iterator.node.Parent
 			if iterator.tree.Comparator(node.Key, iterator.node.Key) >= 0 {
-				return true
+				goto between
 			}
 		}
-		iterator.node = node // fix: if parent didn't satisfy the comparator criteria
 	}
+
+begin:
+	iterator.node = nil
+	iterator.position = begin
 	return false
+
+between:
+	iterator.position = between
+	return true
 }
 
 // Value returns the current element's value.
@@ -361,13 +397,14 @@ func (iterator *Iterator) Key() interface{} {
 // Call Next() to fetch the first element if any.
 func (iterator *Iterator) Begin() {
 	iterator.node = nil
+	iterator.position = begin
 }
 
 // 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() {
-	right := iterator.tree.Right() // ugly hack (TODO: use flags)
+	iterator.node = nil
-	iterator.node = &Node{Parent: right, Key: right.Key}
+	iterator.position = end
 }
 
 // First moves the iterator to the first element and returns true if there was a first element in the container.

trees/redblacktree/redblacktree_test.go

@@ -282,7 +282,6 @@ func TestRedBlackTreeIterator1Prev(t *testing.T) {
 	}
 	countDown := tree.size
 	for it.Prev() {
-		countDown--
 		key := it.Key()
 		switch key {
 		case countDown:
@@ -294,9 +293,9 @@ func TestRedBlackTreeIterator1Prev(t *testing.T) {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		}
+		countDown--
 	}
-	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := countDown, 0; actualValue != expectedValue {
-	if actualValue, expectedValue := countDown, 1; actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }
@@ -337,7 +336,6 @@ func TestRedBlackTreeIterator2Prev(t *testing.T) {
 	}
 	countDown := tree.size
 	for it.Prev() {
-		countDown--
 		key := it.Key()
 		switch key {
 		case countDown:
@@ -349,9 +347,9 @@ func TestRedBlackTreeIterator2Prev(t *testing.T) {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		}
+		countDown--
 	}
-	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := countDown, 0; actualValue != expectedValue {
-	if actualValue, expectedValue := countDown, 1; actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }
@@ -388,7 +386,6 @@ func TestRedBlackTreeIterator3Prev(t *testing.T) {
 	}
 	countDown := tree.size
 	for it.Prev() {
-		countDown--
 		key := it.Key()
 		switch key {
 		case countDown:
@@ -400,9 +397,9 @@ func TestRedBlackTreeIterator3Prev(t *testing.T) {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		}
+		countDown--
 	}
-	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := countDown, 0; actualValue != expectedValue {
-	if actualValue, expectedValue := countDown, 1; actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }
@@ -450,7 +447,7 @@ func TestRedBlackTreeIterator4Next(t *testing.T) {
 	}
 }
 
-func TestRedBlackTreeIterator4(t *testing.T) {
+func TestRedBlackTreeIterator4Prev(t *testing.T) {
 	tree := NewWithIntComparator()
 	tree.Put(13, 5)
 	tree.Put(8, 3)
@@ -477,7 +474,6 @@ func TestRedBlackTreeIterator4(t *testing.T) {
 	for it.Next() {
 	}
 	for it.Prev() {
-		count--
 		value := it.Value()
 		switch value {
 		case count:
@@ -489,9 +485,9 @@ func TestRedBlackTreeIterator4(t *testing.T) {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		}
+		count--
 	}
-	// one less that in Next(), thus "1"
+	if actualValue, expectedValue := count, 0; actualValue != expectedValue {
-	if actualValue, expectedValue := count, 1; actualValue != expectedValue {
 		t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue)
 	}
 }
@@ -502,16 +498,59 @@ func TestRedBlackTreeIteratorBegin(t *testing.T) {
 	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 TestRedBlackTreeIteratorEnd(t *testing.T) {
+	tree := NewWithIntComparator()
+	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 TestRedBlackTreeIteratorFirst(t *testing.T) {
 	tree := NewWithIntComparator()
 	tree.Put(3, "c")