macro: allow flex imports

Author: SwenSchaeferjohann

sdk-libs/macros/src/compressible.rs

@@ -6,7 +6,7 @@ use syn::{
     parse::{Parse, ParseStream},
     punctuated::Punctuated,
     visit_mut, Attribute, Expr, Field, Ident, Item, ItemEnum, ItemFn, ItemMod, ItemStruct, Result,
-    Token,
+    Token, UseTree,
 };
 
 /// Parse a comma-separated list of identifiers
@@ -29,6 +29,88 @@ struct SeedInfo {
     bump_field: Option<Ident>,
 }
 
+/// Information about imported items from use statements
+#[derive(Debug, Clone)]
+struct ImportInfo {
+    /// Map from local name to full path
+    imports: std::collections::HashMap<String, String>,
+}
+
+impl ImportInfo {
+    fn new() -> Self {
+        Self {
+            imports: std::collections::HashMap::new(),
+        }
+    }
+
+    fn add_import(&mut self, local_name: String, full_path: String) {
+        self.imports.insert(local_name, full_path);
+    }
+
+    fn resolve_type(&self, type_name: &str) -> Option<&String> {
+        self.imports.get(type_name)
+    }
+}
+
+/// Parse use statements to understand imported items
+fn parse_use_statements(module_items: &[Item]) -> ImportInfo {
+    let mut import_info = ImportInfo::new();
+
+    for item in module_items {
+        if let Item::Use(item_use) = item {
+            extract_imports_from_use_tree(&item_use.tree, &mut import_info, String::new());
+        }
+    }
+
+    import_info
+}
+
+/// Recursively extract imports from a use tree
+fn extract_imports_from_use_tree(
+    use_tree: &UseTree,
+    import_info: &mut ImportInfo,
+    base_path: String,
+) {
+    match use_tree {
+        UseTree::Path(use_path) => {
+            let new_base = if base_path.is_empty() {
+                use_path.ident.to_string()
+            } else {
+                format!("{}::{}", base_path, use_path.ident)
+            };
+            extract_imports_from_use_tree(&use_path.tree, import_info, new_base);
+        }
+        UseTree::Name(use_name) => {
+            let local_name = use_name.ident.to_string();
+            let full_path = if base_path.is_empty() {
+                local_name.clone()
+            } else {
+                format!("{}::{}", base_path, local_name)
+            };
+            import_info.add_import(local_name, full_path);
+        }
+        UseTree::Rename(use_rename) => {
+            let local_name = use_rename.rename.to_string();
+            let full_path = if base_path.is_empty() {
+                use_rename.ident.to_string()
+            } else {
+                format!("{}::{}", base_path, use_rename.ident)
+            };
+            import_info.add_import(local_name, full_path);
+        }
+        UseTree::Glob(_) => {
+            // For glob imports, we can't easily resolve specific items
+            // In this case, we'll add a special marker for the base path
+            import_info.add_import("*".to_string(), base_path);
+        }
+        UseTree::Group(use_group) => {
+            for tree in &use_group.items {
+                extract_imports_from_use_tree(tree, import_info, base_path.clone());
+            }
+        }
+    }
+}
+
 /// Extract instruction parameter names from #[instruction(...)] attribute
 fn extract_instruction_param_names(attrs: &[Attribute]) -> Vec<String> {
     for attr in attrs {
@@ -79,8 +161,40 @@ fn has_accounts_derive(attrs: &[Attribute]) -> bool {
     })
 }
 
