feat: introduce initialize_or_replace (#6519)

This fixes the issue described in [these
comments](AztecProtocol/aztec-packages#6442 (comment)),
which highlighted that `PrivateMutable` was providing a poor API as it
forced users to call unconstrained functions and leave correctness up to
them. This is a simple replacement to make #6442 cleaner and help guide
future development, but it should only be considered a patch - we'll
eventually do a more thoughtful complete pass over the entire API to
bring coherency to it.

I tried adding tests for this (I wanted to test that regardless what the
return value of the nullifier check oracle was, either an inclusion
proof or a new nullifier would be included in the context), but this
proved to be very hard because notes are sent to the oracle as raw
fields in a custom packed format which `oracle::get_notes` decodes on
the spot. We'd need to have that be a separate library, but that exceeds
the scope of this PR.

aztec/src/state_vars/private_mutable.nr

@@ -18,6 +18,8 @@ struct PrivateMutable<Note, Context> {
 }
 // docs:end:struct
 
+mod test;
+
 impl<T, Context> Storage<T> for PrivateMutable<T, Context> {}
 
 impl<Note, Context> PrivateMutable<Note, Context> {
@@ -28,10 +30,6 @@ impl<Note, Context> PrivateMutable<Note, Context> {
     }
     // docs:end:new
 
-    pub fn to_unconstrained(self) -> PrivateMutable<Note, ()> {
-        PrivateMutable { context: (), storage_slot: self.storage_slot }
-    }
-
     // The following computation is leaky, in that it doesn't hide the storage slot that has been initialized, nor does it hide the contract address of this contract.
     // When this initialization nullifier is emitted, an observer could do a dictionary or rainbow attack to learn the preimage of this nullifier to deduce the storage slot and contract address.
     // For some applications, leaking the details that a particular state variable of a particular contract has been initialized will be unacceptable.
@@ -81,6 +79,31 @@ impl<Note> PrivateMutable<Note, &mut PrivateContext> {
     }
     // docs:end:replace
 
+    pub fn initialize_or_replace<N, M>(
+        self,
+        note: &mut Note,
+        broadcast: bool,
+        ivpk_m: GrumpkinPoint
+    ) where Note: NoteInterface<N, M> {
+        let is_initialized = check_nullifier_exists(self.compute_initialization_nullifier());
+
+        // check_nullifier_exists() is an unconstrained function - we can constrain a true value by providing an 
+        // inclusion proof of the nullifier, but cannot constrain a false value since a non-inclusion proof would only 
+        // be valid if done in public.
+        // Ultimately, this is not an issue ginen that we'll either:
+        //  - initialize the state variable, which would fail if it was already initialized due to the duplicate 
+        //    nullifier, or
+        //  - replace the current value, which would fail if it was not initialized since we wouldn't be able to produce
+        //    an inclusion proof for the current note
+        // This means that an honest oracle will assist the prover to produce a valid proof, while a malicious oracle
+        // (i.e. one that returns an incorrect value for is_initialized) will simply fail to produce a proof.
+        if (!is_initialized) {
+            self.initialize(note, broadcast, ivpk_m);
+        } else {
+            self.replace(note, broadcast, ivpk_m)
+        }
+    }
+
     // docs:start:get_note
     pub fn get_note<N, M>(
         self,

aztec/src/state_vars/private_mutable/test.nr

@@ -0,0 +1,53 @@
+use dep::protocol_types::{address::AztecAddress, grumpkin_point::GrumpkinPoint};
+use crate::{context::PrivateContext, state_vars::private_mutable::PrivateMutable};
+use crate::test::{mocks::mock_note::MockNote, helpers::context_builder::ContextBuilder};
+use dep::std::{unsafe::zeroed, test::OracleMock};
+
+global contract_address = AztecAddress::from_field(13);
+global storage_slot = 17;
+
+fn setup() -> PrivateMutable<MockNote, &mut PrivateContext> {
+    let mut context = ContextBuilder::new().contract_address(contract_address).private();
+    let state_var = PrivateMutable::new(&mut context, storage_slot);
+
+    // This oracle is called for its side effects alone - it's always expected to return 0.
+    OracleMock::mock("notifyCreatedNote").returns(0);
+
+    state_var
+}
+
+#[test]
+fn test_initialize_or_replace_without_nullifier() {
+    let state_var = setup();
+
+    let ivpk_m: GrumpkinPoint = zeroed();
+    let broadcast = false;
+
+    let value = 42;
+    let mut note = MockNote::new(value).contract_address(contract_address).storage_slot(storage_slot).build();
+
+    OracleMock::mock("checkNullifierExists").returns(0);
+    state_var.initialize_or_replace(&mut note, broadcast, ivpk_m);
+
+    // Since we reported there was no nullifier, we should initialize and see the following side-effects:
+    // - a new note being created
+    // - no notes being read
+    // - the initialization nullifier being emitted
+    assert_eq(state_var.context.new_note_hashes.len(), 1);
+    assert_eq(state_var.context.note_hash_read_requests.len(), 0);
+    assert_eq(state_var.context.new_nullifiers.len(), 1);
+
+    // Note that if the oracle was wrong and the initialization nullifier did exist, this attempt to write it again
+    // would cause the sequencer to revert this transaction - we are therefore safe from bad oracles.
+    let nullifier = state_var.context.new_nullifiers.get(0);
+    assert_eq(nullifier.value, state_var.compute_initialization_nullifier());
+    assert_eq(nullifier.note_hash, 0);
+}
+
+#[test]
+fn test_initialize_or_replace_with_nullifier() {
+    // Here we'd want to test a scenario like the one above with the oracle indicating that the initialization
+    // nullifier does exist. Unfortunately that requires us to produce a valid oracle return value for getNotes,
+    // which is fairly involved as it deals with serialization of notes, and is relatively complicated to replicate
+    // in Noir.
+}