From 04d37eb4208209b2d6df64d2b96e9fc7869a457c Mon Sep 17 00:00:00 2001
From: Gregorio Juliana <gregojquiros@gmail.com>
Date: Fri, 2 Feb 2024 15:34:25 +0100
Subject: [PATCH] feat!: init storage macro (#4200)

Closes: https://github.com/AztecProtocol/aztec-packages/issues/3198,
https://github.com/AztecProtocol/aztec-packages/issues/2928

~~Requires https://github.com/AztecProtocol/aztec-packages/pull/4135,
which is blocked by https://github.com/noir-lang/noir/issues/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](https://github.com/AztecProtocol/aztec-packages/pull/4200#issuecomment-1914494507)
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 https://github.com/AztecProtocol/aztec-packages/pull/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>
---
 .../workflows/protocol-circuits-gate-diff.yml |   1 +
 boxes/token/src/contracts/src/main.nr         |  63 +--
 boxes/token/src/contracts/src/types.nr        |   1 -
 .../src/contracts/src/types/balance_set.nr    | 113 -----
 .../src/contracts/src/types/balances_map.nr   | 125 +++++-
 .../src/contracts/src/types/token_note.nr     |  42 +-
 .../contracts/src/types/transparent_note.nr   |  14 +-
 boxes/yarn.lock                               |   2 +-
 .../contracts/syntax/slow_updates_tree.md     |  12 +-
 .../contracts/syntax/storage/main.md          |  20 +-
 .../contracts/syntax/storage/private_state.md |  12 +-
 .../contracts/syntax/storage/public_state.md  |  16 +-
 .../contracts/syntax/storage/storage_slots.md |   4 -
 .../aztecnr-getting-started.md                |  12 -
 docs/docs/developers/testing/cheat_codes.md   |   2 +-
 .../writing_private_voting_contract.md        |  14 +-
 .../tutorials/writing_token_contract.md       |   6 -
 docs/docs/misc/migration_notes.md             |  48 ++-
 noir/aztec_macros/src/lib.rs                  | 385 +++++++++++++++++-
 .../noirc_frontend/src/node_interner.rs       |  30 +-
 yarn-project/aztec-nr/aztec/src/state_vars.nr |   1 +
 .../src/state_vars/immutable_singleton.nr     |   3 +
 .../aztec-nr/aztec/src/state_vars/map.nr      |   3 +
 .../aztec/src/state_vars/public_state.nr      |   3 +
 .../aztec-nr/aztec/src/state_vars/set.nr      |   4 +-
 .../aztec/src/state_vars/singleton.nr         |   3 +
 .../src/state_vars/stable_public_state.nr     |   3 +
 .../aztec-nr/aztec/src/state_vars/storage.nr  |   8 +
 .../src/e2e_blacklist_token_contract.test.ts  |   4 +-
 .../end-to-end/src/e2e_state_vars.test.ts     |   4 +-
 .../benchmarking_contract/src/main.nr         |  10 -
 .../contracts/card_game_contract/src/main.nr  |  45 --
 .../contracts/child_contract/src/main.nr      |  11 -
 .../contracts/counter_contract/src/main.nr    |  16 -
 .../docs_example_contract/src/main.nr         |  48 ++-
 .../easy_private_token_contract/src/main.nr   |  14 -
 .../easy_private_voting_contract/src/main.nr  |  26 +-
 .../ecdsa_account_contract/src/main.nr        |  10 +-
 .../contracts/escrow_contract/src/main.nr     |  10 +-
 .../inclusion_proofs_contract/src/main.nr     |  22 -
 .../contracts/lending_contract/src/main.nr    |  45 --
 .../pending_commitments_contract/src/main.nr  |  14 -
 .../contracts/price_feed_contract/src/main.nr |  17 -
 .../schnorr_account_contract/src/main.nr      |  12 +-
 .../contracts/slow_tree_contract/src/main.nr  |  17 -
 .../stateful_test_contract/src/main.nr        |  24 --
 .../contracts/test_contract/src/main.nr       |  10 +-
 .../token_blacklist_contract/src/main.nr      |  61 +--
 .../token_blacklist_contract/src/types.nr     |   1 -
 .../src/types/balance_set.nr                  | 121 ------
 .../src/types/balances_map.nr                 | 115 +++++-
 .../src/types/token_note.nr                   |  41 +-
 .../token_bridge_contract/src/main.nr         |  11 -
 .../contracts/token_contract/src/main.nr      |  95 +----
 .../contracts/token_contract/src/types.nr     |   1 -
 .../token_contract/src/types/balance_set.nr   | 123 ------
 .../token_contract/src/types/balances_map.nr  | 117 +++++-
 .../token_contract/src/types/token_note.nr    |  17 +-
 .../contracts/uniswap_contract/src/main.nr    |  15 -
 59 files changed, 973 insertions(+), 1054 deletions(-)
 delete mode 100644 boxes/token/src/contracts/src/types/balance_set.nr
 create mode 100644 yarn-project/aztec-nr/aztec/src/state_vars/storage.nr
 delete mode 100644 yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balance_set.nr
 delete mode 100644 yarn-project/noir-contracts/contracts/token_contract/src/types/balance_set.nr

diff --git a/.github/workflows/protocol-circuits-gate-diff.yml b/.github/workflows/protocol-circuits-gate-diff.yml
index 94ca8959b4f..e1e8b96b472 100644
--- a/.github/workflows/protocol-circuits-gate-diff.yml
+++ b/.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:
diff --git a/boxes/token/src/contracts/src/main.nr b/boxes/token/src/contracts/src/main.nr
index 5b85b2e4d74..075bd4d0f13 100644
--- a/boxes/token/src/contracts/src/main.nr
+++ b/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)
diff --git a/boxes/token/src/contracts/src/types.nr b/boxes/token/src/contracts/src/types.nr
index d3b3b1c9e77..62a1bb2a363 100644
--- a/boxes/token/src/contracts/src/types.nr
+++ b/boxes/token/src/contracts/src/types.nr
@@ -1,4 +1,3 @@
 mod transparent_note;
-mod balance_set;
 mod balances_map;
 mod token_note;
diff --git a/boxes/token/src/contracts/src/types/balance_set.nr b/boxes/token/src/contracts/src/types/balance_set.nr
deleted file mode 100644
index 184fd5fe925..00000000000
--- a/boxes/token/src/contracts/src/types/balance_set.nr
+++ /dev/null
@@ -1,113 +0,0 @@
-use dep::std::option::Option;
-use dep::safe_math::SafeU120;
-use dep::aztec::{
-  context::Context,
-  state_vars::set::Set,
-};
-use dep::aztec::protocol_types::{
-    constants::MAX_READ_REQUESTS_PER_CALL,
-    address::AztecAddress,
-};
-use dep::aztec::note::{
-    note_getter::view_notes,
-    note_getter_options::{NoteGetterOptions, SortOrder},
-    note_viewer_options::NoteViewerOptions
-};
-
-use crate::types::token_note::TokenNote;
-
-// A set implementing standard manipulation of balances.
-// Does not require spending key, but only knowledge.
-// Spending key requirement should be enforced by the contract using this.
-struct BalanceSet {
-    owner: AztecAddress,
-    set: Set<TokenNote>
-}
-
-impl BalanceSet {
-  pub fn new(set: Set<TokenNote>, owner: AztecAddress) -> Self {
-    Self {
-        owner,
-        set,
-    }
-  }
-
-  unconstrained pub fn balance_of(self: Self) -> SafeU120 {
-    self.balance_of_with_offset(0)
-  }
-
-  unconstrained pub fn balance_of_with_offset(self: Self, offset: u32) -> SafeU120 {
-      // 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.set.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().amount);
-          }
-      }
-      if (opt_notes[len - 1].is_some()) {
-          balance = balance.add(self.balance_of_with_offset(offset + opt_notes.len() as u32));
-      }
-
-      balance
-  }
-
-    pub fn add(self: Self, addend: SafeU120) {
-        let mut addend_note = TokenNote::new(addend, self.owner);
-
-        // docs:start:insert
-        self.set.insert(&mut addend_note, true);
-        // docs:end:insert
-    }
-
-    pub fn sub(self: Self, subtrahend: SafeU120) {
-        // docs:start:get_notes
-        let options = NoteGetterOptions::with_filter(filter_notes_min_sum, subtrahend);
-        let maybe_notes = self.set.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.set.remove(note);
-                // docs:end:remove
-
-                minuend = minuend.add(note.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(minuend.sub(subtrahend));
-    }
-}
-
-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() {
-        if notes[i].is_some() & sum.lt(min_sum) {
-            let note = notes[i].unwrap_unchecked();
-            selected[i] = Option::some(note);
-            sum = sum.add(note.amount);
-        }
-    }
-    selected
-}
diff --git a/boxes/token/src/contracts/src/types/balances_map.nr b/boxes/token/src/contracts/src/types/balances_map.nr
index e8cb09f62a3..2f5835248a0 100644
--- a/boxes/token/src/contracts/src/types/balances_map.nr
+++ b/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
 }
