Merge pull request #56 from Code-Hex/fix/55

fixed #55

Author: Code-Hex

policy/lfu/lfu.go

@@ -71,8 +71,10 @@ func (c *Cache[K, V]) Set(key K, val V) {
 	}
 
 	if len(c.items) == c.cap {
-		evictedEntry := heap.Pop(c.queue).(*entry[K, V])
-		delete(c.items, evictedEntry.key)
+		evictedEntry := heap.Pop(c.queue)
+		if evictedEntry != nil {
+			delete(c.items, evictedEntry.(*entry[K, V]).key)
+		}
 	}
 
 	e := newEntry(key, val)

policy/lfu/lfu_test.go

@@ -100,3 +100,16 @@ func TestIssue33(t *testing.T) {
 	cache.Delete("foo2")
 	cache.Delete("foo3")
 }
+
+func TestZeroCap(t *testing.T) {
+	cache := lfu.NewCache[string, int](lfu.WithCapacity(0))
+	cache.Set("foo", 1)
+
+	v, ok := cache.Get("foo")
+	if !ok {
+		t.Error(ok)
+	}
+	if v != 1 {
+		t.Error(v)
+	}
+}

policy/lfu/priority_queue.go

@@ -50,6 +50,9 @@ func (q priorityQueue[K, V]) Less(i, j int) bool {
 }
 
 func (q priorityQueue[K, V]) Swap(i, j int) {
+	if len(q) < 2 {
+		return
+	}
 	q[i], q[j] = q[j], q[i]
 	q[i].index = i
 	q[j].index = j
@@ -64,13 +67,13 @@ func (q *priorityQueue[K, V]) Push(x interface{}) {
 func (q *priorityQueue[K, V]) Pop() interface{} {
 	old := *q
 	n := len(old)
+	if n == 0 {
+		return nil // Return nil if the queue is empty to prevent panic
+	}
 	entry := old[n-1]
 	old[n-1] = nil   // avoid memory leak
 	entry.index = -1 // for safety
 	new := old[0 : n-1]
-	for i := 0; i < len(new); i++ {
-		new[i].index = i
-	}
 	*q = new
 	return entry
 }

policy/lfu/priority_queue_test.go

@@ -113,3 +113,56 @@ func Test_priorityQueue_Swap(t *testing.T) {
 		})
 	}
 }
+
+func TestPriorityQueue_Pop(t *testing.T) {
+	t.Run("Pop from empty queue", func(t *testing.T) {
+		pq := newPriorityQueue[int, string](0)
+		if elem := heap.Pop(pq); elem != nil {
+			t.Errorf("Expected nil from empty queue, got %v", elem)
+		}
+	})
+
+	t.Run("Pop from queue with single element", func(t *testing.T) {
+		pq := newPriorityQueue[int, string](10)
+		heap.Push(pq, newEntry(1, "one"))
+		if pq.Len() != 1 {
+			t.Fatalf("Expected queue length of 1, got %d", pq.Len())
+		}
+		elem := heap.Pop(pq).(*entry[int, string])
+		if elem.key != 1 || elem.val != "one" {
+			t.Errorf("Expected to pop element with key=1 and val='one', got key=%d and val='%s'", elem.key, elem.val)
+		}
+		if pq.Len() != 0 {
+			t.Errorf("Expected empty queue after pop, got length %d", pq.Len())
+		}
+	})
+
+	t.Run("Pop from queue with multiple elements", func(t *testing.T) {
+		pq := newPriorityQueue[int, string](10)
+		heap.Push(pq, newEntry(1, "one"))
+		heap.Push(pq, newEntry(2, "two"))
+		heap.Push(pq, newEntry(3, "three"))
+
+		// Pop the first element
+		elem := heap.Pop(pq).(*entry[int, string])
+		if elem.key != 1 || elem.val != "one" {
+			t.Errorf("Expected to pop element with key=1 and val='one', got key=%d and val='%s'", elem.key, elem.val)
+		}
+
+		// Pop the second element
+		elem = heap.Pop(pq).(*entry[int, string])
+		if elem.key != 2 || elem.val != "two" {
+			t.Errorf("Expected to pop element with key=2 and val='two', got key=%d and val='%s'", elem.key, elem.val)
+		}
+
+		// Pop the third element
+		elem = heap.Pop(pq).(*entry[int, string])
+		if elem.key != 3 || elem.val != "three" {
+			t.Errorf("Expected to pop element with key=3 and val='three', got key=%d and val='%s'", elem.key, elem.val)
+		}
+
+		if pq.Len() != 0 {
+			t.Errorf("Expected empty queue after all pops, got length %d", pq.Len())
+		}
+	})
+}