Merge pull request #38 from Code-Hex/add/with-ref-count

added WithReferenceCount option

Author: Code-Hex

cache.go

@@ -38,9 +38,10 @@ var (
 
 // Item is an item
 type Item[K comparable, V any] struct {
-	Key        K
-	Value      V
-	Expiration time.Time
+	Key                   K
+	Value                 V
+	Expiration            time.Time
+	InitialReferenceCount int
 }
 
 // Expired returns true if the item has expired.
@@ -51,13 +52,20 @@ func (item *Item[K, V]) Expired() bool {
 	return nowFunc().After(item.Expiration)
 }
 
+// GetReferenceCount returns reference count to be used when setting
+// the cache item for the first time.
+func (item *Item[K, V]) GetReferenceCount() int {
+	return item.InitialReferenceCount
+}
+
 var nowFunc = time.Now
 
 // ItemOption is an option for cache item.
 type ItemOption func(*itemOptions)
 
 type itemOptions struct {
-	expiration time.Time // default none
+	expiration     time.Time // default none
+	referenceCount int
 }
 
 // WithExpiration is an option to set expiration time for any items.
@@ -68,23 +76,34 @@ func WithExpiration(exp time.Duration) ItemOption {
 	}
 }
 
+// WithReferenceCount is an option to set reference count for any items.
+// This option is only applicable to cache policies that have a reference count (e.g., Clock, LFU).
+// referenceCount specifies the reference count value to set for the cache item.
+//
+// the default is 1.
+func WithReferenceCount(referenceCount int) ItemOption {
+	return func(o *itemOptions) {
+		o.referenceCount = referenceCount
+	}
+}
+
 // newItem creates a new item with specified any options.
 func newItem[K comparable, V any](key K, val V, opts ...ItemOption) *Item[K, V] {
 	o := new(itemOptions)
 	for _, optFunc := range opts {
 		optFunc(o)
 	}
 	return &Item[K, V]{
-		Key:        key,
-		Value:      val,
-		Expiration: o.expiration,
+		Key:                   key,
+		Value:                 val,
+		Expiration:            o.expiration,
+		InitialReferenceCount: o.referenceCount,
 	}
 }
 
 // Cache is a thread safe cache.
 type Cache[K comparable, V any] struct {
 	cache Interface[K, *Item[K, V]]
-	//expirations map[K]chan struct{}
 	// mu is used to do lock in some method process.
 	mu      sync.Mutex
 	janitor *janitor

example_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	cache "github.com/Code-Hex/go-generics-cache"
+	"github.com/Code-Hex/go-generics-cache/policy/lfu"
 )
 
 func ExampleCache() {
@@ -76,6 +77,32 @@ func ExampleWithExpiration() {
 	// 0 false
 }
 
+func ExampleWithReferenceCount() {
+	c := cache.New(cache.AsLFU[string, int](lfu.WithCapacity(2)))
+
+	// set item with reference count
+	c.Set("a", 1, cache.WithReferenceCount(5))
+
+	// check item is set.
+	gota, aok := c.Get("a")
+	fmt.Println(gota, aok)
+
+	c.Set("b", 2)
+	c.Set("c", 3)
+
+	// evicted becauce the lowest reference count.
+	gotb, bok := c.Get("b")
+	fmt.Println(gotb, bok)
+
+	gotc, cok := c.Get("c")
+	fmt.Println(gotc, cok)
+
+	// Output:
+	// 1 true
+	// 0 false
+	// 3 true
+}
+
 func ExampleCache_Delete() {
 	c := cache.New(cache.AsMRU[string, int]())
 	c.Set("a", 1)

policy/clock/clock.go

@@ -2,6 +2,8 @@ package clock
 
 import (
 	"container/ring"
+
+	"github.com/Code-Hex/go-generics-cache/policy/internal/policyutil"
 )
 
 // Cache is used The clock cache replacement policy.
@@ -62,6 +64,9 @@ func NewCache[K comparable, V any](opts ...Option) *Cache[K, V] {
 }
 
 // Set sets any item to the cache. replacing any existing item.
+//
+// If value satisfies "interface{ GetReferenceCount() int }", the value of
+// the GetReferenceCount() method is used to set the initial value of reference count.
 func (c *Cache[K, V]) Set(key K, val V) {
 	if e, ok := c.items[key]; ok {
 		entry := e.Value.(*entry[K, V])
@@ -73,7 +78,7 @@ func (c *Cache[K, V]) Set(key K, val V) {
 	c.hand.Value = &entry[K, V]{
 		key:            key,
 		val:            val,
-		referenceCount: 1,
+		referenceCount: policyutil.GetReferenceCount(val),
 	}
 	c.items[key] = c.hand
 	c.hand = c.hand.Next()

policy/clock/clock_test.go

@@ -8,6 +8,12 @@ import (
 	"github.com/Code-Hex/go-generics-cache/policy/clock"
 )
 
+type tmp struct {
+	i int
+}
+
+func (t *tmp) GetReferenceCount() int { return t.i }
+
 func TestSet(t *testing.T) {
 	// set capacity is 1
 	cache := clock.NewCache[string, int](clock.WithCapacity(1))
@@ -43,6 +49,25 @@ func TestSet(t *testing.T) {
 	if bar != 100 || !ok {
 		t.Fatalf("invalid replacing value bar %d, cachehit %v", bar, ok)
 	}
+
+	t.Run("with initilal reference count", func(t *testing.T) {
+		cache := clock.NewCache[string, *tmp](clock.WithCapacity(2))
+		cache.Set("foo", &tmp{i: 10}) // the highest reference count
+		cache.Set("foo2", &tmp{i: 2}) // expected eviction
+		if got := cache.Len(); got != 2 {
+			t.Fatalf("invalid length: %d", got)
+		}
+
+		cache.Set("foo3", &tmp{i: 3})
+
+		// checks deleted the lowest reference count
+		if _, ok := cache.Get("foo2"); ok {
+			t.Fatalf("invalid delete oldest value foo2 %v", ok)
+		}
+		if _, ok := cache.Get("foo"); !ok {
+			t.Fatalf("invalid value foo is not found")
+		}
+	})
 }
 
 func TestDelete(t *testing.T) {

policy/internal/policyutil/reference_count.go

@@ -0,0 +1,9 @@
+package policyutil
+
+// GetReferenceCount gets reference count from cache value.
+func GetReferenceCount(v any) int {
+	if getter, ok := v.(interface{ GetReferenceCount() int }); ok {
+		return getter.GetReferenceCount()
+	}
+	return 1
+}

policy/internal/policyutil/reference_count_test.go

@@ -0,0 +1,41 @@
+package policyutil
+
+import (
+	"testing"
+)
+
+type refCounter struct {
+	count int
+}
+
+func (r refCounter) GetReferenceCount() int {
+	return r.count
+}
+
+func TestGetReferenceCount(t *testing.T) {
+	tests := []struct {
+		name  string
+		input any
+		want  int
+	}{
+		{
+			name:  "with GetReferenceCount() method",
+			input: refCounter{count: 5},
+			want:  5,
+		},
+		{
+			name:  "without GetReferenceCount() method",
+			input: "sample string",
+			want:  1,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			output := GetReferenceCount(test.input)
+			if output != test.want {
+				t.Errorf("want %d, got %d", test.want, output)
+			}
+		})
+	}
+}

policy/lfu/lfu.go

@@ -61,6 +61,9 @@ func (c *Cache[K, V]) Get(key K) (zero V, _ bool) {
 }
 
 // Set sets a value to the cache with key. replacing any existing value.
+//
+// If value satisfies "interface{ GetReferenceCount() int }", the value of
+// the GetReferenceCount() method is used to set the initial value of reference count.
 func (c *Cache[K, V]) Set(key K, val V) {
 	if e, ok := c.items[key]; ok {
 		c.queue.update(e, val)

policy/lfu/lfu_test.go

@@ -6,6 +6,12 @@ import (
 	"github.com/Code-Hex/go-generics-cache/policy/lfu"
 )
 
+type tmp struct {
+	i int
+}
+
+func (t *tmp) GetReferenceCount() int { return t.i }
+
 func TestSet(t *testing.T) {
 	// set capacity is 1
 	cache := lfu.NewCache[string, int](lfu.WithCapacity(1))
@@ -41,6 +47,25 @@ func TestSet(t *testing.T) {
 	if bar != 100 || !ok {
 		t.Fatalf("invalid replacing value bar %d, cachehit %v", bar, ok)
 	}
+
+	t.Run("with initilal reference count", func(t *testing.T) {
+		cache := lfu.NewCache[string, *tmp](lfu.WithCapacity(2))
+		cache.Set("foo", &tmp{i: 10}) // the highest reference count
+		cache.Set("foo2", &tmp{i: 2}) // expected eviction
+		if got := cache.Len(); got != 2 {
+			t.Fatalf("invalid length: %d", got)
+		}
+
+		cache.Set("foo3", &tmp{i: 3})
+
+		// checks deleted the lowest reference count
+		if _, ok := cache.Get("foo2"); ok {
+			t.Fatalf("invalid delete oldest value foo2 %v", ok)
+		}
+		if _, ok := cache.Get("foo"); !ok {
+			t.Fatalf("invalid value foo is not found")
+		}
+	})
 }
 
 func TestDelete(t *testing.T) {

policy/lfu/priotiry_queue.go

@@ -3,6 +3,8 @@ package lfu
 import (
 	"container/heap"
 	"time"
+
+	"github.com/Code-Hex/go-generics-cache/policy/internal/policyutil"
 )
 
 type entry[K comparable, V any] struct {
@@ -18,7 +20,7 @@ func newEntry[K comparable, V any](key K, val V) *entry[K, V] {
 		index:          0,
 		key:            key,
 		val:            val,
-		referenceCount: 1,
+		referenceCount: policyutil.GetReferenceCount(val),
 		referencedAt:   time.Now(),
 	}
 }