-
-
Notifications
You must be signed in to change notification settings - Fork 400
/
set.go
227 lines (191 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
215
216
217
218
219
220
221
222
223
224
225
226
227
package set
import (
"errors"
"strings"
"sync"
)
// Returned when an added key already exists in the set.
var ErrCollision = errors.New("key already exists")
// Returned when a requested item does not exist in the set.
var ErrMissing = errors.New("item does not exist")
// Returned when a nil item is added. Nil values are considered expired and invalid.
var ErrNil = errors.New("item value must not be nil")
// ZeroValue can be used when we only care about the key, not about the value.
var ZeroValue = struct{}{}
// Interface is the Set interface
type Interface interface {
Clear() int
Each(fn IterFunc) error
// Add only if the item does not already exist
Add(item Item) error
// Set item, override if it already exists
Set(item Item) error
Get(key string) (Item, error)
In(key string) bool
Len() int
ListPrefix(prefix string) []Item
Remove(key string) error
Replace(oldKey string, item Item) error
}
type IterFunc func(key string, item Item) error
type Set struct {
sync.RWMutex
lookup map[string]Item
normalize func(string) string
}
// New creates a new set with case-insensitive keys
func New() *Set {
return &Set{
lookup: map[string]Item{},
normalize: normalize,
}
}
// Clear removes all items and returns the number removed.
func (s *Set) Clear() int {
s.Lock()
n := len(s.lookup)
s.lookup = map[string]Item{}
s.Unlock()
return n
}
// Len returns the size of the set right now.
func (s *Set) Len() int {
s.RLock()
defer s.RUnlock()
return len(s.lookup)
}
// In checks if an item exists in this set.
func (s *Set) In(key string) bool {
key = s.normalize(key)
s.RLock()
item, ok := s.lookup[key]
s.RUnlock()
if ok && item.Value() == nil {
s.cleanup(key)
ok = false
}
return ok
}
// Get returns an item with the given key.
func (s *Set) Get(key string) (Item, error) {
key = s.normalize(key)
s.RLock()
item, ok := s.lookup[key]
s.RUnlock()
if !ok {
return nil, ErrMissing
}
if item.Value() == nil {
s.cleanup(key)
}
return item, nil
}
// Remove potentially expired key (normalized).
func (s *Set) cleanup(key string) {
s.Lock()
item, ok := s.lookup[key]
if ok && item == nil {
delete(s.lookup, key)
}
s.Unlock()
}
// Add item to this set if it does not exist already.
func (s *Set) Add(item Item) error {
if item.Value() == nil {
return ErrNil
}
key := s.normalize(item.Key())
s.Lock()
defer s.Unlock()
oldItem, found := s.lookup[key]
if found && oldItem.Value() != nil {
return ErrCollision
}
s.lookup[key] = item
return nil
}
// Set item to this set, even if it already exists.
func (s *Set) Set(item Item) error {
if item.Value() == nil {
return ErrNil
}
key := s.normalize(item.Key())
s.Lock()
defer s.Unlock()
s.lookup[key] = item
return nil
}
// Remove item from this set.
func (s *Set) Remove(key string) error {
key = s.normalize(key)
s.Lock()
defer s.Unlock()
_, found := s.lookup[key]
if !found {
return ErrMissing
}
delete(s.lookup, key)
return nil
}
// Replace oldKey with a new item, which might be a new key.
// Can be used to rename items.
func (s *Set) Replace(oldKey string, item Item) error {
if item.Value() == nil {
return ErrNil
}
newKey := s.normalize(item.Key())
oldKey = s.normalize(oldKey)
s.Lock()
defer s.Unlock()
if newKey != oldKey {
// Check if it already exists
_, found := s.lookup[newKey]
if found {
return ErrCollision
}
// Remove oldKey
_, found = s.lookup[oldKey]
if !found {
return ErrMissing
}
delete(s.lookup, oldKey)
}
// Add new item
s.lookup[newKey] = item
return nil
}
// Each loops over every item while holding a read lock and applies fn to each
// element.
func (s *Set) Each(fn IterFunc) error {
var err error
s.RLock()
for key, item := range s.lookup {
if item.Value() == nil {
// Expired
defer s.cleanup(key)
continue
}
if err = fn(key, item); err != nil {
// Abort early
break
}
}
s.RUnlock()
return err
}
// ListPrefix returns a list of items with a prefix, normalized.
// TODO: Add trie for efficient prefix lookup
func (s *Set) ListPrefix(prefix string) []Item {
r := []Item{}
prefix = s.normalize(prefix)
s.Each(func(key string, item Item) error {
if strings.HasPrefix(key, prefix) {
r = append(r, item)
}
return nil
})
return r
}
func normalize(key string) string {
return strings.ToLower(key)
}