@layerzerolabs/framework-solana-anchor-trait 0.2.60

package/src/overrides.rs

@@ -0,0 +1,289 @@
+//! Instruction override collection and validation.
+//!
+//! Discovers `#[*_instruction]` handlers inside a module, validates them,
+//! relocates them into wrapper modules, and provides filtered module items for
+//! final assembly.
+
+use itertools::Itertools;
+use proc_macro2::TokenStream;
+use quote::quote;
+use std::collections::BTreeMap;
+use syn::{spanned::Spanned, Attribute, Item, ItemFn, ItemMod, Type};
+
+use crate::InstructionSet;
+
+/// Static configuration for a macro crate's instruction override system.
+#[derive(Debug, Clone, Copy)]
+pub struct OverrideConfig {
+    /// Attribute name used to mark override handlers, e.g. `"oapp_instruction"`.
+    pub attr_name: &'static str,
+    /// Human-readable domain label for diagnostics, e.g. `"OApp"` or `"RBAC"`.
+    pub domain: &'static str,
+}
+
+/// Metadata captured for one overridden handler function.
+pub struct HandlerMeta {
+    /// Wrapper module identifier for the relocated handler.
+    pub wrapper_mod: syn::Ident,
+    /// Original implementor-defined handler function identifier.
+    pub fn_ident: syn::Ident,
+    /// Inner type `T` extracted from `Context<...T>`.
+    pub accounts_type: Type,
+    /// The full `Context<...>` type as written (preserves lifetime args).
+    pub context_type: Type,
+    /// Function-level generics (e.g. `<'info>`) from the override handler.
+    pub fn_generics: syn::Generics,
+}
+
+/// Collection of custom instruction handlers discovered from a module.
+pub struct InstructionOverrides<Ix: InstructionSet> {
+    handlers: BTreeMap<Ix, HandlerMeta>,
+    pub wrapper_defs: TokenStream,
+    pub config: OverrideConfig,
+}
+
+impl<Ix: InstructionSet> InstructionOverrides<Ix> {
+    /// Creates an empty override collection with the given config.
+    fn new(config: OverrideConfig) -> Self {
+        Self { handlers: BTreeMap::new(), wrapper_defs: TokenStream::new(), config }
+    }
+
+    /// Scans a module for `#[<attr>(...)]` handlers and builds the override collection.
+    pub fn collect(module: &ItemMod, config: OverrideConfig) -> syn::Result<Self> {
+        let mut collected = Self::new(config);
+
+        let items = match &module.content {
+            Some((_, items)) => items,
+            None => return Ok(collected),
+        };
+
+        // Walk every function in the module looking for override attributes.
+        for item in items {
+            let Item::Fn(fn_item) = item else { continue };
+            let Some(target) = parse_override_target::<Ix>(fn_item, config)? else {
+                continue;
+            };
+
+            // Reject duplicate overrides for the same instruction.
+            if collected.handlers.contains_key(&target) {
+                return Err(syn::Error::new_spanned(
+                    &fn_item.sig.ident,
+                    format!(
+                        "duplicate #[{}] for instruction `{}` is not allowed",
+                        config.attr_name,
+                        target.canonical_name(),
+                    ),
+                ));
+            }
+
+            // Relocate the handler into a wrapper module and record its metadata.
+            let meta = collected.wrap_handler(target, fn_item)?;
+            collected.handlers.insert(target, meta);
+        }
+
+        // Ensure no plain `fn name()` shadows a known instruction name.
+        collected.reject_name_collisions(module)?;
+        Ok(collected)
+    }
+
+    /// Returns metadata for an overridden instruction, if present.
+    pub fn get(&self, target: Ix) -> Option<&HandlerMeta> {
+        self.handlers.get(&target)
+    }
+
+    /// Returns module items with relocated override handlers removed.
+    pub fn filtered_items(&self, module: &ItemMod) -> TokenStream {
+        let Some((_, items)) = &module.content else {
+            return TokenStream::new();
+        };
+
+        let kept: Vec<_> = items
+            .iter()
+            .filter(|item| match item {
+                Item::Fn(fn_item) => !self.is_overridden(&fn_item.sig.ident),
+                _ => true,
+            })
+            .collect();
+
+        quote!(#(#kept)*)
+    }
+
+    /// Checks whether a function ident belongs to an override handler.
+    fn is_overridden(&self, fn_ident: &syn::Ident) -> bool {
+        self.handlers.values().any(|h| h.fn_ident == *fn_ident)
+    }
+
+    /// Relocates a handler function into an isolated wrapper module
+    /// and returns the corresponding [`HandlerMeta`].
+    fn wrap_handler(&mut self, target: Ix, fn_item: &ItemFn) -> syn::Result<HandlerMeta> {
+        let (accounts_type, context_type) = infer_context_type(fn_item)?;
+        let fn_ident = fn_item.sig.ident.clone();
+        let fn_generics = fn_item.sig.generics.clone();
+
+        // Generate a deterministic module name: `__<attr>__<instruction>_handler`.
+        let wrapper_mod = syn::Ident::new(
+            &format!("__{}__{}_handler", self.config.attr_name, target.canonical_name()),
+            proc_macro2::Span::call_site(),
+        );
+
+        // Clone the function without the instruction attribute.
+        let wrapped_fn = strip_override_attr(fn_item, self.config)?;
+
+        // Emit the wrapper module so the handler lives in its own scope.
+        self.wrapper_defs.extend(quote! {
+            mod #wrapper_mod {
+                use super::*;
+                #wrapped_fn
+            }
+        });
+
+        Ok(HandlerMeta { wrapper_mod, fn_ident, accounts_type, context_type, fn_generics })
+    }
+
+    /// Rejects unattributed functions whose name collides with a known instruction.
+    ///
+    /// A bare `fn lz_receive(...)` without `#[oapp_instruction]`
+    /// would silently shadow the generated entrypoint — catch that early.
+    fn reject_name_collisions(&self, module: &ItemMod) -> syn::Result<()> {
+        let Some((_, items)) = &module.content else {
+            return Ok(());
+        };
+
+        for &ix in Ix::all() {
+            let name = ix.canonical_name();
+            let collides = items.iter().any(|item| match item {
+                Item::Fn(f) => f.sig.ident == name && !has_override_attr(f, self.config),
+                _ => false,
+            });
+
+            if collides {
+                return Err(syn::Error::new(
+                    module.span(),
+                    format!(
+                        "`fn {name}(...)` conflicts with {domain} instruction `{name}`: \
+                         use #[{attr}] to override it",
+                        domain = self.config.domain,
+                        attr = self.config.attr_name,
+                    ),
+                ));
+            }
+        }
+        Ok(())
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Helper functions
+// ---------------------------------------------------------------------------
+
+/// Returns `true` if the function carries the override attribute.
+fn has_override_attr(fn_item: &ItemFn, cfg: OverrideConfig) -> bool {
+    fn_item.attrs.iter().any(|attr| is_override_attr(attr, cfg))
+}
+
+/// Returns `true` if `attr`'s last path segment matches the configured attribute name.
+fn is_override_attr(attr: &Attribute, cfg: OverrideConfig) -> bool {
+    attr.path().segments.last().is_some_and(|seg| seg.ident == cfg.attr_name)
+}
+
+/// Returns a clone of `fn_item` with the override attribute removed.
+fn strip_override_attr(fn_item: &ItemFn, cfg: OverrideConfig) -> syn::Result<ItemFn> {
+    let mut cleaned = fn_item.clone();
+    let before = cleaned.attrs.len();
+    cleaned.attrs.retain(|attr| !is_override_attr(attr, cfg));
+
+    if cleaned.attrs.len() == before {
+        return Err(syn::Error::new_spanned(
+            fn_item,
+            format!("internal error: expected #[{}] attribute on handler function", cfg.attr_name),
+        ));
+    }
+    Ok(cleaned)
+}
+
+/// Parses the override target from a function's override attribute.
+///
+/// The instruction name is inferred from the handler's function name
+/// (e.g. `fn set_paused(...)` → instruction `set_paused`).
+/// The attribute must not contain any arguments.
+///
+/// Returns `Ok(None)` if the function has no such attribute,
+/// `Ok(Some(ix))` on success, or `Err` on malformed attributes.
+fn parse_override_target<Ix: InstructionSet>(
+    fn_item: &ItemFn,
+    cfg: OverrideConfig,
+) -> syn::Result<Option<Ix>> {
+    // Ensure at most one override attribute per function.
+    let attr = fn_item
+        .attrs
+        .iter()
+        .filter(|attr| is_override_attr(attr, cfg))
+        .at_most_one()
+        .map_err(|_| {
+            syn::Error::new_spanned(fn_item, format!("duplicate #[{}] on function", cfg.attr_name))
+        })?;
+
+    // Early return if no override attribute is found.
+    let Some(attr) = attr else {
+        return Ok(None);
+    };
+
+    // Reject arguments — the instruction name is inferred from the function name.
+    if !matches!(attr.meta, syn::Meta::Path(_)) {
+        return Err(syn::Error::new_spanned(
+            attr,
+            format!(
+                "#[{}] does not accept arguments; the instruction name \
+                 is inferred from the handler function name",
+                cfg.attr_name,
+            ),
+        ));
+    }
+
+    // Infer the instruction from the handler's function name.
+    let fn_ident = &fn_item.sig.ident;
+    let ix = Ix::from_ident(fn_ident).ok_or_else(|| {
+        syn::Error::new_spanned(
+            fn_ident,
+            format!("unsupported {} target `{fn_ident}`", cfg.attr_name,),
+        )
+    })?;
+
+    Ok(Some(ix))
+}
+
+/// Extracts the inner `T` from the first `Context<T>` parameter.
+///
+/// Returns `(accounts_type, context_type)` — the inner `T` (for default-context
+/// checks) and the full `Context<...>` type as written (to preserve lifetimes in
+/// the generated entrypoint).
+fn infer_context_type(fn_item: &ItemFn) -> syn::Result<(Type, Type)> {
+    let Some(syn::FnArg::Typed(first)) = fn_item.sig.inputs.first() else {
+        return Err(syn::Error::new_spanned(
+            &fn_item.sig,
+            "handler must have at least one parameter of type `Context<T>`",
+        ));
+    };
+
+    let context_type = (*first.ty).clone();
+    let accounts_type = extract_context_inner_type(&first.ty).ok_or_else(|| {
+        syn::Error::new_spanned(&first.ty, "expected `Context<T>` as the first parameter")
+    })?;
+    Ok((accounts_type, context_type))
+}
+
+/// Unwraps `Context<T>` or `Context<'a, 'b, 'c, 'info, T>` → `T`.
+///
+/// Handles both the simple `Context<T>` form and the lifetime-annotated
+/// `Context<'_, '_, '_, 'info, T>` form by finding the last type argument.
+fn extract_context_inner_type(ty: &Type) -> Option<Type> {
+    let Type::Path(tp) = ty else { return None };
+    let seg = tp.path.segments.last().filter(|s| s.ident == "Context")?;
+    let syn::PathArguments::AngleBracketed(args) = &seg.arguments else {
+        return None;
+    };
+    args.args.iter().rev().find_map(|arg| match arg {
+        syn::GenericArgument::Type(inner) => Some(inner.clone()),
+        _ => None,
+    })
+}

package/tests/assembly.rs

@@ -0,0 +1,339 @@
+mod common;
+
+#[cfg(test)]
+mod tests {
+    use super::common::*;
+    use anchor_trait::{generate_instructions, Assembly};
+    use proc_macro2::TokenStream;
+    use quote::{quote, ToTokens};
+
+    // -----------------------------------------------------------------------
+    // Helper
+    // -----------------------------------------------------------------------
+
+    /// Run the full pipeline (collect → generate → assemble) and parse the output.
+    fn assemble_ok(
+        module: &syn::ItemMod,
+        prelude: TokenStream,
+        extra_mod_attrs: TokenStream,
+    ) -> syn::File {
+        let overrides = collect_overrides(module);
+        let instructions =
+            generate_instructions(&TestCtx, &overrides).expect("generate should succeed");
+        let assembly = Assembly {
+            prelude,
+            overrides,
+            instructions,
+            extra_mod_attrs,
+            source_mod: module.clone(),
+        };
+        syn::parse2(assembly.assemble()).expect("assembly output should parse")
+    }
+
+    /// Find the program module in a parsed file (skips wrapper modules).
+    fn find_program_mod(file: &syn::File) -> Option<(usize, &syn::ItemMod)> {
+        file.items.iter().enumerate().find_map(|(i, item)| match item {
+            syn::Item::Mod(m) if !m.ident.to_string().starts_with("__") => Some((i, m)),
+            _ => None,
+        })
+    }
+
+    /// Extract items from inside a module's body.
+    fn mod_items(m: &syn::ItemMod) -> &Vec<syn::Item> {
+        &m.content.as_ref().expect("module should have a body").1
+    }
+
+    /// A module with only the required lz_receive override.
+    fn module_with_required_only() -> syn::ItemMod {
+        syn::parse_quote! {
+            mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        }
+    }
+
+    /// A module with a full override on set_peer + required lz_receive + non-handler items.
+    fn module_with_override_and_extras() -> syn::ItemMod {
+        syn::parse_quote! {
+            mod my_program {
+                #[test_instruction]
+                pub fn set_peer(ctx: Context<MyCustomAccounts>, params: SetPeerParams) -> Result<()> {
+                    Ok(())
+                }
+
+                #[test_instruction]
+                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+
+                pub fn helper_function() -> u64 {
+                    42
+                }
+
+                struct MyState {
+                    value: u64,
+                }
+            }
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Tests — output structure and ordering
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn prelude_emitted_before_module() {
+        let module = module_with_required_only();
+        let file = assemble_ok(&module, quote! { struct PreludeMarker; }, quote! {});
+
+        let prelude_pos = file
+            .items
+            .iter()
+            .position(|item| matches!(item, syn::Item::Struct(s) if s.ident == "PreludeMarker"))
+            .expect("prelude marker should exist");
+        let (mod_pos, _) = find_program_mod(&file).expect("module should exist");
+
+        assert!(
+            prelude_pos < mod_pos,
+            "prelude (pos {prelude_pos}) should come before module (pos {mod_pos})"
+        );
+    }
+
+    #[test]
+    fn wrapper_defs_and_default_impls_outside_module() {
+        let module = module_with_override_and_extras();
+        let file = assemble_ok(&module, quote! {}, quote! {});
+        let (mod_pos, _) = find_program_mod(&file).expect("module should exist");
+
+        // Wrapper module for the override should be outside and before the program module
+        let wrapper_pos = file
+            .items
+            .iter()
+            .position(|item| {
+                matches!(item, syn::Item::Mod(m) if m.ident.to_string().contains("__test_instruction__"))
+            })
+            .expect("wrapper module should exist");
+        assert!(
+            wrapper_pos < mod_pos,
+            "wrapper defs (pos {wrapper_pos}) should come before module (pos {mod_pos})"
+        );
+
+        // Default impls (e.g. StdGetPeer, StdToken) should be outside the module
+        // set_peer is fully overridden so no StdSetPeer
+        let has_default_impl_outside = file.items[..mod_pos].iter().any(|item| {
+            matches!(item, syn::Item::Struct(s) if s.ident == "StdGetPeer" || s.ident == "StdToken")
+        });
+        assert!(has_default_impl_outside, "default impls should be outside the module");
+    }
+
+    #[test]
+    fn extra_mod_attrs_applied_to_module() {
+        let module = module_with_required_only();
+        let file = assemble_ok(&module, quote! {}, quote! { #[program] });
+        let (_, m) = find_program_mod(&file).expect("module should exist");
+
+        let has_program_attr = m
+            .attrs
+            .iter()
+            .any(|attr| attr.path().segments.last().is_some_and(|seg| seg.ident == "program"));
+        assert!(has_program_attr, "module should have #[program] attribute");
+    }
+
+    #[test]
+    fn original_mod_attrs_preserved() {
+        let module: syn::ItemMod = syn::parse_quote! {
+            #[some_custom_attr]
+            mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        };
+        let file = assemble_ok(&module, quote! {}, quote! { #[program] });
+        let (_, m) = find_program_mod(&file).expect("module should exist");
+
+        let has_custom_attr = m.attrs.iter().any(|attr| {
+            attr.path().segments.last().is_some_and(|seg| seg.ident == "some_custom_attr")
+        });
+        assert!(has_custom_attr, "original module attributes should be preserved");
+    }
+
+    #[test]
+    fn filtered_items_inside_module() {
+        let module = module_with_override_and_extras();
+        let file = assemble_ok(&module, quote! {}, quote! {});
+        let (_, m) = find_program_mod(&file).expect("module should exist");
+        let items = mod_items(m);
+
+        // Override handler should NOT be a function definition inside the module
+        // (it was relocated to wrapper). Note: the name still appears in dispatch
+        // calls inside entrypoints, so we check for function definitions specifically.
+        let has_override_fn = items.iter().any(|item| {
+            matches!(item, syn::Item::Fn(f) if f.sig.ident == "set_peer"
+                && f.block.to_token_stream().to_string().contains("MyCustomAccounts"))
+        });
+        assert!(
+            !has_override_fn,
+            "overridden handler with custom body should be filtered out of module body"
+        );
+
+        // Non-handler items should be preserved inside the module
+        let has_helper = items
+            .iter()
+            .any(|item| matches!(item, syn::Item::Fn(f) if f.sig.ident == "helper_function"));
+        assert!(has_helper, "helper_function should be inside module");
+
+        let has_struct = items
+            .iter()
+            .any(|item| matches!(item, syn::Item::Struct(s) if s.ident == "MyState"));
+        assert!(has_struct, "MyState struct should be inside module");
+    }
+
+    #[test]
+    fn entrypoints_inside_module() {
+        let module = module_with_required_only();
+        let file = assemble_ok(&module, quote! {}, quote! {});
+        let (_, m) = find_program_mod(&file).expect("module should exist");
+        let items = mod_items(m);
+
+        let fns: Vec<_> = items
+            .iter()
+            .filter_map(|item| match item {
+                syn::Item::Fn(f) => Some(f.sig.ident.to_string()),
+                _ => None,
+            })
+            .collect();
+
+        assert!(
+            fns.contains(&"set_peer".to_string()),
+            "set_peer entrypoint should be inside module"
+        );
+        assert!(
+            fns.contains(&"get_peer".to_string()),
+            "get_peer entrypoint should be inside module"
+        );
+        assert!(
+            fns.contains(&"lz_receive".to_string()),
+            "lz_receive entrypoint should be inside module"
+        );
+        assert!(fns.contains(&"token".to_string()), "token entrypoint should be inside module");
+    }
+
+    #[test]
+    fn wrapper_defs_contain_override_handlers() {
+        let module = module_with_override_and_extras();
+        let file = assemble_ok(&module, quote! {}, quote! {});
+
+        // Find the wrapper module for set_peer in the assembled output
+        let wrapper = file
+            .items
+            .iter()
+            .find_map(|item| match item {
+                syn::Item::Mod(m)
+                    if m.ident.to_string().contains("__test_instruction__set_peer") =>
+                {
+                    Some(m)
+                },
+                _ => None,
+            })
+            .expect("wrapper module for set_peer should exist in assembled output");
+
+        // The relocated handler should be inside the wrapper
+        let wrapper_items = mod_items(wrapper);
+        let has_handler = wrapper_items
+            .iter()
+            .any(|item| matches!(item, syn::Item::Fn(f) if f.sig.ident == "set_peer"));
+        assert!(has_handler, "wrapper should contain the relocated override handler");
+    }
+
+    #[test]
+    fn module_visibility_preserved() {
+        let module: syn::ItemMod = syn::parse_quote! {
+            pub mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        };
+        let file = assemble_ok(&module, quote! {}, quote! {});
+        let (_, m) = find_program_mod(&file).expect("module should exist");
+
+        assert!(
+            matches!(m.vis, syn::Visibility::Public(_)),
+            "pub visibility should be preserved on module"
+        );
+    }
+
+    #[test]
+    fn use_super_star_injected_into_module() {
+        let module = module_with_required_only();
+        let file = assemble_ok(&module, quote! {}, quote! {});
+        let (_, m) = find_program_mod(&file).expect("module should exist");
+        let items = mod_items(m);
+
+        let has_use_super = items.iter().any(|item| {
+            matches!(item, syn::Item::Use(u) if u.to_token_stream().to_string().contains("super :: *"))
+        });
+        assert!(has_use_super, "module should contain `use super::*`");
+    }
+
+    #[test]
+    fn full_assembly_integration() {
+        let module = module_with_override_and_extras();
+        let file = assemble_ok(&module, quote! { struct PreludeMarker; }, quote! { #[program] });
+
+        // Verify complete ordering: prelude → wrapper → default impls → module
+        let prelude_pos = file
+            .items
+            .iter()
+            .position(|item| matches!(item, syn::Item::Struct(s) if s.ident == "PreludeMarker"))
+            .expect("prelude should exist");
+
+        let wrapper_pos = file
+            .items
+            .iter()
+            .position(|item| {
+                matches!(item, syn::Item::Mod(m) if m.ident.to_string().contains("__test_instruction__"))
+            })
+            .expect("wrapper should exist");
+
+        let context_pos = file
+            .items
+            .iter()
+            .position(|item| {
+                matches!(item, syn::Item::Struct(s) if s.ident == "StdGetPeer" || s.ident == "StdToken")
+            })
+            .expect("default impl struct should exist");
+
+        let (mod_pos, m) = find_program_mod(&file).expect("module should exist");
+
+        // Ordering: prelude < wrapper < default impls < module
+        assert!(prelude_pos < wrapper_pos, "prelude before wrapper");
+        assert!(wrapper_pos < context_pos, "wrapper before default impls");
+        assert!(context_pos < mod_pos, "default impls before module");
+
+        // Module has #[program]
+        let has_program = m
+            .attrs
+            .iter()
+            .any(|attr| attr.path().segments.last().is_some_and(|seg| seg.ident == "program"));
+        assert!(has_program, "module should have #[program]");
+
+        // Entrypoints inside module
+        let items = mod_items(m);
+        let fn_names: Vec<_> = items
+            .iter()
+            .filter_map(|item| match item {
+                syn::Item::Fn(f) => Some(f.sig.ident.to_string()),
+                _ => None,
+            })
+            .collect();
+        assert!(fn_names.contains(&"set_peer".to_string()), "set_peer inside module");
+        assert!(fn_names.contains(&"token".to_string()), "token inside module");
+    }
+}