diff --git a/pkg/genutil/mapz/set.go b/pkg/genutil/mapz/set.go index 21bd52be99..13fb55a40f 100644 --- a/pkg/genutil/mapz/set.go +++ b/pkg/genutil/mapz/set.go @@ -58,6 +58,23 @@ func (s *Set[T]) Extend(values []T) { } } +// Merge adds all the values from the other set to this set. +func (s *Set[T]) Merge(other *Set[T]) { + for value := range other.values { + s.values[value] = struct{}{} + } +} + +// Union adds all the values from the other set to this set, +// returning a new set. +func (s *Set[T]) Union(other *Set[T]) *Set[T] { + cpy := s.Copy() + for value := range other.values { + cpy.values[value] = struct{}{} + } + return cpy +} + // IntersectionDifference removes any values from this set that // are not shared with the other set. Returns the same set. func (s *Set[T]) IntersectionDifference(other *Set[T]) *Set[T] { diff --git a/pkg/genutil/mapz/set_test.go b/pkg/genutil/mapz/set_test.go index 7bdbef37a7..fb9f821bc9 100644 --- a/pkg/genutil/mapz/set_test.go +++ b/pkg/genutil/mapz/set_test.go @@ -129,6 +129,40 @@ func TestEqual(t *testing.T) { require.False(t, NewSet[string]("1", "2").Equal(NewSet[string]("1", "3"))) } +func TestUnion(t *testing.T) { + u1 := NewSet[string]("1", "2").Union(NewSet[string]("2", "3")).AsSlice() + sort.Strings(u1) + + u2 := NewSet[string]("2", "3").Union(NewSet[string]("1", "2")).AsSlice() + sort.Strings(u2) + + require.Equal(t, []string{"1", "2", "3"}, u1) + require.Equal(t, []string{"1", "2", "3"}, u2) +} + +func TestMerge(t *testing.T) { + u1 := NewSet[string]("1", "2") + u2 := NewSet[string]("2", "3") + + u1.Merge(u2) + + slice := u1.AsSlice() + sort.Strings(slice) + + require.Equal(t, []string{"1", "2", "3"}, slice) + + // Try the reverse. + u1 = NewSet[string]("1", "2") + u2 = NewSet[string]("2", "3") + + u2.Merge(u1) + + slice = u2.AsSlice() + sort.Strings(slice) + + require.Equal(t, []string{"1", "2", "3"}, slice) +} + func TestSetIntersectionDifference(t *testing.T) { tcs := []struct { first []int @@ -271,3 +305,48 @@ func BenchmarkEqual(b *testing.B) { set.Equal(other) } } + +func BenchmarkExtendFromSlice(b *testing.B) { + set := NewSet[int]() + for i := 0; i < b.N; i++ { + set.Add(i) + } + other := NewSet[int]() + for i := 0; i < b.N; i++ { + other.Add(i) + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + set.Extend(other.AsSlice()) + } +} + +func BenchmarkMerge(b *testing.B) { + set := NewSet[int]() + for i := 0; i < b.N; i++ { + set.Add(i) + } + other := NewSet[int]() + for i := 0; i < b.N; i++ { + other.Add(i) + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + set.Merge(other) + } +} + +func BenchmarkUnion(b *testing.B) { + set := NewSet[int]() + for i := 0; i < b.N; i++ { + set.Add(i) + } + other := NewSet[int]() + for i := 0; i < b.N; i++ { + other.Add(i) + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + set.Union(other) + } +}