learning category repo

internal/common/xslices/lru.go

@@ -0,0 +1,103 @@
+package xslices
+
+import (
+	"sync"
+)
+
+func NewLRU[K comparable, T any](capacity int) *lru[K, T] {
+	return &lru[K, T]{
+		items:    make(map[K]*lruNode[K, T], capacity),
+		capacity: capacity,
+	}
+}
+
+type lruNode[K comparable, T any] struct {
+	key   K
+	value T
+
+	next *lruNode[K, T]
+	prev *lruNode[K, T]
+}
+
+type lru[K comparable, T any] struct {
+	items    map[K]*lruNode[K, T]
+	length   int
+	capacity int
+
+	first *lruNode[K, T]
+	last  *lruNode[K, T]
+
+	mu sync.RWMutex
+}
+
+func (l *lru[K, T]) Push(key K, value T) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	node, ok := l.items[key]
+	if ok {
+		l.bumpUnsafe(key, node)
+		return
+	}
+
+	node = &lruNode[K, T]{
+		key:   key,
+		value: value,
+		next:  l.first,
+	}
+
+	if l.first != nil {
+		l.first.prev = node
+	}
+	if l.last == nil {
+		l.last = node
+	}
+
+	l.first = node
+	l.items[key] = node
+
+	if l.length == l.capacity && l.last != nil {
+		deletedNode := l.last
+		delete(l.items, deletedNode.key)
+
+		l.last = l.last.prev
+		return
+	}
+
+	l.length++
+}
+
+func (l *lru[K, T]) Get(key K) (T, bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	node, ok := l.items[key]
+	if !ok {
+		var t T
+		return t, false
+	}
+
+	out := node.value
+	l.bumpUnsafe(key, node)
+
+	return out, true
+}
+
+func (l *lru[K, T]) bumpUnsafe(key K, node *lruNode[K, T]) {
+	if l.first == node {
+		return
+	}
+
+	if node.next != nil {
+		node.next.prev = node.prev
+	}
+	if node.prev != nil {
+		node.prev.next = node.next
+	}
+
+	node.next = l.first
+	l.first.prev = node
+
+	l.first = node
+	node.prev = nil
+}

internal/common/xslices/lru_test.go

@@ -0,0 +1,23 @@
+package xslices
+
+import "testing"
+
+func TestLRU(t *testing.T) {
+	lru := NewLRU[int, string](3)
+	for i := 0; i < 4; i++ {
+		lru.Push(i, "v")
+	}
+
+	for i := 0; i < 4; i++ {
+		_, found := lru.Get(i)
+		if i == 0 {
+			if found {
+				t.Error("expected value to be flushed out of cache")
+			}
+			continue
+		}
+		if !found {
+			t.Errorf("expected value %d to be found", i)
+		}
+	}
+}