feat: add delegations-by-pool-id key

x/restaking/keeper/pool_restaking.go

@@ -9,36 +9,29 @@ import (
 )
 
 // SavePoolDelegation stores the given pool delegation in the store
-func (k *Keeper) SavePoolDelegation(ctx sdk.Context, delegation types.PoolDelegation) error {
+func (k *Keeper) SavePoolDelegation(ctx sdk.Context, delegation types.PoolDelegation) {
 	store := ctx.KVStore(k.storeKey)
 
-	delegationBz, err := k.cdc.Marshal(&delegation)
-	if err != nil {
-		return err
-	}
+	// Marshal and store the delegation
+	delegationBz := types.MustMarshalPoolDelegation(k.cdc, delegation)
+	store.Set(types.UserPoolDelegationStoreKey(delegation.UserAddress, delegation.PoolID), delegationBz)
 
-	store.Set(types.UserPoolDelegationStoreKey(delegation.PoolID, delegation.UserAddress), delegationBz)
-	return nil
+	// Store the delegation in the delegations by pool ID store
+	store.Set(types.DelegationsByPoolIDStoreKey(delegation.PoolID, delegation.UserAddress), []byte{})
 }
 
 // GetPoolDelegation retrieves the delegation for the given user and pool
 // If the delegation does not exist, false is returned instead
-func (k *Keeper) GetPoolDelegation(ctx sdk.Context, poolID uint32, userAddress string) (types.PoolDelegation, bool, error) {
+func (k *Keeper) GetPoolDelegation(ctx sdk.Context, poolID uint32, userAddress string) (types.PoolDelegation, bool) {
 	// Get the delegation amount from the store
 	store := ctx.KVStore(k.storeKey)
-	delegationAmountBz := store.Get(types.UserPoolDelegationStoreKey(poolID, userAddress))
+	delegationAmountBz := store.Get(types.UserPoolDelegationStoreKey(userAddress, poolID))
 	if delegationAmountBz == nil {
-		return types.PoolDelegation{}, false, nil
+		return types.PoolDelegation{}, false
 	}
 
 	// Parse the delegation amount
-	var delegation types.PoolDelegation
-	err := k.cdc.Unmarshal(delegationAmountBz, &delegation)
-	if err != nil {
-		return types.PoolDelegation{}, false, err
-	}
-
-	return delegation, true, nil
+	return types.MustUnmarshalPoolDelegation(k.cdc, delegationAmountBz), true
 }
 
 // AddPoolTokensAndShares adds the given amount of tokens to the pool and returns the added shares
@@ -72,10 +65,7 @@ func (k *Keeper) DelegateToPool(ctx sdk.Context, amount sdk.Coin, delegator stri
 	}
 
 	// Get or create the delegation object and call the appropriate hook if present
-	delegation, found, err := k.GetPoolDelegation(ctx, pool.ID, delegator)
-	if err != nil {
-		return sdkmath.LegacyZeroDec(), err
-	}
+	delegation, found := k.GetPoolDelegation(ctx, pool.ID, delegator)
 
 	if found {
 		// Delegation was found
@@ -120,10 +110,7 @@ func (k *Keeper) DelegateToPool(ctx sdk.Context, amount sdk.Coin, delegator stri
 
 	// Update delegation
 	delegation.Shares = delegation.Shares.Add(newShares)
-	err = k.SavePoolDelegation(ctx, delegation)
-	if err != nil {
-		return newShares, err
-	}
+	k.SavePoolDelegation(ctx, delegation)
 
 	// Call the after-modification hook
 	err = k.AfterPoolDelegationModified(ctx, pool.ID, delegator)

x/restaking/types/keys.go

@@ -1,9 +1,8 @@
 package types
 
 import (
-	"time"
-
-	sdk "github.com/cosmos/cosmos-sdk/types"
+	"bytes"
+	"fmt"
 
 	poolstypes "github.com/milkyway-labs/milkyway/x/pools/types"
 )
@@ -18,6 +17,7 @@ var (
 
 	PoolDelegationPrefix          = []byte{0xa1}
 	UnbondingPoolDelegationPrefix = []byte{0xa2}
+	PoolDelegationsByPoolIDPrefix = []byte{0x71}
 
 	ServiceDelegationPrefix          = []byte{0xb1}
 	UnbondingServiceDelegationPrefix = []byte{0xb2}
@@ -26,30 +26,42 @@ var (
 	UnbondingOperatorDelegationPrefix = []byte{0xc2}
 )
 
-// PoolDelegationsStorePrefix returns the prefix used to store all the delegations to a given pool
-func PoolDelegationsStorePrefix(poolID uint32) []byte {
-	return append(PoolDelegationPrefix, poolstypes.GetPoolIDBytes(poolID)...)
+// UserPoolDelegationsStorePrefix returns the prefix used to store all the delegations to a given pool
+func UserPoolDelegationsStorePrefix(userAddress string) []byte {
+	return append(PoolDelegationPrefix, []byte(userAddress)...)
 }
 
 // UserPoolDelegationStoreKey returns the key used to store the delegation of a user to a given pool
-func UserPoolDelegationStoreKey(poolID uint32, delegator string) []byte {
-	return append(PoolDelegationsStorePrefix(poolID), []byte(delegator)...)
+func UserPoolDelegationStoreKey(delegator string, poolID uint32) []byte {
+	return append(UserPoolDelegationsStorePrefix(delegator), poolstypes.GetPoolIDBytes(poolID)...)
 }
 
-// UnbondingDelegationsByTimeStorePrefix returns the prefix used to store all the unbonding delegations
-// that expire at the given time
-func UnbondingDelegationsByTimeStorePrefix(endTime time.Time) []byte {
-	return append(UnbondingPoolDelegationPrefix, sdk.FormatTimeBytes(endTime)...)
+// DelegationsByPoolIDStorePrefix returns the prefix used to store the delegations to a given pool
+func DelegationsByPoolIDStorePrefix(poolID uint32) []byte {
+	return append(PoolDelegationsByPoolIDPrefix, poolstypes.GetPoolIDBytes(poolID)...)
 }
 
-// UnbondingPoolDelegationsStorePrefix returns the prefix used to store all the unbonding delegations
-// to a given pool that expire at the given time
-func UnbondingPoolDelegationsStorePrefix(poolID uint32, endTime time.Time) []byte {
-	return append(UnbondingDelegationsByTimeStorePrefix(endTime), poolstypes.GetPoolIDBytes(poolID)...)
+// DelegationsByPoolIDStoreKey returns the key used to store the delegations to a given pool
+func DelegationsByPoolIDStoreKey(poolID uint32, delegatorAddress string) []byte {
+	return append(DelegationsByPoolIDStorePrefix(poolID), []byte(delegatorAddress)...)
 }
 
-// UserUnbondingPoolDelegationStoreKey returns the key used to store the unbonding delegation of a user
-// to a given pool that expires at the given time
-func UserUnbondingPoolDelegationStoreKey(poolID uint32, delegator string, endTime time.Time) []byte {
-	return append(UnbondingPoolDelegationsStorePrefix(poolID, endTime), []byte(delegator)...)
+// ParseDelegationsByPoolIDKey parses the pool ID and delegator address from the given key
+func ParseDelegationsByPoolIDKey(bz []byte) (poolID uint32, delegatorAddress string, err error) {
+	prefixLength := len(PoolDelegationsByPoolIDPrefix)
+	if prefix := bz[:prefixLength]; !bytes.Equal(prefix, PoolDelegationsByPoolIDPrefix) {
+		return 0, "", fmt.Errorf("invalid prefix; expected: %X, got: %x", PoolDelegationsByPoolIDPrefix, prefix)
+	}
+
+	// Remove the prefix
+	bz = bz[prefixLength:]
+
+	// Read the pool ID
+	poolID = poolstypes.GetPoolIDFromBytes(bz[:4])
+	bz = bz[4:]
+
+	// Read the delegator address
+	delegatorAddress = string(bz)
+
+	return poolID, delegatorAddress, nil
 }