Replace assert_contains with assert_contains_and_remove.

yarn-project/noir-contracts/src/contracts/escrow_contract/src/address_note.nr

@@ -1,3 +1,5 @@
+mod filter;
+
 use dep::std::hash::pedersen;
 use dep::aztec::note::note_interface::NoteInterface;
 use dep::aztec::note::note_header::NoteHeader;

yarn-project/noir-contracts/src/contracts/escrow_contract/src/address_note/filter.nr

@@ -0,0 +1,22 @@
+use dep::aztec::constants_gen::MAX_READ_REQUESTS_PER_CALL;
+use dep::aztec::types::option::Option;
+
+use crate::address_note::AddressNote;
+
+fn filter_by_owner_and_address(
+    notes: [Option<AddressNote>; MAX_READ_REQUESTS_PER_CALL],
+    params: [Field; 2],
+) -> [Option<AddressNote>; 1] {
+    let address = params[0];
+    let owner = params[1];
+    let mut owner_note = [Option::none(AddressNote::dummy())];
+    for i in 0..notes.len() {
+        if notes[i].is_some() {
+            let note = notes[i].unwrap_unchecked();
+            if (note.address == address) & (note.owner == owner) {
+                owner_note[0] = notes[i];
+            }
+        }
+    }
+    owner_note
+}

yarn-project/noir-contracts/src/contracts/escrow_contract/src/main.nr

@@ -13,14 +13,16 @@ contract Escrow {
     use dep::aztec::oracle::get_public_key::get_public_key;
 
     use dep::aztec::note::{
+        note_getter_options::NoteGetterOptions,
         note_header::{NoteHeader},
         utils as note_utils,
     };
 
     use crate::address_note::{
         AddressNote,
         AddressNoteMethods,
-        ADDRESS_NOTE_LEN
+        ADDRESS_NOTE_LEN,
+        filter::filter_by_owner_and_address,
     };
 
     use crate::storage::Storage;
@@ -59,10 +61,14 @@ contract Escrow {
         let sender = context.msg_sender();
         let storage = Storage::init();
 
-        // TODO: Do we need to manually nullify and recreate this note for access control? Or does Set handle it for us? Or since we have no methods for updating it, we're good?
-        let mut note = AddressNote::new(sender, this);
-        storage.owners.assert_contains(&mut context, &mut note);
-
+        // We don't remove note from the owners set. If a note exists, the owner and recipient are legit.
+        let options = NoteGetterOptions::with_filter(filter_by_owner_and_address, [sender, this]);
+        let notes = storage.owners.get_notes(&mut context, options);
+        let note = notes[0].unwrap_unchecked();
+        // Filter is not constrained. We still need to check if the note is what we expected.
+        assert(note.address == sender);
+        assert(note.owner == this);
+
         // TODO: Can we dynamically get this selector?
         let transfer_selector = 0xdcd4c318;
         let _callStackItem = context.call_private_function(token, transfer_selector, [amount, this, recipient]);

yarn-project/noir-contracts/src/contracts/non_native_token_contract/src/main.nr

@@ -287,23 +287,10 @@ contract NonNativeToken {
             amount, secret, owner
         ]));
 
-        // Assert that the note exists within the tree
-        let mut public_note = TransparentNote::new_from_secret(amount, secret);
-
-        // Ensure that the note exists in the tree
-        pending_shields.assert_contains_note_hash(&mut context, &mut public_note);
-        // The above call to `assert_contains()` also returns a modified note with
-        // the header which is necessary for the next step (remove).
-
-        // Set the nonce to nonzero so that the nullifier is treated as persistable
-        // (non-transient) and so the private kernel does not attempt to match it to
-        // a pending noteHash/commitment and squash them.
-        // TODO(https://github.com/AztecProtocol/aztec-packages/issues/1386): remove
-        // this hack once public kernel injects nonces.
-        public_note.header.nonce = 1;
-
-        // Consume note in secret!
-        pending_shields.remove(&mut context, public_note);
+        let public_note = TransparentNote::new_from_secret(amount, secret);
+
+        // Ensure that the note exists in the tree and remove it.
+        pending_shields.assert_contains_and_remove(&mut context, public_note);
 
         // Mint the tokens
         let balance = storage.balances.at(owner);

