refactor(bank): move bank balances to use collections (#15327)

types/collections.go

@@ -85,6 +85,40 @@ func (a genericAddressKey[T]) SizeNonTerminal(key T) int {
 	return collections.BytesKey.SizeNonTerminal(key)
 }
 
+// Deprecated: genericAddressIndexKey is a special key codec used to retain state backwards compatibility
+// when a generic address key (be: AccAddress, ValAddress, ConsAddress), is used as an index key.
+// More docs can be found in the AddressKeyAsIndexKey function.
+type genericAddressIndexKey[T addressUnion] struct {
+	collcodec.KeyCodec[T]
+}
+
+func (g genericAddressIndexKey[T]) Encode(buffer []byte, key T) (int, error) {
+	return g.EncodeNonTerminal(buffer, key)
+}
+
+func (g genericAddressIndexKey[T]) Decode(buffer []byte) (int, T, error) {
+	return g.DecodeNonTerminal(buffer)
+}
+
+func (g genericAddressIndexKey[T]) Size(key T) int { return g.SizeNonTerminal(key) }
+
+func (g genericAddressIndexKey[T]) KeyType() string { return "index_key/" + g.KeyCodec.KeyType() }
+
+// Deprecated: AddressKeyAsIndexKey implements an SDK backwards compatible indexing key encoder
+// for addresses.
+// The status quo in the SDK is that address keys are length prefixed even when they're the
+// last part of a composite key. This should never be used unless to retain state compatibility.
+// For example, a composite key composed of `[string, address]` in theory would need you only to
+// define a way to understand when the string part finishes, we usually do this by appending a null
+// byte to the string, then when you know when the string part finishes, it's logical that the
+// part which remains is the address key. In the SDK instead we prepend to the address key its
+// length too.
+func AddressKeyAsIndexKey[T addressUnion](keyCodec collcodec.KeyCodec[T]) collcodec.KeyCodec[T] {
+	return genericAddressIndexKey[T]{
+		keyCodec,
+	}
+}
+
 // Collection Codecs
 
 type intValueCodec struct{}

types/collections_test.go

@@ -18,4 +18,8 @@ func TestCollectionsCorrectness(t *testing.T) {
 	t.Run("ConsAddress", func(t *testing.T) {
 		colltest.TestKeyCodec(t, ConsAddressKey, ConsAddress{0x32, 0x0, 0x0, 0x3})
 	})
+
+	t.Run("AddressIndexingKey", func(t *testing.T) {
+		colltest.TestKeyCodec(t, AddressKeyAsIndexKey(AccAddressKey), AccAddress{0x2, 0x5, 0x8})
+	})
 }

types/query/collections_pagination.go

@@ -2,14 +2,22 @@ package query
 
 import (
 	"context"
-	"errors"
-	"fmt"
-
 	"cosmossdk.io/collections"
 	collcodec "cosmossdk.io/collections/codec"
 	storetypes "cosmossdk.io/store/types"
+	"errors"
+	"fmt"
 )
 
+// WithCollectionPaginationPairPrefix applies a prefix to a collection, whose key is a collection.Pair,
+// being paginated that needs prefixing.
+func WithCollectionPaginationPairPrefix[K1, K2 any](prefix K1) func(o *CollectionsPaginateOptions[collections.Pair[K1, K2]]) {
+	return func(o *CollectionsPaginateOptions[collections.Pair[K1, K2]]) {
+		prefix := collections.PairPrefix[K1, K2](prefix)
+		o.Prefix = &prefix
+	}
+}
+
 // CollectionsPaginateOptions provides extra options for pagination in collections.
 type CollectionsPaginateOptions[K any] struct {
 	// Prefix allows to optionally set a prefix for the pagination.
@@ -41,7 +49,7 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 	ctx context.Context,
 	coll C,
 	pageReq *PageRequest,
-	predicateFunc func(key K, value V) (include bool),
+	predicateFunc func(key K, value V) (include bool, err error),
 	opts ...func(opt *CollectionsPaginateOptions[K]),
 ) ([]collections.KeyValue[K, V], *PageResponse, error) {
 	if pageReq == nil {
@@ -89,7 +97,7 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 	}
 	// invalid iter error is ignored to retain Paginate behavior
 	if errors.Is(err, collections.ErrInvalidIterator) {
-		return results, pageRes, nil
+		return results, new(PageResponse), nil
 	}
 	// strip the prefix from next key
 	if len(pageRes.NextKey) != 0 && prefix != nil {
@@ -108,7 +116,7 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 	offset uint64,
 	limit uint64,
 	countTotal bool,
-	predicateFunc func(K, V) bool,
+	predicateFunc func(K, V) (bool, error),
 ) ([]collections.KeyValue[K, V], *PageResponse, error) {
 	iterator, err := getCollIter[K, V](ctx, coll, prefix, nil, reverse)
 	if err != nil {
@@ -137,12 +145,17 @@ 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)
-				count++
 				// if predicate function is defined we check if the result matches the filtering criteria
-			} else if predicateFunc(kv.Key, kv.Value) {
-				results = append(results, kv)
-				count++
+			} else {
+				include, err := predicateFunc(kv.Key, kv.Value)
+				if err != nil {
+					return nil, nil, err
+				}
+				if include {
+					results = append(results, kv)
+				}
 			}
+			count++
 		// second case, we found all the objects specified within the limit
 		case count == limit:
 			key, err := iterator.Key()
@@ -200,7 +213,7 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
 	key []byte,
 	reverse bool,
 	limit uint64,
-	predicateFunc func(K, V) bool,
+	predicateFunc func(K, V) (bool, error),
 ) ([]collections.KeyValue[K, V], *PageResponse, error) {
 	iterator, err := getCollIter[K, V](ctx, coll, prefix, key, reverse)
 	if err != nil {
@@ -237,13 +250,18 @@ 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)
-			count++
 			// if predicate is applied we execute the predicate function
 			// and append only if predicateFunc yields true.
-		} else if predicateFunc(kv.Key, kv.Value) {
-			results = append(results, kv)
-			count++
+		} else {
+			include, err := predicateFunc(kv.Key, kv.Value)
+			if err != nil {
+				return nil, nil, err
+			}
+			if include {
+				results = append(results, kv)
+			}
 		}
+		count++
 	}
 
 	return results, &PageResponse{
@@ -258,14 +276,20 @@ func encodeCollKey[K, V any, C Collection[K, V]](coll C, key K) ([]byte, error)
 	return buffer, err
 }
 
