reduce code duplication

osmoutils/store_helper.go

@@ -8,6 +8,8 @@ import (
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+
 	"github.com/osmosis-labs/osmosis/osmomath"
 
 	"github.com/cosmos/cosmos-sdk/store"
@@ -222,3 +224,61 @@ func DeleteAllKeysFromPrefix(store store.KVStore, prefixKey []byte) {
 		prefixStore.Delete(iter.Key())
 	}
 }
+
+// GetCoinArrayFromPrefix returns all coins from the store that has the given prefix.
+func GetCoinArrayFromPrefix(ctx sdk.Context, storeKey storetypes.StoreKey, storePrefix []byte) []sdk.Coin {
+	coinArray := make([]sdk.Coin, 0)
+
+	store := ctx.KVStore(storeKey)
+	iterator := sdk.KVStorePrefixIterator(store, storePrefix)
+
+	defer iterator.Close()
+	for ; iterator.Valid(); iterator.Next() {
+		bz := iterator.Value()
+		coin := sdk.Coin{}
+		if err := coin.Unmarshal(bz); err == nil {
+			coinArray = append(coinArray, coin)
+		}
+	}
+
+	return coinArray
+}
+
+// GetCoinByDenomFromPrefix returns the coin from the store that has the given prefix and denom.
+// If the denom is not found, a zero coin is returned.
+func GetCoinByDenomFromPrefix(ctx sdk.Context, storeKey storetypes.StoreKey, storePrefix []byte, denom string) (sdk.Coin, error) {
+	store := prefix.NewStore(ctx.KVStore(storeKey), storePrefix)
+	key := append(storePrefix, []byte(denom)...)
+
+	bz := store.Get(key)
+	if len(bz) == 0 {
+		return sdk.NewCoin(denom, osmomath.ZeroInt()), nil
+	}
+
+	coin := sdk.Coin{}
+	if err := coin.Unmarshal(bz); err != nil {
+		return sdk.NewCoin(denom, osmomath.ZeroInt()), err
+	}
+
+	return coin, nil
+}
+
+// IncreaseCoinByDenomFromPrefix increases the coin from the store that has the given prefix and denom by the specified amount.
+func IncreaseCoinByDenomFromPrefix(ctx sdk.Context, storeKey storetypes.StoreKey, storePrefix []byte, denom string, increasedAmt osmomath.Int) error {
+	store := prefix.NewStore(ctx.KVStore(storeKey), storePrefix)
+	key := append(storePrefix, []byte(denom)...)
+
+	coin, err := GetCoinByDenomFromPrefix(ctx, storeKey, storePrefix, denom)
+	if err != nil {
+		return err
+	}
+
+	coin.Amount = coin.Amount.Add(increasedAmt)
+	bz, err := coin.Marshal()
+	if err != nil {
+		return err
+	}
+
+	store.Set(key, bz)
+	return nil
+}

x/poolmanager/protorev.go

@@ -7,118 +7,40 @@ import (
 	"github.com/osmosis-labs/osmosis/osmomath"
 	"github.com/osmosis-labs/osmosis/osmoutils"
 	"github.com/osmosis-labs/osmosis/v21/x/poolmanager/types"
-
-	"github.com/cosmos/cosmos-sdk/store/prefix"
 )
 
 // GetTakerFeeTrackerForStakers returns the taker fee for stakers tracker for all denoms that has been
 // collected since the accounting height.
 func (k Keeper) GetTakerFeeTrackerForStakers(ctx sdk.Context) []sdk.Coin {
-	takerFeesForStakers := make([]sdk.Coin, 0)
-
-	store := ctx.KVStore(k.storeKey)
-	iterator := sdk.KVStorePrefixIterator(store, types.KeyTakerFeeStakersProtoRevArray)
-
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		bz := iterator.Value()
-		takerFeeForStakers := sdk.Coin{}
-		if err := takerFeeForStakers.Unmarshal(bz); err == nil {
-			takerFeesForStakers = append(takerFeesForStakers, takerFeeForStakers)
-		}
-	}
-
-	return takerFeesForStakers
+	return osmoutils.GetCoinArrayFromPrefix(ctx, k.storeKey, types.KeyTakerFeeStakersProtoRevArray)
 }
 
 // GetTakerFeeTrackerForStakersByDenom returns the taker fee for stakers tracker for the specified denom that has been
-// collected since the accounting height.
+// collected since the accounting height. If the denom is not found, a zero coin is returned.
 func (k Keeper) GetTakerFeeTrackerForStakersByDenom(ctx sdk.Context, denom string) (sdk.Coin, error) {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyTakerFeeStakersProtoRevArray)
