From 2397fae0d71c42f146fcd8525dae58d33bc5ce52 Mon Sep 17 00:00:00 2001
From: xiaohuo <xiaohuo200@gmail.com>
Date: Fri, 2 Apr 2021 00:42:52 +0800
Subject: [PATCH] feat(rostta): handle undelegation paypout transactions

---
 hmy/staking.go                   | 13 ++++--
 rosetta/services/tx_operation.go | 71 +++++++++++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/hmy/staking.go b/hmy/staking.go
index eb54a67a8..fd51b749b 100644
--- a/hmy/staking.go
+++ b/hmy/staking.go
@@ -519,7 +519,9 @@ func (hmy *Harmony) GetDelegationsByDelegatorByBlock(
 }
 
 // UndelegationPayouts ..
-type UndelegationPayouts map[common.Address]*big.Int
+// delegator address => validator address => amount
+type UndelegationPayouts map[common.Address]map[common.Address]*big.Int
+
 
 // GetUndelegationPayouts returns the undelegation payouts for each delegator
 //
@@ -556,10 +558,13 @@ func (hmy *Harmony) GetUndelegationPayouts(
 		for _, delegation := range wrapper.Delegations {
 			withdraw := delegation.RemoveUnlockedUndelegations(epoch, wrapper.LastEpochInCommittee, lockingPeriod, noEarlyUnlock)
 			if withdraw.Cmp(bigZero) == 1 {
-				if totalPayout, ok := undelegationPayouts[delegation.DelegatorAddress]; ok {
-					undelegationPayouts[delegation.DelegatorAddress] = new(big.Int).Add(totalPayout, withdraw)
+				if undelegationPayouts[delegation.DelegatorAddress] == nil {
+					undelegationPayouts[delegation.DelegatorAddress] = make(map[common.Address]*big.Int)
+				}
+				if totalPayout, ok := undelegationPayouts[delegation.DelegatorAddress][validator]; ok {
+					undelegationPayouts[delegation.DelegatorAddress][validator] = new(big.Int).Add(totalPayout, withdraw)
 				} else {
-					undelegationPayouts[delegation.DelegatorAddress] = withdraw
+					undelegationPayouts[delegation.DelegatorAddress][validator] = withdraw
 				}
 			}
 		}
diff --git a/rosetta/services/tx_operation.go b/rosetta/services/tx_operation.go
index 82e1cd7c4..cc052feeb 100644
--- a/rosetta/services/tx_operation.go
+++ b/rosetta/services/tx_operation.go
@@ -220,7 +220,7 @@ func GetDelegateOperationForSubAccount(tx *stakingTypes.StakingTransaction, dele
 func GetSideEffectOperationsFromUndelegationPayouts(
 	payouts hmy.UndelegationPayouts, startingOperationIndex *int64,
 ) ([]*types.Operation, *types.Error) {
-	return getSideEffectOperationsFromValueMap(
+	return getSideEffectOperationsFromUndelegateMap(
 		payouts, common.UndelegationPayoutOperation, startingOperationIndex,
 	)
 }
@@ -449,6 +449,75 @@ func getCrossShardSenderTransferNativeOperations(
 	}, nil
 }
 
+// delegator address => validator address => amount
+func getSideEffectOperationsFromUndelegateMap(
+	valueMap map[ethcommon.Address]map[ethcommon.Address]*big.Int, opType string, startingOperationIndex *int64,
+) ([]*types.Operation, *types.Error) {
+	var opIndex int64
+	operations := []*types.Operation{}
+	if startingOperationIndex != nil {
+		opIndex = *startingOperationIndex
+	} else {
+		opIndex = 0
+	}
+
+	for delegator, undelegationMap := range valueMap {
+
+		totalAmount := new(big.Int).SetUint64(0)
+		accID, rosettaError := newAccountIdentifier(delegator)
+		if rosettaError != nil {
+			return nil, rosettaError
+		}
+
+		receiverOp := &types.Operation{
+			OperationIdentifier: &types.OperationIdentifier{
+				Index: opIndex,
+			},
+			Type:    opType,
+			Status:  common.SuccessOperationStatus.Status,
+			Account: accID,
+		}
+		opIndex++
+
+		operations = append(operations, receiverOp)
+		receiverIndex := len(operations) - 1
+		for validator, amount := range undelegationMap {
+			totalAmount = new(big.Int).Add(totalAmount, amount)
+			subAccId, rosettaError := newAccountIdentifierWithSubAccount(delegator, validator, map[string]interface{}{
+				SubAccountMetadataKey: UnDelegation,
+			})
+			if rosettaError != nil {
+				return nil, rosettaError
+			}
+			payoutOp := &types.Operation{
+				OperationIdentifier: &types.OperationIdentifier{
+					Index: opIndex,
+				},
+				RelatedOperations: []*types.OperationIdentifier{
+					receiverOp.OperationIdentifier,
+				},
+				Type:    opType,
+				Status:  common.SuccessOperationStatus.Status,
+				Account: subAccId,
+				Amount: &types.Amount{
+					Value:    negativeBigValue(amount),
+					Currency: &common.NativeCurrency,
+				},
+			}
+			operations = append(operations, payoutOp)
+			opIndex++
+		}
+
+		operations[receiverIndex].Amount = &types.Amount{
+			Value:    totalAmount.String(),
+			Currency: &common.NativeCurrency,
+		}
+
+	}
+
+	return operations, nil
+}
+
 // getSideEffectOperationsFromValueMap is a helper for side effect operation construction from a address to value map.
 func getSideEffectOperationsFromValueMap(
 	valueMap map[ethcommon.Address]*big.Int, opType string, startingOperationIndex *int64,