-func getCollIter[K, V any, C Collection[K, V]](ctx context.Context, coll C, prefix, start []byte, reverse bool) (collections.Iterator[K, V], error) {
+func getCollIter[K, V any, C Collection[K, V]](ctx context.Context, coll C, prefix []byte, start []byte, reverse bool) (collections.Iterator[K, V], error) {
+	// TODO: maybe can be simplified
+	if reverse {
+		var end []byte
+		if prefix != nil {
+			start = storetypes.PrefixEndBytes(append(prefix, start...))
+			end = prefix
+		}
+		return coll.IterateRaw(ctx, end, start, collections.OrderDescending)
+	}
 	var end []byte
 	if prefix != nil {
 		start = append(prefix, start...)
 		end = storetypes.PrefixEndBytes(prefix)
 	}
-	if reverse {
-		return coll.IterateRaw(ctx, nil, start, collections.OrderDescending)
-	}
 	return coll.IterateRaw(ctx, start, end, collections.OrderAscending)
 }

types/query/collections_pagination_test.go

@@ -48,7 +48,7 @@ func TestCollectionPagination(t *testing.T) {
 	type test struct {
 		req        *PageRequest
 		expResp    *PageResponse
-		filter     func(key, value uint64) bool
+		filter     func(key, value uint64) (bool, error)
 		expResults []collections.KeyValue[uint64, uint64]
 		wantErr    error
 	}
@@ -99,15 +99,14 @@ func TestCollectionPagination(t *testing.T) {
 				Limit: 3,
 			},
 			expResp: &PageResponse{
-				NextKey: encodeKey(5),
+				NextKey: encodeKey(3),
 			},
-			filter: func(key, value uint64) bool {
-				return key%2 == 0
+			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": {
@@ -116,15 +115,14 @@ func TestCollectionPagination(t *testing.T) {
 				Limit: 3,
 			},
 			expResp: &PageResponse{
-				NextKey: encodeKey(7),
+				NextKey: encodeKey(5),
 			},
-			filter: func(key, value uint64) bool {
-				return key%2 == 0
+			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},
 			},
 		},
 	}

types/query/pagination_test.go

@@ -231,7 +231,7 @@ func (s *paginationTestSuite) TestReversePagination() {
 	request := types.NewQueryAllBalancesRequest(addr1, pageReq, false)
 	res1, err := queryClient.AllBalances(gocontext.Background(), request)
 	s.Require().NoError(err)
-	s.Require().Equal(res1.Balances.Len(), 2)
+	s.Require().Equal(2, res1.Balances.Len())
 	s.Require().NotNil(res1.Pagination.NextKey)
 
 	s.T().Log("verify paginate with custom limit and countTotal, Reverse false")