-	key := types.GetKeyPrefixTakerFeeStakersProtoRevByDenom(denom)
-
-	bz := store.Get(key)
-	if len(bz) == 0 {
-		return sdk.NewCoin(denom, osmomath.ZeroInt()), nil
-	}
-
-	takerFeeForStakers := sdk.Coin{}
-	if err := takerFeeForStakers.Unmarshal(bz); err != nil {
-		return sdk.NewCoin(denom, osmomath.ZeroInt()), err
-	}
-
-	return takerFeeForStakers, nil
+	return osmoutils.GetCoinByDenomFromPrefix(ctx, k.storeKey, types.KeyTakerFeeStakersProtoRevArray, denom)
 }
 
 // UpdateTakerFeeTrackerForStakersByDenom increases the take fee for stakers tracker for the specified denom by the specified amount.
 func (k Keeper) UpdateTakerFeeTrackerForStakersByDenom(ctx sdk.Context, denom string, increasedAmt osmomath.Int) error {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyTakerFeeStakersProtoRevArray)
-	key := types.GetKeyPrefixTakerFeeStakersProtoRevByDenom(denom)
-
-	takerFeeForStakers, _ := k.GetTakerFeeTrackerForStakersByDenom(ctx, denom)
-	takerFeeForStakers.Amount = takerFeeForStakers.Amount.Add(increasedAmt)
-	bz, err := takerFeeForStakers.Marshal()
-	if err != nil {
-		return err
-	}
-
-	store.Set(key, bz)
-	return nil
+	return osmoutils.IncreaseCoinByDenomFromPrefix(ctx, k.storeKey, types.KeyTakerFeeStakersProtoRevArray, denom, increasedAmt)
 }
 
 // GetTakerFeeTrackerForCommunityPool returns the taker fee for community pool tracker for all denoms that has been
 // collected since the accounting height.
 func (k Keeper) GetTakerFeeTrackerForCommunityPool(ctx sdk.Context) []sdk.Coin {
-	takerFeesForCommunityPool := make([]sdk.Coin, 0)
-
-	store := ctx.KVStore(k.storeKey)
-	iterator := sdk.KVStorePrefixIterator(store, types.KeyTakerFeeCommunityPoolProtoRevArray)
-
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		bz := iterator.Value()
-		takerFeeForCommunityPool := sdk.Coin{}
-		if err := takerFeeForCommunityPool.Unmarshal(bz); err == nil {
-			takerFeesForCommunityPool = append(takerFeesForCommunityPool, takerFeeForCommunityPool)
-		}
-	}
-
-	return takerFeesForCommunityPool
+	return osmoutils.GetCoinArrayFromPrefix(ctx, k.storeKey, types.KeyTakerFeeCommunityPoolProtoRevArray)
 }
 
 // GetTakerFeeTrackerForCommunityPoolByDenom returns the taker fee for community pool tracker for the specified denom that has been
-// collected since the accounting height.
+// collected since the accounting height. If the denom is not found, a zero coin is returned.
 func (k Keeper) GetTakerFeeTrackerForCommunityPoolByDenom(ctx sdk.Context, denom string) (sdk.Coin, error) {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyTakerFeeCommunityPoolProtoRevArray)
-	key := types.GetKeyPrefixTakerFeeCommunityPoolProtoRevByDenom(denom)
-
-	bz := store.Get(key)
-	if len(bz) == 0 {
-		return sdk.NewCoin(denom, osmomath.ZeroInt()), nil
-	}
-
-	takerFeeForCommunityPool := sdk.Coin{}
-	if err := takerFeeForCommunityPool.Unmarshal(bz); err != nil {
-		return sdk.NewCoin(denom, osmomath.ZeroInt()), err
-	}
-
-	return takerFeeForCommunityPool, nil
+	return osmoutils.GetCoinByDenomFromPrefix(ctx, k.storeKey, types.KeyTakerFeeCommunityPoolProtoRevArray, denom)
 }
 
 // UpdateTakerFeeTrackerForCommunityPoolByDenom increases the take fee for community pool tracker for the specified denom by the specified amount.
 func (k Keeper) UpdateTakerFeeTrackerForCommunityPoolByDenom(ctx sdk.Context, denom string, increasedAmt osmomath.Int) error {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyTakerFeeCommunityPoolProtoRevArray)
-	key := types.GetKeyPrefixTakerFeeCommunityPoolProtoRevByDenom(denom)
-
-	takerFeeForCommunityPool, _ := k.GetTakerFeeTrackerForCommunityPoolByDenom(ctx, denom)
-	takerFeeForCommunityPool.Amount = takerFeeForCommunityPool.Amount.Add(increasedAmt)
-	bz, err := takerFeeForCommunityPool.Marshal()
-	if err != nil {
-		return err
-	}
-
-	store.Set(key, bz)
-	return nil
+	return osmoutils.IncreaseCoinByDenomFromPrefix(ctx, k.storeKey, types.KeyTakerFeeCommunityPoolProtoRevArray, denom, increasedAmt)
 }
 
 // GetTakerFeeTrackerStartHeight gets the height from which we started accounting for taker fees.

x/poolmanager/types/keys.go

@@ -108,11 +108,3 @@ func ParseDenomTradePairKey(key []byte) (denom0, denom1 string, err error) {
 
 	return denom0, denom1, nil
 }
