use proc_macro2::TokenStream;
use proc_macro_error::emit_error;
use quote::{quote, ToTokens};
use syn::{
parse_macro_input, parse_quote, punctuated::Punctuated, spanned::Spanned, ItemStruct, Token,
Type,
};
use crate::{
consts::{STYLUS_CONTRACT_ADDRESS_FIELD, STYLUS_HOST_FIELD},
utils::attrs::consume_flag,
};
fn is_phantom_data(ty: &Type) -> bool {
if let Type::Path(type_path) = &ty {
return type_path
.path
.segments
.last()
.map(|s| s.ident == "PhantomData")
.unwrap_or(false);
}
false
}
pub fn storage(
attr: proc_macro::TokenStream,
input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
if !attr.is_empty() {
emit_error!(
TokenStream::from(attr).span(),
"this macro is not configurable"
);
}
let item = parse_macro_input!(input as ItemStruct);
let ItemStruct {
attrs,
vis,
ident,
generics,
fields,
..
} = item;
let expanded_fields = match fields {
syn::Fields::Named(named_fields) => {
let mut original_fields = named_fields.named;
for field in original_fields.iter_mut() {
if field.ident != Some(STYLUS_CONTRACT_ADDRESS_FIELD.as_ident())
&& !is_phantom_data(&field.ty)
{
field.attrs.push(parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
});
}
}
quote! {
#[cfg(not(feature = "contract-client-gen"))]
#STYLUS_HOST_FIELD: stylus_sdk::host::VM,
#original_fields
}
}
syn::Fields::Unnamed(_) => {
emit_error!(
fields.span(),
"Tuple structs are not supported by #[storage]"
);
return fields.to_token_stream().into();
}
syn::Fields::Unit => {
quote! {
#[cfg(not(feature = "contract-client-gen"))]
#STYLUS_HOST_FIELD: stylus_sdk::host::VM,
}
}
};
let mut host_injected_item: syn::ItemStruct = parse_quote! {
#(#attrs)*
#[cfg_attr(feature = "contract-client-gen", derive(Default))]
#vis struct #ident #generics {
#[cfg(feature = "contract-client-gen")]
#STYLUS_CONTRACT_ADDRESS_FIELD: stylus_sdk::alloy_primitives::Address,
#expanded_fields
}
};
let storage = Storage::from(&mut host_injected_item);
let mut output = host_injected_item.into_token_stream();
storage.to_tokens(&mut output);
output.into()
}
#[derive(Debug)]
struct Storage {
name: syn::Ident,
generics: syn::Generics,
fields: Vec<StorageField>,
}
impl Storage {
fn item_impl(&self) -> syn::ItemImpl {
let name = &self.name;
let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
let size = TokenStream::from_iter(self.fields.iter().map(StorageField::size));
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl #impl_generics #name #ty_generics #where_clause {
const fn required_slots() -> usize {
use stylus_sdk::storage;
let mut total: usize = 0;
let mut space: usize = 32;
#size
if space != 32 || total == 0 {
total += 1;
}
total
}
}
}
}
fn impl_storage_type(&self) -> syn::ItemImpl {
let name = &self.name;
let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
let init = Punctuated::<syn::FieldValue, Token![,]>::from_iter(
self.fields.iter().filter_map(StorageField::init),
);
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl #impl_generics stylus_sdk::storage::StorageType for #name #ty_generics #where_clause {
type Wraps<'a> = stylus_sdk::storage::StorageGuard<'a, Self> where Self: 'a;
type WrapsMut<'a> = stylus_sdk::storage::StorageGuardMut<'a, Self> where Self: 'a;
const SLOT_BYTES: usize = 32;
const REQUIRED_SLOTS: usize = Self::required_slots();
unsafe fn new(mut root: stylus_sdk::alloy_primitives::U256, offset: u8, host: stylus_sdk::host::VM) -> Self {
use stylus_sdk::{storage, alloy_primitives};
debug_assert!(offset == 0);
let mut space: usize = 32;
let mut slot: usize = 0;
let accessor = Self {
#STYLUS_HOST_FIELD: host.clone(),
#[cfg(feature = "contract-client-gen")]
#STYLUS_CONTRACT_ADDRESS_FIELD: stylus_sdk::alloy_primitives::Address::ZERO,
#init
};
accessor
}
fn load<'s>(self) -> Self::Wraps<'s> {
stylus_sdk::storage::StorageGuard::new(self)
}
fn load_mut<'s>(self) -> Self::WrapsMut<'s> {
stylus_sdk::storage::StorageGuardMut::new(self)
}
}
}
}
fn impl_host_access(&self) -> syn::ItemImpl {
let name = &self.name;
let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl #impl_generics stylus_sdk::stylus_core::HostAccess for #name #ty_generics #where_clause {
type Host = stylus_sdk::host::VM;
fn vm(&self) -> &Self::Host {
&self.__stylus_host
}
}
}
}
fn impl_value_denier(&self) -> syn::ItemImpl {
let name = &self.name;
let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl #impl_generics stylus_sdk::stylus_core::host::ValueDenier for #name #ty_generics #where_clause {
fn deny_value(&self, method_name: &str) -> Result<(), Vec<u8>> {
if self.vm().msg_value() == stylus_sdk::alloy_primitives::U256::ZERO {
return Ok(());
}
stylus_sdk::console!("method {method_name} not payable");
Err(vec![])
}
}
}
}
fn impl_constructor_guard(&self) -> syn::ItemImpl {
let name = &self.name;
let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl #impl_generics stylus_sdk::stylus_core::host::ConstructorGuard for #name #ty_generics #where_clause {
fn check_constructor_slot(&self) -> Result<(), Vec<u8>> {
let mut slot = unsafe {
stylus_sdk::storage::StorageBool::new(
stylus_sdk::abi::internal::CONSTRUCTOR_EXECUTED_SLOT,
0,
self.__stylus_host.clone()
)
};
if slot.get() {
stylus_sdk::console!("constructor already called");
return Err(alloc::vec![]);
}
slot.set(true);
Ok(())
}
}
}
}
fn impl_from_vm(&self) -> Option<syn::ItemImpl> {
cfg_if::cfg_if! {
if #[cfg(feature = "stylus-test")] {
let name = &self.name;
let (_, ty_generics, where_clause) = self.generics.split_for_impl();
let mut new_generics = self.generics.clone();
let host_param =
parse_quote!(__StylusHostType: stylus_sdk::stylus_core::Host + Clone + 'static);
new_generics.params.push(host_param);
let (impl_generics, _, _) = new_generics.split_for_impl();
Some(parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl #impl_generics From<&__StylusHostType> for #name #ty_generics #where_clause {
fn from(host: &__StylusHostType) -> Self {
unsafe {
Self::new(
stylus_sdk::alloy_primitives::U256::ZERO,
0,
stylus_sdk::host::VM {
host: alloc::boxed::Box::new(host.clone()),
},
)
}
}
}
})
} else {
None
}
}
}
fn impl_vm_access(&self) -> syn::ItemImpl {
let name = &self.name;
let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl #impl_generics stylus_sdk::host::VMAccess for #name #ty_generics #where_clause {
unsafe fn raw_vm(&self) -> stylus_sdk::host::VM {
self.__stylus_host.clone()
}
}
}
}
}
impl From<&mut syn::ItemStruct> for Storage {
fn from(node: &mut syn::ItemStruct) -> Self {
let name = node.ident.clone();
let generics = node.generics.clone();
let fields = node
.fields
.iter_mut()
.enumerate()
.filter_map(|(idx, field)| {
if let syn::Type::Path(..) = &field.ty {
Some(StorageField::new(idx, field))
} else {
None
}
})
.collect();
Self {
name,
generics,
fields,
}
}
}
impl ToTokens for Storage {
fn to_tokens(&self, tokens: &mut TokenStream) {
self.item_impl().to_tokens(tokens);
self.impl_storage_type().to_tokens(tokens);
self.impl_host_access().to_tokens(tokens);
self.impl_value_denier().to_tokens(tokens);
self.impl_constructor_guard().to_tokens(tokens);
self.impl_from_vm().to_tokens(tokens);
self.impl_vm_access().to_tokens(tokens);
for field in &self.fields {
field.impl_borrow(&self.name).to_tokens(tokens);
field.impl_borrow_mut(&self.name).to_tokens(tokens);
}
}
}
#[derive(Debug)]
struct StorageField {
name: Option<syn::Ident>,
ty: syn::Type,
accessor: syn::Member,
borrow: bool,
}
impl StorageField {
fn new(idx: usize, field: &mut syn::Field) -> Self {
check_type(field);
let name = field.ident.clone();
let ty = field.ty.clone();
let accessor = field
.ident
.clone()
.map(syn::Member::from)
.unwrap_or_else(|| idx.into());
let borrow = consume_flag(&mut field.attrs, "borrow");
Self {
name,
ty,
accessor,
borrow,
}
}
fn init(&self) -> Option<syn::FieldValue> {
let Some(ident) = &self.name else {
return None;
};
if *ident == STYLUS_HOST_FIELD.as_ident() {
return None;
}
if *ident == STYLUS_CONTRACT_ADDRESS_FIELD.as_ident() {
return None;
}
let ty = &self.ty;
Some(parse_quote! {
#ident: {
let bytes = <#ty as storage::StorageType>::SLOT_BYTES;
let words = <#ty as storage::StorageType>::REQUIRED_SLOTS;
if space < bytes {
space = 32;
slot += 1;
}
space -= bytes;
let root = root + alloy_primitives::U256::from(slot);
let field = <#ty as storage::StorageType>::new(root, space as u8, host.clone());
if words > 0 {
slot += words;
space = 32;
}
field
}
})
}
fn size(&self) -> TokenStream {
let ty = &self.ty;
let Some(ident) = &self.name else {
return quote! {};
};
if *ident == STYLUS_HOST_FIELD.as_ident() {
return quote! {};
}
if *ident == STYLUS_CONTRACT_ADDRESS_FIELD.as_ident() {
return quote! {};
}
quote! {
let bytes = <#ty as storage::StorageType>::SLOT_BYTES;
let words = <#ty as storage::StorageType>::REQUIRED_SLOTS;
if space < bytes {
space = 32;
total += 1;
}
space -= bytes;
if words > 0 {
total += words;
space = 32;
}
}
}
fn impl_borrow(&self, name: &syn::Ident) -> Option<syn::ItemImpl> {
let Self { ty, accessor, .. } = self;
self.borrow.then(|| {
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl core::borrow::Borrow<#ty> for #name {
fn borrow(&self) -> &#ty {
&self.#accessor
}
}
}
})
}
fn impl_borrow_mut(&self, name: &syn::Ident) -> Option<syn::ItemImpl> {
let Self { ty, accessor, .. } = self;
self.borrow.then(|| {
parse_quote! {
#[cfg(not(feature = "contract-client-gen"))]
impl core::borrow::BorrowMut<#ty> for #name {
fn borrow_mut(&mut self) -> &mut #ty {
&mut self.#accessor
}
}
}
})
}
}
fn check_type(field: &syn::Field) {
let Type::Path(ty) = &field.ty else {
unreachable!();
};
let path = &ty.path.segments.last().unwrap().ident;
let not_supported = format!("Type `{path}` not supported for EVM state storage");
match path.to_string().as_str() {
x @ ("u8" | "u16" | "u32" | "u64" | "u128" | "i8" | "i16" | "i32" | "i64" | "i128"
| "U8" | "U16" | "U32" | "U64" | "U128" | "I8" | "I16" | "I32" | "I64" | "I128") => {
emit_error!(
&field,
"{not_supported}. Instead try `Storage{}`.",
x.to_uppercase()
);
}
"usize" => emit_error!(&field, "{not_supported}."),
"isize" => emit_error!(&field, "{not_supported}."),
"bool" => emit_error!(&field, "{not_supported}. Instead try `StorageBool`."),
"f32" | "f64" => emit_error!(&field, "{not_supported}. Consider fixed-point arithmetic."),
_ => {}
}
}