core: vm: flatten governance

3 files changed, 15 insertions, 9 deletions

diff --git a/consensus/dexcon/dexcon.go b/consensus/dexcon/dexcon.go
index aff5c9c45..c837e3a43 100644
--- a/consensus/dexcon/dexcon.go
+++ b/consensus/dexcon/dexcon.go
@@ -28,7 +28,7 @@ import (
 )
 
 type GovernanceStateFetcher interface {
-	GetGovStateHelperAtRound(round uint64) *vm.GovernanceStateHelper
+	GetStateForConfigAtRound(round uint64) *vm.GovernanceState
 }
 
 // Dexcon is a delegated proof-of-stake consensus engine.
@@ -108,7 +108,7 @@ func (d *Dexcon) Prepare(chain consensus.ChainReader, header *types.Header) erro
 }
 
 func (d *Dexcon) calculateBlockReward(round int64, state *state.StateDB) *big.Int {
-	gs := d.govStateFetcer.GetGovStateHelperAtRound(uint64(round))
+	gs := d.govStateFetcer.GetStateForConfigAtRound(uint64(round))
 	config := gs.Configuration()
 
 	blocksPerRound := config.RoundLength
@@ -134,7 +134,7 @@ func (d *Dexcon) calculateBlockReward(round int64, state *state.StateDB) *big.In
 // Finalize implements consensus.Engine, ensuring no uncles are set, nor block
 // rewards given, and returns the final block.
 func (d *Dexcon) Finalize(chain consensus.ChainReader, header *types.Header, state *state.StateDB, txs []*types.Transaction, uncles []*types.Header, receipts []*types.Receipt) (*types.Block, error) {
-	gs := vm.GovernanceStateHelper{state}
+	gs := vm.GovernanceState{state}
 
 	height := gs.RoundHeight(new(big.Int).SetUint64(header.Round))
 	if header.Round > 0 && height.Uint64() == 0 {
diff --git a/consensus/dexcon/dexcon_test.go b/consensus/dexcon/dexcon_test.go
index 65ed77cc8..f34823570 100644
--- a/consensus/dexcon/dexcon_test.go
+++ b/consensus/dexcon/dexcon_test.go
@@ -35,8 +35,8 @@ type GovStateFetcher struct {
 	statedb *state.StateDB
 }
 
-func (g *GovStateFetcher) GetGovStateHelperAtRound(_ uint64) *vm.GovernanceStateHelper {
-	return &vm.GovernanceStateHelper{g.statedb}
+func (g *GovStateFetcher) GetStateForConfigAtRound(_ uint64) *vm.GovernanceState {
+	return &vm.GovernanceState{g.statedb}
 }
 
 type DexconTestSuite struct {
@@ -45,7 +45,7 @@ type DexconTestSuite struct {
 	config  *params.DexconConfig
 	memDB   *ethdb.MemDatabase
 	stateDB *state.StateDB
-	s       *vm.GovernanceStateHelper
+	s       *vm.GovernanceState
 }
 
 func (d *DexconTestSuite) SetupTest() {
@@ -56,7 +56,7 @@ func (d *DexconTestSuite) SetupTest() {
 	}
 	d.memDB = memDB
 	d.stateDB = stateDB
-	d.s = &vm.GovernanceStateHelper{stateDB}
+	d.s = &vm.GovernanceState{stateDB}
 
 	config := params.TestnetChainConfig.Dexcon
 	config.LockupPeriod = 1000
@@ -73,7 +73,7 @@ func (d *DexconTestSuite) SetupTest() {
 
 	// Genesis CRS.
 	crs := crypto.Keccak256Hash([]byte(config.GenesisCRSText))
-	d.s.PushCRS(crs)
+	d.s.SetCRS(crs)
 
 	// Round 0 height.
 	d.s.PushRoundHeight(big.NewInt(0))
diff --git a/consensus/dexcon/fake_dexcon.go b/consensus/dexcon/fake_dexcon.go
index ca48b5275..ae7bed90e 100644
--- a/consensus/dexcon/fake_dexcon.go
+++ b/consensus/dexcon/fake_dexcon.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	coreCommon "github.com/dexon-foundation/dexon-consensus/common"
+	dexCore "github.com/dexon-foundation/dexon-consensus/core"
 	coreCrypto "github.com/dexon-foundation/dexon-consensus/core/crypto"
 	coreDKG "github.com/dexon-foundation/dexon-consensus/core/crypto/dkg"
 	coreEcdsa "github.com/dexon-foundation/dexon-consensus/core/crypto/ecdsa"
@@ -250,7 +251,12 @@ func (n *NodeSet) Randomness(round uint64, hash common.Hash) []byte {
 }
 
 func (n *NodeSet) SignCRS(round uint64) {
-	signedCRS := n.TSig(round, n.crs[round])
+	var signedCRS []byte
+	if round < dexCore.DKGDelayRound {
+		signedCRS = crypto.Keccak256(n.signedCRS[round])
+	} else {
+		signedCRS = n.TSig(round, n.crs[round])
+	}
 	n.signedCRS[round+1] = signedCRS
 	n.crs[round+1] = crypto.Keccak256Hash(signedCRS)
 }