chore!: define key type in maps (#3841)

Closes #2787

Unfortunately key type with a `ToField` trait binding is only defined in
the `.at` method, due to this blocker:

noir-lang/noir#3954

boxes/token/src/contracts/src/main.nr

@@ -58,7 +58,7 @@ contract Token {
         admin: PublicState<AztecAddress, AZTEC_ADDRESS_SERIALIZED_LEN>,
         // docs:end:storage_admin
         // docs:start:storage_minters
-        minters: Map<PublicState<bool, BOOL_SERIALIZED_LEN>>,
+        minters: Map<AztecAddress, PublicState<bool, BOOL_SERIALIZED_LEN>>,
         // docs:end:storage_minters
         // docs:start:storage_balances
         balances: BalancesMap,
@@ -67,7 +67,7 @@ contract Token {
         // docs:start:storage_pending_shields
         pending_shields: Set<TransparentNote, TRANSPARENT_NOTE_LEN>,
         // docs:end:storage_pending_shields
-        public_balances: Map<PublicState<SafeU120, SAFE_U120_SERIALIZED_LEN>>,
+        public_balances: Map<AztecAddress, PublicState<SafeU120, SAFE_U120_SERIALIZED_LEN>>,
     }
     // docs:end:storage_struct
 
@@ -147,7 +147,7 @@ contract Token {
         assert(storage.admin.read().eq(context.msg_sender()), "caller is not admin");
         // docs:end:read_admin
         // docs:start:write_minter
-        storage.minters.at(minter.to_field()).write(approve);
+        storage.minters.at(minter).write(approve);
         // docs:end:write_minter
     }
     // docs:end:set_minter
@@ -156,21 +156,21 @@ contract Token {
     #[aztec(public)]
     fn mint_public(to: AztecAddress, amount: Field) {
         // docs:start:read_minter
-        assert(storage.minters.at(context.msg_sender().to_field()).read(), "caller is not minter");
+        assert(storage.minters.at(context.msg_sender()).read(), "caller is not minter");
         // docs:end:read_minter
         let amount = SafeU120::new(amount);
-        let new_balance = storage.public_balances.at(to.to_field()).read().add(amount);
+        let new_balance = storage.public_balances.at(to).read().add(amount);
         let supply = storage.total_supply.read().add(amount);
 
-        storage.public_balances.at(to.to_field()).write(new_balance);
+        storage.public_balances.at(to).write(new_balance);
         storage.total_supply.write(supply);
     }
     // docs:end:mint_public
 
     // docs:start:mint_private
     #[aztec(public)]
     fn mint_private(amount: Field, secret_hash: Field) {
-        assert(storage.minters.at(context.msg_sender().to_field()).read(), "caller is not minter");
+        assert(storage.minters.at(context.msg_sender()).read(), "caller is not minter");
         let pending_shields = storage.pending_shields;
         let mut note = TransparentNote::new(amount, secret_hash);
         let supply = storage.total_supply.read().add(SafeU120::new(amount));
@@ -193,12 +193,12 @@ contract Token {
         }
 
         let amount = SafeU120::new(amount);
-        let from_balance = storage.public_balances.at(from.to_field()).read().sub(amount);
+        let from_balance = storage.public_balances.at(from).read().sub(amount);
 
         let pending_shields = storage.pending_shields;
         let mut note = TransparentNote::new(amount.value as Field, secret_hash);
 
-        storage.public_balances.at(from.to_field()).write(from_balance);
+        storage.public_balances.at(from).write(from_balance);
         pending_shields.insert_from_public(&mut note);
     }
     // docs:end:shield
@@ -213,11 +213,11 @@ contract Token {
         }
 
         let amount = SafeU120::new(amount);
-        let from_balance = storage.public_balances.at(from.to_field()).read().sub(amount);
-        storage.public_balances.at(from.to_field()).write(from_balance);
+        let from_balance = storage.public_balances.at(from).read().sub(amount);
+        storage.public_balances.at(from).write(from_balance);
 
-        let to_balance = storage.public_balances.at(to.to_field()).read().add(amount);
-        storage.public_balances.at(to.to_field()).write(to_balance);
+        let to_balance = storage.public_balances.at(to).read().add(amount);
+        storage.public_balances.at(to).write(to_balance);
     }
     // docs:end:transfer_public
 
@@ -233,8 +233,8 @@ contract Token {
         // docs:end:assert_current_call_valid_authwit_public
 
         let amount = SafeU120::new(amount);
-        let from_balance = storage.public_balances.at(from.to_field()).read().sub(amount);
-        storage.public_balances.at(from.to_field()).write(from_balance);
+        let from_balance = storage.public_balances.at(from).read().sub(amount);
+        storage.public_balances.at(from).write(from_balance);
 
         let new_supply = storage.total_supply.read().sub(amount);
         storage.total_supply.write(new_supply);
@@ -315,7 +315,7 @@ contract Token {
     internal fn _initialize(new_admin: AztecAddress) {
         assert(!new_admin.is_zero(), "invalid admin");
         storage.admin.write(new_admin);
-        storage.minters.at(new_admin.to_field()).write(true);
+        storage.minters.at(new_admin).write(true);
     }
     // docs:end:initialize
 
@@ -324,8 +324,8 @@ contract Token {
     // docs:start:increase_public_balance
     #[aztec(public)]
     internal fn _increase_public_balance(to: AztecAddress, amount: Field) {
-        let new_balance = storage.public_balances.at(to.to_field()).read().add(SafeU120::new(amount));
-        storage.public_balances.at(to.to_field()).write(new_balance);
+        let new_balance = storage.public_balances.at(to).read().add(SafeU120::new(amount));
+        storage.public_balances.at(to).write(new_balance);
     }
     // docs:end:increase_public_balance
 
@@ -348,7 +348,7 @@ contract Token {
 
     // docs:start:is_minter
     unconstrained fn is_minter(minter: AztecAddress) -> pub bool {
-        storage.minters.at(minter.to_field()).read()
+        storage.minters.at(minter).read()
     }
     // docs:end:is_minter
 
@@ -366,7 +366,7 @@ contract Token {
 
     // docs:start:balance_of_public
     unconstrained fn balance_of_public(owner: AztecAddress) -> pub u120 {
-        storage.public_balances.at(owner.to_field()).read().value
+        storage.public_balances.at(owner).read().value
     }
     // docs:end:balance_of_public
 

boxes/token/src/contracts/src/types/balance_set.nr

@@ -30,21 +30,13 @@ use crate::types::token_note::{TokenNote, TOKEN_NOTE_LEN, TokenNoteMethods};
 // Does not require spending key, but only knowledge.
 // Spending key requirement should be enforced by the contract using this.
 struct BalanceSet {
-  context: Context,
-  owner: AztecAddress,
-  set: Set<TokenNote, TOKEN_NOTE_LEN>
+    owner: AztecAddress,
+    set: Set<TokenNote, TOKEN_NOTE_LEN>
 }
 
 impl BalanceSet {
-  pub fn new(context: Context, owner: AztecAddress, storage_slot: Field) -> Self {
-    assert(storage_slot != 0, "Storage slot 0 not allowed. Storage slots must start from 1.");
-    let set = Set {
-        context,
-        storage_slot,
-        note_interface: TokenNoteMethods,
-    };
+  pub fn new(set: Set<TokenNote, TOKEN_NOTE_LEN>, owner: AztecAddress) -> Self {
     Self {
-        context,
         owner,
         set,
     }
@@ -114,7 +106,10 @@ impl BalanceSet {
     }
 }
 
-pub fn filter_notes_min_sum(notes: [Option<TokenNote>; MAX_READ_REQUESTS_PER_CALL], min_sum: SafeU120) -> [Option<TokenNote>; MAX_READ_REQUESTS_PER_CALL] {
+pub fn filter_notes_min_sum(
+    notes: [Option<TokenNote>; MAX_READ_REQUESTS_PER_CALL],
+    min_sum: SafeU120
+) -> [Option<TokenNote>; MAX_READ_REQUESTS_PER_CALL] {
     let mut selected = [Option::none(); MAX_READ_REQUESTS_PER_CALL];
     let mut sum = SafeU120::min();
     for i in 0..notes.len() {

boxes/token/src/contracts/src/types/balances_map.nr

@@ -4,25 +4,32 @@ use crate::types::balance_set::BalanceSet;
 use dep::aztec::hash::pedersen_hash;
 use dep::protocol_types::address::AztecAddress;
 
+use crate::types::token_note::{TokenNote, TOKEN_NOTE_LEN, TokenNoteMethods};
+use dep::aztec::state_vars::{map::Map, set::Set};
+
 struct BalancesMap {
-    context: Context,
-    storage_slot: Field,
+    store: Map<AztecAddress, Set<TokenNote, TOKEN_NOTE_LEN>>,
 }
 
 impl BalancesMap {
     pub fn new(
         context: Context,
         storage_slot: Field,
     ) -> Self {
-        assert(storage_slot != 0, "Storage slot 0 not allowed. Storage slots must start from 1.");
+        let store = Map::new(context, storage_slot, |context, storage_slot| {
+            Set {
+                context,
+                storage_slot,
+                note_interface: TokenNoteMethods,
+            }
+        });
         Self {
-            context,
-            storage_slot,
+            store,
         }
     }
 
     pub fn at(self, owner: AztecAddress) -> BalanceSet {
-        let derived_storage_slot = pedersen_hash([self.storage_slot, owner.to_field()], 0);
-        BalanceSet::new(self.context, owner, derived_storage_slot)
+        let set = self.store.at(owner);
+        BalanceSet::new(set, owner)
     }
 }

docs/docs/dev_docs/testing/cheat_codes.md

@@ -462,7 +462,7 @@ The baseSlot is specified in the Aztec.nr contract.
 
 ```rust
 struct Storage {
-    balances: Map<PublicState<Field, FIELD_SERIALIZED_LEN>>,
+    balances: Map<AztecAddress, PublicState<Field, FIELD_SERIALIZED_LEN>>,
 }
 
 impl Storage {
@@ -500,7 +500,7 @@ Note: One Field element occupies a storage slot. Hence, structs with multiple fi
 
 ```rust
 struct Storage {
-    balances: Map<PublicState<Field, FIELD_SERIALIZED_LEN>>,
+    balances: Map<AztecAddress, PublicState<Field, FIELD_SERIALIZED_LEN>>,
 }
 
 impl Storage {

docs/docs/misc/migration_notes.md

@@ -113,3 +113,25 @@ Now, just remove the `src` folder,:
 easy_private_token_contract = {git = "https://github.com/AztecProtocol/aztec-packages/", tag ="v0.17.0", directory = "yarn-project/noir-contracts/contracts/easy_private_token_contract"}
 ```
 
+### [Aztec.nr] key type definition in Map
+
+The `Map` class now requires defining the key type in its declaration which *must* implement the `ToField` trait.
+
+Before:
+```rust
+struct Storage {
+    balances: Map<PublicState<Field, FIELD_SERIALIZED_LEN>>
+}
+
+let user_balance = balances.at(owner.to_field())
+```
+
+Now:
+```rust
+struct Storage {
+    balances: Map<AztecAddress, PublicState<Field, FIELD_SERIALIZED_LEN>>
+}
+
+let user_balance = balances.at(owner)
+```
+

noir/tooling/nargo_fmt/tests/expected/contract.nr

@@ -20,8 +20,8 @@ contract Benchmarking {
     };
 
     struct Storage {
-        notes: Map<Set<ValueNote, VALUE_NOTE_LEN>>,
-        balances: Map<PublicState<Field, FIELD_SERIALIZED_LEN>>,
+        notes: Map<Field, Set<AztecAddress, ValueNote, VALUE_NOTE_LEN>>,
+        balances: Map<Field, PublicState<Field, FIELD_SERIALIZED_LEN>>,
     }
 
     impl Storage {

noir/tooling/nargo_fmt/tests/input/contract.nr

@@ -20,8 +20,8 @@ contract Benchmarking {
     };
 
     struct Storage {
-        notes: Map<Set<ValueNote, VALUE_NOTE_LEN>>,
-        balances: Map<PublicState<Field, FIELD_SERIALIZED_LEN>>,
+        notes: Map<AztecAddress, Set<ValueNote, VALUE_NOTE_LEN>>,
+        balances: Map<AztecAddress, PublicState<Field, FIELD_SERIALIZED_LEN>>,
     }
 
     impl Storage {
@@ -58,7 +58,11 @@ contract Benchmarking {
     fn increment_balance(owner: Field, value: Field) {
         let current = storage.balances.at(owner).read();
         storage.balances.at(owner).write(current + value);
-        let _callStackItem1 = context.call_public_function(context.this_address(), FunctionSelector::from_signature("broadcast(Field)"), [owner]);
+        let _callStackItem1 = context.call_public_function(
+            context.this_address(),
+            FunctionSelector::from_signature("broadcast(Field)"),
+            [owner]
+        );
     }
 
     // Est ultricies integer quis auctor elit sed. In nibh mauris cursus mattis molestie a iaculis.
@@ -67,7 +71,12 @@ contract Benchmarking {
         emit_unencrypted_log(&mut context, storage.balances.at(owner).read());
     }
 
-    unconstrained fn compute_note_hash_and_nullifier(contract_address: AztecAddress, nonce: Field, storage_slot: Field, preimage: [Field; VALUE_NOTE_LEN]) -> [Field; 4] {
+    unconstrained fn compute_note_hash_and_nullifier(
+        contract_address: AztecAddress,
+        nonce: Field,
+        storage_slot: Field,
+        preimage: [Field; VALUE_NOTE_LEN]
+    ) -> [Field; 4] {
         let note_header = NoteHeader::new(contract_address, nonce, storage_slot);
         note_utils::compute_note_hash_and_nullifier(ValueNoteMethods, note_header, preimage)
     }

yarn-project/aztec-nr/authwit/src/account.nr

@@ -8,7 +8,7 @@ use crate::auth::IS_VALID_SELECTOR;
 struct AccountActions {
   context: Context,
   is_valid_impl: fn(&mut PrivateContext, Field) -> bool,
-  approved_action: Map<PublicState<bool, BOOL_SERIALIZED_LEN>>,
+  approved_action: Map<Field, PublicState<bool, BOOL_SERIALIZED_LEN>>,
 }
 
 impl AccountActions {

yarn-project/aztec-nr/aztec/src/state_vars/map.nr

@@ -1,22 +1,25 @@
 use crate::context::{PrivateContext, PublicContext, Context};
 use dep::std::option::Option;
-use dep::protocol_types::hash::pedersen_hash;
+use dep::protocol_types::{
+    hash::pedersen_hash,
+    traits::{ToField}
+};
 
 // docs:start:map
-struct Map<V> {
+struct Map<K, V> {
     context: Context,
     storage_slot: Field,
     state_var_constructor: fn(Context, Field) -> V,
 }
 // docs:end:map
 
-impl<V> Map<V> {
+impl<K, V> Map<K, V> {
     // docs:start:new
     pub fn new(
         context: Context,
         storage_slot: Field,
         state_var_constructor: fn(Context, Field) -> V,
-    ) -> Map<V> {
+    ) -> Self {
         assert(storage_slot != 0, "Storage slot 0 not allowed. Storage slots must start from 1.");
         Map {
             context,
@@ -27,9 +30,9 @@ impl<V> Map<V> {
     // docs:end:new
 
     // docs:start:at
-    pub fn at(self, key: Field) -> V {
+    pub fn at(self, key: K) -> V where K: ToField {
         // TODO(#1204): use a generator index for the storage slot
-        let derived_storage_slot = pedersen_hash([self.storage_slot, key],0);
+        let derived_storage_slot = pedersen_hash([self.storage_slot, key.to_field()],0);
 
         let state_var_constructor = self.state_var_constructor;
         state_var_constructor(self.context, derived_storage_slot)