diff --git a/util/store/iter.go b/util/store/iter.go index 4860b3f5a3..ba6e2072e2 100644 --- a/util/store/iter.go +++ b/util/store/iter.go @@ -62,3 +62,27 @@ func MustLoadAll[TPtr PtrMarshalable[T], T any](s storetypes.KVStore, prefix []b util.Panic(err) return ls } + +// SumCoins aggregates all coins saved as (denom: Int) pairs in store. Use store/prefix.NewStore +// to create a prefix store which will automatically look only at the given prefix. +func SumCoins(s storetypes.KVStore, f StrExtractor) sdk.Coins { + total := sdk.NewCoins() + iter := sdk.KVStorePrefixIterator(s, nil) + defer iter.Close() + for ; iter.Valid(); iter.Next() { + denom := f(iter.Key()) + amount := Int(iter.Value(), "amount") + total = total.Add(sdk.NewCoin(denom, amount)) + } + return total +} + +// StrExtractor is a function type which will take a bytes string value and extracts +// string out of it. +type StrExtractor func([]byte) string + +// NoLastByte returns sub-slice of the key without the last byte. +// Panics if length of key is zero. +func NoLastByte(key []byte) string { + return string(key[:len(key)-1]) +} diff --git a/util/store/iter_test.go b/util/store/iter_test.go index b5b0a49530..84908264be 100644 --- a/util/store/iter_test.go +++ b/util/store/iter_test.go @@ -3,8 +3,11 @@ package store import ( "testing" - "github.com/umee-network/umee/v5/tests/tsdk" + prefixstore "github.com/cosmos/cosmos-sdk/store/prefix" + sdk "github.com/cosmos/cosmos-sdk/types" "gotest.tools/v3/assert" + + "github.com/umee-network/umee/v5/tests/tsdk" ) func TestIterate(t *testing.T) { @@ -38,3 +41,40 @@ func TestIterate(t *testing.T) { assert.DeepEqual(t, pairs[1:4], collected) } + +func TestSumCoins(t *testing.T) { + // test SumCoins using the Prefix Store, which will automatically strip the prefix from + // keys + + prefix := "p1" + pairs := []struct { + K string + V uint64 + }{ + {"atom", 1}, + {"umee", 8}, + {"atom", 8}, // we overwrite + {"ato", 2}, + {"atoma", 3}, + } + expected := sdk.NewCoins( + sdk.NewInt64Coin("ato", 2), + sdk.NewInt64Coin("atom", 8), + sdk.NewInt64Coin("atoma", 3), + sdk.NewInt64Coin("umee", 8)) + + withPrefixAnNull := func(s string) []byte { + return append([]byte(prefix+s), 0) + } + + db := tsdk.KVStore(t) + for i, p := range pairs { + err := SetInt(db, withPrefixAnNull(p.K), sdk.NewIntFromUint64(p.V), "amount") + assert.NilError(t, err, "pairs[%d]", i) + } + + pdb := prefixstore.NewStore(db, []byte(prefix)) + sum := SumCoins(pdb, NoLastByte) + sum.Sort() + assert.DeepEqual(t, expected, sum) +} diff --git a/x/incentive/keeper/grpc_query.go b/x/incentive/keeper/grpc_query.go index aaf39a5bf9..49990f64b6 100644 --- a/x/incentive/keeper/grpc_query.go +++ b/x/incentive/keeper/grpc_query.go @@ -198,11 +198,7 @@ func (q Querier) TotalBonded( if req.Denom != "" { total = sdk.NewCoins(k.getTotalBonded(ctx, req.Denom)) } else { - var err error - total, err = k.getAllTotalBonded(ctx) - if err != nil { - return nil, err - } + total = k.getAllTotalBonded(ctx) } return &incentive.QueryTotalBondedResponse{Bonded: total}, nil @@ -222,11 +218,7 @@ func (q Querier) TotalUnbonding( if req.Denom != "" { total = sdk.NewCoins(k.getTotalUnbonding(ctx, req.Denom)) } else { - var err error - total, err = k.getAllTotalUnbonding(ctx) - if err != nil { - return nil, err - } + total = k.getAllTotalUnbonding(ctx) } return &incentive.QueryTotalUnbondingResponse{Unbonding: total}, nil diff --git a/x/incentive/keeper/iter.go b/x/incentive/keeper/iter.go index 9753e83c76..62c875a686 100644 --- a/x/incentive/keeper/iter.go +++ b/x/incentive/keeper/iter.go @@ -13,8 +13,6 @@ import ( // The status of an incentive program is either Upcoming, Ongoing, or Completed. func (k Keeper) getAllIncentivePrograms(ctx sdk.Context, status incentive.ProgramStatus, ) ([]incentive.IncentiveProgram, error) { - programs := []incentive.IncentiveProgram{} - var prefix []byte switch status { case incentive.ProgramStatusUpcoming: @@ -27,16 +25,7 @@ func (k Keeper) getAllIncentivePrograms(ctx sdk.Context, status incentive.Progra return []incentive.IncentiveProgram{}, incentive.ErrInvalidProgramStatus } - iterator := func(_, val []byte) error { - var p incentive.IncentiveProgram - k.cdc.MustUnmarshal(val, &p) - - programs = append(programs, p) - return nil - } - - err := store.Iterate(k.KVStore(ctx), prefix, iterator) - return programs, err + return store.LoadAll[*incentive.IncentiveProgram](k.KVStore(ctx), prefix) } // getAllBondDenoms gets all uToken denoms for which an account has nonzero bonded amounts. @@ -45,7 +34,7 @@ func (k Keeper) getAllBondDenoms(ctx sdk.Context, addr sdk.AccAddress) ([]string prefix := keyBondAmountNoDenom(addr) bonds := []string{} - iterator := func(key, val []byte) error { + iterator := func(key, _ []byte) error { _, denom, _, err := keys.ExtractAddressAndString(len(keyPrefixBondAmount), key) if err != nil { return err @@ -81,88 +70,27 @@ func (k Keeper) getAllBonds(ctx sdk.Context) ([]incentive.Bond, error) { return bonds, err } -// getAllTotalBonded gets total bonded for all uTokens (used for a query) -func (k Keeper) getAllTotalBonded(ctx sdk.Context) (sdk.Coins, error) { - prefix := keyPrefixTotalBonded - total := sdk.NewCoins() - - iterator := func(key, val []byte) error { - denom, _, err := keys.ExtractString(len(keyPrefixTotalBonded), key) - if err != nil { - return err - } - amount := store.Int(val, "total bonded") - total = total.Add(sdk.NewCoin(denom, amount)) - return nil - } - - err := store.Iterate(k.KVStore(ctx), prefix, iterator) - return total, err -} - // getAllRewardTrackers gets all reward trackers for all accounts (used during export genesis) func (k Keeper) getAllRewardTrackers(ctx sdk.Context) ([]incentive.RewardTracker, error) { - prefix := keyPrefixRewardTracker - rewardTrackers := []incentive.RewardTracker{} - - iterator := func(_, val []byte) error { - tracker := incentive.RewardTracker{} - k.cdc.MustUnmarshal(val, &tracker) - rewardTrackers = append(rewardTrackers, tracker) - return nil - } - - err := store.Iterate(k.KVStore(ctx), prefix, iterator) - return rewardTrackers, err + return store.LoadAll[*incentive.RewardTracker](k.KVStore(ctx), keyPrefixRewardTracker) } // getAllRewardAccumulators gets all reward accumulators for all uTokens (used during export genesis) func (k Keeper) getAllRewardAccumulators(ctx sdk.Context) ([]incentive.RewardAccumulator, error) { - prefix := keyPrefixRewardAccumulator - rewardAccumulators := []incentive.RewardAccumulator{} - - iterator := func(_, val []byte) error { - accumulator := incentive.RewardAccumulator{} - k.cdc.MustUnmarshal(val, &accumulator) - rewardAccumulators = append(rewardAccumulators, accumulator) - return nil - } - - err := store.Iterate(k.KVStore(ctx), prefix, iterator) - return rewardAccumulators, err + return store.LoadAll[*incentive.RewardAccumulator](k.KVStore(ctx), keyPrefixRewardAccumulator) } // getAllAccountUnbondings gets all account unbondings for all accounts (used during export genesis) func (k Keeper) getAllAccountUnbondings(ctx sdk.Context) ([]incentive.AccountUnbondings, error) { - prefix := keyPrefixUnbondings - unbondings := []incentive.AccountUnbondings{} - - iterator := func(key, val []byte) error { - au := incentive.AccountUnbondings{} - k.cdc.MustUnmarshal(val, &au) - unbondings = append(unbondings, au) - return nil - } - - err := store.Iterate(k.KVStore(ctx), prefix, iterator) - return unbondings, err + return store.LoadAll[*incentive.AccountUnbondings](k.KVStore(ctx), keyPrefixUnbondings) } // getAllTotalUnbonding gets total unbonding for all uTokens (used for a query) -func (k Keeper) getAllTotalUnbonding(ctx sdk.Context) (sdk.Coins, error) { - prefix := keyPrefixTotalUnbonding - total := sdk.NewCoins() - - iterator := func(key, val []byte) error { - denom, _, err := keys.ExtractString(len(keyPrefixTotalUnbonding), key) - if err != nil { - return err - } - amount := store.Int(val, "total unbonding") - total = total.Add(sdk.NewCoin(denom, amount)) - return nil - } +func (k Keeper) getAllTotalUnbonding(ctx sdk.Context) sdk.Coins { + return store.SumCoins(k.prefixStore(ctx, keyPrefixTotalUnbonding), store.NoLastByte) +} - err := store.Iterate(k.KVStore(ctx), prefix, iterator) - return total, err +// getAllTotalBonded gets total bonded for all uTokens (used for a query) +func (k Keeper) getAllTotalBonded(ctx sdk.Context) sdk.Coins { + return store.SumCoins(k.prefixStore(ctx, keyPrefixTotalBonded), store.NoLastByte) } diff --git a/x/incentive/keeper/keeper.go b/x/incentive/keeper/keeper.go index 8ac03db163..005b496af7 100644 --- a/x/incentive/keeper/keeper.go +++ b/x/incentive/keeper/keeper.go @@ -2,9 +2,9 @@ package keeper import ( "github.com/cosmos/cosmos-sdk/codec" - sdk "github.com/cosmos/cosmos-sdk/types" - + prefixstore "github.com/cosmos/cosmos-sdk/store/prefix" storetypes "github.com/cosmos/cosmos-sdk/store/types" + sdk "github.com/cosmos/cosmos-sdk/types" authtypes "github.com/cosmos/cosmos-sdk/x/auth/types" "github.com/tendermint/tendermint/libs/log" @@ -46,3 +46,7 @@ func (k Keeper) ModuleBalance(ctx sdk.Context, denom string) sdk.Coin { func (k Keeper) KVStore(ctx sdk.Context) sdk.KVStore { return ctx.KVStore(k.storeKey) } + +func (k Keeper) prefixStore(ctx sdk.Context, prefix []byte) sdk.KVStore { + return prefixstore.NewStore(ctx.KVStore(k.storeKey), prefix) +} diff --git a/x/leverage/keeper/iter.go b/x/leverage/keeper/iter.go index ee888b0e94..97337592cc 100644 --- a/x/leverage/keeper/iter.go +++ b/x/leverage/keeper/iter.go @@ -2,6 +2,8 @@ package keeper import ( sdkmath "cosmossdk.io/math" + prefixstore "github.com/cosmos/cosmos-sdk/store/prefix" + storetypes "github.com/cosmos/cosmos-sdk/store/types" sdk "github.com/cosmos/cosmos-sdk/types" "github.com/umee-network/umee/v5/util" @@ -39,28 +41,6 @@ func (k Keeper) GetAllRegisteredTokens(ctx sdk.Context) []types.Token { return store.MustLoadAll[*types.Token](ctx.KVStore(k.storeKey), types.KeyPrefixRegisteredToken) } -// GetAllReserves returns all reserves. -func (k Keeper) GetAllReserves(ctx sdk.Context) sdk.Coins { - prefix := types.KeyPrefixReserveAmount - reserves := sdk.NewCoins() - - iterator := func(key, val []byte) error { - denom := types.DenomFromKey(key, prefix) - var amount sdkmath.Int - if err := amount.Unmarshal(val); err != nil { - // improperly marshaled reserve amount should never happen - return err - } - - reserves = reserves.Add(sdk.NewCoin(denom, amount)) - return nil - } - - util.Panic(k.iterate(ctx, prefix, iterator)) - - return reserves -} - // GetBorrowerBorrows returns an sdk.Coins object containing all open borrows // associated with an address. func (k Keeper) GetBorrowerBorrows(ctx sdk.Context, borrowerAddr sdk.AccAddress) sdk.Coins { @@ -197,22 +177,14 @@ func (k Keeper) SweepBadDebts(ctx sdk.Context) error { // GetAllUTokenSupply returns total supply of all uToken denoms. func (k Keeper) GetAllUTokenSupply(ctx sdk.Context) sdk.Coins { - prefix := types.KeyPrefixUtokenSupply - supplies := sdk.NewCoins() - - iterator := func(key, val []byte) error { - denom := types.DenomFromKey(key, prefix) - var amount sdkmath.Int - if err := amount.Unmarshal(val); err != nil { - // improperly marshaled utoken supply should never happen - return err - } - - supplies = supplies.Add(sdk.NewCoin(denom, amount)) - return nil - } + return store.SumCoins(k.prefixStore(ctx, types.KeyPrefixUtokenSupply), store.NoLastByte) +} - util.Panic(k.iterate(ctx, prefix, iterator)) +// GetAllReserves returns all reserves. +func (k Keeper) GetAllReserves(ctx sdk.Context) sdk.Coins { + return store.SumCoins(k.prefixStore(ctx, types.KeyPrefixReserveAmount), store.NoLastByte) +} - return supplies +func (k Keeper) prefixStore(ctx sdk.Context, p []byte) storetypes.KVStore { + return prefixstore.NewStore(ctx.KVStore(k.storeKey), p) }