From 697219c9c438a44c69a39dc94fdf35d5de4bb07c Mon Sep 17 00:00:00 2001
From: Akhil Kumar P <36399231+akhilkumarpilli@users.noreply.github.com>
Date: Tue, 31 Dec 2024 22:04:12 +0530
Subject: [PATCH] fix: collection filtered pagination (#23002)

---
 CHANGELOG.md                               |  2 +
 types/query/collections_pagination.go      | 52 ++++++++++++----------
 types/query/collections_pagination_test.go | 17 ++++++-
 3 files changed, 47 insertions(+), 24 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1f664f146d19..86e8a2054b2b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -49,6 +49,8 @@ Every module contains its own CHANGELOG.md. Please refer to the module you are i
 
 ### Bug Fixes
 
+* (query) [23002](https://github.com/cosmos/cosmos-sdk/pull/23002) Fix collection filtered pagination.
+
 ### API Breaking Changes
 
 * (x/params) [#22995](https://github.com/cosmos/cosmos-sdk/pull/22995) Remove `x/params`.  Migrate to the new params system introduced in `v0.47` as demonstrated [here](https://github.com/cosmos/cosmos-sdk/blob/main/UPGRADING.md#xparams).
diff --git a/types/query/collections_pagination.go b/types/query/collections_pagination.go
index 2877df26396c..902d68b1dde3 100644
--- a/types/query/collections_pagination.go
+++ b/types/query/collections_pagination.go
@@ -265,53 +265,59 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V], T any](
 	defer iterator.Close()
 
 	var (
-		count   uint64
-		nextKey []byte
+		count       uint64
+		nextKey     []byte
+		transformed T
 	)
 
 	for ; iterator.Valid(); iterator.Next() {
-		// if we reached the specified limit
-		// then we get the next key, and we exit the iteration.
-		if count == limit {
-			concreteKey, err := iterator.Key()
-			if err != nil {
-				return nil, nil, err
-			}
-
-			nextKey, err = encodeCollKey[K, V](coll, concreteKey)
-			if err != nil {
-				return nil, nil, err
-			}
-			break
-		}
-
 		kv, err := iterator.KeyValue()
 		if err != nil {
 			return nil, nil, err
 		}
+
+		include := false
 		// if no predicate is specified then we just append the result
 		if predicateFunc == nil {
-			transformed, err := transformFunc(kv.Key, kv.Value)
+			transformed, err = transformFunc(kv.Key, kv.Value)
 			if err != nil {
 				return nil, nil, err
 			}
-			results = append(results, transformed)
+			include = true
 			// if predicate is applied we execute the predicate function
 			// and append only if predicateFunc yields true.
 		} else {
-			include, err := predicateFunc(kv.Key, kv.Value)
+			include, err = predicateFunc(kv.Key, kv.Value)
 			if err != nil {
 				return nil, nil, err
 			}
 			if include {
-				transformed, err := transformFunc(kv.Key, kv.Value)
+				transformed, err = transformFunc(kv.Key, kv.Value)
 				if err != nil {
 					return nil, nil, err
 				}
-				results = append(results, transformed)
 			}
 		}
-		count++
+
+		if include {
+			// if we reached the specified limit
+			// then we get the next key, and we exit the iteration.
+			if count == limit {
+				concreteKey, err := iterator.Key()
+				if err != nil {
+					return nil, nil, err
+				}
+
+				nextKey, err = encodeCollKey[K, V](coll, concreteKey)
+				if err != nil {
+					return nil, nil, err
+				}
+				break
+			}
+
+			results = append(results, transformed)
+			count++
+		}
 	}
 
 	return results, &PageResponse{
diff --git a/types/query/collections_pagination_test.go b/types/query/collections_pagination_test.go
index 25dc267a4481..6e4beb8c753c 100644
--- a/types/query/collections_pagination_test.go
+++ b/types/query/collections_pagination_test.go
@@ -127,7 +127,7 @@ func TestCollectionPagination(t *testing.T) {
 				Limit: 3,
 			},
 			expResp: &PageResponse{
-				NextKey: encodeKey(5),
+				NextKey: encodeKey(8),
 			},
 			filter: func(key, value uint64) (bool, error) {
 				return key%2 == 0, nil
@@ -135,6 +135,21 @@ func TestCollectionPagination(t *testing.T) {
 			expResults: []collections.KeyValue[uint64, uint64]{
 				{Key: 2, Value: 2},
 				{Key: 4, Value: 4},
+				{Key: 6, Value: 6},
+			},
+		},
+		"filtered with key and empty next key in response": {
+			req: &PageRequest{
+				Key: encodeKey(295),
+			},
+			expResp: &PageResponse{
+				NextKey: nil,
+			},
+			filter: func(key, value uint64) (bool, error) {
+				return key%5 == 0, nil
+			},
+			expResults: []collections.KeyValue[uint64, uint64]{
+				{Key: 295, Value: 295},
 			},
 		},
 	}