feat!: init storage macro (AztecProtocol#4200)

Closes: AztecProtocol#3198,
AztecProtocol#2928

~~Requires AztecProtocol#4135,
which is blocked by noir-lang/noir#4124

Automatic storage initialization via aztec macro. 

Full support of public and private state from
`dep::aztec::state_vars::*`, including Maps (and nested Maps!)
Limited support for custom types (as long as they have a single
serializable generic and their constructor is `::new(context,
storage_slot`).

~~Pending: better errors, code comments and some cleanup.~~

Hijacking my own
[comment](AztecProtocol#4200 (comment))
for the explanation:

The idea behind this is that in 99% of cases, storage initialization
(that is, the `impl` for a given `struct Storage...` is redundant, and
the only need for its existence was assigning storage slots...which in
turn were necessary because we didn't know how to serialize the data
structures that were used in a given contract or how much space they
used once serialized (relevant for the public state).

After AztecProtocol#4135 is
merged, both of those things don't have to be explicitly provided since
we're using traits, so the aztec macro can infer the implementation of
the Storage struct just by taking hints from the definition. An example:

```rust
    struct Storage {
        // MyAwesomeStuff implements Serialize<2>, so we assign it slot 1 (and remember that it will take 2 slots due to its size)
        public_var: PublicState<MyAwesomeSuff>, 
        // Right after the first one, assign it to slot: current_slot + previous_size = 3
        another_public_var: PublicState<MyAwesomeSuff>,
        // Private and Public state don't share slots since they "live" in different trees, but keeping the slot count simplifies implementation. 
        // Notes also implement Serialize<N>, but they only take up 1 slot anyways because of hashing, assign it slot 5
        a_singleton: Singleton<ANote>,
        // Maps derive slots via hashing, so we can assume they only "take" 1 slot. We assign it slot 6
        balances: Map<AztecAddress, Singleton<ANote>>,
        // Slot 7
        a_set: Set<ANote>,
        // Slot 8
        imm_singleton: ImmutableSingleton<ANote>,
        // Slot 9. 
        profiles: Map<AztecAddress, Map<Singleton<ANote>>>,
    }

```

We have all the info we need in the AST and HIR to build this
automatically:

```rust
impl Storage {
    fn init(context: Context) -> Self {
        Storage {
            public_var: PublicState::new(context, 1), // No need for serialization methods, taken from the the trait impl
            another_public_var: PublicState::new(context, 3),
            a_singleton: Singleton::new(context, 5),
            // Map init lambda always takes the same form for known storage structs
            balances: Map::new(context, 6, |context, slot| { 
                Singleton::new(context, slot) 
            }),
            a_set: Set::new(context, 7),
            imm_singleton: ImmutableSingleton::new(context, 8),
            // A map of maps is just nesting lambdas, we can infer this too
            profiles: Map::new(context, 9, |context, slot| { 
                Map::new(context, slot, |context, slot| { 
                    Singleton::new(context, slot) 
                })
            })
        }
    }
}

```

...as long as we use "canonical" storage implementations. This means
`AStoragePrimitive<SomethingSerializable>` and
`Map<SomethingWithToField, AStoragePrimitive<SomethingSerializable>>`.

**TLDR:** define the Storage struct, in 99% of cases the macro takes
care of the implementation!

Implementing custom storage will look just like it does know, the macro
will skip automatic generation if it finds one.

---------

Co-authored-by: sirasistant <sirasistant@gmail.com>

.github/workflows/protocol-circuits-gate-diff.yml

@@ -34,6 +34,7 @@ jobs:
           sudo cp -r clang+llvm-16.0.0-x86_64-linux-gnu-ubuntu-18.04/include/* /usr/local/include/
           sudo cp -r clang+llvm-16.0.0-x86_64-linux-gnu-ubuntu-18.04/lib/* /usr/local/lib/
           sudo cp -r clang+llvm-16.0.0-x86_64-linux-gnu-ubuntu-18.04/share/* /usr/local/share/
+          rm -rf clang+llvm-16.0.0-x86_64-linux-gnu-ubuntu-18.04.tar.xz clang+llvm-16.0.0-x86_64-linux-gnu-ubuntu-18.04
 
       - uses: actions/cache@v3
         with:

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

@@ -55,7 +55,7 @@ contract Token {
         minters: Map<AztecAddress, PublicState<bool>>,
         // docs:end:storage_minters
         // docs:start:storage_balances
-        balances: BalancesMap,
+        balances: BalancesMap<TokenNote>,
         // docs:end:storage_balances
         total_supply: PublicState<SafeU120>,
         // docs:start:storage_pending_shields
@@ -65,53 +65,6 @@ contract Token {
     }
     // docs:end:storage_struct
 
-    // docs:start:storage_init
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                // docs:start:storage_admin_init
-                admin: PublicState::new(
-                    context,
-                    1,
-                ),
-                // docs:end:storage_admin_init
-                // docs:start:storage_minters_init
-                minters: Map::new(
-                    context,
-                    2,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot,
-                        )
-                    },
-                ),
-                // docs:end:storage_minters_init
-                // docs:start:storage_balances_init
-                balances: BalancesMap::new(context, 3),
-                // docs:end:storage_balances_init
-                total_supply: PublicState::new(
-                    context,
-                    4,
-                ),
-                // docs:start:storage_pending_shields_init
-                pending_shields: Set::new(context, 5),
-                // docs:end:storage_pending_shields_init
-                public_balances: Map::new(
-                    context,
-                    6,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot,
-                        )
-                    },
-                ),
-            }
-        }
-    }
-    // docs:end:storage_init
-
     // docs:start:constructor
     #[aztec(private)]
     fn constructor(admin: AztecAddress) {
@@ -245,7 +198,7 @@ contract Token {
         pending_shields.remove(note);
 
         // Add the token note to user's balances set
-        storage.balances.at(to).add(SafeU120::new(amount));
+        storage.balances.add(to, SafeU120::new(amount));
     }
     // docs:end:redeem_shield
 
@@ -258,7 +211,7 @@ contract Token {
             assert(nonce == 0, "invalid nonce");
         }
 
-        storage.balances.at(from).sub(SafeU120::new(amount));
+        storage.balances.sub(from, SafeU120::new(amount));
 
         let selector = FunctionSelector::from_signature("_increase_public_balance((Field),Field)");
         let _void = context.call_public_function(context.this_address(), selector, [to.to_field(), amount]);
@@ -277,9 +230,9 @@ contract Token {
         // docs:end:assert_current_call_valid_authwit
 
         let amount = SafeU120::new(amount);
-        storage.balances.at(from).sub(amount);
+        storage.balances.sub(from, amount);
         // docs:start:increase_private_balance
-        storage.balances.at(to).add(amount);
+        storage.balances.add(to, amount);
         // docs:end:increase_private_balance
     }
     // docs:end:transfer
@@ -293,7 +246,7 @@ contract Token {
             assert(nonce == 0, "invalid nonce");
         }
 
-        storage.balances.at(from).sub(SafeU120::new(amount));
+        storage.balances.sub(from, SafeU120::new(amount));
 
         let selector = FunctionSelector::from_signature("_reduce_total_supply(Field)");
         let _void = context.call_public_function(context.this_address(), selector, [amount]);
@@ -350,7 +303,7 @@ contract Token {
 
     // docs:start:balance_of_private
     unconstrained fn balance_of_private(owner: AztecAddress) -> pub u120 {
-        storage.balances.at(owner).balance_of().value
+        storage.balances.balance_of(owner).value
     }
     // docs:end:balance_of_private
 
@@ -374,7 +327,7 @@ contract Token {
         serialized_note: [Field; TOKEN_NOTE_LEN]
     ) -> pub [Field; 4] {
         let note_header = NoteHeader::new(contract_address, nonce, storage_slot);
-        if (storage_slot == 5) {
+        if (storage_slot == storage.pending_shields.get_storage_slot()) {
             note_utils::compute_note_hash_and_nullifier(TransparentNote::deserialize, note_header, serialized_note)
         } else {
             note_utils::compute_note_hash_and_nullifier(TokenNote::deserialize, note_header, serialized_note)

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

@@ -1,4 +1,3 @@
 mod transparent_note;
-mod balance_set;
 mod balances_map;
 mod token_note;

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

@@ -1,34 +1,119 @@
-use dep::aztec::context::{PrivateContext, PublicContext, Context};
 use dep::std::option::Option;
-use crate::types::balance_set::BalanceSet;
-use dep::aztec::hash::pedersen_hash;
-use dep::aztec::protocol_types::address::AztecAddress;
-
-use crate::types::token_note::TokenNote;
-use dep::aztec::state_vars::{map::Map, set::Set};
+use dep::safe_math::SafeU120;
+use dep::aztec::{
+    context::{PrivateContext, PublicContext, Context},
+    hash::pedersen_hash,
+    protocol_types::{
+        address::AztecAddress,
+        constants::MAX_READ_REQUESTS_PER_CALL,
+        traits::{Serialize, Deserialize}
+    },
+    state_vars::{
+        set::Set,
+        map::Map
+    },
+    note::{
+        note_getter::view_notes,
+        note_getter_options::{NoteGetterOptions, SortOrder},
+        note_viewer_options::NoteViewerOptions,
+        note_header::NoteHeader,
+        note_interface::NoteInterface,
+    }
+};
+use crate::types::token_note::{TokenNote, OwnedNote};
 
-struct BalancesMap {
-    store: Map<AztecAddress, Set<TokenNote>>,
+struct BalancesMap<T> {
+    map: Map<AztecAddress, Set<T>>
 }
 
-impl BalancesMap {
+impl<T> BalancesMap<T> {
     pub fn new(
         context: Context,
         storage_slot: Field,
     ) -> Self {
-        let store = Map::new(context, storage_slot, |context, storage_slot| {
-            Set {
-                context,
-                storage_slot,
-            }
-        });
+        assert(storage_slot != 0, "Storage slot 0 not allowed. Storage slots must start from 1.");
         Self {
-            store,
+            map: Map::new(context, storage_slot, |context, slot| Set::new(context, slot))
+        }
+    }
+
+    unconstrained pub fn balance_of<T_SERIALIZED_LEN>(self: Self, owner: AztecAddress) -> SafeU120 where T: Deserialize<T_SERIALIZED_LEN> + Serialize<T_SERIALIZED_LEN> + NoteInterface + OwnedNote {
+        self.balance_of_with_offset(owner, 0)
+    }
+
+  unconstrained pub fn balance_of_with_offset<T_SERIALIZED_LEN>(self: Self, owner: AztecAddress, offset: u32) -> SafeU120 where T: Deserialize<T_SERIALIZED_LEN> + Serialize<T_SERIALIZED_LEN> + NoteInterface + OwnedNote {
+      // Same as SafeU120::new(0), but fewer constraints because no check. 
+      let mut balance = SafeU120::min();
+      // docs:start:view_notes
+      let options = NoteViewerOptions::new().set_offset(offset);
+      let opt_notes = self.map.at(owner).view_notes(options);
+      // docs:end:view_notes
+      let len = opt_notes.len();
+      for i in 0..len {
+          if opt_notes[i].is_some() {
+              balance = balance.add(opt_notes[i].unwrap_unchecked().get_amount());
+          }
+      }
+      if (opt_notes[len - 1].is_some()) {
+          balance = balance.add(self.balance_of_with_offset(owner, offset + opt_notes.len() as u32));
+      }
+
+      balance
+  }
+
+    pub fn add<T_SERIALIZED_LEN>(self: Self, owner: AztecAddress, addend: SafeU120) where T: Deserialize<T_SERIALIZED_LEN> + Serialize<T_SERIALIZED_LEN> + NoteInterface + OwnedNote {
+        let mut addend_note = T::new(addend, owner);
+
+        // docs:start:insert
+        self.map.at(owner).insert(&mut addend_note, true);
+        // docs:end:insert
+    }
+
+    pub fn sub<T_SERIALIZED_LEN>(self: Self, owner: AztecAddress, subtrahend: SafeU120) where T: Deserialize<T_SERIALIZED_LEN> + Serialize<T_SERIALIZED_LEN> + NoteInterface + OwnedNote{
+        // docs:start:get_notes
+        let options = NoteGetterOptions::with_filter(filter_notes_min_sum, subtrahend);
+        let maybe_notes = self.map.at(owner).get_notes(options);
+        // docs:end:get_notes
+
+        let mut minuend: SafeU120 = SafeU120::min();
+        for i in 0..maybe_notes.len() {
+            if maybe_notes[i].is_some() {
+                let note = maybe_notes[i].unwrap_unchecked();
+
+                // Removes the note from the owner's set of notes.
+                // This will call the the `compute_nullifer` function of the `token_note`
+                // which require knowledge of the secret key (currently the users encryption key).
+                // The contract logic must ensure that the spending key is used as well.
+                // docs:start:remove
+                self.map.at(owner).remove(note);
+                // docs:end:remove
+
+                minuend = minuend.add(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.ge(subtrahend) == true, "Balance too low");
+
+        self.add(owner, minuend.sub(subtrahend));
     }
 
-    pub fn at(self, owner: AztecAddress) -> BalanceSet {
-        let set = self.store.at(owner);
-        BalanceSet::new(set, owner)
+}
+
+pub fn filter_notes_min_sum<T, T_SERIALIZED_LEN>(
+    notes: [Option<T>; MAX_READ_REQUESTS_PER_CALL],
+    min_sum: SafeU120
+) -> [Option<T>; MAX_READ_REQUESTS_PER_CALL] where T: Deserialize<T_SERIALIZED_LEN> + Serialize<T_SERIALIZED_LEN> + NoteInterface + OwnedNote {
+    let mut selected = [Option::none(); MAX_READ_REQUESTS_PER_CALL];
+    let mut sum = SafeU120::min();
+    for i in 0..notes.len() {
+        if notes[i].is_some() & sum.lt(min_sum) {
+            let note = notes[i].unwrap_unchecked();
+            selected[i] = Option::some(note);
+            sum = sum.add(note.get_amount());
+        }
     }
+    selected
 }