| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296 |
- // Package orderedmap implements an ordered map, i.e. a map that also keeps track of
- // the order in which keys were inserted.
- //
- // All operations are constant-time.
- //
- // Github repo: https://github.com/wk8/go-ordered-map
- //
- package orderedmap
- import (
- "fmt"
- list "github.com/bahlo/generic-list-go"
- )
- type Pair[K comparable, V any] struct {
- Key K
- Value V
- element *list.Element[*Pair[K, V]]
- }
- type OrderedMap[K comparable, V any] struct {
- pairs map[K]*Pair[K, V]
- list *list.List[*Pair[K, V]]
- }
- type initConfig[K comparable, V any] struct {
- capacity int
- initialData []Pair[K, V]
- }
- type InitOption[K comparable, V any] func(config *initConfig[K, V])
- // WithCapacity allows giving a capacity hint for the map, akin to the standard make(map[K]V, capacity).
- func WithCapacity[K comparable, V any](capacity int) InitOption[K, V] {
- return func(c *initConfig[K, V]) {
- c.capacity = capacity
- }
- }
- // WithInitialData allows passing in initial data for the map.
- func WithInitialData[K comparable, V any](initialData ...Pair[K, V]) InitOption[K, V] {
- return func(c *initConfig[K, V]) {
- c.initialData = initialData
- if c.capacity < len(initialData) {
- c.capacity = len(initialData)
- }
- }
- }
- // New creates a new OrderedMap.
- // options can either be one or several InitOption[K, V], or a single integer,
- // which is then interpreted as a capacity hint, à la make(map[K]V, capacity).
- func New[K comparable, V any](options ...any) *OrderedMap[K, V] { //nolint:varnamelen
- orderedMap := &OrderedMap[K, V]{}
- var config initConfig[K, V]
- for _, untypedOption := range options {
- switch option := untypedOption.(type) {
- case int:
- if len(options) != 1 {
- invalidOption()
- }
- config.capacity = option
- case InitOption[K, V]:
- option(&config)
- default:
- invalidOption()
- }
- }
- orderedMap.initialize(config.capacity)
- orderedMap.AddPairs(config.initialData...)
- return orderedMap
- }
- 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
- func invalidOption() { panic(invalidOptionMessage) }
- func (om *OrderedMap[K, V]) initialize(capacity int) {
- om.pairs = make(map[K]*Pair[K, V], capacity)
- om.list = list.New[*Pair[K, V]]()
- }
- // Get looks for the given key, and returns the value associated with it,
- // or V's nil value if not found. The boolean it returns says whether the key is present in the map.
- func (om *OrderedMap[K, V]) Get(key K) (val V, present bool) {
- if pair, present := om.pairs[key]; present {
- return pair.Value, true
- }
- return
- }
- // Load is an alias for Get, mostly to present an API similar to `sync.Map`'s.
- func (om *OrderedMap[K, V]) Load(key K) (V, bool) {
- return om.Get(key)
- }
- // Value returns the value associated with the given key or the zero value.
- func (om *OrderedMap[K, V]) Value(key K) (val V) {
- if pair, present := om.pairs[key]; present {
- val = pair.Value
- }
- return
- }
- // GetPair looks for the given key, and returns the pair associated with it,
- // or nil if not found. The Pair struct can then be used to iterate over the ordered map
- // from that point, either forward or backward.
- func (om *OrderedMap[K, V]) GetPair(key K) *Pair[K, V] {
- return om.pairs[key]
- }
- // Set sets the key-value pair, and returns what `Get` would have returned
- // on that key prior to the call to `Set`.
- func (om *OrderedMap[K, V]) Set(key K, value V) (val V, present bool) {
- if pair, present := om.pairs[key]; present {
- oldValue := pair.Value
- pair.Value = value
- return oldValue, true
- }
- pair := &Pair[K, V]{
- Key: key,
- Value: value,
- }
- pair.element = om.list.PushBack(pair)
- om.pairs[key] = pair
- return
- }
- // AddPairs allows setting multiple pairs at a time. It's equivalent to calling
- // Set on each pair sequentially.
- func (om *OrderedMap[K, V]) AddPairs(pairs ...Pair[K, V]) {
- for _, pair := range pairs {
- om.Set(pair.Key, pair.Value)
- }
- }
- // Store is an alias for Set, mostly to present an API similar to `sync.Map`'s.
- func (om *OrderedMap[K, V]) Store(key K, value V) (V, bool) {
- return om.Set(key, value)
- }
- // Delete removes the key-value pair, and returns what `Get` would have returned
- // on that key prior to the call to `Delete`.
- func (om *OrderedMap[K, V]) Delete(key K) (val V, present bool) {
- if pair, present := om.pairs[key]; present {
- om.list.Remove(pair.element)
- delete(om.pairs, key)
- return pair.Value, true
- }
- return
- }
- // Len returns the length of the ordered map.
- func (om *OrderedMap[K, V]) Len() int {
- if om == nil || om.pairs == nil {
- return 0
- }
- return len(om.pairs)
- }
- // Oldest returns a pointer to the oldest pair. It's meant to be used to iterate on the ordered map's
- // pairs from the oldest to the newest, e.g.:
- // for pair := orderedMap.Oldest(); pair != nil; pair = pair.Next() { fmt.Printf("%v => %v\n", pair.Key, pair.Value) }
- func (om *OrderedMap[K, V]) Oldest() *Pair[K, V] {
- if om == nil || om.list == nil {
- return nil
- }
- return listElementToPair(om.list.Front())
- }
- // Newest returns a pointer to the newest pair. It's meant to be used to iterate on the ordered map's
- // pairs from the newest to the oldest, e.g.:
- // for pair := orderedMap.Oldest(); pair != nil; pair = pair.Next() { fmt.Printf("%v => %v\n", pair.Key, pair.Value) }
- func (om *OrderedMap[K, V]) Newest() *Pair[K, V] {
- if om == nil || om.list == nil {
- return nil
- }
- return listElementToPair(om.list.Back())
- }
- // Next returns a pointer to the next pair.
- func (p *Pair[K, V]) Next() *Pair[K, V] {
- return listElementToPair(p.element.Next())
- }
- // Prev returns a pointer to the previous pair.
- func (p *Pair[K, V]) Prev() *Pair[K, V] {
- return listElementToPair(p.element.Prev())
- }
- func listElementToPair[K comparable, V any](element *list.Element[*Pair[K, V]]) *Pair[K, V] {
- if element == nil {
- return nil
- }
- return element.Value
- }
- // KeyNotFoundError may be returned by functions in this package when they're called with keys that are not present
- // in the map.
- type KeyNotFoundError[K comparable] struct {
- MissingKey K
- }
- func (e *KeyNotFoundError[K]) Error() string {
- return fmt.Sprintf("missing key: %v", e.MissingKey)
- }
- // MoveAfter moves the value associated with key to its new position after the one associated with markKey.
- // Returns an error iff key or markKey are not present in the map. If an error is returned,
- // it will be a KeyNotFoundError.
- func (om *OrderedMap[K, V]) MoveAfter(key, markKey K) error {
- elements, err := om.getElements(key, markKey)
- if err != nil {
- return err
- }
- om.list.MoveAfter(elements[0], elements[1])
- return nil
- }
- // MoveBefore moves the value associated with key to its new position before the one associated with markKey.
- // Returns an error iff key or markKey are not present in the map. If an error is returned,
- // it will be a KeyNotFoundError.
- func (om *OrderedMap[K, V]) MoveBefore(key, markKey K) error {
- elements, err := om.getElements(key, markKey)
- if err != nil {
- return err
- }
- om.list.MoveBefore(elements[0], elements[1])
- return nil
- }
- func (om *OrderedMap[K, V]) getElements(keys ...K) ([]*list.Element[*Pair[K, V]], error) {
- elements := make([]*list.Element[*Pair[K, V]], len(keys))
- for i, k := range keys {
- pair, present := om.pairs[k]
- if !present {
- return nil, &KeyNotFoundError[K]{k}
- }
- elements[i] = pair.element
- }
- return elements, nil
- }
- // MoveToBack moves the value associated with key to the back of the ordered map,
- // i.e. makes it the newest pair in the map.
- // Returns an error iff key is not present in the map. If an error is returned,
- // it will be a KeyNotFoundError.
- func (om *OrderedMap[K, V]) MoveToBack(key K) error {
- _, err := om.GetAndMoveToBack(key)
- return err
- }
- // MoveToFront moves the value associated with key to the front of the ordered map,
- // i.e. makes it the oldest pair in the map.
- // Returns an error iff key is not present in the map. If an error is returned,
- // it will be a KeyNotFoundError.
- func (om *OrderedMap[K, V]) MoveToFront(key K) error {
- _, err := om.GetAndMoveToFront(key)
- return err
- }
- // GetAndMoveToBack combines Get and MoveToBack in the same call. If an error is returned,
- // it will be a KeyNotFoundError.
- func (om *OrderedMap[K, V]) GetAndMoveToBack(key K) (val V, err error) {
- if pair, present := om.pairs[key]; present {
- val = pair.Value
- om.list.MoveToBack(pair.element)
- } else {
- err = &KeyNotFoundError[K]{key}
- }
- return
- }
- // GetAndMoveToFront combines Get and MoveToFront in the same call. If an error is returned,
- // it will be a KeyNotFoundError.
- func (om *OrderedMap[K, V]) GetAndMoveToFront(key K) (val V, err error) {
- if pair, present := om.pairs[key]; present {
- val = pair.Value
- om.list.MoveToFront(pair.element)
- } else {
- err = &KeyNotFoundError[K]{key}
- }
- return
- }
|