diff --git a/internal/kvstore/kvstore.go b/internal/kvstore/kvstore.go index f4fb6d3..9ab9e2f 100644 --- a/internal/kvstore/kvstore.go +++ b/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 } } diff --git a/pkg/stream/set.go b/pkg/stream/set.go index 77de154..6ba60de 100644 --- a/pkg/stream/set.go +++ b/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,7 +214,7 @@ 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: @@ -241,7 +222,7 @@ func SupersetBy[E any](s1, s2 Stream[E], compare cmp.Comparer[E]) bool { // 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 } diff --git a/pkg/stream/set_test.go b/pkg/stream/set_test.go index 6d05603..6f69940 100644 --- a/pkg/stream/set_test.go +++ b/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") + } }