@layerzerolabs/protocol-stellar-v2 0.2.8

package/contracts/common-macros/src/storage.rs

@@ -0,0 +1,439 @@
+//! Storage macro implementation for Stellar smart contracts.
+//!
+//! Generates strongly-typed storage API from enum variants with automatic TTL management.
+
+use heck::ToSnakeCase;
+use itertools::Itertools;
+use proc_macro2::{Ident, TokenStream};
+use quote::{format_ident, quote};
+use syn::{
+    parse::{Parse, ParseStream},
+    Expr, Fields, FieldsNamed, Meta, Result, Type, TypePath, Variant,
+};
+
+// ============================================================================
+// Types & Parsing
+// ============================================================================
+
+/// TTL provider configuration for auto-extending storage TTL.
+/// Parsed from `#[storage(MyTtlProvider)]` attribute.
+#[derive(Debug, Clone)]
+struct TtlProvider(Option<Type>);
+
+impl Parse for TtlProvider {
+    fn parse(input: ParseStream) -> Result<Self> {
+        if input.is_empty() {
+            return Ok(Self(None));
+        }
+
+        let meta: Meta = input.parse()?;
+        let Meta::Path(path) = meta else {
+            return Err(syn::Error::new(input.span(), "TTL provider must be a path"));
+        };
+
+        Ok(Self(Some(Type::Path(TypePath { qself: None, path }))))
+    }
+}
+
+/// Storage type parsed from `#[instance(Type)]`, `#[persistent(Type)]`, or `#[temporary(Type)]`.
+#[derive(Debug, Clone)]
+enum StorageType {
+    Instance(Type),
+    Persistent(Type),
+    Temporary(Type),
+}
+
+impl StorageType {
+    fn value_type(&self) -> &Type {
+        match self {
+            Self::Instance(t) | Self::Persistent(t) | Self::Temporary(t) => t,
+        }
+    }
+
+    fn is_instance(&self) -> bool {
+        matches!(self, Self::Instance(_))
+    }
+
+    /// Returns the storage accessor method name.
+    fn accessor_name(&self) -> &'static str {
+        match self {
+            Self::Instance(_) => "instance",
+            Self::Persistent(_) => "persistent",
+            Self::Temporary(_) => "temporary",
+        }
+    }
+}
+
+impl TryFrom<&syn::Attribute> for StorageType {
+    type Error = String;
+
+    fn try_from(attr: &syn::Attribute) -> std::result::Result<Self, Self::Error> {
+        let ident = attr.path().get_ident().ok_or("identifier expected to be instance, persistent, or temporary")?;
+        let value_type = attr.parse_args::<Type>().map_err(|e| format!("failed to parse value type: {}", e))?;
+
+        match ident.to_string().as_str() {
+            "instance" => Ok(Self::Instance(value_type)),
+            "persistent" => Ok(Self::Persistent(value_type)),
+            "temporary" => Ok(Self::Temporary(value_type)),
+            other => Err(format!("unknown storage type: {}", other)),
+        }
+    }
+}
+
+/// Parsed configuration for a single storage enum variant.
+#[derive(Debug, Clone)]
+pub(crate) struct StorageVariantConfig {
+    name: String,
+    storage_type: StorageType,
+    default_value: Option<Expr>,
+    no_auto_ttl_extension: bool,
+}
+
+impl StorageVariantConfig {
+    /// Returns generated function names for this storage variant.
+    fn function_names(&self) -> FunctionNames {
+        let base = &self.name;
+        FunctionNames {
+            getter: format_ident!("{}", base),
+            setter: format_ident!("set_{}", base),
+            remover: format_ident!("remove_{}", base),
+            set_or_remove: format_ident!("set_or_remove_{}", base),
+            ttl_extender: format_ident!("extend_ttl_{}", base),
+            has: format_ident!("has_{}", base),
+        }
+    }
+}
+
+impl TryFrom<&Variant> for StorageVariantConfig {
+    type Error = String;
+
+    fn try_from(variant: &Variant) -> std::result::Result<Self, Self::Error> {
+        let attrs = &variant.attrs;
+
+        // Known attribute identifiers for this macro ("doc" allows /// comments on variants)
+        const KNOWN_ATTRS: &[&str] =
+            &["doc", "instance", "persistent", "temporary", "default", "no_auto_ttl_extension", "name"];
+
+        // Check for unsupported attributes
+        if let Some(unknown) = attrs
+            .iter()
+            .filter_map(|attr| attr.path().get_ident())
+            .find(|ident| !KNOWN_ATTRS.contains(&ident.to_string().as_str()))
+        {
+            return Err(format!(
+                "unknown attribute '{}' on variant '{}'. Supported attributes are: {}",
+                unknown,
+                variant.ident,
+                KNOWN_ATTRS.join(", ")
+            ));
+        }
+
+        // Parse storage type (exactly one required)
+        let storage_type =
+            attrs.iter().filter_map(|attr| StorageType::try_from(attr).ok()).exactly_one().map_err(|e| {
+                format!(
+                    "storage type must be specified exactly once as \
+                     '#[instance(Type)]', '#[persistent(Type)]', or '#[temporary(Type)]': {}",
+                    e
+                )
+            })?;
+
+        // Parse optional default value
+        let default_value = attrs
+            .iter()
+            .filter(|attr| attr.path().is_ident("default"))
+            .at_most_one()
+            .map_err(|_| "multiple default values specified")?
+            .map(|attr| attr.parse_args::<Expr>())
+            .transpose()
+            .map_err(|e| format!("failed to parse default value: {}", e))?;
+
+        // Check for no_auto_ttl_extension flag
+        let no_auto_ttl_extension = attrs
+            .iter()
+            .filter(|attr| attr.path().is_ident("no_auto_ttl_extension"))
+            .map(|attr| {
+                attr.meta.require_path_only().map_err(|_| "no_auto_ttl_extension does not accept arguments".to_string())
+            })
+            .at_most_one()
+            .map_err(|_| "multiple no_auto_ttl_extension attributes specified")?
+            .transpose()?
+            .is_some();
+
+        // Parse optional name override for storage key generation
+        // If no #[name("custom_name")] attribute is provided, defaults to snake_case of variant name
+        let name = attrs
+            .iter()
+            .filter(|attr| attr.path().is_ident("name"))
+            .at_most_one()
+            .map_err(|_| "multiple name attributes specified")?
+            .map(|attr| attr.parse_args::<syn::LitStr>().map(|lit| lit.value()))
+            .transpose()
+            .map_err(|e| format!("failed to parse name attribute: {}", e))?
+            .unwrap_or_else(|| variant.ident.to_string().to_snake_case());
+
+        Ok(Self { storage_type, default_value, no_auto_ttl_extension, name })
+    }
+}
+
+/// Generated function names for a storage variant.
+struct FunctionNames {
+    getter: Ident,
+    setter: Ident,
+    remover: Ident,
+    set_or_remove: Ident,
+    ttl_extender: Ident,
+    has: Ident,
+}
+
+// ============================================================================
+// Code Generation
+// ============================================================================
+
+/// Generates the storage API from the `#[storage]` attribute macro.
+pub(crate) fn generate_storage(attr: TokenStream, input: TokenStream) -> TokenStream {
+    let item_enum: syn::ItemEnum = syn::parse2(input).unwrap_or_else(|e| panic!("failed to parse enum: {}", e));
+    let ttl_provider: TtlProvider = syn::parse2(attr).unwrap_or_else(|e| panic!("failed to parse attribute: {}", e));
+
+    let enum_name = &item_enum.ident;
+    let contracttype_variants: Vec<_> = item_enum.variants.iter().map(gen_contracttype_variant).collect();
+    let inherent_impls: Vec<_> =
+        item_enum.variants.iter().map(|v| gen_storage_functions(enum_name, v, &ttl_provider)).collect();
+
+    let visibility = &item_enum.vis;
+    quote! {
+        use utils::ttl::TtlConfigProvider as _;
+
+        #[soroban_sdk::contracttype]
+        #visibility enum #enum_name {
+            #(#contracttype_variants,)*
+        }
+
+        impl #enum_name {
+            #(#inherent_impls)*
+        }
+    }
+}
+
+/// Generates contracttype enum variant: `Variant` or `Variant(Type1, Type2, ...)`.
+fn gen_contracttype_variant(variant: &Variant) -> TokenStream {
+    let name = &variant.ident;
+    match &variant.fields {
+        Fields::Unit => quote! { #name },
+        Fields::Named(FieldsNamed { named, .. }) => {
+            let types = named.iter().map(|f| &f.ty);
+            quote! { #name(#(#types),*) }
+        }
+        _ => panic!("only unit variants or named fields are supported in storage enums"),
+    }
+}
+
+/// Generates all storage accessor functions for a variant.
+fn gen_storage_functions(enum_name: &Ident, variant: &Variant, ttl_provider: &TtlProvider) -> TokenStream {
+    let config = StorageVariantConfig::try_from(variant)
+        .unwrap_or_else(|e| panic!("failed to parse storage variant for {}: {}", variant.ident, e));
+
+    let FunctionNames { getter, setter, remover, set_or_remove, has, ttl_extender } = config.function_names();
+
+    let fn_params = gen_fn_params(variant);
+    let fn_args = gen_fn_args(variant);
+    let storage_key = gen_storage_key(variant, enum_name);
+    let storage_acc = gen_storage_accessor(&config.storage_type);
+    let value_type = config.storage_type.value_type();
+
+    // Auto TTL extension (if provider exists and not disabled)
+    let auto_ttl = ttl_provider
+        .0
+        .as_ref()
+        .filter(|_| !config.no_auto_ttl_extension)
+        .map(|p| gen_ttl_extension(&config.storage_type, Some(p)))
+        .unwrap_or_default();
+
+    let getter_impl =
+        gen_getter_fn(&getter, &fn_params, &storage_key, &storage_acc, value_type, &config.default_value, &auto_ttl);
+
+    // Manual TTL extension (user provides threshold/extend_to)
+    let manual_ttl = gen_ttl_extension(&config.storage_type, None);
+
+    quote! {
+        #getter_impl
+
+        pub fn #setter(#fn_params, value: &#value_type) {
+            let key = #storage_key;
+            #storage_acc.set(&key, value);
+            #auto_ttl
+        }
+
+        pub fn #remover(#fn_params) {
+            let key = #storage_key;
+            #storage_acc.remove(&key);
+        }
+
+        pub fn #set_or_remove(#fn_params, value: &Option<#value_type>) {
+            match value.as_ref() {
+                Some(v) => Self::#setter(#fn_args, v),
+                None => Self::#remover(#fn_args),
+            }
+        }
+
+        pub fn #has(#fn_params) -> bool {
+            let key = #storage_key;
+            let exists = #storage_acc.has(&key);
+            if exists { #auto_ttl }
+            exists
+        }
+
+        pub fn #ttl_extender(#fn_params, threshold: u32, extend_to: u32) {
+            let key = #storage_key;
+            #manual_ttl
+        }
+    }
+}
+
+/// Generates the storage accessor: `env.storage().instance()` etc.
+fn gen_storage_accessor(storage_type: &StorageType) -> TokenStream {
+    let method = format_ident!("{}", storage_type.accessor_name());
+    quote! { env.storage().#method() }
+}
+
+/// Generates TTL extension code.
+/// - With provider: uses provider to get threshold/extend_to
+/// - Without provider (None): uses local variables `threshold` and `extend_to`
+fn gen_ttl_extension(storage_type: &StorageType, provider: Option<&Type>) -> TokenStream {
+    let get_ttl_values = provider.map(|p| {
+        let method = format_ident!("{}", storage_type.accessor_name());
+        quote! { let utils::ttl::TtlConfig { threshold, extend_to } = utils::ttl::ttl_config::#method::<#p>(env); }
+    });
+
+    let accessor = gen_storage_accessor(storage_type);
+    let extend_call = if storage_type.is_instance() {
+        quote! { #accessor.extend_ttl(threshold, extend_to); }
+    } else {
+        quote! { #accessor.extend_ttl(&key, threshold, extend_to); }
+    };
+
+    quote! {
+        #get_ttl_values
+        #extend_call
+    }
+}
+
+/// Checks if a type is a primitive type
+fn is_primitive_type(ty: &Type) -> bool {
+    if let Type::Path(type_path) = ty {
+        if type_path.path.segments.len() == 1 {
+            let segment = &type_path.path.segments[0];
+            let ident = &segment.ident;
+            matches!(ident.to_string().as_str(), "u32" | "i32" | "u64" | "i64" | "u128" | "i128" | "bool")
+        } else {
+            false
+        }
+    } else {
+        false
+    }
+}
+
+/// Generates function parameters: `env: &Env` or `env: &Env, field1: Type1, ...`.
+/// For primitive types (u32, i32, u64, i64, u128, i128, bool), uses pass-by-value.
+/// For other types, uses pass-by-reference.
+fn gen_fn_params(variant: &Variant) -> TokenStream {
+    let fields = extract_named_fields(&variant.fields);
+    if fields.is_empty() {
+        quote! { env: &soroban_sdk::Env }
+    } else {
+        let params = fields.iter().map(|(name, ty)| {
+            if is_primitive_type(ty) {
+                quote! { #name: #ty }
+            } else {
+                quote! { #name: &#ty }
+            }
+        });
+        quote! { env: &soroban_sdk::Env, #(#params),* }
+    }
+}
+
+/// Generates function argument list for calls: `env, field1, field2, ...`.
+fn gen_fn_args(variant: &Variant) -> TokenStream {
+    let fields = extract_named_fields(&variant.fields);
+    if fields.is_empty() {
+        quote! { env }
+    } else {
+        let field_names = fields.iter().map(|(name, _)| name);
+        quote! { env, #(#field_names),* }
+    }
+}
+
+/// Generates key constructor: `Enum::Variant` or `Enum::Variant(field1, field2)`.
+/// For non-primitive types that are passed by reference, clones them.
+fn gen_storage_key(variant: &Variant, enum_name: &Ident) -> TokenStream {
+    let fields = extract_named_fields(&variant.fields);
+    let variant_ident = &variant.ident;
+    if fields.is_empty() {
+        quote! { #enum_name::#variant_ident }
+    } else {
+        let field_names = fields.iter().map(|(name, ty)| {
+            if is_primitive_type(ty) {
+                quote! { #name }
+            } else {
+                quote! { #name.clone() }
+            }
+        });
+        quote! { #enum_name::#variant_ident(#(#field_names),*) }
+    }
+}
+
+/// Extracts (name, type) pairs from named fields.
+fn extract_named_fields(fields: &Fields) -> Vec<(&Ident, &Type)> {
+    match fields {
+        Fields::Unit => vec![],
+        Fields::Named(named) => named.named.iter().map(|f| (f.ident.as_ref().unwrap(), &f.ty)).collect(),
+        _ => panic!("only unit variants or named fields are supported in storage enums"),
+    }
+}
+
+/// Generates the getter function with appropriate return type and TTL handling.
+fn gen_getter_fn(
+    name: &Ident,
+    params: &TokenStream,
+    key: &TokenStream,
+    accessor: &TokenStream,
+    value_type: &Type,
+    has_default: &Option<Expr>,
+    ttl_extend: &TokenStream,
+) -> TokenStream {
+    let has_ttl = !ttl_extend.is_empty();
+    match (has_default, has_ttl) {
+        // Has default + has TTL extension
+        (Some(default_expr), true) => quote! {
+            pub fn #name(#params) -> #value_type {
+                let key = #key;
+                let value = #accessor.get::<_, #value_type>(&key);
+                if value.is_some() { #ttl_extend }
+                value.unwrap_or(#default_expr)
+            }
+        },
+        // Has default, no TTL extension
+        (Some(default_expr), false) => quote! {
+            pub fn #name(#params) -> #value_type {
+                let key = #key;
+                #accessor.get::<_, #value_type>(&key).unwrap_or(#default_expr)
+            }
+        },
+        // No default + has TTL extension
+        (None, true) => quote! {
+            pub fn #name(#params) -> Option<#value_type> {
+                let key = #key;
+                let value = #accessor.get::<_, #value_type>(&key);
+                if value.is_some() { #ttl_extend }
+                value
+            }
+        },
+        // No default, no TTL extension
+        (None, false) => quote! {
+            pub fn #name(#params) -> Option<#value_type> {
+                let key = #key;
+                #accessor.get::<_, #value_type>(&key)
+            }
+        },
+    }
+}