From 5342907e859c56ab541fccec8631120314b75602 Mon Sep 17 00:00:00 2001
From: Rongjian Lan <rongjian.lan@gmail.com>
Date: Thu, 14 Nov 2019 23:17:37 -0800
Subject: [PATCH] Add undelegated token withdrawal

---
 consensus/engine/consensus_engine.go |  3 +++
 core/chain_makers.go                 |  1 +
 internal/chain/engine.go             | 34 ++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/consensus/engine/consensus_engine.go b/consensus/engine/consensus_engine.go
index 88f6dd887..980b160c6 100644
--- a/consensus/engine/consensus_engine.go
+++ b/consensus/engine/consensus_engine.go
@@ -39,6 +39,9 @@ type ChainReader interface {
 	// This api reads the shard state cached or saved on the chaindb.
 	// Thus, only should be used to read the shard state of the current chain.
 	ReadShardState(epoch *big.Int) (shard.State, error)
+
+	// CurrentValidatorAddresses retrieves the current list of validators
+	CurrentValidatorAddresses() []common.Address
 }
 
 // Engine is an algorithm agnostic consensus engine.
diff --git a/core/chain_makers.go b/core/chain_makers.go
index 0acfa0d2d..f09bc8028 100644
--- a/core/chain_makers.go
+++ b/core/chain_makers.go
@@ -270,3 +270,4 @@ func (cr *fakeChainReader) GetHeaderByHash(hash common.Hash) *block.Header
 func (cr *fakeChainReader) GetHeader(hash common.Hash, number uint64) *block.Header { return nil }
 func (cr *fakeChainReader) GetBlock(hash common.Hash, number uint64) *types.Block   { return nil }
 func (cr *fakeChainReader) ReadShardState(epoch *big.Int) (shard.State, error)      { return nil, nil }
+func (cr *fakeChainReader) CurrentValidatorAddresses() []common.Address             { return nil }
diff --git a/internal/chain/engine.go b/internal/chain/engine.go
index 20b96e432..b1092d654 100644
--- a/internal/chain/engine.go
+++ b/internal/chain/engine.go
@@ -2,6 +2,7 @@ package chain
 
 import (
 	"encoding/binary"
+	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/rlp"
@@ -169,9 +170,42 @@ func (e *engineImpl) Finalize(
 	incxs []*types.CXReceiptsProof, stks []*staking.StakingTransaction) (*types.Block, error) {
 	// Accumulate any block and uncle rewards and commit the final state root
 	// Header seems complete, assemble into a block and return
+	// TODO: Block rewards should be done only in beacon chain based on cross-links
 	if err := AccumulateRewards(chain, state, header, e.Rewarder()); err != nil {
 		return nil, ctxerror.New("cannot pay block reward").WithCause(err)
 	}
+
+	// Withdraw unlocked tokens to the delegators' accounts
+	// Only do such at the last block of an epoch
+	if len(header.ShardState()) > 0 {
+		// TODO: make sure we are using the correct validator list
+		validators := chain.CurrentValidatorAddresses()
+		for _, validator := range validators {
+			wrapper := state.GetStakingInfo(validator)
+			if wrapper != nil {
+				for i := range wrapper.Delegations {
+					delegation := wrapper.Delegations[i]
+					totalWithdraw := big.NewInt(0)
+					count := 0
+					// TODO: need ot make sure the entries are ordered by epoch
+					for j := range delegation.Entries {
+						if delegation.Entries[j].Epoch.Cmp(header.Epoch()) > 14 { // need to wait at least 14 epochs to withdraw;
+							totalWithdraw.Add(totalWithdraw, delegation.Entries[j].Amount)
+							count++
+						} else {
+							break
+						}
+
+					}
+					state.AddBalance(delegation.DelegatorAddress, totalWithdraw)
+					delegation.Entries = delegation.Entries[count:]
+				}
+				if err := state.UpdateStakingInfo(validator, wrapper); err != nil {
+					return nil, ctxerror.New("failed update validator info").WithCause(err)
+				}
+			}
+		}
+	}
 	header.SetRoot(state.IntermediateRoot(chain.Config().IsS3(header.Epoch())))
 	return types.NewBlock(header, txs, receipts, outcxs, incxs, stks), nil
 }