From 93d88c04437a3053f1de8a761d2504a80b00a512 Mon Sep 17 00:00:00 2001
From: Dhruba Basu <7675102+dhrubabasu@users.noreply.github.com>
Date: Wed, 8 Nov 2023 15:02:10 -0500
Subject: [PATCH] Return if element was deleted from `Hashmap` (#2271)

Co-authored-by: Dan Laine <daniel.laine@avalabs.org>
---
 utils/linkedhashmap/linkedhashmap.go      | 12 +++++++-----
 utils/linkedhashmap/linkedhashmap_test.go |  8 ++++----
 vms/avm/txs/mempool/mempool.go            | 14 +++++---------
 3 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/utils/linkedhashmap/linkedhashmap.go b/utils/linkedhashmap/linkedhashmap.go
index 12e9569c3391..fa5a123a0942 100644
--- a/utils/linkedhashmap/linkedhashmap.go
+++ b/utils/linkedhashmap/linkedhashmap.go
@@ -17,7 +17,7 @@ var _ LinkedHashmap[int, struct{}] = (*linkedHashmap[int, struct{}])(nil)
 type Hashmap[K, V any] interface {
 	Put(key K, val V)
 	Get(key K) (val V, exists bool)
-	Delete(key K)
+	Delete(key K) (deleted bool)
 	Len() int
 }
 
@@ -63,11 +63,11 @@ func (lh *linkedHashmap[K, V]) Get(key K) (V, bool) {
 	return lh.get(key)
 }
 
-func (lh *linkedHashmap[K, V]) Delete(key K) {
+func (lh *linkedHashmap[K, V]) Delete(key K) bool {
 	lh.lock.Lock()
 	defer lh.lock.Unlock()
 
-	lh.delete(key)
+	return lh.delete(key)
 }
 
 func (lh *linkedHashmap[K, V]) Len() int {
@@ -114,11 +114,13 @@ func (lh *linkedHashmap[K, V]) get(key K) (V, bool) {
 	return utils.Zero[V](), false
 }
 
-func (lh *linkedHashmap[K, V]) delete(key K) {
-	if e, ok := lh.entryMap[key]; ok {
+func (lh *linkedHashmap[K, V]) delete(key K) bool {
+	e, ok := lh.entryMap[key]
+	if ok {
 		lh.entryList.Remove(e)
 		delete(lh.entryMap, key)
 	}
+	return ok
 }
 
 func (lh *linkedHashmap[K, V]) len() int {
diff --git a/utils/linkedhashmap/linkedhashmap_test.go b/utils/linkedhashmap/linkedhashmap_test.go
index 8bd7239ed5d9..0c95c30b24a8 100644
--- a/utils/linkedhashmap/linkedhashmap_test.go
+++ b/utils/linkedhashmap/linkedhashmap_test.go
@@ -62,7 +62,7 @@ func TestLinkedHashmap(t *testing.T) {
 	require.Equal(key1, rkey1, "wrong key")
 	require.Equal(1, val1, "wrong value")
 
-	lh.Delete(key0)
+	require.True(lh.Delete(key0))
 	require.Equal(1, lh.Len(), "wrong hashmap length")
 
 	_, exists = lh.Get(key0)
@@ -132,7 +132,7 @@ func TestIterator(t *testing.T) {
 		// Should be empty
 		require.False(iter.Next())
 		// Delete id1
-		lh.Delete(id1)
+		require.True(lh.Delete(id1))
 		iter = lh.NewIterator()
 		require.NotNil(iter)
 		// Should immediately be exhausted
@@ -169,8 +169,8 @@ func TestIterator(t *testing.T) {
 		iter := lh.NewIterator()
 		require.True(iter.Next())
 		require.True(iter.Next())
-		lh.Delete(id1)
-		lh.Delete(id2)
+		require.True(lh.Delete(id1))
+		require.True(lh.Delete(id2))
 		require.True(iter.Next())
 		require.Equal(id3, iter.Key())
 		require.Equal(3, iter.Value())
diff --git a/vms/avm/txs/mempool/mempool.go b/vms/avm/txs/mempool/mempool.go
index b64002e8f39d..84a3583781ef 100644
--- a/vms/avm/txs/mempool/mempool.go
+++ b/vms/avm/txs/mempool/mempool.go
@@ -48,8 +48,7 @@ type Mempool interface {
 	Get(txID ids.ID) *txs.Tx
 	Remove(txs []*txs.Tx)
 
-	// Peek returns the next first tx that was added to the mempool whose size
-	// is less than or equal to maxTxSize.
+	// Peek returns the first tx in the mempool whose size is <= [maxTxSize].
 	Peek(maxTxSize int) *txs.Tx
 
 	// RequestBuildBlock notifies the consensus engine that a block should be
@@ -162,23 +161,20 @@ func (m *mempool) Has(txID ids.ID) bool {
 }
 
 func (m *mempool) Get(txID ids.ID) *txs.Tx {
-	unissuedTxs, _ := m.unissuedTxs.Get(txID)
-	return unissuedTxs
+	tx, _ := m.unissuedTxs.Get(txID)
+	return tx
 }
 
 func (m *mempool) Remove(txsToRemove []*txs.Tx) {
 	for _, tx := range txsToRemove {
 		txID := tx.ID()
-		if _, ok := m.unissuedTxs.Get(txID); !ok {
-			// If tx isn't in the mempool, there is nothing to do.
+		if !m.unissuedTxs.Delete(txID) {
 			continue
 		}
 
-		txBytes := tx.Bytes()
-		m.bytesAvailable += len(txBytes)
+		m.bytesAvailable += len(tx.Bytes())
 		m.bytesAvailableMetric.Set(float64(m.bytesAvailable))
 
-		m.unissuedTxs.Delete(txID)
 		m.numTxs.Dec()
 
 		inputs := tx.Unsigned.InputIDs()