@layerzerolabs/framework-solana-anchor-trait 0.2.61 → 0.2.63

package/tests/assembly.rs

@@ -3,9 +3,9 @@ mod common;
 #[cfg(test)]
 mod tests {
     use super::common::*;
-    use anchor_trait::{generate_instructions, Assembly};
+    use anchor_trait::{generate_instructions, Assembly, GeneratedInstructions};
     use proc_macro2::TokenStream;
-    use quote::{quote, ToTokens};
+    use quote::{format_ident, quote, ToTokens};
 
     // -----------------------------------------------------------------------
     // Helper
@@ -26,6 +26,7 @@ mod tests {
             instructions,
             extra_mod_attrs,
             source_mod: module.clone(),
+            extend_macro_name: None,
         };
         syn::parse2(assembly.assemble()).expect("assembly output should parse")
     }
@@ -48,7 +49,7 @@ mod tests {
         syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -60,12 +61,12 @@ mod tests {
         syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn set_peer(ctx: Context<MyCustomAccounts>, params: SetPeerParams) -> Result<()> {
+                pub fn set_peer(ctx: &mut Context<MyCustomAccounts>, params: &SetPeerParams) -> Result<()> {
                     Ok(())
                 }
 
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
 
@@ -108,7 +109,6 @@ mod tests {
         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()
@@ -121,10 +121,9 @@ mod tests {
             "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
+        // Default structs are always emitted for provided instructions, even when overridden.
         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")
+            matches!(item, syn::Item::Struct(s) if s.ident == "StdGetPeer" || s.ident == "StdToken" || s.ident == "StdSetPeer")
         });
         assert!(has_default_impl_outside, "default impls should be outside the module");
     }
@@ -132,7 +131,7 @@ mod tests {
     #[test]
     fn extra_mod_attrs_applied_to_module() {
         let module = module_with_required_only();
-        let file = assemble_ok(&module, quote! {}, quote! { #[program] });
+        let file = assemble_ok(&module, quote! {}, quote! {});
         let (_, m) = find_program_mod(&file).expect("module should exist");
 
         let has_program_attr = m
@@ -148,12 +147,12 @@ mod tests {
             #[some_custom_attr]
             mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
         };
-        let file = assemble_ok(&module, quote! {}, quote! { #[program] });
+        let file = assemble_ok(&module, quote! {}, quote! {});
         let (_, m) = find_program_mod(&file).expect("module should exist");
 
         let has_custom_attr = m.attrs.iter().any(|attr| {
@@ -169,9 +168,9 @@ mod tests {
         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.
+        // The override handler is relocated to a wrapper module, so it should not appear
+        // as a function definition inside the program module. The name still appears in
+        // dispatch calls inside entrypoints, so we match on 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"))
@@ -181,7 +180,6 @@ mod tests {
             "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"));
@@ -228,7 +226,6 @@ mod tests {
         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()
@@ -242,7 +239,6 @@ mod tests {
             })
             .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()
@@ -255,7 +251,7 @@ mod tests {
         let module: syn::ItemMod = syn::parse_quote! {
             pub mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -282,10 +278,148 @@ mod tests {
         assert!(has_use_super, "module should contain `use super::*`");
     }
 
+    // -----------------------------------------------------------------------
+    // Tests — extend macro emission
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn extend_macro_emitted_when_configured() {
+        let module = module_with_required_only();
+        let overrides = collect_overrides(&module);
+        let instructions =
+            generate_instructions(&TestCtx, &overrides).expect("generate should succeed");
+        let assembly = Assembly {
+            prelude: quote! {},
+            overrides,
+            instructions,
+            extra_mod_attrs: quote! {},
+            source_mod: module.clone(),
+            extend_macro_name: Some(format_ident!("test_extend_accounts")),
+        };
+        let file: syn::File =
+            syn::parse2(assembly.assemble()).expect("assembly output should parse");
+
+        let has_macro = file.items.iter().any(|item| {
+            matches!(item, syn::Item::Macro(m) if m.ident.as_ref().is_some_and(|id| id == "test_extend_accounts"))
+        });
+        assert!(has_macro, "should emit macro_rules! test_extend_accounts");
+    }
+
+    #[test]
+    fn program_attr_already_present_is_not_duplicated() {
+        // If the source module already carries `#[program]` (from a prior macro pass),
+        // assembly must strip it before re-emitting it as the innermost attribute.
+        let module: syn::ItemMod = syn::parse_quote! {
+            #[program]
+            mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        };
+        let file = assemble_ok(&module, quote! {}, quote! {});
+        let (_, m) = find_program_mod(&file).expect("module should exist");
+
+        let program_count = m.attrs.iter().filter(|a| a.path().is_ident("program")).count();
+        assert_eq!(
+            program_count, 1,
+            "#[program] should appear exactly once, not duplicated: {program_count}"
+        );
+    }
+
+    #[test]
+    fn extend_macro_omitted_when_extend_rules_empty() {
+        // `extend_macro_name = Some(name)` but no extend rules → no macro_rules! emitted.
+        let module = module_with_required_only();
+        let overrides = collect_overrides(&module);
+        let assembly = Assembly {
+            prelude: quote! {},
+            overrides,
+            instructions: GeneratedInstructions::default(),
+            extra_mod_attrs: quote! {},
+            source_mod: module.clone(),
+            extend_macro_name: Some(format_ident!("test_extend_accounts")),
+        };
+        let file: syn::File = syn::parse2(assembly.assemble()).expect("should parse");
+
+        let has_any_macro = file.items.iter().any(|item| matches!(item, syn::Item::Macro(_)));
+        assert!(
+            !has_any_macro,
+            "no macro_rules! should be emitted when extend_rules is empty, even when extend_macro_name is Some"
+        );
+    }
+
+    #[test]
+    fn extend_macro_omitted_when_none() {
+        let module = module_with_required_only();
+        let file = assemble_ok(&module, quote! {}, quote! {});
+
+        let has_any_macro = file.items.iter().any(|item| matches!(item, syn::Item::Macro(_)));
+        assert!(!has_any_macro, "no macro_rules! should be emitted when extend_macro_name is None");
+    }
+
+    #[test]
+    fn extend_macro_positioned_before_program_module() {
+        let module = module_with_required_only();
+        let overrides = collect_overrides(&module);
+        let instructions =
+            generate_instructions(&TestCtx, &overrides).expect("generate should succeed");
+        let assembly = Assembly {
+            prelude: quote! {},
+            overrides,
+            instructions,
+            extra_mod_attrs: quote! {},
+            source_mod: module.clone(),
+            extend_macro_name: Some(format_ident!("test_extend_accounts")),
+        };
+        let file: syn::File =
+            syn::parse2(assembly.assemble()).expect("assembly output should parse");
+
+        let macro_pos = file
+            .items
+            .iter()
+            .position(|item| matches!(item, syn::Item::Macro(_)))
+            .expect("extend macro should exist");
+        let (mod_pos, _) = find_program_mod(&file).expect("module should exist");
+        assert!(
+            macro_pos < mod_pos,
+            "extend macro (pos {macro_pos}) should come before program module (pos {mod_pos})"
+        );
+    }
+
+    #[test]
+    fn extend_macro_contains_rules_for_each_provided_instruction() {
+        let module = module_with_required_only();
+        let overrides = collect_overrides(&module);
+        let instructions =
+            generate_instructions(&TestCtx, &overrides).expect("generate should succeed");
+        let assembly = Assembly {
+            prelude: quote! {},
+            overrides,
+            instructions,
+            extra_mod_attrs: quote! {},
+            source_mod: module.clone(),
+            extend_macro_name: Some(format_ident!("test_extend_accounts")),
+        };
+        let output = assembly.assemble().to_string();
+
+        // Check for the rule-arm matcher pattern `(set_peer ,` rather than just the bare name.
+        // The bare name also appears in entrypoints and wrapper modules, making it a trivially
+        // weak check. The `(name ,` prefix is unique to the macro_rules! arm head.
+        assert!(output.contains("(set_peer ,"), "macro should contain set_peer rule arm: {output}");
+        assert!(output.contains("(get_peer ,"), "macro should contain get_peer rule arm: {output}");
+        assert!(output.contains("(token ,"), "macro should contain token rule arm: {output}");
+    }
+
+    // -----------------------------------------------------------------------
+    // Tests — full integration
+    // -----------------------------------------------------------------------
+
     #[test]
     fn full_assembly_integration() {
         let module = module_with_override_and_extras();
-        let file = assemble_ok(&module, quote! { struct PreludeMarker; }, quote! { #[program] });
+        let file = assemble_ok(&module, quote! { struct PreludeMarker; }, quote! {});
 
         // Verify complete ordering: prelude → wrapper → default impls → module
         let prelude_pos = file
@@ -312,19 +446,16 @@ mod tests {
 
         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()

package/tests/common/mod.rs

@@ -3,12 +3,14 @@
 #![allow(dead_code)]
 
 use anchor_trait::{
-    generate_instructions, DefaultImpl, GeneratedInstructions, InstructionOverrides,
-    InstructionSet, InstructionSpec, OverrideConfig,
+    generate_instructions, DefaultAccounts, DefaultImpl, GeneratedInstructions,
+    InstructionOverrides, InstructionSet, InstructionSpec, OverrideConfig,
 };
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote, ToTokens};
-use syn::{FnArg, GenericArgument, ItemFn, ItemStruct, Pat, PatIdent, PathArguments, Type};
+use syn::{
+    parse_quote, FnArg, GenericArgument, ItemFn, ItemStruct, Pat, PatIdent, PathArguments, Type,
+};
 
 // ---------------------------------------------------------------------------
 // Test instruction set
@@ -31,11 +33,66 @@ pub enum TestIx {
 /// Dummy codegen context.
 pub struct TestCtx;
 
-/// Generates a placeholder default impl (struct only) named after the instruction's
-/// default accounts type, mirroring real codegen which emits instruction-specific structs.
-pub fn gen_default_impl(_ctx: &TestCtx, spec: &InstructionSpec<TestCtx>) -> TokenStream {
-    let ident = &spec.default_impl.as_ref().unwrap().accounts_ident;
-    quote! { struct #ident; }
+fn gen_set_peer(_ctx: &TestCtx, spec: &InstructionSpec<TestCtx>) -> DefaultAccounts {
+    let handler_ident = &spec.default_impl.as_ref().unwrap().handler_ident;
+    let params_type = spec.params_type.as_ref().unwrap();
+    let return_type = &spec.return_type;
+    DefaultAccounts {
+        struct_type: parse_quote! {
+            #[derive(Accounts)]
+            pub struct StdSetPeer<'info> {
+                pub placeholder: UncheckedAccount<'info>,
+            }
+        },
+        impl_block: parse_quote! {
+            impl StdSetPeer<'_> {
+                pub fn #handler_ident(_ctx: &mut Context<Self>, _params: &#params_type) -> Result<#return_type> {
+                    unimplemented!()
+                }
+            }
+        },
+    }
+}
+
+fn gen_get_peer(_ctx: &TestCtx, spec: &InstructionSpec<TestCtx>) -> DefaultAccounts {
+    let handler_ident = &spec.default_impl.as_ref().unwrap().handler_ident;
+    let params_type = spec.params_type.as_ref().unwrap();
+    let return_type = &spec.return_type;
+    DefaultAccounts {
+        struct_type: parse_quote! {
+            #[derive(Accounts)]
+            pub struct StdGetPeer<'info> {
+                pub placeholder: UncheckedAccount<'info>,
+            }
+        },
+        impl_block: parse_quote! {
+            impl StdGetPeer<'_> {
+                pub fn #handler_ident(_ctx: &Context<Self>, _params: &#params_type) -> Result<#return_type> {
+                    unimplemented!()
+                }
+            }
+        },
+    }
+}
+
+fn gen_token(_ctx: &TestCtx, spec: &InstructionSpec<TestCtx>) -> DefaultAccounts {
+    let handler_ident = &spec.default_impl.as_ref().unwrap().handler_ident;
+    let return_type = &spec.return_type;
+    DefaultAccounts {
+        struct_type: parse_quote! {
+            #[derive(Accounts)]
+            pub struct StdToken<'info> {
+                pub placeholder: UncheckedAccount<'info>,
+            }
+        },
+        impl_block: parse_quote! {
+            impl StdToken<'_> {
+                pub fn #handler_ident(_ctx: &Context<Self>) -> Result<#return_type> {
+                    unimplemented!()
+                }
+            }
+        },
+    }
 }
 
 impl InstructionSet for TestIx {
@@ -71,9 +128,8 @@ impl InstructionSet for TestIx {
                 return_type: quote!(()),
                 is_view: false,
                 default_impl: Some(DefaultImpl {
-                    accounts_ident: format_ident!("StdSetPeer"),
                     handler_ident: format_ident!("apply"),
-                    generator: gen_default_impl,
+                    generator: gen_set_peer,
                 }),
             },
             Self::GetPeer => InstructionSpec {
@@ -81,9 +137,8 @@ impl InstructionSet for TestIx {
                 return_type: quote!([u8; 32]),
                 is_view: true,
                 default_impl: Some(DefaultImpl {
-                    accounts_ident: format_ident!("StdGetPeer"),
                     handler_ident: format_ident!("apply"),
-                    generator: gen_default_impl,
+                    generator: gen_get_peer,
                 }),
             },
             Self::LzReceive => InstructionSpec {
@@ -97,9 +152,8 @@ impl InstructionSet for TestIx {
                 return_type: quote!(Pubkey),
                 is_view: true,
                 default_impl: Some(DefaultImpl {
-                    accounts_ident: format_ident!("StdToken"),
                     handler_ident: format_ident!("apply"),
-                    generator: gen_default_impl,
+                    generator: gen_token,
                 }),
             },
         }