diff --git a/util/set/set.go b/util/set/set.go index c3d2350a7..38bf676f8 100644 --- a/util/set/set.go +++ b/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] diff --git a/util/set/set_test.go b/util/set/set_test.go index 2188cbb4d..d3d118bee 100644 --- a/util/set/set_test.go +++ b/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) {