Something went wrong. Try again.
forked from tailscale/tailscale
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199// Copyright (c) Tailscale Inc & contributors// 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 := range bits.UintSize { 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) }}