Skip to content

Commit

Permalink
Merge pull request #28 from initia-labs/fix/avoid-map-iteration
Browse files Browse the repository at this point in the history 
feat: remove map iterations
  • Loading branch information
@beer-1
beer-1 authored Dec 5, 2023
2 parents 7c56f5e + 628a96e commit c6252bb
Show file tree
Hide file tree
Showing 4 changed files with 64 additions and 25 deletions.
39 changes: 25 additions & 14 deletions x/distribution/keeper/allocation.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,20 +60,23 @@ func (k Keeper) AllocateTokens(ctx sdk.Context, totalPreviousPower int64, bonded
communityTax := k.GetCommunityTax(ctx)
voteMultiplier := sdk.OneDec().Sub(communityTax)

rewardWeights, weightsSum := k.LoadRewardWeights(ctx)
validators, bondedTokens, bondedTokensSum := k.LoadBondedTokens(ctx, bondedVotes, rewardWeights)
// map iteration not guarantee the ordering,
// so we have to use array for iteration.
rewardWeights, rewardWeightMap, weightsSum := k.LoadRewardWeights(ctx)
validators, bondedTokens, bondedTokensSum := k.LoadBondedTokens(ctx, bondedVotes, rewardWeightMap)

// allocate rewards proportionally to reward power
for rewardDenom, rewardWeight := range rewardWeights {
poolFraction := rewardWeight.Quo(weightsSum)
for _, rewardWeight := range rewardWeights {
poolFraction := rewardWeight.Weight.Quo(weightsSum)
poolReward := feesCollected.MulDecTruncate(voteMultiplier).MulDecTruncate(poolFraction)

poolDenom := rewardDenom
poolDenom := rewardWeight.Denom
poolSize := bondedTokensSum[poolDenom]
for valAddr, amount := range bondedTokens[poolDenom] {
validator := validators[valAddr]

amountFraction := math.LegacyNewDecFromInt(amount).QuoInt(poolSize)
for _, bondedTokens := range bondedTokens[poolDenom] {
validator := validators[bondedTokens.ValAddr]

amountFraction := math.LegacyNewDecFromInt(bondedTokens.Amount).QuoInt(poolSize)
reward := poolReward.MulDecTruncate(amountFraction)

k.AllocateTokensToValidatorPool(ctx, validator, poolDenom, reward)
Expand All @@ -88,7 +91,7 @@ func (k Keeper) AllocateTokens(ctx sdk.Context, totalPreviousPower int64, bonded

// LoadRewardWeights load reward weights with its sum
func (k Keeper) LoadRewardWeights(ctx sdk.Context) (
map[string]sdk.Dec, sdk.Dec,
[]customtypes.RewardWeight, map[string]sdk.Dec, sdk.Dec,
) {
rewardWeights := k.GetRewardWeights(ctx)

Expand All @@ -100,18 +103,23 @@ func (k Keeper) LoadRewardWeights(ctx sdk.Context) (
weightsMap[rewardWeight.Denom] = rewardWeight.Weight
}

return weightsMap, weightsSum
return rewardWeights, weightsMap, weightsSum
}

type validatorBondedToken struct {
ValAddr string
Amount math.Int
}

// LoadBondedTokens build denom:(validator:amount) map
func (k Keeper) LoadBondedTokens(ctx sdk.Context, bondedVotes []abci.VoteInfo, rewardWeights map[string]sdk.Dec) (
map[string]stakingtypes.ValidatorI, map[string]map[string]math.Int, map[string]math.Int,
map[string]stakingtypes.ValidatorI, map[string][]validatorBondedToken, map[string]math.Int,
) {
numOfValidators := len(bondedVotes)
numOfDenoms := len(rewardWeights)

validators := make(map[string]stakingtypes.ValidatorI, numOfValidators)
bondedTokens := make(map[string]map[string]math.Int, numOfDenoms)
bondedTokens := make(map[string][]validatorBondedToken, numOfDenoms)
bondedTokensSum := make(map[string]math.Int, numOfDenoms)

for _, vote := range bondedVotes {
Expand All @@ -127,13 +135,16 @@ func (k Keeper) LoadBondedTokens(ctx sdk.Context, bondedVotes []abci.VoteInfo, r
}

if _, found := bondedTokens[token.Denom]; !found {
bondedTokens[token.Denom] = make(map[string]math.Int, numOfValidators)
bondedTokens[token.Denom] = make([]validatorBondedToken, 0, numOfValidators)
}
if _, found := bondedTokensSum[token.Denom]; !found {
bondedTokensSum[token.Denom] = sdk.ZeroInt()
}

bondedTokens[token.Denom][valAddr] = token.Amount
bondedTokens[token.Denom] = append(bondedTokens[token.Denom], validatorBondedToken{
ValAddr: valAddr,
Amount: token.Amount,
})
bondedTokensSum[token.Denom] = bondedTokensSum[token.Denom].Add(token.Amount)
}
}
Expand Down
23 changes: 17 additions & 6 deletions x/distribution/keeper/allocation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ func TestLoadRewardWeights(t *testing.T) {
}
input.DistKeeper.SetRewardWeights(ctx, weights)

loadedWeights, sum := input.DistKeeper.LoadRewardWeights(ctx)
_, loadedWeights, sum := input.DistKeeper.LoadRewardWeights(ctx)
require.Equal(t, sdk.NewDecWithPrec(3, 1), loadedWeights["aaa"])
require.Equal(t, sdk.NewDecWithPrec(4, 1), loadedWeights["bar"])
require.Equal(t, sdk.NewDecWithPrec(3, 1), loadedWeights["foo"])
Expand Down Expand Up @@ -90,14 +90,25 @@ func TestLoadBondedTokens(t *testing.T) {
},
}

rewardWeight, _ := input.DistKeeper.LoadRewardWeights(ctx)
_, rewardWeight, _ := input.DistKeeper.LoadRewardWeights(ctx)
validators, bondedTokens, bondedTokensSum := input.DistKeeper.LoadBondedTokens(ctx, votes, rewardWeight)
require.Equal(t, validator1, validators[string(valConsPk1.Address())])
require.Equal(t, validator2, validators[string(valConsPk2.Address())])
require.Equal(t, sdk.NewInt(3_000_000), bondedTokens["foo"][string(valConsPk1.Address())])
require.Equal(t, sdk.NewInt(5_000_000), bondedTokens["foo"][string(valConsPk2.Address())])
require.Equal(t, sdk.NewInt(5_000_000), bondedTokens["bar"][string(valConsPk1.Address())])
require.Equal(t, sdk.NewInt(3_000_000), bondedTokens["bar"][string(valConsPk2.Address())])
for _, val := range bondedTokens["foo"] {
if val.ValAddr == string(valConsPk1.Address()) {
require.Equal(t, sdk.NewInt(3_000_000), val.Amount)
} else {
sdk.NewInt(5_000_000)
}
}

for _, val := range bondedTokens["bar"] {
if val.ValAddr == string(valConsPk1.Address()) {
require.Equal(t, sdk.NewInt(5_000_000), val.Amount)
} else {
sdk.NewInt(3_000_000)
}
}
require.Equal(t, sdk.NewInt(8_000_000), bondedTokensSum["foo"])
require.Equal(t, sdk.NewInt(8_000_000), bondedTokensSum["bar"])
}
Expand Down
6 changes: 4 additions & 2 deletions x/move/abci.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@ func BeginBlocker(ctx sdk.Context, k keeper.Keeper) {

// get rewards from active validators
activeValidators := k.StakingKeeper.GetBondedValidatorsByPower(ctx)
rewardDenom := k.RewardKeeper.GetParams(ctx).RewardDenom

denoms := []string{}
rewardVecMap := make(map[string][]uint64)
valAddrVecMap := make(map[string][][]byte)
for _, activeValidator := range activeValidators {
Expand All @@ -37,11 +39,11 @@ func BeginBlocker(ctx sdk.Context, k keeper.Keeper) {
for _, pool := range rewardPools {
poolDenom := pool.Denom
if _, found := rewardVecMap[poolDenom]; !found {
denoms = append(denoms, poolDenom)
rewardVecMap[poolDenom] = make([]uint64, 0, len(activeValidators))
valAddrVecMap[poolDenom] = make([][]byte, 0, len(activeValidators))
}

rewardDenom := k.RewardKeeper.GetParams(ctx).RewardDenom
rewardAmount := pool.Coins.AmountOf(rewardDenom)
if rewardAmount.IsZero() {
continue
Expand All @@ -52,7 +54,7 @@ func BeginBlocker(ctx sdk.Context, k keeper.Keeper) {
}
}

for poolDenom := range rewardVecMap {
for _, poolDenom := range denoms {
rewardArg, err := vmtypes.SerializeUint64Vector(rewardVecMap[poolDenom])
if err != nil {
panic(err)
Expand Down
21 changes: 18 additions & 3 deletions x/move/keeper/staking.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -161,15 +161,20 @@ func (k Keeper) ApplyStakingDeltas(
ctx sdk.Context,
stakingDeltas []vmtypes.StakingDelta,
) error {
// keep the array to avoid map iteration.
delegationValAddrs := []string{}
undelegationValAddrs := []string{}
delegations := make(map[string]sdk.Coins)
undelegations := make(map[string]sdk.DecCoins)
for _, delta := range stakingDeltas {
valAddrStr := string(delta.Validator)
if _, found := delegations[valAddrStr]; !found {
delegations[valAddrStr] = sdk.NewCoins()
delegationValAddrs = append(delegationValAddrs, valAddrStr)
}
if _, found := undelegations[valAddrStr]; !found {
undelegations[valAddrStr] = sdk.NewDecCoins()
undelegationValAddrs = append(undelegationValAddrs, valAddrStr)
}

denom, err := types.DenomFromMetadataAddress(ctx, NewMoveBankKeeper(&k), delta.Metadata)
Expand All @@ -188,7 +193,8 @@ func (k Keeper) ApplyStakingDeltas(
}
}

for valAddrStr, delegationCoins := range delegations {
for _, valAddrStr := range delegationValAddrs {
delegationCoins := delegations[valAddrStr]
if !delegationCoins.IsZero() {
valAddr, err := sdk.ValAddressFromBech32(valAddrStr)
if err != nil {
Expand All @@ -202,9 +208,12 @@ func (k Keeper) ApplyStakingDeltas(
}
}

// keep denoms array to avoid map iteration.
denoms := []string{}
amountVecMap := make(map[string][]uint64)
valAddrVecMap := make(map[string][][]byte)
for valAddrStr, undelegationShares := range undelegations {
for _, valAddrStr := range undelegationValAddrs {
undelegationShares := undelegations[valAddrStr]
if !undelegationShares.IsZero() {
valAddr, err := sdk.ValAddressFromBech32(valAddrStr)
if err != nil {
Expand All @@ -222,13 +231,19 @@ func (k Keeper) ApplyStakingDeltas(
continue
}

if _, found := amountVecMap[amount.Denom]; !found {
denoms = append(denoms, amount.Denom)
amountVecMap[amount.Denom] = []uint64{}
valAddrVecMap[amount.Denom] = [][]byte{}
}

amountVecMap[amount.Denom] = append(amountVecMap[amount.Denom], amount.Amount.Uint64())
valAddrVecMap[amount.Denom] = append(valAddrVecMap[amount.Denom], []byte(valAddrStr))
}
}
}

for unbondingDenom := range amountVecMap {
for _, unbondingDenom := range denoms {
err := k.DepositUnbondingCoins(ctx, unbondingDenom, amountVecMap[unbondingDenom], valAddrVecMap[unbondingDenom])
if err != nil {
return err
Expand Down

0 comments on commit c6252bb

Please sign in to comment.