util/set: add iterator support to Set[T] (#20159)

This patch adds:

- Set.All which returns an iter.Seq to complement Set.Slice.

- Set.AddSeq which adds an iter.Seq.

- Set.DeleteSeq which deletes an iter.Seq to complement Set.AddSeq
  and provide the missing method for deleting multiple elements.

- Set.DeleteSlice and Set.DeleteSet to complement AddSlice and AddSet.

Updates #cleanup

Signed-off-by: Simon Law <sfllaw@tailscale.com>

util/set/set.go

@@ -6,6 +6,7 @@ package set
 
 import (
 	"encoding/json"
+	"iter"
 	"maps"
 	"reflect"
 	"sort"
@@ -31,9 +32,23 @@ func (s Set[T]) Clone() Set[T] {
 	return maps.Clone(s)
 }
 
+// All returns an iterator over all elements in s.
+// The iteration order is not specified
+// and is not guaranteed to be the same from one call to the next.
+func (s Set[T]) All() iter.Seq[T] {
+	return maps.Keys(s)
+}
+
 // Add adds e to s.
 func (s Set[T]) Add(e T) { s[e] = struct{}{} }
 
+// AddSeq adds each element of es to s.
+func (s Set[T]) AddSeq(es iter.Seq[T]) {
+	for e := range es {
+		s.Add(e)
+	}
+}
+
 // AddSlice adds each element of es to s.
 func (s Set[T]) AddSlice(es []T) {
 	for _, e := range es {
@@ -105,6 +120,27 @@ func genOrderedSwapper(rt reflect.Type) func(reflect.Value) func(i, j int) bool
 // Delete removes e from the set.
 func (s Set[T]) Delete(e T) { delete(s, e) }
 
+// DeleteSeq removes all elements in es from the set.
+func (s Set[T]) DeleteSeq(es iter.Seq[T]) {
+	for e := range es {
+		s.Delete(e)
+	}
+}
+
+// DeleteSlice removes all elements in es from the set.
+func (s Set[T]) DeleteSlice(es []T) {
+	for _, e := range es {
+		s.Delete(e)
+	}
+}
+
+// DeleteSet removes all elements in es from the set.
+func (s Set[T]) DeleteSet(es Set[T]) {
+	for e := range es {
+		s.Delete(e)
+	}
+}
+
 // Contains reports whether s contains e.
 func (s Set[T]) Contains(e T) bool {
 	_, ok := s[e]

util/set/set_test.go

@@ -50,6 +50,40 @@ func TestSet(t *testing.T) {
 			t.Errorf("slice missing %d (%#v)", e, es)
 		}
 	}
+
+	s.Delete(1)
+	if s.Contains(1) {
+		t.Error("shouldn't have 1")
+	}
+	if !s.Contains(2) {
+		t.Error("missing 2")
+	}
+	if !s.Contains(3) {
+		t.Error("missing 3")
+	}
+	if !s.Contains(4) {
+		t.Error("missing 4")
+	}
+	if s.Len() != 3 {
+		t.Errorf("wrong len %d; want 3", s.Len())
+	}
+
+	s.DeleteSeq(slices.Values([]int{2, 3}))
+	if s.Contains(1) {
+		t.Error("shouldn't have 1")
+	}
+	if s.Contains(2) {
+		t.Error("shouldn't have 2")
+	}
+	if s.Contains(3) {
+		t.Error("shouldn't have 3")
+	}
+	if !s.Contains(4) {
+		t.Error("missing 4")
+	}
+	if s.Len() != 1 {
+		t.Errorf("wrong len %d; want 1", s.Len())
+	}
 }
 
 func TestSetOf(t *testing.T) {