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,16 +76,26 @@ 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.
+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,
 	}
 }
 

example_test.go

@@ -76,6 +76,24 @@ func ExampleWithExpiration() {
 	// 0 false
 }
 
+func ExampleWithReferenceCount() {
+	c := cache.New(cache.AsLFU[string, int]())
+	c.Set("a", 1, cache.WithReferenceCount(2))
+
+	// check item is set.
+	gota, aok := c.Get("a")
+	fmt.Println(gota, aok)
+
+	c.Set("b", 2, cache.WithReferenceCount(3))
+
+	gotb, bok := c.Get("b")
+	fmt.Println(gotb, bok)
+
+	// Output:
+	// 1 true
+	// 2 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.
@@ -73,7 +75,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/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/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(),
 	}
 }