refactor: nicer way to fetch slots in TokenWithRefunds (#7797)

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

@@ -1,6 +1,5 @@
 contract BoxReact {
-    use dep::aztec::prelude::{AztecAddress, PrivateMutable, Map, NoteInterface, NoteHeader};
-    use dep::aztec::protocol_types::point::Point;
+    use dep::aztec::prelude::{AztecAddress, PrivateMutable, Map, NoteInterface, NoteHeader, Point};
     use dep::aztec::encrypted_logs::encrypted_note_emission::encode_and_encrypt_note_with_keys;
     use dep::value_note::value_note::{ValueNote, VALUE_NOTE_LEN};
 

boxes/boxes/vanilla/src/contracts/src/main.nr

@@ -1,6 +1,5 @@
 contract Vanilla {
-    use dep::aztec::prelude::{AztecAddress, PrivateMutable, Map, NoteInterface, NoteHeader};
-    use dep::aztec::protocol_types::point::Point;
+    use dep::aztec::prelude::{AztecAddress, PrivateMutable, Map, NoteInterface, NoteHeader, Point};
     use dep::aztec::encrypted_logs::encrypted_note_emission::encode_and_encrypt_note_with_keys;
     use dep::value_note::value_note::{ValueNote, VALUE_NOTE_LEN};
 

noir-projects/aztec-nr/aztec/src/prelude.nr

@@ -1,6 +1,6 @@
 // docs:start:prelude
 use dep::protocol_types::{
-    address::{AztecAddress, EthAddress}, abis::function_selector::FunctionSelector,
+    address::{AztecAddress, EthAddress}, abis::function_selector::FunctionSelector, point::Point,
     traits::{Serialize, Deserialize}
 };
 use crate::{

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

@@ -424,10 +424,7 @@ contract TokenWithRefunds {
     // docs:end:balance_of_private
 
     // REFUNDS SPECIFIC FUNCTIONALITY FOLLOWS
-    use dep::aztec::{
-        prelude::{FunctionSelector, NoteHeader},
-        protocol_types::{storage::map::derive_storage_slot_in_map, point::Point}
-    };
+    use dep::aztec::prelude::{FunctionSelector, NoteHeader, Point};
     use crate::types::token_note::TokenNoteHidingPoint;
 
     /// We need to use different randomness for the user and for the fee payer notes because if the randomness values
@@ -469,18 +466,13 @@ contract TokenWithRefunds {
         // to the user in the `complete_refund(...)` function.
         storage.balances.sub(user, U128::from_integer(funded_amount)).emit(encode_and_encrypt_note_with_keys(&mut context, user_ovpk, user_ivpk, user));
 
-        // 4. Now we "manually" compute the slots (by setting the slots we insert the notes to the balances map under
-        // the correct keys)
-        let fee_payer_balances_slot = derive_storage_slot_in_map(TokenWithRefunds::storage().balances.slot, fee_payer);
-        let user_balances_slot = derive_storage_slot_in_map(TokenWithRefunds::storage().balances.slot, user);
-
-        // 5. We create the partial notes for the fee payer and the user.
+        // 4. We create the partial notes for the fee payer and the user.
         // --> Called "partial" because they don't have the amount set yet (that will be done in `complete_refund(...)`).
         let fee_payer_partial_note = TokenNote {
             header: NoteHeader {
                 contract_address: AztecAddress::zero(),
                 nonce: 0,
-                storage_slot: fee_payer_balances_slot,
+                storage_slot: storage.balances.map.at(fee_payer).storage_slot,
                 note_hash_counter: 0
             },
             amount: U128::zero(),
@@ -491,19 +483,19 @@ contract TokenWithRefunds {
             header: NoteHeader {
                 contract_address: AztecAddress::zero(),
                 nonce: 0,
-                storage_slot: user_balances_slot,
+                storage_slot: storage.balances.map.at(user).storage_slot,
                 note_hash_counter: 0
             },
             amount: U128::zero(),
             npk_m_hash: user_npk_m_hash,
             randomness: user_randomness
         };
 
-        // 6. Now we get the note hiding points.
+        // 5. Now we get the note hiding points.
         let mut fee_payer_point = fee_payer_partial_note.to_note_hiding_point();
         let mut user_point = user_partial_note.to_note_hiding_point();
 
-        // 7. Set the public teardown function to `complete_refund(...)`. Public teardown is the only time when a public
+        // 6. Set the public teardown function to `complete_refund(...)`. Public teardown is the only time when a public
         // function has access to the final transaction fee, which is needed to compute the actual refund amount.
         context.set_public_teardown_function(
             context.this_address(),
@@ -535,6 +527,7 @@ contract TokenWithRefunds {
         let tx_fee = U128::from_integer(context.transaction_fee());
 
         // 1. We check that user funded the fee payer contract with at least the transaction fee.
+        // TODO(#7796): we should try to prevent reverts here
         assert(funded_amount >= tx_fee, "funded amount not enough to cover tx fee");
 
         // 2. We compute the refund amount as the difference between funded amount and tx fee.