feat!: update private_mutable replace functions to closure style (#7065)

# Motivation
Previously, updating a note required reading it first via `get_note` (which nullified and recreated it) and then calling `replace` — effectively proving the note twice.  Now, `replace` accepts a callback that transforms the current note directly, and `initialize_or_replace` uses this updated `replace` internally.  This reduces circuit cost while maintaining exactly one current note.

# 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`.
3. The previous note is automatically nullified before the new note is inserted.
4. `NoteEmission<Note>` still requires `.emit()` or `.discard()`.
5. Example contracts and docs have been updated to reflect this new usage pattern.

# Breaking change
Code relying on the old `get_note() + replace()` pattern may need to be updated.

# Addresses
Closes #7065

Author: wildjos

boxes/boxes/react/src/contracts/src/main.nr

@@ -31,15 +31,19 @@ contract BoxReact {
     #[private]
     fn setNumber(number: Field, owner: AztecAddress) {
         let numbers = storage.numbers;
-        let new_number = ValueNote::new(number, owner);
 
-        numbers.at(owner).replace(new_number).emit(
-            &mut context,
-            owner,
-            MessageDelivery.CONSTRAINED_ONCHAIN,
-        );
+        numbers.at(owner)
+            .replace(|_old| {
+                ValueNote::new(number, owner)
+            }) 
+            .emit(
+                &mut context,
+                owner,
+                MessageDelivery.CONSTRAINED_ONCHAIN,
+            );
     }
 
+
     #[utility]
     unconstrained fn getNumber(owner: AztecAddress) -> ValueNote {
         let numbers = storage.numbers;

boxes/boxes/vite/src/contracts/src/main.nr

@@ -31,9 +31,12 @@ contract BoxReact {
     #[private]
     fn setNumber(number: Field, owner: AztecAddress) {
         let numbers = storage.numbers;
-        let mut new_number = ValueNote::new(number, owner);
 
-        numbers.at(owner).replace(new_number).emit(
+
+    numbers.at(owner)
+        .replace(|_old| {
+            ValueNote::new(number, owner)
+        }).emit(
             &mut context,
             owner,
             MessageDelivery.CONSTRAINED_ONCHAIN,

docs/docs/developers/guides/smart_contracts/storage_types.md

@@ -105,9 +105,14 @@ An unconstrained method to check whether the PrivateMutable has been initialized
 
 #### `replace`
 
-To update the value of a `PrivateMutable`, we can use the `replace` method. The method takes a new note as input and replaces the current note with the new one. It emits a nullifier for the old value, and inserts the new note into the data tree.
+To update the value of a `PrivateMutable`, we can use the `replace` method. The method takes a function (or closure) that transforms the current note into a new one.
 
-An example of this is seen in a example card game, where we create a new note (a `CardNote`) containing some new data, and replace the current note with it:
+When called, the method will:
+- Nullify the old note
+- Apply the transform function to produce a new note
+- Insert the new note into the data tree
+
+An example of this is seen in an example card game, where an update function is passed in to transform the current note into a new one (in this example, updating a `CardNote` data):
 
 #include_code state_vars-PrivateMutableReplace /noir-projects/noir-contracts/contracts/docs/docs_example_contract/src/main.nr rust
 

docs/docs/migration_notes.md

@@ -45,6 +45,36 @@ As we prepare for a bigger `Wallet` interface refactor and the upcoming `WalletS
 
 ## [Aztec.nr]
 
+
+### PrivateMutable: replace / initialize_or_replace behaviour change
+
+**Motivation:**
+Updating a note used to require reading it first (via `get_note`, which nullifies and recreates it) and then calling `replace` — effectively proving a note twice. Now, `replace` accepts a callback that transforms the current note directly, and `initialize_or_replace` simply uses this updated `replace` internally. This reduces circuit cost while maintaining exactly one current note.
+
+**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`.
+3. Previous note is automatically nullified before the new note is inserted.
+4. `NoteEmission<Note>` still requires `.emit()` or `.discard()`.
+
+**Example Migration:**
+
+```diff
+- storage.my_var.replace(new_note);
++ let x: Field = 5;
++ storage.my_var.replace(|note| UintNote::new(note.value + x));
+```
+
+```diff
+- storage.my_var.initialize_or_replace(init_note);
++ storage.my_var.initialize_or_replace(init_note, |note| UintNote::new(note.value + 5));
+```
+
+
+- The callback can be a closure (inline) or a named function.
+- Any previous assumptions that replace simply inserts a new_note directly must be updated.
+
 ### Unified oracles into single get_utility_context oracle
 
 The following oracles:

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

@@ -251,11 +251,12 @@ where
     }
     // docs:end:initialize
 
-    /// Replaces the current note of a PrivateMutable state variable with a new note.
+    /// 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 the typical "nullify-and-create" pattern for updating
-    /// private state in Aztec. It first retrieves the current note, nullifies it
-    /// (marking it as "spent"), and then inserts a `new_note` with the updated data.
+    /// 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,
     /// except that in Aztec we achieve this by consuming the old note and creating a
@@ -267,8 +268,8 @@ where
     ///
     /// ## Arguments
     ///
-    /// * `new_note` - The new note that will replace the current note. This becomes
-    ///                the new "current value" of the PrivateMutable.
+    /// * `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.
     ///
     /// ## Returns
     ///
@@ -284,15 +285,16 @@ where
     /// - Retrieves the current note from the PXE via an oracle call
     /// - Validates that the current note exists and belongs to this storage slot
     /// - Computes the nullifier for the current note and pushes it to the context
-    /// - Inserts the provided `new_note` into the Note Hash Tree
+    /// - Calls the user-provided function `f` to produce a new note
+    /// - Inserts the resulting `new_note` into the Note Hash Tree using 'create_note'
     /// - Returns a NoteEmission type for the `new_note`, that allows the caller to
     ///   decide how to encrypt and deliver this note to its intended recipient.
     ///
     /// The nullification of the previous note ensures that it cannot be used again,
     /// maintaining the invariant that a PrivateMutable has exactly one current note.
     ///
     // docs:start:replace
-    pub fn replace(self, new_note: Note) -> NoteEmission<Note>
+    pub fn replace<Env>(self, f: fn[Env](Note) -> Note) -> NoteEmission<Note>
     where
         Note: Packable,
     {
@@ -306,23 +308,23 @@ where
             note_hash_for_read_request,
         );
 
+        let new_note = f(prev_retrieved_note.note);
+
         // Add replacement note.
         create_note(self.context, self.storage_slot, new_note)
     }
     // docs:end:replace
 
-    /// Initializes the PrivateMutable if it's uninitialized, or replaces the current
-    /// note if it's already initialized.
+    /// Initializes the PrivateMutable if it's uninitialized, or replaces the current note
+    /// using a transform function.
     ///
-    /// This is a convenience function that automatically chooses between `initialize`
-    /// and `replace` based on whether the PrivateMutable has been previously initialized.
-    /// This is useful when you don't know the initialization state beforehand, such
-    /// as in functions that may be called multiple times.
+    /// If uninitialized, `init_note` is used to initialize. If already initialized, the `transform_fn`
+    /// is passed to `replace`, which retrieves the current note, nullifies it, and inserts the transformed note.
     ///
     /// ## Arguments
     ///
-    /// * `note` - The note to store. If initializing, this becomes the first note.
-    ///            If replacing, this becomes the new current note.
+    /// * `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.
     ///
     /// ## Returns
     ///
@@ -332,7 +334,11 @@ where
     ///                          the note, or `.discard()` to skip emission.
     ///                          See NoteEmission documentation for more details.
     ///
-    pub fn initialize_or_replace(self, note: Note) -> NoteEmission<Note>
+    pub fn initialize_or_replace<Env>(
+        self,
+        init_note: Note,
+        transform_fn: fn[Env](Note) -> Note,
+    ) -> NoteEmission<Note>
     where
         Note: Packable,
     {
@@ -350,10 +356,10 @@ where
         let is_initialized =
             unsafe { check_nullifier_exists(self.compute_initialization_nullifier()) };
 
-        if (!is_initialized) {
-            self.initialize(note)
+        if !is_initialized {
+            self.initialize(init_note)
         } else {
-            self.replace(note)
+            self.replace(transform_fn)
         }
     }
 

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

@@ -27,8 +27,63 @@ unconstrained fn replace_uninitialized() {
     env.private_context(|context| {
         let state_var = in_private(context);
 
-        let note = MockNote::new(VALUE).build_note();
-        let _ = state_var.replace(note);
+        let _ = state_var.replace(|_old_note| {
+            let note = MockNote::new(VALUE).build_note();
+            note
+        });
+    });
+}
+
+// Named function to use as a callback for replace
+fn plus_one(note: MockNote) -> MockNote {
+    MockNote::new(note.value + 1).build_note()
+}
+
+#[test]
+unconstrained fn test_replace_plus_one() {
+    let env = TestEnvironment::new();
+
+    env.private_context(|context| {
+        let mut state_var = in_private(context);
+
+        // Initialize with a known value
+        let INIT_VALUE: Field = 7;
+        let EXPECTED_VALUE: Field = 8;
+        let initial_note = MockNote::new(INIT_VALUE).build_note();
+        let _ = state_var.initialize(initial_note);
+
+        // run read_and_replace with helper function
+        let emission = state_var.replace(plus_one);
+
+        // emission should contain new note with value 8
+        let expected_note = MockNote::new(EXPECTED_VALUE).build_note();
+        assert_eq(emission.note, expected_note);
+    });
+}
+
+#[test]
+unconstrained fn test_replace_capture_variable() {
+    let env = TestEnvironment::new();
+
+    env.private_context(|context| {
+        let mut state_var = in_private(context);
+
+        // Initialize with a known value
+        let INIT_VALUE: Field = 10;
+        let initial_note = MockNote::new(INIT_VALUE).build_note();
+        let _ = state_var.initialize(initial_note);
+
+        // Local variable to capture
+        let x: Field = 5;
+
+        // Use a closure to increment the note by x
+        let emission = state_var.replace(|note| MockNote::new(note.value + x).build_note());
+
+        // Expected value is initial + x
+        let expected_value: Field = INIT_VALUE + x;
+        let expected_note = MockNote::new(expected_value).build_note();
+
+        assert_eq(emission.note, expected_note);
     });
 }
 
@@ -99,10 +154,12 @@ unconstrained fn initialize_and_replace_pending() {
         let note_hashes_pre_replace = context.note_hashes.len();
 
         let replacement_value = VALUE + 1;
-        let replacement_note = MockNote::new(replacement_value).build_note();
-        let emission = state_var.replace(replacement_note);
 
-        assert_eq(emission.note, replacement_note);
+        let emission =
+            state_var.replace(|_old_note| MockNote::new(replacement_value).build_note());
+
+        let expected_note = MockNote::new(replacement_value).build_note();
+        assert_eq(emission.note, expected_note);
         assert_eq(emission.storage_slot, STORAGE_SLOT);
 
         // Replacing a PrivateMutable results in:
@@ -122,10 +179,13 @@ unconstrained fn initialize_or_replace_uninitialized() {
     env.private_context(|context| {
         let state_var = in_private(context);
 
-        let note = MockNote::new(VALUE).build_note();
-        let emission = state_var.initialize_or_replace(note);
+        let init_note = MockNote::new(VALUE).build_note();
 
-        assert_eq(emission.note, note);
+        let emission = state_var.initialize_or_replace(init_note, |_old_note: MockNote| {
+            panic(f"Unexpected call to replacement closure") // This should not be called
+        });
+
+        assert_eq(emission.note, init_note);
         assert_eq(emission.storage_slot, STORAGE_SLOT);
 
         // During initialization we both create the new note and emit the initialization nullifier. This will only
@@ -144,27 +204,23 @@ unconstrained fn initialize_or_replace_uninitialized() {
 //     env.private_context(|context| {
 //         let state_var = in_private(context);
 
-//         let note = MockNote::new(VALUE).build_note();
-
-//         let _ = state_var.initialize(note);
+//         // Initialize with a known value
+//         let init_note = MockNote::new(VALUE).build_note();
+//         let _ = state_var.initialize(init_note);
 
+//         // Record context lengths before replacement
 //         let note_hash_read_requests_pre_replace = context.note_hash_read_requests.len();
 //         let nullifiers_pre_replace = context.nullifiers.len();
 //         let note_hashes_pre_replace = context.note_hashes.len();
 
-//         let replacement_value = VALUE + 1;
-//         let replacement_note = MockNote::new(replacement_value).build_note();
-//         let emission = state_var.initialize_or_replace(replacement_note);
+//         // Use initialize_or_replace with the helper function
+//         let emission = state_var.initialize_or_replace(init_note, plus_one);
 
-//         assert_eq(emission.note, replacement_note);
-//         assert_eq(emission.storage_slot, STORAGE_SLOT);
+//         let expected_note = MockNote::new(VALUE + 1).build_note();
+//         assert_eq(emission.note, expected_note);
+//         // assert_eq(emission.storage_slot, STORAGE_SLOT);
 
-//         // Replacing a PrivateMutable results in:
-//         // - a read request for the read value
-//         // - a nullifier for the read note
-//         // - a new note for the replacement note
-//         // This would only succeed if the variable had already been initialized, as otherwise the read request would
-//         // fail.
+//         // Verify context updates: read request, nullifier, and new note
 //         assert_eq(context.note_hash_read_requests.len(), note_hash_read_requests_pre_replace + 1);
 //         assert_eq(context.nullifiers.len(), nullifiers_pre_replace + 1);
 //         assert_eq(context.note_hashes.len(), note_hashes_pre_replace + 1);

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

@@ -79,16 +79,12 @@ pub contract AppSubscription {
         // the parent takes. See aztec-nr/src/authwit/auth.nr for more details.
         assert_current_call_valid_authwit::<2>(&mut context, user_address);
 
-        let mut note = storage.subscriptions.at(user_address).get_note().note;
-        assert(note.remaining_txs > 0, "you're out of txs");
-
-        note.remaining_txs -= 1;
-
-        storage.subscriptions.at(user_address).replace(note).emit(
-            &mut context,
-            user_address,
-            MessageDelivery.CONSTRAINED_ONCHAIN,
-        );
+        let emission = storage.subscriptions.at(user_address).replace(|mut note| {
+            assert(note.remaining_txs > 0, "you're out of txs");
+            note.remaining_txs -= 1;
+            note
+        });
+        emission.emit(&mut context, user_address, MessageDelivery.CONSTRAINED_ONCHAIN);
 
         context.set_as_fee_payer();
 
@@ -101,7 +97,11 @@ pub contract AppSubscription {
 
         // We check that the note is not expired. We do that via the router contract to conceal which contract
         // is performing the check.
-        privately_check_block_number(Comparator.LT, note.expiry_block_number, &mut context);
+        privately_check_block_number(
+            Comparator.LT,
+            emission.note.expiry_block_number,
+            &mut context,
+        );
 
         payload.execute_calls(&mut context, config.target_address);
     }
@@ -155,11 +155,11 @@ pub contract AppSubscription {
         );
 
         let subscription_note = SubscriptionNote::new(subscriber, expiry_block_number, tx_count);
-        storage.subscriptions.at(subscriber).initialize_or_replace(subscription_note).emit(
-            &mut context,
-            subscriber,
-            MessageDelivery.CONSTRAINED_ONCHAIN,
-        );
+        storage
+            .subscriptions
+            .at(subscriber)
+            .initialize_or_replace(subscription_note, |_old| subscription_note)
+            .emit(&mut context, subscriber, MessageDelivery.CONSTRAINED_ONCHAIN);
     }
 
     #[utility]