Improving set operations

internal/kvstore/kvstore.go

@@ -9,9 +9,11 @@ import (
 // Store represents a container of key-value pairs.
 // Used internally for key-grouping and key-joining operations.
 type Store[K, V any] interface {
+	Size() int
 	Get(key K) opt.Optional[V]
 	Put(key K, value V)
 	ForEach(func(key K, value V) bool) bool
+	ForEachKey(func(key K) bool) bool
 }
 
 // NewMapped creates a new Store backed by a map.
@@ -49,6 +51,10 @@ func SortedMaker[K any, V any](compare cmp.Comparer[K]) Maker[K, V] {
 // The key type K must be comparable.
 type mappedStore[K comparable, V any] map[K]V
 
+func (s mappedStore[K, V]) Size() int {
+	return len(s)
+}
+
 func (s mappedStore[K, V]) Get(key K) opt.Optional[V] {
 	if v, ok := s[key]; ok {
 		return opt.Of(v)
@@ -69,6 +75,15 @@ func (s mappedStore[K, V]) ForEach(yield func(K, V) bool) bool {
 	return true
 }
 
+func (s mappedStore[K, V]) ForEachKey(yield func(K) bool) bool {
+	for k := range s {
+		if !yield(k) {
+			return false
+		}
+	}
+	return true
+}
+
 // sortedStore provides an implementation of Store using sorted slices and binary-search.
 // The keys are ordered using the given cmp.Comparer.
 type sortedStore[K any, V any] struct {
@@ -77,6 +92,10 @@ type sortedStore[K any, V any] struct {
 	values  []V
 }
 
+func (s *sortedStore[K, V]) Size() int {
+	return len(s.keys)
+}
+
 func (s *sortedStore[K, V]) Get(key K) opt.Optional[V] {
 	if i, ok := s.indexOf(key); ok {
 		return opt.Of(s.values[i])
@@ -102,9 +121,18 @@ func (s *sortedStore[K, V]) indexOf(key K) (int, bool) {
 	return slices.BinarySearchFunc(s.keys, key, s.compare)
 }
 
-func (s *sortedStore[K, V]) ForEach(f func(K, V) bool) bool {
+func (s *sortedStore[K, V]) ForEach(yield func(K, V) bool) bool {
 	for i, k := range s.keys {
-		if !f(k, s.values[i]) {
+		if !yield(k, s.values[i]) {
+			return false
+		}
+	}
+	return true
+}
+
+func (s *sortedStore[K, V]) ForEachKey(yield func(K) bool) bool {
+	for _, k := range s.keys {
+		if !yield(k) {
 			return false
 		}
 	}

pkg/stream/set.go

@@ -91,17 +91,11 @@ func IntersectionAllBy[E any](compare cmp.Comparer[E], ss ...Stream[E]) Stream[E
 }
 
 func intersection[E any](s1, s2 Stream[E], kv kvstore.Maker[E, struct{}]) Stream[E] {
-	// Exactly 2 streams; intersect ss[0] and ss[1].
 	return func(yield Consumer[E]) bool {
-		// Index elements of the first stream into a set.
-		seen := kv()
-		s1(func(e E) bool {
-			seen.Put(e, struct{}{})
-			return true
-		})
-		// Yield elements of the second stream that are in the set.
-		return s2(func(e E) bool {
-			if seen.Get(e).Present() {
+		set2 := toSet(s2, kv)
+		// Yield elements of the first stream that are in the set.
+		return s1(func(e E) bool {
+			if set2.Get(e).Present() {
 				return yield(e)
 			}
 			return true
@@ -134,15 +128,10 @@ func DifferenceBy[E any](s1, s2 Stream[E], compare cmp.Comparer[E]) Stream[E] {
 
 func difference[E any](s1, s2 Stream[E], kv kvstore.Maker[E, struct{}]) Stream[E] {
 	return func(yield Consumer[E]) bool {
-		// Index elements of the second stream into a set.
-		seen := kv()
-		s2(func(e E) bool {
-			seen.Put(e, struct{}{})
-			return true
-		})
+		set2 := toSet(s2, kv)
 		// Yield elements of the first stream that are not in the set.
 		return s1(func(e E) bool {
-			if !seen.Get(e).Present() {
+			if !set2.Get(e).Present() {
 				return yield(e)
 			}
 			return true
@@ -181,16 +170,7 @@ func SymmetricDifferenceBy[E any](s1, s2 Stream[E], compare cmp.Comparer[E]) Str
 //	ok := stream.Subset(stream.Of(1, 2), stream.Of(1, 2, 3, 4))
 //	fmt.Println(ok) // "true"
 func Subset[E comparable](s1, s2 Stream[E]) bool {
-	// Index elements of the second stream into a set.
-	seen := kvstore.MappedMaker[E, struct{}]()()
-	s2(func(e E) bool {
-		seen.Put(e, struct{}{})
-		return true
-	})
-	// Check if all elements of the first stream are in the set.
-	return s1(func(e E) bool {
-		return seen.Get(e).Present()
-	})
+	return subset(s1, s2, kvstore.MappedMaker[E, struct{}]())
 }
 
 // SubsetBy returns true if all elements of the first stream are in the second stream, compared by the given cmp.Comparer.
@@ -200,15 +180,16 @@ func Subset[E comparable](s1, s2 Stream[E]) bool {
 //	ok := stream.SubsetBy(stream.Of(1, 2), stream.Of(1, 2, 3, 4), cmp.Natural[int]())
 //	fmt.Println(ok) // "true"
 func SubsetBy[E any](s1, s2 Stream[E], compare cmp.Comparer[E]) bool {
+	return subset(s1, s2, kvstore.SortedMaker[E, struct{}](compare))
+}
+
+func subset[E any](s1, s2 Stream[E], kv kvstore.Maker[E, struct{}]) bool {
 	// Index elements of the second stream into a set.
-	seen := kvstore.SortedMaker[E, struct{}](compare)()
-	s2(func(e E) bool {
-		seen.Put(e, struct{}{})
-		return true
-	})
+	set2 := toSet(s2, kv)
 	// Check if all elements of the first stream are in the set.
+	// If an element is not in the set, the result is false
 	return s1(func(e E) bool {
-		return seen.Get(e).Present()
+		return set2.Get(e).Present()
 	})
 }
 
@@ -233,15 +214,15 @@ func SupersetBy[E any](s1, s2 Stream[E], compare cmp.Comparer[E]) bool {
 	return SubsetBy(s2, s1, compare)
 }
 
-// SetEqual returns true if the two streams contain the same elements (in any order).
+// SetEqual returns true if the two streams contain the same elements, ignoring order and duplicates (ie: set equality).
 // The element type E must be comparable.
 //
 // Example usage:
 //
 //	ok := stream.SetEqual(stream.Of(1, 2, 3), stream.Of(3, 2, 1))
 //	fmt.Println(ok) // "true"
 func SetEqual[E comparable](s1, s2 Stream[E]) bool {
-	return Subset(s1, s2) && Subset(s2, s1)
+	return setEqual(s1, s2, kvstore.MappedMaker[E, struct{}]())
 }
 
 // SetEqualBy returns true if the two streams contain the same elements (in any order), compared by the given cmp.Comparer.
@@ -251,5 +232,25 @@ func SetEqual[E comparable](s1, s2 Stream[E]) bool {
 //	ok := stream.SetEqualBy(stream.Of(1, 2, 3), stream.Of(3, 2, 1), cmp.Natural[int]())
 //	fmt.Println(ok) // "true"
 func SetEqualBy[E any](s1, s2 Stream[E], compare cmp.Comparer[E]) bool {
-	return SubsetBy(s1, s2, compare) && SubsetBy(s2, s1, compare)
+	return setEqual(s1, s2, kvstore.SortedMaker[E, struct{}](compare))
+}
+
+func setEqual[E any](s1, s2 Stream[E], kv kvstore.Maker[E, struct{}]) bool {
+	set1 := toSet(s1, kv)
+	set2 := toSet(s2, kv)
+	if set1.Size() != set2.Size() {
+		return false
+	}
+	return set1.ForEachKey(func(e E) bool {
+		return set2.Get(e).Present()
+	})
+}
+
+func toSet[E any](s Stream[E], kv kvstore.Maker[E, struct{}]) kvstore.Store[E, struct{}] {
+	set := kv()
+	s(func(e E) bool {
+		set.Put(e, struct{}{})
+		return true
+	})
+	return set
 }

pkg/stream/set_test.go

@@ -204,16 +204,28 @@ func TestSetEqual(t *testing.T) {
 	if SetEqual(Of(1, 2, 3), Of(2, 3, 4)) {
 		t.Errorf("SetEqual(Of(1, 2, 3), Of(2, 3, 4)) = true; want false")
 	}
+	if SetEqual(Of(1, 2, 3), Of(1, 2)) {
+		t.Errorf("SetEqual(Of(1, 2, 3), Of(1, 2) = true; want false")
+	}
 	if !SetEqual(Of(1, 2, 3), Of(3, 2, 1)) {
 		t.Errorf("SetEqual(Of(1, 2, 3), Of(3, 2, 1)) = false; want true")
 	}
+	if !SetEqual(Of(1, 2, 3, 3), Of(3, 2, 1, 1)) {
+		t.Errorf("SetEqual(Of(1, 2, 3, 3), Of(3, 2, 1, 1)) = false; want true")
+	}
 }
 
 func TestSetEqualBy(t *testing.T) {
 	if SetEqualBy(Of(1, 2, 3), Of(2, 3, 4), cmp.Natural[int]()) {
 		t.Errorf("SetEqualBy(Of(1, 2, 3), Of(2, 3, 4), cmp.Natural[int]()) = true; want false")
 	}
+	if SetEqualBy(Of(1, 2, 3), Of(1, 2), cmp.Natural[int]()) {
+		t.Errorf("SetEqualBy(Of(1, 2, 3), Of(1, 2), cmp.Natural[int]()) = true; want false")
+	}
 	if !SetEqualBy(Of(1, 2, 3), Of(3, 2, 1), cmp.Natural[int]()) {
 		t.Errorf("SetEqualBy(Of(1, 2, 3), Of(3, 2, 1), cmp.Natural[int]()) = false; want true")
 	}
+	if !SetEqualBy(Of(1, 2, 3, 3), Of(3, 2, 1, 1), cmp.Natural[int]()) {
+		t.Errorf("SetEqualBy(Of(1, 2, 3, 3), Of(3, 2, 1, 1), cmp.Natural[int]()) = false; want true")
+	}
 }