Update

vendor/tailscale.com/util/set/handle.go

@@ -9,20 +9,28 @@ package set
 type HandleSet[T any] map[Handle]T
 
 // Handle is an opaque comparable value that's used as the map key in a
-// HandleSet. The only way to get one is to call HandleSet.Add.
+// HandleSet.
 type Handle struct {
 	v *byte
 }
 
+// NewHandle returns a new handle value.
+func NewHandle() Handle {
+	return Handle{new(byte)}
+}
+
 // Add adds the element (map value) e to the set.
 //
-// It returns the handle (map key) with which e can be removed, using a map
-// delete.
+// It returns a new handle (map key) with which e can be removed, using a map
+// delete or the [HandleSet.Delete] method.
 func (s *HandleSet[T]) Add(e T) Handle {
-	h := Handle{new(byte)}
+	h := NewHandle()
 	if *s == nil {
 		*s = make(HandleSet[T])
 	}
 	(*s)[h] = e
 	return h
 }
+
+// Delete removes the element with handle h from the set.
+func (s HandleSet[T]) Delete(h Handle) { delete(s, h) }

vendor/tailscale.com/util/set/intset.go

@@ -0,0 +1,198 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package set
+
+import (
+	"iter"
+	"maps"
+	"math/bits"
+	"math/rand/v2"
+
+	"golang.org/x/exp/constraints"
+	"tailscale.com/util/mak"
+)
+
+// IntSet is a set optimized for integer values close to zero
+// or set of integers that are close in value.
+type IntSet[T constraints.Integer] struct {
+	// bits is a [bitSet] for numbers less than [bits.UintSize].
+	bits bitSet
+
+	// extra is a mapping of [bitSet] for numbers not in bits,
+	// where the key is a number modulo [bits.UintSize].
+	extra map[uint64]bitSet
+
+	// extraLen is the count of numbers in extra since len(extra)
+	// does not reflect that each bitSet may have multiple numbers.
+	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] {
+	return func(yield func(T) bool) {
+		if s.bits != 0 {
+			for i := range s.bits.values() {
+				if !yield(decodeZigZag[T](i)) {
+					return
+				}
+			}
+		}
+		if s.extra != nil {
+			for hi, bs := range s.extra {
+				for lo := range bs.values() {
+					if !yield(decodeZigZag[T](hi*bits.UintSize + lo)) {
+						return
+					}
+				}
+			}
+		}
+	}
+}
+
+// Contains reports whether e is in the set.
+func (s IntSet[T]) Contains(e T) bool {
+	if v := encodeZigZag(e); v < bits.UintSize {
+		return s.bits.contains(v)
+	} else {
+		hi, lo := v/uint64(bits.UintSize), v%uint64(bits.UintSize)
+		return s.extra[hi].contains(lo)
+	}
+}
+
+// Add adds e to the set.
+//
+// When storing a IntSet in a map as a value type,
+// it is important to re-assign the map entry after calling Add or Delete,
+// as the IntSet's representation may change.
+func (s *IntSet[T]) Add(e T) {
+	if v := encodeZigZag(e); v < bits.UintSize {
+		s.bits.add(v)
+	} else {
+		hi, lo := v/uint64(bits.UintSize), v%uint64(bits.UintSize)
+		if bs := s.extra[hi]; !bs.contains(lo) {
+			bs.add(lo)
+			mak.Set(&s.extra, hi, bs)
+			s.extra[hi] = bs
+			s.extraLen++
+		}
+	}
+}
+
+// AddSeq adds the values from seq to the set.
+func (s *IntSet[T]) AddSeq(seq iter.Seq[T]) {
+	for e := range seq {
+		s.Add(e)
+	}
+}
+
+// Len reports the number of elements in the set.
+func (s IntSet[T]) Len() int {
+	return s.bits.len() + s.extraLen
+}
+
+// Delete removes e from the set.
+//
+// When storing a IntSet in a map as a value type,
+// it is important to re-assign the map entry after calling Add or Delete,
+// as the IntSet's representation may change.
+func (s *IntSet[T]) Delete(e T) {
+	if v := encodeZigZag(e); v < bits.UintSize {
+		s.bits.delete(v)
+	} else {
+		hi, lo := v/uint64(bits.UintSize), v%uint64(bits.UintSize)
+		if bs := s.extra[hi]; bs.contains(lo) {
+			bs.delete(lo)
+			mak.Set(&s.extra, hi, bs)
+			s.extra[hi] = bs
+			s.extraLen--
+		}
+	}
+}
+
+// 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]{
+		bits:     s.bits,
+		extra:    maps.Clone(s.extra),
+		extraLen: s.extraLen,
+	}
+}
+
+type bitSet uint
+
+func (s bitSet) values() iter.Seq[uint64] {
+	return func(yield func(uint64) bool) {
+		// Hyrum-proofing: randomly iterate in forwards or reverse.
+		if rand.Uint64()%2 == 0 {
+			for i := 0; i < bits.UintSize; i++ {
+				if s.contains(uint64(i)) && !yield(uint64(i)) {
+					return
+				}
+			}
+		} else {
+			for i := bits.UintSize; i >= 0; i-- {
+				if s.contains(uint64(i)) && !yield(uint64(i)) {
+					return
+				}
+			}
+		}
+	}
+}
+func (s bitSet) len() int               { return bits.OnesCount(uint(s)) }
+func (s bitSet) contains(i uint64) bool { return s&(1<<i) > 0 }
+func (s *bitSet) add(i uint64)          { *s |= 1 << i }
+func (s *bitSet) delete(i uint64)       { *s &= ^(1 << i) }
+
+// encodeZigZag encodes an integer as an unsigned integer ensuring that
+// negative integers near zero still have a near zero positive value.
+// For unsigned integers, it returns the value verbatim.
+func encodeZigZag[T constraints.Integer](v T) uint64 {
+	var zero T
+	if ^zero >= 0 { // must be constraints.Unsigned
+		return uint64(v)
+	} else { // must be constraints.Signed
+		// See [google.golang.org/protobuf/encoding/protowire.EncodeZigZag]
+		return uint64(int64(v)<<1) ^ uint64(int64(v)>>63)
+	}
+}
+
+// decodeZigZag decodes an unsigned integer as an integer ensuring that
+// negative integers near zero still have a near zero positive value.
+// For unsigned integers, it returns the value verbatim.
+func decodeZigZag[T constraints.Integer](v uint64) T {
+	var zero T
+	if ^zero >= 0 { // must be constraints.Unsigned
+		return T(v)
+	} else { // must be constraints.Signed
+		// See [google.golang.org/protobuf/encoding/protowire.DecodeZigZag]
+		return T(int64(v>>1) ^ int64(v)<<63>>63)
+	}
+}

