From b1235e07009e70351652cd50a2666fce23f359d2 Mon Sep 17 00:00:00 2001
From: vasayxtx <vasayxtx@gmail.com>
Date: Sun, 27 Dec 2020 21:22:13 +0300
Subject: [PATCH] Add method GetAndAdd to LRU Cache

---
 lru.go                     | 24 +++++++++++++++---------
 lru_test.go                | 19 +++++++++++++++++++
 simplelru/lru.go           | 17 ++++++++++++-----
 simplelru/lru_interface.go |  3 +++
 simplelru/lru_test.go      | 19 +++++++++++++++++++
 5 files changed, 68 insertions(+), 14 deletions(-)

diff --git a/lru.go b/lru.go
index 9af03d4..8525484 100644
--- a/lru.go
+++ b/lru.go
@@ -71,9 +71,23 @@ func (c *Cache) Purge() {
 
 // Add adds a value to the cache. Returns true if an eviction occurred.
 func (c *Cache) Add(key, value interface{}) (evicted bool) {
+	_, _, evicted = c.GetAndAdd(key, value)
+	return
+}
+
+// Get looks up a key's value from the cache.
+func (c *Cache) Get(key interface{}) (value interface{}, ok bool) {
+	c.lock.Lock()
+	value, ok = c.lru.Get(key)
+	c.lock.Unlock()
+	return value, ok
+}
+
+// GetAndAdd returns the previous key's value from the cache and adds a new one.
+func (c *Cache) GetAndAdd(key, value interface{}) (previous interface{}, ok, evicted bool) {
 	var k, v interface{}
 	c.lock.Lock()
-	evicted = c.lru.Add(key, value)
+	previous, ok, evicted = c.lru.GetAndAdd(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]
@@ -85,14 +99,6 @@ func (c *Cache) Add(key, value interface{}) (evicted bool) {
 	return
 }
 
-// Get looks up a key's value from the cache.
-func (c *Cache) Get(key interface{}) (value interface{}, ok bool) {
-	c.lock.Lock()
-	value, ok = c.lru.Get(key)
-	c.lock.Unlock()
-	return value, ok
-}
-
 // Contains checks if a key is in the cache, without updating the
 // recent-ness or deleting it for being stale.
 func (c *Cache) Contains(key interface{}) bool {
diff --git a/lru_test.go b/lru_test.go
index e4750d9..9d3b505 100644
--- a/lru_test.go
+++ b/lru_test.go
@@ -152,6 +152,25 @@ func TestLRUAdd(t *testing.T) {
 	}
 }
 
+func TestLRU_GetAndAdd(t *testing.T) {
+	l, err := New(1)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if previous, ok, evicted := l.GetAndAdd(1, 1); ok || evicted {
+		t.Errorf("1 should not be found, eviction should not be: %v, %v, %v", previous, ok, evicted)
+	}
+	if previous, ok, evicted := l.GetAndAdd(1, 10); !ok || previous != 1 || evicted {
+		t.Errorf("1 should be returned as previous value, eviction should not be: %v, %v, %v", previous, ok, evicted)
+	}
+	if v, ok := l.Get(1); !ok || v != 10 {
+		t.Errorf("1 should be set to 10: %v, %v", v, ok)
+	}
+	if previous, ok, evicted := l.GetAndAdd(2, 2); ok || !evicted {
+		t.Errorf("2 should not be found, eviction should be: %v, %v, %v", previous, ok, evicted)
+	}
+}
+
 // test that Contains doesn't update recent-ness
 func TestLRUContains(t *testing.T) {
 	l, err := New(2)
diff --git a/simplelru/lru.go b/simplelru/lru.go
index 9233583..d589449 100644
--- a/simplelru/lru.go
+++ b/simplelru/lru.go
@@ -47,13 +47,20 @@ func (c *LRU) Purge() {
 	c.evictList.Init()
 }
 
-// Add adds a value to the cache.  Returns true if an eviction occurred.
+// Add adds a value to the cache. Returns true if an eviction occurred.
 func (c *LRU) Add(key, value interface{}) (evicted bool) {
+	_, _, evicted = c.GetAndAdd(key, value)
+	return
+}
+
+// GetAndAdd returns the previous key's value from the cache and adds a new one.
+func (c *LRU) GetAndAdd(key, value interface{}) (previous interface{}, ok, evicted bool) {
 	// Check for existing item
 	if ent, ok := c.items[key]; ok {
 		c.evictList.MoveToFront(ent)
+		previous = ent.Value.(*entry).value
 		ent.Value.(*entry).value = value
-		return false
+		return previous, true, false
 	}
 
 	// Add new item
@@ -61,12 +68,12 @@ func (c *LRU) Add(key, value interface{}) (evicted bool) {
 	entry := c.evictList.PushFront(ent)
 	c.items[key] = entry
 
-	evict := c.evictList.Len() > c.size
+	evicted = c.evictList.Len() > c.size
 	// Verify size not exceeded
-	if evict {
+	if evicted {
 		c.removeOldest()
 	}
-	return evict
+	return nil, false, evicted
 }
 
 // Get looks up a key's value from the cache.
diff --git a/simplelru/lru_interface.go b/simplelru/lru_interface.go
index cb7f8ca..b532157 100644
--- a/simplelru/lru_interface.go
+++ b/simplelru/lru_interface.go
@@ -11,6 +11,9 @@ type LRUCache interface {
 	// updates the "recently used"-ness of the key. #value, isFound
 	Get(key interface{}) (value interface{}, ok bool)
 
+	// Returns the previous key's value from the cache and adds a new one.
+	GetAndAdd(key, value interface{}) (previous interface{}, ok, evicted bool)
+
 	// Checks if a key exists in cache without updating the recent-ness.
 	Contains(key interface{}) (ok bool)
 
diff --git a/simplelru/lru_test.go b/simplelru/lru_test.go
index a9ffcd5..5150356 100644
--- a/simplelru/lru_test.go
+++ b/simplelru/lru_test.go
@@ -128,6 +128,25 @@ func TestLRU_Add(t *testing.T) {
 	}
 }
 
+func TestLRU_GetAndAdd(t *testing.T) {
+	l, err := NewLRU(1, nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if previous, ok, evicted := l.GetAndAdd(1, 1); ok || evicted {
+		t.Errorf("1 should not be found, eviction should not be: %v, %v, %v", previous, ok, evicted)
+	}
+	if previous, ok, evicted := l.GetAndAdd(1, 10); !ok || previous != 1 || evicted {
+		t.Errorf("1 should be returned as previous value, eviction should not be: %v, %v, %v", previous, ok, evicted)
+	}
+	if v, ok := l.Get(1); !ok || v != 10 {
+		t.Errorf("1 should be set to 10: %v, %v", v, ok)
+	}
+	if previous, ok, evicted := l.GetAndAdd(2, 2); ok || !evicted {
+		t.Errorf("2 should not be found, eviction should be: %v, %v, %v", previous, ok, evicted)
+	}
+}
+
 // Test that Contains doesn't update recent-ness
 func TestLRU_Contains(t *testing.T) {
 	l, err := NewLRU(2, nil)