lru.go (3537B)
1 package maps 2 3 import ( 4 "fmt" 5 "reflect" 6 ) 7 8 // LRU provides an ordered hashmap implementation that keeps elements ordered according to last recently used (hence, LRU). 9 type LRUMap[K comparable, V any] struct { 10 ordered[K, V] 11 size int 12 } 13 14 // NewLRU returns a new instance of LRUMap with given initializing length and maximum capacity. 15 func NewLRU[K comparable, V any](len, cap int) *LRUMap[K, V] { 16 m := new(LRUMap[K, V]) 17 m.Init(len, cap) 18 return m 19 } 20 21 // Init will initialize this map with initial length and maximum capacity. 22 func (m *LRUMap[K, V]) Init(len, cap int) { 23 if cap <= 0 { 24 panic("lru cap must be greater than zero") 25 } else if m.pool != nil { 26 panic("lru map already initialized") 27 } 28 m.ordered.hmap = make(map[K]*elem[K, V], len) 29 m.ordered.pool = allocElems[K, V](len) 30 m.size = cap 31 } 32 33 // Get will fetch value for given key from map, in the process pushing it to the front of the map. Returns false if not found. 34 func (m *LRUMap[K, V]) Get(key K) (V, bool) { 35 if elem, ok := m.hmap[key]; ok { 36 // Ensure safe 37 m.write_check() 38 39 // Unlink elem from list 40 m.list.Unlink(elem) 41 42 // Push to front of list 43 m.list.PushFront(elem) 44 45 return elem.V, true 46 } 47 var z V // zero value 48 return z, false 49 } 50 51 // Add will add the given key-value pair to the map, pushing them to the front of the map. Returns false if already exists. Evicts old at maximum capacity. 52 func (m *LRUMap[K, V]) Add(key K, value V) bool { 53 return m.AddWithHook(key, value, nil) 54 } 55 56 // AddWithHook performs .Add() but passing any evicted entry to given hook function. 57 func (m *LRUMap[K, V]) AddWithHook(key K, value V, evict func(K, V)) bool { 58 // Ensure safe 59 m.write_check() 60 61 // Look for existing elem 62 elem, ok := m.hmap[key] 63 if ok { 64 return false 65 } 66 67 if m.list.len >= m.size { 68 // We're at capacity, sir! 69 // Pop current tail elem 70 elem = m.list.PopTail() 71 72 if evict != nil { 73 // Pass to evict hook 74 evict(elem.K, elem.V) 75 } 76 77 // Delete key from map 78 delete(m.hmap, elem.K) 79 } else { 80 // Allocate elem 81 elem = m.alloc() 82 } 83 84 // Set elem 85 elem.K = key 86 elem.V = value 87 88 // Add element map entry 89 m.hmap[key] = elem 90 91 // Push to front of list 92 m.list.PushFront(elem) 93 return true 94 } 95 96 // Set will ensure that given key-value pair exists in the map, by either adding new or updating existing, pushing them to the front of the map. Evicts old at maximum capacity. 97 func (m *LRUMap[K, V]) Set(key K, value V) { 98 m.SetWithHook(key, value, nil) 99 } 100 101 // SetWithHook performs .Set() but passing any evicted entry to given hook function. 102 func (m *LRUMap[K, V]) SetWithHook(key K, value V, evict func(K, V)) { 103 // Ensure safe 104 m.write_check() 105 106 // Look for existing elem 107 elem, ok := m.hmap[key] 108 109 if ok { 110 // Unlink elem from list 111 m.list.Unlink(elem) 112 113 // Update existing 114 elem.V = value 115 } else { 116 if m.list.len >= m.size { 117 // We're at capacity, sir! 118 // Pop current tail elem 119 elem = m.list.PopTail() 120 121 if evict != nil { 122 // Pass to evict hook 123 evict(elem.K, elem.V) 124 } 125 126 // Delete key from map 127 delete(m.hmap, elem.K) 128 } else { 129 // Allocate elem 130 elem = m.alloc() 131 } 132 133 // Set elem 134 elem.K = key 135 elem.V = value 136 137 // Add element map entry 138 m.hmap[key] = elem 139 } 140 141 // Push to front of list 142 m.list.PushFront(elem) 143 } 144 145 // Cap returns the maximum capacity of this LRU map. 146 func (m *LRUMap[K, V]) Cap() int { 147 return m.size 148 } 149 150 // Format implements fmt.Formatter, allowing performant string formatting of map. 151 func (m *LRUMap[K, V]) Format(state fmt.State, verb rune) { 152 m.format(reflect.TypeOf(m), state, verb) 153 }