fix: address review feedback and fix adding stakeouts

Author: erictaylor

src/vms/pvm/etna-builder/spend-reducers/useSpendableLockedUTXOs.test.ts

@@ -162,30 +162,6 @@ describe('useSpendableLockedUTXOs', () => {
     expect(stakeOutputs).toHaveLength(0);
   });
 
-  it('should add stake outputs even with no UTXOs', () => {
-    const toBurn = new Map();
-    const toStake = new Map([[testContext.avaxAssetID, 1_000n]]);
-
-    const initialState = getInitialReducerState({
-      excessAVAX: 0n,
-      spendOptions: {
-        minIssuanceTime: 100n,
-      },
-      toBurn,
-      toStake,
-      utxos: [],
-    });
-
-    const spendHelper = getSpendHelper({ toBurn, toStake });
-
-    useSpendableLockedUTXOs(initialState, spendHelper, testContext);
-    const { stakeOutputs } = spendHelper.getInputsOutputs();
-
-    expect(stakeOutputs).toHaveLength(1);
-    expect(stakeOutputs[0].assetId.toString()).toEqual(testContext.avaxAssetID);
-    expect(stakeOutputs[0].amount()).toEqual(1_000n);
-  });
-
   it('should add spendable locked UTXO with change', () => {
     const toBurn = new Map();
     const toStake = new Map([[testAvaxAssetID.toString(), 1_000n]]);

src/vms/pvm/etna-builder/spend-reducers/useSpendableLockedUTXOs.ts

@@ -33,7 +33,7 @@ export const getUsableUTXOsFilter =
     }
 
     // 1b. Ensure UTXO is stakeable.
-    if (!(state.spendOptions.minIssuanceTime < utxo.output.getLocktime())) {
+    if (state.spendOptions.minIssuanceTime >= utxo.output.getLocktime()) {
       return false;
     }
 
@@ -43,7 +43,7 @@ export const getUsableUTXOsFilter =
     }
 
     // 1d. Filter out UTXOs that aren't needed for staking.
-    if ((state.toStake.get(utxo.assetId.value()) ?? 0n) === 0n) {
+    if (!state.toStake.has(utxo.assetId.value())) {
       return false;
     }
 
@@ -85,19 +85,14 @@ export const useSpendableLockedUTXOs: SpendReducerFunction = (
         utxo.utxoId,
         utxo.assetId,
         new StakeableLockIn(
-          // StakeableLockOut
           new BigIntPr(utxoInfo.stakeableLocktime),
-          TransferInput.fromNative(
-            // TransferOutput
-            utxoInfo.amount,
-            sigData.sigIndicies,
-          ),
+          TransferInput.fromNative(utxoInfo.amount, sigData.sigIndicies),
         ),
       ),
     );
 
     // 3c. Consume the locked asset and get the remaining amount.
-    const remainingAmount = spendHelper.consumeLockedAsset(
+    const [remainingAmount] = spendHelper.consumeLockedAsset(
       utxoInfo.assetId,
       utxoInfo.amount,
     );
@@ -133,20 +128,5 @@ export const useSpendableLockedUTXOs: SpendReducerFunction = (
     }
   }
 
-  // 4. Add all remaining stake amounts assuming they are unlocked.
-  for (const [assetId, amount] of state.toStake) {
-    if (amount === 0n) {
-      continue;
-    }
-
-    spendHelper.addStakedOutput(
-      TransferableOutput.fromNative(
-        assetId,
-        amount,
-        state.spendOptions.changeAddresses,
-      ),
-    );
-  }
-
   return state;
 };

src/vms/pvm/etna-builder/spend-reducers/useUnlockedUTXOs.ts

@@ -90,6 +90,12 @@ export const useUnlockedUTXOs: SpendReducerFunction = (
       { otherVerifiedUsableUTXOs: [], avaxVerifiedUsableUTXOs: [] },
     );
 