vendor/tailscale.com/util/set/smallset.go

@@ -0,0 +1,148 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package set
+
+import (
+	"iter"
+	"maps"
+
+	"tailscale.com/types/structs"
+)
+
+// SmallSet is a set that is optimized for reducing memory overhead when the
+// expected size of the set is 0 or 1 elements.
+//
+// The zero value of SmallSet is a usable empty set.
+//
+// When storing a SmallSet in a map as a value type, it is important to re-assign
+// the map entry after calling Add or Delete, as the SmallSet's representation
+// may change.
+//
+// Copying a SmallSet by value may alias the previous value. Use the Clone method
+// to create a new SmallSet with the same contents.
+type SmallSet[T comparable] struct {
+	_   structs.Incomparable // to prevent == mistakes
+	one T                    // if non-zero, then single item in set
+	m   Set[T]               // if non-nil, the set of items, which might be size 1 if it's the zero value of T
+}
+
+// Values returns an iterator over the elements of the set.
+// The iterator will yield the elements in no particular order.
+func (s SmallSet[T]) Values() iter.Seq[T] {
+	if s.m != nil {
+		return maps.Keys(s.m)
+	}
+	var zero T
+	return func(yield func(T) bool) {
+		if s.one != zero {
+			yield(s.one)
+		}
+	}
+}
+
+// Contains reports whether e is in the set.
+func (s SmallSet[T]) Contains(e T) bool {
+	if s.m != nil {
+		return s.m.Contains(e)
+	}
+	var zero T
+	return e != zero && s.one == e
+}
+
+// SoleElement returns the single value in the set, if the set has exactly one
+// element.
+//
+// If the set is empty or has more than one element, ok will be false and e will
+// be the zero value of T.
+func (s SmallSet[T]) SoleElement() (e T, ok bool) {
+	return s.one, s.Len() == 1
+}
+
+// Add adds e to the set.
+//
+// When storing a SmallSet in a map as a value type, it is important to
+// re-assign the map entry after calling Add or Delete, as the SmallSet's
+// representation may change.
+func (s *SmallSet[T]) Add(e T) {
+	var zero T
+	if s.m != nil {
+		s.m.Add(e)
+		return
+	}
+	// Non-zero elements can go into s.one.
+	if e != zero {
+		if s.one == zero {
+			s.one = e // Len 0 to Len 1
+			return
+		}
+		if s.one == e {
+			return // dup
+		}
+	}
+	// Need to make a multi map, either
+	// because we now have two items, or
+	// because e is the zero value.
+	s.m = Set[T]{}
+	if s.one != zero {
+		s.m.Add(s.one) // move single item to multi
+	}
+	s.m.Add(e) // add new item, possibly zero
+	s.one = zero
+}
+
+// Len reports the number of elements in the set.
+func (s SmallSet[T]) Len() int {
+	var zero T
+	if s.m != nil {
+		return s.m.Len()
+	}
+	if s.one != zero {
+		return 1
+	}
+	return 0
+}
+
+// Delete removes e from the set.
+//
+// When storing a SmallSet in a map as a value type, it is important to
+// re-assign the map entry after calling Add or Delete, as the SmallSet's
+// representation may change.
+func (s *SmallSet[T]) Delete(e T) {
+	var zero T
+	if s.m == nil {
+		if s.one == e {
+			s.one = zero
+		}
+		return
+	}
+	s.m.Delete(e)
+
+	// If the map size drops to zero, that means
+	// it only contained the zero value of T.
+	if s.m.Len() == 0 {
+		s.m = nil
+		return
+	}
+
+	// If the map size drops to one element and doesn't
+	// contain the zero value, we can switch back to the
+	// single-item representation.
+	if s.m.Len() == 1 {
+		for v := range s.m {
+			if v != zero {
+				s.one = v
+				s.m = nil
+			}
+		}
+	}
+	return
+}
+
+// Clone returns a copy of s that doesn't alias the original.
+func (s SmallSet[T]) Clone() SmallSet[T] {
+	return SmallSet[T]{
+		one: s.one,
+		m:   maps.Clone(s.m), // preserves nilness
+	}
+}