Merge pull request #7 from Code-Hex/add/lfu

added lfu cache

Author: Code-Hex

README.md

@@ -10,6 +10,9 @@ go-generics-cache is an in-memory key:value store/cache that is suitable for app
   - See [examples](https://github.com/Code-Hex/go-generics-cache/blob/main/simple/example_test.go)
 - LRU cache
   - See [examples](https://github.com/Code-Hex/go-generics-cache/blob/main/lru/example_test.go)
+- LFU cache
+  - See [examples](https://github.com/Code-Hex/go-generics-cache/blob/main/lfu/example_test.go)
+  - [An O(1) algorithm for implementing the LFU cache eviction scheme](http://dhruvbird.com/lfu.pdf)
 
 ## Requirements
 

cache.go

@@ -16,10 +16,17 @@ type Cache[K comparable, V any] interface {
 
 // Item is an item
 type Item[K comparable, V any] struct {
-	Key        K
-	Value      V
-	Expiration time.Duration
-	CreatedAt  time.Time
+	Key            K
+	Value          V
+	ReferenceCount int
+	Expiration     time.Duration
+	CreatedAt      time.Time
+	ReferencedAt   time.Time
+}
+
+func (i *Item[K, V]) Referenced() {
+	i.ReferenceCount++
+	i.ReferencedAt = nowFunc()
 }
 
 var nowFunc = time.Now
@@ -53,11 +60,14 @@ func NewItem[K comparable, V any](key K, val V, opts ...ItemOption) *Item[K, V]
 	for _, optFunc := range opts {
 		optFunc(o)
 	}
+	now := nowFunc()
 	return &Item[K, V]{
-		Key:        key,
-		Value:      val,
-		Expiration: o.expiration,
-		CreatedAt:  nowFunc(),
+		Key:            key,
+		Value:          val,
+		ReferenceCount: 1,
+		Expiration:     o.expiration,
+		CreatedAt:      now,
+		ReferencedAt:   now,
 	}
 }
 

cache_test.go

@@ -101,3 +101,31 @@ func TestHasExpired(t *testing.T) {
 		})
 	}
 }
+
+func TestReferenced(t *testing.T) {
+	tm := time.Now()
+	reset := cache.SetNowFunc(tm)
+
+	item := cache.NewItem("hello", "world")
+	if item.ReferenceCount != 1 {
+		t.Errorf("want ref count is 1 but got %d", item.ReferenceCount)
+	}
+	if !item.ReferencedAt.Equal(tm) {
+		t.Errorf("unexpected referenced_at got %v", item.ReferencedAt)
+	}
+
+	reset()
+	// update referenced_at
+	wantReferencedAt := tm.Add(time.Hour)
+	reset = cache.SetNowFunc(wantReferencedAt)
+	defer reset()
+
+	item.Referenced()
+
+	if item.ReferenceCount != 2 {
+		t.Errorf("want ref count is 2 but got %d", item.ReferenceCount)
+	}
+	if !item.ReferencedAt.Equal(wantReferencedAt) {
+		t.Errorf("unexpected referenced_at got %v", item.ReferencedAt)
+	}
+}

lfu/example_test.go

@@ -0,0 +1,38 @@
+package lfu_test
+
+import (
+	"fmt"
+
+	"github.com/Code-Hex/go-generics-cache/lfu"
+)
+
+func ExampleLFUCache() {
+	c := lfu.NewCache[string, int]()
+	c.Set("a", 1)
+	c.Set("b", 2)
+	av, aok := c.Get("a")
+	bv, bok := c.Get("b")
+	cv, cok := c.Get("c")
+	fmt.Println(av, aok)
+	fmt.Println(bv, bok)
+	fmt.Println(cv, cok)
+	// Output:
+	// 1 true
+	// 2 true
+	// 0 false
+}
+
+func ExampleCacheKeys() {
+	c := lfu.NewCache[string, int]()
+	c.Set("a", 1)
+	c.Set("b", 2)
+	c.Set("c", 3)
+	keys := c.Keys()
+	for _, key := range keys {
+		fmt.Println(key)
+	}
+	// Output:
+	// a
+	// b
+	// c
+}

lfu/lfu.go

@@ -0,0 +1,108 @@
+package lfu
+
+import (
+	"container/heap"
+	"sync"
+
+	cache "github.com/Code-Hex/go-generics-cache"
+)
+
+// Cache is a thread safe LRU cache
+type Cache[K comparable, V any] struct {
+	cap   int
+	queue *priorityQueue[K, V]
+	items map[K]*entry[K, V]
+	mu    sync.RWMutex
+}
+
+var _ cache.Cache[interface{}, any] = (*Cache[interface{}, any])(nil)
+
+// NewCache creates a new LFU cache whose capacity is the default size (128).
+func NewCache[K comparable, V any]() *Cache[K, V] {
+	return NewCacheWithCap[K, V](128)
+}
+
+// NewCacheWithCap creates a new LFU cache whose capacity is the specified size.
+func NewCacheWithCap[K comparable, V any](cap int) *Cache[K, V] {
+	return &Cache[K, V]{
+		cap:   cap,
+		queue: newPriorityQueue[K, V](cap),
+		items: make(map[K]*entry[K, V], cap),
+	}
+}
+
+// Get looks up a key's value from the cache.
+func (c *Cache[K, V]) Get(key K) (zero V, _ bool) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	e, ok := c.items[key]
+	if !ok {
+		return
+	}
+	if e.item.HasExpired() {
+		return
+	}
+	e.item.Referenced()
+	heap.Fix(c.queue, e.index)
+	return e.item.Value, true
+}
+
+// Set sets a value to the cache with key. replacing any existing value.
+func (c *Cache[K, V]) Set(key K, val V, opts ...cache.ItemOption) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if e, ok := c.items[key]; ok {
+		c.queue.update(e, val)
+		return
+	}
+
+	if len(c.items) == c.cap {
+		evictedEntry := heap.Pop(c.queue).(*entry[K, V])
+		delete(c.items, evictedEntry.item.Key)
+	}
+
+	e := newEntry(key, val, opts...)
+	heap.Push(c.queue, e)
+	c.items[key] = e
+}
+
+// Keys returns the keys of the cache. the order is from oldest to newest.
+func (c *Cache[K, V]) Keys() []K {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	keys := make([]K, 0, len(c.items))
+	for _, entry := range *c.queue {
+		keys = append(keys, entry.item.Key)
+	}
+	return keys
+}
+
+// Delete deletes the item with provided key from the cache.
+func (c *Cache[K, V]) Delete(key K) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if e, ok := c.items[key]; ok {
+		heap.Remove(c.queue, e.index)
+		delete(c.items, key)
+	}
+}
+
+// Contains reports whether key is within cache.
+func (c *Cache[K, V]) Contains(key K) bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	e, ok := c.items[key]
+	if !ok {
+		return false
+	}
+	return !e.item.HasExpired()
+}
+
+// Len returns the number of items in the cache.
+func (c *Cache[K, V]) Len() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.queue.Len()
+}

