TopLevelStorage dependency injection

examples/erc20/src/erc20.rs

@@ -14,7 +14,7 @@ use alloy_sol_types::sol;
 use core::marker::PhantomData;
 use stylus_sdk::{evm, msg, prelude::*};
 
-pub trait Erc20Params {
+pub trait Erc20Params: 'static {
     /// Immutable token name
     const NAME: &'static str;
 
@@ -27,7 +27,7 @@ pub trait Erc20Params {
 
 sol_storage! {
     /// Erc20 implements all ERC-20 methods.
-    pub struct Erc20<T> {
+    pub struct Erc20<T: Erc20Params> {
         /// Maps users to balances
         mapping(address => uint256) balances;
         /// Maps users to a mapping of each spender's allowance

stylus-proc/src/methods/external.rs

@@ -123,7 +123,7 @@ pub fn external(_attr: TokenStream, input: TokenStream) -> TokenStream {
         let storage = if needed_purity == Pure {
             quote!()
         } else if has_self {
-            quote! { core::borrow::BorrowMut::borrow_mut(storage), }
+            quote! { storage.inner_mut(), }
         } else {
             quote! { storage, }
         };
@@ -241,13 +241,6 @@ pub fn external(_attr: TokenStream, input: TokenStream) -> TokenStream {
         }
     });
 
-    // ensure we can actually borrow the things we inherit
-    let borrow_clauses = inherits.iter().map(|ty| {
-        quote! {
-            S: core::borrow::BorrowMut<#ty>
-        }
-    });
-
     let self_ty = &input.self_ty;
     let generic_params = &input.generics.params;
     let where_clauses = input
@@ -263,13 +256,9 @@ pub fn external(_attr: TokenStream, input: TokenStream) -> TokenStream {
 
         impl<S, #generic_params> stylus_sdk::abi::Router<S> for #self_ty
         where
-            S: stylus_sdk::storage::TopLevelStorage + core::borrow::BorrowMut<Self>,
-            #(#borrow_clauses,)*
+            S: stylus_sdk::storage::TopLevelStorage,
             #where_clauses
         {
-            // TODO: this should be configurable
-            type Storage = Self;
-
             #[inline(always)]
             fn route(storage: &mut S, selector: u32, input: &[u8]) -> Option<stylus_sdk::ArbResult> {
                 use stylus_sdk::{function_selector, alloy_sol_types::SolType};

stylus-proc/src/storage/mod.rs

@@ -18,13 +18,20 @@ pub fn solidity_storage(_attr: TokenStream, input: TokenStream) -> TokenStream {
     let mut init = quote! {};
     let mut size = quote! {};
     let mut borrows = quote! {};
+    let mut inner_storage_accessors = vec![];
 
     for (field_index, field) in input.fields.iter_mut().enumerate() {
         // deny complex types
         let Type::Path(ty) = &field.ty else {
             error!(&field, "Type not supported for EVM state storage");
         };
 
+        let accessor = match field.ident.as_ref() {
+            Some(accessor) => accessor.into_token_stream(),
+            None => Index::from(field_index).into_token_stream(),
+        };
+        inner_storage_accessors.push(accessor.clone());
+
         // implement borrows (TODO: use drain_filter when stable)
         let attrs = mem::take(&mut field.attrs);
         for attr in attrs {
@@ -36,10 +43,7 @@ pub fn solidity_storage(_attr: TokenStream, input: TokenStream) -> TokenStream {
                 error!(attr.tokens, "borrow attribute does not take parameters");
             }
             let ty = &field.ty;
-            let accessor = match field.ident.as_ref() {
-                Some(accessor) => accessor.into_token_stream(),
-                None => Index::from(field_index).into_token_stream(),
-            };
+
             borrows.extend(quote! {
                 impl core::borrow::Borrow<#ty> for #name {
                     fn borrow(&self) -> &#ty {
@@ -115,6 +119,43 @@ pub fn solidity_storage(_attr: TokenStream, input: TokenStream) -> TokenStream {
             }
         });
     }
+
+    let inner_storage_calls = inner_storage_accessors.iter().map(|accessor|{
+        quote! {
+            .or(self.#accessor.try_get_inner())
+        }
+    });    
+
+    let inner_mut_storage_calls = inner_storage_accessors.iter().map(|accessor|{
+        quote! {
+            .or(self.#accessor.try_get_inner_mut())
+        }
+    });
+
+    let storage_level_impl = quote! {
+        unsafe impl #impl_generics stylus_sdk::storage::StorageLevel for #name #ty_generics #where_clause {
+
+            unsafe fn try_get_inner<S: stylus_sdk::storage::StorageLevel + 'static>(&self) -> Option<&S> {
+                use core::any::TypeId;
+                use stylus_sdk::storage::StorageLevel;
+                if TypeId::of::<S>() == TypeId::of::<Self>() {
+                    Some(unsafe { &*(self as *const Self as *const S) })
+                } else {
+                    None #(#inner_storage_calls)*
+                }
+            }
+
+            unsafe fn try_get_inner_mut<S: stylus_sdk::storage::StorageLevel + 'static>(&mut self) -> Option<&mut S> {
+                use core::any::TypeId;
+                use stylus_sdk::storage::StorageLevel;
+                if TypeId::of::<S>() == TypeId::of::<Self>() {
+                    Some(unsafe { &mut *(self as *mut Self as *mut S) })
+                } else {
+                    None #(#inner_mut_storage_calls)*
+                }
+            }
+        }
+    };
 
     // TODO: add mechanism for struct assignment
     let expanded = quote! {
@@ -163,6 +204,8 @@ pub fn solidity_storage(_attr: TokenStream, input: TokenStream) -> TokenStream {
         }
 
         #borrows
+
+        #storage_level_impl
     };
     expanded.into()
 }

stylus-sdk/src/abi/mod.rs

@@ -41,11 +41,8 @@ pub mod internal;
 /// Composition with other routers is possible via `#[inherit]`.
 pub trait Router<S>
 where
-    S: TopLevelStorage + BorrowMut<Self::Storage>,
+    S: TopLevelStorage
 {
-    /// The type the [`TopLevelStorage`] borrows into. Usually just `Self`.
-    type Storage;
-
     /// Tries to find and execute a method for the given selector, returning `None` if none is found.
     /// Routes add via `#[inherit]` will only execute if no match is found among `Self`.
     /// This means that it is possible to override a method by redefining it in `Self`.

stylus-sdk/src/storage/array.rs

@@ -1,7 +1,7 @@
 // Copyright 2023, Offchain Labs, Inc.
 // For licensing, see https://github.com/OffchainLabs/stylus-sdk-rs/blob/stylus/licenses/COPYRIGHT.md
 
-use super::{Erase, StorageGuard, StorageGuardMut, StorageType};
+use super::{Erase, StorageGuard, StorageGuardMut, StorageLevel, StorageType};
 use alloy_primitives::U256;
 use core::marker::PhantomData;
 
@@ -11,6 +11,8 @@ pub struct StorageArray<S: StorageType, const N: usize> {
     marker: PhantomData<S>,
 }
 
+unsafe impl<S: StorageType, const N: usize> StorageLevel for StorageArray<S, N> {}
+
 impl<S: StorageType, const N: usize> StorageType for StorageArray<S, N> {
     type Wraps<'a> = StorageGuard<'a, StorageArray<S, N>> where Self: 'a;
     type WrapsMut<'a> = StorageGuardMut<'a, StorageArray<S, N>> where Self: 'a;

stylus-sdk/src/storage/bytes.rs

@@ -1,7 +1,7 @@
 // Copyright 2022-2023, Offchain Labs, Inc.
 // For licensing, see https://github.com/OffchainLabs/stylus-sdk-rs/blob/stylus/licenses/COPYRIGHT.md
 
-use super::{Erase, GlobalStorage, Storage, StorageB8, StorageGuard, StorageGuardMut, StorageType};
+use super::{Erase, GlobalStorage, Storage, StorageArray, StorageB8, StorageGuard, StorageGuardMut, StorageLevel, StorageType};
 use crate::crypto;
 use alloc::{
     string::{String, ToString},
@@ -16,6 +16,8 @@ pub struct StorageBytes {
     base: OnceCell<U256>,
 }
 
+unsafe impl StorageLevel for StorageBytes {}
+
 impl StorageType for StorageBytes {
     type Wraps<'a> = StorageGuard<'a, StorageBytes> where Self: 'a;
     type WrapsMut<'a> = StorageGuardMut<'a, StorageBytes> where Self: 'a;
@@ -264,6 +266,8 @@ impl<'a> Extend<&'a u8> for StorageBytes {
 /// Accessor for storage-backed bytes
 pub struct StorageString(pub StorageBytes);
 
+unsafe impl StorageLevel for StorageString {}
+
 impl StorageType for StorageString {
     type Wraps<'a> = StorageGuard<'a, StorageString> where Self: 'a;
     type WrapsMut<'a> = StorageGuardMut<'a, StorageString> where Self: 'a;

stylus-sdk/src/storage/map.rs

@@ -3,7 +3,7 @@
 
 use crate::crypto;
 
-use super::{Erase, SimpleStorageType, StorageGuard, StorageGuardMut, StorageType};
+use super::{Erase, SimpleStorageType, StorageBytes, StorageGuard, StorageGuardMut, StorageLevel, StorageType};
 use alloc::{string::String, vec::Vec};
 use alloy_primitives::{Address, FixedBytes, Signed, Uint, B256, U160, U256};
 use core::marker::PhantomData;
@@ -14,6 +14,8 @@ pub struct StorageMap<K: StorageKey, V: StorageType> {
     marker: PhantomData<(K, V)>,
 }
 
+unsafe impl<K: StorageKey, S: StorageType> StorageLevel for StorageMap<K, S> {}
+
 impl<K, V> StorageType for StorageMap<K, V>
 where
     K: StorageKey,

stylus-sdk/src/storage/mod.rs

@@ -32,7 +32,7 @@ pub use bytes::{StorageBytes, StorageString};
 pub use map::{StorageKey, StorageMap};
 pub use traits::{
     Erase, GlobalStorage, SimpleStorageType, StorageGuard, StorageGuardMut, StorageType,
-    TopLevelStorage,
+    TopLevelStorage, StorageLevel,
 };
 pub use vec::StorageVec;
 
@@ -647,6 +647,8 @@ impl From<StorageBlockHash> for BlockHash {
     }
 }
 
+unsafe impl<T> StorageLevel for PhantomData<T> {}
+
 /// We implement `StorageType` for `PhantomData` so that storage types can be generic.
 impl<T> StorageType for PhantomData<T> {
     type Wraps<'a> = Self where Self: 'a;

stylus-sdk/src/storage/traits.rs

@@ -95,7 +95,63 @@ where
 /// # Safety
 ///
 /// The type must be top-level to prevent storage aliasing.
-pub unsafe trait TopLevelStorage {}
+pub unsafe trait TopLevelStorage : StorageLevel {
+
+    /// Retrieve immutable reference to inner storage of type [`S`] or panic.
+    fn inner<S: StorageLevel + 'static>(&self) -> &S {
+        unsafe {
+            self.try_get_inner().unwrap_or_else(|| {
+                panic!(
+                    "type does not exist inside TopLevelStorage - type is {}",
+                    core::any::type_name::<S>()
+                )})
+        }
+    }
+
+    /// Retrieve mutable reference to inner storage of type [`S`] or panic.
+    fn inner_mut<S: StorageLevel + 'static>(&mut self) -> &mut S {
+        unsafe {
+            self.try_get_inner_mut().unwrap_or_else(|| {
+                panic!(
+                    "type does not exist inside TopLevelStorage - type is {}",
+                    core::any::type_name::<S>()
+                )})
+        }
+    }
+}
+
+/// Trait for all-level storage types, usually implemented by proc macros.
+///
+/// # Safety
+///
+/// For simple types like (StorageMap, StorageBool, ..) should have default implementation.
+/// Since it is pointless to querry for a type that can exists in many contracts.
+pub unsafe trait StorageLevel {
+
+    /// Try to retrieve immutable reference to inner laying type [`S`].
+    /// [`Option::None`] result usually means we should panic.
+    ///
+    /// # Safety
+    ///
+    /// To be able to retrieve type that contain current type (parrent) you should
+    /// call [`TopLevelStorage::inner`].
+    unsafe fn try_get_inner<S: StorageLevel + 'static>(&self) -> Option<&S>{
+        None
+    }
+
+    /// Try to retrieve mutable reference to inner laying type [`S`].
+    /// [`Option::None`] result usually means we should panic.
+    ///
+    /// # Safety
+    ///
+    /// To be able to retrieve type that contain current type (parrent) you should
+    /// call [`TopLevelStorage::inner_mut`].
+    unsafe fn try_get_inner_mut<S: StorageLevel + 'static>(&mut self) -> Option<&mut S>{
+        None
+    }
+}
+
+unsafe impl<'a, T: SimpleStorageType<'a>> StorageLevel for T {}
 
 /// Binds a storage accessor to a lifetime to prevent aliasing.
 /// Because this type doesn't implement `DerefMut`, mutable methods on the accessor aren't available.

stylus-sdk/src/storage/vec.rs

@@ -1,9 +1,7 @@
 // Copyright 2023, Offchain Labs, Inc.
 // For licensing, see https://github.com/OffchainLabs/stylus-sdk-rs/blob/stylus/licenses/COPYRIGHT.md
 
-use super::{
-    Erase, GlobalStorage, SimpleStorageType, Storage, StorageGuard, StorageGuardMut, StorageType,
-};
+use super::{Erase, GlobalStorage, SimpleStorageType, Storage, StorageGuard, StorageGuardMut, StorageKey, StorageLevel, StorageMap, StorageType};
 use crate::crypto;
 use alloy_primitives::U256;
 use core::{cell::OnceCell, marker::PhantomData};
@@ -15,6 +13,8 @@ pub struct StorageVec<S: StorageType> {
     marker: PhantomData<S>,
 }
 
+unsafe impl<S: StorageType> StorageLevel for StorageVec<S> {}
+
 impl<S: StorageType> StorageType for StorageVec<S> {
     type Wraps<'a> = StorageGuard<'a, StorageVec<S>> where Self: 'a;
     type WrapsMut<'a> = StorageGuardMut<'a, StorageVec<S>> where Self: 'a;