diff --git a/set/hashset.go b/set/hashset.go index 0201119..2d007ef 100644 --- a/set/hashset.go +++ b/set/hashset.go @@ -53,6 +53,15 @@ func (s *HashSet) Foreach(f func(interface{})) { } } +// Call f for each item in the set, set result as new key +func (s *HashSet) Map(f func(interface{}) interface{}) Set { + n := make(map[interface{}]nothing) + for k := range s.hash { + n[f(k)] = nothing{} + } + return &HashSet{n} +} + // Returns true if this set contains no elements. func (s *HashSet) IsEmpty() bool { return len(s.hash) == 0 diff --git a/set/hashset_test.go b/set/hashset_test.go index 5557f2d..02d72dc 100644 --- a/set/hashset_test.go +++ b/set/hashset_test.go @@ -70,6 +70,22 @@ func TestSet_Foreach(t *testing.T) { s.Foreach(f) } +func TestHashSet_Map(t *testing.T) { + oldSet := NewHashSet(1, 2, 3) + f := func(x int) int { return x * x } + fWrapper := func(x interface{}) interface{} { + xInt := x.(int) + return f(xInt) + } + s := oldSet.Map(fWrapper) + if s.Len() != oldSet.Len() { + t.Errorf("Length should be %d", oldSet.Len()) + } + if !s.ContainsAll(1, 4, 9) { + t.Error("Set should be 1, 4, 9") + } +} + func TestSet_Len(t *testing.T) { s := NewHashSet() if s.Len() != 0 { diff --git a/set/set.go b/set/set.go index 9f9649b..5f3247e 100755 --- a/set/set.go +++ b/set/set.go @@ -19,6 +19,9 @@ type Set interface { // Call f for each item in the set Foreach(f func(interface{})) + // Map f for each item + Map(f func(interface{}) interface{}) Set + // Returns true if this set contains no elements. IsEmpty() bool