feat!: init storage macro

.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
 }