diff --git a/yarn-project/end-to-end/src/e2e_lending_contract.test.ts b/yarn-project/end-to-end/src/e2e_lending_contract.test.ts index efc5b8843d4..d9b31c0315e 100644 --- a/yarn-project/end-to-end/src/e2e_lending_contract.test.ts +++ b/yarn-project/end-to-end/src/e2e_lending_contract.test.ts @@ -4,21 +4,30 @@ import { Account, AuthWitnessAccountContract, AuthWitnessEntrypointWallet, - AztecAddress, CheatCodes, Fr, IAuthWitnessAccountEntrypoint, + SentTx, computeMessageSecretHash, } from '@aztec/aztec.js'; import { CircuitsWasm, CompleteAddress, FunctionSelector, GeneratorIndex, GrumpkinScalar } from '@aztec/circuits.js'; -import { pedersenPlookupCommitInputs, pedersenPlookupCompressWithHashIndex } from '@aztec/circuits.js/barretenberg'; +import { pedersenPlookupCompressWithHashIndex } from '@aztec/circuits.js/barretenberg'; import { DebugLogger } from '@aztec/foundation/log'; -import { LendingContract, NativeTokenContract, PriceFeedContract } from '@aztec/noir-contracts/types'; +import { + LendingContract, + PriceFeedContract, + SchnorrAuthWitnessAccountContract, + TokenContract, +} from '@aztec/noir-contracts/types'; import { AztecRPC, TxStatus } from '@aztec/types'; +import { jest } from '@jest/globals'; + import { setup } from './fixtures/utils.js'; +import { LendingAccount, LendingSimulator, TokenSimulator } from './simulators/index.js'; describe('e2e_lending_contract', () => { + jest.setTimeout(100_000); let aztecNode: AztecNodeService | undefined; let aztecRpcServer: AztecRPC; let wallet: AuthWitnessEntrypointWallet; @@ -26,83 +35,105 @@ describe('e2e_lending_contract', () => { let logger: DebugLogger; let cc: CheatCodes; + const TIME_JUMP = 100; + + let lendingContract: LendingContract; + let priceFeedContract: PriceFeedContract; + let collateralAsset: TokenContract; + let stableCoin: TokenContract; - const WAD = 10n ** 18n; - const BASE = 10n ** 9n; + let lendingAccount: LendingAccount; + let lendingSim: LendingSimulator; + + const waitForSuccess = async (tx: SentTx) => { + const receipt = await tx.wait(); + expect(receipt.status).toBe(TxStatus.MINED); + return receipt; + }; - const deployContracts = async (owner: AztecAddress) => { + const deployContracts = async () => { let lendingContract: LendingContract; let priceFeedContract: PriceFeedContract; - let collateralAsset: NativeTokenContract; - let stableCoin: NativeTokenContract; + let collateralAsset: TokenContract; + let stableCoin: TokenContract; { logger(`Deploying price feed contract...`); - const tx = PriceFeedContract.deploy(wallet).send(); - logger(`Tx sent with hash ${await tx.getTxHash()}`); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); + const receipt = await waitForSuccess(PriceFeedContract.deploy(wallet).send()); logger(`Price feed deployed to ${receipt.contractAddress}`); priceFeedContract = await PriceFeedContract.at(receipt.contractAddress!, wallet); } { logger(`Deploying collateral asset feed contract...`); - const tx = NativeTokenContract.deploy(wallet, 10000n, owner).send(); - logger(`Tx sent with hash ${await tx.getTxHash()}`); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); + const receipt = await waitForSuccess(TokenContract.deploy(wallet).send()); logger(`Collateral asset deployed to ${receipt.contractAddress}`); - collateralAsset = await NativeTokenContract.at(receipt.contractAddress!, wallet); + collateralAsset = await TokenContract.at(receipt.contractAddress!, wallet); } { logger(`Deploying stable coin contract...`); - const tx = NativeTokenContract.deploy(wallet, 0n, owner).send(); - logger(`Tx sent with hash ${await tx.getTxHash()}`); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); + const receipt = await waitForSuccess(TokenContract.deploy(wallet).send()); logger(`Stable coin asset deployed to ${receipt.contractAddress}`); - stableCoin = await NativeTokenContract.at(receipt.contractAddress!, wallet); + stableCoin = await TokenContract.at(receipt.contractAddress!, wallet); } { logger(`Deploying L2 public contract...`); - const tx = LendingContract.deploy(wallet).send(); - logger(`Tx sent with hash ${await tx.getTxHash()}`); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); + const receipt = await waitForSuccess(LendingContract.deploy(wallet).send()); logger(`CDP deployed at ${receipt.contractAddress}`); lendingContract = await LendingContract.at(receipt.contractAddress!, wallet); } + + await waitForSuccess(collateralAsset.methods._initialize(accounts[0]).send()); + await waitForSuccess(collateralAsset.methods.set_minter({ address: lendingContract.address }, 1).send()); + await waitForSuccess(stableCoin.methods._initialize(accounts[0]).send()); + await waitForSuccess(stableCoin.methods.set_minter({ address: lendingContract.address }, 1).send()); + return { priceFeedContract, lendingContract, collateralAsset, stableCoin }; }; - beforeEach(async () => { + beforeAll(async () => { ({ aztecNode, aztecRpcServer, logger, cheatCodes: cc } = await setup(0)); - { - const privateKey = GrumpkinScalar.random(); - const account = new Account(aztecRpcServer, privateKey, new AuthWitnessAccountContract(privateKey)); - const deployTx = await account.deploy(); - await deployTx.wait({ interval: 0.1 }); - wallet = new AuthWitnessEntrypointWallet( - aztecRpcServer, - (await account.getEntrypoint()) as unknown as IAuthWitnessAccountEntrypoint, - await account.getCompleteAddress(), - ); - accounts = await wallet.getAccounts(); - } - }, 100_000); + const privateKey = GrumpkinScalar.random(); + const account = new Account(aztecRpcServer, privateKey, new AuthWitnessAccountContract(privateKey)); + const deployTx = await account.deploy(); + await deployTx.wait({ interval: 0.1 }); + wallet = new AuthWitnessEntrypointWallet( + aztecRpcServer, + (await account.getEntrypoint()) as unknown as IAuthWitnessAccountEntrypoint, + await account.getCompleteAddress(), + ); + accounts = await wallet.getAccounts(); - afterEach(async () => { + ({ lendingContract, priceFeedContract, collateralAsset, stableCoin } = await deployContracts()); + lendingAccount = new LendingAccount(accounts[0].address, new Fr(42)); + + // Also specified in `noir-contracts/src/contracts/lending_contract/src/main.nr` + const rate = 1268391679n; + lendingSim = new LendingSimulator( + cc, + lendingAccount, + rate, + lendingContract, + new TokenSimulator(collateralAsset, logger, [lendingContract.address, ...accounts.map(a => a.address)]), + new TokenSimulator(stableCoin, logger, [lendingContract.address, ...accounts.map(a => a.address)]), + ); + }, 200_000); + + afterAll(async () => { await aztecNode?.stop(); if (aztecRpcServer instanceof AztecRPCServer) { await aztecRpcServer?.stop(); } }); + afterEach(async () => { + await lendingSim.check(); + }); + const hashPayload = async (payload: Fr[]) => { return pedersenPlookupCompressWithHashIndex( await CircuitsWasm.get(), @@ -111,265 +142,59 @@ describe('e2e_lending_contract', () => { ); }; - // Fetch a storage snapshot from the contract that we can use to compare between transitions. - const getStorageSnapshot = async ( - lendingContract: LendingContract, - collateralAsset: NativeTokenContract, - stableCoin: NativeTokenContract, - account: LendingAccount, - ) => { - logger('Fetching storage snapshot 📸 '); - const accountKey = await account.key(); - - const tot = await lendingContract.methods.get_asset(0).view(); - const privatePos = await lendingContract.methods.get_position(accountKey).view(); - const publicPos = await lendingContract.methods.get_position(account.address.toField()).view(); - const totalCollateral = await collateralAsset.methods.public_balance_of(lendingContract.address).view(); - - return { - interestAccumulator: new Fr(tot['interest_accumulator']), - lastUpdatedTs: new Fr(tot['last_updated_ts']), - privateCollateral: new Fr(privatePos['collateral']), - privateStaticDebt: new Fr(privatePos['static_debt']), - privateDebt: new Fr(privatePos['debt']), - publicCollateral: new Fr(publicPos['collateral']), - publicStaticDebt: new Fr(publicPos['static_debt']), - publicDebt: new Fr(publicPos['debt']), - totalCollateral: new Fr(totalCollateral), - stableCoinLending: new Fr(await stableCoin.methods.public_balance_of(lendingContract.address).view()), - stableCoinPublic: new Fr(await stableCoin.methods.public_balance_of(account.address).view()), - stableCoinPrivate: new Fr(await stableCoin.methods.balance_of(account.address).view()), - stableCoinSupply: new Fr(await stableCoin.methods.total_supply().view()), - }; - }; - - // Convenience struct to hold an account's address and secret that can easily be passed around. - // Contains utilities to compute the "key" for private holdings in the public state. - class LendingAccount { - public readonly address: AztecAddress; - public readonly secret: Fr; - - constructor(address: AztecAddress, secret: Fr) { - this.address = address; - this.secret = secret; - } - - public async key(): Promise { - return Fr.fromBuffer( - pedersenPlookupCommitInputs( - await CircuitsWasm.get(), - [this.address, this.secret].map(f => f.toBuffer()), - ), - ); - } - } - - const muldivDown = (a: bigint, b: bigint, c: bigint) => (a * b) / c; - - const muldivUp = (a: bigint, b: bigint, c: bigint) => { - const adder = (a * b) % c > 0n ? 1n : 0n; - return muldivDown(a, b, c) + adder; - }; - - const computeMultiplier = (rate: bigint, dt: bigint) => { - if (dt == 0n) { - return BASE; - } - - const expMinusOne = dt - 1n; - const expMinusTwo = dt > 2 ? dt - 2n : 0n; - - const basePowerTwo = muldivDown(rate, rate, WAD); - const basePowerThree = muldivDown(basePowerTwo, rate, WAD); - - const temp = dt * expMinusOne; - const secondTerm = muldivDown(temp, basePowerTwo, 2n); - const thirdTerm = muldivDown(temp * expMinusTwo, basePowerThree, 6n); - - const offset = (dt * rate + secondTerm + thirdTerm) / (WAD / BASE); - - return BASE + offset; - }; - - // Helper class that emulates the logic of the lending contract. Used to have a "twin" to check values against. - class LendingSimulator { - public accumulator: bigint = BASE; - public time: number = 0; - - private collateral: { [key: string]: Fr } = {}; - private staticDebt: { [key: string]: Fr } = {}; - private stableBalance: { [key: string]: Fr } = {}; - private repaid: bigint = 0n; - - private key: Fr = Fr.ZERO; - - constructor(private cc: CheatCodes, private account: LendingAccount, private rate: bigint) {} - - async prepare() { - this.key = await this.account.key(); - const ts = await this.cc.eth.timestamp(); - this.time = ts + 10 + (ts % 10); - await this.cc.aztec.warp(this.time); - } - - async progressTime(diff: number) { - this.time = this.time + diff; - await this.cc.aztec.warp(this.time); - this.accumulator = muldivDown(this.accumulator, computeMultiplier(this.rate, BigInt(diff)), BASE); - } - - mintStable(to: Fr, amount: bigint) { - const balance = this.stableBalance[to.toString()] ?? Fr.ZERO; - this.stableBalance[to.toString()] = new Fr(balance.value + amount); - } - - deposit(onBehalfOf: Fr, amount: bigint) { - const coll = this.collateral[onBehalfOf.toString()] ?? Fr.ZERO; - this.collateral[onBehalfOf.toString()] = new Fr(coll.value + amount); - } - - withdraw(owner: Fr, amount: bigint) { - const coll = this.collateral[owner.toString()] ?? Fr.ZERO; - this.collateral[owner.toString()] = new Fr(coll.value - amount); - } - - borrow(owner: Fr, recipient: Fr, amount: bigint) { - const staticDebtBal = this.staticDebt[owner.toString()] ?? Fr.ZERO; - const increase = muldivUp(amount, BASE, this.accumulator); - this.staticDebt[owner.toString()] = new Fr(staticDebtBal.value + increase); - - const balance = this.stableBalance[recipient.toString()] ?? Fr.ZERO; - this.stableBalance[recipient.toString()] = new Fr(balance.value + amount); - } - - repay(owner: Fr, onBehalfOf: Fr, amount: bigint) { - const staticDebtBal = this.staticDebt[onBehalfOf.toString()] ?? Fr.ZERO; - const decrease = muldivDown(amount, BASE, this.accumulator); - this.staticDebt[onBehalfOf.toString()] = new Fr(staticDebtBal.value - decrease); - - const balance = this.stableBalance[owner.toString()] ?? Fr.ZERO; - this.stableBalance[owner.toString()] = new Fr(balance.value - amount); - this.repaid += amount; - } - - check(storage: { [key: string]: Fr }) { - expect(storage['interestAccumulator']).toEqual(new Fr(this.accumulator)); - expect(storage['lastUpdatedTs']).toEqual(new Fr(this.time)); - - // Private values - const keyPriv = this.key.toString(); - expect(storage['privateCollateral']).toEqual(this.collateral[keyPriv] ?? Fr.ZERO); - expect(storage['privateStaticDebt']).toEqual(this.staticDebt[keyPriv] ?? Fr.ZERO); - expect(storage['privateDebt'].value).toEqual( - muldivUp((this.staticDebt[keyPriv] ?? Fr.ZERO).value, this.accumulator, BASE), - ); - - // Public values - const keyPub = this.account.address.toString(); - expect(storage['publicCollateral']).toEqual(this.collateral[keyPub] ?? Fr.ZERO); - expect(storage['publicStaticDebt']).toEqual(this.staticDebt[keyPub] ?? Fr.ZERO); - expect(storage['publicDebt'].value).toEqual( - muldivUp((this.staticDebt[keyPub] ?? Fr.ZERO).value, this.accumulator, BASE), - ); - - const totalCollateral = Object.values(this.collateral).reduce((a, b) => new Fr(a.value + b.value), Fr.ZERO); - expect(storage['totalCollateral']).toEqual(totalCollateral); - - expect(storage['stableCoinLending'].value).toEqual(this.repaid); - expect(storage['stableCoinPublic']).toEqual(this.stableBalance[keyPub] ?? Fr.ZERO); - - // Abusing notation and using the `keyPriv` as if an address for private holdings of stable_coin while it has the same owner in reality. - expect(storage['stableCoinPrivate']).toEqual(this.stableBalance[keyPriv] ?? Fr.ZERO); - - const totalStableSupply = Object.values(this.stableBalance).reduce((a, b) => new Fr(a.value + b.value), Fr.ZERO); - // @todo @lherskind To be updated such that we burn assets on repay instead. - expect(storage['stableCoinSupply'].value).toEqual(totalStableSupply.value + this.repaid); - } - } - - it('Full lending run-through', async () => { - // Gotta use the actual auth witness account here and not the standard wallet. - const recipientFull = accounts[0]; - const recipient = recipientFull.address; - - const { lendingContract, priceFeedContract, collateralAsset, stableCoin } = await deployContracts(recipient); - - const lendingAccount = new LendingAccount(recipient, new Fr(42)); - - const storageSnapshots: { [key: string]: { [key: string]: Fr } } = {}; - - const setPrice = async (newPrice: bigint) => { - const tx = priceFeedContract.methods.set_price(0n, newPrice).send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - }; - - await setPrice(2n * 10n ** 9n); + it('Mint assets for later usage', async () => { + await waitForSuccess(priceFeedContract.methods.set_price(0n, 2n * 10n ** 9n).send()); { - // Minting some collateral in public so we got it at hand. - const tx = collateralAsset.methods.owner_mint_pub(lendingAccount.address, 10000n).send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - - const tx2 = collateralAsset.methods.approve(lendingContract.address, 10000n).send(); - const receipt2 = await tx2.wait(); - expect(receipt2.status).toBe(TxStatus.MINED); - - // Minting some collateral in private so we got it at hand. - const secret = Fr.random(); - const secretHash = await computeMessageSecretHash(secret); - const shieldAmount = 10000n; - const tx3 = stableCoin.methods.owner_mint_priv(shieldAmount, secretHash).send(); - const receipt3 = await tx3.wait(); - expect(receipt3.status).toBe(TxStatus.MINED); - - const tx4 = stableCoin.methods.redeemShield(shieldAmount, secret, recipient).send(); - const receipt4 = await tx4.wait(); - expect(receipt4.status).toBe(TxStatus.MINED); - - const tx5 = stableCoin.methods.approve(lendingContract.address, 10000n).send(); - const receipt5 = await tx5.wait(); - expect(receipt5.status).toBe(TxStatus.MINED); + const assets = [collateralAsset, stableCoin]; + const mintAmount = 10000n; + for (const asset of assets) { + const secret = Fr.random(); + const secretHash = await computeMessageSecretHash(secret); + + const a = asset.methods.mint_public({ address: lendingAccount.address }, mintAmount).send(); + const b = asset.methods.mint_private(mintAmount, secretHash).send(); + + await Promise.all([a, b].map(waitForSuccess)); + await waitForSuccess( + asset.methods.redeem_shield({ address: lendingAccount.address }, mintAmount, secret).send(), + ); + } } - // Also specified in `noir-contracts/src/contracts/lending_contract/src/main.nr` - const rate = 1268391679n; - const lendingSim = new LendingSimulator(cc, lendingAccount, rate); - await lendingSim.prepare(); - // To handle initial mint (we use these funds to refund privately without shielding first). - lendingSim.mintStable(await lendingAccount.key(), 10000n); + lendingSim.mintStableCoinOutsideLoan(lendingAccount.address, 10000n, true); + lendingSim.stableCoin.redeemShield(lendingAccount.address, 10000n); + lendingSim.mintStableCoinOutsideLoan(lendingAccount.address, 10000n, false); - { - // Initialize the contract values, setting the interest accumulator to 1e9 and the last updated timestamp to now. - logger('Initializing contract'); - const tx = lendingContract.methods - .init(priceFeedContract.address, 8000, collateralAsset.address, stableCoin.address) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['initial'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, - ); + lendingSim.collateralAsset.mintPrivate(10000n); + lendingSim.collateralAsset.redeemShield(lendingAccount.address, 10000n); + lendingSim.collateralAsset.mintPublic(lendingAccount.address, 10000n); + }); - lendingSim.check(storageSnapshots['initial']); - } + it('Initialize the contract', async () => { + await lendingSim.prepare(); + logger('Initializing contract'); + await waitForSuccess( + lendingContract.methods.init(priceFeedContract.address, 8000, collateralAsset.address, stableCoin.address).send(), + ); + }); - { + describe('Deposits', () => { + it('Depositing 🥸 : 💰 -> 🏦', async () => { const depositAmount = 420n; - + const nonce = Fr.random(); const messageHash = await hashPayload([ - FunctionSelector.fromSignature('unshieldTokens(Field,Field,Field)').toField(), - recipientFull.address.toField(), + lendingContract.address.toField(), + collateralAsset.address.toField(), + FunctionSelector.fromSignature('unshield((Field),(Field),Field,Field)').toField(), + lendingAccount.address.toField(), lendingContract.address.toField(), new Fr(depositAmount), + nonce, ]); await wallet.signAndAddAuthWitness(messageHash); - await lendingSim.progressTime(10); - lendingSim.deposit(await lendingAccount.key(), depositAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.depositPrivate(lendingAccount.address, await lendingAccount.key(), depositAmount); // Make a private deposit of funds into own account. // This should: @@ -377,58 +202,76 @@ describe('e2e_lending_contract', () => { // - increase last updated timestamp. // - increase the private collateral. logger('Depositing 🥸 : 💰 -> 🏦'); - const tx = lendingContract.methods - .deposit_private(lendingAccount.secret, lendingAccount.address, 0n, depositAmount, collateralAsset.address) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['private_deposit'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods + .deposit_private( + lendingAccount.address, + depositAmount, + nonce, + lendingAccount.secret, + 0n, + collateralAsset.address, + ) + .send(), ); + }); - lendingSim.check(storageSnapshots['private_deposit']); - } - - { + it('Depositing 🥸 on behalf of recipient: 💰 -> 🏦', async () => { const depositAmount = 421n; + const nonce = Fr.random(); const messageHash = await hashPayload([ - FunctionSelector.fromSignature('unshieldTokens(Field,Field,Field)').toField(), - recipientFull.address.toField(), + lendingContract.address.toField(), + collateralAsset.address.toField(), + FunctionSelector.fromSignature('unshield((Field),(Field),Field,Field)').toField(), + lendingAccount.address.toField(), lendingContract.address.toField(), new Fr(depositAmount), + nonce, ]); await wallet.signAndAddAuthWitness(messageHash); - await lendingSim.progressTime(10); - lendingSim.deposit(recipient.toField(), depositAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.depositPrivate(lendingAccount.address, lendingAccount.address.toField(), depositAmount); // Make a private deposit of funds into another account, in this case, a public account. // This should: // - increase the interest accumulator // - increase last updated timestamp. // - increase the public collateral. logger('Depositing 🥸 on behalf of recipient: 💰 -> 🏦'); - const tx = lendingContract.methods - .deposit_private(0n, lendingAccount.address, recipient.toField(), depositAmount, collateralAsset.address) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['private_deposit_on_behalf'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods + .deposit_private( + lendingAccount.address, + depositAmount, + nonce, + 0n, + lendingAccount.address, + collateralAsset.address, + ) + .send(), ); + }); - lendingSim.check(storageSnapshots['private_deposit_on_behalf']); - } - - { + it('Depositing: 💰 -> 🏦', async () => { const depositAmount = 211n; - await lendingSim.progressTime(10); - lendingSim.deposit(recipient.toField(), depositAmount); + + const nonce = Fr.random(); + const messageHash = await hashPayload([ + lendingContract.address.toField(), + collateralAsset.address.toField(), + FunctionSelector.fromSignature('transfer_public((Field),(Field),Field,Field)').toField(), + lendingAccount.address.toField(), + lendingContract.address.toField(), + new Fr(depositAmount), + nonce, + ]); + + // Add it to the wallet as approved + const me = await SchnorrAuthWitnessAccountContract.at(accounts[0].address, wallet); + await waitForSuccess(me.methods.set_is_valid_storage(messageHash, 1).send()); + + await lendingSim.progressTime(TIME_JUMP); + lendingSim.depositPublic(lendingAccount.address, lendingAccount.address.toField(), depositAmount); // Make a public deposit of funds into self. // This should: @@ -437,24 +280,31 @@ describe('e2e_lending_contract', () => { // - increase the public collateral. logger('Depositing: 💰 -> 🏦'); - const tx = lendingContract.methods - .deposit_public(lendingAccount.address, depositAmount, collateralAsset.address) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['public_deposit'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods + .deposit_public(depositAmount, nonce, lendingAccount.address, collateralAsset.address) + .send(), ); - lendingSim.check(storageSnapshots['public_deposit']); - } + }); + describe('failure cases', () => { + it('calling internal _deposit function directly', async () => { + // Try to call the internal `_deposit` function directly + // This should: + // - not change any storage values. + // - fail + + await expect( + lendingContract.methods._deposit(lendingAccount.address.toField(), 42n, collateralAsset.address).simulate(), + ).rejects.toThrow(); + }); + }); + }); - { + describe('Borrow', () => { + it('Borrow 🥸 : 🏦 -> 🍌', async () => { const borrowAmount = 69n; - await lendingSim.progressTime(10); - lendingSim.borrow(await lendingAccount.key(), lendingAccount.address.toField(), borrowAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.borrow(await lendingAccount.key(), lendingAccount.address, borrowAmount); // Make a private borrow using the private account // This should: @@ -463,25 +313,15 @@ describe('e2e_lending_contract', () => { // - increase the private debt. logger('Borrow 🥸 : 🏦 -> 🍌'); - const tx = lendingContract.methods - .borrow_private(lendingAccount.secret, lendingAccount.address, borrowAmount) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['private_borrow'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods.borrow_private(lendingAccount.secret, lendingAccount.address, borrowAmount).send(), ); + }); - lendingSim.check(storageSnapshots['private_borrow']); - } - - { + it('Borrow: 🏦 -> 🍌', async () => { const borrowAmount = 69n; - await lendingSim.progressTime(10); - lendingSim.borrow(recipient.toField(), lendingAccount.address.toField(), borrowAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.borrow(lendingAccount.address.toField(), lendingAccount.address, borrowAmount); // Make a public borrow using the private account // This should: @@ -490,31 +330,26 @@ describe('e2e_lending_contract', () => { // - increase the public debt. logger('Borrow: 🏦 -> 🍌'); - const tx = lendingContract.methods.borrow_public(lendingAccount.address, borrowAmount).send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['public_borrow'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, - ); - - lendingSim.check(storageSnapshots['public_borrow']); - } + await waitForSuccess(lendingContract.methods.borrow_public(lendingAccount.address, borrowAmount).send()); + }); + }); - { + describe('Repay', () => { + it('Repay 🥸 : 🍌 -> 🏦', async () => { const repayAmount = 20n; + const nonce = Fr.random(); const messageHash = await hashPayload([ - FunctionSelector.fromSignature('unshieldTokens(Field,Field,Field)').toField(), - recipientFull.address.toField(), lendingContract.address.toField(), + stableCoin.address.toField(), + FunctionSelector.fromSignature('burn((Field),Field,Field)').toField(), + lendingAccount.address.toField(), new Fr(repayAmount), + nonce, ]); await wallet.signAndAddAuthWitness(messageHash); - await lendingSim.progressTime(10); - lendingSim.repay(await lendingAccount.key(), await lendingAccount.key(), repayAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.repayPrivate(lendingAccount.address, await lendingAccount.key(), repayAmount); // Make a private repay of the debt in the private account // This should: @@ -523,33 +358,28 @@ describe('e2e_lending_contract', () => { // - decrease the private debt. logger('Repay 🥸 : 🍌 -> 🏦'); - const tx = lendingContract.methods - .repay_private(lendingAccount.secret, lendingAccount.address, 0n, repayAmount, stableCoin.address) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['private_repay'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods + .repay_private(lendingAccount.address, repayAmount, nonce, lendingAccount.secret, 0n, stableCoin.address) + .send(), ); + }); - lendingSim.check(storageSnapshots['private_repay']); - } - - { + it('Repay 🥸 on behalf of public: 🍌 -> 🏦', async () => { const repayAmount = 21n; + const nonce = Fr.random(); const messageHash = await hashPayload([ - FunctionSelector.fromSignature('unshieldTokens(Field,Field,Field)').toField(), - recipientFull.address.toField(), lendingContract.address.toField(), + stableCoin.address.toField(), + FunctionSelector.fromSignature('burn((Field),Field,Field)').toField(), + lendingAccount.address.toField(), new Fr(repayAmount), + nonce, ]); await wallet.signAndAddAuthWitness(messageHash); - await lendingSim.progressTime(10); - lendingSim.repay(await lendingAccount.key(), lendingAccount.address.toField(), repayAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.repayPrivate(lendingAccount.address, lendingAccount.address.toField(), repayAmount); // Make a private repay of the debt in the public account // This should: @@ -558,25 +388,32 @@ describe('e2e_lending_contract', () => { // - decrease the public debt. logger('Repay 🥸 on behalf of public: 🍌 -> 🏦'); - const tx = lendingContract.methods - .repay_private(0n, lendingAccount.address, recipient.toField(), repayAmount, stableCoin.address) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['private_repay_on_behalf'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods + .repay_private(lendingAccount.address, repayAmount, nonce, 0n, lendingAccount.address, stableCoin.address) + .send(), ); + }); - lendingSim.check(storageSnapshots['private_repay_on_behalf']); - } - - { + it('Repay: 🍌 -> 🏦', async () => { const repayAmount = 20n; - await lendingSim.progressTime(10); - lendingSim.repay(lendingAccount.address.toField(), lendingAccount.address.toField(), repayAmount); + + const nonce = Fr.random(); + const messageHash = await hashPayload([ + lendingContract.address.toField(), + stableCoin.address.toField(), + FunctionSelector.fromSignature('burn_public((Field),Field,Field)').toField(), + lendingAccount.address.toField(), + new Fr(repayAmount), + nonce, + ]); + + // Add it to the wallet as approved + const me = await SchnorrAuthWitnessAccountContract.at(accounts[0].address, wallet); + await waitForSuccess(me.methods.set_is_valid_storage(messageHash, 1).send()); + + await lendingSim.progressTime(TIME_JUMP); + lendingSim.repayPublic(lendingAccount.address, lendingAccount.address.toField(), repayAmount); // Make a public repay of the debt in the public account // This should: @@ -585,32 +422,17 @@ describe('e2e_lending_contract', () => { // - decrease the public debt. logger('Repay: 🍌 -> 🏦'); - const tx = lendingContract.methods.repay_public(recipient.toField(), 20n, stableCoin.address).send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['public_repay'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods.repay_public(repayAmount, nonce, lendingAccount.address, stableCoin.address).send(), ); + }); + }); - lendingSim.check(storageSnapshots['public_repay']); - } - - { - // Withdraw more than possible to test the revert. - logger('Withdraw: trying to withdraw more than possible'); - const tx = lendingContract.methods.withdraw_public(recipient, 10n ** 9n).send({ skipPublicSimulation: true }); - await tx.isMined({ interval: 0.1 }); - const receipt = await tx.getReceipt(); - expect(receipt.status).toBe(TxStatus.DROPPED); - } - - { + describe('Withdraw', () => { + it('Withdraw: 🏦 -> 💰', async () => { const withdrawAmount = 42n; - await lendingSim.progressTime(10); - lendingSim.withdraw(recipient.toField(), withdrawAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.withdraw(lendingAccount.address.toField(), lendingAccount.address, withdrawAmount); // Withdraw funds from the public account // This should: @@ -619,23 +441,13 @@ describe('e2e_lending_contract', () => { // - decrease the public collateral. logger('Withdraw: 🏦 -> 💰'); - const tx = lendingContract.methods.withdraw_public(recipient, withdrawAmount).send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['public_withdraw'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, - ); - - lendingSim.check(storageSnapshots['public_withdraw']); - } + await waitForSuccess(lendingContract.methods.withdraw_public(lendingAccount.address, withdrawAmount).send()); + }); - { + it('Withdraw 🥸 : 🏦 -> 💰', async () => { const withdrawAmount = 42n; - await lendingSim.progressTime(10); - lendingSim.withdraw(await lendingAccount.key(), withdrawAmount); + await lendingSim.progressTime(TIME_JUMP); + lendingSim.withdraw(await lendingAccount.key(), lendingAccount.address, withdrawAmount); // Withdraw funds from the private account // This should: @@ -644,41 +456,19 @@ describe('e2e_lending_contract', () => { // - decrease the private collateral. logger('Withdraw 🥸 : 🏦 -> 💰'); - const tx = lendingContract.methods - .withdraw_private(lendingAccount.secret, lendingAccount.address, withdrawAmount) - .send(); - const receipt = await tx.wait(); - expect(receipt.status).toBe(TxStatus.MINED); - storageSnapshots['private_withdraw'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, + await waitForSuccess( + lendingContract.methods.withdraw_private(lendingAccount.secret, lendingAccount.address, withdrawAmount).send(), ); - - lendingSim.check(storageSnapshots['private_withdraw']); - } - - { - // Try to call the internal `_deposit` function directly - // This should: - // - not change any storage values. - // - fail - - const tx = lendingContract.methods - ._deposit(recipient.toField(), 42n, collateralAsset.address) - .send({ skipPublicSimulation: true }); - await tx.isMined({ interval: 0.1 }); - const receipt = await tx.getReceipt(); - expect(receipt.status).toBe(TxStatus.DROPPED); - logger('Rejected call directly to internal function 🧚 '); - storageSnapshots['attempted_internal_deposit'] = await getStorageSnapshot( - lendingContract, - collateralAsset, - stableCoin, - lendingAccount, - ); - expect(storageSnapshots['private_withdraw']).toEqual(storageSnapshots['attempted_internal_deposit']); - } - }, 650_000); + }); + + describe('failure cases', () => { + it('withdraw more than possible to revert', async () => { + // Withdraw more than possible to test the revert. + logger('Withdraw: trying to withdraw more than possible'); + await expect( + lendingContract.methods.withdraw_public(lendingAccount.address, 10n ** 9n).simulate(), + ).rejects.toThrow(); + }); + }); + }); }); diff --git a/yarn-project/end-to-end/src/e2e_token_contract.test.ts b/yarn-project/end-to-end/src/e2e_token_contract.test.ts index 649a53cdee6..7e8defafae4 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract.test.ts @@ -4,7 +4,6 @@ import { Account, AuthWitnessAccountContract, AuthWitnessEntrypointWallet, - AztecAddress, IAuthWitnessAccountEntrypoint, computeMessageSecretHash, } from '@aztec/aztec.js'; @@ -24,6 +23,7 @@ import { AztecRPC, TxStatus } from '@aztec/types'; import { jest } from '@jest/globals'; import { setup } from './fixtures/utils.js'; +import { TokenSimulator } from './simulators/token_simulator.js'; const hashPayload = async (payload: Fr[]) => { return pedersenPlookupCompressWithHashIndex( @@ -35,95 +35,6 @@ const hashPayload = async (payload: Fr[]) => { const TIMEOUT = 60_000; -class TokenSimulator { - private balancesPrivate: Map = new Map(); - private balancePublic: Map = new Map(); - public totalSupply: bigint = 0n; - - constructor(protected token: TokenContract, protected logger: DebugLogger, protected accounts: CompleteAddress[]) {} - - public mintPrivate(to: AztecAddress, amount: bigint) { - this.totalSupply += amount; - } - - public mintPublic(to: AztecAddress, amount: bigint) { - this.totalSupply += amount; - const value = this.balancePublic.get(to) || 0n; - this.balancePublic.set(to, value + amount); - } - - public transferPublic(from: AztecAddress, to: AztecAddress, amount: bigint) { - const fromBalance = this.balancePublic.get(from) || 0n; - this.balancePublic.set(from, fromBalance - amount); - expect(fromBalance).toBeGreaterThanOrEqual(amount); - - const toBalance = this.balancePublic.get(to) || 0n; - this.balancePublic.set(to, toBalance + amount); - } - - public transferPrivate(from: AztecAddress, to: AztecAddress, amount: bigint) { - const fromBalance = this.balancesPrivate.get(from) || 0n; - expect(fromBalance).toBeGreaterThanOrEqual(amount); - this.balancesPrivate.set(from, fromBalance - amount); - - const toBalance = this.balancesPrivate.get(to) || 0n; - this.balancesPrivate.set(to, toBalance + amount); - } - - public shield(from: AztecAddress, amount: bigint) { - const fromBalance = this.balancePublic.get(from) || 0n; - expect(fromBalance).toBeGreaterThanOrEqual(amount); - this.balancePublic.set(from, fromBalance - amount); - } - - public redeemShield(to: AztecAddress, amount: bigint) { - const toBalance = this.balancesPrivate.get(to) || 0n; - this.balancesPrivate.set(to, toBalance + amount); - } - - public unshield(from: AztecAddress, to: AztecAddress, amount: bigint) { - const fromBalance = this.balancesPrivate.get(from) || 0n; - const toBalance = this.balancePublic.get(to) || 0n; - expect(fromBalance).toBeGreaterThanOrEqual(amount); - this.balancesPrivate.set(from, fromBalance - amount); - this.balancePublic.set(to, toBalance + amount); - } - - public burnPrivate(from: AztecAddress, amount: bigint) { - const fromBalance = this.balancesPrivate.get(from) || 0n; - expect(fromBalance).toBeGreaterThanOrEqual(amount); - this.balancesPrivate.set(from, fromBalance - amount); - - this.totalSupply -= amount; - } - - public burnPublic(from: AztecAddress, amount: bigint) { - const fromBalance = this.balancePublic.get(from) || 0n; - expect(fromBalance).toBeGreaterThanOrEqual(amount); - this.balancePublic.set(from, fromBalance - amount); - - this.totalSupply -= amount; - } - - public balanceOfPublic(address: AztecAddress) { - return this.balancePublic.get(address) || 0n; - } - - public balanceOfPrivate(address: AztecAddress) { - return this.balancesPrivate.get(address) || 0n; - } - - public async check() { - expect(await this.token.methods.total_supply().view()).toEqual(this.totalSupply); - - // Check that all our public matches - for (const { address } of this.accounts) { - expect(await this.token.methods.balance_of_public({ address }).view()).toEqual(this.balanceOfPublic(address)); - expect(await this.token.methods.balance_of_private({ address }).view()).toEqual(this.balanceOfPrivate(address)); - } - } -} - describe('e2e_token_contract', () => { jest.setTimeout(TIMEOUT); @@ -173,7 +84,11 @@ describe('e2e_token_contract', () => { asset = await TokenContract.at(receipt.contractAddress!, wallets[0]); } - tokenSim = new TokenSimulator(asset, logger, accounts); + tokenSim = new TokenSimulator( + asset, + logger, + accounts.map(account => account.address), + ); { const initializeTx = asset.methods._initialize({ address: accounts[0].address }).send(); @@ -299,7 +214,7 @@ describe('e2e_token_contract', () => { const tx = asset.methods.mint_private(amount, secretHash).send(); const receipt = await tx.wait(); expect(receipt.status).toBe(TxStatus.MINED); - tokenSim.mintPrivate(accounts[0].address, amount); + tokenSim.mintPrivate(amount); }); it('redeem as recipient', async () => { @@ -976,8 +891,8 @@ describe('e2e_token_contract', () => { caller.address.toField(), asset.address.toField(), FunctionSelector.fromSignature('unshield((Field),(Field),Field,Field)').toField(), - accounts[0].address.toField(), - accounts[1].address.toField(), + from.address.toField(), + to.address.toField(), new Fr(amount), nonce, ]); diff --git a/yarn-project/end-to-end/src/simulators/index.ts b/yarn-project/end-to-end/src/simulators/index.ts new file mode 100644 index 00000000000..bf023483a1b --- /dev/null +++ b/yarn-project/end-to-end/src/simulators/index.ts @@ -0,0 +1,2 @@ +export * from './lending_simulator.js'; +export * from './token_simulator.js'; diff --git a/yarn-project/end-to-end/src/simulators/lending_simulator.ts b/yarn-project/end-to-end/src/simulators/lending_simulator.ts new file mode 100644 index 00000000000..37333ab4572 --- /dev/null +++ b/yarn-project/end-to-end/src/simulators/lending_simulator.ts @@ -0,0 +1,187 @@ +// Convenience struct to hold an account's address and secret that can easily be passed around. +import { CheatCodes } from '@aztec/aztec.js'; +import { AztecAddress, CircuitsWasm, Fr } from '@aztec/circuits.js'; +import { pedersenPlookupCommitInputs } from '@aztec/circuits.js/barretenberg'; +import { LendingContract } from '@aztec/noir-contracts/types'; + +import { TokenSimulator } from './token_simulator.js'; + +/** + * Contains utilities to compute the "key" for private holdings in the public state. + */ +export class LendingAccount { + /** The address that owns this account */ + public readonly address: AztecAddress; + /** The secret used for private deposits */ + public readonly secret: Fr; + + constructor(address: AztecAddress, secret: Fr) { + this.address = address; + this.secret = secret; + } + + /** + * Computes the key for the private holdings of this account. + * @returns Key in public space + */ + public async key(): Promise { + return Fr.fromBuffer( + pedersenPlookupCommitInputs( + await CircuitsWasm.get(), + [this.address, this.secret].map(f => f.toBuffer()), + ), + ); + } +} + +const WAD = 10n ** 18n; +const BASE = 10n ** 9n; + +const muldivDown = (a: bigint, b: bigint, c: bigint) => (a * b) / c; + +const muldivUp = (a: bigint, b: bigint, c: bigint) => { + const adder = (a * b) % c > 0n ? 1n : 0n; + return muldivDown(a, b, c) + adder; +}; + +const computeMultiplier = (rate: bigint, dt: bigint) => { + if (dt == 0n) { + return BASE; + } + + const expMinusOne = dt - 1n; + const expMinusTwo = dt > 2 ? dt - 2n : 0n; + + const basePowerTwo = muldivDown(rate, rate, WAD); + const basePowerThree = muldivDown(basePowerTwo, rate, WAD); + + const temp = dt * expMinusOne; + const secondTerm = muldivDown(temp, basePowerTwo, 2n); + const thirdTerm = muldivDown(temp * expMinusTwo, basePowerThree, 6n); + + const offset = (dt * rate + secondTerm + thirdTerm) / (WAD / BASE); + + return BASE + offset; +}; + +/** + * Helper class that emulates the logic of the lending contract. Used to have a "twin" to check values against. + */ +export class LendingSimulator { + /** interest rate accumulator */ + public accumulator: bigint = 0n; + /** the timestamp of the simulator*/ + public time: number = 0; + + private collateral: { [key: string]: Fr } = {}; + private staticDebt: { [key: string]: Fr } = {}; + private borrowed: bigint = 0n; + private mintedOutside: bigint = 0n; + + constructor( + private cc: CheatCodes, + private account: LendingAccount, + private rate: bigint, + /** the lending contract */ + public lendingContract: LendingContract, + /** the collateral asset used in the lending contract */ + public collateralAsset: TokenSimulator, + /** the stable-coin borrowed in the lending contract */ + public stableCoin: TokenSimulator, + ) {} + + async prepare() { + this.accumulator = BASE; + const ts = await this.cc.eth.timestamp(); + this.time = ts + 10 + (ts % 10); + await this.cc.aztec.warp(this.time); + } + + async progressTime(diff: number) { + this.time = this.time + diff; + await this.cc.aztec.warp(this.time); + this.accumulator = muldivDown(this.accumulator, computeMultiplier(this.rate, BigInt(diff)), BASE); + } + + depositPrivate(from: AztecAddress, onBehalfOf: Fr, amount: bigint) { + this.collateralAsset.unshield(from, this.lendingContract.address, amount); + this.deposit(onBehalfOf, amount); + } + + depositPublic(from: AztecAddress, onBehalfOf: Fr, amount: bigint) { + this.collateralAsset.transferPublic(from, this.lendingContract.address, amount); + this.deposit(onBehalfOf, amount); + } + + private deposit(onBehalfOf: Fr, amount: bigint) { + const coll = this.collateral[onBehalfOf.toString()] ?? Fr.ZERO; + this.collateral[onBehalfOf.toString()] = new Fr(coll.value + amount); + } + + withdraw(owner: Fr, recipient: AztecAddress, amount: bigint) { + const coll = this.collateral[owner.toString()] ?? Fr.ZERO; + this.collateral[owner.toString()] = new Fr(coll.value - amount); + this.collateralAsset.transferPublic(this.lendingContract.address, recipient, amount); + } + + borrow(owner: Fr, recipient: AztecAddress, amount: bigint) { + const staticDebtBal = this.staticDebt[owner.toString()] ?? Fr.ZERO; + const increase = muldivUp(amount, BASE, this.accumulator); + this.staticDebt[owner.toString()] = new Fr(staticDebtBal.value + increase); + + this.stableCoin.mintPublic(recipient, amount); + this.borrowed += amount; + } + + repayPrivate(from: AztecAddress, onBehalfOf: Fr, amount: bigint) { + this.stableCoin.burnPrivate(from, amount); + this.repay(onBehalfOf, onBehalfOf, amount); + } + + repayPublic(from: AztecAddress, onBehalfOf: Fr, amount: bigint) { + this.stableCoin.burnPublic(from, amount); + this.repay(onBehalfOf, onBehalfOf, amount); + } + + private repay(from: Fr, onBehalfOf: Fr, amount: bigint) { + const staticDebtBal = this.staticDebt[onBehalfOf.toString()] ?? Fr.ZERO; + const decrease = muldivDown(amount, BASE, this.accumulator); + this.staticDebt[onBehalfOf.toString()] = new Fr(staticDebtBal.value - decrease); + + this.borrowed -= amount; + } + + mintStableCoinOutsideLoan(recipient: AztecAddress, amount: bigint, priv = false) { + if (priv) { + this.stableCoin.mintPrivate(amount); + } else { + this.stableCoin.mintPublic(recipient, amount); + } + this.mintedOutside += amount; + } + + async check() { + // Run checks on both underlying assets + await this.collateralAsset.check(); + await this.stableCoin.check(); + + // Check that total collateral equals total holdings by contract. + const totalCollateral = Object.values(this.collateral).reduce((a, b) => new Fr(a.value + b.value), Fr.ZERO); + expect(totalCollateral).toEqual(new Fr(this.collateralAsset.balanceOfPublic(this.lendingContract.address))); + + expect(this.borrowed).toEqual(this.stableCoin.totalSupply - this.mintedOutside); + + const asset = await this.lendingContract.methods.get_asset(0).view(); + expect(asset['interest_accumulator']).toEqual(this.accumulator); + expect(asset['last_updated_ts']).toEqual(BigInt(this.time)); + + for (const key of [this.account.address, await this.account.key()]) { + const privatePos = await this.lendingContract.methods.get_position(key).view(); + expect(new Fr(privatePos['collateral'])).toEqual(this.collateral[key.toString()] ?? Fr.ZERO); + expect(new Fr(privatePos['static_debt'])).toEqual(this.staticDebt[key.toString()] ?? Fr.ZERO); + expect(privatePos['debt']).toEqual( + muldivUp((this.staticDebt[key.toString()] ?? Fr.ZERO).value, this.accumulator, BASE), + ); + } + } +} diff --git a/yarn-project/end-to-end/src/simulators/token_simulator.ts b/yarn-project/end-to-end/src/simulators/token_simulator.ts new file mode 100644 index 00000000000..b28c5480d6e --- /dev/null +++ b/yarn-project/end-to-end/src/simulators/token_simulator.ts @@ -0,0 +1,93 @@ +/* eslint-disable jsdoc/require-jsdoc */ +import { DebugLogger } from '@aztec/aztec.js'; +import { AztecAddress } from '@aztec/circuits.js'; +import { TokenContract } from '@aztec/noir-contracts/types'; + +export class TokenSimulator { + private balancesPrivate: Map = new Map(); + private balancePublic: Map = new Map(); + public totalSupply: bigint = 0n; + + constructor(protected token: TokenContract, protected logger: DebugLogger, protected accounts: AztecAddress[]) {} + + public mintPrivate(amount: bigint) { + this.totalSupply += amount; + } + + public mintPublic(to: AztecAddress, amount: bigint) { + this.totalSupply += amount; + const value = this.balancePublic.get(to) || 0n; + this.balancePublic.set(to, value + amount); + } + + public transferPublic(from: AztecAddress, to: AztecAddress, amount: bigint) { + const fromBalance = this.balancePublic.get(from) || 0n; + this.balancePublic.set(from, fromBalance - amount); + expect(fromBalance).toBeGreaterThanOrEqual(amount); + + const toBalance = this.balancePublic.get(to) || 0n; + this.balancePublic.set(to, toBalance + amount); + } + + public transferPrivate(from: AztecAddress, to: AztecAddress, amount: bigint) { + const fromBalance = this.balancesPrivate.get(from) || 0n; + expect(fromBalance).toBeGreaterThanOrEqual(amount); + this.balancesPrivate.set(from, fromBalance - amount); + + const toBalance = this.balancesPrivate.get(to) || 0n; + this.balancesPrivate.set(to, toBalance + amount); + } + + public shield(from: AztecAddress, amount: bigint) { + const fromBalance = this.balancePublic.get(from) || 0n; + expect(fromBalance).toBeGreaterThanOrEqual(amount); + this.balancePublic.set(from, fromBalance - amount); + } + + public redeemShield(to: AztecAddress, amount: bigint) { + const toBalance = this.balancesPrivate.get(to) || 0n; + this.balancesPrivate.set(to, toBalance + amount); + } + + public unshield(from: AztecAddress, to: AztecAddress, amount: bigint) { + const fromBalance = this.balancesPrivate.get(from) || 0n; + const toBalance = this.balancePublic.get(to) || 0n; + expect(fromBalance).toBeGreaterThanOrEqual(amount); + this.balancesPrivate.set(from, fromBalance - amount); + this.balancePublic.set(to, toBalance + amount); + } + + public burnPrivate(from: AztecAddress, amount: bigint) { + const fromBalance = this.balancesPrivate.get(from) || 0n; + expect(fromBalance).toBeGreaterThanOrEqual(amount); + this.balancesPrivate.set(from, fromBalance - amount); + + this.totalSupply -= amount; + } + + public burnPublic(from: AztecAddress, amount: bigint) { + const fromBalance = this.balancePublic.get(from) || 0n; + expect(fromBalance).toBeGreaterThanOrEqual(amount); + this.balancePublic.set(from, fromBalance - amount); + + this.totalSupply -= amount; + } + + public balanceOfPublic(address: AztecAddress) { + return this.balancePublic.get(address) || 0n; + } + + public balanceOfPrivate(address: AztecAddress) { + return this.balancesPrivate.get(address) || 0n; + } + + public async check() { + expect(await this.token.methods.total_supply().view()).toEqual(this.totalSupply); + + // Check that all our public matches + for (const address of this.accounts) { + expect(await this.token.methods.balance_of_public({ address }).view()).toEqual(this.balanceOfPublic(address)); + expect(await this.token.methods.balance_of_private({ address }).view()).toEqual(this.balanceOfPrivate(address)); + } + } +} diff --git a/yarn-project/noir-contracts/src/contracts/lending_contract/src/interfaces.nr b/yarn-project/noir-contracts/src/contracts/lending_contract/src/interfaces.nr index c79378d7ecf..f8b7a9a41b5 100644 --- a/yarn-project/noir-contracts/src/contracts/lending_contract/src/interfaces.nr +++ b/yarn-project/noir-contracts/src/contracts/lending_contract/src/interfaces.nr @@ -36,36 +36,44 @@ impl Token { Self { address } } - fn transfer_pub(self: Self, context: PublicContext, to: Field, amount: Field) { + fn transfer_public(self: Self, context: PublicContext, from: Field, to: Field, amount: Field, nonce: Field) { let _transfer_return_values = context.call_public_function( self.address, - compute_selector("transfer_pub(Field,Field)"), - [to, amount] + compute_selector("transfer_public((Field),(Field),Field,Field)"), + [from, to, amount, nonce] ); } - fn transfer_from_pub(self: Self, context: PublicContext, from: Field, to: Field, amount: Field) { - let _transfer_return_values = context.call_public_function( + fn mint_public(self: Self, context: PublicContext, to: Field, amount: Field) { + let _return_values = context.call_public_function( self.address, - compute_selector("transfer_from_pub(Field,Field,Field)"), - [from, to, amount] + compute_selector("mint_public((Field),Field)"), + [to, amount] ); } - fn owner_mint_pub(self: Self, context: PublicContext, to: Field, amount: Field) { - let _transfer_return_values = context.call_public_function( - self.address, - compute_selector("owner_mint_pub(Field,Field)"), - [to, amount] + fn burn_public(self: Self, context: PublicContext, from: Field, amount: Field, nonce: Field){ + let _return_values = context.call_public_function( + self.address, + compute_selector("burn_public((Field),Field,Field)"), + [from, amount, nonce] ); } // Private - fn unshield(self: Self, context: &mut PrivateContext, from: Field, to: Field, amount: Field) -> [Field; RETURN_VALUES_LENGTH] { + fn unshield(self: Self, context: &mut PrivateContext, from: Field, to: Field, amount: Field, nonce: Field) -> [Field; RETURN_VALUES_LENGTH] { + context.call_private_function( + self.address, + compute_selector("unshield((Field),(Field),Field,Field)"), + [from, to, amount, nonce] + ) + } + + fn burn(self: Self, context: &mut PrivateContext, from: Field, amount: Field, nonce: Field) -> [Field; RETURN_VALUES_LENGTH] { context.call_private_function( self.address, - compute_selector("unshieldTokens(Field,Field,Field)"), - [from, to, amount] + compute_selector("burn((Field),Field,Field)"), + [from, amount, nonce] ) } } diff --git a/yarn-project/noir-contracts/src/contracts/lending_contract/src/main.nr b/yarn-project/noir-contracts/src/contracts/lending_contract/src/main.nr index 3683aafe96d..c26af11e340 100644 --- a/yarn-project/noir-contracts/src/contracts/lending_contract/src/main.nr +++ b/yarn-project/noir-contracts/src/contracts/lending_contract/src/main.nr @@ -155,17 +155,17 @@ contract Lending { asset } - // This don't need to be on behalf of self. We should be able to repay on behalf of someone else. #[aztec(private)] fn deposit_private( + from: Field, + amount: Field, + nonce: Field, secret: Field, - asset_owner: Field, on_behalf_of: Field, - amount: Field, collateral_asset: Field, ) { let on_behalf_of = compute_identifier(secret, on_behalf_of, context.msg_sender()); - let _res = Token::at(collateral_asset).unshield(&mut context, asset_owner, context.this_address(), amount); + let _res = Token::at(collateral_asset).unshield(&mut context, from, context.this_address(), amount, nonce); // _deposit(on_behalf_of, amount, collateral_asset) let selector = compute_selector("_deposit(Field,Field,Field)"); let _callStackItem2 = context.call_public_function(context.this_address(), selector, [on_behalf_of, amount, collateral_asset]); @@ -173,13 +173,14 @@ contract Lending { #[aztec(public)] fn deposit_public( - owner: Field, amount: Field, + nonce: Field, + on_behalf_of: Field, collateral_asset: Field, ) -> Field { - Token::at(collateral_asset).transfer_from_pub(context, context.msg_sender(), context.this_address(), amount); + Token::at(collateral_asset).transfer_public(context, context.msg_sender(), context.this_address(), amount, nonce); let selector = compute_selector("_deposit(Field,Field,Field)"); - let return_values = context.call_public_function(context.this_address(), selector, [owner, amount, collateral_asset]); + let return_values = context.call_public_function(context.this_address(), selector, [on_behalf_of, amount, collateral_asset]); return_values[0] } @@ -253,7 +254,7 @@ contract Lending { // @todo @LHerskind Support both shielding and transfers (for now just transfer) let collateral_asset = storage.collateral_asset.read(); - Token::at(collateral_asset).transfer_pub(context, recipient, amount); + Token::at(collateral_asset).transfer_public(context, context.this_address(), recipient, amount, 0); 1 } @@ -304,33 +305,34 @@ contract Lending { // @todo @LHerskind Need to support both private and public minting. let stable_coin = storage.stable_coin.read(); - Token::at(stable_coin).owner_mint_pub(context, to, amount); + Token::at(stable_coin).mint_public(context, to, amount); 1 } #[aztec(private)] fn repay_private( + from: Field, + amount: Field, + nonce: Field, secret: Field, - asset_owner: Field, on_behalf_of: Field, - amount: Field, stable_coin: Field, ) { let on_behalf_of = compute_identifier(secret, on_behalf_of, context.msg_sender()); - let _res = Token::at(stable_coin).unshield(&mut context, asset_owner, context.this_address(), amount); + let _res = Token::at(stable_coin).burn(&mut context, from, amount, nonce); let selector = compute_selector("_repay(Field,Field,Field)"); let _callStackItem = context.call_public_function(context.this_address(), selector, [on_behalf_of, amount, stable_coin]); } #[aztec(public)] fn repay_public( - owner: Field, amount: Field, + nonce: Field, + owner: Field, stable_coin: Field, ) -> Field { - // Should probably just burn the tokens actually :thinking: - Token::at(stable_coin).transfer_from_pub(context, context.msg_sender(), context.this_address(), amount); + Token::at(stable_coin).burn_public(context, context.msg_sender(), amount, nonce); let selector = compute_selector("_repay(Field,Field,Field)"); let return_values = context.call_public_function(context.this_address(), selector, [owner, amount, stable_coin]);