Merge pull request #1 from fortio/more_ops

Adding Subset, Equals, Minus, Plus, Len, Clear and JSON serialization/deserialization

Author: ldemailly

sets.go

@@ -8,6 +8,7 @@
 package sets // import "fortio.org/sets"
 
 import (
+	"encoding/json"
 	"fmt"
 	"sort"
 	"strings"
@@ -98,6 +99,48 @@ func (s Set[T]) Elements() []T {
 	return res
 }
 
+// Subset returns true if all elements of s are in the passed in set.
+func (s Set[T]) Subset(bigger Set[T]) bool {
+	for k := range s {
+		if !bigger.Has(k) {
+			return false
+		}
+	}
+	return true
+}
+
+// Minus mutates the receiver to remove all the elements of the passed in set.
+// If you want a copy use s.Clone().Minus(other). Returns the receiver for chaining.
+func (s Set[T]) Minus(other Set[T]) Set[T] {
+	for k := range other {
+		s.Remove(k)
+	}
+	return s
+}
+
+// Plus is similar to Union but mutates the receiver. Added for symmetry with Minus.
+// Returns the receiver for chaining.
+func (s Set[T]) Plus(others ...Set[T]) Set[T] {
+	for _, o := range others {
+		s.Add(o.Elements()...)
+	}
+	return s
+}
+
+func (s Set[T]) Equals(other Set[T]) bool {
+	return s.Subset(other) && other.Subset(s)
+}
+
+func (s Set[T]) Len() int {
+	return len(s)
+}
+
+func (s Set[T]) Clear() {
+	for k := range s {
+		delete(s, k)
+	}
+}
+
 // String() returns a coma separated list of the elements in the set.
 // This is mostly for troubleshooting/debug output unless the [T] serializes
 // to a string that doesn't contain commas.
@@ -131,6 +174,37 @@ func XOR[T comparable](a, b Set[T]) {
 	RemoveCommon(a, b)
 }
 
+// -- Serialization
+
+// MarshalJSON implements the json.Marshaler interface and only gets the elements as an array.
+func (s Set[T]) MarshalJSON() ([]byte, error) {
+	// How to handle all ordered at once??
+	switch v := any(s).(type) {
+	case Set[string]:
+		return json.Marshal(Sort(v))
+	case Set[int]:
+		return json.Marshal(Sort(v))
+	case Set[int8]:
+		return json.Marshal(Sort(v))
+	case Set[int64]:
+		return json.Marshal(Sort(v))
+	case Set[float64]:
+		return json.Marshal(Sort(v))
+	default:
+		return json.Marshal(s.Elements())
+	}
+}
+
+// UnmarshalJSON implements the json.Unmarshaler interface turns the slice back to a Set.
+func (s *Set[T]) UnmarshalJSON(data []byte) error {
+	var items []T
+	if err := json.Unmarshal(data, &items); err != nil {
+		return err
+	}
+	*s = New[T](items...)
+	return nil
+}
+
 // -- Additional operations on sets of ordered types
 
 func Sort[Q constraints.Ordered](s Set[Q]) []Q {

sets_test.go

@@ -4,6 +4,7 @@
 package sets_test
 
 import (
+	"encoding/json"
 	"testing"
 
 	"fortio.org/assert"
@@ -13,6 +14,10 @@ import (
 func TestSetToString(t *testing.T) {
 	s := sets.Set[string]{"z": {}, "a": {}, "c": {}, "b": {}}
 	assert.Equal(t, "a,b,c,z", s.String())
+	assert.Equal(t, s.Len(), 4)
+	s.Clear()
+	assert.Equal(t, "", s.String())
+	assert.Equal(t, s.Len(), 0)
 }
 
 func TestArrayToSet(t *testing.T) {
@@ -34,6 +39,10 @@ func TestRemoveCommon(t *testing.T) {
 	// also check clone is not modifying the original etc
 	setAA = setB.Clone() // putting B in AA on purpose and vice versa
 	setBB = setA.Clone()
+	assert.True(t, setAA.Equals(setB))
+	assert.True(t, setB.Equals(setAA))
+	assert.False(t, setAA.Equals(setA))
+	assert.False(t, setB.Equals(setBB))
 	sets.XOR(setAA, setBB)
 	assert.Equal(t, "a,c", setBB.String())
 	assert.Equal(t, "e,f,g", setAA.String())
@@ -42,6 +51,24 @@ func TestRemoveCommon(t *testing.T) {
 	assert.False(t, setBB.Has("c"))
 }
 
+func TestMinus(t *testing.T) {
+	setA := sets.New("a", "b", "c", "d")
+	setB := sets.New("b", "d", "e", "f", "g")
+	setAB := setA.Clone().Minus(setB)
+	setBA := setB.Clone().Minus(setA)
+	assert.Equal(t, "a,c", setAB.String())
+	assert.Equal(t, "e,f,g", setBA.String())
+}
+
+func TestPlus(t *testing.T) {
+	setA := sets.New("a", "b", "c", "d")
+	setB := sets.New("b", "d", "e", "f", "g")
+	setAB := setA.Clone().Plus(setB)
+	setBA := setB.Clone().Plus(setA)
+	assert.Equal(t, "a,b,c,d,e,f,g", setAB.String())
+	assert.Equal(t, "a,b,c,d,e,f,g", setBA.String())
+}
+
 func TestUnion(t *testing.T) {
 	setA := sets.New("a", "b", "c", "d")
 	setB := sets.New("b", "d", "e", "f", "g")
@@ -64,3 +91,83 @@ func TestIntersection2(t *testing.T) {
 	setC := sets.Intersection(setA, setB, setA)
 	assert.Equal(t, "", setC.String())
 }
+
+func TestSubset(t *testing.T) {
+	setA := sets.New("a", "b", "c", "d")
+	setB := sets.New("b", "d", "e", "f", "g")
+	setC := sets.New("b", "d")
+	assert.True(t, setC.Subset(setA))
+	assert.True(t, setA.Subset(setA))
+	assert.False(t, setA.Subset(setC))
+	assert.False(t, setA.Subset(setB))
+	assert.False(t, setB.Subset(setA))
+}
+
+func TestJSON(t *testing.T) {
+	setA := sets.New("c,d", "a b", "y\000z", "mno")
+	b, err := json.Marshal(setA)
+	assert.NoError(t, err)
+	assert.Equal(t, `["a b","c,d","mno","y\u0000z"]`, string(b))
+	jsonStr := `[
+		"a,b",
+		"c,d"
+	]`
+	setB := sets.New[string]()
+	err = json.Unmarshal([]byte(jsonStr), &setB)
+	assert.NoError(t, err)
+	assert.Equal(t, setB.Len(), 2)
+	assert.True(t, setB.Has("a,b"))
+	assert.True(t, setB.Has("c,d"))
+	setI := sets.New(3, 42, 7, 10)
+	b, err = json.Marshal(setI)
+	assert.NoError(t, err)
+	assert.Equal(t, `[3,7,10,42]`, string(b))
+	smallIntSet := sets.New[int8](66, 65, 67) // if using byte, aka uint8, one gets base64("ABC")
+	b, err = json.Marshal(smallIntSet)
+	assert.NoError(t, err)
+	t.Logf("smallIntSet: %q", string(b))
+	assert.Equal(t, `[65,66,67]`, string(b))
+	floatSet := sets.New[float64](2.3, 1.1, -7.6, 42)
+	b, err = json.Marshal(floatSet)
+	assert.NoError(t, err)
+	t.Logf("floatSet: %q", string(b))
+	assert.Equal(t, `[-7.6,1.1,2.3,42]`, string(b))
+	i64Set := sets.New[int64](2, 1, -7, 42)
+	b, err = json.Marshal(i64Set)
+	assert.NoError(t, err)
+	t.Logf("i64Set: %q", string(b))
+	assert.Equal(t, `[-7,1,2,42]`, string(b))
+}
+
+type foo struct {
+	X int
+}
+
+func TestNonOrderedJSON(t *testing.T) {
+	s := sets.New(
+		foo{3},
+		foo{1},
+		foo{2},
+		foo{4},
+	)
+	b, err := json.Marshal(s)
+	t.Logf("b: %s", string(b))
+	assert.NoError(t, err)
+	// though I guess given it could be in any order it could be accidentally sorted too
+	assert.NotEqual(t, `[{"X":1},{"X":2},{"X":3},{"X":4}]`, string(b))
+	u := sets.New[foo]()
+	json.Unmarshal(b, &u)
+	assert.NoError(t, err)
+	assert.Equal(t, 4, u.Len())
+	assert.True(t, s.Equals(u))
+}
+
+func TestBadJson(t *testing.T) {
+	jsonStr := `[
+		"a,b",
+		"c,d"
+	]`
+	s := sets.New[int]()
+	err := json.Unmarshal([]byte(jsonStr), &s)
+	assert.Error(t, err)
+}