From 15bc0eeb31fd8976916a79bcc3ad161b33ff5ca5 Mon Sep 17 00:00:00 2001
From: Edgar Aroutiounian <edgar.factorial@gmail.com>
Date: Mon, 11 Nov 2019 16:54:07 -0800
Subject: [PATCH] [beacon] Only change epoch when beaconchain changed to higher
 epoch (#1823)

* [beacon] Only change epoch when beaconchain changed to higher epoch

* [beacon] Address PR comments
---
 node/node_newblock.go         | 47 ++++++-----------------------------
 node/worker/worker.go         | 37 +++++++++++++++++++++------
 shard/committee/assignment.go |  5 ++--
 3 files changed, 38 insertions(+), 51 deletions(-)

diff --git a/node/node_newblock.go b/node/node_newblock.go
index 40e374785..8f93ca5f4 100644
--- a/node/node_newblock.go
+++ b/node/node_newblock.go
@@ -1,7 +1,6 @@
 package node
 
 import (
-	"math/big"
 	"sort"
 	"time"
 
@@ -11,7 +10,6 @@ import (
 	"github.com/harmony-one/harmony/core/types"
 	"github.com/harmony-one/harmony/internal/utils"
 	"github.com/harmony-one/harmony/shard"
-	"github.com/harmony-one/harmony/shard/committee"
 )
 
 // Constants of proposing a new block
@@ -118,7 +116,13 @@ func (node *Node) proposeNewBlock() (*types.Block, error) {
 	}
 
 	// Prepare shard state
-	shardState := node.Worker.ProposeShardStateWithoutBeaconSync()
+	shardState, err := node.Worker.SuperCommitteeForNextEpoch(
+		node.Consensus.ShardID, node.Beaconchain(),
+	)
+
+	if err != nil {
+		return nil, err
+	}
 
 	// Prepare last commit signatures
 	sig, mask, err := node.Consensus.LastCommitSig()
@@ -129,43 +133,6 @@ func (node *Node) proposeNewBlock() (*types.Block, error) {
 	return node.Worker.FinalizeNewBlock(sig, mask, node.Consensus.GetViewID(), coinbase, crossLinks, shardState)
 }
 
-func (node *Node) proposeShardStateWithoutBeaconSync(block *types.Block) shard.State {
-	if block == nil || !shard.Schedule.IsLastBlock(block.Number().Uint64()) {
-		return nil
-	}
-	shardState, _ := committee.WithStakingEnabled.Compute(
-		new(big.Int).Add(block.Header().Epoch(), common.Big1), node.chainConfig, nil,
-	)
-	return shardState
-}
-
-func (node *Node) proposeShardState(block *types.Block) error {
-	switch node.Consensus.ShardID {
-	case 0:
-		return node.proposeBeaconShardState(block)
-	default:
-		node.proposeLocalShardState(block)
-		return nil
-	}
-}
-
-func (node *Node) proposeBeaconShardState(block *types.Block) error {
-	// TODO ek - replace this with variable epoch logic.
-	if !shard.Schedule.IsLastBlock(block.Number().Uint64()) {
-		// We haven't reached the end of this epoch; don't propose yet.
-		return nil
-	}
-	// TODO Use ReadFromComputation
-	prevEpoch := new(big.Int).Sub(block.Header().Epoch(), common.Big1)
-	shardState, err := committee.WithStakingEnabled.ReadFromDB(
-		prevEpoch, node.Blockchain(),
-	)
-	if err != nil {
-		return err
-	}
-	return block.AddShardState(shardState)
-}
-
 func (node *Node) proposeLocalShardState(block *types.Block) {
 	logger := block.Logger(utils.Logger())
 	// TODO ek – read this from beaconchain once BC sync is fixed
diff --git a/node/worker/worker.go b/node/worker/worker.go
index 841feef60..6a8a2b827 100644
--- a/node/worker/worker.go
+++ b/node/worker/worker.go
@@ -281,15 +281,36 @@ func (w *Worker) IncomingReceipts() []*types.CXReceiptsProof {
 	return w.current.incxs
 }
 
-// ProposeShardStateWithoutBeaconSync proposes the next shard state for next epoch.
-func (w *Worker) ProposeShardStateWithoutBeaconSync() shard.State {
-	if !shard.Schedule.IsLastBlock(w.current.header.Number().Uint64()) {
-		return nil
-	}
-	shardState, _ := committee.WithStakingEnabled.Compute(
-		new(big.Int).Add(w.current.header.Epoch(), common.Big1), *w.config, nil,
+// SuperCommitteeForNextEpoch assumes only called by consensus leader
+func (w *Worker) SuperCommitteeForNextEpoch(
+	shardID uint32,
+	beacon *core.BlockChain,
+) (shard.State, error) {
+	var (
+		nextCommittee shard.State
+		oops          error
 	)
-	return shardState
+
+	switch shardID {
+	case shard.BeaconChainShardID:
+		if shard.Schedule.IsLastBlock(w.current.header.Number().Uint64()) {
+			nextCommittee, oops = committee.WithStakingEnabled.Compute(
+				new(big.Int).Add(w.current.header.Epoch(), common.Big1),
+				*w.config,
+				beacon,
+			)
+		}
+	default:
+		// WARN When we first enable staking, this condition may not be robust by itself.
+		switch beacon.CurrentHeader().Epoch().Cmp(w.current.header.Epoch()) {
+		case 1:
+			nextCommittee, oops = committee.WithStakingEnabled.ReadFromDB(
+				beacon.CurrentHeader().Epoch(), beacon,
+			)
+		}
+
+	}
+	return nextCommittee, oops
 }
 
 // FinalizeNewBlock generate a new block for the next consensus round.
diff --git a/shard/committee/assignment.go b/shard/committee/assignment.go
index 11143dd27..7d4b47e92 100644
--- a/shard/committee/assignment.go
+++ b/shard/committee/assignment.go
@@ -12,7 +12,6 @@ import (
 	"github.com/harmony-one/harmony/internal/ctxerror"
 	"github.com/harmony-one/harmony/internal/params"
 	"github.com/harmony-one/harmony/internal/utils"
-	"github.com/harmony-one/harmony/numeric"
 	"github.com/harmony-one/harmony/shard"
 	staking "github.com/harmony-one/harmony/staking/types"
 )
@@ -51,7 +50,7 @@ type Reader interface {
 // StakingCandidatesReader ..
 type StakingCandidatesReader interface {
 	ValidatorInformation(addr common.Address) (*staking.Validator, error)
-	ValidatorStakingWithDelegation(addr common.Address) numeric.Dec
+	ValidatorStakingWithDelegation(addr common.Address) *big.Int
 	ValidatorCandidates() []common.Address
 }
 
@@ -203,7 +202,7 @@ func (def partialStakingEnabled) ReadPublicKeysFromDB(
 	return nil, nil
 }
 
-// ReadPublicKeysFromChain produces publicKeys of entire supercommittee per epoch, optionally providing a
+// ComputePublicKeys produces publicKeys of entire supercommittee per epoch, optionally providing a
 // shard specific subcommittee
 func (def partialStakingEnabled) ComputePublicKeys(
 	epoch *big.Int, reader ChainReader, shardID int,