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