Update RestakeRewards boolean to a RewardMethod enum (#612)

* Update RestakeRewards boolean to a RewardMethod enum

* Update TxIndex

* Maming consistency

lib/block_view_stake.go

@@ -32,10 +32,17 @@ import (
 // TYPES: StakeEntry
 //
 
+type StakingRewardMethod = uint8
+
+const (
+	StakingRewardMethodPayToBalance StakingRewardMethod = iota
+	StakingRewardMethodRestake      StakingRewardMethod = 1
+)
+
 type StakeEntry struct {
 	StakerPKID       *PKID
 	ValidatorPKID    *PKID
-	RestakeRewards   bool
+	RewardMethod     StakingRewardMethod
 	StakeAmountNanos *uint256.Int
 	ExtraData        map[string][]byte
 	isDeleted        bool
@@ -50,7 +57,7 @@ func (stakeEntry *StakeEntry) Copy() *StakeEntry {
 	return &StakeEntry{
 		StakerPKID:       stakeEntry.StakerPKID.NewPKID(),
 		ValidatorPKID:    stakeEntry.ValidatorPKID.NewPKID(),
-		RestakeRewards:   stakeEntry.RestakeRewards,
+		RewardMethod:     stakeEntry.RewardMethod,
 		StakeAmountNanos: stakeEntry.StakeAmountNanos.Clone(),
 		ExtraData:        copyExtraData(stakeEntry.ExtraData),
 		isDeleted:        stakeEntry.isDeleted,
@@ -72,7 +79,7 @@ func (stakeEntry *StakeEntry) RawEncodeWithoutMetadata(blockHeight uint64, skipM
 	var data []byte
 	data = append(data, EncodeToBytes(blockHeight, stakeEntry.StakerPKID, skipMetadata...)...)
 	data = append(data, EncodeToBytes(blockHeight, stakeEntry.ValidatorPKID, skipMetadata...)...)
-	data = append(data, BoolToByte(stakeEntry.RestakeRewards))
+	data = append(data, stakeEntry.RewardMethod)
 	data = append(data, VariableEncodeUint256(stakeEntry.StakeAmountNanos)...)
 	data = append(data, EncodeExtraData(stakeEntry.ExtraData)...)
 	return data
@@ -93,10 +100,10 @@ func (stakeEntry *StakeEntry) RawDecodeWithoutMetadata(blockHeight uint64, rr *b
 		return errors.Wrapf(err, "StakeEntry.Decode: Problem reading ValidatorPKID: ")
 	}
 
-	// RestakeRewards
-	stakeEntry.RestakeRewards, err = ReadBoolByte(rr)
+	// RewardMethod
+	stakeEntry.RewardMethod, err = rr.ReadByte()
 	if err != nil {
-		return errors.Wrapf(err, "StakeEntry.Decode: Problem reading RestakeRewards")
+		return errors.Wrapf(err, "StakeEntry.Decode: Problem reading RewardMethod")
 	}
 
 	// StakeAmountNanos
@@ -224,7 +231,7 @@ func (lockedStakeEntry *LockedStakeEntry) GetEncoderType() EncoderType {
 
 type StakeMetadata struct {
 	ValidatorPublicKey *PublicKey
-	RestakeRewards     bool
+	RewardMethod       StakingRewardMethod
 	StakeAmountNanos   *uint256.Int
 }
 
@@ -235,7 +242,7 @@ func (txnData *StakeMetadata) GetTxnType() TxnType {
 func (txnData *StakeMetadata) ToBytes(preSignature bool) ([]byte, error) {
 	var data []byte
 	data = append(data, EncodeByteArray(txnData.ValidatorPublicKey.ToBytes())...)
-	data = append(data, BoolToByte(txnData.RestakeRewards))
+	data = append(data, txnData.RewardMethod)
 	data = append(data, VariableEncodeUint256(txnData.StakeAmountNanos)...)
 	return data, nil
 }
@@ -250,10 +257,10 @@ func (txnData *StakeMetadata) FromBytes(data []byte) error {
 	}
 	txnData.ValidatorPublicKey = NewPublicKey(validatorPublicKeyBytes)
 
-	// RestakeRewards
-	txnData.RestakeRewards, err = ReadBoolByte(rr)
+	// RewardMethod
+	txnData.RewardMethod, err = rr.ReadByte()
 	if err != nil {
-		return errors.Wrapf(err, "StakeMetadata.FromBytes: Problem reading RestakeRewards: ")
+		return errors.Wrapf(err, "StakeMetadata.FromBytes: Problem reading RewardMethod: ")
 	}
 
 	// StakeAmountNanos
@@ -370,15 +377,15 @@ func (txnData *UnlockStakeMetadata) New() DeSoTxnMetadata {
 type StakeTxindexMetadata struct {
 	StakerPublicKeyBase58Check    string
 	ValidatorPublicKeyBase58Check string
-	RestakeRewards                bool
+	RewardMethod                  StakingRewardMethod
 	StakeAmountNanos              *uint256.Int
 }
 
 func (txindexMetadata *StakeTxindexMetadata) RawEncodeWithoutMetadata(blockHeight uint64, skipMetadata ...bool) []byte {
 	var data []byte
 	data = append(data, EncodeByteArray([]byte(txindexMetadata.StakerPublicKeyBase58Check))...)
 	data = append(data, EncodeByteArray([]byte(txindexMetadata.ValidatorPublicKeyBase58Check))...)
-	data = append(data, BoolToByte(txindexMetadata.RestakeRewards))
+	data = append(data, txindexMetadata.RewardMethod)
 	data = append(data, VariableEncodeUint256(txindexMetadata.StakeAmountNanos)...)
 	return data
 }
@@ -400,10 +407,10 @@ func (txindexMetadata *StakeTxindexMetadata) RawDecodeWithoutMetadata(blockHeigh
 	}
 	txindexMetadata.ValidatorPublicKeyBase58Check = string(validatorPublicKeyBase58CheckBytes)
 
-	// RestakeRewards
-	txindexMetadata.RestakeRewards, err = ReadBoolByte(rr)
+	// RewardMethod
+	txindexMetadata.RewardMethod, err = rr.ReadByte()
 	if err != nil {
-		return errors.Wrapf(err, "StakeTxindexMetadata.Decode: Problem reading RestakeRewards: ")
+		return errors.Wrapf(err, "StakeTxindexMetadata.Decode: Problem reading RewardMethod: ")
 	}
 
 	// StakeAmountNanos
@@ -1354,7 +1361,7 @@ func (bav *UtxoView) _connectStake(
 	currentStakeEntry := &StakeEntry{
 		StakerPKID:       transactorPKIDEntry.PKID,
 		ValidatorPKID:    prevValidatorEntry.ValidatorPKID,
-		RestakeRewards:   txMeta.RestakeRewards,
+		RewardMethod:     txMeta.RewardMethod,
 		StakeAmountNanos: stakeAmountNanos,
 		ExtraData:        mergeExtraData(prevExtraData, txn.ExtraData),
 	}
@@ -1979,6 +1986,11 @@ func (bav *UtxoView) IsValidStakeMetadata(transactorPkBytes []byte, metadata *St
 		return errors.Wrapf(RuleErrorInvalidStakeValidatorDisabledDelegatedStake, "UtxoView.IsValidStakeMetadata: ")
 	}
 
+	// Validate RewardMethod.
+	if metadata.RewardMethod != StakingRewardMethodPayToBalance && metadata.RewardMethod != StakingRewardMethodRestake {
+		return errors.Wrapf(RuleErrorInvalidStakingRewardMethod, "UtxoView.IsValidStakeMetadata: ")
+	}
+
 	// Validate 0 <= StakeAmountNanos <= transactor's DESO Balance. We ignore
 	// the txn fees in this check. The StakeAmountNanos will be validated to
 	// be less than the transactor's DESO balance net of txn fees in the call
@@ -3092,6 +3104,7 @@ func (bav *UtxoView) IsValidStakeLimitKey(transactorPublicKeyBytes []byte, stake
 //
 
 const RuleErrorInvalidStakerPKID RuleError = "RuleErrorInvalidStakerPKID"
+const RuleErrorInvalidStakingRewardMethod RuleError = "RuleErrorInvalidStakingRewardMethod"
 const RuleErrorInvalidStakeAmountNanos RuleError = "RuleErrorInvalidStakeAmountNanos"
 const RuleErrorInvalidStakeInsufficientBalance RuleError = "RuleErrorInvalidStakeInsufficientBalance"
 const RuleErrorInvalidStakeValidatorDisabledDelegatedStake RuleError = "RuleErrorInvalidStakeValidatorDisabledDelegatedStake"

lib/block_view_stake_test.go

@@ -125,7 +125,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
-			RestakeRewards:     false,
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   uint256.NewInt().SetUint64(100),
 		}
 		_, err = _submitStakeTxn(
@@ -142,7 +142,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		// RuleErrorInvalidValidatorPKID
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m2PkBytes),
-			RestakeRewards:     false,
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   uint256.NewInt(),
 		}
 		_, err = _submitStakeTxn(
@@ -151,11 +151,24 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		require.Error(t, err)
 		require.Contains(t, err.Error(), RuleErrorInvalidValidatorPKID)
 	}
+	{
+		// RuleErrorInvalidStakingRewardMethod
+		stakeMetadata := &StakeMetadata{
+			ValidatorPublicKey: NewPublicKey(m0PkBytes),
+			RewardMethod:       99,
+			StakeAmountNanos:   uint256.NewInt().SetUint64(1),
+		}
+		_, err = _submitStakeTxn(
+			testMeta, m1Pub, m1Priv, stakeMetadata, nil, flushToDB,
+		)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), RuleErrorInvalidStakingRewardMethod)
+	}
 	{
 		// RuleErrorInvalidStakeAmountNanos
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
-			RestakeRewards:     false,
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   nil,
 		}
 		_, err = _submitStakeTxn(
@@ -168,6 +181,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		// RuleErrorInvalidStakeAmountNanos
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   uint256.NewInt(),
 		}
 		_, err = _submitStakeTxn(
@@ -180,6 +194,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		// RuleErrorInvalidStakeAmountNanos
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   MaxUint256,
 		}
 		_, err = _submitStakeTxn(
@@ -192,7 +207,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		// RuleErrorInvalidStakeInsufficientBalance
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
-			RestakeRewards:     false,
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   uint256.NewInt().SetUint64(math.MaxUint64),
 		}
 		_, err = _submitStakeTxn(
@@ -206,7 +221,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		m1OldDESOBalanceNanos := getDESOBalanceNanos(m1PkBytes)
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
-			RestakeRewards:     false,
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   uint256.NewInt().SetUint64(100),
 		}
 		extraData := map[string][]byte{"TestKey": []byte("TestValue")}
@@ -219,7 +234,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		stakeEntry, err := utxoView().GetStakeEntry(m0PKID, m1PKID)
 		require.NoError(t, err)
 		require.NotNil(t, stakeEntry)
-		require.False(t, stakeEntry.RestakeRewards)
+		require.Equal(t, stakeEntry.RewardMethod, StakingRewardMethodPayToBalance)
 		require.Equal(t, stakeEntry.StakeAmountNanos, uint256.NewInt().SetUint64(100))
 		require.Equal(t, stakeEntry.ExtraData["TestKey"], []byte("TestValue"))
 
@@ -238,7 +253,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		m1OldDESOBalanceNanos := getDESOBalanceNanos(m1PkBytes)
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
-			RestakeRewards:     false,
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   uint256.NewInt().SetUint64(50),
 		}
 		extraData := map[string][]byte{"TestKey": []byte("TestValue2")}
@@ -251,7 +266,7 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		stakeEntry, err := utxoView().GetStakeEntry(m0PKID, m1PKID)
 		require.NoError(t, err)
 		require.NotNil(t, stakeEntry)
-		require.False(t, stakeEntry.RestakeRewards)
+		require.Equal(t, stakeEntry.RewardMethod, StakingRewardMethodPayToBalance)
 		require.Equal(t, stakeEntry.StakeAmountNanos, uint256.NewInt().SetUint64(150))
 		require.Equal(t, stakeEntry.ExtraData["TestKey"], []byte("TestValue2"))
 
@@ -266,11 +281,11 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		require.Equal(t, m1OldDESOBalanceNanos-feeNanos-stakeMetadata.StakeAmountNanos.Uint64(), m1NewDESOBalanceNanos)
 	}
 	{
-		// m1 changes the RestakeRewards flag on their stake with m0.
+		// m1 changes the RewardMethod value on their stake with m0.
 		m1OldDESOBalanceNanos := getDESOBalanceNanos(m1PkBytes)
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(m0PkBytes),
-			RestakeRewards:     true,
+			RewardMethod:       StakingRewardMethodRestake,
 			StakeAmountNanos:   uint256.NewInt(),
 		}
 		extraData := map[string][]byte{"TestKey": []byte("TestValue2")}
@@ -286,8 +301,8 @@ func _testStaking(t *testing.T, flushToDB bool) {
 		require.Equal(t, stakeEntry.StakeAmountNanos, uint256.NewInt().SetUint64(150))
 		require.Equal(t, stakeEntry.ExtraData["TestKey"], []byte("TestValue2"))
 
-		// Verify the StakeEntry.RestakeRewards flag is updated to true.
-		require.True(t, stakeEntry.RestakeRewards)
+		// Verify the StakeEntry.RewardMethod has changed to StakingRewardMethodRestake.
+		require.Equal(t, stakeEntry.RewardMethod, StakingRewardMethodRestake)
 
 		// Verify the ValidatorEntry.TotalStakeAmountNanos does not change.
 		validatorEntry, err := utxoView().GetValidatorByPKID(m0PKID)
@@ -1725,7 +1740,7 @@ func _testGetTopStakesByStakeAmount(t *testing.T, flushToDB bool) {
 	constructAndSubmitStakeTxn := func(stakerPk string, stakerPriv string, validatorPkBytes []byte, amountNanos uint64) {
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(validatorPkBytes),
-			RestakeRewards:     false,
+			RewardMethod:       StakingRewardMethodPayToBalance,
 			StakeAmountNanos:   uint256.NewInt().SetUint64(amountNanos),
 		}
 		_, err := _submitStakeTxn(testMeta, stakerPk, stakerPriv, stakeMetadata, nil, flushToDB)

lib/pos_epoch_complete_hook_test.go

@@ -125,9 +125,14 @@ func TestRunEpochCompleteHook(t *testing.T) {
 		_, err = _submitRegisterAsValidatorTxn(testMeta, publicKey, privateKey, registerMetadata, nil, true)
 		require.NoError(t, err)
 
+		rewardMethod := StakingRewardMethodPayToBalance
+		if restakeRewards {
+			rewardMethod = StakingRewardMethodRestake
+		}
+
 		stakeMetadata := &StakeMetadata{
 			ValidatorPublicKey: NewPublicKey(pkBytes),
-			RestakeRewards:     restakeRewards,
+			RewardMethod:       rewardMethod,
 			StakeAmountNanos:   uint256.NewInt().SetUint64(stakeAmountNanos),
 		}
 		_, err = _submitStakeTxn(testMeta, publicKey, privateKey, stakeMetadata, nil, true)
@@ -462,7 +467,7 @@ func TestRunEpochCompleteHook(t *testing.T) {
 		require.Len(t, snapshotStakeEntries, 6)
 	}
 	{
-		// Test staking rewards distribution with RestakeRewards enabled.
+		// Test staking rewards distribution with restaking enabled.
 
 		// m6 now has a 14333333578 nano balance from staking rewards so far.
 		balance, err := utxoView().GetDeSoBalanceNanosForPublicKey(m6PkBytes)
@@ -478,14 +483,14 @@ func TestRunEpochCompleteHook(t *testing.T) {
 		require.Equal(t, balance, uint64(16747126681))
 	}
 	{
-		// Test staking rewards distribution with RestakeRewards enabled.
+		// Test staking rewards distribution with restaking enabled.
 
 		// m6 has 700 nanos staked.
 		stakeEntry, err := utxoView().GetStakeEntry(m6PKID, m6PKID)
 		require.NoError(t, err)
 		require.Equal(t, stakeEntry.StakeAmountNanos, uint256.NewInt().SetUint64(700))
 
-		// m6 sets their RestakeRewards flag to true.
+		// m6 enables restaking.
 		_registerAndStake(m6Pub, m6Priv, 0, true)
 
 		// m6's wallet balance is 16747126627 after they submit their stake transaction.

lib/pos_snapshot_entries.go

@@ -688,7 +688,7 @@ type SnapshotStakeMapKey struct {
 // This is a bare bones in-memory only construct used to capture the ValidatorPKID,
 // StakerPKID, and StakeAmountNanos from a StakeEntry that has been snapshotted. We
 // define a new type here rather than re-using the StakeEntry type to reduce the risk
-// of bugs. The StakeEntry type has additional fields (ex: RestakeRewards, ExtraData)
+// of bugs. The StakeEntry type has additional fields (ex: RewardMethod, ExtraData)
 // that are not snapshotted.
 type SnapshotStakeEntry struct {
 	SnapshotAtEpochNumber uint64

lib/pos_staking_rewards.go

@@ -79,7 +79,7 @@ func (bav *UtxoView) DistributeStakingRewardsToSnapshotStakes(blockHeight uint64
 		// StakeEntry. Their stake is currently in lockup.
 
 		// For case 1, we distribute the rewards by adding them to the staker's staked amount.
-		if stakeEntry != nil && stakeEntry.RestakeRewards {
+		if stakeEntry != nil && stakeEntry.RewardMethod == StakingRewardMethodRestake {
 			stakeEntry.StakeAmountNanos.Add(stakeEntry.StakeAmountNanos, rewardAmount)
 			bav._setStakeEntryMappings(stakeEntry)