-
-func GetKeyPrefixTakerFeeStakersProtoRevByDenom(denom string) []byte {
-	return append(KeyTakerFeeStakersProtoRevArray, []byte(denom)...)
-}
-
-func GetKeyPrefixTakerFeeCommunityPoolProtoRevByDenom(denom string) []byte {
-	return append(KeyTakerFeeCommunityPoolProtoRevArray, []byte(denom)...)
-}

x/protorev/keeper/statistics.go

@@ -49,41 +49,20 @@ func (k Keeper) IncrementNumberOfTrades(ctx sdk.Context) error {
 	return nil
 }
 
-// GetProfitsByDenom returns the profits made by the ProtoRev module for the given denom
-func (k Keeper) GetProfitsByDenom(ctx sdk.Context, denom string) (sdk.Coin, error) {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixProfitByDenom)
-	key := types.GetKeyPrefixProfitByDenom(denom)
-
-	bz := store.Get(key)
-	if len(bz) == 0 {
-		return sdk.NewCoin(denom, osmomath.ZeroInt()), fmt.Errorf("no profits for denom %s", denom)
-	}
-
-	profits := sdk.Coin{}
-	if err := profits.Unmarshal(bz); err != nil {
-		return sdk.NewCoin(denom, osmomath.ZeroInt()), err
-	}
-
-	return profits, nil
-}
-
 // GetAllProfits returns all of the profits made by the ProtoRev module.
 func (k Keeper) GetAllProfits(ctx sdk.Context) []sdk.Coin {
-	profits := make([]sdk.Coin, 0)
-
-	store := ctx.KVStore(k.storeKey)
-	iterator := sdk.KVStorePrefixIterator(store, types.KeyPrefixProfitByDenom)
+	return osmoutils.GetCoinArrayFromPrefix(ctx, k.storeKey, types.KeyPrefixProfitByDenom)
+}
 
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		bz := iterator.Value()
-		profit := sdk.Coin{}
-		if err := profit.Unmarshal(bz); err == nil {
-			profits = append(profits, profit)
-		}
-	}
+// GetProfitsByDenom returns the profits made by the ProtoRev module for the given denom.
+// If the denom is not found, a zero coin is returned.
+func (k Keeper) GetProfitsByDenom(ctx sdk.Context, denom string) (sdk.Coin, error) {
+	return osmoutils.GetCoinByDenomFromPrefix(ctx, k.storeKey, types.KeyPrefixProfitByDenom, denom)
+}
 
-	return profits
+// UpdateProfitsByDenom updates the profits made by the ProtoRev module for the given denom
+func (k Keeper) UpdateProfitsByDenom(ctx sdk.Context, denom string, tradeProfit osmomath.Int) error {
+	return osmoutils.IncreaseCoinByDenomFromPrefix(ctx, k.storeKey, types.KeyPrefixProfitByDenom, denom, tradeProfit)
 }
 
 func (k Keeper) SetCyclicArbProfitTrackerValue(ctx sdk.Context, cyclicArbProfits sdk.Coins) {
@@ -126,22 +105,6 @@ func (k Keeper) SetCyclicArbProfitTrackerStartHeight(ctx sdk.Context, startHeigh
 	osmoutils.MustSet(ctx.KVStore(k.storeKey), types.KeyCyclicArbTrackerStartHeight, &gogotypes.Int64Value{Value: startHeight})
 }
 
-// UpdateProfitsByDenom updates the profits made by the ProtoRev module for the given denom
-func (k Keeper) UpdateProfitsByDenom(ctx sdk.Context, denom string, tradeProfit osmomath.Int) error {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixProfitByDenom)
-	key := types.GetKeyPrefixProfitByDenom(denom)
-
-	profits, _ := k.GetProfitsByDenom(ctx, denom)
-	profits.Amount = profits.Amount.Add(tradeProfit)
-	bz, err := profits.Marshal()
-	if err != nil {
-		return err
-	}
-
-	store.Set(key, bz)
-	return nil
-}
-
 // GetAllRoutes returns all of the routes that the ProtoRev module has traded on
 func (k Keeper) GetAllRoutes(ctx sdk.Context) ([][]uint64, error) {
 	routes := make([][]uint64, 0)

x/protorev/types/keys.go

@@ -114,11 +114,6 @@ func GetKeyPrefixRouteForTokenPair(tokenA, tokenB string) []byte {
 	return append(KeyPrefixTokenPairRoutes, []byte(tokenA+"|"+tokenB)...)
 }
 
-// Returns the key needed to fetch the profit by coin
-func GetKeyPrefixProfitByDenom(denom string) []byte {
-	return append(KeyPrefixProfitByDenom, []byte(denom)...)
-}
-
 // Returns the key needed to fetch the number of trades by route
 func GetKeyPrefixTradesByRoute(route []uint64) []byte {
 	return append(KeyPrefixTradesByRoute, CreateRouteKey(route)...)