LRU cache using Go 1.18 generics

lrucache.go

package main

type lruNode[K comparable, V any] struct {
	prev  *lruNode[K, V]
	next  *lruNode[K, V]
	key   K
	value V
}

type LRUCache[K comparable, V any] struct {
	capacity int
	list     *lruNode[K, V]
	dict     map[K]*lruNode[K, V]
}

func (c *LRUCache[K, V]) removeNode(node *lruNode[K, V]) {
	node.prev.next = node.next
	node.next.prev = node.prev
	node.next = nil
	node.prev = nil
}

func (c *LRUCache[K, V]) addNode(node *lruNode[K, V]) {
	node.next = c.list.next
	node.prev = c.list
	node.prev.next = node
	node.next.prev = node
}

func (c *LRUCache[K, V]) Put(key K, value V) {
	if len(c.dict) == c.capacity {
		evict := c.list.prev
		c.removeNode(evict)
		delete(c.dict, evict.key)
	}

	node := &lruNode[K, V]{
		key:   key,
		value: value,
	}
	c.dict[key] = node
	c.addNode(node)
}

func (c *LRUCache[K, V]) Get(key K) *V {
	node := c.dict[key]
	if node == nil {
		return nil
	}
	c.removeNode(node)
	c.addNode(node)
	return &node.value
}

func New[K comparable, V any](capacity int) *LRUCache[K, V] {
	sentinel := &lruNode[K, V]{}
	sentinel.prev = sentinel
	sentinel.next = sentinel
	return &LRUCache[K, V]{
		capacity: capacity,
		list:     sentinel,
		dict:     map[K]*lruNode[K, V]{},
	}
}

func assert(cond bool) {
	if !cond {
		panic("assertion failed")
	}
}

func main() {
	x := New[string, int](1)
	assert(x.Get("wat") == nil)
	x.Put("foo", 1)
	assert(*x.Get("foo") == 1)
	x.Put("bar", 2)
	assert(*x.Get("bar") == 2)
	assert(x.Get("foo") == nil)

	y := New[string, int](2)
	y.Put("foo", 1)
	y.Put("bar", 2)
	y.Put("baz", 3)
	assert(y.Get("foo") == nil)
	assert(*y.Get("bar") == 2)
	assert(*y.Get("baz") == 3)
	y.Get("bar")
	y.Put("qux", 4)
	assert(y.Get("baz") == nil)
}