From 4ccae573e2523f35ba7d62c9623eaec52a91c1b1 Mon Sep 17 00:00:00 2001
From: Yiling-J <njjyl723@gmail.com>
Date: Tue, 29 Oct 2024 11:00:22 +0800
Subject: [PATCH] fix policy updateCost

---
 internal/tlfu.go      |  6 +-----
 internal/tlfu_test.go | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/internal/tlfu.go b/internal/tlfu.go
index 622274d..1d6450f 100644
--- a/internal/tlfu.go
+++ b/internal/tlfu.go
@@ -223,11 +223,7 @@ func (t *TinyLfu[K, V]) UpdateCost(entry *Entry[K, V], weightChange int64) {
 
 	// entry's policy weigh already updated
 	// so update weightedSize to keep sync
-	if weightChange > 0 {
-		t.weightedSize += uint(weightChange)
-	} else {
-		t.weightedSize -= uint(weightChange)
-	}
+	t.weightedSize += uint(weightChange)
 
 	// update window/slru
 	// if entry new weight > max size
diff --git a/internal/tlfu_test.go b/internal/tlfu_test.go
index 075a77f..4d22577 100644
--- a/internal/tlfu_test.go
+++ b/internal/tlfu_test.go
@@ -578,3 +578,24 @@ func TestTlfu_SketchResize(t *testing.T) {
 		require.Equal(t, size, len(tlfu.sketch.Table))
 	}
 }
+
+func TestTlfu_UpdateCost(t *testing.T) {
+	hasher := NewHasher[int](nil)
+	tlfu := NewTinyLfu[int, int](100, hasher)
+	e1 := &Entry[int, int]{key: 1, value: 1, policyWeight: 1}
+	e2 := &Entry[int, int]{key: 2, value: 1, policyWeight: 2}
+	e3 := &Entry[int, int]{key: 3, value: 1, policyWeight: 3}
+
+	tlfu.Set(e1)
+	tlfu.Set(e2)
+	tlfu.Set(e3)
+	require.Equal(t, 6, int(tlfu.weightedSize))
+
+	e1.policyWeight = 3
+	e2.policyWeight = 2
+	e3.policyWeight = 1
+	tlfu.UpdateCost(e1, 2)
+	tlfu.UpdateCost(e2, 0)
+	tlfu.UpdateCost(e3, -2)
+	require.Equal(t, 6, int(tlfu.weightedSize))
+}