💸 cache: new package

This commit adds a simple bounded LRU cache implementation.

Author: database64128

cache/cache.go

@@ -0,0 +1,212 @@
+package cache
+
+import (
+	"iter"
+	"math"
+)
+
+// Entry represents a key-value pair in the cache.
+type Entry[K comparable, V any] struct {
+	Key   K
+	Value V
+}
+
+type boundedNode[K comparable, V any] struct {
+	prev *boundedNode[K, V]
+	next *boundedNode[K, V]
+	Entry[K, V]
+}
+
+// BoundedCache is a simple in-memory cache with a fixed upper bound on the number of entries.
+// It uses a doubly linked list to maintain the order of access, where the tail is the most recently used node.
+// When the cache reaches its capacity, insertions will evict the least recently used node (the head of the list).
+type BoundedCache[K comparable, V any] struct {
+	nodeByKey map[K]*boundedNode[K, V]
+	capacity  int
+
+	// head is the least recently used node.
+	head *boundedNode[K, V]
+	// tail is the most recently used node.
+	tail *boundedNode[K, V]
+}
+
+// NewBoundedCache returns a new bounded cache with the given capacity.
+// If capacity is not positive, the cache will be effectively unbounded.
+func NewBoundedCache[K comparable, V any](capacity int) *BoundedCache[K, V] {
+	if capacity <= 0 {
+		capacity = math.MaxInt
+	}
+	return &BoundedCache[K, V]{
+		nodeByKey: make(map[K]*boundedNode[K, V]),
+		capacity:  capacity,
+	}
+}
+
+// Len returns the number of entries in the cache.
+func (c *BoundedCache[K, V]) Len() int {
+	return len(c.nodeByKey)
+}
+
+// Capacity returns the maximum number of entries the cache can hold.
+func (c *BoundedCache[K, V]) Capacity() int {
+	return c.capacity
+}
+
+// Contains returns whether the cache contains the given key.
+//
+// Unlike Get, this method does not update the access order of the cache.
+func (c *BoundedCache[K, V]) Contains(key K) bool {
+	_, ok := c.nodeByKey[key]
+	return ok
+}
+
+// Get returns the value associated with key.
+func (c *BoundedCache[K, V]) Get(key K) (value V, ok bool) {
+	node, ok := c.nodeByKey[key]
+	if !ok {
+		return value, false
+	}
+	c.moveToTail(node)
+	return node.Value, true
+}
+
+// GetEntry returns the entry associated with key.
+func (c *BoundedCache[K, V]) GetEntry(key K) (entry *Entry[K, V], ok bool) {
+	node, ok := c.nodeByKey[key]
+	if !ok {
+		return nil, false
+	}
+	c.moveToTail(node)
+	return &node.Entry, true
+}
+
+// Set inserts or updates the value associated with key.
+func (c *BoundedCache[K, V]) Set(key K, value V) {
+	node, ok := c.nodeByKey[key]
+	if !ok {
+		c.insert(key, value)
+		return
+	}
+	node.Value = value
+	c.moveToTail(node)
+}
+
+// Insert adds a new key-value pair to the cache if the key does not already exist.
+// It returns true if the insertion was successful, false if the key already exists.
+func (c *BoundedCache[K, V]) Insert(key K, value V) bool {
+	if _, ok := c.nodeByKey[key]; ok {
+		return false
+	}
+	c.insert(key, value)
+	return true
+}
+
+// InsertUnchecked adds a new key-value pair to the cache without checking if the key already exists.
+//
+// WARNING: This is undefined behavior if the key already exists in the cache.
+func (c *BoundedCache[K, V]) InsertUnchecked(key K, value V) {
+	c.insert(key, value)
+}
+
+// Remove deletes the value associated with key and returns whether the key was found.
+func (c *BoundedCache[K, V]) Remove(key K) bool {
+	node, ok := c.nodeByKey[key]
+	if !ok {
+		return false
+	}
+	c.remove(node)
+	return true
+}
+
+// Clear removes all entries from the cache.
+func (c *BoundedCache[K, V]) Clear() {
+	clear(c.nodeByKey)
+	c.head = nil
+	c.tail = nil
+}
+
+// All returns an iterator over all entries in the cache,
+// starting from the least recently used (head) to the most recently used (tail).
+func (c *BoundedCache[K, V]) All() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		for node := c.head; node != nil; node = node.next {
+			if !yield(node.Key, node.Value) {
+				break
+			}
+		}
+	}
+}
+
+// Backward returns an iterator over all entries in the cache,
+// starting from the most recently used (tail) to the least recently used (head).
+func (c *BoundedCache[K, V]) Backward() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		for node := c.tail; node != nil; node = node.prev {
+			if !yield(node.Key, node.Value) {
+				break
+			}
+		}
+	}
+}
+
+// insert inserts the key-value pair as a new node at the tail of the list.
+func (c *BoundedCache[K, V]) insert(key K, value V) {
+	// If the cache is at capacity, remove the least recently used item.
+	if len(c.nodeByKey) == c.capacity {
+		c.remove(c.head)
+	}
+
+	node := &boundedNode[K, V]{
+		prev:  c.tail,
+		Entry: Entry[K, V]{Key: key, Value: value},
+	}
+
+	c.nodeByKey[key] = node
+
+	if c.tail != nil {
+		c.tail.next = node
+	} else {
+		c.head = node
+	}
+	c.tail = node
+}
+
+// moveToTail promotes the given node to the tail of the list,
+// indicating that it was recently accessed.
+func (c *BoundedCache[K, V]) moveToTail(node *boundedNode[K, V]) {
+	// Check if the node is already at the tail.
+	if node.next == nil {
+		return
+	}
+
+	// Detach from the list.
+	node.next.prev = node.prev
+	if node.prev != nil {
+		node.prev.next = node.next
+	} else {
+		c.head = node.next
+	}
+
+	// Attach to the tail.
+	node.prev = c.tail
+	node.next = nil
+	c.tail.next = node
+	c.tail = node
+}
+
+// remove deletes the given node from the cache.
+func (c *BoundedCache[K, V]) remove(node *boundedNode[K, V]) {
+	delete(c.nodeByKey, node.Key)
+
+	if node.prev != nil {
+		node.prev.next = node.next
+	} else {
+		c.head = node.next
+	}
+
+	if node.next != nil {
+		node.next.prev = node.prev
+	} else {
+		c.tail = node.prev
+	}
+}

