- add iterator to tree map

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()

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()