fix(query): filtered collections pagination

CHANGELOG.md

@@ -102,6 +102,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (baseapp) [#16700](https://github.com/cosmos/cosmos-sdk/pull/16700) Fix consensus failure in returning no response to malformed transactions.
 * (baseapp) [#16596](https://github.com/cosmos/cosmos-sdk/pull/16596) Return error during ExtendVote and VerifyVoteExtension if the request height is earlier than `VoteExtensionsEnableHeight`.
 * (x/slashing) [#16784](https://github.com/cosmos/cosmos-sdk/pull/16784) Emit event with the correct reason in SlashWithInfractionReason.
+* (types/query) [#16905](https://github.com/cosmos/cosmos-sdk/pull/16905) – Collections Pagination now applies proper count when filtering results.
 
 ### API Breaking Changes
 

types/query/collections_pagination.go

@@ -36,6 +36,7 @@ type Collection[K, V any] interface {
 }
 
 // CollectionPaginate follows the same behavior as Paginate but works on a Collection.
+// Returns a list of the key-value pairs and a PageResponse.
 func CollectionPaginate[K, V any, C Collection[K, V]](
 	ctx context.Context,
 	coll C,
@@ -44,15 +45,62 @@ func CollectionPaginate[K, V any, C Collection[K, V]](
 	return CollectionFilteredPaginate[K, V](ctx, coll, pageReq, nil)
 }
 
-// CollectionFilteredPaginate works in the same way as FilteredPaginate but for collection types.
-// A nil predicateFunc means no filtering is applied and results are collected as is.
+// CollectionFilteredPaginate works in the same way as CollectionPaginate but allows to filter
+// the results using the predicateFunc.
+// NOTE: results should not be collected by the predicateFunc as they might iterate
+// over results which are not in the pagination result collection range.
 func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 	ctx context.Context,
 	coll C,
 	pageReq *PageRequest,
 	predicateFunc func(key K, value V) (include bool, err error),
 	opts ...func(opt *CollectionsPaginateOptions[K]),
 ) ([]collections.KeyValue[K, V], *PageResponse, error) {
+	return CollectionFilteredPaginateTransform(
+		ctx,
+		coll,
+		pageReq,
+		predicateFunc,
+		func(key K, value V) (collections.KeyValue[K, V], error) {
+			return collections.KeyValue[K, V]{Key: key, Value: value}, nil
+		},
+		opts...,
+	)
+}
+
+// CollectionPaginateTransform works like CollectionsPaginate but allows to transform the result
+// to a different type.
+func CollectionPaginateTransform[K, V any, C Collection[K, V], T any](
+	ctx context.Context,
+	coll C,
+	pageReq *PageRequest,
+	transformFunc func(key K, value V) (T, error),
+	opts ...func(opt *CollectionsPaginateOptions[K]),
+) ([]T, *PageResponse, error) {
+	return CollectionFilteredPaginateTransform(
+		ctx,
+		coll,
+		pageReq,
+		nil,
+		transformFunc,
+		opts...,
+	)
+}
+
+// CollectionFilteredPaginateTransform works in the same way as FilteredPaginate but for collection types.
+// A nil predicateFunc means no filtering is applied and results are collected as is.
+// TransformFunc is applied only to results which are in range of the pagination and allow
+// to convert the result to a different type.
+// NOTE: do not collect results using the values/keys passed to predicateFunc as they are not
+// guaranteed to be in the pagination range requested.
+func CollectionFilteredPaginateTransform[K, V any, C Collection[K, V], T any](
+	ctx context.Context,
+	coll C,
+	pageReq *PageRequest,
+	predicateFunc func(key K, value V) (include bool, err error),
+	transformFunc func(key K, value V) (T, error),
+	opts ...func(opt *CollectionsPaginateOptions[K]),
+) (results []T, pageRes *PageResponse, err error) {
 	pageReq = initPageRequestDefaults(pageReq)
 
 	offset := pageReq.Offset
@@ -65,12 +113,6 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 		return nil, nil, fmt.Errorf("invalid request, either offset or key is expected, got both")
 	}
 
-	var (
-		results []collections.KeyValue[K, V]
-		pageRes *PageResponse
-		err     error
-	)
-
 	opt := new(CollectionsPaginateOptions[K])
 	for _, o := range opts {
 		o(opt)
@@ -85,9 +127,9 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 	}
 
 	if len(key) != 0 {
-		results, pageRes, err = collFilteredPaginateByKey(ctx, coll, prefix, key, reverse, limit, predicateFunc)
+		results, pageRes, err = collFilteredPaginateByKey(ctx, coll, prefix, key, reverse, limit, predicateFunc, transformFunc)
 	} else {
-		results, pageRes, err = collFilteredPaginateNoKey(ctx, coll, prefix, reverse, offset, limit, countTotal, predicateFunc)
+		results, pageRes, err = collFilteredPaginateNoKey(ctx, coll, prefix, reverse, offset, limit, countTotal, predicateFunc, transformFunc)
 	}
 	// invalid iter error is ignored to retain Paginate behavior
 	if errors.Is(err, collections.ErrInvalidIterator) {
@@ -102,7 +144,7 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 
 // collFilteredPaginateNoKey applies the provided pagination on the collection when the starting key is not set.
 // If predicateFunc is nil no filtering is applied.
-func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
+func collFilteredPaginateNoKey[K, V any, C Collection[K, V], T any](
 	ctx context.Context,
 	coll C,
 	prefix []byte,
@@ -111,7 +153,8 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 	limit uint64,
 	countTotal bool,
 	predicateFunc func(K, V) (bool, error),
-) ([]collections.KeyValue[K, V], *PageResponse, error) {
+	transformFunc func(K, V) (T, error),
+) ([]T, *PageResponse, error) {
 	iterator, err := getCollIter[K, V](ctx, coll, prefix, nil, reverse)
 	if err != nil {
 		return nil, nil, err
@@ -125,7 +168,7 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 	var (
 		count   uint64
 		nextKey []byte
-		results []collections.KeyValue[K, V]
+		results []T
 	)
 
 	for ; iterator.Valid(); iterator.Next() {
@@ -138,18 +181,28 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 			}
 			// if no predicate function is specified then we just include the result
 			if predicateFunc == nil {
-				results = append(results, kv)
+				transformed, err := transformFunc(kv.Key, kv.Value)
+				if err != nil {
+					return nil, nil, err
+				}
+				results = append(results, transformed)
+				count++
+
 				// if predicate function is defined we check if the result matches the filtering criteria
 			} else {
 				include, err := predicateFunc(kv.Key, kv.Value)
 				if err != nil {
 					return nil, nil, err
 				}
 				if include {
-					results = append(results, kv)
+					transformed, err := transformFunc(kv.Key, kv.Value)
+					if err != nil {
+						return nil, nil, err
+					}
+					results = append(results, transformed)
+					count++
 				}
 			}
-			count++
 		// second case, we found all the objects specified within the limit
 		case count == limit:
 			key, err := iterator.Key()
@@ -172,12 +225,31 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 		// but we need to count how many possible results exist in total.
 		// so we keep increasing the count until the iterator is fully consumed.
 		case count > limit:
-			count++
+			if predicateFunc == nil {
+				count++
+
+				// if predicate function is defined we check if the result matches the filtering criteria
+			} else {
+				kv, err := iterator.KeyValue()
+				if err != nil {
+					return nil, nil, err
+				}
+
+				include, err := predicateFunc(kv.Key, kv.Value)
+				if err != nil {
+					return nil, nil, err
+				}
+				if include {
+					count++
+				}
+			}
 		}
 	}
+
 	resp := &PageResponse{
 		NextKey: nextKey,
 	}
+
 	if countTotal {
 		resp.Total = count + offset
 	}
@@ -200,15 +272,16 @@ func advanceIter[I interface {
 
 // collFilteredPaginateByKey paginates a collection when a starting key
 // is provided in the PageRequest. Predicate is applied only if not nil.
-func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
+func collFilteredPaginateByKey[K, V any, C Collection[K, V], T any](
 	ctx context.Context,
 	coll C,
 	prefix []byte,
 	key []byte,
 	reverse bool,
 	limit uint64,
-	predicateFunc func(K, V) (bool, error),
-) ([]collections.KeyValue[K, V], *PageResponse, error) {
+	predicateFunc func(key K, value V) (bool, error),
+	transformFunc func(key K, value V) (transformed T, err error),
+) (results []T, pageRes *PageResponse, err error) {
 	iterator, err := getCollIter[K, V](ctx, coll, prefix, key, reverse)
 	if err != nil {
 		return nil, nil, err
@@ -218,7 +291,6 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
 	var (
 		count   uint64
 		nextKey []byte
-		results []collections.KeyValue[K, V]
 	)
 
 	for ; iterator.Valid(); iterator.Next() {
@@ -243,7 +315,11 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
 		}
 		// if no predicate is specified then we just append the result
 		if predicateFunc == nil {
-			results = append(results, kv)
+			transformed, err := transformFunc(kv.Key, kv.Value)
+			if err != nil {
+				return nil, nil, err
+			}
+			results = append(results, transformed)
 			// if predicate is applied we execute the predicate function
 			// and append only if predicateFunc yields true.
 		} else {
@@ -252,7 +328,11 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
 				return nil, nil, err
 			}
 			if include {
-				results = append(results, kv)
+				transformed, err := transformFunc(kv.Key, kv.Value)
+				if err != nil {
+					return nil, nil, err
+				}
+				results = append(results, transformed)
 			}
 		}
 		count++

types/query/collections_pagination_test.go

@@ -100,14 +100,15 @@ func TestCollectionPagination(t *testing.T) {
 				Limit: 3,
 			},
 			expResp: &PageResponse{
-				NextKey: encodeKey(3),
+				NextKey: encodeKey(5),
 			},
 			filter: func(key, value uint64) (bool, error) {
 				return key%2 == 0, nil
 			},
 			expResults: []collections.KeyValue[uint64, uint64]{
 				{Key: 0, Value: 0},
 				{Key: 2, Value: 2},
+				{Key: 4, Value: 4},
 			},
 		},
 		"filtered with key": {

x/auth/keeper/grpc_query.go

@@ -47,15 +47,14 @@ func (s queryServer) Accounts(ctx context.Context, req *types.QueryAccountsReque
 		return nil, status.Error(codes.InvalidArgument, "empty request")
 	}
 
-	var accounts []*codectypes.Any
-	_, pageRes, err := query.CollectionFilteredPaginate(ctx, s.k.Accounts, req.Pagination, func(_ sdk.AccAddress, value sdk.AccountI) (include bool, err error) {
-		accountAny, err := codectypes.NewAnyWithValue(value)
-		if err != nil {
-			return false, err
-		}
-		accounts = append(accounts, accountAny)
-		return false, nil // we don't include it since we're already appending the account
-	})
+	accounts, pageRes, err := query.CollectionPaginateTransform(
+		ctx,
+		s.k.Accounts,
+		req.Pagination,
+		func(_ sdk.AccAddress, value sdk.AccountI) (*codectypes.Any, error) {
+			return codectypes.NewAnyWithValue(value)
+		},
+	)
 
 	return &types.QueryAccountsResponse{Accounts: accounts, Pagination: pageRes}, err
 }

x/bank/keeper/grpc_query.go

@@ -60,19 +60,20 @@ func (k BaseKeeper) AllBalances(ctx context.Context, req *types.QueryAllBalances
 	}
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
-
-	balances := sdk.NewCoins()
-
-	_, pageRes, err := query.CollectionFilteredPaginate(ctx, k.Balances, req.Pagination, func(key collections.Pair[sdk.AccAddress, string], value math.Int) (include bool, err error) {
-		denom := key.K2()
-		if req.ResolveDenom {
-			if metadata, ok := k.GetDenomMetaData(sdkCtx, denom); ok {
-				denom = metadata.Display
+	balances, pageRes, err := query.CollectionPaginateTransform(
+		ctx,
+		k.Balances,
+		req.Pagination,
+		func(key collections.Pair[sdk.AccAddress, string], value math.Int) (sdk.Coin, error) {
+			if req.ResolveDenom {
+				if metadata, ok := k.GetDenomMetaData(sdkCtx, key.K2()); ok {
+					return sdk.NewCoin(metadata.Display, value), nil
+				}
 			}
-		}
-		balances = append(balances, sdk.NewCoin(denom, value))
-		return false, nil // we don't include results because we're appending them here.
-	}, query.WithCollectionPaginationPairPrefix[sdk.AccAddress, string](addr))
+			return sdk.NewCoin(key.K2(), value), nil
+		},
+		query.WithCollectionPaginationPairPrefix[sdk.AccAddress, string](addr),
+	)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "paginate: %v", err)
 	}
@@ -94,12 +95,10 @@ func (k BaseKeeper) SpendableBalances(ctx context.Context, req *types.QuerySpend
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
-	balances := sdk.NewCoins()
 	zeroAmt := math.ZeroInt()
 
-	_, pageRes, err := query.CollectionFilteredPaginate(ctx, k.Balances, req.Pagination, func(key collections.Pair[sdk.AccAddress, string], _ math.Int) (include bool, err error) {
-		balances = append(balances, sdk.NewCoin(key.K2(), zeroAmt))
-		return false, nil // not including results as they're appended here
+	balances, pageRes, err := query.CollectionPaginateTransform(ctx, k.Balances, req.Pagination, func(key collections.Pair[sdk.AccAddress, string], _ math.Int) (coin sdk.Coin, err error) {
+		return sdk.NewCoin(key.K2(), zeroAmt), nil
 	}, query.WithCollectionPaginationPairPrefix[sdk.AccAddress, string](addr))
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "paginate: %v", err)
@@ -280,19 +279,16 @@ func (k BaseKeeper) DenomOwners(
 		return nil, status.Error(codes.InvalidArgument, err.Error())
 	}
 
-	var denomOwners []*types.DenomOwner
-
-	_, pageRes, err := query.CollectionFilteredPaginate(goCtx, k.Balances.Indexes.Denom, req.Pagination,
-		func(key collections.Pair[string, sdk.AccAddress], value collections.NoValue) (include bool, err error) {
+	denomOwners, pageRes, err := query.CollectionPaginateTransform(
+		goCtx,
+		k.Balances.Indexes.Denom,
+		req.Pagination,
+		func(key collections.Pair[string, sdk.AccAddress], value collections.NoValue) (*types.DenomOwner, error) {
 			amt, err := k.Balances.Get(goCtx, collections.Join(key.K2(), req.Denom))
 			if err != nil {
-				return false, err
+				return nil, err
 			}
-			denomOwners = append(denomOwners, &types.DenomOwner{
-				Address: key.K2().String(),
-				Balance: sdk.NewCoin(req.Denom, amt),
-			})
-			return false, nil
+			return &types.DenomOwner{Address: key.K2().String(), Balance: sdk.NewCoin(req.Denom, amt)}, nil
 		},
 		query.WithCollectionPaginationPairPrefix[string, sdk.AccAddress](req.Denom),
 	)