add vote reviser (iotexproject#2370)

* add vote reviser

* fix unit test

* address comment

* fix a bug

* remove reviseHeights from staking config

* revert correction for staking receipt log

Co-authored-by: zhi <[email protected]>

Author: Yutong Pei

action/protocol/staking/handlers_test.go

@@ -70,7 +70,7 @@ func TestProtocol_HandleCreateStake(t *testing.T) {
 	require.NoError(err)
 
 	// create protocol
-	p, err := NewProtocol(depositGas, genesis.Default.Staking, nil)
+	p, err := NewProtocol(depositGas, genesis.Default.Staking, nil, genesis.Default.GreenlandBlockHeight)
 	require.NoError(err)
 
 	// set up candidate
@@ -2582,7 +2582,7 @@ func initAll(t *testing.T, ctrl *gomock.Controller) (protocol.StateManager, *Pro
 	require.NoError(err)
 
 	// create protocol
-	p, err := NewProtocol(depositGas, genesis.Default.Staking, nil)
+	p, err := NewProtocol(depositGas, genesis.Default.Staking, nil, genesis.Default.GreenlandBlockHeight)
 	require.NoError(err)
 
 	// set up candidate

action/protocol/staking/protocol.go

@@ -71,6 +71,7 @@ type (
 		config             Configuration
 		hu                 config.HeightUpgrade
 		candBucketsIndexer *CandidatesBucketsIndexer
+		voteReviser        *VoteReviser
 	}
 
 	// Configuration is the staking protocol configuration.
@@ -87,7 +88,7 @@ type (
 )
 
 // NewProtocol instantiates the protocol of staking
-func NewProtocol(depositGas DepositGas, cfg genesis.Staking, candBucketsIndexer *CandidatesBucketsIndexer) (*Protocol, error) {
+func NewProtocol(depositGas DepositGas, cfg genesis.Staking, candBucketsIndexer *CandidatesBucketsIndexer, reviseHeights ...uint64) (*Protocol, error) {
 	h := hash.Hash160b([]byte(protocolID))
 	addr, err := address.FromBytes(h[:])
 	if err != nil {
@@ -109,6 +110,9 @@ func NewProtocol(depositGas DepositGas, cfg genesis.Staking, candBucketsIndexer
 		return nil, ErrInvalidAmount
 	}
 
+	// new vote reviser, revise ate greenland
+	voteReviser := NewVoteReviser(cfg.VoteWeightCalConsts, reviseHeights...)
+
 	return &Protocol{
 		addr: addr,
 		config: Configuration{
@@ -123,6 +127,7 @@ func NewProtocol(depositGas DepositGas, cfg genesis.Staking, candBucketsIndexer
 		},
 		depositGas:         depositGas,
 		candBucketsIndexer: candBucketsIndexer,
+		voteReviser:        voteReviser,
 	}, nil
 }
 
@@ -219,6 +224,15 @@ func (p *Protocol) CreatePreStates(ctx context.Context, sm protocol.StateManager
 		}
 	}
 
+	if p.voteReviser.NeedRevise(blkCtx.BlockHeight) {
+		csm, err := NewCandidateStateManager(sm, p.hu.IsPost(config.Greenland, blkCtx.BlockHeight))
+		if err != nil {
+			return err
+		}
+		if err := p.voteReviser.Revise(csm, blkCtx.BlockHeight); err != nil {
+			return err
+		}
+	}
 	if p.candBucketsIndexer == nil {
 		return nil
 	}

action/protocol/staking/protocol_test.go

@@ -83,7 +83,7 @@ func TestProtocol(t *testing.T) {
 	}
 
 	// test loading with no candidate in stateDB
-	stk, err := NewProtocol(nil, genesis.Default.Staking, nil)
+	stk, err := NewProtocol(nil, genesis.Default.Staking, nil, genesis.Default.GreenlandBlockHeight)
 	r.NotNil(stk)
 	r.NoError(err)
 	buckets, _, err := getAllBuckets(sm)
@@ -188,7 +188,7 @@ func TestCreatePreStates(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	sm := testdb.NewMockStateManager(ctrl)
-	p, err := NewProtocol(nil, genesis.Default.Staking, nil)
+	p, err := NewProtocol(nil, genesis.Default.Staking, nil, genesis.Default.GreenlandBlockHeight)
 	require.NoError(err)
 	ctx := protocol.WithBlockCtx(
 		protocol.WithBlockchainCtx(

action/protocol/staking/validations_test.go

@@ -64,7 +64,7 @@ func TestIsValidCandidateName(t *testing.T) {
 
 func initTestProtocol(t *testing.T) (*Protocol, []*Candidate) {
 	require := require.New(t)
-	p, err := NewProtocol(nil, genesis.Default.Staking, nil)
+	p, err := NewProtocol(nil, genesis.Default.Staking, nil, genesis.Default.GreenlandBlockHeight)
 	require.NoError(err)
 
 	var cans []*Candidate

action/protocol/staking/vote_reviser.go

@@ -0,0 +1,120 @@
+package staking
+
+import (
+	"math/big"
+	"sort"
+
+	"github.com/pkg/errors"
+	"go.uber.org/zap"
+
+	"github.com/iotexproject/iotex-core/action/protocol"
+	"github.com/iotexproject/iotex-core/blockchain/genesis"
+	"github.com/iotexproject/iotex-core/pkg/log"
+	"github.com/iotexproject/iotex-core/state"
+)
+
+// VoteReviser is used to recalculate candidate votes.
+type VoteReviser struct {
+	reviseHeights []uint64
+	cache         map[uint64]CandidateList
+	c             genesis.VoteWeightCalConsts
+}
+
+// NewVoteReviser creates a VoteReviser.
+func NewVoteReviser(c genesis.VoteWeightCalConsts, reviseHeights ...uint64) *VoteReviser {
+	return &VoteReviser{
+		reviseHeights: reviseHeights,
+		cache:         make(map[uint64]CandidateList),
+		c:             c,
+	}
+}
+
+// Revise recalculate candidate votes on preset revising height.
+func (vr *VoteReviser) Revise(csm CandidateStateManager, height uint64) error {
+	if !vr.isCacheExist(height) {
+		cands, err := vr.calculateVoteWeight(csm)
+		if err != nil {
+			return err
+		}
+		vr.storeToCache(height, cands)
+	}
+	return vr.flush(height, csm)
+}
+
+func (vr *VoteReviser) storeToCache(height uint64, cands CandidateList) {
+	vr.cache[height] = cands
+}
+
+func (vr *VoteReviser) isCacheExist(height uint64) bool {
+	_, ok := vr.cache[height]
+	return ok
+}
+
+// NeedRevise returns true if height needs revise
+func (vr *VoteReviser) NeedRevise(height uint64) bool {
+	for _, h := range vr.reviseHeights {
+		if height == h {
+			return true
+		}
+	}
+	return false
+}
+
+func (vr *VoteReviser) calculateVoteWeight(sm protocol.StateManager) (CandidateList, error) {
+	cands, _, err := getAllCandidates(sm)
+	switch {
+	case errors.Cause(err) == state.ErrStateNotExist:
+	case err != nil:
+		return nil, err
+	}
+	candm := make(map[string]*Candidate)
+	for _, cand := range cands {
+		candm[cand.Owner.String()] = cand.Clone()
+		candm[cand.Owner.String()].Votes = new(big.Int)
+		candm[cand.Owner.String()].SelfStake = new(big.Int)
+	}
+	buckets, _, err := getAllBuckets(sm)
+	switch {
+	case errors.Cause(err) == state.ErrStateNotExist:
+	case err != nil:
+		return nil, err
+	}
+
+	for _, bucket := range buckets {
+		if bucket.isUnstaked() {
+			continue
+		}
+		cand, ok := candm[bucket.Candidate.String()]
+		if !ok {
+			log.L().Error("invalid bucket candidate", zap.Uint64("bucket index", bucket.Index), zap.String("candidate", bucket.Candidate.String()))
+			continue
+		}
+
+		if cand.SelfStakeBucketIdx == bucket.Index {
+			cand.AddVote(calculateVoteWeight(vr.c, bucket, true))
+			cand.SelfStake = bucket.StakedAmount
+		} else {
+			cand.AddVote(calculateVoteWeight(vr.c, bucket, false))
+		}
+	}
+
+	cands = make(CandidateList, 0, len(candm))
+	for _, cand := range candm {
+		cands = append(cands, cand)
+	}
+	return cands, nil
+}
+
+func (vr *VoteReviser) flush(height uint64, csm CandidateStateManager) error {
+	cands, ok := vr.cache[height]
+	if !ok {
+		return nil
+	}
+	sort.Sort(cands)
+	for _, cand := range cands {
+		if err := csm.Upsert(cand); err != nil {
+			return err
+		}
+	}
+	return nil
+}

chainservice/chainservice.go

@@ -240,7 +240,7 @@ func New(
 	)
 	// staking protocol need to be put in registry before poll protocol when enabling
 	if cfg.Chain.EnableStakingProtocol {
-		stakingProtocol, err = staking.NewProtocol(rewarding.DepositGas, cfg.Genesis.Staking, candBucketsIndexer)
+		stakingProtocol, err = staking.NewProtocol(rewarding.DepositGas, cfg.Genesis.Staking, candBucketsIndexer, cfg.Genesis.GreenlandBlockHeight)
 		if err != nil {
 			return nil, err
 		}