-/// Scan module items to find account structs that initialize the given account type
-fn find_account_seeds_for_type(
+/// Enhanced function to find seeds that can handle imports and re-exports
+fn find_account_seeds_for_type_enhanced(
+    module_items: &[Item],
+    account_type: &Ident,
+    import_info: &ImportInfo,
+) -> Result<Option<SeedInfo>> {
+    // First, try the original approach (look for directly defined structs)
+    if let Some(seeds_info) = find_account_seeds_for_type_original(module_items, account_type)? {
+        return Ok(Some(seeds_info));
+    }
+
+    // Then, try to find imported or re-exported structs
+    for item in module_items {
+        if let Item::Struct(item_struct) = item {
+            if has_accounts_derive(&item_struct.attrs) {
+                if let syn::Fields::Named(fields) = &item_struct.fields {
+                    for field in &fields.named {
+                        // Try to match field types with account_type, considering imports
+                        if let Some(seeds_info) =
+                            extract_seeds_from_field_enhanced(field, account_type, import_info)?
+                        {
+                            return Ok(Some(seeds_info));
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    Ok(None)
+}
+
+/// Original seed finding function (for direct definitions)
+fn find_account_seeds_for_type_original(
     module_items: &[Item],
     account_type: &Ident,
 ) -> Result<Option<SeedInfo>> {
@@ -138,6 +252,65 @@ fn matches_account_type(ty: &syn::Type, target_type: &Ident) -> bool {
     }
 }
 
+/// Enhanced type matching that considers imports
+fn matches_account_type_enhanced(
+    ty: &syn::Type,
+    target_type: &Ident,
+    import_info: &ImportInfo,
+) -> bool {
+    // First try the original approach
+    if matches_account_type(ty, target_type) {
+        return true;
+    }
+
+    // Then try with import resolution
+    match ty {
+        syn::Type::Path(type_path) => {
+            if let Some(last_segment) = type_path.path.segments.last() {
+                let account_type_names = ["Account", "AccountLoader", "InterfaceAccount"];
+                if account_type_names.contains(&&*last_segment.ident.to_string()) {
+                    if let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments {
+                        for arg in &args.args {
+                            if let syn::GenericArgument::Type(syn::Type::Path(inner_type)) = arg {
+                                // Try to resolve the inner type through imports
+                                if let Some(inner_segment) = inner_type.path.segments.last() {
+                                    let inner_type_name = inner_segment.ident.to_string();
+
+                                    // Direct match
+                                    if inner_segment.ident == *target_type {
+                                        return true;
+                                    }
+
+                                    // Check if it matches through imports
+                                    if let Some(resolved_path) =
+                                        import_info.resolve_type(&inner_type_name)
+                                    {
+                                        if resolved_path.ends_with(&target_type.to_string()) {
+                                            return true;
+                                        }
+                                    }
+
+                                    // Check if target_type matches through imports
+                                    let target_type_name = target_type.to_string();
+                                    if let Some(resolved_target) =
+                                        import_info.resolve_type(&target_type_name)
+                                    {
+                                        if resolved_target.ends_with(&inner_type_name) {
+                                            return true;
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+            false
+        }
+        _ => false,
+    }
+}
+
 /// Parse account attribute to extract init, seeds, and bump information using proper AST parsing
 fn parse_account_attribute(attr: &Attribute) -> Result<Option<(bool, Vec<Expr>, bool)>> {
     if !attr.path().is_ident("account") {
@@ -262,6 +435,48 @@ fn extract_seeds_from_field(field: &Field, target_type: &Ident) -> Result<Option
     Ok(None)
 }
 
+/// Enhanced version of extract_seeds_from_field that handles imports
+fn extract_seeds_from_field_enhanced(
+    field: &Field,
+    target_type: &Ident,
+    import_info: &ImportInfo,
+) -> Result<Option<SeedInfo>> {
+    // First try the original approach
+    if let Some(seeds_info) = extract_seeds_from_field(field, target_type)? {
+        return Ok(Some(seeds_info));
+    }
+
+    // Then try with import resolution
+    let field_type_matches = matches_account_type_enhanced(&field.ty, target_type, import_info);
+
+    if !field_type_matches {
+        return Ok(None);
+    }
+
+    // Look for account attribute with init and seeds
+    for attr in &field.attrs {
+        if let Some((has_init, seeds, has_bump)) = parse_account_attribute(attr)? {
+            if has_init && !seeds.is_empty() {
+                let bump_field = if has_bump {
+                    Some(format_ident!("bump"))
+                } else {
+                    None
+                };
+
+                // Convert instruction parameter references to account field references
+                let converted_seeds = convert_seed_parameters(seeds, target_type)?;
+
+                return Ok(Some(SeedInfo {
+                    seeds: converted_seeds,
+                    bump_field,
+                }));
+            }
+        }
+    }
+
+    Ok(None)
+}
+
 /// Generate compress instructions for the specified account types (Anchor version)
 pub(crate) fn add_compressible_instructions(
     args: TokenStream,
@@ -277,6 +492,9 @@ pub(crate) fn add_compressible_instructions(
     // Get the module content
     let content = module.content.as_mut().unwrap();
 
+    // Parse import information to handle multi-file structures
+    let import_info = parse_use_statements(&content.1);
+
     // Collect all struct names for the enum
     let struct_names: Vec<_> = ident_list.idents.iter().cloned().collect();
 
@@ -604,13 +822,14 @@ pub(crate) fn add_compressible_instructions(
         let compress_accounts_name = format_ident!("Compress{}", struct_name);
 
         // Find seeds for this account type from existing account structs
-        let seeds_info = find_account_seeds_for_type(&content.1, &struct_name)?
+        let seeds_info = find_account_seeds_for_type_enhanced(&content.1, &struct_name, &import_info)?
             .ok_or_else(|| syn::Error::new_spanned(
                 &struct_name,
                 format!(
                     "No account struct found with 'init' constraint and seeds for type '{}'. \
                     Please ensure you have an account struct (with #[derive(Accounts)]) that initializes \
-                    this account type with seeds specified in the #[account(init, seeds = [...], ...)] attribute.",
+                    this account type with seeds specified in the #[account(init, seeds = [...], ...)] attribute. \
+                    This can be in the same module or imported via 'use' statements.",
                     struct_name
                 )
             ))?;

sdk-tests/anchor-compressible-derived/Cargo.toml

@@ -24,17 +24,17 @@ light-sdk = { workspace = true, features = ["anchor", "idl-build", "anchor-discr
 light-sdk-types = { workspace = true }
 light-sdk-macros = { workspace = true }
 light-hasher = { workspace = true, features = ["solana"] }
-solana-program = { workspace = true }
 light-macros = { workspace = true, features = ["solana"] }
+solana-program = { workspace = true }
 borsh = { workspace = true }
 light-compressed-account = { workspace = true, features = ["solana"] }
 anchor-lang = { workspace = true, features = ["idl-build"] }
 
 [dev-dependencies]
-light-program-test = { workspace = true, features = ["devenv"] }
-light-client = { workspace = true, features = ["devenv"] }
+light-program-test = { workspace = true, features = ["v2"] }
+light-client = { workspace = true, features = ["v2"] }
 light-compressible-client = { workspace = true, features = ["anchor"] }
-light-test-utils = { workspace = true, features = ["devenv"] }
+light-test-utils = { workspace = true }
 tokio = { workspace = true }
 solana-sdk = { workspace = true }
 

sdk-tests/anchor-compressible-derived/src/constraints.rs

@@ -0,0 +1,26 @@
+use anchor_lang::prelude::*;
+
+use crate::state::UserRecord;
+
+#[derive(Accounts)]
+pub struct CreateRecord<'info> {
+    #[account(mut)]
+    pub user: Signer<'info>,
+    #[account(
+        init,
+        payer = user,
+        // Manually add 10 bytes! Discriminator + owner + string len + name +
+        // score + option<compression_info>
+        space = 8 + 32 + 4 + 32 + 8 + 10,
+        seeds = [b"user_record", user.key().as_ref()],
+        bump,
+    )]
+    pub user_record: Account<'info, UserRecord>,
+    /// UNCHECKED: checked via config.
+    #[account(mut)]
+    pub rent_recipient: AccountInfo<'info>,
+    /// The global config account
+    /// UNCHECKED: checked via load_checked.
+    pub config: AccountInfo<'info>,
+    pub system_program: Program<'info, System>,
+}