fix: contracts with public keys txe (#11910)

Please read [contributing guidelines](CONTRIBUTING.md) and remove this
line.

noir-projects/noir-contracts/contracts/escrow_contract/src/main.nr

@@ -1,6 +1,8 @@
 // Sample escrow contract that stores a balance of a private token on behalf of an owner.
 use dep::aztec::macros::aztec;
 
+mod test;
+
 #[aztec]
 pub contract Escrow {
     use dep::aztec::{

noir-projects/noir-contracts/contracts/escrow_contract/src/test.nr

@@ -0,0 +1,120 @@
+// TODO(ek): Clean up this test
+use crate::Escrow;
+use dep::token::Token;
+
+use aztec::{
+    oracle::{execution::{get_block_number, get_contract_address}, storage::storage_read},
+    prelude::AztecAddress,
+    protocol_types::storage::map::derive_storage_slot_in_map,
+    test::helpers::{cheatcodes, test_environment::TestEnvironment},
+};
+
+pub unconstrained fn get_public_balance(
+    token_contract_address: AztecAddress,
+    address: AztecAddress,
+) -> U128 {
+    let current_contract_address = get_contract_address();
+    cheatcodes::set_contract_address(token_contract_address);
+    let block_number = get_block_number();
+
+    let balances_slot = Token::storage_layout().public_balances.slot;
+    let address_slot = derive_storage_slot_in_map(balances_slot, address);
+    let amount: U128 = storage_read(token_contract_address, address_slot, block_number);
+    cheatcodes::set_contract_address(current_contract_address);
+    amount
+}
+
+pub unconstrained fn get_private_balance(
+    token_contract_address: AztecAddress,
+    address: AztecAddress,
+) -> U128 {
+    let current_contract_address = get_contract_address();
+    cheatcodes::set_contract_address(token_contract_address);
+    // Direct call to unconstrained
+    let amt = Token::balance_of_private(address);
+    cheatcodes::set_contract_address(current_contract_address);
+    amt
+}
+
+global MINT_AMOUNT: U128 = U128::from_integer(200000);
+
+unconstrained fn deploy_contracts(
+    env: &mut TestEnvironment,
+    admin_and_owner: AztecAddress,
+) -> (AztecAddress, AztecAddress) {
+    env.impersonate(admin_and_owner);
+
+    // Deploy token contract
+    let donation_token_initializer_call_interface = Token::interface().constructor(
+        admin_and_owner,
+        "Token00000000000000000000000000",
+        "TKN0000000000000000000000000000",
+        18,
+    );
+    let donation_token_contract = env
+        .deploy("./@token_contract", "Token")
+        .with_public_void_initializer(donation_token_initializer_call_interface);
+    let token_contract_address = donation_token_contract.to_address();
+    env.advance_block_by(1);
+
+    // Deploy Escrow contract with public keys
+    let escrow_contract_initializer_call_interface =
+        Escrow::interface().constructor(admin_and_owner);
+    let escrow_contract = env
+        .deploy_with_public_keys("./@escrow_contract", "Escrow", 6969)
+        .with_private_initializer(escrow_contract_initializer_call_interface);
+    let escrow_contract_address = escrow_contract.to_address();
+
+    env.advance_block_by(1);
+
+    Token::at(token_contract_address)
+        .mint_to_private(admin_and_owner, admin_and_owner, MINT_AMOUNT)
+        .call(&mut env.private());
+
+    env.advance_block_by(1);
+
+    let private_balance_after_mint = get_private_balance(token_contract_address, admin_and_owner);
+    assert(private_balance_after_mint == MINT_AMOUNT);
+
+    (token_contract_address, escrow_contract_address)
+}
+
+#[test]
+unconstrained fn main() {
+    let mut env = TestEnvironment::new();
+
+    let (account_1, account_2) = (env.create_account_contract(1), env.create_account_contract(2));
+
+    let (token_contract_address, escrow_contract_address) = deploy_contracts(&mut env, account_1);
+
+    // We transfer tokens to the escrow contract
+    let TRANSFER_AMOUNT = U128::from_integer(20000);
+    Token::at(token_contract_address).transfer(escrow_contract_address, TRANSFER_AMOUNT).call(
+        &mut env.private(),
+    );
+    env.advance_block_by(1);
+
+    let balance_of_escrow_after_transfer =
+        get_private_balance(token_contract_address, escrow_contract_address);
+
+    assert_eq(balance_of_escrow_after_transfer, TRANSFER_AMOUNT);
+
+    // We then withdraw some escrowed funds to account_2
+    let balance_of_account_2_before_withdrawal =
+        get_private_balance(token_contract_address, account_2);
+    assert(balance_of_account_2_before_withdrawal == U128::zero());
+
+    let WITHDRAWAL_AMOUNT = U128::from_integer(69);
+    Escrow::at(escrow_contract_address)
+        .withdraw(token_contract_address, WITHDRAWAL_AMOUNT, account_2)
+        .call(&mut env.private());
+    env.advance_block_by(1);
+
+    let balance_of_account_2_after_withdrawal =
+        get_private_balance(token_contract_address, account_2);
+    assert(balance_of_account_2_after_withdrawal == WITHDRAWAL_AMOUNT);
+
+    let balance_of_escrow_after_withdrawal =
+        get_private_balance(token_contract_address, escrow_contract_address);
+    assert(balance_of_escrow_after_withdrawal == TRANSFER_AMOUNT - WITHDRAWAL_AMOUNT);
+}

yarn-project/txe/src/txe_service/txe_service.ts

@@ -95,12 +95,13 @@ export class TXEService {
     ]);
 
     if (!fromSingle(secret).equals(Fr.ZERO)) {
-      await this.createAccount(secret);
+      await this.addAccount(artifact, instance, secret);
+    } else {
+      await (this.typedOracle as TXE).addContractInstance(instance);
+      await (this.typedOracle as TXE).addContractArtifact(instance.contractClassId, artifact);
+      this.logger.debug(`Deployed ${artifact.name} at ${instance.address}`);
     }
 
-    this.logger.debug(`Deployed ${artifact.name} at ${instance.address}`);
-    await (this.typedOracle as TXE).addContractInstance(instance);
-    await (this.typedOracle as TXE).addContractArtifact(instance.contractClassId, artifact);
     return toForeignCallResult([
       toArray([
         instance.salt,
@@ -137,6 +138,7 @@ export class TXEService {
   async createAccount(secret: ForeignCallSingle) {
     const keyStore = (this.typedOracle as TXE).getKeyStore();
     const secretFr = fromSingle(secret);
+    // This is a footgun !
     const completeAddress = await keyStore.addAccount(secretFr, secretFr);
     const accountStore = (this.typedOracle as TXE).getTXEDatabase();
     await accountStore.setAccount(completeAddress.address, completeAddress);