Skip to content

Commit

Permalink
add Each method
Browse files Browse the repository at this point in the history
  • Loading branch information
bluele authored and deckarep committed Oct 13, 2017
1 parent b3af78e commit 1d4478f
Show file tree
Hide file tree
Showing 6 changed files with 121 additions and 0 deletions.
38 changes: 38 additions & 0 deletions bench_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
4 changes: 4 additions & 0 deletions set.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down
31 changes: 31 additions & 0 deletions set_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
10 changes: 10 additions & 0 deletions threadsafe.go
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
30 changes: 30 additions & 0 deletions threadsafe_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import (
"math/rand"
"runtime"
"sync"
"sync/atomic"
"testing"
)

Expand Down Expand Up @@ -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)

Expand Down
8 changes: 8 additions & 0 deletions threadunsafe.go
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down

0 comments on commit 1d4478f

Please sign in to comment.