From 8833dc51f1e9c66442911a32e477e34cd77df767 Mon Sep 17 00:00:00 2001
From: Chris Palmer <cpalmer@tailscale.com>
Date: Fri, 29 Sep 2023 14:31:02 -0700
Subject: [PATCH] util/set: add some useful utility functions for Set (#9535)

Also give each type of set its own file.

Updates #cleanup

Signed-off-by: Chris Palmer <cpalmer@tailscale.com>
---
 util/set/handle.go   | 28 ++++++++++++++++++++++++++
 util/set/set.go      | 48 ++++++++++++++++++++++----------------------
 util/set/set_test.go | 42 +++++++++++++++++++++++++++++++++++++-
 util/set/slice.go    |  2 +-
 4 files changed, 94 insertions(+), 26 deletions(-)
 create mode 100644 util/set/handle.go

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) }