yarn-project/noir-contracts/src/contracts/private_token_airdrop_contract/src/main.nr

@@ -136,9 +136,8 @@ contract PrivateTokenAirdrop {
         let mut context = PrivateContext::new(inputs, abi::hash_args([amount, secret, owner]));
 
         // Remove the claim note if it exists in the set.
-        let mut note = ClaimNote::new(amount, secret);
-        storage.claims.assert_contains(&mut context, &mut note);
-        storage.claims.remove(&mut context, note);
+        let note = ClaimNote::new(amount, secret);
+        storage.claims.assert_contains_note_and_remove(&mut context, note);
 
         // Send the value note.
         let balance = storage.balances.at(owner);

yarn-project/noir-libs/noir-aztec/src/state_vars/immutable_singleton.nr

@@ -1,10 +1,7 @@
 use dep::std::hash::pedersen_with_separator;
 use crate::context::PrivateContext;
 use crate::note::lifecycle::create_note;
-use crate::note::note_getter::{
-    get_note,
-    ensure_note_exists,
-};
+use crate::note::note_getter::get_note;
 use crate::note::note_interface::NoteInterface;
 use crate::oracle;
 use crate::constants_gen::{
@@ -43,8 +40,4 @@ impl<Note, N> ImmutableSingleton<Note, N> {
         let storage_slot = self.storage_slot;
         get_note(context, storage_slot, self.note_interface)
     }
-
-    fn assert_contains(self, context: &mut PrivateContext, note: &mut Note) {
-        ensure_note_exists(context, self.storage_slot, self.note_interface, note);
-    }
 }

yarn-project/noir-libs/noir-aztec/src/state_vars/set.nr

@@ -30,15 +30,34 @@ impl<Note, N> Set<Note, N> {
         create_note_hash_from_public(context, self.storage_slot, note, self.note_interface);
     }
 
-    fn assert_contains(self, context: &mut PrivateContext, note: &mut Note) {
-        ensure_note_exists(context, self.storage_slot, self.note_interface, note);
+    // TODO(#1386)
+    // Should be replaced by `assert_contains_and_remove`.
+    fn assert_contains_note_and_remove(self, context: &mut PrivateContext, note: Note) {
+        let mut note_with_header = note;
+        ensure_note_exists(context, self.storage_slot, self.note_interface, &mut note_with_header);
+        self.remove(context, note_with_header);
     }
 
-    // TODO(https://github.com/AztecProtocol/aztec-packages/issues/1386): rename to
-    // `assert_contains` and replace function above ^ once public kernel injects
+    // TODO(https://github.com/AztecProtocol/aztec-packages/issues/1386):
+    // replace function above ^ once public kernel injects
     // nonces to note hashes.
-    fn assert_contains_note_hash(self, context: &mut PrivateContext, note: &mut Note) {
-        ensure_note_hash_exists(context, self.storage_slot, self.note_interface, note)
+    fn assert_contains_and_remove(self, context: &mut PrivateContext, note: Note) {
+        let mut note_with_header = note;
+        // Modifies note with the header which is necessary for the next step (remove).
+        ensure_note_hash_exists(context, self.storage_slot, self.note_interface, &mut note_with_header);
+
+        let get_header = self.note_interface.get_header;
+        let set_header = self.note_interface.set_header;
+        let mut header = get_header(note);
+        // Set the nonce to nonzero so that the nullifier is treated as persistable
+        // (non-transient) and so the private kernel does not attempt to match it to
+        // a pending noteHash/commitment and squash them.
+        // TODO(https://github.com/AztecProtocol/aztec-packages/issues/1386): remove
+        // this hack once public kernel injects nonces.
+        header.nonce = 1;
+        set_header(&mut note_with_header, header);
+
+        self.remove(context, note_with_header);
     }
 
     fn remove(self, context: &mut PrivateContext, note: Note) {