1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
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
}
}
|
