Repository: deckarep/golang-set

Issue:
Suggestion: Add PopN(n int) ([]T, int)
When using threadSafeSet for batch Pop, the overhead caused by RWMutex can be reduced.

```golang
func (s *threadUnsafeSet[T])PopN(n int)([]T, int) {
  suc := 0
  ss := make([]T, 0, n)
  for item := range *s {
    ss = append(ss, item)
    delete(*s, item)
    suc += 1
    if suc >= n {
      return  ss, suc
    }
  }
  return ss, suc
}
```

Discussion:
@deckarep: Thank you @lazychanger,

Looks like a good idea to me. Is this something you’d like to contribute a patch for?@lazychanger: Thanks! Yes, I'd be happy to contribute a patch for this.

Solution patch:
diff --git a/go.mod b/go.mod
index ae7f41e..0a507bf 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module github.com/deckarep/golang-set/v2
 
 go 1.18
+
+require go.mongodb.org/mongo-driver v1.17.4
diff --git a/go.sum b/go.sum
index e69de29..a669c19 100644
--- a/go.sum
+++ b/go.sum
@@ -0,0 +1,6 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+go.mongodb.org/mongo-driver v1.16.0 h1:tpRsfBJMROVHKpdGyc1BBEzzjDUWjItxbVSZ8Ls4BQ4=
+go.mongodb.org/mongo-driver v1.16.0/go.mod h1:oB6AhJQvFQL4LEHyXi6aJzQJtBiTQHiAd83l0GdFaiw=
+go.mongodb.org/mongo-driver v1.17.4 h1:jUorfmVzljjr0FLzYQsGP8cgN/qzzxlY9Vh0C9KFXVw=
+go.mongodb.org/mongo-driver v1.17.4/go.mod h1:Hy04i7O2kC4RS06ZrhPRqj/u4DTYkFDAAccj+rVKqgQ=
diff --git a/set.go b/set.go
index e9409aa..92b67bf 100644
--- a/set.go
+++ b/set.go
@@ -35,6 +35,8 @@ SOFTWARE.
 // that can enforce mutual exclusion through other means.
 package mapset
 
+import "go.mongodb.org/mongo-driver/bson/bsontype"
+
 // Set is the primary interface provided by the mapset package.  It
 // represents an unordered set of data and a large number of
 // operations that can be applied to that set.
@@ -189,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
 
@@ -196,8 +204,15 @@ type Set[T comparable] interface {
 	MarshalJSON() ([]byte, error)
 
 	// UnmarshalJSON will unmarshal a JSON-based byte slice into a full Set datastructure.
-	// For this to work, set subtypes must implemented the Marshal/Unmarshal interface.
+	// For this to work, set subtypes must implement the Marshal/Unmarshal interface.
 	UnmarshalJSON(b []byte) error
+
+	// MarshalBSONValue will marshal the set into a BSON-based representation.
+	MarshalBSONValue() (bsontype.Type, []byte, error)
+
+	// UnmarshalBSONValue will unmarshal a BSON-based byte slice into a full Set datastructure.
+	// For this to work, set subtypes must implement the Marshal/Unmarshal interface.
+	UnmarshalBSONValue(bt bsontype.Type, b []byte) error
 }
 
 // NewSet creates and returns a new set with the given elements.
diff --git a/threadsafe.go b/threadsafe.go
index 0f3e593..835797e 100644
--- a/threadsafe.go
+++ b/threadsafe.go
@@ -25,7 +25,11 @@ SOFTWARE.
 
 package mapset
 
-import "sync"
+import (
+	"sync"
+
+	"go.mongodb.org/mongo-driver/bson/bsontype"
+)
 
 type threadSafeSet[T comparable] struct {
 	sync.RWMutex
@@ -285,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)
@@ -305,9 +315,25 @@ func (t *threadSafeSet[T]) MarshalJSON() ([]byte, error) {
 }
 
 func (t *threadSafeSet[T]) UnmarshalJSON(p []byte) error {
-	t.RLock()
+	t.Lock()
 	err := t.uss.UnmarshalJSON(p)
+	t.Unlock()
+
+	return err
+}
+
+func (t *threadSafeSet[T]) MarshalBSONValue() (bsontype.Type, []byte, error) {
+	t.RLock()
+	bt, b, err := t.uss.MarshalBSONValue()
 	t.RUnlock()
 
+	return bt, b, err
+}
+
+func (t *threadSafeSet[T]) UnmarshalBSONValue(bt bsontype.Type, p []byte) error {
+	t.Lock()
+	err := t.uss.UnmarshalBSONValue(bt, p)
+	t.Unlock()
+
 	return err
 }
