move slot limit code to prepareOrders() and fix RawStake of validator

core/blockchain.go

@@ -2568,29 +2568,7 @@ func (bc *BlockChain) UpdateValidatorVotingPower(
 		if snapshot, err := bc.ReadValidatorSnapshotAtEpoch(
 			newEpochSuperCommittee.Epoch, key,
 		); err == nil {
-			wrapper := snapshot.Validator
-			spread := numeric.ZeroDec()
-			if len(wrapper.SlotPubKeys) > 0 {
-				spread = numeric.NewDecFromBigInt(wrapper.TotalDelegation()).
-					QuoInt64(int64(len(wrapper.SlotPubKeys)))
-			}
-			instance := shard.Schedule.InstanceForEpoch(newEpochSuperCommittee.Epoch)
-			limitedSlotsCount := 0 // limited slots count for HIP16
-			slotsLimit := instance.SlotsLimit()
-			if slotsLimit > 0 {
-				shardCount := big.NewInt(int64(instance.NumShards()))
-				shardSlotsCount := make([]int, shardCount.Uint64()) // number slots keys on each shard
-				for _, pubkey := range wrapper.SlotPubKeys {
-					shardIndex := new(big.Int).Mod(pubkey.Big(), shardCount).Uint64()
-					shardSlotsCount[shardIndex] += 1
-					if shardSlotsCount[shardIndex] > slotsLimit {
-						continue
-					}
-					limitedSlotsCount += 1
-				}
-				spread = numeric.NewDecFromBigInt(wrapper.TotalDelegation()).
-					QuoInt64(int64(limitedSlotsCount))
-			}
+			spread := snapshot.RawStake()
 			for i := range stats.MetricsPerShard {
 				stats.MetricsPerShard[i].Vote.RawStake = spread
 			}

hmy/staking.go

@@ -46,9 +46,7 @@ func (hmy *Harmony) readAndUpdateRawStakes(
 			if err != nil {
 				continue
 			}
-			wrapper := snapshot.Validator
-			spread = numeric.NewDecFromBigInt(wrapper.TotalDelegation()).
-				QuoInt64(int64(len(wrapper.SlotPubKeys)))
+			spread = snapshot.RawStake()
 			validatorSpreads[slotAddr] = spread
 		}
 

shard/committee/assignment.go

@@ -88,7 +88,7 @@ func (p CandidateOrder) MarshalJSON() ([]byte, error) {
 func NewEPoSRound(epoch *big.Int, stakedReader StakingCandidatesReader, isExtendedBound bool, slotsLimit, shardCount int) (
 	*CompletedEPoSRound, error,
 ) {
-	eligibleCandidate, err := prepareOrders(stakedReader)
+	eligibleCandidate, err := prepareOrders(stakedReader, slotsLimit, shardCount)
 	if err != nil {
 		return nil, err
 	}
@@ -96,7 +96,7 @@ func NewEPoSRound(epoch *big.Int, stakedReader StakingCandidatesReader, isExtend
 		epoch,
 	)
 	median, winners := effective.Apply(
-		eligibleCandidate, maxExternalSlots, isExtendedBound, slotsLimit, shardCount,
+		eligibleCandidate, maxExternalSlots, isExtendedBound,
 	)
 	auctionCandidates := make([]*CandidateOrder, len(eligibleCandidate))
 
@@ -131,6 +131,7 @@ func NewEPoSRound(epoch *big.Int, stakedReader StakingCandidatesReader, isExtend
 
 func prepareOrders(
 	stakedReader StakingCandidatesReader,
+	slotsLimit, shardCount int,
 ) (map[common.Address]*effective.SlotOrder, error) {
 	candidates := stakedReader.ValidatorCandidates()
 	blsKeys := map[bls.SerializedPublicKey]struct{}{}
@@ -174,12 +175,19 @@ func prepareOrders(
 			continue
 		}
 
+		slotPubKeysLimited := make([]bls.SerializedPublicKey, 0, len(validator.SlotPubKeys))
 		found := false
+		shardSlotsCount := make([]int, shardCount)
 		for _, key := range validator.SlotPubKeys {
 			if _, ok := blsKeys[key]; ok {
 				found = true
 			} else {
 				blsKeys[key] = struct{}{}
+				shard := new(big.Int).Mod(key.Big(), big.NewInt(int64(shardCount))).Int64()
+				if slotsLimit == 0 || shardSlotsCount[shard] < slotsLimit {
+					slotPubKeysLimited = append(slotPubKeysLimited, key)
+				}
+				shardSlotsCount[shard]++
 			}
 		}
 
@@ -198,7 +206,7 @@ func prepareOrders(
 
 		essentials[validator.Address] = &effective.SlotOrder{
 			Stake:       validatorStake,
-			SpreadAmong: validator.SlotPubKeys,
+			SpreadAmong: slotPubKeysLimited,
 			Percentage:  tempZero,
 		}
 	}

staking/effective/calculate.go

@@ -79,10 +79,11 @@ func Median(stakes []SlotPurchase) numeric.Dec {
 
 // Compute ..
 func Compute(
-	shortHand map[common.Address]*SlotOrder, pull, slotsLimit, shardCount int,
+	shortHand map[common.Address]*SlotOrder, pull int,
 ) (numeric.Dec, []SlotPurchase) {
+	eposedSlots := []SlotPurchase{}
 	if len(shortHand) == 0 {
-		return numeric.ZeroDec(), []SlotPurchase{}
+		return numeric.ZeroDec(), eposedSlots
 	}
 
 	type t struct {
@@ -90,13 +91,10 @@ func Compute(
 		slot *SlotOrder
 	}
 
-	totalSlots := 0
 	shorter := []t{}
 	for key, value := range shortHand {
-		totalSlots += len(value.SpreadAmong)
 		shorter = append(shorter, t{key, value})
 	}
-	eposedSlots := make([]SlotPurchase, 0, totalSlots)
 
 	sort.SliceStable(
 		shorter,
@@ -113,36 +111,16 @@ func Compute(
 		if slotsCount == 0 {
 			continue
 		}
-		shardSlotsCount := make([]int, shardCount)
-		// may changed spread later
 		spread := numeric.NewDecFromBigInt(staker.slot.Stake).
 			QuoInt64(int64(slotsCount))
-		startIndex := len(eposedSlots)
 		for i := 0; i < slotsCount; i++ {
-			slot := SlotPurchase{
+			eposedSlots = append(eposedSlots, SlotPurchase{
 				Addr: staker.addr,
 				Key:  staker.slot.SpreadAmong[i],
 				// NOTE these are same because later the .EPoSStake mutated
 				RawStake:  spread,
 				EPoSStake: spread,
-			}
-			shard := new(big.Int).Mod(slot.Key.Big(), big.NewInt(int64(shardCount))).Int64()
-			shardSlotsCount[int(shard)]++
-			// skip if count of slots in this shard exceeds the limit
-			if slotsLimit > 0 && shardSlotsCount[int(shard)] > slotsLimit {
-				continue
-			}
-			eposedSlots = append(eposedSlots, slot)
-		}
-		// recalculate the effectiveSpread if slots exceed the limit
-		if limitedSlotsCount := len(eposedSlots) - startIndex; limitedSlotsCount != slotsCount {
-			effectiveSpread := numeric.NewDecFromBigInt(staker.slot.Stake).QuoInt64(int64(limitedSlotsCount))
-			// spread is wrapped pointer of big.Int, when we set it's value, releated slot.RawStake and slot.EPoSStake will also 'changed'
-			spread.Int.Set(effectiveSpread.Int)
-			//for _, slot := range eposedSlots[startIndex:] {
-			//	slot.RawStake = effectiveSpread
-			//	slot.EPoSStake = effectiveSpread
-			//}
+			})
 		}
 	}
 
@@ -167,10 +145,10 @@ func Compute(
 }
 
 // Apply ..
-func Apply(shortHand map[common.Address]*SlotOrder, pull int, isExtendedBound bool, slotsLimit int, shardCount int) (
+func Apply(shortHand map[common.Address]*SlotOrder, pull int, isExtendedBound bool) (
 	numeric.Dec, []SlotPurchase,
 ) {
-	median, picks := Compute(shortHand, pull, slotsLimit, shardCount)
+	median, picks := Compute(shortHand, pull)
 	max := onePlusC.Mul(median)
 	min := oneMinusC.Mul(median)
 	if isExtendedBound {

staking/types/validator.go

@@ -387,6 +387,32 @@ func (w *ValidatorWrapper) SanityCheck() error {
 	return nil
 }
 
+func (snapshot *ValidatorSnapshot) RawStake() numeric.Dec {
+	wrapper := snapshot.Validator
+	instance := shard.Schedule.InstanceForEpoch(snapshot.Epoch)
+	slotsLimit := instance.SlotsLimit()
+	if slotsLimit > 0 {
+		limitedSlotsCount := 0 // limited slots count for HIP16
+		shardCount := big.NewInt(int64(instance.NumShards()))
+		shardSlotsCount := make([]int, shardCount.Uint64()) // number slots keys on each shard
+		for _, pubkey := range wrapper.SlotPubKeys {
+			shardIndex := new(big.Int).Mod(pubkey.Big(), shardCount).Uint64()
+			shardSlotsCount[shardIndex] += 1
+			if shardSlotsCount[shardIndex] > slotsLimit {
+				continue
+			}
+			limitedSlotsCount += 1
+		}
+		return numeric.NewDecFromBigInt(wrapper.TotalDelegation()).
+			QuoInt64(int64(limitedSlotsCount))
+	}
+	if len(wrapper.SlotPubKeys) > 0 {
+		return numeric.NewDecFromBigInt(wrapper.TotalDelegation()).
+			QuoInt64(int64(len(wrapper.SlotPubKeys)))
+	}
+	return numeric.ZeroDec()
+}
+
 // Description - some possible IRL connections
 type Description struct {
 	Name            string `json:"name"`             // name