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

Author: lazychanger

set.go

@@ -191,6 +191,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
 

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]()
 

threadsafe.go

@@ -289,6 +289,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)

threadunsafe.go

@@ -262,6 +262,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)
 }