lfu/lfu_internal_test.go

@@ -0,0 +1,69 @@
+package lfu
+
+import (
+	"testing"
+	"time"
+
+	cache "github.com/Code-Hex/go-generics-cache"
+)
+
+func TestContains(t *testing.T) {
+	t.Run("without expiration", func(t *testing.T) {
+		cache := NewCache[string, int]()
+		cache.Set("foo", 1)
+		cache.Set("bar", 2)
+		cache.Set("baz", 3)
+		for _, key := range []string{
+			"foo",
+			"bar",
+			"baz",
+		} {
+			if !cache.Contains(key) {
+				t.Errorf("not found: %s", key)
+			}
+		}
+		if cache.Contains("not found") {
+			t.Errorf("found")
+		}
+	})
+
+	t.Run("with expiration", func(t *testing.T) {
+		c := NewCache[string, int]()
+		key := "foo"
+		exp := time.Hour
+		c.Set(key, 1, cache.WithExpiration(exp))
+		// modify directly
+		e, ok := c.items[key]
+		if !ok {
+			t.Fatal("unexpected not found key")
+		}
+		item := e.item
+		item.CreatedAt = time.Now().Add(-2 * exp)
+
+		if c.Contains(key) {
+			t.Errorf("found")
+		}
+	})
+}
+
+func TestGet(t *testing.T) {
+	c := NewCache[string, int]()
+	key := "foo"
+	exp := time.Hour
+	c.Set(key, 1, cache.WithExpiration(exp))
+	_, ok := c.Get(key)
+	if !ok {
+		t.Fatal("unexpected not found")
+	}
+	// modify directly
+	e, ok := c.items[key]
+	if !ok {
+		t.Fatal("unexpected not found key")
+	}
+	item := e.item
+	item.CreatedAt = time.Now().Add(-2 * exp)
+	_, ok2 := c.Get(key)
+	if ok2 {
+		t.Fatal("unexpected found (expired)")
+	}
+}

lfu/lfu_test.go

@@ -0,0 +1,65 @@
+package lfu_test
+
+import (
+	"testing"
+
+	"github.com/Code-Hex/go-generics-cache/lfu"
+)
+
+func TestSet(t *testing.T) {
+	// set capacity is 1
+	cache := lfu.NewCacheWithCap[string, int](1)
+	cache.Set("foo", 1)
+	if got := cache.Len(); got != 1 {
+		t.Fatalf("invalid length: %d", got)
+	}
+	if got, ok := cache.Get("foo"); got != 1 || !ok {
+		t.Fatalf("invalid value got %d, cachehit %v", got, ok)
+	}
+
+	// if over the cap
+	cache.Set("bar", 2)
+	if got := cache.Len(); got != 1 {
+		t.Fatalf("invalid length: %d", got)
+	}
+	bar, ok := cache.Get("bar")
+	if bar != 2 || !ok {
+		t.Fatalf("invalid value bar %d, cachehit %v", bar, ok)
+	}
+
+	// checks deleted oldest
+	if _, ok := cache.Get("foo"); ok {
+		t.Fatalf("invalid delete oldest value foo %v", ok)
+	}
+
+	// valid: if over the cap but same key
+	cache.Set("bar", 100)
+	if got := cache.Len(); got != 1 {
+		t.Fatalf("invalid length: %d", got)
+	}
+	bar, ok = cache.Get("bar")
+	if bar != 100 || !ok {
+		t.Fatalf("invalid replacing value bar %d, cachehit %v", bar, ok)
+	}
+}
+
+func TestDelete(t *testing.T) {
+	cache := lfu.NewCacheWithCap[string, int](1)
+	cache.Set("foo", 1)
+	if got := cache.Len(); got != 1 {
+		t.Fatalf("invalid length: %d", got)
+	}
+
+	cache.Delete("foo2")
+	if got := cache.Len(); got != 1 {
+		t.Fatalf("invalid length after deleted does not exist key: %d", got)
+	}
+
+	cache.Delete("foo")
+	if got := cache.Len(); got != 0 {
+		t.Fatalf("invalid length after deleted: %d", got)
+	}
+	if _, ok := cache.Get("foo"); ok {
+		t.Fatalf("invalid get after deleted %v", ok)
+	}
+}