-
Notifications
You must be signed in to change notification settings - Fork 664
/
set.go
214 lines (185 loc) · 4.27 KB
/
set.go
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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package set
import (
"bytes"
stdjson "encoding/json"
"golang.org/x/exp/maps"
"github.com/ava-labs/avalanchego/utils"
"github.com/ava-labs/avalanchego/utils/json"
"github.com/ava-labs/avalanchego/utils/wrappers"
)
// The minimum capacity of a set
const minSetSize = 16
var _ stdjson.Marshaler = (*Set[int])(nil)
// Set is a set of elements.
type Set[T comparable] map[T]struct{}
// Of returns a Set initialized with [elts]
func Of[T comparable](elts ...T) Set[T] {
s := NewSet[T](len(elts))
s.Add(elts...)
return s
}
// Return a new set with initial capacity [size].
// More or less than [size] elements can be added to this set.
// Using NewSet() rather than Set[T]{} is just an optimization that can
// be used if you know how many elements will be put in this set.
func NewSet[T comparable](size int) Set[T] {
if size < 0 {
return Set[T]{}
}
return make(map[T]struct{}, size)
}
func (s *Set[T]) resize(size int) {
if *s == nil {
if minSetSize > size {
size = minSetSize
}
*s = make(map[T]struct{}, size)
}
}
// Add all the elements to this set.
// If the element is already in the set, nothing happens.
func (s *Set[T]) Add(elts ...T) {
s.resize(2 * len(elts))
for _, elt := range elts {
(*s)[elt] = struct{}{}
}
}
// Union adds all the elements from the provided set to this set.
func (s *Set[T]) Union(set Set[T]) {
s.resize(2 * set.Len())
for elt := range set {
(*s)[elt] = struct{}{}
}
}
// Difference removes all the elements in [set] from [s].
func (s *Set[T]) Difference(set Set[T]) {
for elt := range set {
delete(*s, elt)
}
}
// Contains returns true iff the set contains this element.
func (s *Set[T]) Contains(elt T) bool {
_, contains := (*s)[elt]
return contains
}
// Overlaps returns true if the intersection of the set is non-empty
func (s *Set[T]) Overlaps(big Set[T]) bool {
small := *s
if small.Len() > big.Len() {
small, big = big, small
}
for elt := range small {
if _, ok := big[elt]; ok {
return true
}
}
return false
}
// Len returns the number of elements in this set.
func (s Set[_]) Len() int {
return len(s)
}
// Remove all the given elements from this set.
// If an element isn't in the set, it's ignored.
func (s *Set[T]) Remove(elts ...T) {
for _, elt := range elts {
delete(*s, elt)
}
}
// Clear empties this set
func (s *Set[_]) Clear() {
maps.Clear(*s)
}
// List converts this set into a list
func (s Set[T]) List() []T {
return maps.Keys(s)
}
// CappedList returns a list of length at most [size].
// Size should be >= 0. If size < 0, returns nil.
func (s Set[T]) CappedList(size int) []T {
if size < 0 {
return nil
}
if l := s.Len(); l < size {
size = l
}
i := 0
elts := make([]T, size)
for elt := range s {
if i >= size {
break
}
elts[i] = elt
i++
}
return elts
}
// Equals returns true if the sets contain the same elements
func (s Set[T]) Equals(other Set[T]) bool {
return maps.Equal(s, other)
}
// Removes and returns an element.
// If the set is empty, does nothing and returns false.
func (s *Set[T]) Pop() (T, bool) {
for elt := range *s {
delete(*s, elt)
return elt, true
}
return utils.Zero[T](), false
}
func (s *Set[T]) UnmarshalJSON(b []byte) error {
str := string(b)
if str == json.Null {
return nil
}
var elts []T
if err := stdjson.Unmarshal(b, &elts); err != nil {
return err
}
s.Clear()
s.Add(elts...)
return nil
}
func (s *Set[_]) MarshalJSON() ([]byte, error) {
var (
eltBytes = make([][]byte, len(*s))
i int
err error
)
for elt := range *s {
eltBytes[i], err = stdjson.Marshal(elt)
if err != nil {
return nil, err
}
i++
}
// Sort for determinism
utils.SortBytes(eltBytes)
// Build the JSON
var (
jsonBuf = bytes.Buffer{}
errs = wrappers.Errs{}
)
_, err = jsonBuf.WriteString("[")
errs.Add(err)
for i, elt := range eltBytes {
_, err := jsonBuf.Write(elt)
errs.Add(err)
if i != len(eltBytes)-1 {
_, err := jsonBuf.WriteString(",")
errs.Add(err)
}
}
_, err = jsonBuf.WriteString("]")
errs.Add(err)
return jsonBuf.Bytes(), errs.Err
}
// Returns an element. If the set is empty, returns false
func (s *Set[T]) Peek() (T, bool) {
for elt := range *s {
return elt, true
}
return utils.Zero[T](), false
}