add Each method

Author: bluele

bench_test.go

@@ -501,6 +501,44 @@ func BenchmarkUnion100Unsafe(b *testing.B) {
 	benchUnion(b, 100, NewThreadUnsafeSet(), NewThreadUnsafeSet())
 }
 
+func benchEach(b *testing.B, n int, s Set) {
+	nums := nrand(n)
+	for _, v := range nums {
+		s.Add(v)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		s.Each(func(elem interface{}) bool {
+			return false
+		})
+	}
+}
+
+func BenchmarkEach1Safe(b *testing.B) {
+	benchEach(b, 1, NewSet())
+}
+
+func BenchmarkEach1Unsafe(b *testing.B) {
+	benchEach(b, 1, NewThreadUnsafeSet())
+}
+
+func BenchmarkEach10Safe(b *testing.B) {
+	benchEach(b, 10, NewSet())
+}
+
+func BenchmarkEach10Unsafe(b *testing.B) {
+	benchEach(b, 10, NewThreadUnsafeSet())
+}
+
+func BenchmarkEach100Safe(b *testing.B) {
+	benchEach(b, 100, NewSet())
+}
+
+func BenchmarkEach100Unsafe(b *testing.B) {
+	benchEach(b, 100, NewThreadUnsafeSet())
+}
+
 func benchIter(b *testing.B, n int, s Set) {
 	nums := nrand(n)
 	for _, v := range nums {

set.go

@@ -126,6 +126,10 @@ type Set interface {
 	// panic.
 	IsSuperset(other Set) bool
 
+	// Iterates over elements and executes the passed func against each element.
+	// If passed func returns true, stop iteration at the time.
+	Each(func(interface{}) bool)
+
 	// Returns a channel of elements that you can
 	// range over.
 	Iter() <-chan interface{}

set_test.go

@@ -849,6 +849,37 @@ func Test_UnsafeSetClone(t *testing.T) {
 	}
 }
 
+func Test_Each(t *testing.T) {
+	a := NewSet()
+
+	a.Add("Z")
+	a.Add("Y")
+	a.Add("X")
+	a.Add("W")
+
+	b := NewSet()
+	a.Each(func(elem interface{}) bool {
+		b.Add(elem)
+		return false
+	})
+
+	if !a.Equal(b) {
+		t.Error("The sets are not equal after iterating (Each) through the first set")
+	}
+
+	var count int
+	a.Each(func(elem interface{}) bool {
+		if count == 2 {
+			return true
+		}
+		count++
+		return false
+	})
+	if count != 2 {
+		t.Error("Iteration should stop on the way")
+	}
+}
+
 func Test_Iter(t *testing.T) {
 	a := NewSet()
 

threadsafe.go

@@ -151,6 +151,16 @@ func (set *threadSafeSet) Cardinality() int {
 	return len(set.s)
 }
 
+func (set *threadSafeSet) Each(cb func(interface{}) bool) {
+	set.RLock()
+	for elem := range set.s {
+		if cb(elem) {
+			break
+		}
+	}
+	set.RUnlock()
+}
+
 func (set *threadSafeSet) Iter() <-chan interface{} {
 	ch := make(chan interface{})
 	go func() {

threadsafe_test.go

@@ -30,6 +30,7 @@ import (
 	"math/rand"
 	"runtime"
 	"sync"
+	"sync/atomic"
 	"testing"
 )
 
@@ -294,6 +295,35 @@ func Test_IsProperSupersetConcurrent(t *testing.T) {
 	wg.Wait()
 }
 
+func Test_EachConcurrent(t *testing.T) {
+	runtime.GOMAXPROCS(2)
+	concurrent := 10
+
+	s := NewSet()
+	ints := rand.Perm(N)
+	for _, v := range ints {
+		s.Add(v)
+	}
+
+	var count int64
+	wg := new(sync.WaitGroup)
+	wg.Add(concurrent)
+	for n := 0; n < concurrent; n++ {
+		go func() {
+			defer wg.Done()
+			s.Each(func(elem interface{}) bool {
+				atomic.AddInt64(&count, 1)
+				return false
+			})
+		}()
+	}
+	wg.Wait()
+
+	if count != int64(N*concurrent) {
+		t.Errorf("%v != %v", count, int64(N*concurrent))
+	}
+}
+
 func Test_IterConcurrent(t *testing.T) {
 	runtime.GOMAXPROCS(2)
 

threadunsafe.go

@@ -159,6 +159,14 @@ func (set *threadUnsafeSet) Cardinality() int {
 	return len(*set)
 }
 
+func (set *threadUnsafeSet) Each(cb func(interface{}) bool) {
+	for elem := range *set {
+		if cb(elem) {
+			break
+		}
+	}
+}
+
 func (set *threadUnsafeSet) Iter() <-chan interface{} {
 	ch := make(chan interface{})
 	go func() {