Implement BSON Marshaler support

Author: Ryan Clarke

go.mod

@@ -1,3 +1,5 @@
 module github.com/deckarep/golang-set/v2
 
 go 1.18
+
+require go.mongodb.org/mongo-driver v1.16.0

go.sum

@@ -0,0 +1,4 @@
+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=

set.go

@@ -192,8 +192,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
+
+	// MarshalBSON will marshal the set into a BSON-based representation.
+	MarshalBSON() ([]byte, error)
+
+	// UnmarshalBSON will unmarshal a BSON-based byte slice into a full Set datastructure.
+	// For this to work, set subtypes must implement the Marshal/Unmarshal interface.
+	UnmarshalBSON(b []byte) error
 }
 
 // NewSet creates and returns a new set with the given elements.

threadsafe.go

@@ -291,9 +291,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]) MarshalBSON() ([]byte, error) {
+	t.RLock()
+	b, err := t.uss.MarshalBSON()
 	t.RUnlock()
 
+	return b, err
+}
+
+func (t *threadSafeSet[T]) UnmarshalBSON(p []byte) error {
+	t.Lock()
+	err := t.uss.UnmarshalBSON(p)
+	t.Unlock()
+
 	return err
 }

threadsafe_test.go

@@ -32,6 +32,8 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+
+	"go.mongodb.org/mongo-driver/bson"
 )
 
 const N = 1000
@@ -625,3 +627,79 @@ func Test_MarshalJSON(t *testing.T) {
 		t.Errorf("Expected no difference, got: %v", expected.Difference(actual))
 	}
 }
+
+func Test_UnmarshalBSON(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_UnmarshalBSON(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.UnmarshalBSON([]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_MarshalBSON(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))
+	}
+}

threadunsafe.go

@@ -29,6 +29,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+
+	"go.mongodb.org/mongo-driver/bson"
 )
 
 type threadUnsafeSet[T comparable] map[T]struct{}
@@ -328,3 +330,22 @@ func (s threadUnsafeSet[T]) UnmarshalJSON(b []byte) error {
 
 	return nil
 }
+
+// MarshalBSON creates a BSON array from the set.
+func (s threadUnsafeSet[T]) MarshalBSON() ([]byte, error) {
+	_, data, err := bson.MarshalValue(s.ToSlice())
+
+	return data, err
+}
+
+// UnmarshalBSON recreates a set from a BSON array.
+func (s threadUnsafeSet[T]) UnmarshalBSON(b []byte) error {
+	var i []T
+	err := bson.UnmarshalValue(bson.TypeArray, b, &i)
+	if err != nil {
+		return err
+	}
+	s.Append(i...)
+
+	return nil
+}