diff --git a/sdk-libs/macros/src/light_pdas/program/parsing.rs b/sdk-libs/macros/src/light_pdas/program/parsing.rs index c740be621e..e61cb1ca43 100644 --- a/sdk-libs/macros/src/light_pdas/program/parsing.rs +++ b/sdk-libs/macros/src/light_pdas/program/parsing.rs @@ -617,8 +617,8 @@ fn is_delegation_body(block: &syn::Block, ctx_name: &str) -> bool { syn::Stmt::Expr(expr, _) => { // Check if it's a function call that takes ctx as an argument match expr { - syn::Expr::Call(call) => call_has_ctx_arg(&call.args, ctx_name), - syn::Expr::MethodCall(call) => call_has_ctx_arg(&call.args, ctx_name), + syn::Expr::Call(call) => call_moves_ctx_arg(&call.args, ctx_name), + syn::Expr::MethodCall(call) => call_moves_ctx_arg(&call.args, ctx_name), _ => false, } } @@ -626,6 +626,29 @@ fn is_delegation_body(block: &syn::Block, ctx_name: &str) -> bool { } } +/// Check if any argument consumes the context param. +/// +/// Passing `&ctx` or `&mut ctx` lets the wrapper continue using `ctx` for +/// finalization after the handler returns, so those are intentionally excluded. +fn call_moves_ctx_arg( + args: &syn::punctuated::Punctuated, + ctx_name: &str, +) -> bool { + for arg in args { + match arg { + syn::Expr::Path(path) if path.path.is_ident(ctx_name) => return true, + syn::Expr::MethodCall(method_call) + if method_call.method == "into" + && matches!(&*method_call.receiver, syn::Expr::Path(path) if path.path.is_ident(ctx_name)) => + { + return true; + } + _ => {} + } + } + false +} + /// Check if any argument in the call is the context param (moving the context). /// Detects: ctx, &ctx, &mut ctx, ctx.clone(), ctx.into(), etc. /// `ctx_name` is the context parameter name to look for (e.g., "ctx", "context"). diff --git a/sdk-libs/macros/src/light_pdas_tests/parsing_tests.rs b/sdk-libs/macros/src/light_pdas_tests/parsing_tests.rs index b1135412a9..bbae92cd9a 100644 --- a/sdk-libs/macros/src/light_pdas_tests/parsing_tests.rs +++ b/sdk-libs/macros/src/light_pdas_tests/parsing_tests.rs @@ -2,10 +2,11 @@ //! //! Extracted from `light_pdas/program/parsing.rs`. +use quote::quote; use syn::punctuated::Punctuated; use crate::light_pdas::program::parsing::{ - call_has_ctx_arg, extract_context_and_params, ExtractResult, + call_has_ctx_arg, extract_context_and_params, wrap_function_with_light, ExtractResult, }; fn parse_args(code: &str) -> Punctuated { @@ -216,3 +217,37 @@ fn test_extract_context_and_params_multiple_args_detected() { _ => panic!("Expected ExtractResult::MultipleParams"), } } + +#[test] +fn test_wrap_single_expression_mut_context_runs_finalize() { + let fn_item: syn::ItemFn = syn::parse_quote! { + pub fn handler(ctx: Context, params: Params) -> Result<()> { + instructions::handler(&mut ctx) + } + }; + let params_ident: syn::Ident = syn::parse_quote!(params); + let ctx_ident: syn::Ident = syn::parse_quote!(ctx); + + let wrapped = wrap_function_with_light(&fn_item, ¶ms_ident, &ctx_ident); + let wrapped_tokens = quote!(#wrapped).to_string(); + + assert!(wrapped_tokens.contains("__user_result")); + assert!(wrapped_tokens.contains("light_finalize")); +} + +#[test] +fn test_wrap_moved_context_delegation_skips_finalize() { + let fn_item: syn::ItemFn = syn::parse_quote! { + pub fn handler(ctx: Context, params: Params) -> Result<()> { + instructions::handler(ctx) + } + }; + let params_ident: syn::Ident = syn::parse_quote!(params); + let ctx_ident: syn::Ident = syn::parse_quote!(ctx); + + let wrapped = wrap_function_with_light(&fn_item, ¶ms_ident, &ctx_ident); + let wrapped_tokens = quote!(#wrapped).to_string(); + + assert!(!wrapped_tokens.contains("__user_result")); + assert!(!wrapped_tokens.contains("light_finalize")); +}