- add enumerable to tree map

containers/enumerable.go

@@ -31,7 +31,7 @@ package containers
 
 // Enumerable function for ordered containers whose values can be fetched by an index.
 type EnumerableWithIndex interface {
-	// Calls the given function once for each element, passing that element's index(key) and value.
+	// Calls the given function once for each element, passing that element's index and value.
 	Each(func(index int, value interface{}))
 
 	// Invokes the given function once for each element and returns a
@@ -50,33 +50,33 @@ type EnumerableWithIndex interface {
 	All(func(index int, value interface{}) bool) bool
 
 	// Passes each element of the collection to the given function and returns
-	// the first for which the function is true or nil,nil otherwise if no element
+	// the first for which the function is true or -1,nil otherwise if no element
 	// matches the criteria.
 	Find(func(index int, value interface{}) bool) (index int, value interface{})
 }
 
 // Enumerable function for ordered containers whose values whose elements are key value pairs.
 type EnumerableWithKey interface {
-	// Calls the given function once for each element, passing that element's index(key) and value.
+	// Calls the given function once for each element, passing that element's key and value.
-	Each(func(index interface{}, value interface{}))
+	Each(func(key interface{}, value interface{}))
 
 	// Invokes the given function once for each element and returns a
 	// container containing the values returned by the given function.
-	Map(func(index interface{}, value interface{}) interface{}) Container
+	Map(func(key1 interface{}, value1 interface{}) (key2 interface{}, value2 interface{})) Container
 
 	// Returns a new container containing all elements for which the given function returns a true value.
-	Select(func(index interface{}, value interface{}) bool) Container
+	Select(func(key interface{}, value interface{}) bool) Container
 
 	// Passes each element of the collection to the given function and
 	// returns true if the function ever returns true for any element.
-	Any(func(index interface{}, value interface{}) bool) bool
+	Any(func(key interface{}, value interface{}) bool) bool
 
 	// Passes each element of the collection to the given function and
 	// returns true if the function returns true for all elements.
-	All(func(index interface{}, value interface{}) bool) bool
+	All(func(key interface{}, value interface{}) bool) bool
 
 	// Passes each element of the collection to the given function and returns
 	// the first for which the function is true or nil,nil otherwise if no element
 	// matches the criteria.
-	Find(func(index interface{}, value interface{}) bool) (index interface{}, value interface{})
+	Find(func(key interface{}, value interface{}) bool) (key interface{}, value interface{})
 }

maps/treemap/treemap.go

@@ -40,6 +40,7 @@ import (
 
 func assertInterfaceImplementation() {
 	var _ maps.Map = (*Map)(nil)
+	var _ containers.EnumerableWithKey = (*Map)(nil)
 	var _ containers.IteratorWithKey = (*Iterator)(nil)
 }
 
@@ -144,6 +145,64 @@ func (iterator *Iterator) Key() interface{} {
 	return iterator.iterator.Key()
 }
 
+func (m *Map) Each(f func(key interface{}, value interface{})) {
+	iterator := m.Iterator()
+	for iterator.Next() {
+		f(iterator.Key(), iterator.Value())
+	}
+}
+
+func (m *Map) Map(f func(key1 interface{}, value1 interface{}) (key2 interface{}, value2 interface{})) containers.Container {
+	newMap := &Map{tree: rbt.NewWith(m.tree.Comparator)}
+	iterator := m.Iterator()
+	for iterator.Next() {
+		key2, value2 := f(iterator.Key(), iterator.Value())
+		newMap.Put(key2, value2)
+	}
+	return newMap
+}
+
+func (m *Map) Select(f func(key interface{}, value interface{}) bool) containers.Container {
+	newMap := &Map{tree: rbt.NewWith(m.tree.Comparator)}
+	iterator := m.Iterator()
+	for iterator.Next() {
+		if f(iterator.Key(), iterator.Value()) {
+			newMap.Put(iterator.Key(), iterator.Value())
+		}
+	}
+	return newMap
+}
+
+func (m *Map) Any(f func(key interface{}, value interface{}) bool) bool {
+	iterator := m.Iterator()
+	for iterator.Next() {
+		if f(iterator.Key(), iterator.Value()) {
+			return true
+		}
+	}
+	return false
+}
+
+func (m *Map) All(f func(key interface{}, value interface{}) bool) bool {
+	iterator := m.Iterator()
+	for iterator.Next() {
+		if !f(iterator.Key(), iterator.Value()) {
+			return false
+		}
+	}
+	return true
+}
+
+func (m *Map) Find(f func(key interface{}, value interface{}) bool) (key interface{}, value interface{}) {
+	iterator := m.Iterator()
+	for iterator.Next() {
+		if f(iterator.Key(), iterator.Value()) {
+			return iterator.Key(), iterator.Value()
+		}
+	}
+	return nil, nil
+}
+
 func (m *Map) String() string {
 	str := "TreeMap\n"
 	str += m.tree.String()

maps/treemap/treemap_test.go

@@ -179,21 +179,126 @@ func TestTreeMap(t *testing.T) {
 	}
 }
 
-func TestTreeMapIterator(t *testing.T) {
+func TestTreeMapEnumerableAndIterator(t *testing.T) {
 	m := NewWithStringComparator()
 	m.Put("c", 3)
 	m.Put("a", 1)
 	m.Put("b", 2)
 
+	// Each
+	count := 0
+	m.Each(func(key interface{}, value interface{}) {
+		count += 1
+		if actualValue, expectedValue := count, value; actualValue != expectedValue {
+			t.Errorf("Got %v expected %v", actualValue, expectedValue)
+		}
+		switch value {
+		case 1:
+			if actualValue, expectedValue := key, "a"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case 2:
+			if actualValue, expectedValue := key, "b"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case 3:
+			if actualValue, expectedValue := key, "c"; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		default:
+			t.Errorf("Too many")
+		}
+	})
+
+	// Map
+	mappedMap := m.Map(func(key1 interface{}, value1 interface{}) (key2 interface{}, value2 interface{}) {
+		return key1, "mapped: " + key1.(string)
+	}).(*Map)
+	if actualValue, _ := mappedMap.Get("a"); actualValue != "mapped: a" {
+		t.Errorf("Got %v expected %v", actualValue, "mapped: a")
+	}
+	if actualValue, _ := mappedMap.Get("b"); actualValue != "mapped: b" {
+		t.Errorf("Got %v expected %v", actualValue, "mapped: b")
+	}
+	if actualValue, _ := mappedMap.Get("c"); actualValue != "mapped: c" {
+		t.Errorf("Got %v expected %v", actualValue, "mapped: c")
+	}
+	if mappedMap.Size() != 3 {
+		t.Errorf("Got %v expected %v", mappedMap.Size(), 3)
+	}
+
+	// Select
+	selectedMap := m.Select(func(key interface{}, value interface{}) bool {
+		return key.(string) >= "a" && key.(string) <= "b"
+	}).(*Map)
+	if actualValue, _ := selectedMap.Get("a"); actualValue != 1 {
+		t.Errorf("Got %v expected %v", actualValue, "value: a")
+	}
+	if actualValue, _ := selectedMap.Get("b"); actualValue != 2 {
+		t.Errorf("Got %v expected %v", actualValue, "value: b")
+	}
+	if selectedMap.Size() != 2 {
+		t.Errorf("Got %v expected %v", selectedMap.Size(), 3)
+	}
+
+	// Any
+	any := m.Any(func(key interface{}, value interface{}) bool {
+		return value.(int) == 3
+	})
+	if any != true {
+		t.Errorf("Got %v expected %v", any, true)
+	}
+	any = m.Any(func(key interface{}, value interface{}) bool {
+		return value.(int) == 4
+	})
+	if any != false {
+		t.Errorf("Got %v expected %v", any, false)
+	}
+
+	// All
+	all := m.All(func(key interface{}, value interface{}) bool {
+		return key.(string) >= "a" && key.(string) <= "c"
+	})
+	if all != true {
+		t.Errorf("Got %v expected %v", all, true)
+	}
+	all = m.All(func(key interface{}, value interface{}) bool {
+		return key.(string) >= "a" && key.(string) <= "b"
+	})
+	if all != false {
+		t.Errorf("Got %v expected %v", all, false)
+	}
+
+	// Find
+	foundKey, foundValue := m.Find(func(key interface{}, value interface{}) bool {
+		return key.(string) == "c"
+	})
+	if foundKey != "c" || foundValue != 3 {
+		t.Errorf("Got %v -> %v expected %v -> %v", foundKey, foundValue, "c", 3)
+	}
+	foundKey, foundValue = m.Find(func(key interface{}, value interface{}) bool {
+		return key.(string) == "x"
+	})
+	if foundKey != nil || foundValue != nil {
+		t.Errorf("Got %v at %v expected %v at %v", foundValue, foundKey, nil, nil)
+	}
+
 	// Iterator
 	it := m.Iterator()
-	count := 0
 	for it.Next() {
-		count += 1
+		key := it.Key()
 		value := it.Value()
-		switch value {
+		switch key {
-		case count:
+		case "a":
-			if actualValue, expectedValue := value, count; actualValue != expectedValue {
+			if actualValue, expectedValue := value, 1; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case "b":
+			if actualValue, expectedValue := value, 2; actualValue != expectedValue {
+				t.Errorf("Got %v expected %v", actualValue, expectedValue)
+			}
+		case "c":
+			if actualValue, expectedValue := value, 3; actualValue != expectedValue {
 				t.Errorf("Got %v expected %v", actualValue, expectedValue)
 			}
 		default: