orderedmap.go 8.6 KB

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