+  const changeOwner = OutputOwners.fromNative(
+    state.spendOptions.changeAddresses,
+    0n,
+    1,
+  );
+
   // 4. Handle all the non-AVAX asset UTXOs first.
   for (const { sigData, data: utxo } of otherVerifiedUsableUTXOs) {
     const utxoInfo = getUtxoInfo(utxo);
@@ -114,30 +120,27 @@ export const useUnlockedUTXOs: SpendReducerFunction = (
     );
 
     // 4c. Consume the asset and get the remaining amount.
-    const remainingAmount = spendHelper.consumeAsset(
+    const [remainingAmount, amountToStake] = spendHelper.consumeAsset(
       utxoInfo.assetId,
       utxoInfo.amount,
     );
 
     // 4d. If "amountToStake" is greater than 0, add the stake output.
-    // TODO: Implement or determine if needed.
+    if (amountToStake > 0n) {
+      spendHelper.addStakedOutput(
+        new TransferableOutput(
+          utxo.assetId,
+          new TransferOutput(new BigIntPr(amountToStake), changeOwner),
+        ),
+      );
+    }
 
     // 4e. Add the change output if there is any remaining amount.
     if (remainingAmount > 0n) {
       spendHelper.addChangeOutput(
         new TransferableOutput(
           utxo.assetId,
-          new TransferableOutput(
-            utxo.assetId,
-            new TransferOutput(
-              new BigIntPr(remainingAmount),
-              OutputOwners.fromNative(
-                state.spendOptions.changeAddresses,
-                0n,
-                1,
-              ),
-            ),
-          ),
+          new TransferOutput(new BigIntPr(remainingAmount), changeOwner),
         ),
       );
     }
@@ -170,11 +173,20 @@ export const useUnlockedUTXOs: SpendReducerFunction = (
       ),
     );
 
-    const remainingAmount = spendHelper.consumeAsset(
+    const [remainingAmount, amountToStake] = spendHelper.consumeAsset(
       context.avaxAssetID,
       utxoInfo.amount,
     );
 
+    if (amountToStake > 0n) {
+      spendHelper.addStakedOutput(
+        new TransferableOutput(
+          utxo.assetId,
+          new TransferOutput(new BigIntPr(amountToStake), changeOwner),
+        ),
+      );
+    }
+
     excessAVAX += remainingAmount;
 
     // The ownerOverride is no longer needed. Clear it.

src/vms/pvm/etna-builder/spendHelper.test.ts

@@ -167,65 +167,65 @@ describe('src/vms/pvm/etna-builder/spendHelper', () => {
 
       expect(spendHelper.shouldConsumeAsset('asset')).toBe(false);
     });
+  });
 
-    describe('SpendHelper.consumeLockedAsset', () => {
-      const testCases = [
-        {
-          description: 'consumes the full amount',
-          toStake: new Map([['asset', 1n]]),
-          asset: 'asset',
-          amount: 1n,
-          expected: 0n,
-        },
-        {
-          description: 'consumes a partial amount',
-          toStake: new Map([['asset', 1n]]),
-          asset: 'asset',
-          amount: 2n,
-          expected: 1n,
-        },
-        {
-          description: 'consumes nothing',
-          toStake: new Map([['asset', 1n]]),
-          asset: 'asset',
-          amount: 0n,
-          expected: 0n,
-        },
-        {
-          description: 'consumes nothing when asset not in toStake',
-          toStake: new Map(),
-          asset: 'asset',
-          amount: 1n,
-          expected: 1n,
-        },
-        {
-          description: 'consumes nothing when asset in toStake with 0 value',
-          toStake: new Map([['asset', 0n]]),
-          asset: 'asset',
-          amount: 1n,
-          expected: 1n,
-        },
-      ];
-
-      test.each(testCases)(
-        '$description',
-        ({ toStake, asset, amount, expected }) => {
-          const spendHelper = new SpendHelper({
-            ...DEFAULT_PROPS,
-            toStake,
-          });
-
-          expect(spendHelper.consumeLockedAsset(asset, amount)).toBe(expected);
-        },
-      );
+  describe('SpendHelper.consumeLockedAsset', () => {
+    const testCases = [
+      {
+        description: 'consumes the full amount',
+        toStake: new Map([['asset', 1n]]),
+        asset: 'asset',
+        amount: 1n,
+        expected: 0n,
+      },
+      {
+        description: 'consumes a partial amount',
+        toStake: new Map([['asset', 1n]]),
+        asset: 'asset',
+        amount: 2n,
+        expected: 1n,
+      },
+      {
+        description: 'consumes nothing',
+        toStake: new Map([['asset', 1n]]),
+        asset: 'asset',
+        amount: 0n,
+        expected: 0n,
+      },
+      {
+        description: 'consumes nothing when asset not in toStake',
+        toStake: new Map(),
+        asset: 'asset',
+        amount: 1n,
+        expected: 1n,
+      },
+      {
+        description: 'consumes nothing when asset in toStake with 0 value',
+        toStake: new Map([['asset', 0n]]),
+        asset: 'asset',
+        amount: 1n,
+        expected: 1n,
+      },
+    ];
 
-      test('throws an error when amount is negative', () => {
-        const spendHelper = new SpendHelper(DEFAULT_PROPS);
+    test.each(testCases)(
+      '$description',
+      ({ toStake, asset, amount, expected }) => {
+        const spendHelper = new SpendHelper({
+          ...DEFAULT_PROPS,
+          toStake,
+        });
 
-        expect(() => {
-          spendHelper.consumeLockedAsset('asset', -1n);
-        }).toThrow('Amount to consume must be greater than or equal to 0');
-      });
+        expect(spendHelper.consumeLockedAsset(asset, amount)[0]).toBe(expected);
+      },
+    );
+
+    test('throws an error when amount is negative', () => {
+      const spendHelper = new SpendHelper(DEFAULT_PROPS);
+
+      expect(() => {
+        spendHelper.consumeLockedAsset('asset', -1n);
+      }).toThrow('Amount to consume must be greater than or equal to 0');
     });
   });
 
@@ -284,7 +284,7 @@ describe('src/vms/pvm/etna-builder/spendHelper', () => {
           toBurn,
         });
 
-        expect(spendHelper.consumeAsset(asset, amount)).toBe(expected);
+        expect(spendHelper.consumeAsset(asset, amount)[0]).toBe(expected);
       },
     );
 

