From 921c86a78bd1037042831cccab146bad55975a35 Mon Sep 17 00:00:00 2001 From: Alisander Qoshqosh Date: Tue, 23 Apr 2024 19:06:19 +0400 Subject: [PATCH] add try_get_inner and try_get_inner_mut functions and implementations --- examples/erc20/src/erc20.rs | 4 +- stylus-proc/src/methods/entrypoint.rs | 13 +------ stylus-proc/src/methods/external.rs | 15 +------ stylus-proc/src/storage/mod.rs | 50 ++++++++++++++++-------- stylus-sdk/src/abi/mod.rs | 5 +-- stylus-sdk/src/storage/array.rs | 4 +- stylus-sdk/src/storage/bytes.rs | 6 ++- stylus-sdk/src/storage/map.rs | 4 +- stylus-sdk/src/storage/mod.rs | 18 ++++++++- stylus-sdk/src/storage/traits.rs | 56 +++++++++++++++++++++++---- stylus-sdk/src/storage/vec.rs | 6 +-- 11 files changed, 119 insertions(+), 62 deletions(-) diff --git a/examples/erc20/src/erc20.rs b/examples/erc20/src/erc20.rs index 9f237378..ae3d4dde 100644 --- a/examples/erc20/src/erc20.rs +++ b/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 { + pub struct Erc20 { /// Maps users to balances mapping(address => uint256) balances; /// Maps users to a mapping of each spender's allowance diff --git a/stylus-proc/src/methods/entrypoint.rs b/stylus-proc/src/methods/entrypoint.rs index a5fbf2a7..e7d8c52e 100644 --- a/stylus-proc/src/methods/entrypoint.rs +++ b/stylus-proc/src/methods/entrypoint.rs @@ -22,18 +22,7 @@ pub fn entrypoint(attr: TokenStream, input: TokenStream) -> TokenStream { let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl(); output.extend(quote!{ - unsafe impl #impl_generics stylus_sdk::storage::TopLevelStorage for #name #ty_generics #where_clause { - fn get_storage(&mut self) -> &mut S { - use stylus_sdk::storage::InnerStorage; - unsafe { - self.try_get_storage().unwrap_or_else(|| { - panic!( - "storage for type doesn't exist - type name is {}", - core::any::type_name::() - )}) - } - } - } + unsafe impl #impl_generics stylus_sdk::storage::TopLevelStorage for #name #ty_generics #where_clause {} fn entrypoint(input: alloc::vec::Vec) -> stylus_sdk::ArbResult { use stylus_sdk::{abi::Router, alloy_primitives::U256, console, hex, storage::StorageType}; diff --git a/stylus-proc/src/methods/external.rs b/stylus-proc/src/methods/external.rs index 9ec4e98b..b3fc23c6 100644 --- a/stylus-proc/src/methods/external.rs +++ b/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 stylus_sdk::abi::Router for #self_ty where - S: stylus_sdk::storage::TopLevelStorage + core::borrow::BorrowMut, - #(#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 { use stylus_sdk::{function_selector, alloy_sol_types::SolType}; diff --git a/stylus-proc/src/storage/mod.rs b/stylus-proc/src/storage/mod.rs index 2361d170..8ce4bdd0 100644 --- a/stylus-proc/src/storage/mod.rs +++ b/stylus-proc/src/storage/mod.rs @@ -26,6 +26,12 @@ pub fn solidity_storage(_attr: TokenStream, input: TokenStream) -> TokenStream { 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 { @@ -37,13 +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(), - }; - - inner_storage_accessors.push(accessor.clone()); - + borrows.extend(quote! { impl core::borrow::Borrow<#ty> for #name { fn borrow(&self) -> &#ty { @@ -120,24 +120,40 @@ pub fn solidity_storage(_attr: TokenStream, input: TokenStream) -> TokenStream { }); } - let inner_storage_calls = inner_storage_accessors.into_iter().map(|accessor|{ + let inner_storage_calls = inner_storage_accessors.iter().map(|accessor|{ quote! { - .or(self.#accessor.try_get_storage()) + .or(self.#accessor.try_get_inner()) } - }); + }); - let storage_impl = quote! { - #[allow(clippy::transmute_ptr_to_ptr)] - unsafe impl #impl_generics stylus_sdk::storage::InnerStorage for #name #ty_generics #where_clause { - unsafe fn try_get_storage(&mut self) -> Option<&mut S> { + 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(&self) -> Option<&S> { use core::any::TypeId; - use stylus_sdk::storage::InnerStorage; + use stylus_sdk::storage::StorageLevel; if TypeId::of::() == TypeId::of::() { - Some(unsafe { core::mem::transmute::<_, _>(self) }) + Some(unsafe { &*(self as *const Self as *const S) }) } else { None #(#inner_storage_calls)* } } + + unsafe fn try_get_inner_mut(&mut self) -> Option<&mut S> { + use core::any::TypeId; + use stylus_sdk::storage::StorageLevel; + if TypeId::of::() == TypeId::of::() { + Some(unsafe { &mut *(self as *mut Self as *mut S) }) + } else { + None #(#inner_mut_storage_calls)* + } + } } }; @@ -189,7 +205,7 @@ pub fn solidity_storage(_attr: TokenStream, input: TokenStream) -> TokenStream { #borrows - #storage_impl + #storage_level_impl }; expanded.into() } diff --git a/stylus-sdk/src/abi/mod.rs b/stylus-sdk/src/abi/mod.rs index 0545ff1a..07722383 100644 --- a/stylus-sdk/src/abi/mod.rs +++ b/stylus-sdk/src/abi/mod.rs @@ -41,11 +41,8 @@ pub mod internal; /// Composition with other routers is possible via `#[inherit]`. pub trait Router where - S: TopLevelStorage + BorrowMut, + 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`. diff --git a/stylus-sdk/src/storage/array.rs b/stylus-sdk/src/storage/array.rs index 697ec494..389a97e7 100644 --- a/stylus-sdk/src/storage/array.rs +++ b/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 { marker: PhantomData, } +unsafe impl StorageLevel for StorageArray {} + impl StorageType for StorageArray { type Wraps<'a> = StorageGuard<'a, StorageArray> where Self: 'a; type WrapsMut<'a> = StorageGuardMut<'a, StorageArray> where Self: 'a; diff --git a/stylus-sdk/src/storage/bytes.rs b/stylus-sdk/src/storage/bytes.rs index c142df41..c9581afe 100644 --- a/stylus-sdk/src/storage/bytes.rs +++ b/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, } +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; diff --git a/stylus-sdk/src/storage/map.rs b/stylus-sdk/src/storage/map.rs index 660ae6b1..5aae4a6e 100644 --- a/stylus-sdk/src/storage/map.rs +++ b/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 { marker: PhantomData<(K, V)>, } +unsafe impl StorageLevel for StorageMap {} + impl StorageType for StorageMap where K: StorageKey, diff --git a/stylus-sdk/src/storage/mod.rs b/stylus-sdk/src/storage/mod.rs index 46c2e9e8..1c645cdf 100644 --- a/stylus-sdk/src/storage/mod.rs +++ b/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, InnerStorage, + TopLevelStorage, StorageLevel, }; pub use vec::StorageVec; @@ -160,6 +160,8 @@ impl StorageUint { } } +unsafe impl StorageLevel for StorageUint {} + impl StorageType for StorageUint { type Wraps<'a> = Uint; type WrapsMut<'a> = StorageGuardMut<'a, Self>; @@ -235,6 +237,8 @@ impl StorageSigned { } } +unsafe impl StorageLevel for StorageSigned {} + impl StorageType for StorageSigned { type Wraps<'a> = Signed; type WrapsMut<'a> = StorageGuardMut<'a, Self>; @@ -306,6 +310,8 @@ impl StorageFixedBytes { } } +unsafe impl StorageLevel for StorageFixedBytes {} + impl StorageType for StorageFixedBytes where ByteCount: SupportedFixedBytes, @@ -386,6 +392,8 @@ impl StorageBool { } } +unsafe impl StorageLevel for StorageBool {} + impl StorageType for StorageBool { type Wraps<'a> = bool; type WrapsMut<'a> = StorageGuardMut<'a, Self>; @@ -459,6 +467,8 @@ impl StorageAddress { } } +unsafe impl StorageLevel for StorageAddress {} + impl StorageType for StorageAddress { type Wraps<'a> = Address; type WrapsMut<'a> = StorageGuardMut<'a, Self>; @@ -531,6 +541,8 @@ impl StorageBlockNumber { } } +unsafe impl StorageLevel for StorageBlockNumber {} + impl StorageType for StorageBlockNumber { type Wraps<'a> = BlockNumber; type WrapsMut<'a> = StorageGuardMut<'a, Self>; @@ -603,6 +615,8 @@ impl StorageBlockHash { } } +unsafe impl StorageLevel for StorageBlockHash {} + impl StorageType for StorageBlockHash { type Wraps<'a> = BlockHash; type WrapsMut<'a> = StorageGuardMut<'a, Self>; @@ -647,6 +661,8 @@ impl From for BlockHash { } } +unsafe impl StorageLevel for PhantomData {} + /// We implement `StorageType` for `PhantomData` so that storage types can be generic. impl StorageType for PhantomData { type Wraps<'a> = Self where Self: 'a; diff --git a/stylus-sdk/src/storage/traits.rs b/stylus-sdk/src/storage/traits.rs index 86597c69..34ca610f 100644 --- a/stylus-sdk/src/storage/traits.rs +++ b/stylus-sdk/src/storage/traits.rs @@ -95,18 +95,60 @@ 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(&self) -> &S { + unsafe { + self.try_get_inner().unwrap_or_else(|| { + panic!( + "type does not exist inside TopLevelStorage - type is {}", + core::any::type_name::() + )}) + } + } - /// Retrieve mutable reference to inner storage of type [`S`] - fn get_storage(&mut self) -> &mut S { - panic!("arbitrary storage access is not implemented") + /// Retrieve mutable reference to inner storage of type [`S`] or panic. + fn inner_mut(&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::() + )}) + } } } -pub unsafe trait InnerStorage { +/// 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 etrieve mutable reference to inner storage of type [`S`] - unsafe fn try_get_storage(&mut self) -> Option<&mut S>; + /// 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(&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(&mut self) -> Option<&mut S>{ + None + } } /// Binds a storage accessor to a lifetime to prevent aliasing. diff --git a/stylus-sdk/src/storage/vec.rs b/stylus-sdk/src/storage/vec.rs index 8c97dc6f..7c2e4389 100644 --- a/stylus-sdk/src/storage/vec.rs +++ b/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 { marker: PhantomData, } +unsafe impl StorageLevel for StorageVec {} + impl StorageType for StorageVec { type Wraps<'a> = StorageGuard<'a, StorageVec> where Self: 'a; type WrapsMut<'a> = StorageGuardMut<'a, StorageVec> where Self: 'a;