Skip to content

feat: add PopN method to remove multiple arbitrary items from set #166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions set.go
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,12 @@ type Set[T comparable] interface {
// Pop removes and returns an arbitrary item from the set.
Pop() (T, bool)

// PopN removes and returns up to n arbitrary items from the set.
// It returns a slice of the removed items and the actual number of items removed.
// If the set is empty or n is less than or equal to 0s, it returns an empty slice and 0.
// If n is greater than the set's size, all items are
PopN(n int) ([]T, int)

// ToSlice returns the members of the set as a slice.
ToSlice() []T

Expand Down
112 changes: 112 additions & 0 deletions set_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1242,6 +1242,118 @@ func Test_PopUnsafe(t *testing.T) {
}
}

func Test_PopNSafe(t *testing.T) {
a := NewSet[string]()
a.Add("a")
a.Add("b")
a.Add("c")
a.Add("d")

// Test pop with n <= 0
items, count := a.PopN(0)
if count != 0 {
t.Errorf("expected 0 items popped, got %d", count)
}
items, count = a.PopN(-1)
if count != 0 {
t.Errorf("expected 0 items popped, got %d", count)
}

captureSet := NewSet[string]()

// Test pop 2 items
items, count = a.PopN(2)
if count != 2 {
t.Errorf("expected 2 items popped, got %d", count)
}
if len(items) != 2 {
t.Errorf("expected 2 items in slice, got %d", len(items))
}
if a.Cardinality() != 2 {
t.Errorf("expected 2 items remaining, got %d", a.Cardinality())
}
captureSet.Append(items...)

// Test pop more than remaining
items, count = a.PopN(3)
if count != 2 {
t.Errorf("expected 2 items popped, got %d", count)
}
if a.Cardinality() != 0 {
t.Errorf("expected 0 items remaining, got %d", a.Cardinality())
}
captureSet.Append(items...)

// Test pop from empty set
items, count = a.PopN(1)
if count != 0 {
t.Errorf("expected 0 items popped, got %d", count)
}
if len(items) != 0 {
t.Errorf("expected empty slice, got %d items", len(items))
}

if !captureSet.Contains("c", "a", "d", "b") {
t.Error("unexpected result set; should be a,b,c,d (any order is fine")
}
}

func Test_PopNUnsafe(t *testing.T) {
a := NewThreadUnsafeSet[string]()
a.Add("a")
a.Add("b")
a.Add("c")
a.Add("d")

// Test pop with n <= 0
items, count := a.PopN(0)
if count != 0 {
t.Errorf("expected 0 items popped, got %d", count)
}
items, count = a.PopN(-1)
if count != 0 {
t.Errorf("expected 0 items popped, got %d", count)
}

captureSet := NewThreadUnsafeSet[string]()

// Test pop 2 items
items, count = a.PopN(2)
if count != 2 {
t.Errorf("expected 2 items popped, got %d", count)
}
if len(items) != 2 {
t.Errorf("expected 2 items in slice, got %d", len(items))
}
if a.Cardinality() != 2 {
t.Errorf("expected 2 items remaining, got %d", a.Cardinality())
}
captureSet.Append(items...)

// Test pop more than remaining
items, count = a.PopN(3)
if count != 2 {
t.Errorf("expected 2 items popped, got %d", count)
}
if a.Cardinality() != 0 {
t.Errorf("expected 0 items remaining, got %d", a.Cardinality())
}
captureSet.Append(items...)

// Test pop from empty set
items, count = a.PopN(1)
if count != 0 {
t.Errorf("expected 0 items popped, got %d", count)
}
if len(items) != 0 {
t.Errorf("expected empty slice, got %d items", len(items))
}

if !captureSet.Contains("c", "a", "d", "b") {
t.Error("unexpected result set; should be a,b,c,d (any order is fine")
}
}

func Test_EmptySetProperties(t *testing.T) {
empty := NewSet[string]()

Expand Down
6 changes: 6 additions & 0 deletions threadsafe.go
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,12 @@ func (t *threadSafeSet[T]) Pop() (T, bool) {
return t.uss.Pop()
}

func (t *threadSafeSet[T]) PopN(n int) ([]T, int) {
t.Lock()
defer t.Unlock()
return t.uss.PopN(n)
}

func (t *threadSafeSet[T]) ToSlice() []T {
t.RLock()
l := len(*t.uss)
Expand Down
21 changes: 21 additions & 0 deletions threadunsafe.go
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,27 @@ func (s *threadUnsafeSet[T]) Pop() (v T, ok bool) {
return v, false
}

func (s *threadUnsafeSet[T]) PopN(n int) (items []T, count int) {
if n <= 0 || len(*s) == 0 {
return make([]T, 0), 0
}
sn := s.Cardinality()
if n > sn {
n = sn
}

items = make([]T, 0, sn)
for item := range *s {
if count >= n {
break
}
delete(*s, item)
items = append(items, item)
count++
}
return items, count
}

func (s threadUnsafeSet[T]) Remove(v T) {
delete(s, v)
}
Expand Down