chore: updating token with refunds (#7969)

noir-projects/noir-contracts/contracts/token_with_refunds_contract/src/main.nr

@@ -11,7 +11,7 @@ contract TokenWithRefunds {
     use dep::compressed_string::FieldCompressedString;
 
     use dep::aztec::{
-        hash::compute_secret_hash,
+        context::{PrivateContext, PrivateCallInterface}, hash::compute_secret_hash,
         prelude::{NoteGetterOptions, Map, PublicMutable, SharedImmutable, PrivateSet, AztecAddress},
         encrypted_logs::{
         encrypted_note_emission::{
@@ -29,7 +29,16 @@ contract TokenWithRefunds {
     use crate::types::{transparent_note::TransparentNote, token_note::{TokenNote, TOKEN_NOTE_LEN}, balances_map::BalancesMap};
     // docs:end::imports
 
-    // TODO(#7425): Rename back to `Transfer`
+    // In the first transfer iteration we are computing a lot of additional information (validating inputs, retrieving
+    // keys, etc.), so the gate count is already relatively high. We therefore only read a few notes to keep the happy
+    // case with few constraints.
+    global INITIAL_TRANSFER_CALL_MAX_NOTES = 2;
+    // All the recursive call does is nullify notes, meaning the gate count is low, but it is all constant overhead. We
+    // therefore read more notes than in the base case to increase the efficiency of the overhead, since this results in
+    // an overall small circuit regardless.
+    global RECURSIVE_TRANSFER_CALL_MAX_NOTES = 8;
+
+    // TODO(#7425): Rename back to `Transfer2`
     #[aztec(event)]
     struct Transfer2 {
         from: AztecAddress,
@@ -330,13 +339,86 @@ contract TokenWithRefunds {
         let to_ivpk = header.get_ivpk_m(&mut context, to);
 
         let amount = U128::from_integer(amount);
-        storage.balances.sub(from, amount).emit(encode_and_encrypt_note_with_keys_unconstrained(&mut context, from_ovpk, from_ivpk, from));
+
+        // We reduce `from`'s balance by amount by recursively removing notes over potentially multiple calls. This
+        // method keeps the gate count for each individual call low - reading too many notes at once could result in
+        // circuits in which proving is not feasible.
+        // Since the sum of the amounts in the notes we nullified was potentially larger than amount, we create a new
+        // note for `from` with the change amount, e.g. if `amount` is 10 and two notes are nullified with amounts 8 and
+        // 5, then the change will be 3 (since 8 + 5 - 10 = 3).
+        let change = subtract_balance(
+            &mut context,
+            storage,
+            from,
+            amount,
+            INITIAL_TRANSFER_CALL_MAX_NOTES
+        );
+
+        storage.balances.add(from, change).emit(encode_and_encrypt_note_with_keys_unconstrained(&mut context, from_ovpk, from_ivpk, from));
+
         storage.balances.add(to, amount).emit(encode_and_encrypt_note_with_keys_unconstrained(&mut context, from_ovpk, to_ivpk, to));
 
+        // We don't constrain encryption of the note log in `transfer` (unlike in `transfer_from`) because the transfer
+        // function is only designed to be used in situations where the event is not strictly necessary (e.g. payment to
+        // another person where the payment is considered to be successful when the other party successfully decrypts a
+        // note).
         Transfer2 { from, to, amount: amount.to_field() }.emit(encode_and_encrypt_event_with_keys_unconstrained(&mut context, from_ovpk, to_ivpk, to));
     }
     // docs:end:transfer
 
+    #[contract_library_method]
+    fn subtract_balance(
+        context: &mut PrivateContext,
+        storage: Storage<&mut PrivateContext>,
+        account: AztecAddress,
+        amount: U128,
+        max_notes: u32
+    ) -> U128 {
+        let subtracted = storage.balances.try_sub(account, amount, max_notes);
+
+        // Failing to subtract any amount means that the owner was unable to produce more notes that could be nullified.
+        // We could in some cases fail early inside try_sub if we detected that fewer notes than the maximum were
+        // returned and we were still unable to reach the target amount, but that'd make the code more complicated, and
+        // optimizing for the failure scenario is not as important.
+        assert(subtracted > U128::from_integer(0), "Balance too low");
+
+        if subtracted >= amount {
+            // We have achieved our goal of nullifying notes that add up to more than amount, so we return the change
+            subtracted - amount
+        } else {
+            // try_sub failed to nullify enough notes to reach the target amount, so we compute the amount remaining
+            // and try again.
+            let remaining = amount - subtracted;
+            compute_recurse_subtract_balance_call(*context, account, remaining).call(context)
+        }
+    }
+
+    // TODO(#7729): apply no_predicates to the contract interface method directly instead of having to use a wrapper
+    // like we do here.
+    #[no_predicates]
+    #[contract_library_method]
+    fn compute_recurse_subtract_balance_call(
+        context: PrivateContext,
+        account: AztecAddress,
+        remaining: U128
+    ) -> PrivateCallInterface<25, U128, (AztecAddress, Field)> {
+        TokenWithRefunds::at(context.this_address())._recurse_subtract_balance(account, remaining.to_field())
+    }
+
+    // TODO(#7728): even though the amount should be a U128, we can't have that type in a contract interface due to
+    // serialization issues.
+    #[aztec(internal)]
+    #[aztec(private)]
+    fn _recurse_subtract_balance(account: AztecAddress, amount: Field) -> U128 {
+        subtract_balance(
+            &mut context,
+            storage,
+            account,
+            U128::from_integer(amount),
+            RECURSIVE_TRANSFER_CALL_MAX_NOTES
+        )
+    }
+
     /**
      * Cancel a private authentication witness.
      * @param inner_hash The inner hash of the authwit to cancel.

noir-projects/noir-contracts/contracts/token_with_refunds_contract/src/types/balances_map.nr

@@ -82,35 +82,67 @@ impl<T> BalancesMap<T, &mut PrivateContext> {
     pub fn sub<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
         self: Self,
         owner: AztecAddress,
-        subtrahend: U128
+        amount: U128
     ) -> OuterNoteEmission<T> where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote + Eq {
-        let options = NoteGetterOptions::with_filter(filter_notes_min_sum, subtrahend);
+        let subtracted = self.try_sub(owner, amount, MAX_NOTE_HASH_READ_REQUESTS_PER_CALL);
+
+        // try_sub may have substracted more or less than amount. We must ensure that we subtracted at least as much as
+        // we needed, and then create a new note for the owner for the change (if any).
+        assert(subtracted >= amount, "Balance too low");
+        self.add(owner, subtracted - amount)
+    }
+
+    // Attempts to remove 'target_amount' from the owner's balance. try_sub returns how much was actually subtracted
+    // (i.e. the sum of the value of nullified notes), but this subtracted amount may be more or less than the target
+    // amount.
+    // This may seem odd, but is unfortunately unavoidable due to the number of notes available and their amounts being
+    // unknown. What try_sub does is a best-effort attempt to consume as few notes as possible that add up to more than 
+    // `target_amount`.
+    // The `max_notes` parameter is used to fine-tune the number of constraints created by this function. The gate count
+    // scales relatively linearly with `max_notes`, but a lower `max_notes` parameter increases the likelihood of 
+    // `try_sub` subtracting an amount smaller than `target_amount`.
+    pub fn try_sub<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
+        self: Self,
+        owner: AztecAddress,
+        target_amount: U128,
+        max_notes: u32
+    ) -> U128 where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote + Eq {
+        // We are using a preprocessor here (filter applied in an unconstrained context) instead of a filter because
+        // we do not need to prove correct execution of the preprocessor.
+        // Because the `min_sum` notes is not constrained, users could choose to e.g. not call it. However, all this
+        // might result in is simply higher DA costs due to more nullifiers being emitted. Since we don't care
+        // about proving optimal note usage, we can save these constraints and make the circuit smaller.
+        let options = NoteGetterOptions::with_preprocessor(preprocess_notes_min_sum, target_amount).set_limit(max_notes);
         let notes = self.map.at(owner).pop_notes(options);
 
-        let mut minuend: U128 = U128::from_integer(0);
+        let mut subtracted = U128::from_integer(0);
         for i in 0..options.limit {
             if i < notes.len() {
                 let note = notes.get_unchecked(i);
-                minuend = minuend + note.get_amount();
+                subtracted = subtracted + note.get_amount();
             }
         }
 
-        // This is to provide a nicer error msg,
-        // without it minuend-subtrahend would still catch it, but more generic error then.
-        // without the == true, it includes 'minuend.ge(subtrahend)' as part of the error.
-        assert(minuend >= subtrahend, "Balance too low");
-
-        self.add(owner, minuend - subtrahend)
+        subtracted
     }
 }
 
-pub fn filter_notes_min_sum<T, T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
+// Computes the partial sum of the notes array, stopping once 'min_sum' is reached. This can be used to minimize the
+// number of notes read that add to some value, e.g. when transferring some amount of tokens.
+// The preprocessor (a filter applied in an unconstrained context) does not check if total sum is larger or equal to
+// 'min_sum' - all it does is remove extra notes if it does reach that value.
+// Note that proper usage of this preprocessor requires for notes to be sorted in descending order.
+pub fn preprocess_notes_min_sum<T, T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
     notes: [Option<T>; MAX_NOTE_HASH_READ_REQUESTS_PER_CALL],
     min_sum: U128
 ) -> [Option<T>; MAX_NOTE_HASH_READ_REQUESTS_PER_CALL] where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote {
     let mut selected = [Option::none(); MAX_NOTE_HASH_READ_REQUESTS_PER_CALL];
     let mut sum = U128::from_integer(0);
     for i in 0..notes.len() {
+        // Because we process notes in retrieved order, notes need to be sorted in descending amount order for this
+        // filter to be useful. Consider a 'min_sum' of 4, and a set of notes with amounts [3, 2, 1, 1, 1, 1, 1]. If
+        // sorted in descending order, the filter will only choose the notes with values 3 and 2, but if sorted in
+        // ascending order it will choose 4 notes of value 1.
         if notes[i].is_some() & sum < min_sum {
             let note = notes[i].unwrap_unchecked();
             selected[i] = Option::some(note);