feat: use more store functions in incentive (#2101)

* feat: use more store functions in incentive * Update x/incentive/keeper/iter.go Co-authored-by: Adam Moser <63419657+toteki@users.noreply.github.com> * add SumCoins to store/iter * add prefixStore method * use store.SumCoins * use store.SumCoins in incentive * add tests --------- Co-authored-by: Adam Moser <63419657+toteki@users.noreply.github.com>
umee-network · Jun 22, 2023 · fc4fd4d · fc4fd4d
1 parent 5c40f6f
commit fc4fd4d
Show file tree

Hide file tree

Showing 6 changed files with 94 additions and 134 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)
 }