diff --git a/threadunsafe.go b/threadunsafe.go
index c95d32b..2ed0e06 100644
--- a/threadunsafe.go
+++ b/threadunsafe.go
@@ -29,6 +29,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/bsontype"
 )
 
 type threadUnsafeSet[T comparable] map[T]struct{}
@@ -259,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)
 }
@@ -350,3 +374,24 @@ func (s *threadUnsafeSet[T]) UnmarshalJSON(b []byte) error {
 
 	return nil
 }
+
+// MarshalBSON creates a BSON array from the set.
+func (s threadUnsafeSet[T]) MarshalBSONValue() (bsontype.Type, []byte, error) {
+	return bson.MarshalValue(s.ToSlice())
+}
+
+// UnmarshalBSON recreates a set from a BSON array.
+func (s threadUnsafeSet[T]) UnmarshalBSONValue(bt bsontype.Type, b []byte) error {
+	if bt != bson.TypeArray {
+		return fmt.Errorf("must use BSON Array to unmarshal Set")
+	}
+
+	var i []T
+	err := bson.UnmarshalValue(bt, b, &i)
+	if err != nil {
+		return err
+	}
+	s.Append(i...)
+
+	return nil
+}



Test patch:
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_test.go b/threadsafe_test.go
index ed15d02..7fbeba5 100644
--- a/threadsafe_test.go
+++ b/threadsafe_test.go
@@ -33,6 +33,8 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+
+	"go.mongodb.org/mongo-driver/bson"
 )
 
 const N = 1000
@@ -683,3 +685,81 @@ func Test_DeadlockOnEachCallbackWhenPanic(t *testing.T) {
 		t.Errorf("Expected widgets to have 5 elements, but has %d", card)
 	}
 }
+
+func Test_UnmarshalBSONValue(t *testing.T) {
+	tp, s, initErr := bson.MarshalValue(
+		bson.A{"1", "2", "3", "test"},
+	)
+
+	if initErr != nil {
+		t.Errorf("Init Error should be nil: %v", initErr)
+
+		return
+	}
+
+	if tp != bson.TypeArray {
+		t.Errorf("Encoded Type should be bson.Array, got: %v", tp)
+
+		return
+	}
+
+	expected := NewSet("1", "2", "3", "test")
+	actual := NewSet[string]()
+	err := bson.UnmarshalValue(bson.TypeArray, s, actual)
+	if err != nil {
+		t.Errorf("Error should be nil: %v", err)
+	}
+
+	if !expected.Equal(actual) {
+		t.Errorf("Expected no difference, got: %v", expected.Difference(actual))
+	}
+}
+
+func TestThreadUnsafeSet_UnmarshalBSONValue(t *testing.T) {
+	tp, s, initErr := bson.MarshalValue(
+		bson.A{int64(1), int64(2), int64(3)},
+	)
+
+	if initErr != nil {
+		t.Errorf("Init Error should be nil: %v", initErr)
+
+		return
+	}
+
+	if tp != bson.TypeArray {
+		t.Errorf("Encoded Type should be bson.Array, got: %v", tp)
+
+		return
+	}
+
+	expected := NewThreadUnsafeSet[int64](1, 2, 3)
+	actual := NewThreadUnsafeSet[int64]()
+	err := actual.UnmarshalBSONValue(bson.TypeArray, []byte(s))
+	if err != nil {
+		t.Errorf("Error should be nil: %v", err)
+	}
+	if !expected.Equal(actual) {
+		t.Errorf("Expected no difference, got: %v", expected.Difference(actual))
+	}
+}
+
+func Test_MarshalBSONValue(t *testing.T) {
+	expected := NewSet("1", "test")
+
+	_, b, err := bson.MarshalValue(
+		NewSet("1", "test"),
+	)
+	if err != nil {
+		t.Errorf("Error should be nil: %v", err)
+	}
+
+	actual := NewSet[string]()
+	err = bson.UnmarshalValue(bson.TypeArray, b, actual)
+	if err != nil {
+		t.Errorf("Error should be nil: %v", err)
+	}
+
+	if !expected.Equal(actual) {
+		t.Errorf("Expected no difference, got: %v", expected.Difference(actual))
+	}
+}