From 3d42325d98766a231e53546c279a196b8425e159 Mon Sep 17 00:00:00 2001
From: Rongjian Lan <rongjian.lan@gmail.com>
Date: Thu, 7 May 2020 16:39:19 -0700
Subject: [PATCH] Fix both total stake counting and delegation index for
 validator itself (#2952)

---
 core/blockchain.go               | 12 ++++++++----
 core/rawdb/accessors_offchain.go |  1 +
 shard/committee/assignment.go    | 19 ++++++++++---------
 3 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/core/blockchain.go b/core/blockchain.go
index 0cb385d58..982a9c23d 100644
--- a/core/blockchain.go
+++ b/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:
diff --git a/core/rawdb/accessors_offchain.go b/core/rawdb/accessors_offchain.go
index b3c9294c0..1dd2c368e 100644
--- a/core/rawdb/accessors_offchain.go
+++ b/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 {
diff --git a/shard/committee/assignment.go b/shard/committee/assignment.go
index 012efbeb3..5d07421aa 100644
--- a/shard/committee/assignment.go
+++ b/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()