Add Union and Merge to Set

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] {

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)
+	}
+}