refactor: rename TotalDelegatorShares and implement genesis things

milkyway-labs · Nov 14, 2024 · 2431179 · 2431179
1 parent 78c9d61
commit 2431179
Show file tree

Hide file tree

Showing 13 changed files with 284 additions and 232 deletions.
diff --git a/proto/milkyway/rewards/v1/genesis.proto b/proto/milkyway/rewards/v1/genesis.proto
@@ -270,7 +270,8 @@ message GenesisState {
       operator_accumulated_commissions = 9
       [ (gogoproto.nullable) = false, (amino.dont_omitempty) = true ];
 
-  // total_delegator_shares defines the total delegator shares at genesis.
-  repeated TotalDelegatorShares total_delegator_shares = 10
-      [ (gogoproto.nullable) = false, (amino.dont_omitempty) = true ];
+  // pool_service_total_delegator_shares defines the total delegator shares at
+  // genesis.
+  repeated PoolServiceTotalDelegatorShares pool_service_total_delegator_shares =
+      10 [ (gogoproto.nullable) = false, (amino.dont_omitempty) = true ];
 }
diff --git a/proto/milkyway/rewards/v1/models.proto b/proto/milkyway/rewards/v1/models.proto
@@ -230,9 +230,9 @@ message DelegationDelegatorReward {
 
 // ---------------------------------------------------------------------------
 
-// TotalDelegatorShares represents the total delegator shares for a pool-service
-// pair.
-message TotalDelegatorShares {
+// PoolServiceTotalDelegatorShares represents the total delegator shares for a
+// pool-service pair.
+message PoolServiceTotalDelegatorShares {
   uint32 pool_id = 1 [ (gogoproto.customname) = "PoolID" ];
   uint32 service_id = 2 [ (gogoproto.customname) = "ServiceID" ];
   repeated cosmos.base.v1beta1.DecCoin shares = 3 [

diff --git a/x/rewards/keeper/alias_functions.go b/x/rewards/keeper/alias_functions.go
@@ -420,9 +420,9 @@ func (k *Keeper) GetServiceDelegationRewards(
 
 // --------------------------------------------------------------------------------------------------------------------
 
-// GetPoolTotalDelegatorShares returns the total delegator shares for a pool-service pair.
-func (k *Keeper) GetPoolTotalDelegatorShares(ctx context.Context, poolID, serviceID uint32) (sdk.DecCoins, error) {
-	shares, err := k.PoolTotalDelegatorShares.Get(ctx, collections.Join(poolID, serviceID))
+// GetPoolServiceTotalDelegatorShares returns the total delegator shares for a pool-service pair.
+func (k *Keeper) GetPoolServiceTotalDelegatorShares(ctx context.Context, poolID, serviceID uint32) (sdk.DecCoins, error) {
+	shares, err := k.PoolServiceTotalDelegatorShares.Get(ctx, collections.Join(poolID, serviceID))
 	if err != nil {
 		if errors.IsOf(err, collections.ErrNotFound) {
 			return nil, nil
@@ -432,33 +432,33 @@ func (k *Keeper) GetPoolTotalDelegatorShares(ctx context.Context, poolID, servic
 	return shares.Shares, nil
 }
 
-// SetPoolTotalDelegatorShares sets the total delegator shares for a pool-service pair.
-func (k *Keeper) SetPoolTotalDelegatorShares(ctx context.Context, poolID, serviceID uint32, shares sdk.DecCoins) error {
-	return k.PoolTotalDelegatorShares.Set(
+// SetPoolServiceTotalDelegatorShares sets the total delegator shares for a pool-service pair.
+func (k *Keeper) SetPoolServiceTotalDelegatorShares(ctx context.Context, poolID, serviceID uint32, shares sdk.DecCoins) error {
+	return k.PoolServiceTotalDelegatorShares.Set(
 		ctx,
 		collections.Join(poolID, serviceID),
-		types.NewTotalDelegatorShares(poolID, serviceID, shares),
+		types.NewPoolServiceTotalDelegatorShares(poolID, serviceID, shares),
 	)
 }
 
-// IncrementPoolTotalDelegatorShares increments the total delegator shares for a pool-service pair.
-func (k *Keeper) IncrementPoolTotalDelegatorShares(
+// IncrementPoolServiceTotalDelegatorShares increments the total delegator shares for a pool-service pair.
+func (k *Keeper) IncrementPoolServiceTotalDelegatorShares(
 	ctx context.Context, poolID, serviceID uint32, shares sdk.DecCoins,
 ) error {
-	prevShares, err := k.GetPoolTotalDelegatorShares(ctx, poolID, serviceID)
+	prevShares, err := k.GetPoolServiceTotalDelegatorShares(ctx, poolID, serviceID)
 	if err != nil {
 		return err
 	}
-	return k.SetPoolTotalDelegatorShares(ctx, poolID, serviceID, prevShares.Add(shares...))
+	return k.SetPoolServiceTotalDelegatorShares(ctx, poolID, serviceID, prevShares.Add(shares...))
 }
 
-// DecrementPoolTotalDelegatorShares decrements the total delegator shares for a pool-service pair.
-func (k *Keeper) DecrementPoolTotalDelegatorShares(
+// DecrementPoolServiceTotalDelegatorShares decrements the total delegator shares for a pool-service pair.
+func (k *Keeper) DecrementPoolServiceTotalDelegatorShares(
 	ctx context.Context, poolID, serviceID uint32, shares sdk.DecCoins,
 ) error {
-	prevShares, err := k.GetPoolTotalDelegatorShares(ctx, poolID, serviceID)
+	prevShares, err := k.GetPoolServiceTotalDelegatorShares(ctx, poolID, serviceID)
 	if err != nil {
 		return err
 	}
-	return k.SetPoolTotalDelegatorShares(ctx, poolID, serviceID, prevShares.Sub(shares))
+	return k.SetPoolServiceTotalDelegatorShares(ctx, poolID, serviceID, prevShares.Sub(shares))
 }
diff --git a/x/rewards/keeper/genesis.go b/x/rewards/keeper/genesis.go
@@ -107,6 +107,15 @@ func (k *Keeper) ExportGenesis(ctx sdk.Context) (*types.GenesisState, error) {
 		return nil, err
 	}
 
+	var totalDelShares []types.PoolServiceTotalDelegatorShares
+	err = k.PoolServiceTotalDelegatorShares.Walk(ctx, nil, func(_ collections.Pair[uint32, uint32], shares types.PoolServiceTotalDelegatorShares) (stop bool, err error) {
+		totalDelShares = append(totalDelShares, shares)
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
 	return types.NewGenesisState(
 		params,
 		nextRewardsPlanID,
@@ -117,6 +126,7 @@ func (k *Keeper) ExportGenesis(ctx sdk.Context) (*types.GenesisState, error) {
 		operatorRecords,
 		serviceRecords,
 		operatorAccumulatedCommissions,
+		totalDelShares,
 	), nil
 }
 
@@ -409,8 +419,12 @@ func (k *Keeper) InitGenesis(ctx sdk.Context, state types.GenesisState) error {
 			totalOutstandingRewardsTruncated)
 	}
 
-	// TODO: total delegator shares
-
+	for _, shares := range state.PoolServiceTotalDelegatorShares {
+		err = k.SetPoolServiceTotalDelegatorShares(ctx, shares.PoolID, shares.ServiceID, shares.Shares)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 

diff --git a/x/rewards/keeper/hooks.go b/x/rewards/keeper/hooks.go
@@ -72,7 +72,7 @@ func (h Hooks) BeforePoolDelegationSharesModified(ctx sdk.Context, poolID uint32
 			return err
 		}
 
-		err = h.k.DecrementPoolTotalDelegatorShares(ctx, poolID, serviceID, del.Shares)
+		err = h.k.DecrementPoolServiceTotalDelegatorShares(ctx, poolID, serviceID, del.Shares)
 		if err != nil {
 			return err
 		}
@@ -102,7 +102,7 @@ func (h Hooks) AfterPoolDelegationModified(ctx sdk.Context, poolID uint32, deleg
 		if err != nil {
 			return err
 		}
-		err = h.k.IncrementPoolTotalDelegatorShares(ctx, poolID, serviceID, del.Shares)
+		err = h.k.IncrementPoolServiceTotalDelegatorShares(ctx, poolID, serviceID, del.Shares)
 		if err != nil {
 			return err
 		}

diff --git a/x/rewards/keeper/keeper.go b/x/rewards/keeper/keeper.go
@@ -45,7 +45,7 @@ type Keeper struct {
 	// (poolID, serviceID) -> outstanding rewards
 	PoolOutstandingRewards collections.Map[collections.Pair[uint32, uint32], types.OutstandingRewards]
 	// (poolID, serviceID) -> total delegator shares
-	PoolTotalDelegatorShares collections.Map[collections.Pair[uint32, uint32], types.TotalDelegatorShares]
+	PoolServiceTotalDelegatorShares collections.Map[collections.Pair[uint32, uint32], types.PoolServiceTotalDelegatorShares]
 
 	OperatorAccumulatedCommissions collections.Map[uint32, types.AccumulatedCommission]
 	OperatorDelegatorStartingInfos collections.Map[collections.Pair[uint32, sdk.AccAddress], types.DelegatorStartingInfo]
@@ -148,12 +148,12 @@ func NewKeeper(
 			collections.PairKeyCodec(collections.Uint32Key, collections.Uint32Key),
 			codec.CollValue[types.OutstandingRewards](cdc),
 		),
-		PoolTotalDelegatorShares: collections.NewMap(
+		PoolServiceTotalDelegatorShares: collections.NewMap(
 			sb,
-			types.PoolTotalDelegatorSharesKeyPrefix,
-			"pool_total_delegator_shares",
+			types.PoolServiceTotalDelegatorSharesKeyPrefix,
+			"pool_service_total_delegator_shares",
 			collections.PairKeyCodec(collections.Uint32Key, collections.Uint32Key),
-			codec.CollValue[types.TotalDelegatorShares](cdc),
+			codec.CollValue[types.PoolServiceTotalDelegatorShares](cdc),
 		),
 		OperatorAccumulatedCommissions: collections.NewMap(
 			sb,

diff --git a/x/rewards/keeper/target.go b/x/rewards/keeper/target.go
@@ -103,7 +103,7 @@ func (k *Keeper) IncrementPoolPeriod(
 	// calculate current ratio
 	var current types.DecPools
 
-	totalShares, err := k.GetPoolTotalDelegatorShares(ctx, pool.GetID(), serviceID)
+	totalShares, err := k.GetPoolServiceTotalDelegatorShares(ctx, pool.GetID(), serviceID)
 	if err != nil {
 		return 0, err
 	}

diff --git a/x/rewards/types/genesis.go b/x/rewards/types/genesis.go
@@ -19,6 +19,7 @@ func NewGenesisState(
 	operatorsRecords,
 	servicesRecords DelegationTypeRecords,
 	operatorAccumulatedCommissionRecords []OperatorAccumulatedCommissionRecord,
+	poolServiceTotalDelShares []PoolServiceTotalDelegatorShares,
 ) *GenesisState {
 	return &GenesisState{
 		Params:                         params,
@@ -30,6 +31,7 @@ func NewGenesisState(
 		OperatorsRecords:               operatorsRecords,
 		ServicesRecords:                servicesRecords,
 		OperatorAccumulatedCommissions: operatorAccumulatedCommissionRecords,
+		PoolServiceTotalDelegatorShares:           poolServiceTotalDelShares,
 	}
 }
 
@@ -47,6 +49,7 @@ func DefaultGenesis() *GenesisState {
 			[]OutstandingRewardsRecord{}, []HistoricalRewardsRecord{}, []CurrentRewardsRecord{},
 			[]DelegatorStartingInfoRecord{}),
 		[]OperatorAccumulatedCommissionRecord{},
+		[]PoolServiceTotalDelegatorShares{},
 	)
 }
 
@@ -73,6 +76,19 @@ func (genState *GenesisState) Validate(unpacker codectypes.AnyUnpacker) error {
 			return fmt.Errorf("invalid rewards plan at index %d: %w", i, err)
 		}
 	}
+
+	for _, shares := range genState.PoolServiceTotalDelegatorShares {
+		if shares.PoolID == 0 {
+			return fmt.Errorf("pool ID must not be 0")
+		}
+		if shares.ServiceID == 0 {
+			return fmt.Errorf("service ID must not be 0")
+		}
+		err = shares.Shares.Validate()
+		if err != nil {
+			return fmt.Errorf("invalid pool service total delegator shares: %w", err)
+		}
+	}
 	return nil
 }