fix: collection filtered pagination (backport #23002) (#23139)

Co-authored-by: Akhil Kumar P <36399231+akhilkumarpilli@users.noreply.github.com>
Co-authored-by: akhilkumarpilli <akhilkumar7947@gmail.com>
Co-authored-by: Julien Robert <julien@rbrt.fr>

CHANGELOG.md

@@ -46,6 +46,7 @@ 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.
 * (x/auth/tx) [#23148](https://github.com/cosmos/cosmos-sdk/pull/23148) Avoid panic from intoAnyV2 when v1.PublicKey is optional.
 
 ### Deprecated

tests/integration/slashing/keeper/slash_redelegation_test.go

@@ -149,16 +149,6 @@ func TestSlashRedelegation(t *testing.T) {
 	err = slashKeeper.Slash(ctx, evilValConsAddr, math.LegacyMustNewDecFromStr("0.9"), evilPower, 3)
 	require.NoError(t, err)
 
-	// assert invariant to make sure we conduct slashing correctly
-	_, stop := stakingkeeper.AllInvariants(stakingKeeper)(ctx)
-	require.False(t, stop)
-
-	_, stop = bankkeeper.AllInvariants(bankKeeper)(ctx)
-	require.False(t, stop)
-
-	_, stop = distributionkeeper.AllInvariants(distrKeeper)(ctx)
-	require.False(t, stop)
-
 	// one eternity later
 	ctx, err = simtestutil.NextBlock(app, ctx, time.Duration(1000000000000000000))
 	require.NoError(t, err)
@@ -335,14 +325,6 @@ func TestOverSlashing(t *testing.T) {
 	err = slashKeeper.Slash(ctx, evilValConsAddr, math.LegacyMustNewDecFromStr(slashFraction), evilPower, misbehaveHeight)
 	require.NoError(t, err)
 
-	// assert invariants
-	_, stop := stakingkeeper.AllInvariants(stakingKeeper)(ctx)
-	require.False(t, stop)
-	_, stop = bankkeeper.AllInvariants(bankKeeper)(ctx)
-	require.False(t, stop)
-	_, stop = distributionkeeper.AllInvariants(distrKeeper)(ctx)
-	require.False(t, stop)
-
 	// fastforward 2 blocks to complete redelegations and unbondings
 	for i := 0; i < 2; i++ {
 		ctx, err = simtestutil.NextBlock(app, ctx, time.Duration(1000000000000000000))
@@ -489,14 +471,4 @@ func TestSlashRedelegation_ValidatorLeftWithNoTokens(t *testing.T) {
 
 	err = slashKeeper.Slash(ctx, srcConsAddr, math.LegacyMustNewDecFromStr("0.5"), srcPower, srcInfractionHeight)
 	require.NoError(t, err)
-
-	// assert invariants to ensure correctness
-	_, stop := stakingkeeper.AllInvariants(stakingKeeper)(ctx)
-	require.False(t, stop)
-
-	_, stop = bankkeeper.AllInvariants(bankKeeper)(ctx)
-	require.False(t, stop)
-
-	_, stop = distributionkeeper.AllInvariants(distrKeeper)(ctx)
-	require.False(t, stop)
 }

tests/integration/staking/keeper/slash_test.go

@@ -742,13 +742,6 @@ func TestFixAvoidFullSlashPenalty(t *testing.T) {
 	// next block
 	ctx, err = simtestutil.NextBlock(app, ctx, time.Duration(1))
 	require.NoError(t, err)
-	// assert invariants
-	_, stop := keeper.AllInvariants(stakingKeeper)(ctx)
-	require.False(t, stop)
-	_, stop = bankkeeper.AllInvariants(bankKeeper)(ctx)
-	require.False(t, stop)
-	_, stop = distributionkeeper.AllInvariants(distrKeeper)(ctx)
-	require.False(t, stop)
 	// evilValAddr1 is bad!
 	// lets slash evilValAddr1 with a 100% penalty
 	evilVal, err := stakingKeeper.GetValidator(ctx, evilValAddr1)

types/query/collections_pagination.go

@@ -266,53 +266,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{

types/query/collections_pagination_test.go

@@ -127,14 +127,29 @@ 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
 			},
 			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},
 			},
 		},
 	}