From b13a3c6f5f9a89f13121640f64dd820cf646655b Mon Sep 17 00:00:00 2001
From: Pavel Buchinchik <p.buchinchik@gmail.com>
Date: Thu, 7 Nov 2024 12:36:17 +0100
Subject: [PATCH] feat: add retrun value for Remove() method

---
 kpq/keyed_priority_queue.go      |  7 +++---
 kpq/keyed_priority_queue_test.go | 42 ++++++++++++++++++++------------
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/kpq/keyed_priority_queue.go b/kpq/keyed_priority_queue.go
index 1627c45..bf98021 100644
--- a/kpq/keyed_priority_queue.go
+++ b/kpq/keyed_priority_queue.go
@@ -229,14 +229,14 @@ func (pq *KeyedPriorityQueue[K, V]) ValueOf(k K) (V, bool) {
 }
 
 // Remove removes the priority value associated with the given key k from the priority queue.
-// It's a no-op if there's no such key k in the priority queue.
-func (pq *KeyedPriorityQueue[K, V]) Remove(k K) {
+// It's a no-op if there's no such key k in the priority queue and it will return false.
+func (pq *KeyedPriorityQueue[K, V]) Remove(k K) bool {
 	pq.mu.Lock()
 	defer pq.mu.Unlock()
 
 	i, ok := pq.im[k]
 	if !ok {
-		return
+		return false
 	}
 	n := len(pq.pm) - 1
 	if i != n {
@@ -247,6 +247,7 @@ func (pq *KeyedPriorityQueue[K, V]) Remove(k K) {
 	pq.pm = pq.pm[:n]
 	delete(pq.im, k)
 	delete(pq.vals, k)
+	return true
 }
 
 // Len returns the size of the priority queue.
diff --git a/kpq/keyed_priority_queue_test.go b/kpq/keyed_priority_queue_test.go
index 9b1b91d..100ade1 100644
--- a/kpq/keyed_priority_queue_test.go
+++ b/kpq/keyed_priority_queue_test.go
@@ -286,39 +286,60 @@ func TestKeyedPriorityQueue_Remove(t *testing.T) {
 	t.Run("Keys", func(t *testing.T) {
 		testCases := []struct {
 			key           string
+			wantStatus    bool
 			wantPeekKey   string
 			wantPeekValue int
 			wantLen       int
 		}{
 			{
 				key:           "first",
+				wantStatus:    true,
 				wantPeekKey:   "second",
 				wantPeekValue: 8,
 				wantLen:       4,
 			},
 			{
 				key:           "third",
+				wantStatus:    true,
 				wantPeekKey:   "second",
 				wantPeekValue: 8,
 				wantLen:       3,
 			},
 			{
 				key:           "second",
+				wantStatus:    true,
 				wantPeekKey:   "fourth",
 				wantPeekValue: 10,
 				wantLen:       2,
 			},
 			{
 				key:           "last",
+				wantStatus:    true,
 				wantPeekKey:   "fourth",
 				wantPeekValue: 10,
 				wantLen:       1,
 			},
+			{
+				key:        "nonexistent",
+				wantStatus: false,
+				wantLen:    1,
+			},
 		}
 
 		for _, tc := range testCases {
 			t.Run(tc.key, func(t *testing.T) {
-				pq.Remove(tc.key)
+				gotStatus := pq.Remove(tc.key)
+				if gotStatus != tc.wantStatus {
+					t.Errorf("pq.Remove(): got status %t; want %t", gotStatus, tc.wantStatus)
+				}
+
+				if got := pq.Len(); got != tc.wantLen {
+					t.Errorf("pq.Len(): got %d; want %d", got, tc.wantLen)
+				}
+
+				if !gotStatus {
+					return
+				}
 
 				gotPeekKey, gotPeekValue, ok := pq.Peek()
 				if !ok {
@@ -332,22 +353,9 @@ func TestKeyedPriorityQueue_Remove(t *testing.T) {
 				if gotPeekValue != tc.wantPeekValue {
 					t.Errorf("pq.PeekValue(): got value %d; want %d", gotPeekValue, tc.wantPeekValue)
 				}
-
-				if got := pq.Len(); got != tc.wantLen {
-					t.Errorf("pq.Len(): got %d; want %d", got, tc.wantLen)
-				}
 			})
 		}
 	})
-
-	t.Run("NonExistingKey", func(t *testing.T) {
-		want := pq.Len()
-		pq.Remove("non-existing-key")
-
-		if got := pq.Len(); got != want {
-			t.Errorf("pq.Len(): got %d; want %d", got, want)
-		}
-	})
 }
 
 func TestKeyedPriorityQueue_Peek_EmptyQueue(t *testing.T) {
@@ -479,7 +487,6 @@ func TestKeyedPriorityQueue_Set(t *testing.T) {
 			}
 		})
 	}
-
 }
 
 func benchmarkKeyedPriorityQueue_PushPop(b *testing.B, n int) {
@@ -500,18 +507,23 @@ func benchmarkKeyedPriorityQueue_PushPop(b *testing.B, n int) {
 func BenchmarkKeyedPriorityQueue_PushPop_10(b *testing.B) {
 	benchmarkKeyedPriorityQueue_PushPop(b, 10)
 }
+
 func BenchmarkKeyedPriorityQueue_PushPop_100(b *testing.B) {
 	benchmarkKeyedPriorityQueue_PushPop(b, 100)
 }
+
 func BenchmarkKeyedPriorityQueue_PushPop_1000(b *testing.B) {
 	benchmarkKeyedPriorityQueue_PushPop(b, 1000)
 }
+
 func BenchmarkKeyedPriorityQueue_PushPop_10000(b *testing.B) {
 	benchmarkKeyedPriorityQueue_PushPop(b, 10000)
 }
+
 func BenchmarkKeyedPriorityQueue_PushPop_100000(b *testing.B) {
 	benchmarkKeyedPriorityQueue_PushPop(b, 100000)
 }
+
 func BenchmarkKeyedPriorityQueue_PushPop_1000000(b *testing.B) {
 	benchmarkKeyedPriorityQueue_PushPop(b, 1000000)
 }