-
Notifications
You must be signed in to change notification settings - Fork 661
/
node_id_set.go
159 lines (140 loc) · 3.49 KB
/
node_id_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
// Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package ids
import "strings"
// NodeIDSet is a set of NodeIDs
type NodeIDSet map[NodeID]struct{}
// Return a new NodeIDSet with initial capacity [size].
// More or less than [size] elements can be added to this set.
// Using NewNodeIDSet() rather than ids.NodeIDSet{} is just an optimization that can
// be used if you know how many elements will be put in this set.
func NewNodeIDSet(size int) NodeIDSet {
if size < 0 {
return NodeIDSet{}
}
return make(map[NodeID]struct{}, size)
}
func (ids *NodeIDSet) init(size int) {
if *ids == nil {
if minShortSetSize > size {
size = minShortSetSize
}
*ids = make(map[NodeID]struct{}, size)
}
}
// Add all the ids to this set, if the id is already in the set, nothing happens
func (ids *NodeIDSet) Add(idList ...NodeID) {
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 *NodeIDSet) Union(idSet NodeIDSet) {
ids.init(2 * idSet.Len())
for id := range idSet {
(*ids)[id] = struct{}{}
}
}
// Difference removes all the ids from the provided set to this set.
func (ids *NodeIDSet) Difference(idSet NodeIDSet) {
for id := range idSet {
delete(*ids, id)
}
}
// Contains returns true if the set contains this id, false otherwise
func (ids *NodeIDSet) Contains(id NodeID) bool {
ids.init(1)
_, contains := (*ids)[id]
return contains
}
// Len returns the number of ids in this set
func (ids NodeIDSet) Len() int { return len(ids) }
// Remove all the id from this set, if the id isn't in the set, nothing happens
func (ids *NodeIDSet) Remove(idList ...NodeID) {
ids.init(1)
for _, id := range idList {
delete(*ids, id)
}
}
// Clear empties this set
func (ids *NodeIDSet) Clear() { *ids = nil }
// CappedList returns a list of length at most [size].
// Size should be >= 0. If size < 0, returns nil.
func (ids NodeIDSet) CappedList(size int) []NodeID {
if size < 0 {
return nil
}
if l := ids.Len(); l < size {
size = l
}
i := 0
idList := make([]NodeID, size)
for id := range ids {
if i >= size {
break
}
idList[i] = id
i++
}
return idList
}
// List converts this set into a list
func (ids NodeIDSet) List() []NodeID {
idList := make([]NodeID, len(ids))
i := 0
for id := range ids {
idList[i] = id
i++
}
return idList
}
// SortedList returns this set as a sorted list
func (ids NodeIDSet) SortedList() []NodeID {
lst := ids.List()
SortNodeIDs(lst)
return lst
}
// Equals returns true if the sets contain the same elements
func (ids NodeIDSet) Equals(oIDs NodeIDSet) 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 NodeIDSet) 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()
}
// Returns an element. If the set is empty, returns false
func (ids *NodeIDSet) Peek() (NodeID, bool) {
for id := range *ids {
return id, true
}
return NodeID{}, false
}
// Removes and returns an element. If the set is empty, does nothing and returns
// false
func (ids *NodeIDSet) Pop() (NodeID, bool) {
for id := range *ids {
delete(*ids, id)
return id, true
}
return NodeID{}, false
}