util/set: add more functionality to IntSet

This adds IntsOf, IntSet.DelSeq, IntSet.Equal.

Updates tailscale/corp#29809
Signed-off-by: Joe Tsai <joetsai@digital-static.net>
This commit is contained in:
Joe Tsai 2025-07-23 15:50:51 -07:00
parent c87f44b687
commit 3088585a68
2 changed files with 32 additions and 0 deletions

View File

@ -28,6 +28,15 @@ type IntSet[T constraints.Integer] struct {
extraLen int
}
// IntsOf constructs an [IntSet] with the provided elements.
func IntsOf[T constraints.Integer](slice ...T) IntSet[T] {
var s IntSet[T]
for _, e := range slice {
s.Add(e)
}
return s
}
// Values returns an iterator over the elements of the set.
// The iterator will yield the elements in no particular order.
func (s IntSet[T]) Values() iter.Seq[T] {
@ -111,6 +120,23 @@ func (s *IntSet[T]) Delete(e T) {
}
}
// DeleteSeq deletes the values in seq from the set.
func (s *IntSet[T]) DeleteSeq(seq iter.Seq[T]) {
for e := range seq {
s.Delete(e)
}
}
// Equal reports whether s is equal to other.
func (s IntSet[T]) Equal(other IntSet[T]) bool {
for hi, bits := range s.extra {
if other.extra[hi] != bits {
return false
}
}
return s.extraLen == other.extraLen && s.bits == other.bits
}
// Clone returns a copy of s that doesn't alias the original.
func (s IntSet[T]) Clone() IntSet[T] {
return IntSet[T]{

View File

@ -47,6 +47,9 @@ func TestIntSet(t *testing.T) {
deleteInt(t, ss, &si, math.MinInt64)
deleteInt(t, ss, &si, math.MaxInt64)
intValues(t, ss, si)
if !si.Equal(IntsOf(ss.Slice()...)) {
t.Errorf("{%v}.Equal({%v}) = false, want true", si, ss)
}
})
t.Run("Uint64", func(t *testing.T) {
@ -80,6 +83,9 @@ func TestIntSet(t *testing.T) {
intValues(t, ss, si)
deleteInt(t, ss, &si, math.MaxInt64)
intValues(t, ss, si)
if !si.Equal(IntsOf(ss.Slice()...)) {
t.Errorf("{%v}.Equal({%v}) = false, want true", si, ss)
}
})
}