refactor: use TokenAmount in erc20 transfer (#11)

Author: olehmisar

packages/contracts/noir/erc20/src/lib.nr

@@ -52,22 +52,17 @@ pub mod Token {
         from_secret_key: Field,
         from_note_inputs: common::Erc20NoteConsumptionInputs,
         to: common::WaAddress,
-        amount: common::U256,
+        amount: common::TokenAmount,
         to_randomness: Field,
         change_randomness: Field,
     ) {
         common::Erc20Note::sub_and_emit_change(
             context,
             [from_note_inputs],
-            common::TokenAmount { token: from_note_inputs.note.amount.token, amount },
+            amount,
             change_randomness,
             from_secret_key,
         );
-        common::Erc20Note {
-            owner: to,
-            amount: common::TokenAmount { token: from_note_inputs.note.amount.token, amount },
-            randomness: to_randomness,
-        }
-            .emit(context);
+        common::Erc20Note { owner: to, amount, randomness: to_randomness }.emit(context);
     }
 }

packages/contracts/noir/erc20_transfer/src/main.nr

@@ -5,7 +5,7 @@ fn main(
     from_note_inputs: erc20::Erc20NoteConsumptionInputs,
     // transfer params
     to: common::WaAddress,
-    amount: common::U256,
+    amount: common::TokenAmount,
     to_randomness: Field,
     // change
     change_randomness: Field,

packages/contracts/sdk/RollupService.ts

@@ -218,7 +218,7 @@ export class PoolErc20Service {
     secretKey: string;
     fromNote: Erc20Note;
     to: CompleteWaAddress;
-    amount: bigint;
+    amount: TokenAmount;
   }) {
     const { Fr } = await import("@aztec/aztec.js");
 
@@ -231,25 +231,19 @@ export class PoolErc20Service {
       from_note_inputs: await this.toNoteConsumptionInputs(secretKey, fromNote),
       from_secret_key: secretKey,
       to: { inner: to.address },
-      amount: toNoirU256(amount),
+      amount: await amount.toNoir(),
       to_randomness,
       change_randomness,
     };
     const changeNote = await Erc20Note.from({
       owner: fromNote.owner,
-      amount: await TokenAmount.from({
-        token: fromNote.amount.token,
-        amount: fromNote.amount.amount - amount,
-      }),
+      amount: fromNote.amount.sub(amount),
       randomness: change_randomness,
     });
     assert(changeNote.amount.amount >= 0n, "invalid change note");
     const toNote = await Erc20Note.from({
       owner: to,
-      amount: await TokenAmount.from({
-        token: fromNote.amount.token,
-        amount,
-      }),
+      amount,
       randomness: to_randomness,
     });
     // console.log("input\n", JSON.stringify(input));
@@ -496,6 +490,12 @@ export class TokenAmount {
       amount: toNoirU256(this.amount),
     };
   }
+
+  sub(other: TokenAmount): TokenAmount {
+    const result = this.amount - other.amount;
+    assert(result >= 0n, "TokenAmount.sub: underflow");
+    return new TokenAmount(this.token, result);
+  }
 }
 
 export type WaAddress = string;

packages/contracts/test/PoolERC20.test.ts

@@ -28,6 +28,7 @@ describe("PoolERC20", () => {
   let sdk: ReturnType<typeof interfaceSdk.createInterfaceSdk>;
   let backendSdk: ReturnType<typeof createBackendSdk>;
   let CompleteWaAddress: typeof import("../sdk").sdk.CompleteWaAddress;
+  let TokenAmount: typeof import("../sdk").sdk.TokenAmount;
   snapshottedBeforeEach(async () => {
     [alice, bob, charlie] = await ethers.getSigners();
     await typedDeployments.fixture();
@@ -42,8 +43,9 @@ describe("PoolERC20", () => {
     await usdc.mintForTests(alice, await parseUnits(usdc, "1000000"));
     await usdc.connect(alice).approve(pool, ethers.MaxUint256);
 
-    CompleteWaAddress = (await tsImport("../sdk", __filename)).sdk
-      .CompleteWaAddress;
+    ({ CompleteWaAddress, TokenAmount } = (
+      await tsImport("../sdk", __filename)
+    ).sdk);
   });
 
   before(async () => {
@@ -203,7 +205,10 @@ describe("PoolERC20", () => {
 
     // interact
     const [note] = await sdk.poolErc20.getBalanceNotesOf(usdc, aliceSecretKey);
-    const transferAmount = 123n;
+    const transferAmount = await TokenAmount.from({
+      token: await usdc.getAddress(),
+      amount: 123n,
+    });
     const { nullifier, changeNote, toNote } = await sdk.poolErc20.transfer({
       secretKey: aliceSecretKey,
       fromNote: note,
@@ -230,10 +235,10 @@ describe("PoolERC20", () => {
     await backendSdk.rollup.rollup();
 
     expect(await sdk.poolErc20.balanceOf(usdc, aliceSecretKey)).to.equal(
-      amount - transferAmount,
+      amount - transferAmount.amount,
     );
     expect(await sdk.poolErc20.balanceOf(usdc, bobSecretKey)).to.equal(
-      transferAmount,
+      transferAmount.amount,
     );
   });
 
@@ -254,7 +259,10 @@ describe("PoolERC20", () => {
       secretKey: aliceSecretKey,
       fromNote: shieldedNote,
       to: await CompleteWaAddress.fromSecretKey(bobSecretKey),
-      amount: 30n,
+      amount: await TokenAmount.from({
+        token: await usdc.getAddress(),
+        amount: 30n,
+      }),
     });
     // TODO: split notes even if they are not rolled up
     // const {  } = await sdk.poolErc20.transfer({
@@ -270,7 +278,10 @@ describe("PoolERC20", () => {
       secretKey: bobSecretKey,
       fromNote: bobNote,
       to: await CompleteWaAddress.fromSecretKey(charlieSecretKey),
-      amount: 10n,
+      amount: await TokenAmount.from({
+        token: await usdc.getAddress(),
+        amount: 10n,
+      }),
     });
     await backendSdk.rollup.rollup();
 
@@ -284,11 +295,14 @@ describe("PoolERC20", () => {
   });
 
   it("can't double spend a note", async () => {
-    const amount = 100n;
+    const amount = await TokenAmount.from({
+      token: await usdc.getAddress(),
+      amount: 100n,
+    });
     await sdk.poolErc20.shield({
       account: alice,
       token: usdc,
-      amount,
+      amount: amount.amount,
       secretKey: aliceSecretKey,
     });
     await backendSdk.rollup.rollup();
@@ -338,7 +352,10 @@ describe("PoolERC20", () => {
       secretKey: aliceSecretKey,
       fromNote: note,
       to: await CompleteWaAddress.fromSecretKey(bobSecretKey),
-      amount: 100n,
+      amount: await TokenAmount.from({
+        token: await usdc.getAddress(),
+        amount: 100n,
+      }),
     });
     expect(
       await sdk.poolErc20.getBalanceNotesOf(usdc, aliceSecretKey),