1package set 2 3import ( 4 "errors" 5 "strings" 6 "sync" 7) 8 9// Returned when an added key already exists in the set. 10var ErrCollision = errors.New("key already exists") 11 12// Returned when a requested item does not exist in the set. 13var ErrMissing = errors.New("item does not exist") 14 15// Returned when a nil item is added. Nil values are considered expired and invalid. 16var ErrNil = errors.New("item value must not be nil") 17 18// ZeroValue can be used when we only care about the key, not about the value. 19var ZeroValue = struct{}{} 20 21// Interface is the Set interface 22type Interface interface { 23 Clear() int 24 Each(fn IterFunc) error 25 // Add only if the item does not already exist 26 Add(item Item) error 27 // Set item, override if it already exists 28 Set(item Item) error 29 Get(key string) (Item, error) 30 In(key string) bool 31 Len() int 32 ListPrefix(prefix string) []Item 33 Remove(key string) error 34 Replace(oldKey string, item Item) error 35} 36 37type IterFunc func(key string, item Item) error 38 39type Set struct { 40 sync.RWMutex 41 lookup map[string]Item 42 normalize func(string) string 43} 44 45// New creates a new set with case-insensitive keys 46func New() *Set { 47 return &Set{ 48 lookup: map[string]Item{}, 49 normalize: normalize, 50 } 51} 52 53// Clear removes all items and returns the number removed. 54func (s *Set) Clear() int { 55 s.Lock() 56 n := len(s.lookup) 57 s.lookup = map[string]Item{} 58 s.Unlock() 59 return n 60} 61 62// Len returns the size of the set right now. 63func (s *Set) Len() int { 64 s.RLock() 65 defer s.RUnlock() 66 return len(s.lookup) 67} 68 69// In checks if an item exists in this set. 70func (s *Set) In(key string) bool { 71 key = s.normalize(key) 72 s.RLock() 73 item, ok := s.lookup[key] 74 s.RUnlock() 75 if ok && item.Value() == nil { 76 s.cleanup(key) 77 ok = false 78 } 79 return ok 80} 81 82// Get returns an item with the given key. 83func (s *Set) Get(key string) (Item, error) { 84 key = s.normalize(key) 85 s.RLock() 86 item, ok := s.lookup[key] 87 s.RUnlock() 88 89 if !ok { 90 return nil, ErrMissing 91 } 92 if item.Value() == nil { 93 s.cleanup(key) 94 } 95 96 return item, nil 97} 98 99// Remove potentially expired key (normalized). 100func (s *Set) cleanup(key string) { 101 s.Lock() 102 item, ok := s.lookup[key] 103 if ok && item == nil { 104 delete(s.lookup, key) 105 } 106 s.Unlock() 107} 108 109// Add item to this set if it does not exist already. 110func (s *Set) Add(item Item) error { 111 if item.Value() == nil { 112 return ErrNil 113 } 114 key := s.normalize(item.Key()) 115 116 s.Lock() 117 defer s.Unlock() 118 119 oldItem, found := s.lookup[key] 120 if found && oldItem.Value() != nil { 121 return ErrCollision 122 } 123 124 s.lookup[key] = item 125 return nil 126} 127 128// Set item to this set, even if it already exists. 129func (s *Set) Set(item Item) error { 130 if item.Value() == nil { 131 return ErrNil 132 } 133 key := s.normalize(item.Key()) 134 135 s.Lock() 136 defer s.Unlock() 137 s.lookup[key] = item 138 return nil 139} 140 141// Remove item from this set. 142func (s *Set) Remove(key string) error { 143 key = s.normalize(key) 144 145 s.Lock() 146 defer s.Unlock() 147 148 _, found := s.lookup[key] 149 if !found { 150 return ErrMissing 151 } 152 delete(s.lookup, key) 153 return nil 154} 155 156// Replace oldKey with a new item, which might be a new key. 157// Can be used to rename items. 158func (s *Set) Replace(oldKey string, item Item) error { 159 if item.Value() == nil { 160 return ErrNil 161 } 162 newKey := s.normalize(item.Key()) 163 oldKey = s.normalize(oldKey) 164 165 s.Lock() 166 defer s.Unlock() 167 168 if newKey != oldKey { 169 // Check if it already exists 170 _, found := s.lookup[newKey] 171 if found { 172 return ErrCollision 173 } 174 175 // Remove oldKey 176 _, found = s.lookup[oldKey] 177 if !found { 178 return ErrMissing 179 } 180 delete(s.lookup, oldKey) 181 } 182 183 // Add new item 184 s.lookup[newKey] = item 185 186 return nil 187} 188 189// Each loops over every item while holding a read lock and applies fn to each 190// element. 191func (s *Set) Each(fn IterFunc) error { 192 var err error 193 s.RLock() 194 for key, item := range s.lookup { 195 if item.Value() == nil { 196 // Expired 197 defer s.cleanup(key) 198 continue 199 } 200 if err = fn(key, item); err != nil { 201 // Abort early 202 break 203 } 204 } 205 s.RUnlock() 206 return err 207} 208 209// ListPrefix returns a list of items with a prefix, normalized. 210// TODO: Add trie for efficient prefix lookup 211func (s *Set) ListPrefix(prefix string) []Item { 212 r := []Item{} 213 prefix = s.normalize(prefix) 214 215 s.Each(func(key string, item Item) error { 216 if strings.HasPrefix(key, prefix) { 217 r = append(r, item) 218 } 219 return nil 220 }) 221 222 return r 223} 224 225func normalize(key string) string { 226 return strings.ToLower(key) 227} 228