add rewards to TotalManaIn

iotaledger · Oct 31, 2023 · 3fcd9e2 · 3fcd9e2
1 parent 2898851
commit 3fcd9e2
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 8 deletions.
diff --git a/vm/nova/stvf_test.go b/vm/nova/stvf_test.go
@@ -3727,6 +3727,7 @@ func createWorkingSet(t *testing.T, input *vm.ChainOutputWithIDs, workingSet *vm
 			tpkg.TestAPI.StorageScoreStructure(),
 			workingSet.Tx.CreationSlot,
 			workingSet.UTXOInputsSet,
+			workingSet.Rewards,
 		)
 		require.NoError(t, err)
 		workingSet.TotalManaIn = totalManaIn

diff --git a/vm/nova/vm.go b/vm/nova/vm.go
@@ -70,6 +70,7 @@ func NewVMParamsWorkingSet(api iotago.API, t *iotago.Transaction, inputs vm.Reso
 		api.StorageScoreStructure(),
 		workingSet.Tx.CreationSlot,
 		workingSet.UTXOInputsSet,
+		workingSet.Rewards,
 	)
 	if err != nil {
 		return nil, ierrors.Join(iotago.ErrManaAmountInvalid, err)

diff --git a/vm/nova/vm_test.go b/vm/nova/vm_test.go
@@ -1437,7 +1437,7 @@ func runNovaTransactionExecutionTest(t *testing.T, test *txExecTest) {
 		// HINT: all outputs are created at slot 0 and the transaction is executed at slot 10000
 		var txCreationSlot iotago.SlotIndex = 10000
 
-		totalInputMana, err := vm.TotalManaIn(testAPI.ManaDecayProvider(), testAPI.StorageScoreStructure(), txCreationSlot, inputSet)
+		totalInputMana, err := vm.TotalManaIn(testAPI.ManaDecayProvider(), testAPI.StorageScoreStructure(), txCreationSlot, inputSet, vm.RewardsInputSet{})
 		require.NoError(t, err)
 
 		outputs := iotago.TxEssenceOutputs{

diff --git a/vm/vm.go b/vm/vm.go
@@ -71,7 +71,7 @@ func (workingSet *WorkingSet) UTXOInputAtIndex(inputIndex uint16) *iotago.UTXOIn
 	return workingSet.Tx.TransactionEssence.Inputs[inputIndex].(*iotago.UTXOInput)
 }
 
-func TotalManaIn(manaDecayProvider *iotago.ManaDecayProvider, storageScoreStructure *iotago.StorageScoreStructure, txCreationSlot iotago.SlotIndex, inputSet InputSet) (iotago.Mana, error) {
+func TotalManaIn(manaDecayProvider *iotago.ManaDecayProvider, storageScoreStructure *iotago.StorageScoreStructure, txCreationSlot iotago.SlotIndex, inputSet InputSet, rewards RewardsInputSet) (iotago.Mana, error) {
 	var totalIn iotago.Mana
 	for outputID, input := range inputSet {
 		// stored Mana
@@ -83,7 +83,6 @@ func TotalManaIn(manaDecayProvider *iotago.ManaDecayProvider, storageScoreStruct
 		if err != nil {
 			return 0, ierrors.Wrapf(iotago.ErrManaOverflow, "%w", err)
 		}
-
 		// potential Mana
 		// the storage deposit does not generate potential mana, so we only use the excess base tokens to calculate the potential mana
 		minDeposit, err := storageScoreStructure.MinDeposit(input)
@@ -103,6 +102,14 @@ func TotalManaIn(manaDecayProvider *iotago.ManaDecayProvider, storageScoreStruct
 			return 0, ierrors.Wrapf(iotago.ErrManaOverflow, "%w", err)
 		}
 	}
+	// rewards
+	for _, reward := range rewards {
+		var err error
+		totalIn, err = safemath.SafeAdd(totalIn, reward)
+		if err != nil {
+			return 0, ierrors.Wrapf(iotago.ErrManaOverflow, "%w", err)
+		}
+	}
 
 	return totalIn, nil
 }
@@ -583,11 +590,6 @@ func ExecFuncBalancedMana() ExecFunc {
 		manaIn := vmParams.WorkingSet.TotalManaIn
 		manaOut := vmParams.WorkingSet.TotalManaOut
 
-		// Whether it's valid to claim rewards is checked in the delegation and staking STVFs.
-		for _, reward := range vmParams.WorkingSet.Rewards {
-			manaIn += reward
-		}
-
 		if manaIn < manaOut {
 			// less mana on input side than on output side => not allowed
 			return ierrors.Wrapf(iotago.ErrInputOutputManaMismatch, "Mana in %d, Mana out %d", manaIn, manaOut)