From 0418a59aaf6584daa5fb27b0894789a4a7756f5a Mon Sep 17 00:00:00 2001
From: Emir Pasic <emirpasic@hotmail.com>
Date: Wed, 22 Jun 2016 20:18:37 +0200
Subject: [PATCH] - add iterator to tree set

---
 sets/treeset/treeset.go      | 24 ++++++++++++++++++++++++
 sets/treeset/treeset_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/sets/treeset/treeset.go b/sets/treeset/treeset.go
index da6d7d9..42d0eee 100644
--- a/sets/treeset/treeset.go
+++ b/sets/treeset/treeset.go
@@ -24,6 +24,7 @@ package treeset
 
 import (
 	"fmt"
+	"github.com/emirpasic/gods/containers"
 	"github.com/emirpasic/gods/sets"
 	rbt "github.com/emirpasic/gods/trees/redblacktree"
 	"github.com/emirpasic/gods/utils"
@@ -32,6 +33,7 @@ import (
 
 func assertInterfaceImplementation() {
 	var _ sets.Set = (*Set)(nil)
+	var _ containers.IteratorWithIndex = (*Iterator)(nil)
 }
 
 type Set struct {
@@ -101,6 +103,28 @@ func (set *Set) Values() []interface{} {
 	return set.tree.Keys()
 }
 
+type Iterator struct {
+	index    int
+	iterator rbt.Iterator
+}
+
+func (set *Set) Iterator() Iterator {
+	return Iterator{index: -1, iterator: set.tree.Iterator()}
+}
+
+func (iterator *Iterator) Next() bool {
+	iterator.index += 1
+	return iterator.iterator.Next()
+}
+
+func (iterator *Iterator) Value() interface{} {
+	return iterator.iterator.Key()
+}
+
+func (iterator *Iterator) Index() int {
+	return iterator.index
+}
+
 func (set *Set) String() string {
 	str := "TreeSet\n"
 	items := []string{}
diff --git a/sets/treeset/treeset_test.go b/sets/treeset/treeset_test.go
index b033a7b..780a9db 100644
--- a/sets/treeset/treeset_test.go
+++ b/sets/treeset/treeset_test.go
@@ -82,7 +82,41 @@ func TestTreeSet(t *testing.T) {
 	if actualValue := set.Empty(); actualValue != true {
 		t.Errorf("Got %v expected %v", actualValue, true)
 	}
+}
 
+func TestTreeSetIterator(t *testing.T) {
+	set := NewWithStringComparator()
+	set.Add("c")
+	set.Add("a")
+	set.Add("b")
+
+	// Iterator
+	it := set.Iterator()
+	for it.Next() {
+		index := it.Index()
+		value := it.Value()
+		switch index {
+		case 0:
+			if actualValue, expectedValue := value, "a"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case 1:
+			if actualValue, expectedValue := value, "b"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case 2:
+			if actualValue, expectedValue := value, "c"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			t.Errorf("Too many")
+		}
+	}
+	set.Clear()
+	it = set.Iterator()
+	for it.Next() {
+		t.Errorf("Shouldn't iterate on empty set")
+	}
 }
 
 func BenchmarkTreeSet(b *testing.B) {