stacksafe_macro/
lib.rs

1//! Procedural macro implementation for the `stacksafe` crate.
2//!
3//! This crate provides the `#[stacksafe]` attribute macro that transforms functions
4//! to use automatic stack growth, preventing stack overflow in deeply recursive scenarios.
5
6use proc_macro::TokenStream;
7use proc_macro_error2::abort;
8use proc_macro_error2::abort_call_site;
9use proc_macro_error2::proc_macro_error;
10use quote::ToTokens;
11use quote::quote;
12use syn::ItemFn;
13use syn::Path;
14use syn::ReturnType;
15use syn::Type;
16use syn::parse_macro_input;
17use syn::parse_quote;
18
19#[proc_macro_attribute]
20#[proc_macro_error]
21pub fn stacksafe(args: TokenStream, item: TokenStream) -> TokenStream {
22    let mut crate_path: Option<Path> = None;
23
24    let arg_parser = syn::meta::parser(|meta| {
25        if meta.path.is_ident("crate") {
26            crate_path = Some(meta.value()?.parse()?);
27            Ok(())
28        } else {
29            Err(meta.error(format!(
30                "unknown attribute parameter `{}`",
31                meta.path
32                    .get_ident()
33                    .map_or("unknown".to_string(), |i| i.to_string())
34            )))
35        }
36    });
37    parse_macro_input!(args with arg_parser);
38
39    let item_fn: ItemFn = match syn::parse(item.clone()) {
40        Ok(item) => item,
41        Err(_) => abort_call_site!("#[stacksafe] can only be applied to functions"),
42    };
43
44    if item_fn.sig.asyncness.is_some() {
45        abort!(
46            item_fn.sig.asyncness,
47            "#[stacksafe] does not support async functions"
48        );
49    }
50
51    let mut item_fn = item_fn;
52    let block = item_fn.block;
53    let ret = match &item_fn.sig.output {
54        // impl trait is not supported in closure return type, override with
55        // default, which is inferring.
56        ReturnType::Type(_, ty) if matches!(**ty, Type::ImplTrait(_)) => ReturnType::Default,
57        _ => item_fn.sig.output.clone(),
58    };
59
60    let stacksafe_crate = crate_path.unwrap_or_else(|| parse_quote!(::stacksafe));
61
62    let wrapped_block = quote! {
63        {
64            #stacksafe_crate::internal::stacker::maybe_grow(
65                #stacksafe_crate::get_minimum_stack_size(),
66                #stacksafe_crate::get_stack_allocation_size(),
67                #stacksafe_crate::internal::with_protected(move || #ret { #block })
68            )
69        }
70    };
71    item_fn.block = Box::new(syn::parse(wrapped_block.into()).unwrap());
72    item_fn.into_token_stream().into()
73}
stacksafe_macro/lib.rs

stacksafe_macro/
lib.rs