diff --git a/boxes/token/src/contracts/src/types/token_note.nr b/boxes/token/src/contracts/src/types/token_note.nr
index a0283d4a95a..214f942cec8 100644
--- a/boxes/token/src/contracts/src/types/token_note.nr
+++ b/boxes/token/src/contracts/src/types/token_note.nr
@@ -1,10 +1,10 @@
-use dep::aztec::{
-    protocol_types::{
-        address::AztecAddress,
-        constants::{
-            MAX_READ_REQUESTS_PER_CALL
-        },
+use dep::aztec::protocol_types::{
+    address::AztecAddress,
+    constants::{
+        MAX_READ_REQUESTS_PER_CALL
     },
+};
+use dep::aztec::{
     note::{
         note_header::NoteHeader,
         note_interface::NoteInterface,
@@ -14,7 +14,10 @@ use dep::aztec::{
     state_vars::set::Set,
     log::emit_encrypted_log,
     hash::pedersen_hash,
-    protocol_types::traits::{Serialize, Deserialize},
+    protocol_types::traits::{
+        Serialize,
+        Deserialize
+    },
 };
 use dep::aztec::oracle::{
     rand::rand,
@@ -24,6 +27,12 @@ use dep::aztec::oracle::{
 use dep::safe_math::SafeU120;
 use dep::std::option::Option;
 
+trait OwnedNote {
+    fn new(amount: SafeU120, owner: AztecAddress) -> Self;
+    fn get_amount(self) -> SafeU120;
+    fn get_owner(self) -> AztecAddress;
+}
+
 global TOKEN_NOTE_LEN: Field = 3; // 3 plus a header.
 
 struct TokenNote {
@@ -58,7 +67,7 @@ impl Deserialize<TOKEN_NOTE_LEN> for TokenNote {
 }
 
 impl NoteInterface for TokenNote {
-      fn compute_note_content_hash(self) -> Field {
+    fn compute_note_content_hash(self) -> Field {
         // TODO(#1205) Should use a non-zero generator index.
         pedersen_hash(self.serialize(), 0)
     }
@@ -91,8 +100,8 @@ impl NoteInterface for TokenNote {
         self.header = header;
     }
 
-    fn get_header(self) -> NoteHeader {
-        self.header
+    fn get_header(note: TokenNote) -> NoteHeader {
+        note.header
     }
 
     // Broadcasts the note as an encrypted log on L1.
@@ -111,8 +120,8 @@ impl NoteInterface for TokenNote {
   }
 }
 
-impl TokenNote {
-    pub fn new(amount: SafeU120, owner: AztecAddress) -> Self {
+impl OwnedNote for TokenNote {
+    fn new(amount: SafeU120, owner: AztecAddress) -> Self {
         Self {
             amount,
             owner,
@@ -120,4 +129,13 @@ impl TokenNote {
             header: NoteHeader::empty(),
         }
     }
+
+    fn get_amount(self) -> SafeU120 {
+        self.amount
+    }
+
+    fn get_owner(self) -> AztecAddress {
+        self.owner
+    }
+   
 }
diff --git a/boxes/token/src/contracts/src/types/transparent_note.nr b/boxes/token/src/contracts/src/types/transparent_note.nr
index b8d6812330d..4ee641794c4 100644
--- a/boxes/token/src/contracts/src/types/transparent_note.nr
+++ b/boxes/token/src/contracts/src/types/transparent_note.nr
@@ -7,7 +7,7 @@ use dep::aztec::{
     },
     hash::{compute_secret_hash, pedersen_hash},
     context::PrivateContext,
-    protocol_types::traits::{Serialize, Deserialize}
+    protocol_types::traits::{Serialize, Deserialize, Empty}
 };
 
 global TRANSPARENT_NOTE_LEN: Field = 2;
@@ -40,6 +40,12 @@ impl Deserialize<TRANSPARENT_NOTE_LEN> for TransparentNote {
     }
 }
 
+impl Empty for TransparentNote {
+    fn empty() -> Self {
+        TransparentNote::new(0, 0)
+    }
+}
+
 impl NoteInterface for TransparentNote {
 
     fn compute_note_content_hash(self) -> Field {
@@ -57,12 +63,13 @@ impl NoteInterface for TransparentNote {
         pedersen_hash([self.secret, siloed_note_hash],0)
     }
 
+
     fn set_header(&mut self, header: NoteHeader) {
         self.header = header;
     }
 
-    fn get_header(self) -> NoteHeader {
-        self.header
+    fn get_header(note: TransparentNote) -> NoteHeader {
+        note.header
     }
 
     fn broadcast(self, context: &mut PrivateContext, slot: Field) {
@@ -82,7 +89,6 @@ impl TransparentNote {
             header: NoteHeader::empty(),
         }
     }
-
     // new oracle call primitive
     // get me the secret corresponding to this hash
     pub fn new_from_secret(amount: Field, secret: Field) -> Self {
diff --git a/boxes/yarn.lock b/boxes/yarn.lock
index 0f764e1dd06..2b37f7139af 100644
--- a/boxes/yarn.lock
+++ b/boxes/yarn.lock
@@ -253,9 +253,9 @@ __metadata:
   dependencies:
     "@aztec/bb.js": "portal:../../barretenberg/ts"
     "@aztec/foundation": "workspace:^"
+    "@aztec/types": "workspace:^"
     eslint: "npm:^8.35.0"
     lodash.chunk: "npm:^4.2.0"
-    lodash.times: "npm:^4.3.2"
     tslib: "npm:^2.4.0"
   languageName: node
   linkType: soft
diff --git a/docs/docs/developers/contracts/syntax/slow_updates_tree.md b/docs/docs/developers/contracts/syntax/slow_updates_tree.md
index 3f8538baeb9..4f1e79608a1 100644
--- a/docs/docs/developers/contracts/syntax/slow_updates_tree.md
+++ b/docs/docs/developers/contracts/syntax/slow_updates_tree.md
@@ -93,23 +93,19 @@ graph TD
 
 #include_code interface yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr rust
 
-5. Create a storage init function for the same value in both public and private storage
-
-#include_code slow_updates_storage yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr rust
-
-6. Store the SlowTree address in private storage as a FieldNote
+5. Store the SlowTree address in private storage as a FieldNote
 
 #include_code constructor yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr rust
 
-7. Store the SlowTree address in public storage and initialize an instance of SlowMap using this address
+6. Store the SlowTree address in public storage and initialize an instance of SlowMap using this address
 
 #include_code write_slow_update_public yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr rust
 
-8. Now you can read and update from private functions:
+7. Now you can read and update from private functions:
 
 #include_code get_and_update_private yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr rust
 
-9. Or from public functions:
+8. Or from public functions:
 
 #include_code get_public yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr rust
 
diff --git a/docs/docs/developers/contracts/syntax/storage/main.md b/docs/docs/developers/contracts/syntax/storage/main.md
index 7e9b382f82d..1d233536235 100644
--- a/docs/docs/developers/contracts/syntax/storage/main.md
+++ b/docs/docs/developers/contracts/syntax/storage/main.md
@@ -46,7 +46,15 @@ If you don't yet have any private state variables defined you can use this place
 #include_code compute_note_hash_and_nullifier_placeholder /yarn-project/noir-contracts/contracts/token_bridge_contract/src/main.nr rust
 :::
 
-Since Aztec.nr is written in Noir, which is state-less, we need to specify how the storage struct should be initialized to read and write data correctly. This is done by specifying an `init` function that is run in functions that rely on reading or altering the state variables. This `init` function must declare the Storage struct with an instantiation defining how variables are accessed and manipulated. The function MUST be called `init` for the Aztec.nr library to properly handle it (this will be relaxed in the future).
+Since Aztec.nr is written in Noir, which is state-less, we need to specify how the storage struct should be initialized to read and write data correctly. In most cases, the initialization of the storage is handled by the Aztec.nr library, which creates a default implementation based on the provided `Storage` struct.
+
+:::danger
+The automatic initialization has some limitations. Using any combination of the provided storage primitives (`PublicState`, `Singleton`, `ImmutableSingleton`, `Set`, `Map`, including nested ones) will ensure that the library will be able to handle it on its own. 
+
+Custom structures are also supported as long as they are generic over `T`, where `T` implements `Serialize<N>` and `Deserialize<N>` in the case it represents public state and additionally `NoteInterface` for private state.
+:::
+
+Custom initialization of the storage is also possible. This is done by specifying an `init` function that is run in functions that rely on reading or altering the state variables. This `init` function must declare the Storage struct with an instantiation defining how variables are accessed and manipulated. The function MUST be called `init` for the Aztec.nr library to properly handle it (this will be relaxed in the future).
 
 ```rust
 impl Storage {
@@ -79,7 +87,7 @@ You can view the implementation in the Aztec.nr library [here](https://github.co
 
 ### `new`
 
-When declaring the storage for a map, we use the `Map::new()` constructor. As seen below, this takes the `storage_slot` and the `start_var_constructor` along with the [`Context`](../context.mdx).
+When declaring the storage for a map, we use the `Map::new()` constructor. This takes the `storage_slot` and the `start_var_constructor` along with the [`Context`](../context.mdx).
 
 We will see examples of map constructors for public and private variables in later sections.
 
@@ -91,10 +99,6 @@ In the Storage struct:
 
 #include_code storage-map-singleton-declaration /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
-In the `Storage::init` function:
-
-#include_code state_vars-MapSingleton /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
-
 #### Public Example
 
 When declaring a public mapping in Storage, we have to specify that the type is public by declaring it as `PublicState` instead of specifying a note type like with private storage above.
@@ -103,10 +107,6 @@ In the Storage struct:
 
 #include_code storage_minters /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
 
-In the `Storage::init` function:
-
-#include_code storage_minters_init /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
-
 ### `at`
 
 When dealing with a Map, we can access the value at a given key using the `::at` method. This takes the key as an argument and returns the value at that key.
diff --git a/docs/docs/developers/contracts/syntax/storage/private_state.md b/docs/docs/developers/contracts/syntax/storage/private_state.md
index d636459720a..c60590612c1 100644
--- a/docs/docs/developers/contracts/syntax/storage/private_state.md
+++ b/docs/docs/developers/contracts/syntax/storage/private_state.md
@@ -138,9 +138,9 @@ Functionally similar to [`get_note`](#get_note), but executed in unconstrained f
 
 ### `new`
 
-As part of the initialization of the `Storage` struct, the `Singleton` is created as follows, here at storage slot 1 and with the `NoteInterface` for `PublicKeyNote`.
+As part of the initialization of the `Storage` struct, the `Singleton` is created as follows, here at storage slot 1.
 
-#include_code storage_init /yarn-project/noir-contracts/contracts/schnorr_account_contract/src/main.nr rust
+#include_code storage-immutable-singleton-declaration /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 ### `initialize`
 
@@ -154,7 +154,7 @@ For example, if the storage slot depends on the an address then it is possible t
 
 Set the value of an ImmutableSingleton by calling the `initialize` method:
 
-#include_code initialize /yarn-project/noir-contracts/contracts/schnorr_account_contract/src/main.nr rust
+#include_code initialize-immutable-singleton /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 Once initialized, an ImmutableSingleton's value remains unchangeable. This method can only be called once.
 
@@ -168,7 +168,7 @@ Similar to the `Singleton`, we can use the `get_note` method to read the value o
 
 Use this method to retrieve the value of an initialized ImmutableSingleton.
 
-#include_code get_note /yarn-project/noir-contracts/contracts/schnorr_account_contract/src/main.nr rust
+#include_code get_note-immutable-singleton /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 Unlike a `Singleton`, the `get_note` function for an ImmutableSingleton doesn't nullify the current note in the background. This means that multiple accounts can concurrently call this function to read the value.
 
@@ -186,7 +186,7 @@ You can view the implementation [here](https://github.com/AztecProtocol/aztec-pa
 
 And can be added to the `Storage` struct as follows. Here adding a set for a custom note, the TransparentNote (useful for [public -> private communication](../functions/calling_functions.md#public---private)).
 
-#include_code storage_pending_shields /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-set-declaration /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 ### `new`
 
@@ -194,7 +194,7 @@ The `new` method tells the contract how to operate on the underlying storage.
 
 We can initialize the set as follows:
 
-#include_code storage_pending_shields_init /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-set-init /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 ### `insert`
 
diff --git a/docs/docs/developers/contracts/syntax/storage/public_state.md b/docs/docs/developers/contracts/syntax/storage/public_state.md
index 1b53e9dbaf9..8c2a16bfe41 100644
--- a/docs/docs/developers/contracts/syntax/storage/public_state.md
+++ b/docs/docs/developers/contracts/syntax/storage/public_state.md
@@ -29,11 +29,11 @@ When declaring the storage for `T` as a persistent public storage variable, we u
 
 Say that we wish to add `admin` public state variable into our storage struct. In the struct we can define it as:
 
-#include_code storage_admin /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-leader-declaration /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 And then when initializing it in the `Storage::init` function we can do:
 
-#include_code storage_admin_init /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-leader-init /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 We have specified that we are storing a `Field` that should be placed in storage slot `1`. This is just a single value, and is similar to the following in solidity:
 
@@ -41,19 +41,15 @@ We have specified that we are storing a `Field` that should be placed in storage
 address internal admin;
 ```
 
-:::info
-We know its verbose, and are working on making it less so.
-:::
-
 #### Mapping example
 
 Say we want to have a group of `minters` that are able to mint assets in our contract, and we want them in public storage, because [access control in private is quite cumbersome](../../../../learn/concepts/communication/cross_chain_calls.md#a-note-on-l2-access-control). In the `Storage` struct we can add it as follows:
 
-#include_code storage_minters /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-minters-declaration /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 And then when initializing it in the `Storage::init` function we can do it as follows:
 
-#include_code storage_minters_init /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-minters-init /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 In this case, specifying that we are dealing with a map of Fields, and that it should be put at slot 2.
 
@@ -108,9 +104,9 @@ You can find the details of `StablePublicState` in the implementation [here](htt
 ### `new`
 Is done exactly like the `PublicState` struct, but with the `StablePublicState` struct.
 
-#include_code storage_decimals /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-stable-declaration /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
-#include_code storage_decimals_init /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
+#include_code storage-stable-init /yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr rust
 
 ### `initialize`
 
diff --git a/docs/docs/developers/contracts/syntax/storage/storage_slots.md b/docs/docs/developers/contracts/syntax/storage/storage_slots.md
index 62386330e7f..7888d77aae7 100644
--- a/docs/docs/developers/contracts/syntax/storage/storage_slots.md
+++ b/docs/docs/developers/contracts/syntax/storage/storage_slots.md
@@ -49,10 +49,6 @@ siloed_note_hash = H(contract_address, H(H(map_slot, to), note_hash))
 siloed_note_hash = H(contract_address, H(H(map_slot, to), H(amount, to, randomness)))
 ```
 
-Where the `map_slot` is the slot specified in `Storage::init`, recall:
-
-#include_code storage_balances_init yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
-
 And `to` is the actor who receives the note, `amount` of the note and `randomness` is the randomness used to make the note hiding. Without the `randomness` the note could just as well be plaintext (computational cost of a preimage attack would be trivial in such a case).
 
 :::info
diff --git a/docs/docs/developers/getting_started/aztecnr-getting-started.md b/docs/docs/developers/getting_started/aztecnr-getting-started.md
index 32be3a1e2fa..5235e2942e4 100644
--- a/docs/docs/developers/getting_started/aztecnr-getting-started.md
+++ b/docs/docs/developers/getting_started/aztecnr-getting-started.md
@@ -92,18 +92,6 @@ We are also using `balance_utils` from this import, a useful library that allows
 
 This allows us to store our counter in a way that acts as an integer, abstracting the note logic.
 
-## Implement a Storage struct
-
-In this step, we will initiate a `Storage` struct to store balances in a private way. The vast majority Aztec.nr smart contracts will need this.
-
-#include_code storage_struct /yarn-project/noir-contracts/contracts/counter_contract/src/main.nr rust
-
-We are only storing one variable - `counts` as a `Map` of `EasyPrivateUint`. This means our `count` will act as a private integer, and we can map it to an address.
-
-#include_code storage_init /yarn-project/noir-contracts/contracts/counter_contract/src/main.nr rust
-
-This `init` method is creating and initializing a `Storage` instance. This instance includes a `Map` named `counters`. Each entry in this `Map` represents an account's counter.
-
 ## Keep the counter private
 
 Now we’ve got a mechanism for storing our private state, we can start using it to ensure the privacy of balances.
diff --git a/docs/docs/developers/testing/cheat_codes.md b/docs/docs/developers/testing/cheat_codes.md
index 982020d4ae1..834685210f8 100644
--- a/docs/docs/developers/testing/cheat_codes.md
+++ b/docs/docs/developers/testing/cheat_codes.md
@@ -462,7 +462,7 @@ The baseSlot is specified in the Aztec.nr contract.
 
 ```rust
 struct Storage {
-    balances: Map<AztecAddress, PublicState<Field, FIELD_SERIALIZED_LEN>>,
+    balances: Map<AztecAddress, PublicState<Field>>,
 }
 
 impl Storage {
diff --git a/docs/docs/developers/tutorials/writing_private_voting_contract.md b/docs/docs/developers/tutorials/writing_private_voting_contract.md
index e24eecfa367..6c107d3ee2d 100644
--- a/docs/docs/developers/tutorials/writing_private_voting_contract.md
+++ b/docs/docs/developers/tutorials/writing_private_voting_contract.md
@@ -80,11 +80,7 @@ We are using various utils within the Aztec library:
 
 ## Set up storage
 
-Under these imports, we need to set up our contract storage. This is done in two steps:
-
-1. Storage struct
-2. Storage impl block with init function
-
+Under these imports, we need to set up our contract storage. 
 Define the storage struct like so:
 
 #include_code storage_struct yarn-project/noir-contracts/contracts/easy_private_voting_contract/src/main.nr rust
@@ -95,14 +91,6 @@ In this contract, we will store three vars:
 2. tally, as a map with key as the persona and value as the number (in Field) held in public state
 3. voteEnded, as a boolean held in public state
 
-Under the struct, define the impl block like this:
-
-#include_code storage_impl yarn-project/noir-contracts/contracts/easy_private_voting_contract/src/main.nr rust
-
-The `impl` block must define one function `init` that explains how to access and manipulate our variables. We pass context, a storage slot, and serialization methods we imported earlier.
-
-This `init` function will be called every time we access `storage` in our functions.
-
 ## Constructor
 
 The next step is to initialize the contract with a constructor. The constructor will take an address as a parameter and set the admin.
diff --git a/docs/docs/developers/tutorials/writing_token_contract.md b/docs/docs/developers/tutorials/writing_token_contract.md
index 3a867ab7a70..cde563647ad 100644
--- a/docs/docs/developers/tutorials/writing_token_contract.md
+++ b/docs/docs/developers/tutorials/writing_token_contract.md
@@ -244,12 +244,6 @@ Reading through the storage variables:
 
 You can read more about it [here](../contracts/syntax/storage/main.md).
 
-### Initializing Storage
-
-Once we have Storage defined, we need to specify how to initialize it. The `init` method creates and initializes an instance of `Storage`. We define an initialization method for each of the storage variables defined above. Storage initialization is generic and can largely be reused for similar types, across different contracts, but it is important to note that each storage variable specifies it's storage slot, starting at 1.
-
-#include_code storage_init /yarn-project/noir-contracts/contracts/token_contract/src/main.nr rust
-
 ## Functions
 
 Copy and paste the body of each function into the appropriate place in your project if you are following along.
diff --git a/docs/docs/misc/migration_notes.md b/docs/docs/misc/migration_notes.md
index 62168759c5b..0c6533a34ea 100644
--- a/docs/docs/misc/migration_notes.md
+++ b/docs/docs/misc/migration_notes.md
@@ -87,27 +87,6 @@ struct Storage {
     test: Set<CardNote>,
     imm_singleton: ImmutableSingleton<CardNote>,
 }
-
-impl Storage {
-    fn init(context: Context) -> Self {
-        Storage {
-            leader: PublicState::new(
-                context,
-                1
-            ),
-            legendary_card: Singleton::new(context, 2),
-            profiles: Map::new(
-                context,
-                3,
-                |context, slot| {
-                    Singleton::new(context, slot)
-                },
-            ),
-            test: Set::new(context, 4),
-            imm_singleton: ImmutableSingleton::new(context, 4),
-        }
-    }
-}
 ```
 
 For this to work, Notes must implement Serialize, Deserialize and NoteInterface Traits. Previously:
@@ -335,7 +314,32 @@ impl NoteInterface for CardNote {
 }
 ```
 
-Public state must implement Serialize and Deserialize traits.
+Public state must implement Serialize and Deserialize traits. 
+
+It is still possible to manually implement the storage initialization (for custom storage wrappers or internal types that don't implement the required traits). For the above example, the `impl Storage` section would look like this:
+
+```rust
+impl Storage {
+    fn init(context: Context) -> Self {
+        Storage {
+            leader: PublicState::new(
+                context,
+                1
+            ),
+            legendary_card: Singleton::new(context, 2),
+            profiles: Map::new(
+                context,
+                3,
+                |context, slot| {
+                    Singleton::new(context, slot)
+                },
+            ),
+            test: Set::new(context, 4),
+            imm_singleton: ImmutableSingleton::new(context, 4),
+        }
+    }
+}
+```
 
 ## 0.20.0
 
diff --git a/noir/aztec_macros/src/lib.rs b/noir/aztec_macros/src/lib.rs
index 1dbe7631388..00666488455 100644
--- a/noir/aztec_macros/src/lib.rs
+++ b/noir/aztec_macros/src/lib.rs
@@ -1,5 +1,7 @@
-use iter_extended::vecmap;
+use std::borrow::{Borrow, BorrowMut};
+use std::vec;
 
+use iter_extended::vecmap;
 use noirc_frontend::macros_api::FieldElement;
 use noirc_frontend::macros_api::{
     BlockExpression, CallExpression, CastExpression, Distinctness, Expression, ExpressionKind,
@@ -13,6 +15,8 @@ use noirc_frontend::macros_api::{
 use noirc_frontend::macros_api::{CrateId, FileId};
 use noirc_frontend::macros_api::{MacroError, MacroProcessor};
 use noirc_frontend::macros_api::{ModuleDefId, NodeInterner, SortedModule, StructId};
+use noirc_frontend::node_interner::{TraitId, TraitImplKind};
+use noirc_frontend::Lambda;
 
 pub struct AztecMacro;
 
@@ -45,6 +49,8 @@ pub enum AztecMacroError {
     ContractHasTooManyFunctions { span: Span },
     ContractConstructorMissing { span: Span },
     UnsupportedFunctionArgumentType { span: Span, typ: UnresolvedTypeData },
+    UnsupportedStorageType { span: Option<Span>, typ: UnresolvedTypeData },
+    CouldNotAssignStorageSlots { secondary_message: Option<String> },
     EventError { span: Span, message: String },
 }
 
@@ -76,6 +82,16 @@ impl From<AztecMacroError> for MacroError {
                 secondary_message: None,
                 span: Some(span),
             },
+            AztecMacroError::UnsupportedStorageType { span, typ } => MacroError {
+                primary_message: format!("Provided storage type `{typ:?}` is not directly supported in Aztec. Please provide a custom storage implementation"),
+                secondary_message: None,
+                span,
+            },
+            AztecMacroError::CouldNotAssignStorageSlots { secondary_message } => MacroError {
+                primary_message: "Could not assign storage slots, please provide a custom storage implementation".to_string(),
+                secondary_message,
+                span: None,
+            },
             AztecMacroError::EventError { span, message } => MacroError {
                 primary_message: message,
                 secondary_message: None,
@@ -166,6 +182,22 @@ fn member_access(lhs: &str, rhs: &str) -> Expression {
     })))
 }
 
+fn return_type(path: Path) -> FunctionReturnType {
+    let ty = make_type(UnresolvedTypeData::Named(path, vec![]));
+    FunctionReturnType::Ty(ty)
+}
+
+fn lambda(parameters: Vec<(Pattern, UnresolvedType)>, body: Expression) -> Expression {
+    expression(ExpressionKind::Lambda(Box::new(Lambda {
+        parameters,
+        return_type: UnresolvedType {
+            typ: UnresolvedTypeData::Unspecified,
+            span: Some(Span::default()),
+        },
+        body,
+    })))
+}
+
 macro_rules! chained_path {
     ( $base:expr $(, $tail:expr)* ) => {
         {
@@ -196,7 +228,7 @@ fn cast(lhs: Expression, ty: UnresolvedTypeData) -> Expression {
 }
 
 fn make_type(typ: UnresolvedTypeData) -> UnresolvedType {
-    UnresolvedType { typ, span: None }
+    UnresolvedType { typ, span: Some(Span::default()) }
 }
 
 fn index_array(array: Ident, index: &str) -> Expression {
@@ -251,7 +283,8 @@ fn transform_hir(
     crate_id: &CrateId,
     context: &mut HirContext,
 ) -> Result<(), (AztecMacroError, FileId)> {
-    transform_events(crate_id, context)
+    transform_events(crate_id, context)?;
+    assign_storage_slots(crate_id, context)
 }
 
 /// Includes an import to the aztec library if it has not been included yet
@@ -288,6 +321,16 @@ fn check_for_storage_definition(module: &SortedModule) -> bool {
     module.types.iter().any(|r#struct| r#struct.name.0.contents == "Storage")
 }
 
+// Check to see if the user has defined a storage struct
+fn check_for_storage_implementation(module: &SortedModule) -> bool {
+    module.impls.iter().any(|r#impl| match &r#impl.object_type.typ {
+        UnresolvedTypeData::Named(path, _) => {
+            path.segments.last().is_some_and(|segment| segment.0.contents == "Storage")
+        }
+        _ => false,
+    })
+}
+
 // Check if "compute_note_hash_and_nullifier(AztecAddress,Field,Field,[Field; N]) -> [Field; 4]" is defined
 fn check_for_compute_note_hash_and_nullifier_definition(module: &SortedModule) -> bool {
     module.functions.iter().any(|func| {
@@ -343,9 +386,15 @@ fn transform_module(
 
     // Check for a user defined storage struct
     let storage_defined = check_for_storage_definition(module);
+    let storage_implemented = check_for_storage_implementation(module);
+
+    let crate_graph = &context.crate_graph[crate_id];
+
+    if storage_defined && !storage_implemented {
+        generate_storage_implementation(module).map_err(|err| (err, crate_graph.root_file_id))?;
+    }
 
     if storage_defined && !check_for_compute_note_hash_and_nullifier_definition(module) {
-        let crate_graph = &context.crate_graph[crate_id];
         return Err((
             AztecMacroError::ComputeNoteHashAndNullifierNotFound { span: Span::default() },
             crate_graph.root_file_id,
@@ -403,6 +452,134 @@ fn transform_module(
     Ok(has_transformed_module)
 }
 
+/// Auxiliary function to generate the storage constructor for a given field, using
+/// the Storage definition as a reference. Supports nesting.
+fn generate_storage_field_constructor(
+    (type_ident, unresolved_type): &(Ident, UnresolvedType),
+    slot: Expression,
+) -> Result<Expression, AztecMacroError> {
+    let typ = &unresolved_type.typ;
+    match typ {
+        UnresolvedTypeData::Named(path, generics) => {
+            let mut new_path = path.clone().to_owned();
+            new_path.segments.push(ident("new"));
+            match path.segments.last().unwrap().0.contents.as_str() {
+                "Map" => Ok(call(
+                    variable_path(new_path),
+                    vec![
+                        variable("context"),
+                        slot,
+                        lambda(
+                            vec![
+                                (
+                                    pattern("context"),
+                                    make_type(UnresolvedTypeData::Named(
+                                        chained_path!("aztec", "context", "Context"),
+                                        vec![],
+                                    )),
+                                ),
+                                (
+                                    Pattern::Identifier(ident("slot")),
+                                    make_type(UnresolvedTypeData::FieldElement),
+                                ),
+                            ],
+                            generate_storage_field_constructor(
+                                &(type_ident.clone(), generics.iter().last().unwrap().clone()),
+                                variable("slot"),
+                            )?,
+                        ),
+                    ],
+                )),
+                _ => Ok(call(variable_path(new_path), vec![variable("context"), slot])),
+            }
+        }
+        _ => Err(AztecMacroError::UnsupportedStorageType {
+            typ: typ.clone(),
+            span: Some(type_ident.span()),
+        }),
+    }
+}
+
+// Generates the Storage implementation block from the Storage struct definition if it does not exist
+/// From:
+///
+/// struct Storage {
+///     a_map: Map<Field, SomeStoragePrimitive<ASerializableType>>,
+///     a_nested_map: Map<Field, Map<Field, SomeStoragePrimitive<ASerializableType>>>,
+///     a_field: SomeStoragePrimitive<ASerializableType>,
+/// }
+///
+/// To:
+///
+/// impl Storage {
+///    fn init(context: Context) -> Self {
+///        Storage {
+///             a_map: Map::new(context, 0, |context, slot| {
+///                 SomeStoragePrimitive::new(context, slot)
+///             }),
+///             a_nested_map: Map::new(context, 0, |context, slot| {
+///                 Map::new(context, slot, |context, slot| {
+///                     SomeStoragePrimitive::new(context, slot)
+///                })
+///            }),
+///            a_field: SomeStoragePrimitive::new(context, 0),
+///         }
+///    }
+/// }
+///
+/// Storage slots are generated as 0 and will be populated using the information from the HIR
+/// at a later stage.
+fn generate_storage_implementation(module: &mut SortedModule) -> Result<(), AztecMacroError> {
+    let definition =
+        module.types.iter().find(|r#struct| r#struct.name.0.contents == "Storage").unwrap();
+
+    let slot_zero = expression(ExpressionKind::Literal(Literal::Integer(
+        FieldElement::from(i128::from(0)),
+        false,
+    )));
+
+    let field_constructors = definition
+        .fields
+        .iter()
+        .flat_map(|field| {
+            generate_storage_field_constructor(field, slot_zero.clone())
+                .map(|expression| (field.0.clone(), expression))
+        })
+        .collect();
+
+    let storage_constructor_statement = make_statement(StatementKind::Expression(expression(
+        ExpressionKind::constructor((chained_path!("Storage"), field_constructors)),
+    )));
+
+    let init = NoirFunction::normal(FunctionDefinition::normal(
+        &ident("init"),
+        &vec![],
+        &[(
+            ident("context"),
+            make_type(UnresolvedTypeData::Named(
+                chained_path!("aztec", "context", "Context"),
+                vec![],
+            )),
+        )],
+        &BlockExpression(vec![storage_constructor_statement]),
+        &[],
+        &return_type(chained_path!("Self")),
+    ));
+
+    let storage_impl = TypeImpl {
+        object_type: UnresolvedType {
+            typ: UnresolvedTypeData::Named(chained_path!("Storage"), vec![]),
+            span: Some(Span::default()),
+        },
+        type_span: Span::default(),
+        generics: vec![],
+        methods: vec![init],
+    };
+    module.impls.push(storage_impl);
+
+    Ok(())
+}
+
 /// If it does, it will insert the following things:
 /// - A new Input that is provided for a kernel app circuit, named: {Public/Private}ContextInputs
 /// - Hashes all of the function input variables
@@ -484,6 +661,23 @@ fn collect_crate_structs(crate_id: &CrateId, context: &HirContext) -> Vec<Struct
         .collect()
 }
 
+fn collect_traits(context: &HirContext) -> Vec<TraitId> {
+    let crates = context.crates();
+    crates
+        .flat_map(|crate_id| context.def_map(&crate_id).map(|def_map| def_map.modules()))
+        .flatten()
+        .flat_map(|(_, module)| {
+            module.type_definitions().filter_map(|typ| {
+                if let ModuleDefId::TraitId(struct_id) = typ {
+                    Some(struct_id)
+                } else {
+                    None
+                }
+            })
+        })
+        .collect()
+}
+
 /// Substitutes the signature literal that was introduced in the selector method previously with the actual signature.
 fn transform_event(
     struct_id: StructId,
@@ -576,6 +770,185 @@ fn transform_events(
     Ok(())
 }
 
+/// Obtains the serialized length of a type that implements the Serialize trait.
+fn get_serialized_length(
+    traits: &[TraitId],
+    typ: &Type,
+    interner: &NodeInterner,
+) -> Result<u64, AztecMacroError> {
+    let (struct_name, maybe_stored_in_state) = match typ {
+        Type::Struct(struct_type, generics) => {
+            Ok((struct_type.borrow().name.0.contents.clone(), generics.get(0)))
+        }
+        _ => Err(AztecMacroError::CouldNotAssignStorageSlots {
+            secondary_message: Some("State storage variable must be a struct".to_string()),
+        }),
+    }?;
+    let stored_in_state =
+        maybe_stored_in_state.ok_or(AztecMacroError::CouldNotAssignStorageSlots {
+            secondary_message: Some("State storage variable must be generic".to_string()),
+        })?;
+
+    let is_note = traits.iter().any(|&trait_id| {
+        let r#trait = interner.get_trait(trait_id);
+        r#trait.name.0.contents == "NoteInterface"
+            && !interner.lookup_all_trait_implementations(stored_in_state, trait_id).is_empty()
+    });
+
+    // Maps and (private) Notes always occupy a single slot. Someone could store a Note in PublicState for whatever reason though.
+    if struct_name == "Map" || (is_note && struct_name != "PublicState") {
+        return Ok(1);
+    }
+
+    let serialized_trait_impl_kind = traits
+        .iter()
+        .filter_map(|&trait_id| {
+            let r#trait = interner.get_trait(trait_id);
+            if r#trait.borrow().name.0.contents == "Serialize"
+                && r#trait.borrow().generics.len() == 1
+            {
+                interner
+                    .lookup_all_trait_implementations(stored_in_state, trait_id)
+                    .into_iter()
+                    .next()
+            } else {
+                None
+            }
+        })
+        .next()
+        .ok_or(AztecMacroError::CouldNotAssignStorageSlots {
+            secondary_message: Some("Stored data must implement Serialize trait".to_string()),
+        })?;
+
+    let serialized_trait_impl_id = match serialized_trait_impl_kind {
+        TraitImplKind::Normal(trait_impl_id) => Ok(trait_impl_id),
+        _ => Err(AztecMacroError::CouldNotAssignStorageSlots { secondary_message: None }),
+    }?;
+
+    let serialized_trait_impl_shared = interner.get_trait_implementation(*serialized_trait_impl_id);
+    let serialized_trait_impl = serialized_trait_impl_shared.borrow();
+
+    match serialized_trait_impl.trait_generics.get(0).unwrap() {
+        Type::Constant(value) => Ok(*value),
+        _ => Err(AztecMacroError::CouldNotAssignStorageSlots { secondary_message: None }),
+    }
+}
+
+/// Assigns storage slots to the storage struct fields based on the serialized length of the types. This automatic assignment
+/// will only trigger if the assigned storage slot is invalid (0 as generated by generate_storage_implementation)
+fn assign_storage_slots(
+    crate_id: &CrateId,
+    context: &mut HirContext,
+) -> Result<(), (AztecMacroError, FileId)> {
+    let traits: Vec<_> = collect_traits(context);
+    for struct_id in collect_crate_structs(crate_id, context) {
+        let interner: &mut NodeInterner = context.def_interner.borrow_mut();
+        let r#struct = interner.get_struct(struct_id);
+        let file_id = r#struct.borrow().location.file;
+        if r#struct.borrow().name.0.contents == "Storage" && r#struct.borrow().id.krate().is_root()
+        {
+            let init_id = interner
+                .lookup_method(
+                    &Type::Struct(interner.get_struct(struct_id), vec![]),
+                    struct_id,
+                    "init",
+                    false,
+                )
+                .ok_or((
+                    AztecMacroError::CouldNotAssignStorageSlots {
+                        secondary_message: Some(
+                            "Storage struct must have an init function".to_string(),
+                        ),
+                    },
+                    file_id,
+                ))?;
+            let init_function = interner.function(&init_id).block(interner);
+            let init_function_statement_id = init_function.statements().first().ok_or((
+                AztecMacroError::CouldNotAssignStorageSlots {
+                    secondary_message: Some("Init storage statement not found".to_string()),
+                },
+                file_id,
+            ))?;
+            let storage_constructor_statement = interner.statement(init_function_statement_id);
+
+            let storage_constructor_expression = match storage_constructor_statement {
+                HirStatement::Expression(expression_id) => {
+                    match interner.expression(&expression_id) {
+                        HirExpression::Constructor(hir_constructor_expression) => {
+                            Ok(hir_constructor_expression)
+                        }
+                        _ => Err((AztecMacroError::CouldNotAssignStorageSlots {
+                            secondary_message: Some(
+                                "Storage constructor statement must be a constructor expression"
+                                    .to_string(),
+                            ),
+                        }, file_id))
+                    }
+                }
+                _ => Err((
+                    AztecMacroError::CouldNotAssignStorageSlots {
+                        secondary_message: Some(
+                            "Storage constructor statement must be an expression".to_string(),
+                        ),
+                    },
+                    file_id,
+                )),
+            }?;
+
+            let mut storage_slot: u64 = 1;
+            for (index, (_, expr_id)) in storage_constructor_expression.fields.iter().enumerate() {
+                let fields = r#struct.borrow().get_fields(&[]);
+                let (_, field_type) = fields.get(index).unwrap();
+                let new_call_expression = match interner.expression(expr_id) {
+                    HirExpression::Call(hir_call_expression) => Ok(hir_call_expression),
+                    _ => Err((
+                        AztecMacroError::CouldNotAssignStorageSlots {
+                            secondary_message: Some(
+                                "Storage field initialization expression is not a call expression"
+                                    .to_string(),
+                            ),
+                        },
+                        file_id,
+                    )),
+                }?;
+
+                let slot_arg_expression = interner.expression(&new_call_expression.arguments[1]);
+
+                let current_storage_slot = match slot_arg_expression {
+                    HirExpression::Literal(HirLiteral::Integer(slot, _)) => {
+                        Ok(slot.borrow().to_u128())
+                    }
+                    _ => Err((
+                        AztecMacroError::CouldNotAssignStorageSlots {
+                            secondary_message: Some(
+                                "Storage slot argument expression must be a literal integer"
+                                    .to_string(),
+                            ),
+                        },
+                        file_id,
+                    )),
+                }?;
+
+                if current_storage_slot != 0 {
+                    continue;
+                }
+
+                let type_serialized_len = get_serialized_length(&traits, field_type, interner)
+                    .map_err(|err| (err, file_id))?;
+                interner.update_expression(new_call_expression.arguments[1], |expr| {
+                    *expr = HirExpression::Literal(HirLiteral::Integer(
+                        FieldElement::from(u128::from(storage_slot)),
+                        false,
+                    ));
+                });
+
+                storage_slot += type_serialized_len;
+            }
+        }
+    }
+    Ok(())
+}
+
 const SIGNATURE_PLACEHOLDER: &str = "SIGNATURE_PLACEHOLDER";
 
 /// Generates the impl for an event selector
@@ -969,9 +1342,7 @@ fn make_castable_return_type(expression: Expression) -> Statement {
 /// }
 fn create_return_type(ty: &str) -> FunctionReturnType {
     let return_path = chained_path!("aztec", "abi", ty);
-
-    let ty = make_type(UnresolvedTypeData::Named(return_path, vec![]));
-    FunctionReturnType::Ty(ty)
+    return_type(return_path)
 }
 
 /// Create Context Finish
diff --git a/noir/compiler/noirc_frontend/src/node_interner.rs b/noir/compiler/noirc_frontend/src/node_interner.rs
index b856b54f6ca..11ef12ef83e 100644
--- a/noir/compiler/noirc_frontend/src/node_interner.rs
+++ b/noir/compiler/noirc_frontend/src/node_interner.rs
@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::ops::Deref;
 
 use arena::{Arena, Index};
 use fm::FileId;
@@ -1042,6 +1043,34 @@ impl NodeInterner {
         Ok(impl_kind)
     }
 
+    /// Given a `ObjectType: TraitId` pair, find all implementations without taking constraints into account or
+    /// applying any type bindings. Useful to look for a specific trait in a type that is used in a macro.
+    pub fn lookup_all_trait_implementations(
+        &self,
+        object_type: &Type,
+        trait_id: TraitId,
+    ) -> Vec<&TraitImplKind> {
+        let trait_impl = self.trait_implementation_map.get(&trait_id);
+
+        trait_impl
+            .map(|trait_impl| {
+                trait_impl
+                    .iter()
+                    .filter_map(|(typ, impl_kind)| match &typ {
+                        Type::Forall(_, typ) => {
+                            if typ.deref() == object_type {
+                                Some(impl_kind)
+                            } else {
+                                None
+                            }
+                        }
+                        _ => None,
+                    })
+                    .collect()
+            })
+            .unwrap_or(vec![])
+    }
+
     /// Similar to `lookup_trait_implementation` but does not apply any type bindings on success.
     pub fn try_lookup_trait_implementation(
         &self,
@@ -1263,7 +1292,6 @@ impl NodeInterner {
         force_type_check: bool,
     ) -> Option<FuncId> {
         let methods = self.struct_methods.get(&(id, method_name.to_owned()));
-
         // If there is only one method, just return it immediately.
         // It will still be typechecked later.
         if !force_type_check {
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars.nr b/yarn-project/aztec-nr/aztec/src/state_vars.nr
index d8213bb1ef0..177844f75c1 100644
--- a/yarn-project/aztec-nr/aztec/src/state_vars.nr
+++ b/yarn-project/aztec-nr/aztec/src/state_vars.nr
@@ -4,3 +4,4 @@ mod public_state;
 mod set;
 mod singleton;
 mod stable_public_state;
+mod storage;
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars/immutable_singleton.nr b/yarn-project/aztec-nr/aztec/src/state_vars/immutable_singleton.nr
index 187feca45fc..9d595081c20 100644
--- a/yarn-project/aztec-nr/aztec/src/state_vars/immutable_singleton.nr
+++ b/yarn-project/aztec-nr/aztec/src/state_vars/immutable_singleton.nr
@@ -16,6 +16,7 @@ use crate::note::{
     note_viewer_options::NoteViewerOptions,
 };
 use crate::oracle::notes::check_nullifier_exists;
+use crate::state_vars::storage::Storage;
 
 // docs:start:struct
 struct ImmutableSingleton<Note> {
@@ -24,6 +25,8 @@ struct ImmutableSingleton<Note> {
 }
 // docs:end:struct
 
+impl<T> Storage<T> for ImmutableSingleton<T> {}
+
 impl<Note> ImmutableSingleton<Note> {
     // docs:start:new
     pub fn new(
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars/map.nr b/yarn-project/aztec-nr/aztec/src/state_vars/map.nr
index fefa84ea00e..7e38552e2ca 100644
--- a/yarn-project/aztec-nr/aztec/src/state_vars/map.nr
+++ b/yarn-project/aztec-nr/aztec/src/state_vars/map.nr
@@ -4,6 +4,7 @@ use dep::protocol_types::{
     hash::pedersen_hash,
     traits::{ToField}
 };
+use crate::state_vars::storage::Storage;
 
 // docs:start:map
 struct Map<K, V> {
@@ -13,6 +14,8 @@ struct Map<K, V> {
 }
 // docs:end:map
 
+impl<K, T> Storage<T> for Map<K, T> {}
+
 impl<K, V> Map<K, V> {
     // docs:start:new
     pub fn new(
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars/public_state.nr b/yarn-project/aztec-nr/aztec/src/state_vars/public_state.nr
index 9309f3d9d49..081b0d59653 100644
--- a/yarn-project/aztec-nr/aztec/src/state_vars/public_state.nr
+++ b/yarn-project/aztec-nr/aztec/src/state_vars/public_state.nr
@@ -3,6 +3,7 @@ use crate::oracle::storage::storage_read;
 use crate::oracle::storage::storage_write;
 use dep::std::option::Option;
 use dep::protocol_types::traits::{Deserialize, Serialize};
+use crate::state_vars::storage::Storage;
 
 // docs:start:public_state_struct
 struct PublicState<T> {
@@ -11,6 +12,8 @@ struct PublicState<T> {
 }
 // docs:end:public_state_struct
 
+impl<T> Storage<T> for PublicState<T> {}
+
 impl<T> PublicState<T> {
     // docs:start:public_state_struct_new
     pub fn new(
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars/set.nr b/yarn-project/aztec-nr/aztec/src/state_vars/set.nr
index 03e53b33d7a..6813a9fa98b 100644
--- a/yarn-project/aztec-nr/aztec/src/state_vars/set.nr
+++ b/yarn-project/aztec-nr/aztec/src/state_vars/set.nr
@@ -15,6 +15,7 @@ use crate::note::{
     note_viewer_options::NoteViewerOptions,
     utils::compute_note_hash_for_consumption,
 };
+use crate::state_vars::storage::Storage;
 
 // docs:start:struct
 struct Set<Note> {
@@ -23,6 +24,8 @@ struct Set<Note> {
 }
 // docs:end:struct
 
+impl<T> Storage<T> for Set<T> {}
+
 impl<Note> Set<Note> {
     // docs:start:new
     pub fn new(
@@ -36,7 +39,6 @@ impl<Note> Set<Note> {
         }
     }
     // docs:end:new
-
     // docs:start:insert
     pub fn insert<N>(self,
         note: &mut Note,
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars/singleton.nr b/yarn-project/aztec-nr/aztec/src/state_vars/singleton.nr
index 19b4e91bb5c..47d8aef5a9f 100644
--- a/yarn-project/aztec-nr/aztec/src/state_vars/singleton.nr
+++ b/yarn-project/aztec-nr/aztec/src/state_vars/singleton.nr
@@ -20,6 +20,7 @@ use crate::oracle::{
     nullifier_key::get_nullifier_secret_key,
     notes::check_nullifier_exists,
 };
+use crate::state_vars::storage::Storage;
 
 // docs:start:struct
 struct Singleton<Note> {
@@ -28,6 +29,8 @@ struct Singleton<Note> {
 }
 // docs:end:struct
 
+impl<T> Storage<T> for Singleton<T> {}
+
 impl<Note> Singleton<Note> {
     // docs:start:new
     pub fn new(
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars/stable_public_state.nr b/yarn-project/aztec-nr/aztec/src/state_vars/stable_public_state.nr
index 013f059aef0..af4925f47ff 100644
--- a/yarn-project/aztec-nr/aztec/src/state_vars/stable_public_state.nr
+++ b/yarn-project/aztec-nr/aztec/src/state_vars/stable_public_state.nr
@@ -5,12 +5,15 @@ use crate::oracle::{
 use crate::history::public_value_inclusion::prove_public_value_inclusion;
 use dep::std::option::Option;
 use dep::protocol_types::traits::{Deserialize, Serialize};
+use crate::state_vars::storage::Storage;
 
 struct StablePublicState<T>{
     context: Context,
     storage_slot: Field,
 }
 
+impl<T> Storage<T> for StablePublicState<T> {}
+
 impl<T> StablePublicState<T> {
     pub fn new(
         // Note: Passing the contexts to new(...) just to have an interface compatible with a Map.
diff --git a/yarn-project/aztec-nr/aztec/src/state_vars/storage.nr b/yarn-project/aztec-nr/aztec/src/state_vars/storage.nr
new file mode 100644
index 00000000000..5f28b55f9f7
--- /dev/null
+++ b/yarn-project/aztec-nr/aztec/src/state_vars/storage.nr
@@ -0,0 +1,8 @@
+use crate::context::{Context};
+use dep::protocol_types::traits::{Deserialize, Serialize};
+
+trait Storage<T> where T: Serialize<N> + Deserialize<N> {
+    fn get_storage_slot(self) -> Field {
+        self.storage_slot
+    }
+}
diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts
index f39c1284f81..d6f41b87384 100644
--- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts
+++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts
@@ -87,7 +87,7 @@ describe('e2e_blacklist_token_contract', () => {
   };
 
   const addPendingShieldNoteToPXE = async (accountIndex: number, amount: bigint, secretHash: Fr, txHash: TxHash) => {
-    const storageSlot = new Fr(5); // The storage slot of `pending_shields` is 5.
+    const storageSlot = new Fr(4); // The storage slot of `pending_shields` is 4.
     const note = new Note([new Fr(amount), secretHash]);
     const extendedNote = new ExtendedNote(note, accounts[accountIndex].address, asset.address, storageSlot, txHash);
     await wallets[accountIndex].addNote(extendedNote);
@@ -112,7 +112,7 @@ describe('e2e_blacklist_token_contract', () => {
 
     // Add the note
     const note = new Note([slowTree.address.toField()]);
-    const storageSlot = new Fr(7);
+    const storageSlot = new Fr(6);
 
     for (const wallet of wallets) {
       const extendedNote = new ExtendedNote(
diff --git a/yarn-project/end-to-end/src/e2e_state_vars.test.ts b/yarn-project/end-to-end/src/e2e_state_vars.test.ts
index 0ad3a75a04e..dcf98641561 100644
--- a/yarn-project/end-to-end/src/e2e_state_vars.test.ts
+++ b/yarn-project/end-to-end/src/e2e_state_vars.test.ts
@@ -127,7 +127,7 @@ describe('e2e_state_vars', () => {
   describe('Immutable Singleton', () => {
     it('fail to read uninitialized singleton', async () => {
       expect(await contract.methods.is_imm_initialized().view()).toEqual(false);
-      await expect(contract.methods.get_imm_card().view()).rejects.toThrowError();
+      await expect(contract.methods.view_imm_card().view()).rejects.toThrowError();
     });
 
     it('initialize singleton', async () => {
@@ -152,7 +152,7 @@ describe('e2e_state_vars', () => {
 
     it('read initialized singleton', async () => {
       expect(await contract.methods.is_imm_initialized().view()).toEqual(true);
-      const { points, randomness } = await contract.methods.get_imm_card().view();
+      const { points, randomness } = await contract.methods.view_imm_card().view();
       expect(points).toEqual(POINTS);
       expect(randomness).toEqual(RANDOMNESS);
     });
diff --git a/yarn-project/noir-contracts/contracts/benchmarking_contract/src/main.nr b/yarn-project/noir-contracts/contracts/benchmarking_contract/src/main.nr
index 8b9bd1df476..6b4502e83bc 100644
--- a/yarn-project/noir-contracts/contracts/benchmarking_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/benchmarking_contract/src/main.nr
@@ -14,7 +14,6 @@ contract Benchmarking {
         protocol_types::{
             abis::function_selector::FunctionSelector,
             address::AztecAddress,
-            type_serialization::FIELD_SERIALIZED_LEN
         },
         context::{Context},
         note::{utils as note_utils, note_getter_options::NoteGetterOptions, note_header::NoteHeader},
@@ -27,15 +26,6 @@ contract Benchmarking {
         balances: Map<AztecAddress, PublicState<Field>>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                notes: Map::new(context, 1, |context, slot| { Set::new(context, slot) }),
-                balances: Map::new(context, 2, |context, slot| { PublicState::new(context, slot) }),
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor() {}
 
diff --git a/yarn-project/noir-contracts/contracts/card_game_contract/src/main.nr b/yarn-project/noir-contracts/contracts/card_game_contract/src/main.nr
index df143229104..d046e92d82f 100644
--- a/yarn-project/noir-contracts/contracts/card_game_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/card_game_contract/src/main.nr
@@ -58,51 +58,6 @@ contract CardGame {
         games: Map<Field, PublicState<Game>>,
     }
 
-    impl Storage {
-        fn init(
-            context: Context,
-        ) -> Self {
-            Storage {
-                collections: Map::new(
-                    context,
-                    1,
-                    |context, slot| {
-                        Deck::new(
-                            context,
-                            slot,
-                        )
-                    },
-                ),
-                game_decks: Map::new(
-                    context,
-                    2,
-                    |context, slot| {
-                        Map::new(
-                            context,
-                            slot,
-                            |context, slot|{
-                                Deck::new(
-                                    context,
-                                    slot,
-                                )
-                            }
-                        )
-                    },
-                ),
-                games: Map::new(
-                    context,
-                    3,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot
-                        )
-                    },
-                )
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor() {}
 
diff --git a/yarn-project/noir-contracts/contracts/child_contract/src/main.nr b/yarn-project/noir-contracts/contracts/child_contract/src/main.nr
index ee7213c4687..1298819f326 100644
--- a/yarn-project/noir-contracts/contracts/child_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/child_contract/src/main.nr
@@ -17,17 +17,6 @@ contract Child {
         current_value: PublicState<Field>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                current_value: PublicState::new(
-                    context,
-                    1,
-                ),
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor() {}
 
diff --git a/yarn-project/noir-contracts/contracts/counter_contract/src/main.nr b/yarn-project/noir-contracts/contracts/counter_contract/src/main.nr
index e424fbda693..e07b793c403 100644
--- a/yarn-project/noir-contracts/contracts/counter_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/counter_contract/src/main.nr
@@ -27,22 +27,6 @@ contract Counter {
 
     // docs:end:storage_struct
 
-    // docs:start:storage_init
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                counters: Map::new(
-                    context,
-                    1,
-                    |context, slot| {
-                        EasyPrivateUint::new(context, slot)
-                    },
-                ),
-            }
-        }
-    }
-    // docs:end:storage_init
-
     // docs:start:constructor
     #[aztec(private)]
     fn constructor(headstart: u120, owner: AztecAddress) {
diff --git a/yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr b/yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr
index ee28961aefd..cd184bd27a4 100644
--- a/yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/docs_example_contract/src/main.nr
@@ -35,8 +35,10 @@ contract DocsExample {
     };
 
     struct Storage {
-        // Shows how to create a custom struct in Public
+        // Shows how to create a custom struct in PublicState
+        // docs:start:storage-leader-declaration
         leader: PublicState<Leader>,
+        // docs:end:storage-leader-declaration
         // docs:start:storage-singleton-declaration
         legendary_card: Singleton<CardNote>,
         // docs:end:storage-singleton-declaration
@@ -44,20 +46,29 @@ contract DocsExample {
         // docs:start:storage-map-singleton-declaration
         profiles: Map<AztecAddress, Singleton<CardNote>>,
         // docs:end:storage-map-singleton-declaration
-        test: Set<CardNote>,
+        // docs:start:storage-set-declaration
+        set: Set<CardNote>,
+        // docs:end:storage-set-declaration
+        // docs:start:storage-immutable-singleton-declaration
         imm_singleton: ImmutableSingleton<CardNote>,
-        // docs:start:start_vars_stable
+        // docs:end:storage-immutable-singleton-declaration
+        // docs:start:storage-stable-declaration
         stable_value: StablePublicState<Leader>,
-        // docs:end:start_vars_stable
+        // docs:end:storage-stable-declaration
+        // docs:start:storage-minters-declaration
+        minters: Map<AztecAddress, PublicState<bool>>, 
+        // docs:end:storage-minters-declaration
     }
 
     impl Storage {
         fn init(context: Context) -> Self {
             Storage {
+                // docs:start:storage-leader-init
                 leader: PublicState::new(
                     context,
                     1
                 ),
+                // docs:end:storage-leader-init
                 // docs:start:start_vars_singleton
                 legendary_card: Singleton::new(context, 3),
                 // docs:end:start_vars_singleton
@@ -71,9 +82,22 @@ contract DocsExample {
                     },
                 ),
                 // docs:end:state_vars-MapSingleton
-                test: Set::new(context, 5),
+                // docs:start:storage-set-init
+                set: Set::new(context, 5),
+                // docs:end:storage-set-init
                 imm_singleton: ImmutableSingleton::new(context, 6),
+                // docs:start:storage-stable-init
                 stable_value: StablePublicState::new(context, 7),
+                // docs:end:storage-stable-init
+                // docs:start:storage-minters-init
+                minters: Map::new(
+                    context,
+                    8,
+                    |context, slot| {
+                        PublicState::new(context, slot)
+                    },
+                ),
+                // docs:end:storage-minters-init
             }
         }
     }
@@ -100,11 +124,13 @@ contract DocsExample {
         storage.stable_value.read_public()
     }
 
+    // docs:start:initialize-immutable-singleton
     #[aztec(private)]
     fn initialize_immutable_singleton(randomness: Field, points: u8) {
         let mut new_card = CardNote::new(points, randomness, context.msg_sender());
         storage.imm_singleton.initialize(&mut new_card, true);
     }
+    // docs:end:initialize-immutable-singleton
 
     #[aztec(private)]
     // msg_sender() is 0 at deploy time. So created another function
@@ -118,20 +144,20 @@ contract DocsExample {
     fn insert_notes(amounts: [u8; 10]) {
         for i in 0..amounts.len() {
             let mut note = CardNote::new(amounts[i], 1, context.msg_sender());
-            storage.test.insert(&mut note, true);
+            storage.set.insert(&mut note, true);
         }
     }
 
     #[aztec(private)]
     fn insert_note(amount: u8, randomness: Field) {
         let mut note = CardNote::new(amount, randomness, context.msg_sender());
-        storage.test.insert(&mut note, true);
+        storage.set.insert(&mut note, true);
     }
 
     // docs:start:state_vars-NoteGetterOptionsComparatorExampleNoir
     unconstrained fn read_note(amount: Field, comparator: u3) -> pub [Option<CardNote>; 10] {
         let options = NoteViewerOptions::new().select(0, amount, Option::some(comparator));
-        let notes = storage.test.view_notes(options);
+        let notes = storage.set.view_notes(options);
 
         notes
     }
@@ -192,7 +218,13 @@ contract DocsExample {
     }
     // docs:end:singleton_is_initialized
 
+    // docs:start:get_note-immutable-singleton
     unconstrained fn get_imm_card() -> pub CardNote {
+        storage.imm_singleton.get_note()
+    }
+    // docs:end:get_note-immutable-singleton
+
+    unconstrained fn view_imm_card() -> pub CardNote {
         storage.imm_singleton.view_note()
     }
 
diff --git a/yarn-project/noir-contracts/contracts/easy_private_token_contract/src/main.nr b/yarn-project/noir-contracts/contracts/easy_private_token_contract/src/main.nr
index c43aae116ea..26dbdc4ae2f 100644
--- a/yarn-project/noir-contracts/contracts/easy_private_token_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/easy_private_token_contract/src/main.nr
@@ -23,20 +23,6 @@ contract EasyPrivateToken {
         balances: Map<AztecAddress, EasyPrivateUint>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                balances: Map::new(
-                    context,
-                    1,
-                    |context, slot| {
-                        EasyPrivateUint::new(context, slot)
-                    },
-                ),
-            }
-        }
-    }
-
     /**
      * initialize the contract's initial state variables.
      */
diff --git a/yarn-project/noir-contracts/contracts/easy_private_voting_contract/src/main.nr b/yarn-project/noir-contracts/contracts/easy_private_voting_contract/src/main.nr
index 5cc3606d39e..731976633cc 100644
--- a/yarn-project/noir-contracts/contracts/easy_private_voting_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/easy_private_voting_contract/src/main.nr
@@ -16,31 +16,7 @@ contract EasyPrivateVoting {
         voteEnded: PublicState<bool>, // voteEnded is boolean
     }
     // docs:end:storage_struct
-    // docs:start:storage_impl
-    impl Storage {
-	fn init(context: Context) -> Self {
-	    Storage {
-		 admin: PublicState::new(
-		    context, 
-            1, // storage slot. this can be anything except 0. it is hashed, and hash on 0 = 0
-		), 
-        tally: Map::new(
-		    context, 
-            2, 
-            |context, slot| {
-			PublicState::new(
-			    context,
-                slot,
-			)
-		    },
-		 ),
-		  voteEnded: PublicState::new(
-		    context, 
-            3
-		)
-	    } }
-	}
-    // docs:end:storage_impl
+
     // docs:start:constructor
     #[aztec(private)] // annotation to mark function as private and expose private context
     fn constructor(admin: AztecAddress) { // called when contract is deployed
diff --git a/yarn-project/noir-contracts/contracts/ecdsa_account_contract/src/main.nr b/yarn-project/noir-contracts/contracts/ecdsa_account_contract/src/main.nr
index 53c26c8b3c3..d74016d2f2d 100644
--- a/yarn-project/noir-contracts/contracts/ecdsa_account_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/ecdsa_account_contract/src/main.nr
@@ -30,14 +30,6 @@ contract EcdsaAccount {
         public_key: ImmutableSingleton<EcdsaPublicKeyNote>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                public_key: ImmutableSingleton::new(context, 1),
-            }
-        }
-    }
-
     global ACCOUNT_ACTIONS_STORAGE_SLOT = 2;
 
     // Creates a new account out of an ECDSA public key to use for signature verification
@@ -104,7 +96,7 @@ contract EcdsaAccount {
         storage_slot: Field,
         serialized_note: [Field; ECDSA_PUBLIC_KEY_NOTE_LEN]
     ) -> pub [Field; 4] {
-        assert(storage_slot == 1);
+        assert(storage_slot == storage.public_key.get_storage_slot());
         let note_header = NoteHeader::new(contract_address, nonce, storage_slot);
         note_utils::compute_note_hash_and_nullifier(EcdsaPublicKeyNote::deserialize, note_header, serialized_note)
     }
diff --git a/yarn-project/noir-contracts/contracts/escrow_contract/src/main.nr b/yarn-project/noir-contracts/contracts/escrow_contract/src/main.nr
index 0aaf5426fda..b479e6593ca 100644
--- a/yarn-project/noir-contracts/contracts/escrow_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/escrow_contract/src/main.nr
@@ -27,14 +27,6 @@ contract Escrow {
         owners: Set<AddressNote>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                owners: Set::new(context, 1),
-            }
-        }
-    }
-
     // Creates a new instance
     // docs:start:constructor
     #[aztec(private)]
@@ -75,7 +67,7 @@ contract Escrow {
         serialized_note: [Field; ADDRESS_NOTE_LEN]
     ) -> pub [Field; 4] {
         let note_header = NoteHeader::new(contract_address, nonce, storage_slot);
-        assert(storage_slot == 1);
+        assert(storage_slot == storage.owners.get_storage_slot());
         note_utils::compute_note_hash_and_nullifier(AddressNote::deserialize, note_header, serialized_note)
     }
 }
diff --git a/yarn-project/noir-contracts/contracts/inclusion_proofs_contract/src/main.nr b/yarn-project/noir-contracts/contracts/inclusion_proofs_contract/src/main.nr
index 305176fd3df..218a2898d90 100644
--- a/yarn-project/noir-contracts/contracts/inclusion_proofs_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/inclusion_proofs_contract/src/main.nr
@@ -55,28 +55,6 @@ contract InclusionProofs {
         public_unused_value: PublicState<Field>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                private_values: Map::new(
-                    context,
-                    1, // Storage slot
-                    |context, slot| {
-                        Set::new(context, slot)
-                    },
-                ),
-                public_value: PublicState::new(
-                    context,
-                    2, // Storage slot
-                ),
-                public_unused_value: PublicState::new(
-                    context,
-                    3, // Storage slot
-                ),
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor(public_value: Field) {
         let selector = FunctionSelector::from_signature("_initialize(Field)");
diff --git a/yarn-project/noir-contracts/contracts/lending_contract/src/main.nr b/yarn-project/noir-contracts/contracts/lending_contract/src/main.nr
index 9811f0919d1..2dd9bc1c16a 100644
--- a/yarn-project/noir-contracts/contracts/lending_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/lending_contract/src/main.nr
@@ -38,51 +38,6 @@ contract Lending {
         static_debt: Map<AztecAddress, PublicState<Field>>, // abusing keys very heavily
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                collateral_asset: PublicState::new(
-                    context,
-                    1
-                ),
-                stable_coin: PublicState::new(
-                    context,
-                    2
-                ),
-                assets: Map::new(
-                    context,
-                    3,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot
-                        )
-                    },
-                ),
-                collateral: Map::new(
-                    context,
-                    4,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot
-                        )
-                    },
-                ),
-                static_debt: Map::new(
-                    context,
-                    5,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot
-                        )
-                    },
-                ),
-            }
-        }
-    }
-
     struct Position {
         collateral: Field,
         static_debt: Field,
diff --git a/yarn-project/noir-contracts/contracts/pending_commitments_contract/src/main.nr b/yarn-project/noir-contracts/contracts/pending_commitments_contract/src/main.nr
index 3878ae5134c..2c163a14184 100644
--- a/yarn-project/noir-contracts/contracts/pending_commitments_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/pending_commitments_contract/src/main.nr
@@ -29,20 +29,6 @@ contract PendingCommitments {
         balances: Map<AztecAddress, Set<ValueNote>>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                balances: Map::new(
-                    context,
-                    1, // Storage slot
-                    |context, slot| {
-                        Set::new(context, slot)
-                    },
-                ),
-            }
-        }
-    }
-
     // TODO(dbanks12): consolidate code into internal helper functions
     // (once Noir's support for this is more robust)
 
diff --git a/yarn-project/noir-contracts/contracts/price_feed_contract/src/main.nr b/yarn-project/noir-contracts/contracts/price_feed_contract/src/main.nr
index a72caab5e83..b1ee56c8991 100644
--- a/yarn-project/noir-contracts/contracts/price_feed_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/price_feed_contract/src/main.nr
@@ -17,23 +17,6 @@ contract PriceFeed {
         assets: Map<Field, PublicState<Asset>>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> pub Self {
-            Storage {
-                assets: Map::new(
-                    context,
-                    1,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot,
-                        )
-                    },
-                ),
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor() {}
 
diff --git a/yarn-project/noir-contracts/contracts/schnorr_account_contract/src/main.nr b/yarn-project/noir-contracts/contracts/schnorr_account_contract/src/main.nr
index ef0402fcc01..c412146f82e 100644
--- a/yarn-project/noir-contracts/contracts/schnorr_account_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/schnorr_account_contract/src/main.nr
@@ -28,16 +28,6 @@ contract SchnorrAccount {
         // docs:end:storage
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                // docs:start:storage_init
-                signing_public_key: ImmutableSingleton::new(context, 1),
-                // docs:end:storage_init
-            }
-        }
-    }
-
     global ACCOUNT_ACTIONS_STORAGE_SLOT = 2;
 
     // Constructs the contract
@@ -111,7 +101,7 @@ contract SchnorrAccount {
         storage_slot: Field,
         serialized_note: [Field; PUBLIC_KEY_NOTE_LEN]
     ) -> pub [Field; 4] {
-        assert(storage_slot == 1);
+        assert(storage_slot == storage.signing_public_key.get_storage_slot());
         let note_header = NoteHeader::new(contract_address, nonce, storage_slot);
         note_utils::compute_note_hash_and_nullifier(PublicKeyNote::deserialize, note_header, serialized_note)
     }
diff --git a/yarn-project/noir-contracts/contracts/slow_tree_contract/src/main.nr b/yarn-project/noir-contracts/contracts/slow_tree_contract/src/main.nr
index e0986fbd238..ca25b3c193b 100644
--- a/yarn-project/noir-contracts/contracts/slow_tree_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/slow_tree_contract/src/main.nr
@@ -46,23 +46,6 @@ contract SlowTree {
     }
     // docs:end:constants_and_storage
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                trees: Map::new(
-                    context,
-                    1,
-                    |context, slot| {
-                        SlowMap::new(
-                            context,
-                            slot,
-                        )
-                    }
-                ),
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor() {}
     // docs:start:initialize
diff --git a/yarn-project/noir-contracts/contracts/stateful_test_contract/src/main.nr b/yarn-project/noir-contracts/contracts/stateful_test_contract/src/main.nr
index b1e22e2eb42..96932ffa102 100644
--- a/yarn-project/noir-contracts/contracts/stateful_test_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/stateful_test_contract/src/main.nr
@@ -24,30 +24,6 @@ contract StatefulTest {
         public_values: Map<AztecAddress, PublicState<Field>>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                notes: Map::new(
-                    context,
-                    1, // Storage slot
-                    |context, slot| {
-                        Set::new(context, slot)
-                    },
-                ),
-                public_values: Map::new(
-                    context,
-                    2,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot,
-                        )
-                    },
-                ),
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor(owner: AztecAddress, value: Field) {
         let selector = FunctionSelector::from_signature("create_note((Field),Field)");
diff --git a/yarn-project/noir-contracts/contracts/test_contract/src/main.nr b/yarn-project/noir-contracts/contracts/test_contract/src/main.nr
index 791016b7a8d..d1a18eea77b 100644
--- a/yarn-project/noir-contracts/contracts/test_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/test_contract/src/main.nr
@@ -55,14 +55,6 @@ contract Test {
         example_constant: ImmutableSingleton<FieldNote>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                example_constant: ImmutableSingleton::new(context, 1),
-            }
-        }
-    }
-
     #[aztec(private)]
     // docs:start:empty-constructor
     fn constructor() {}
@@ -333,7 +325,7 @@ contract Test {
         storage_slot: Field,
         serialized_note: [Field; VALUE_NOTE_LEN] // must fit either a FieldNote or a ValueNote
     ) -> pub [Field; 4] {
-        if (storage_slot == 1) {
+        if (storage_slot == storage.example_constant.get_storage_slot()) {
             let note_header = NoteHeader::new(contract_address, nonce, storage_slot);
             note_utils::compute_note_hash_and_nullifier(FieldNote::deserialize, note_header, serialized_note)
         } else {
diff --git a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr
index b8b02f80817..f6558115fc3 100644
--- a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/main.nr
@@ -52,50 +52,15 @@ contract TokenBlacklist {
     // docs:end:interface
 
     struct Storage {
-        admin: PublicState<AztecAddress>,
-        balances: BalancesMap,
-        total_supply: PublicState<SafeU120>,
-        pending_shields: Set<TransparentNote>,
-        public_balances: Map<AztecAddress, PublicState<SafeU120>>,
+        admin: PublicState<AztecAddress>, 
+        balances: BalancesMap<TokenNote>, 
+        total_supply: PublicState<SafeU120>, 
+        pending_shields: Set<TransparentNote>, 
+        public_balances: Map<AztecAddress, PublicState<SafeU120>>, 
         slow_update: ImmutableSingleton<FieldNote>,
         public_slow_update: PublicState<AztecAddress>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> pub Self {
-            Storage {
-                admin: PublicState::new(
-                    context,
-                    1,
-                ),
-                balances: BalancesMap::new(context, 3),
-                total_supply: PublicState::new(
-                    context,
-                    4,
-                ),
-                pending_shields: Set::new(context, 5),
-                public_balances: Map::new(
-                    context,
-                    6,
-                    |context, slot| {
-                        PublicState::new(
-                            context,
-                            slot
-                        )
-                    },
-                ),
-                // Below is an abomination to have same value in private and public (immutable in solidity).
-                // docs:start:slow_updates_storage
-                slow_update: ImmutableSingleton::new(context, 7),
-                public_slow_update: PublicState::new(
-                    context,
-                    8
-                ),
-                // docs:end:slow_updates_storage
-
-            }
-        }
-    }
     // docs:start:constructor
     #[aztec(private)]
     fn constructor(admin: AztecAddress, slow_updates_contract: AztecAddress) {
@@ -278,7 +243,7 @@ contract TokenBlacklist {
         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));
     }
 
     #[aztec(private)]
@@ -295,7 +260,7 @@ contract TokenBlacklist {
             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)");
         context.call_public_function(context.this_address(), selector, [to.to_field(), amount]);
@@ -318,8 +283,8 @@ contract TokenBlacklist {
         }
 
         let amount = SafeU120::new(amount);
-        storage.balances.at(from).sub(amount);
-        storage.balances.at(to).add(amount);
+        storage.balances.sub(from, amount);
+        storage.balances.add(to, amount);
     }
 
     #[aztec(private)]
@@ -334,7 +299,7 @@ contract TokenBlacklist {
             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)");
         context.call_public_function(context.this_address(), selector, [amount]);
@@ -362,7 +327,7 @@ contract TokenBlacklist {
     }
 
     unconstrained fn balance_of_private(owner: AztecAddress) -> pub u120 {
-        storage.balances.at(owner).balance_of().value
+        storage.balances.balance_of(owner).value
     }
 
     unconstrained fn balance_of_public(owner: AztecAddress) -> pub u120 {
@@ -382,9 +347,9 @@ contract TokenBlacklist {
         preimage: [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, preimage)
-        } else if (storage_slot == 7) {
+        } else if (storage_slot == storage.slow_update.get_storage_slot()) {
             note_utils::compute_note_hash_and_nullifier(FieldNote::deserialize, note_header, preimage)
         } else {
             note_utils::compute_note_hash_and_nullifier(TokenNote::deserialize, note_header, preimage)
diff --git a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types.nr b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types.nr
index 3ba03658390..149d823a339 100644
--- a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types.nr
+++ b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types.nr
@@ -1,5 +1,4 @@
 mod transparent_note;
-mod balance_set;
 mod balances_map;
 mod token_note;
 mod roles;
diff --git a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balance_set.nr b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balance_set.nr
deleted file mode 100644
index a6952f87c1d..00000000000
--- a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balance_set.nr
+++ /dev/null
@@ -1,121 +0,0 @@
-use dep::std::option::Option;
-use dep::safe_math::SafeU120;
-use dep::aztec::protocol_types::constants::MAX_READ_REQUESTS_PER_CALL;
-use dep::aztec::protocol_types::address::AztecAddress;
-use dep::aztec::{
-  context::Context,
-  state_vars::set::Set,
-};
-use dep::aztec::note::{
-    note_getter::view_notes,
-    note_getter_options::{NoteGetterOptions, SortOrder},
-    note_viewer_options::NoteViewerOptions
-};
-use dep::aztec::note::{
-    note_header::NoteHeader,
-    note_interface::NoteInterface,
-};
-use crate::types::token_note::TokenNote;
-
-// A set implementing standard manipulation of balances.
-// 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>
-}
-
-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,
-    };
-    Self {
-        context,
-        owner,
-        set,
-    }
-  }
-
-  unconstrained pub fn balance_of(self: Self) -> SafeU120 {
-    self.balance_of_with_offset(0)
-  }
-
-  unconstrained pub fn balance_of_with_offset(self: Self, offset: u32) -> SafeU120 {
-      // 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.set.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().amount);
-          }
-      }
-      if (opt_notes[len - 1].is_some()) {
-          balance = balance.add(self.balance_of_with_offset(offset + opt_notes.len() as u32));
-      }
-
-      balance
-  }
-
-    pub fn add(self: Self, addend: SafeU120) {
-        let mut addend_note = TokenNote::new(addend, self.owner);
-
-        // docs:start:insert
-        self.set.insert(&mut addend_note, true);
-        // docs:end:insert
-    }
-
-    pub fn sub(self: Self, subtrahend: SafeU120) {
-        // docs:start:get_notes
-        let options = NoteGetterOptions::with_filter(filter_notes_min_sum, subtrahend);
-        let maybe_notes = self.set.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.set.remove(note);
-                // docs:end:remove
-
-                minuend = minuend.add(note.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(minuend.sub(subtrahend));
-    }
-}
-
-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() {
-        if notes[i].is_some() & sum.lt(min_sum) {
-            let note = notes[i].unwrap_unchecked();
-            selected[i] = Option::some(note);
-            sum = sum.add(note.amount);
-        }
-    }
-    selected
-}
diff --git a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balances_map.nr b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balances_map.nr
index ed3c8c3fd93..a1e41d16bcf 100644
--- a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balances_map.nr
+++ b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/balances_map.nr
@@ -1,15 +1,33 @@
-use dep::aztec::protocol_types::address::AztecAddress;
-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::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 {
+struct BalancesMap<T> {
     context: Context,
-    storage_slot: Field,
+    map: Map<AztecAddress, Set<T>>
 }
 
-impl BalancesMap {
+impl<T> BalancesMap<T> {
     pub fn new(
         context: Context,
         storage_slot: Field,
@@ -17,12 +35,87 @@ impl BalancesMap {
         assert(storage_slot != 0, "Storage slot 0 not allowed. Storage slots must start from 1.");
         Self {
             context,
-            storage_slot,
+            map: Map::new(context, storage_slot, |context, slot| Set::new(context, slot))
         }
     }
 
-    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)
+    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 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
 }
diff --git a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/token_note.nr b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/token_note.nr
index e405a600131..efd2a4b857d 100644
--- a/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/token_note.nr
+++ b/yarn-project/noir-contracts/contracts/token_blacklist_contract/src/types/token_note.nr
@@ -1,5 +1,9 @@
-use dep::aztec::protocol_types::constants::MAX_READ_REQUESTS_PER_CALL;
-use dep::aztec::protocol_types::address::AztecAddress;
+use dep::aztec::protocol_types::{
+    address::AztecAddress,
+    constants::{
+        MAX_READ_REQUESTS_PER_CALL
+    },
+};
 use dep::aztec::{
     note::{
         note_header::NoteHeader,
@@ -10,7 +14,10 @@ use dep::aztec::{
     state_vars::set::Set,
     log::emit_encrypted_log,
     hash::pedersen_hash,
-    protocol_types::traits::{Serialize, Deserialize}
+    protocol_types::traits::{
+        Serialize,
+        Deserialize
+    },
 };
 use dep::aztec::oracle::{
     rand::rand,
@@ -20,6 +27,12 @@ use dep::aztec::oracle::{
 use dep::safe_math::SafeU120;
 use dep::std::option::Option;
 
+trait OwnedNote {
+    fn new(amount: SafeU120, owner: AztecAddress) -> Self;
+    fn get_amount(self) -> SafeU120;
+    fn get_owner(self) -> AztecAddress;
+}
+
 global TOKEN_NOTE_LEN: Field = 3; // 3 plus a header.
 
 struct TokenNote {
@@ -43,11 +56,11 @@ impl Serialize<TOKEN_NOTE_LEN> for TokenNote {
 }
 
 impl Deserialize<TOKEN_NOTE_LEN> for TokenNote {
-    fn deserialize(preimage: [Field; TOKEN_NOTE_LEN]) -> Self {
+    fn deserialize(serialized_note: [Field; TOKEN_NOTE_LEN]) -> Self {
         Self {
-            amount: SafeU120::new(preimage[0]),
-            owner: AztecAddress::from_field(preimage[1]),
-            randomness: preimage[2],
+            amount: SafeU120::new(serialized_note[0]),
+            owner: AztecAddress::from_field(serialized_note[1]),
+            randomness: serialized_note[2],
             header: NoteHeader::empty(),
         }
     }
@@ -87,7 +100,6 @@ impl NoteInterface for TokenNote {
         self.header = header;
     }
 
-
     fn get_header(note: TokenNote) -> NoteHeader {
         note.header
     }
@@ -108,8 +120,8 @@ impl NoteInterface for TokenNote {
   }
 }
 
-impl TokenNote {
-    pub fn new(amount: SafeU120, owner: AztecAddress) -> Self {
+impl OwnedNote for TokenNote {
+    fn new(amount: SafeU120, owner: AztecAddress) -> Self {
         Self {
             amount,
             owner,
@@ -117,4 +129,13 @@ impl TokenNote {
             header: NoteHeader::empty(),
         }
     }
+
+    fn get_amount(self) -> SafeU120 {
+        self.amount
+    }
+
+    fn get_owner(self) -> AztecAddress {
+        self.owner
+    }
+   
 }
diff --git a/yarn-project/noir-contracts/contracts/token_bridge_contract/src/main.nr b/yarn-project/noir-contracts/contracts/token_bridge_contract/src/main.nr
index e7893ef8ceb..489c95cb26c 100644
--- a/yarn-project/noir-contracts/contracts/token_bridge_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/token_bridge_contract/src/main.nr
@@ -32,17 +32,6 @@ contract TokenBridge {
         token: PublicState<AztecAddress>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> pub Self {
-            Storage {
-                token: PublicState::new(
-                    context,
-                    1
-                ),
-            }
-        }
-    }
-
     // Constructs the contract.
     #[aztec(private)]
     fn constructor(token: AztecAddress) {
diff --git a/yarn-project/noir-contracts/contracts/token_contract/src/main.nr b/yarn-project/noir-contracts/contracts/token_contract/src/main.nr
index 9815719c26f..aa6ad831fcb 100644
--- a/yarn-project/noir-contracts/contracts/token_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/token_contract/src/main.nr
@@ -13,8 +13,8 @@ contract Token {
     // Libs
     use dep::std::option::Option;
 
-    use dep::safe_math::{SafeU120, SAFE_U120_SERIALIZED_LEN};
-    use dep::compressed_string::{FieldCompressedString};
+    use dep::safe_math::SafeU120;
+    use dep::compressed_string::FieldCompressedString;
 
     use dep::aztec::{
         note::{
@@ -26,12 +26,6 @@ contract Token {
         hash::{compute_secret_hash},
         state_vars::{map::Map, public_state::PublicState, stable_public_state::StablePublicState,  set::Set},
         protocol_types::{
-            type_serialization::{
-                FIELD_SERIALIZED_LEN, 
-                BOOL_SERIALIZED_LEN, 
-                U8_SERIALIZED_LEN,
-                AZTEC_ADDRESS_SERIALIZED_LEN
-            },
             abis::function_selector::FunctionSelector,
             address::AztecAddress
         }
@@ -49,7 +43,7 @@ contract Token {
     use crate::types::{
         transparent_note::TransparentNote,
         token_note::{TokenNote, TOKEN_NOTE_LEN},
-        balances_map::{BalancesMap}
+        balances_map::BalancesMap
     };
     // docs:end::imports
 
@@ -59,14 +53,14 @@ contract Token {
         admin: PublicState<AztecAddress>,
         // docs:end:storage_admin
         // docs:start:storage_minters
-        minters: Map<AztecAddress, PublicState<bool>>,
+        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>,
+        total_supply: PublicState<SafeU120>, 
         // docs:start:storage_pending_shields
-        pending_shields: Set<TransparentNote>,
+        pending_shields: Set<TransparentNote>, 
         // docs:end:storage_pending_shields
         public_balances: Map<AztecAddress, PublicState<SafeU120>>,
         symbol: StablePublicState<FieldCompressedString>,
@@ -77,67 +71,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,
-                        )
-                    },
-                ),
-                symbol: StablePublicState::new(
-                    context,
-                    7,
-                ),
-                name: StablePublicState::new(
-                    context,
-                    8,
-                ),
-                // docs:start:storage_decimals_init
-                decimals: StablePublicState::new(
-                    context,
-                    9,
-                ),
-                // docs:end:storage_decimals_init
-            }
-        }
-    }
-    // docs:end:storage_init
-
     // docs:start:constructor
     #[aztec(private)]
     fn constructor(admin: AztecAddress, name: str<31>, symbol: str<31>, decimals: u8) {
@@ -323,7 +256,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
 
@@ -336,7 +269,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]);
@@ -355,9 +288,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
@@ -371,7 +304,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]);
@@ -438,7 +371,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
 
@@ -462,7 +395,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)
diff --git a/yarn-project/noir-contracts/contracts/token_contract/src/types.nr b/yarn-project/noir-contracts/contracts/token_contract/src/types.nr
index d3b3b1c9e77..62a1bb2a363 100644
--- a/yarn-project/noir-contracts/contracts/token_contract/src/types.nr
+++ b/yarn-project/noir-contracts/contracts/token_contract/src/types.nr
@@ -1,4 +1,3 @@
 mod transparent_note;
-mod balance_set;
 mod balances_map;
 mod token_note;
diff --git a/yarn-project/noir-contracts/contracts/token_contract/src/types/balance_set.nr b/yarn-project/noir-contracts/contracts/token_contract/src/types/balance_set.nr
deleted file mode 100644
index e84b4cfd110..00000000000
--- a/yarn-project/noir-contracts/contracts/token_contract/src/types/balance_set.nr
+++ /dev/null
@@ -1,123 +0,0 @@
-use dep::std::option::Option;
-use dep::safe_math::SafeU120;
-use dep::aztec::protocol_types::{
-    constants::MAX_READ_REQUESTS_PER_CALL,
-    address::AztecAddress,
-};
-use dep::aztec::{
-  context::Context,
-  state_vars::set::Set,
-};
-use dep::aztec::note::{
-    note_getter::view_notes,
-    note_getter_options::{NoteGetterOptions, SortOrder},
-    note_viewer_options::NoteViewerOptions
-};
-use dep::aztec::note::{
-    note_header::NoteHeader,
-    note_interface::NoteInterface
-};
-use crate::types::token_note::TokenNote;
-
-// A set implementing standard manipulation of balances.
-// 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>
-}
-
-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,
-    };
-    Self {
-        context,
-        owner,
-        set,
-    }
-  }
-
-  unconstrained pub fn balance_of(self: Self) -> SafeU120 {
-    self.balance_of_with_offset(0)
-  }
-
-  unconstrained pub fn balance_of_with_offset(self: Self, offset: u32) -> SafeU120 {
-      // 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.set.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().amount);
-          }
-      }
-      if (opt_notes[len - 1].is_some()) {
-          balance = balance.add(self.balance_of_with_offset(offset + opt_notes.len() as u32));
-      }
-
-      balance
-  }
-
-    pub fn add(self: Self, addend: SafeU120) {
-        let mut addend_note = TokenNote::new(addend, self.owner);
-
-        // docs:start:insert
-        self.set.insert(&mut addend_note, true);
-        // docs:end:insert
-    }
-
-    pub fn sub(self: Self, subtrahend: SafeU120) {
-        // docs:start:get_notes
-        let options = NoteGetterOptions::with_filter(filter_notes_min_sum, subtrahend);
-        let maybe_notes = self.set.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.set.remove(note);
-                // docs:end:remove
-
-                minuend = minuend.add(note.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(minuend.sub(subtrahend));
-    }
-}
-
-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() {
-        if notes[i].is_some() & sum.lt(min_sum) {
-            let note = notes[i].unwrap_unchecked();
-            selected[i] = Option::some(note);
-            sum = sum.add(note.amount);
-        }
-    }
-    selected
-}
diff --git a/yarn-project/noir-contracts/contracts/token_contract/src/types/balances_map.nr b/yarn-project/noir-contracts/contracts/token_contract/src/types/balances_map.nr
index c8a43f11d55..2f5835248a0 100644
--- a/yarn-project/noir-contracts/contracts/token_contract/src/types/balances_map.nr
+++ b/yarn-project/noir-contracts/contracts/token_contract/src/types/balances_map.nr
@@ -1,28 +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 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 {
-    context: Context,
-    storage_slot: Field,
+struct BalancesMap<T> {
+    map: Map<AztecAddress, Set<T>>
 }
 
-impl BalancesMap {
+impl<T> BalancesMap<T> {
     pub fn new(
         context: Context,
         storage_slot: Field,
     ) -> Self {
         assert(storage_slot != 0, "Storage slot 0 not allowed. Storage slots must start from 1.");
         Self {
-            context,
-            storage_slot,
+            map: Map::new(context, storage_slot, |context, slot| Set::new(context, slot))
         }
     }
 
-    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)
+    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 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
 }
diff --git a/yarn-project/noir-contracts/contracts/token_contract/src/types/token_note.nr b/yarn-project/noir-contracts/contracts/token_contract/src/types/token_note.nr
index 17e77155a28..214f942cec8 100644
--- a/yarn-project/noir-contracts/contracts/token_contract/src/types/token_note.nr
+++ b/yarn-project/noir-contracts/contracts/token_contract/src/types/token_note.nr
@@ -27,6 +27,12 @@ use dep::aztec::oracle::{
 use dep::safe_math::SafeU120;
 use dep::std::option::Option;
 
+trait OwnedNote {
+    fn new(amount: SafeU120, owner: AztecAddress) -> Self;
+    fn get_amount(self) -> SafeU120;
+    fn get_owner(self) -> AztecAddress;
+}
+
 global TOKEN_NOTE_LEN: Field = 3; // 3 plus a header.
 
 struct TokenNote {
@@ -114,8 +120,8 @@ impl NoteInterface for TokenNote {
   }
 }
 
-impl TokenNote {
-    pub fn new(amount: SafeU120, owner: AztecAddress) -> Self {
+impl OwnedNote for TokenNote {
+    fn new(amount: SafeU120, owner: AztecAddress) -> Self {
         Self {
             amount,
             owner,
@@ -124,5 +130,12 @@ impl TokenNote {
         }
     }
 
+    fn get_amount(self) -> SafeU120 {
+        self.amount
+    }
+
+    fn get_owner(self) -> AztecAddress {
+        self.owner
+    }
    
 }
diff --git a/yarn-project/noir-contracts/contracts/uniswap_contract/src/main.nr b/yarn-project/noir-contracts/contracts/uniswap_contract/src/main.nr
index c06630b6f92..65f9a8ea590 100644
--- a/yarn-project/noir-contracts/contracts/uniswap_contract/src/main.nr
+++ b/yarn-project/noir-contracts/contracts/uniswap_contract/src/main.nr
@@ -32,21 +32,6 @@ contract Uniswap {
         nonce_for_burn_approval: PublicState<Field>,
     }
 
-    impl Storage {
-        fn init(context: Context) -> Self {
-            Storage {
-                approved_action: Map::new(
-                    context,
-                    1,
-                    |context, slot| {
-                        PublicState::new(context, slot)
-                    },
-                ),
-                nonce_for_burn_approval: PublicState::new(context, 2),
-            }
-        }
-    }
-
     #[aztec(private)]
     fn constructor() {}
     // docs:end:uniswap_setup