Skip to content

Commit

Permalink
feat: add ContainsOne method (#134)
Browse files Browse the repository at this point in the history
* feat: add ContainsOne method

* chore: removed unnecessary word
  • Loading branch information
vladvalkov authored Dec 26, 2023
1 parent 5ff637d commit 4a7d053
Show file tree
Hide file tree
Showing 6 changed files with 182 additions and 0 deletions.
93 changes: 93 additions & 0 deletions bench_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,99 @@ func BenchmarkContains100Unsafe(b *testing.B) {
benchContains(b, 100, NewThreadUnsafeSet[int]())
}

func benchContainsOne(b *testing.B, n int, s Set[int]) {
nums := nrand(n)
for _, v := range nums {
s.Add(v)
}

b.ResetTimer()
for i := 0; i < b.N; i++ {
s.ContainsOne(-1)
}
}

func BenchmarkContainsOne1Safe(b *testing.B) {
benchContainsOne(b, 1, NewSet[int]())
}

func BenchmarkContainsOne1Unsafe(b *testing.B) {
benchContainsOne(b, 1, NewThreadUnsafeSet[int]())
}

func BenchmarkContainsOne10Safe(b *testing.B) {
benchContainsOne(b, 10, NewSet[int]())
}

func BenchmarkContainsOne10Unsafe(b *testing.B) {
benchContainsOne(b, 10, NewThreadUnsafeSet[int]())
}

func BenchmarkContainsOne100Safe(b *testing.B) {
benchContainsOne(b, 100, NewSet[int]())
}

func BenchmarkContainsOne100Unsafe(b *testing.B) {
benchContainsOne(b, 100, NewThreadUnsafeSet[int]())
}

// In this scenario, Contains argument escapes to the heap, while ContainsOne does not.
func benchContainsComparison(b *testing.B, n int, s Set[int]) {
nums := nrand(n)
for _, v := range nums {
s.Add(v)
}

b.Run("Contains", func(b *testing.B) {
b.ReportAllocs()
for i := 0; i < b.N; i++ {
for _, v := range nums {
s.Contains(v) // 1 allocation, v is moved to the heap
}
}
})
b.Run("Contains slice", func(b *testing.B) {
b.ReportAllocs()
for i := 0; i < b.N; i++ {
for i := range nums {
s.Contains(nums[i : i+1]...) // no allocations, using heap-allocated slice
}
}
})
b.Run("ContainsOne", func(b *testing.B) {
b.ReportAllocs()
for i := 0; i < b.N; i++ {
for _, v := range nums {
s.ContainsOne(v) // no allocations, using stack-allocated v
}
}
})
}

func BenchmarkContainsComparison1Unsafe(b *testing.B) {
benchContainsComparison(b, 1, NewThreadUnsafeSet[int]())
}

func BenchmarkContainsComparison1Safe(b *testing.B) {
benchContainsComparison(b, 1, NewSet[int]())
}

func BenchmarkContainsComparison10Unsafe(b *testing.B) {
benchContainsComparison(b, 10, NewThreadUnsafeSet[int]())
}

func BenchmarkContainsComparison10Safe(b *testing.B) {
benchContainsComparison(b, 10, NewSet[int]())
}

func BenchmarkContainsComparison100Unsafe(b *testing.B) {
benchContainsComparison(b, 100, NewThreadUnsafeSet[int]())
}

func BenchmarkContainsComparison100Safe(b *testing.B) {
benchContainsComparison(b, 100, NewSet[int]())
}

func benchEqual(b *testing.B, n int, s, t Set[int]) {
nums := nrand(n)
for _, v := range nums {
Expand Down
7 changes: 7 additions & 0 deletions set.go
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,13 @@ type Set[T comparable] interface {
// are all in the set.
Contains(val ...T) bool

// ContainsOne returns whether the given item
// is in the set.
//
// Contains may cause the argument to escape to the heap.
// See: https://github.com/deckarep/golang-set/issues/118
ContainsOne(val T) bool

// ContainsAny returns whether at least one of the
// given items are in the set.
ContainsAny(val ...T) bool
Expand Down
48 changes: 48 additions & 0 deletions set_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,54 @@ func Test_ContainsMultipleUnsafeSet(t *testing.T) {
}
}

func Test_ContainsOneSet(t *testing.T) {
a := NewSet[int]()

a.Add(71)

if !a.ContainsOne(71) {
t.Error("ContainsSet should contain 71")
}

a.Remove(71)

if a.ContainsOne(71) {
t.Error("ContainsSet should not contain 71")
}

a.Add(13)
a.Add(7)
a.Add(1)

if !(a.ContainsOne(13) && a.ContainsOne(7) && a.ContainsOne(1)) {
t.Error("ContainsSet should contain 13, 7, 1")
}
}

func Test_ContainsOneUnsafeSet(t *testing.T) {
a := NewThreadUnsafeSet[int]()

a.Add(71)

if !a.ContainsOne(71) {
t.Error("ContainsSet should contain 71")
}

a.Remove(71)

if a.ContainsOne(71) {
t.Error("ContainsSet should not contain 71")
}

a.Add(13)
a.Add(7)
a.Add(1)

if !(a.ContainsOne(13) && a.ContainsOne(7) && a.ContainsOne(1)) {
t.Error("ContainsSet should contain 13, 7, 1")
}
}

func Test_ContainsAnySet(t *testing.T) {
a := NewSet[int]()

Expand Down
8 changes: 8 additions & 0 deletions threadsafe.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,14 @@ func (t *threadSafeSet[T]) Contains(v ...T) bool {
return ret
}

func (t *threadSafeSet[T]) ContainsOne(v T) bool {
t.RLock()
ret := t.uss.ContainsOne(v)
t.RUnlock()

return ret
}

func (t *threadSafeSet[T]) ContainsAny(v ...T) bool {
t.RLock()
ret := t.uss.ContainsAny(v...)
Expand Down
21 changes: 21 additions & 0 deletions threadsafe_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,27 @@ func Test_ContainsConcurrent(t *testing.T) {
wg.Wait()
}

func Test_ContainsOneConcurrent(t *testing.T) {
runtime.GOMAXPROCS(2)

s := NewSet[int]()
ints := rand.Perm(N)
for _, v := range ints {
s.Add(v)
}

var wg sync.WaitGroup
for _, v := range ints {
number := v
wg.Add(1)
go func() {
s.ContainsOne(number)
wg.Done()
}()
}
wg.Wait()
}

func Test_ContainsAnyConcurrent(t *testing.T) {
runtime.GOMAXPROCS(2)

Expand Down
5 changes: 5 additions & 0 deletions threadunsafe.go
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,11 @@ func (s threadUnsafeSet[T]) Contains(v ...T) bool {
return true
}

func (s threadUnsafeSet[T]) ContainsOne(v T) bool {
_, ok := s[v]
return ok
}

func (s threadUnsafeSet[T]) ContainsAny(v ...T) bool {
for _, val := range v {
if _, ok := s[val]; ok {
Expand Down

0 comments on commit 4a7d053

Please sign in to comment.