fix: txe partial note support (#11414)

For the note hashes to work from public, we need a deterministic way to
derive note hashes. This does so by using the txHash or first
non-revertible nullifier. This is okay because TXe we never get to the
revertible state.

noir-projects/noir-contracts/contracts/nft_contract/src/test/transfer_to_private.nr

@@ -48,21 +48,7 @@ unconstrained fn transfer_to_private_external_orchestration() {
         &mut env.public(),
     );
 
-    // TODO(#8771): We need to manually add the note because in the partial notes flow `notify_created_note_oracle`
-    // is not called and we don't have a `NoteProcessor` in TXE.
-    let private_nfts_recipient_slot =
-        derive_storage_slot_in_map(NFT::storage_layout().private_nfts.slot, recipient);
-
-    env.add_note(
-        &mut NFTNote {
-            token_id,
-            owner: recipient,
-            randomness: note_randomness,
-            header: NoteHeader::empty(),
-        },
-        private_nfts_recipient_slot,
-        nft_contract_address,
-    );
+    env.advance_block_by(1);
 
     // Recipient should have the note in their private nfts
     utils::assert_owns_private_nft(nft_contract_address, recipient, token_id);

noir-projects/noir-contracts/contracts/nft_contract/src/test/utils.nr

@@ -70,21 +70,7 @@ pub unconstrained fn setup_mint_and_transfer_to_private(
         &mut env.private(),
     );
 
-    // TODO(#8771): We need to manually add the note because in the partial notes flow `notify_created_note_oracle`
-    // is not called and we don't have a `NoteProcessor` in TXE.
-    let private_nfts_owner_slot =
-        derive_storage_slot_in_map(NFT::storage_layout().private_nfts.slot, owner);
-
-    env.add_note(
-        &mut NFTNote {
-            token_id: minted_token_id,
-            owner,
-            randomness: note_randomness,
-            header: NoteHeader::empty(),
-        },
-        private_nfts_owner_slot,
-        nft_contract_address,
-    );
+    env.advance_block_by(1);
 
     (env, nft_contract_address, owner, recipient, minted_token_id)
 }

noir-projects/noir-contracts/contracts/token_contract/src/test/refunds.nr

@@ -25,13 +25,7 @@ unconstrained fn setup_refund_success() {
     let user = owner;
     let fee_payer = recipient;
 
-    // We use the same randomness for both the fee payer, the user and the nonce as we currently don't have
-    // `OracleMock::mock_once()`
-    let fee_payer_randomness = 123;
-    let user_randomness = fee_payer_randomness;
-    let nonce = fee_payer_randomness;
-
-    let _ = OracleMock::mock("getRandomField").returns(fee_payer_randomness);
+    let nonce = 123;
 
     let setup_refund_from_call_interface =
         Token::at(token_contract_address).setup_refund(user, funded_amount, nonce);
@@ -46,27 +40,10 @@ unconstrained fn setup_refund_success() {
 
     setup_refund_from_call_interface.call(&mut env.private());
 
-    // When the refund was set up, we would've spent the note worth mint_amount, and inserted a note worth
-    //`mint_amount - funded_amount`. When completing the refund, we would've constructed a hash corresponding to a note
-    // worth  `funded_amount - transaction_fee`. We "know" the transaction fee was 1 (it is hardcoded in
-    // `executePublicFunction` TXE oracle) but we need to notify TXE of the note (preimage).
-    utils::add_token_note(
-        env,
-        token_contract_address,
-        fee_payer,
-        expected_tx_fee,
-        fee_payer_randomness,
-    );
-    utils::add_token_note(
-        env,
-        token_contract_address,
-        user,
-        funded_amount - expected_tx_fee,
-        user_randomness,
-    );
+    env.advance_block_by(1);
 
+    utils::check_public_balance(token_contract_address, fee_payer, expected_tx_fee);
     utils::check_private_balance(token_contract_address, user, mint_amount - expected_tx_fee);
-    utils::check_private_balance(token_contract_address, fee_payer, expected_tx_fee)
 }
 
 // This test should be reworked when final support for partial notes is in

noir-projects/noir-contracts/contracts/token_contract/src/test/transfer_to_private.nr

@@ -41,14 +41,7 @@ unconstrained fn transfer_to_private_external_orchestration() {
         &mut env.public(),
     );
 
-    // We need to manually add the note because #8771 has not yet been implemented
-    utils::add_token_note(
-        env,
-        token_contract_address,
-        recipient,
-        amount,
-        note_randomness,
-    );
+    env.advance_block_by(1);
 
     // Recipient's private balance should be equal to the amount
     utils::check_private_balance(token_contract_address, recipient, amount);

noir-projects/noir-contracts/contracts/token_contract/src/test/utils.nr

@@ -90,13 +90,7 @@ pub unconstrained fn mint_to_private(
         &mut env.private(),
     );
 
-    add_token_note(
-        env,
-        token_contract_address,
-        recipient,
-        amount,
-        note_randomness,
-    );
+    env.advance_block_by(1);
 }
 
 // docs:start:txe_test_read_public
@@ -117,6 +111,21 @@ pub unconstrained fn check_public_balance(
 }
 // docs:end:txe_test_read_public
 
