diff --git a/set.go b/set.go index e9409aa..068b38d 100644 --- a/set.go +++ b/set.go @@ -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 diff --git a/set_test.go b/set_test.go index d1eadf8..57b5be0 100644 --- a/set_test.go +++ b/set_test.go @@ -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]() diff --git a/threadsafe.go b/threadsafe.go index 0f3e593..1f23720 100644 --- a/threadsafe.go +++ b/threadsafe.go @@ -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) diff --git a/threadunsafe.go b/threadunsafe.go index c95d32b..b7f10fb 100644 --- a/threadunsafe.go +++ b/threadunsafe.go @@ -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) }