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/// Attribute macro for automatic stack overflow prevention in recursive functions.
20///
21/// This macro transforms functions to automatically check available stack space
22/// and allocate new stack segments when needed, preventing stack overflow in
23/// deeply recursive scenarios.
24///
25/// # Parameters
26///
27/// The macro accepts an optional `crate` parameter to specify the path to the
28/// stacksafe crate:
29///
30/// ```rust
31/// use stacksafe::stacksafe;
32///
33/// #[stacksafe(crate = stacksafe)]
34/// fn my_function() {
35///     // function body
36/// }
37/// ```
38///
39/// # Examples
40///
41/// ```rust
42/// use stacksafe::stacksafe;
43///
44/// #[stacksafe]
45/// fn factorial(n: u64) -> u64 {
46///     if n <= 1 { 1 } else { n * factorial(n - 1) }
47/// }
48///
49/// #[stacksafe]
50/// fn tree_depth<T>(node: &Option<Box<TreeNode<T>>>) -> usize {
51///     match node {
52///         None => 0,
53///         Some(n) => 1 + tree_depth(&n.left).max(tree_depth(&n.right)),
54///     }
55/// }
56/// ```
57///
58/// # Limitations
59///
60/// - Cannot be applied to `async` functions
61/// - Functions with `impl Trait` return types may need type annotations
62/// - Adds small runtime overhead for stack size checking
63#[proc_macro_attribute]
64#[proc_macro_error]
65pub fn stacksafe(args: TokenStream, item: TokenStream) -> TokenStream {
66    let mut crate_path: Option<Path> = None;
67
68    let arg_parser = syn::meta::parser(|meta| {
69        if meta.path.is_ident("crate") {
70            crate_path = Some(meta.value()?.parse()?);
71            Ok(())
72        } else {
73            Err(meta.error(format!(
74                "unknown attribute parameter `{}`",
75                meta.path
76                    .get_ident()
77                    .map_or("unknown".to_string(), |i| i.to_string())
78            )))
79        }
80    });
81    parse_macro_input!(args with arg_parser);
82
83    let item_fn: ItemFn = match syn::parse(item.clone()) {
84        Ok(item) => item,
85        Err(_) => abort_call_site!("#[stacksafe] can only be applied to functions"),
86    };
87
88    if item_fn.sig.asyncness.is_some() {
89        abort!(
90            item_fn.sig.asyncness,
91            "#[stacksafe] does not support async functions"
92        );
93    }
94
95    let mut item_fn = item_fn;
96    let block = item_fn.block;
97    let ret = match &item_fn.sig.output {
98        // impl trait is not supported in closure return type, override with
99        // default, which is inferring.
100        ReturnType::Type(_, typ) if matches!(**typ, Type::ImplTrait(_)) => ReturnType::Default,
101        _ => item_fn.sig.output.clone(),
102    };
103
104    let stacksafe_crate = crate_path.unwrap_or_else(|| parse_quote!(::stacksafe));
105
106    let wrapped_block = quote! {
107        {
108            #stacksafe_crate::internal::stacker::maybe_grow(
109                #stacksafe_crate::get_minimum_stack_size(),
110                #stacksafe_crate::get_stack_allocation_size(),
111                #stacksafe_crate::internal::with_protected(move || #ret { #block })
112            )
113        }
114    };
115    item_fn.block = Box::new(syn::parse(wrapped_block.into()).unwrap());
116    item_fn.into_token_stream().into()
117}
stacksafe_macro/lib.rs

stacksafe_macro/
lib.rs
