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

---
 maps/treemap/treemap.go      | 22 ++++++++++++++++++++++
 maps/treemap/treemap_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/maps/treemap/treemap.go b/maps/treemap/treemap.go
index d0fc742..20e4ba2 100644
--- a/maps/treemap/treemap.go
+++ b/maps/treemap/treemap.go
@@ -32,6 +32,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 package treemap
 
 import (
+	"github.com/emirpasic/gods/containers"
 	"github.com/emirpasic/gods/maps"
 	rbt "github.com/emirpasic/gods/trees/redblacktree"
 	"github.com/emirpasic/gods/utils"
@@ -39,6 +40,7 @@ import (
 
 func assertInterfaceImplementation() {
 	var _ maps.Map = (*Map)(nil)
+	var _ containers.IteratorWithKey = (*Iterator)(nil)
 }
 
 type Map struct {
@@ -122,6 +124,26 @@ func (m *Map) Max() (key interface{}, value interface{}) {
 	return nil, nil
 }
 
+type Iterator struct {
+	iterator rbt.Iterator
+}
+
+func (m *Map) Iterator() Iterator {
+	return Iterator{iterator: m.tree.Iterator()}
+}
+
+func (iterator *Iterator) Next() bool {
+	return iterator.iterator.Next()
+}
+
+func (iterator *Iterator) Value() interface{} {
+	return iterator.iterator.Value()
+}
+
+func (iterator *Iterator) Key() interface{} {
+	return iterator.iterator.Key()
+}
+
 func (m *Map) String() string {
 	str := "TreeMap\n"
 	str += m.tree.String()
diff --git a/maps/treemap/treemap_test.go b/maps/treemap/treemap_test.go
index c48c454..95ee459 100644
--- a/maps/treemap/treemap_test.go
+++ b/maps/treemap/treemap_test.go
@@ -179,6 +179,34 @@ func TestTreeMap(t *testing.T) {
 	}
 }
 
+func TestTreeMapIterator(t *testing.T) {
+	m := NewWithStringComparator()
+	m.Put("c", 3)
+	m.Put("a", 1)
+	m.Put("b", 2)
+
+	// Iterator
+	it := m.Iterator()
+	count := 0
+	for it.Next() {
+		count += 1
+		value := it.Value()
+		switch value {
+		case count:
+			if actualValue, expectedValue := value, count; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			t.Errorf("Too many")
+		}
+	}
+	m.Clear()
+	it = m.Iterator()
+	for it.Next() {
+		t.Errorf("Shouldn't iterate on empty map")
+	}
+}
+
 func BenchmarkTreeMap(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		m := NewWithIntComparator()