-
Notifications
You must be signed in to change notification settings - Fork 660
/
set.go
190 lines (167 loc) · 3.78 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
// Copyright (C) 2019-2022, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package ids
import (
"encoding/json"
"strings"
)
const (
// The minimum capacity of a set
minSetSize = 16
// If a set has more than this many keys, it will be cleared by setting the map to nil
// rather than iteratively deleting
clearSizeThreshold = 512
)
// Set is a set of IDs
type Set map[ID]struct{}
// Return a new set with initial capacity [size].
// More or less than [size] elements can be added to this set.
// Using NewSet() rather than ids.Set{} is just an optimization that can
// be used if you know how many elements will be put in this set.
func NewSet(size int) Set {
if size < 0 {
return Set{}
}
return make(map[ID]struct{}, size)
}
func (ids *Set) init(size int) {
if *ids == nil {
if minSetSize > size {
size = minSetSize
}
*ids = make(map[ID]struct{}, size)
}
}
// Add all the ids to this set, if the id is already in the set, nothing happens
func (ids *Set) Add(idList ...ID) {
ids.init(2 * len(idList))
for _, id := range idList {
(*ids)[id] = struct{}{}
}
}
// Union adds all the ids from the provided set to this set.
func (ids *Set) Union(set Set) {
ids.init(2 * set.Len())
for id := range set {
(*ids)[id] = struct{}{}
}
}
// Difference removes all the ids from the provided set to this set.
func (ids *Set) Difference(set Set) {
for id := range set {
delete(*ids, id)
}
}
// Contains returns true if the set contains this id, false otherwise
func (ids *Set) Contains(id ID) bool {
_, contains := (*ids)[id]
return contains
}
// Overlaps returns true if the intersection of the set is non-empty
func (ids *Set) Overlaps(big Set) bool {
small := *ids
if small.Len() > big.Len() {
small, big = big, small
}
for id := range small {
if _, ok := big[id]; ok {
return true
}
}
return false
}
// Len returns the number of ids in this set
func (ids Set) Len() int { return len(ids) }
// Remove all the id from this set, if the id isn't in the set, nothing happens
func (ids *Set) Remove(idList ...ID) {
for _, id := range idList {
delete(*ids, id)
}
}
// Clear empties this set
func (ids *Set) Clear() {
if len(*ids) > clearSizeThreshold {
*ids = nil
return
}
for key := range *ids {
delete(*ids, key)
}
}
// List converts this set into a list
func (ids Set) List() []ID {
idList := make([]ID, ids.Len())
i := 0
for id := range ids {
idList[i] = id
i++
}
return idList
}
// SortedList returns this set as a sorted list
func (ids Set) SortedList() []ID {
lst := ids.List()
SortIDs(lst)
return lst
}
// CappedList returns a list of length at most [size].
// Size should be >= 0. If size < 0, returns nil.
func (ids Set) CappedList(size int) []ID {
if size < 0 {
return nil
}
if l := ids.Len(); l < size {
size = l
}
i := 0
idList := make([]ID, size)
for id := range ids {
if i >= size {
break
}
idList[i] = id
i++
}
return idList
}
// Equals returns true if the sets contain the same elements
func (ids Set) Equals(oIDs Set) bool {
if ids.Len() != oIDs.Len() {
return false
}
for key := range oIDs {
if _, contains := ids[key]; !contains {
return false
}
}
return true
}
// String returns the string representation of a set
func (ids Set) String() string {
sb := strings.Builder{}
sb.WriteString("{")
first := true
for id := range ids {
if !first {
sb.WriteString(", ")
}
first = false
sb.WriteString(id.String())
}
sb.WriteString("}")
return sb.String()
}
// Removes and returns an element. If the set is empty, does nothing and returns
// false.
func (ids *Set) Pop() (ID, bool) {
for id := range *ids {
delete(*ids, id)
return id, true
}
return ID{}, false
}
func (ids *Set) MarshalJSON() ([]byte, error) {
idsList := ids.List()
SortIDs(idsList)
return json.Marshal(idsList)
}