Fix both total stake counting and delegation index for validator itself (#2952)

core/blockchain.go

@@ -2651,11 +2651,15 @@ func (bc *BlockChain) prepareStakingMetaData(
 				blockNum,
 			}
 			delegations, ok := newDelegations[createValidator.ValidatorAddress]
-			if ok {
-				delegations = append(delegations, selfIndex)
-			} else {
-				delegations = staking.DelegationIndexes{selfIndex}
+			if !ok {
+				// If the cache doesn't have it, load it from DB for the first time.
+				delegations, err = bc.ReadDelegationsByDelegator(createValidator.ValidatorAddress)
+				if err != nil {
+					return nil, nil, err
+				}
 			}
+
+			delegations = append(delegations, selfIndex)
 			newDelegations[createValidator.ValidatorAddress] = delegations
 		case staking.DirectiveEditValidator:
 		case staking.DirectiveDelegate:

core/rawdb/accessors_offchain.go

@@ -242,6 +242,7 @@ func WriteValidatorList(
 }
 
 // ReadDelegationsByDelegator retrieves the list of validators delegated by a delegator
+// Returns empty results instead of error if there is not data found.
 func ReadDelegationsByDelegator(db DatabaseReader, delegator common.Address) (staking.DelegationIndexes, error) {
 	data, err := db.Get(delegatorValidatorListKey(delegator))
 	if err != nil || len(data) == 0 {

shard/committee/assignment.go

@@ -146,15 +146,6 @@ func prepareOrders(
 			continue
 		}
 
-		validatorStake := big.NewInt(0)
-		for i := range validator.Delegations {
-			validatorStake.Add(
-				validatorStake, validator.Delegations[i].Amount,
-			)
-		}
-
-		totalStaked.Add(totalStaked, validatorStake)
-
 		found := false
 		for _, key := range validator.SlotPubKeys {
 			if _, ok := blsKeys[key]; ok {
@@ -168,6 +159,15 @@ func prepareOrders(
 			continue
 		}
 
+		validatorStake := big.NewInt(0)
+		for i := range validator.Delegations {
+			validatorStake.Add(
+				validatorStake, validator.Delegations[i].Amount,
+			)
+		}
+
+		totalStaked.Add(totalStaked, validatorStake)
+
 		essentials[validator.Address] = &effective.SlotOrder{
 			validatorStake,
 			validator.SlotPubKeys,
@@ -240,6 +240,7 @@ var (
 	ErrComputeForEpochInPast = errors.New("cannot compute for epoch in past")
 )
 
+// This is the shard state computation logic before staking epoch.
 func preStakingEnabledCommittee(s shardingconfig.Instance) *shard.State {
 	shardNum := int(s.NumShards())
 	shardHarmonyNodes := s.NumHarmonyOperatedNodesPerShard()