feat!: make initialize_or_replace take an option

Follow up on #17017, this change allows for potentially cheaper circuits by not forcing the caller to compute the initial value (which might be expensive) unless needed. The multiple `Option` methods such as `unwrap_or`, `unwrap_or_else`, etc., make handling this simple.

E.g. a counter becomes
```noir
storage.counter.initialize_or_replace(|opt_current| UintNote::new(
  opt_current
    .map(|note| note.value)
    .unwrap_or(0) + 1
));
```

Author: nventuro

docs/docs/migration_notes.md

@@ -58,21 +58,28 @@ Updating a note used to require reading it first (via `get_note`, which nullifie
 **Key points:**
 
 1. `replace(self, new_note)` (old) → `replace(self, f)` (new), where `f` takes the current note and returns a transformed note.
-2. `initialize_or_replace(self, note)` (old) → `initialize_or_replace(self, init_note, f)` (new). Uninitialized variables use `init_note`, initialized ones pass the transform function to `replace`.
+2. `initialize_or_replace(self, note)` (old) → `initialize_or_replace(self, f)` (new), where `f` takes an `Option` with the current none, or `none` if uninitialized.
 3. Previous note is automatically nullified before the new note is inserted.
 4. `NoteEmission<Note>` still requires `.emit()` or `.discard()`.
 
 **Example Migration:**
 
 ```diff
+- let current_note = storage.my_var.get_note();
+- let new_note = f(current_note);
 - storage.my_var.replace(new_note);
-+ let x: Field = 5;
-+ storage.my_var.replace(|note| UintNote::new(note.value + x));
++ storage.my_var.replace(|current_note| f(current_note));
 ```
 
 ```diff
-- storage.my_var.initialize_or_replace(init_note);
-+ storage.my_var.initialize_or_replace(init_note, |note| UintNote::new(note.value + 5));
+- storage.my_var.initialize_or_replace(new_note);
++ storage.my_var.initialize_or_replace(|_| new_note);
+```
+
+This makes it easy and efficient to handle both initialization and current value mutation via `initialize_or_replace`, e.g. if implementing a note that simply counts how many times it has been read:
+
+```diff
++ storage.my_var.initialize_or_replace(|opt_current: Option<Note>| opt_current.unwrap_or(0 /* initial value */) + 1);
 ```
 
 

noir-projects/aztec-nr/aztec/src/state_vars/private_mutable.nr

@@ -253,8 +253,8 @@ where
     /// Reads the current note of a PrivateMutable state variable, nullifies it,
     /// and inserts a new note produced by a user-provided function.
     ///
-    /// This function implements a "read-and-replace" pattern for updating private state in Aztec.
-    /// It first retrieves the current note, then nullifies it (marking it as spent),
+    /// This function implements a "read-and-replace" pattern for updating private state
+    /// in Aztec. It first retrieves the current note, then nullifies it (marking it as spent),
     /// and finally inserts a `new_note` produced by the user-provided function `f`.
     ///
     /// This is conceptually similar to updating a variable in Ethereum smart contracts,
@@ -263,12 +263,13 @@ where
     ///
     /// This function can only be called after the PrivateMutable has been initialized.
     /// If called on an uninitialized PrivateMutable, it will fail because there is
-    /// no current note to replace.
+    /// no current note to replace. If you don't know if the state variable has been
+    /// initialized already, you can use `initialize_or_replace` to handle both cases.
     ///
     /// ## Arguments
     ///
-    /// * `f` - A function that takes the current `Note` and returns a new `Note`.
-    ///         This allows you to transform the current note before it is reinserted.
+    /// * `f` - A function that takes the current `Note` and returns a new `Note` that
+    ///         will replace it and become the "current value".
     ///
     /// ## Returns
     ///
@@ -322,8 +323,10 @@ where
     ///
     /// ## Arguments
     ///
-    /// * `init_note`    - The note to use if the PrivateMutable is uninitialized.
-    /// * `transform_fn` - A function that takes the current `Note` and returns a new `Note` to replace it.
+    /// * `f` - A function that takes an `Option` with the current `Note` and returns the `Note` to insert. This allows
+    ///         you to transform the current note before it is reinserted. The `Option` is `none` if the state variable
+    ///         was not initialized.
+    ///
     ///
     /// ## Returns
     ///
@@ -333,11 +336,7 @@ where
     ///                          the note, or `.discard()` to skip emission.
     ///                          See NoteEmission documentation for more details.
     ///
-    pub fn initialize_or_replace<Env>(
-        self,
-        init_note: Note,
-        transform_fn: fn[Env](Note) -> Note,
-    ) -> NoteEmission<Note>
+    pub fn initialize_or_replace<Env>(self, f: fn[Env](Option<Note>) -> Note) -> NoteEmission<Note>
     where
         Note: Packable,
     {
@@ -356,9 +355,9 @@ where
             unsafe { check_nullifier_exists(self.compute_initialization_nullifier()) };
 
         if !is_initialized {
-            self.initialize(init_note)
+            self.initialize(f(Option::none()))
         } else {
-            self.replace(transform_fn)
+            self.replace(|note| f(Option::some(note)))
         }
     }
 

noir-projects/aztec-nr/aztec/src/state_vars/private_mutable/test.nr

@@ -234,8 +234,9 @@ unconstrained fn initialize_or_replace_uninitialized() {
 
         let init_note = MockNote::new(VALUE).build_note();
 
-        let emission = state_var.initialize_or_replace(init_note, |_: MockNote| {
-            panic(f"") // This should not be called
+        let emission = state_var.initialize_or_replace(|opt_current_note| {
+            assert(opt_current_note.is_none());
+            init_note
         });
 
         assert_eq(emission.note, init_note);
@@ -303,8 +304,8 @@ unconstrained fn initialize_or_replace_initialized_settled() {
 
         let replacement_value = VALUE + 1;
         let replacement_note = MockNote::new(replacement_value).build_note();
-        let emission = state_var.initialize_or_replace(std::mem::zeroed(), |current_note| {
-            assert_eq(current_note, note);
+        let emission = state_var.initialize_or_replace(|opt_current_note| {
+            assert_eq(opt_current_note.unwrap(), note);
             replacement_note
         });
 

noir-projects/noir-contracts/contracts/app/app_subscription_contract/src/main.nr

@@ -154,11 +154,12 @@ pub contract AppSubscription {
             &mut context,
         );
 
-        let subscription_note = SubscriptionNote::new(subscriber, expiry_block_number, tx_count);
         storage
             .subscriptions
             .at(subscriber)
-            .initialize_or_replace(subscription_note, |_old| subscription_note)
+            .initialize_or_replace(|_| {
+                SubscriptionNote::new(subscriber, expiry_block_number, tx_count)
+            })
             .emit(&mut context, subscriber, MessageDelivery.CONSTRAINED_ONCHAIN);
     }