cache/cache_test.go

@@ -0,0 +1,138 @@
+package cache_test
+
+import (
+	"math"
+	"slices"
+	"testing"
+
+	"github.com/database64128/shadowsocks-go/cache"
+)
+
+func TestBoundedCache(t *testing.T) {
+	c := cache.NewBoundedCache[int, int](3)
+	assertBoundedCacheLenCapacityContent(t, c, nil, 3)
+	c.Set(1, -1)
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{1, -1}}, 3)
+	if !c.Insert(2, -2) {
+		t.Error("c.Insert(2, -2) = false, want true")
+	}
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{1, -1}, {2, -2}}, 3)
+	c.InsertUnchecked(3, -3)
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{1, -1}, {2, -2}, {3, -3}}, 3)
+	c.Set(4, -4)
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{2, -2}, {3, -3}, {4, -4}}, 3)
+	if value, ok := c.Get(2); value != -2 || !ok {
+		t.Errorf("c.Get(2) = %d, %v, want -2, true", value, ok)
+	}
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{3, -3}, {4, -4}, {2, -2}}, 3)
+	if entry, ok := c.GetEntry(4); entry == nil || entry.Key != 4 || entry.Value != -4 || !ok {
+		t.Errorf("c.GetEntry(4) = %v, %v, want {Key: 4, Value: -4}, true", entry, ok)
+	}
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{3, -3}, {2, -2}, {4, -4}}, 3)
+	if !c.Remove(4) {
+		t.Error("c.Remove(4) = false, want true")
+	}
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{3, -3}, {2, -2}}, 3)
+	c.Set(2, 2)
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{3, -3}, {2, 2}}, 3)
+	if value, ok := c.Get(3); value != -3 || !ok {
+		t.Errorf("c.Get(3) = %d, %v, want -3, true", value, ok)
+	}
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{{2, 2}, {3, -3}}, 3)
+	c.Clear()
+	assertBoundedCacheLenCapacityContent(t, c, nil, 3)
+}
+
+func TestBoundedCacheUnboundedCapacity(t *testing.T) {
+	c := cache.NewBoundedCache[int, int](0)
+	assertBoundedCacheLenCapacityContent(t, c, nil, math.MaxInt)
+	c.Set(1, -1)
+	c.Set(2, -2)
+	c.Set(3, -3)
+	for range c.All() {
+		for range c.Backward() {
+			break
+		}
+		break
+	}
+	c.Set(4, -4)
+	c.Set(5, -5)
+	c.Set(6, -6)
+	assertBoundedCacheLenCapacityContent(t, c, []cache.Entry[int, int]{
+		{1, -1}, {2, -2}, {3, -3}, {4, -4}, {5, -5}, {6, -6},
+	}, math.MaxInt)
+}
+
+func assertBoundedCacheLenCapacityContent(t *testing.T, c *cache.BoundedCache[int, int], want []cache.Entry[int, int], expectedCapacity int) {
+	t.Helper()
+
+	if got := c.Len(); got != len(want) {
+		t.Errorf("c.Len() = %d, want %d", got, len(want))
+	}
+	if got := c.Capacity(); got != expectedCapacity {
+		t.Errorf("c.Capacity() = %d, want %d", got, expectedCapacity)
+	}
+
+	got := make([]cache.Entry[int, int], 0, len(want))
+	for key, value := range c.All() {
+		got = append(got, cache.Entry[int, int]{Key: key, Value: value})
+	}
+	if !slices.Equal(got, want) {
+		t.Errorf("c.All() = %v, want %v", got, want)
+	}
+
+	got = got[:0]
+	for key, value := range c.Backward() {
+		got = append(got, cache.Entry[int, int]{Key: key, Value: value})
+	}
+	if !slicesReverseEqual(got, want) {
+		t.Errorf("c.Backward() = %v, want %v", got, want)
+	}
+
+	for key := range 10 {
+		if index := slices.IndexFunc(want, func(e cache.Entry[int, int]) bool {
+			return e.Key == key
+		}); index != -1 {
+			expectedEntry := want[index]
+			expectedValue := expectedEntry.Value
+
+			if !c.Contains(key) {
+				t.Errorf("c.Contains(%d) = false, want true", key)
+			}
+
+			if c.Insert(key, expectedValue) {
+				t.Errorf("c.Insert(%d, %d) = true, want false", key, expectedValue)
+			}
+		} else {
+			if c.Contains(key) {
+				t.Errorf("c.Contains(%d) = true, want false", key)
+			}
+
+			value, ok := c.Get(key)
+			if value != 0 || ok {
+				t.Errorf("c.Get(%d) = %d, %v, want 0, false", key, value, ok)
+			}
+
+			entry, ok := c.GetEntry(key)
+			if entry != nil || ok {
+				t.Errorf("c.GetEntry(%d) = %v, %v, want nil, false", key, entry, ok)
+			}
+
+			if c.Remove(key) {
+				t.Errorf("c.Remove(%d) = true, want false", key)
+			}
+		}
+	}
+}
+
+func slicesReverseEqual[S ~[]E, E comparable](s1, s2 S) bool {
+	if len(s1) != len(s2) {
+		return false
+	}
+	for i := range s1 {
+		if s1[i] != s2[len(s2)-1-i] {
+			return false
+		}
+	}
+	return true
+}