From 2e27fb89cea0684f6b84c8c4ccc2e285de2c785f Mon Sep 17 00:00:00 2001
From: raf555 <irsoekarno888@gmail.com>
Date: Sat, 3 Feb 2024 19:27:03 +0700
Subject: [PATCH] implement GetOrAdd to LRU

---
 lru.go      |  24 ++++++++++++
 lru_test.go | 104 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 128 insertions(+)

diff --git a/lru.go b/lru.go
index 2bb07fd..5d82ee5 100644
--- a/lru.go
+++ b/lru.go
@@ -117,6 +117,30 @@ func (c *Cache[K, V]) Peek(key K) (value V, ok bool) {
 	return value, ok
 }
 
+// GetOrAdd checks if a key is in the cache while
+// updating the recent-ness, and if not, adds the value.
+// Returns the value if found, whether found, and whether an eviction occurred.
+func (c *Cache[K, V]) GetOrAdd(key K, value V) (previous V, ok, evicted bool) {
+	var k K
+	var v V
+	c.lock.Lock()
+	previous, ok = c.lru.Get(key)
+	if ok {
+		c.lock.Unlock()
+		return previous, true, false
+	}
+	evicted = c.lru.Add(key, value)
+	if c.onEvictedCB != nil && evicted {
+		k, v = c.evictedKeys[0], c.evictedVals[0]
+		c.evictedKeys, c.evictedVals = c.evictedKeys[:0], c.evictedVals[:0]
+	}
+	c.lock.Unlock()
+	if c.onEvictedCB != nil && evicted {
+		c.onEvictedCB(k, v)
+	}
+	return
+}
+
 // ContainsOrAdd checks if a key is in the cache without updating the
 // recent-ness or deleting it for being stale, and if not, adds the value.
 // Returns whether found and whether an eviction occurred.
diff --git a/lru_test.go b/lru_test.go
index 7ecc7ae..5b65c23 100644
--- a/lru_test.go
+++ b/lru_test.go
@@ -298,6 +298,59 @@ func TestLRUResize(t *testing.T) {
 	}
 }
 
+// test that GetOrAdd update recent-ness
+func TestLRUGetkOrAdd(t *testing.T) {
+	l, err := New[int, int](2)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	l.Add(1, 1)
+	l.Add(2, 2)
+	// recent-ness keys: [1,2]
+	previous, contains, evict := l.GetOrAdd(1, 1)
+	if !contains {
+		t.Errorf("1 should be contained")
+	}
+	if evict {
+		t.Errorf("nothing should be evicted here")
+	}
+	if previous != 1 {
+		t.Errorf("previous is not equal to 1")
+	}
+	// recent-ness keys: [2,1]
+
+	l.Add(3, 3)
+	// recent-ness keys: [1, 3]
+	contains, evict = l.ContainsOrAdd(2, 2)
+	if contains {
+		t.Errorf("2 should not have been contained")
+	}
+	if !evict {
+		t.Errorf("an eviction should have occurred")
+	}
+	if l.Contains(1) {
+		t.Errorf("now 1 should be evicted")
+	}
+	// recent-ness keys: [3,2]
+
+	previous, contains, evict = l.GetOrAdd(1, 1)
+	if contains {
+		t.Errorf("1 should not be contained")
+	}
+	if !evict {
+		t.Errorf("eviction should have occurred")
+	}
+	if previous == 1 {
+		t.Errorf("previous should be zero")
+	}
+	// recent-ness keys: [2,1]
+
+	if !l.Contains(1) {
+		t.Errorf("now 1 should be contained")
+	}
+}
+
 func (c *Cache[K, V]) wantKeys(t *testing.T, want []K) {
 	t.Helper()
 	got := c.Keys()
@@ -443,4 +496,55 @@ func TestCache_EvictionSameKey(t *testing.T) {
 			t.Errorf("evictedKeys got: %v want: %v", evictedKeys, want)
 		}
 	})
+
+	t.Run("GetOrAdd", func(t *testing.T) {
+		var evictedKeys []int
+
+		cache, _ := NewWithEvict(
+			2,
+			func(key int, _ struct{}) {
+				evictedKeys = append(evictedKeys, key)
+			})
+
+		_, contained, evicted := cache.GetOrAdd(1, struct{}{})
+		if contained {
+			t.Error("First 1: got unexpected contained")
+		}
+		if evicted {
+			t.Error("First 1: got unexpected eviction")
+		}
+		cache.wantKeys(t, []int{1})
+
+		_, contained, evicted = cache.GetOrAdd(2, struct{}{})
+		if contained {
+			t.Error("2: got unexpected contained")
+		}
+		if evicted {
+			t.Error("2: got unexpected eviction")
+		}
+		cache.wantKeys(t, []int{1, 2})
+
+		_, contained, evicted = cache.GetOrAdd(1, struct{}{})
+		if !contained {
+			t.Error("Second 1: did not get expected contained")
+		}
+		if evicted {
+			t.Error("Second 1: got unexpected eviction")
+		}
+		cache.wantKeys(t, []int{2, 1})
+
+		_, contained, evicted = cache.GetOrAdd(3, struct{}{})
+		if contained {
+			t.Error("3: got unexpected contained")
+		}
+		if !evicted {
+			t.Error("3: did not get expected eviction")
+		}
+		cache.wantKeys(t, []int{1, 3})
+
+		want := []int{2}
+		if !reflect.DeepEqual(evictedKeys, want) {
+			t.Errorf("evictedKeys got: %v want: %v", evictedKeys, want)
+		}
+	})
 }