1// Package orderedmap implements an ordered map, i.e. a map that also keeps track of
2// the order in which keys were inserted.
3//
4// All operations are constant-time.
5//
6// Github repo: https://github.com/wk8/go-ordered-map
7//
8package orderedmap
9
10import (
11 "fmt"
12
13 list "github.com/bahlo/generic-list-go"
14)
15
16type Pair[K comparable, V any] struct {
17 Key K
18 Value V
19
20 element *list.Element[*Pair[K, V]]
21}
22
23type OrderedMap[K comparable, V any] struct {
24 pairs map[K]*Pair[K, V]
25 list *list.List[*Pair[K, V]]
26}
27
28type initConfig[K comparable, V any] struct {
29 capacity int
30 initialData []Pair[K, V]
31}
32
33type InitOption[K comparable, V any] func(config *initConfig[K, V])
34
35// WithCapacity allows giving a capacity hint for the map, akin to the standard make(map[K]V, capacity).
36func WithCapacity[K comparable, V any](capacity int) InitOption[K, V] {
37 return func(c *initConfig[K, V]) {
38 c.capacity = capacity
39 }
40}
41
42// WithInitialData allows passing in initial data for the map.
43func WithInitialData[K comparable, V any](initialData ...Pair[K, V]) InitOption[K, V] {
44 return func(c *initConfig[K, V]) {
45 c.initialData = initialData
46 if c.capacity < len(initialData) {
47 c.capacity = len(initialData)
48 }
49 }
50}
51
52// New creates a new OrderedMap.
53// options can either be one or several InitOption[K, V], or a single integer,
54// which is then interpreted as a capacity hint, à la make(map[K]V, capacity).
55func New[K comparable, V any](options ...any) *OrderedMap[K, V] { //nolint:varnamelen
56 orderedMap := &OrderedMap[K, V]{}
57
58 var config initConfig[K, V]
59 for _, untypedOption := range options {
60 switch option := untypedOption.(type) {
61 case int:
62 if len(options) != 1 {
63 invalidOption()
64 }
65 config.capacity = option
66
67 case InitOption[K, V]:
68 option(&config)
69
70 default:
71 invalidOption()
72 }
73 }
74
75 orderedMap.initialize(config.capacity)
76 orderedMap.AddPairs(config.initialData...)
77
78 return orderedMap
79}
80
81const invalidOptionMessage = `when using orderedmap.New[K,V]() with options, either provide one or several InitOption[K, V]; or a single integer which is then interpreted as a capacity hint, à la make(map[K]V, capacity).` //nolint:lll
82
83func invalidOption() { panic(invalidOptionMessage) }
84
85func (om *OrderedMap[K, V]) initialize(capacity int) {
86 om.pairs = make(map[K]*Pair[K, V], capacity)
87 om.list = list.New[*Pair[K, V]]()
88}
89
90// Get looks for the given key, and returns the value associated with it,
91// or V's nil value if not found. The boolean it returns says whether the key is present in the map.
92func (om *OrderedMap[K, V]) Get(key K) (val V, present bool) {
93 if pair, present := om.pairs[key]; present {
94 return pair.Value, true
95 }
96
97 return
98}
99
100// Load is an alias for Get, mostly to present an API similar to `sync.Map`'s.
101func (om *OrderedMap[K, V]) Load(key K) (V, bool) {
102 return om.Get(key)
103}
104
105// Value returns the value associated with the given key or the zero value.
106func (om *OrderedMap[K, V]) Value(key K) (val V) {
107 if pair, present := om.pairs[key]; present {
108 val = pair.Value
109 }
110 return
111}
112
113// GetPair looks for the given key, and returns the pair associated with it,
114// or nil if not found. The Pair struct can then be used to iterate over the ordered map
115// from that point, either forward or backward.
116func (om *OrderedMap[K, V]) GetPair(key K) *Pair[K, V] {
117 return om.pairs[key]
118}
119
120// Set sets the key-value pair, and returns what `Get` would have returned
121// on that key prior to the call to `Set`.
122func (om *OrderedMap[K, V]) Set(key K, value V) (val V, present bool) {
123 if pair, present := om.pairs[key]; present {
124 oldValue := pair.Value
125 pair.Value = value
126 return oldValue, true
127 }
128
129 pair := &Pair[K, V]{
130 Key: key,
131 Value: value,
132 }
133 pair.element = om.list.PushBack(pair)
134 om.pairs[key] = pair
135
136 return
137}
138
139// AddPairs allows setting multiple pairs at a time. It's equivalent to calling
140// Set on each pair sequentially.
141func (om *OrderedMap[K, V]) AddPairs(pairs ...Pair[K, V]) {
142 for _, pair := range pairs {
143 om.Set(pair.Key, pair.Value)
144 }
145}
146
147// Store is an alias for Set, mostly to present an API similar to `sync.Map`'s.
148func (om *OrderedMap[K, V]) Store(key K, value V) (V, bool) {
149 return om.Set(key, value)
150}
151
152// Delete removes the key-value pair, and returns what `Get` would have returned
153// on that key prior to the call to `Delete`.
154func (om *OrderedMap[K, V]) Delete(key K) (val V, present bool) {
155 if pair, present := om.pairs[key]; present {
156 om.list.Remove(pair.element)
157 delete(om.pairs, key)
158 return pair.Value, true
159 }
160 return
161}
162
163// Len returns the length of the ordered map.
164func (om *OrderedMap[K, V]) Len() int {
165 if om == nil || om.pairs == nil {
166 return 0
167 }
168 return len(om.pairs)
169}
170
171// Oldest returns a pointer to the oldest pair. It's meant to be used to iterate on the ordered map's
172// pairs from the oldest to the newest, e.g.:
173// for pair := orderedMap.Oldest(); pair != nil; pair = pair.Next() { fmt.Printf("%v => %v\n", pair.Key, pair.Value) }
174func (om *OrderedMap[K, V]) Oldest() *Pair[K, V] {
175 if om == nil || om.list == nil {
176 return nil
177 }
178 return listElementToPair(om.list.Front())
179}
180
181// Newest returns a pointer to the newest pair. It's meant to be used to iterate on the ordered map's
182// pairs from the newest to the oldest, e.g.:
183// for pair := orderedMap.Oldest(); pair != nil; pair = pair.Next() { fmt.Printf("%v => %v\n", pair.Key, pair.Value) }
184func (om *OrderedMap[K, V]) Newest() *Pair[K, V] {
185 if om == nil || om.list == nil {
186 return nil
187 }
188 return listElementToPair(om.list.Back())
189}
190
191// Next returns a pointer to the next pair.
192func (p *Pair[K, V]) Next() *Pair[K, V] {
193 return listElementToPair(p.element.Next())
194}
195
196// Prev returns a pointer to the previous pair.
197func (p *Pair[K, V]) Prev() *Pair[K, V] {
198 return listElementToPair(p.element.Prev())
199}
200
201func listElementToPair[K comparable, V any](element *list.Element[*Pair[K, V]]) *Pair[K, V] {
202 if element == nil {
203 return nil
204 }
205 return element.Value
206}
207
208// KeyNotFoundError may be returned by functions in this package when they're called with keys that are not present
209// in the map.
210type KeyNotFoundError[K comparable] struct {
211 MissingKey K
212}
213
214func (e *KeyNotFoundError[K]) Error() string {
215 return fmt.Sprintf("missing key: %v", e.MissingKey)
216}
217
218// MoveAfter moves the value associated with key to its new position after the one associated with markKey.
219// Returns an error iff key or markKey are not present in the map. If an error is returned,
220// it will be a KeyNotFoundError.
221func (om *OrderedMap[K, V]) MoveAfter(key, markKey K) error {
222 elements, err := om.getElements(key, markKey)
223 if err != nil {
224 return err
225 }
226 om.list.MoveAfter(elements[0], elements[1])
227 return nil
228}
229
230// MoveBefore moves the value associated with key to its new position before the one associated with markKey.
231// Returns an error iff key or markKey are not present in the map. If an error is returned,
232// it will be a KeyNotFoundError.
233func (om *OrderedMap[K, V]) MoveBefore(key, markKey K) error {
234 elements, err := om.getElements(key, markKey)
235 if err != nil {
236 return err
237 }
238 om.list.MoveBefore(elements[0], elements[1])
239 return nil
240}
241
242func (om *OrderedMap[K, V]) getElements(keys ...K) ([]*list.Element[*Pair[K, V]], error) {
243 elements := make([]*list.Element[*Pair[K, V]], len(keys))
244 for i, k := range keys {
245 pair, present := om.pairs[k]
246 if !present {
247 return nil, &KeyNotFoundError[K]{k}
248 }
249 elements[i] = pair.element
250 }
251 return elements, nil
252}
253
254// MoveToBack moves the value associated with key to the back of the ordered map,
255// i.e. makes it the newest pair in the map.
256// Returns an error iff key is not present in the map. If an error is returned,
257// it will be a KeyNotFoundError.
258func (om *OrderedMap[K, V]) MoveToBack(key K) error {
259 _, err := om.GetAndMoveToBack(key)
260 return err
261}
262
263// MoveToFront moves the value associated with key to the front of the ordered map,
264// i.e. makes it the oldest pair in the map.
265// Returns an error iff key is not present in the map. If an error is returned,
266// it will be a KeyNotFoundError.
267func (om *OrderedMap[K, V]) MoveToFront(key K) error {
268 _, err := om.GetAndMoveToFront(key)
269 return err
270}
271
272// GetAndMoveToBack combines Get and MoveToBack in the same call. If an error is returned,
273// it will be a KeyNotFoundError.
274func (om *OrderedMap[K, V]) GetAndMoveToBack(key K) (val V, err error) {
275 if pair, present := om.pairs[key]; present {
276 val = pair.Value
277 om.list.MoveToBack(pair.element)
278 } else {
279 err = &KeyNotFoundError[K]{key}
280 }
281
282 return
283}
284
285// GetAndMoveToFront combines Get and MoveToFront in the same call. If an error is returned,
286// it will be a KeyNotFoundError.
287func (om *OrderedMap[K, V]) GetAndMoveToFront(key K) (val V, err error) {
288 if pair, present := om.pairs[key]; present {
289 val = pair.Value
290 om.list.MoveToFront(pair.element)
291 } else {
292 err = &KeyNotFoundError[K]{key}
293 }
294
295 return
296}