src/vms/pvm/etna-builder/spendHelper.ts

@@ -160,9 +160,12 @@ export class SpendHelper {
    *
    * @param {string} assetId - The ID of the asset to consume.
    * @param {bigint} amount - The amount of the asset to consume.
-   * @returns {bigint} The remaining amount of the asset after consumption.
+   * @returns A tuple of the remaining amount in the first position and the amount to stake in the second position.
    */
-  consumeLockedAsset(assetId: string, amount: bigint): bigint {
+  consumeLockedAsset(
+    assetId: string,
+    amount: bigint,
+  ): [remainingAmount: bigint, amountToStake: bigint] {
     if (amount < 0n) {
       throw new Error('Amount to consume must be greater than or equal to 0');
     }
@@ -179,17 +182,20 @@ export class SpendHelper {
 
     this.toStake.set(assetId, remainingAmountToStake - amountToStake);
 
-    return amount - amountToStake;
+    return [amount - amountToStake, amountToStake];
   }
 
   /**
    * Consumes an asset based on its asset ID and amount.
    *
    * @param {string} assetId - The ID of the asset to consume.
    * @param {bigint} amount - The amount of the asset to consume.
-   * @returns {bigint} The remaining amount of the asset after consumption.
+   * @returns A tuple of the remaining amount in the first position and the amount to stake in the second position.
    */
-  consumeAsset(assetId: string, amount: bigint): bigint {
+  consumeAsset(
+    assetId: string,
+    amount: bigint,
+  ): [remainingAmount: bigint, amountToStake: bigint] {
     if (amount < 0n) {
       throw new Error('Amount to consume must be greater than or equal to 0');
     }

src/vms/pvm/txs/fee/calculator.ts

@@ -10,8 +10,6 @@ import { getTxComplexity } from './complexity';
  * transaction must pay for valid inclusion into a block.
  */
 export const calculateFee = (
-  // TODO: Do we need this to be UnsignedTx?
-  // If so, we can use .getTx() to get the Transaction.
   tx: Transaction,
   weights: Dimensions,
   price: bigint,

src/vms/pvm/txs/fee/complexity.ts

@@ -188,7 +188,6 @@ export const getOwnerComplexity = (outputOwners: OutputOwners): Dimensions => {
  * It does not include the typeID of the credential.
  */
 export const getAuthComplexity = (input: Serializable): Dimensions => {
-  // TODO: Not a fan of this. May be better to re-type `subnetAuth` as `Input` in `AddSubnetValidatorTx`?
   if (!(input instanceof Input)) {
     throw new Error(
       'Unable to calculate auth complexity of transaction. Expected Input as subnet auth.',