+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 check_total_supply(
     token_contract_address: AztecAddress,
     expected_total_supply: U128,
@@ -146,8 +155,19 @@ pub unconstrained fn check_private_balance(
 }
 // docs:end:txe_test_call_unconstrained
 
-// TODO(#8771): We need to manually add the note because in the partial notes flow `notify_created_note_oracle`
-// is not called and we don't have a `NoteProcessor` in TXE.
+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
+}
+
+// This is used if we need to add a token note manually, in the case where the note is not emitted in logs.
 pub unconstrained fn add_token_note(
     env: &mut TestEnvironment,
     token_contract_address: AztecAddress,

yarn-project/simulator/src/public/fixtures/index.ts

@@ -96,7 +96,12 @@ export async function simulateAvmTestContractGenerateCircuitInputs(
     teardownExecutionRequest = new PublicExecutionRequest(callContext, fnArgs);
   }
 
-  const tx: Tx = createTxForPublicCalls(setupExecutionRequests, appExecutionRequests, teardownExecutionRequest);
+  const tx: Tx = createTxForPublicCalls(
+    setupExecutionRequests,
+    appExecutionRequests,
+    Fr.random(),
+    teardownExecutionRequest,
+  );
 
   const avmResult = await simulator.simulate(tx);
 
@@ -165,6 +170,7 @@ export async function simulateAvmTestContractCall(
 export function createTxForPublicCalls(
   setupExecutionRequests: PublicExecutionRequest[],
   appExecutionRequests: PublicExecutionRequest[],
+  firstNullifier: Fr,
   teardownExecutionRequest?: PublicExecutionRequest,
   gasUsedByPrivate: Gas = Gas.empty(),
 ): Tx {
@@ -179,7 +185,7 @@ export function createTxForPublicCalls(
 
   const forPublic = PartialPrivateTailPublicInputsForPublic.empty();
   // TODO(#9269): Remove this fake nullifier method as we move away from 1st nullifier as hash.
-  forPublic.nonRevertibleAccumulatedData.nullifiers[0] = Fr.random(); // fake tx nullifier
+  forPublic.nonRevertibleAccumulatedData.nullifiers[0] = firstNullifier; // fake tx nullifier
 
   // We reverse order because the simulator expects it to be like a "stack" of calls to pop from
   for (let i = setupCallRequests.length - 1; i >= 0; i--) {

yarn-project/txe/src/oracle/txe_oracle.ts

@@ -828,19 +828,30 @@ export class TXE implements TypedOracle {
       globalVariables,
     );
 
+    const { usedTxRequestHashForNonces } = this.noteCache.finish();
+    const firstNullifier = usedTxRequestHashForNonces ? this.getTxRequestHash() : this.noteCache.getAllNullifiers()[0];
+
     // When setting up a teardown call, we tell it that
     // private execution used Gas(1, 1) so it can compute a tx fee.
     const gasUsedByPrivate = isTeardown ? new Gas(1, 1) : Gas.empty();
     const tx = createTxForPublicCalls(
       /*setupExecutionRequests=*/ [],
       /*appExecutionRequests=*/ isTeardown ? [] : [executionRequest],
+      firstNullifier,
       /*teardownExecutionRequests=*/ isTeardown ? executionRequest : undefined,
       gasUsedByPrivate,
     );
 
     const result = await simulator.simulate(tx);
+    const noteHashes = result.avmProvingRequest.inputs.output.accumulatedData.noteHashes.filter(s => !s.isEmpty());
+
+    await this.addUniqueNoteHashesFromPublic(noteHashes);
 
-    this.addPublicLogs(result.avmProvingRequest.inputs.publicInputs.publicLogs);
+    this.addPublicLogs(
+      result.avmProvingRequest.inputs.output.accumulatedData.publicLogs.filter(
+        log => !log.contractAddress.equals(AztecAddress.ZERO),
+      ),
+    );
 
     return Promise.resolve(result);
   }
@@ -892,7 +903,11 @@ export class TXE implements TypedOracle {
     const sideEffects = executionResult.avmProvingRequest.inputs.output.accumulatedData;
     const publicDataWrites = sideEffects.publicDataWrites.filter(s => !s.isEmpty());
     const noteHashes = sideEffects.noteHashes.filter(s => !s.isEmpty());
-    const nullifiers = sideEffects.nullifiers.filter(s => !s.isEmpty());
+
+    const { usedTxRequestHashForNonces } = this.noteCache.finish();
+    const firstNullifier = usedTxRequestHashForNonces ? this.getTxRequestHash() : this.noteCache.getAllNullifiers()[0];
+    const nullifiers = sideEffects.nullifiers.filter(s => !s.isEmpty()).filter(s => !s.equals(firstNullifier));
+
     await this.addPublicDataWrites(publicDataWrites);
     await this.addUniqueNoteHashesFromPublic(noteHashes);
     await this.addSiloedNullifiers(nullifiers);
@@ -1006,7 +1021,11 @@ export class TXE implements TypedOracle {
       const sideEffects = executionResult.avmProvingRequest.inputs.output.accumulatedData;
       const publicDataWrites = sideEffects.publicDataWrites.filter(s => !s.isEmpty());
       const noteHashes = sideEffects.noteHashes.filter(s => !s.isEmpty());
-      const nullifiers = sideEffects.nullifiers.filter(s => !s.isEmpty());
+      const { usedTxRequestHashForNonces } = this.noteCache.finish();
+      const firstNullifier = usedTxRequestHashForNonces
+        ? this.getTxRequestHash()
+        : this.noteCache.getAllNullifiers()[0];
+      const nullifiers = sideEffects.nullifiers.filter(s => !s.isEmpty()).filter(s => !s.equals(firstNullifier));
       await this.addPublicDataWrites(publicDataWrites);
       await this.addUniqueNoteHashes(noteHashes);
       await this.addSiloedNullifiers(nullifiers);