util/set: add Of variant of SetOf that takes variadic parameter

set.Of(1, 2, 3) is prettier than set.SetOf([]int{1, 2, 3}).

I was going to change the signature of SetOf but then I noticed its
name has stutter anyway, so I kept it for compatibility. People can
prefer to use set.Of for new code or slowly migrate.

Also add a lazy Make method, which I often find myself wanting,
without having to resort to uglier mak.Set(&set, k, struct{}{}).

Updates #cleanup

Change-Id: Ic6f3870115334efcbd65e79c437de2ad3edb7625
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
pull/12019/head
Brad Fitzpatrick 7 months ago committed by Brad Fitzpatrick
parent 80decd83c1
commit 1a963342c7

@ -14,6 +14,11 @@ type Set[T comparable] map[T]struct{}
// SetOf returns a new set constructed from the elements in slice. // SetOf returns a new set constructed from the elements in slice.
func SetOf[T comparable](slice []T) Set[T] { func SetOf[T comparable](slice []T) Set[T] {
return Of(slice...)
}
// Of returns a new set constructed from the elements in slice.
func Of[T comparable](slice ...T) Set[T] {
s := make(Set[T]) s := make(Set[T])
s.AddSlice(slice) s.AddSlice(slice)
return s return s
@ -41,6 +46,13 @@ func (s Set[T]) AddSet(es Set[T]) {
} }
} }
// Make lazily initializes the map pointed to by s to be non-nil.
func (s *Set[T]) Make() {
if *s == nil {
*s = make(Set[T])
}
}
// Slice returns the elements of the set as a slice. The elements will not be // Slice returns the elements of the set as a slice. The elements will not be
// in any particular order. // in any particular order.
func (s Set[T]) Slice() []T { func (s Set[T]) Slice() []T {

@ -53,7 +53,7 @@ func TestSet(t *testing.T) {
} }
func TestSetOf(t *testing.T) { func TestSetOf(t *testing.T) {
s := SetOf[int]([]int{1, 2, 3, 4, 4, 1}) s := Of(1, 2, 3, 4, 4, 1)
if s.Len() != 4 { if s.Len() != 4 {
t.Errorf("wrong len %d; want 4", s.Len()) t.Errorf("wrong len %d; want 4", s.Len())
} }
@ -74,20 +74,20 @@ func TestEqual(t *testing.T) {
tests := []test{ tests := []test{
{ {
"equal", "equal",
SetOf([]int{1, 2, 3, 4}), Of(1, 2, 3, 4),
SetOf([]int{1, 2, 3, 4}), Of(1, 2, 3, 4),
true, true,
}, },
{ {
"not equal", "not equal",
SetOf([]int{1, 2, 3, 4}), Of(1, 2, 3, 4),
SetOf([]int{1, 2, 3, 5}), Of(1, 2, 3, 5),
false, false,
}, },
{ {
"different lengths", "different lengths",
SetOf([]int{1, 2, 3, 4, 5}), Of(1, 2, 3, 4, 5),
SetOf([]int{1, 2, 3, 5}), Of(1, 2, 3, 5),
false, false,
}, },
} }
@ -100,7 +100,7 @@ func TestEqual(t *testing.T) {
} }
func TestClone(t *testing.T) { func TestClone(t *testing.T) {
s := SetOf[int]([]int{1, 2, 3, 4, 4, 1}) s := Of(1, 2, 3, 4, 4, 1)
if s.Len() != 4 { if s.Len() != 4 {
t.Errorf("wrong len %d; want 4", s.Len()) t.Errorf("wrong len %d; want 4", s.Len())
} }
@ -122,8 +122,8 @@ func TestSetJSONRoundTrip(t *testing.T) {
}{ }{
{"empty", make(Set[string]), make(Set[int])}, {"empty", make(Set[string]), make(Set[int])},
{"nil", nil, nil}, {"nil", nil, nil},
{"one-item", SetOf([]string{"one"}), SetOf([]int{1})}, {"one-item", Of("one"), Of(1)},
{"multiple-items", SetOf([]string{"one", "two", "three"}), SetOf([]int{1, 2, 3})}, {"multiple-items", Of("one", "two", "three"), Of(1, 2, 3)},
} }
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.desc, func(t *testing.T) { t.Run(tt.desc, func(t *testing.T) {
@ -158,3 +158,12 @@ func TestSetJSONRoundTrip(t *testing.T) {
}) })
} }
} }
func TestMake(t *testing.T) {
var s Set[int]
s.Make()
s.Add(1)
if !s.Contains(1) {
t.Error("missing 1")
}
}

Loading…
Cancel
Save