diff --git a/util/set/handle.go b/util/set/handle.go new file mode 100644 index 000000000..471ceeba2 --- /dev/null +++ b/util/set/handle.go @@ -0,0 +1,28 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +package set + +// HandleSet is a set of T. +// +// It is not safe for concurrent use. +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. +type Handle struct { + v *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. +func (s *HandleSet[T]) Add(e T) Handle { + h := Handle{new(byte)} + if *s == nil { + *s = make(HandleSet[T]) + } + (*s)[h] = e + return h +} diff --git a/util/set/set.go b/util/set/set.go index e6f3ef1f0..987747892 100644 --- a/util/set/set.go +++ b/util/set/set.go @@ -7,9 +7,33 @@ // Set is a set of T. type Set[T comparable] map[T]struct{} +// SetOf returns a new set constructed from the elements in slice. +func SetOf[T comparable](slice []T) Set[T] { + s := make(Set[T]) + s.AddSlice(slice) + return s +} + // Add adds e to the set. func (s Set[T]) Add(e T) { s[e] = struct{}{} } +// AddSlice adds each element of es to the set. +func (s Set[T]) AddSlice(es []T) { + for _, e := range es { + s.Add(e) + } +} + +// Slice returns the elements of the set as a slice. The elements will not be +// in any particular order. +func (s Set[T]) Slice() []T { + es := make([]T, 0, s.Len()) + for k := range s { + es = append(es, k) + } + return es +} + // Delete removes e from the set. func (s Set[T]) Delete(e T) { delete(s, e) } @@ -21,27 +45,3 @@ func (s Set[T]) Contains(e T) bool { // Len reports the number of items in s. func (s Set[T]) Len() int { return len(s) } - -// HandleSet is a set of T. -// -// It is not safe for concurrent use. -type HandleSet[T any] map[Handle]T - -// Handle is a opaque comparable value that's used as the map key -// in a HandleSet. The only way to get one is to call HandleSet.Add. -type Handle struct { - v *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. -func (s *HandleSet[T]) Add(e T) Handle { - h := Handle{new(byte)} - if *s == nil { - *s = make(HandleSet[T]) - } - (*s)[h] = e - return h -} diff --git a/util/set/set_test.go b/util/set/set_test.go index 7a920ed88..e898f4f69 100644 --- a/util/set/set_test.go +++ b/util/set/set_test.go @@ -3,7 +3,10 @@ package set -import "testing" +import ( + "slices" + "testing" +) func TestSet(t *testing.T) { s := Set[int]{} @@ -21,4 +24,41 @@ func TestSet(t *testing.T) { if s.Len() != 2 { t.Errorf("wrong len %d; want 2", s.Len()) } + + more := []int{3, 4} + s.AddSlice(more) + if !s.Contains(3) { + t.Error("missing 3") + } + if !s.Contains(4) { + t.Error("missing 4") + } + if s.Contains(5) { + t.Error("shouldn't have 5") + } + if s.Len() != 4 { + t.Errorf("wrong len %d; want 4", s.Len()) + } + + es := s.Slice() + if len(es) != 4 { + t.Errorf("slice has wrong len %d; want 4", len(es)) + } + for _, e := range []int{1, 2, 3, 4} { + if !slices.Contains(es, e) { + t.Errorf("slice missing %d (%#v)", e, es) + } + } +} + +func TestSetOf(t *testing.T) { + s := SetOf[int]([]int{1, 2, 3, 4, 4, 1}) + if s.Len() != 4 { + t.Errorf("wrong len %d; want 2", s.Len()) + } + for _, n := range []int{1, 2, 3, 4} { + if !s.Contains(n) { + t.Errorf("should contain %d", n) + } + } } diff --git a/util/set/slice.go b/util/set/slice.go index fe764b550..9c300cc28 100644 --- a/util/set/slice.go +++ b/util/set/slice.go @@ -15,7 +15,7 @@ type Slice[T comparable] struct { set map[T]bool // nil until/unless slice is large enough } -// Slice returns the a view of the underlying slice. +// Slice returns a view of the underlying slice. // The elements are in order of insertion. // The returned value is only valid until ss is modified again. func (ss *Slice[T]) Slice() views.Slice[T] { return views.SliceOf(ss.slice) }