stacksafe_macro/
lib.rs

1// Copyright 2025 FastLabs Developers
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Procedural macro implementation for the `stacksafe` crate.
16//!
17//! This crate provides the `#[stacksafe]` attribute macro that transforms functions
18//! to use automatic stack growth, preventing stack overflow in deeply recursive scenarios.
19
20use proc_macro::TokenStream;
21use proc_macro_error2::abort;
22use proc_macro_error2::abort_call_site;
23use proc_macro_error2::proc_macro_error;
24use quote::ToTokens;
25use quote::quote;
26use syn::ItemFn;
27use syn::Path;
28use syn::ReturnType;
29use syn::Type;
30use syn::parse_macro_input;
31use syn::parse_quote;
32
33#[proc_macro_attribute]
34#[proc_macro_error]
35pub fn stacksafe(args: TokenStream, item: TokenStream) -> TokenStream {
36    let mut crate_path: Option<Path> = None;
37
38    let arg_parser = syn::meta::parser(|meta| {
39        if meta.path.is_ident("crate") {
40            crate_path = Some(meta.value()?.parse()?);
41            Ok(())
42        } else {
43            Err(meta.error(format!(
44                "unknown attribute parameter `{}`",
45                meta.path
46                    .get_ident()
47                    .map_or("unknown".to_string(), |i| i.to_string())
48            )))
49        }
50    });
51    parse_macro_input!(args with arg_parser);
52
53    let item_fn: ItemFn = match syn::parse(item.clone()) {
54        Ok(item) => item,
55        Err(_) => abort_call_site!("#[stacksafe] can only be applied to functions"),
56    };
57
58    if item_fn.sig.asyncness.is_some() {
59        abort!(
60            item_fn.sig.asyncness,
61            "#[stacksafe] does not support async functions"
62        );
63    }
64
65    let mut item_fn = item_fn;
66    let ret = match &item_fn.sig.output {
67        // impl trait is not supported in closure return type, override with
68        // default, which is inferring.
69        ReturnType::Type(_, ty) if matches!(**ty, Type::ImplTrait(_)) => ReturnType::Default,
70        _ => item_fn.sig.output.clone(),
71    };
72
73    let stacksafe_crate = crate_path.unwrap_or_else(|| parse_quote!(::stacksafe));
74    let block = &item_fn.block;
75    let wrapped_block = quote! {
76        {
77            #stacksafe_crate::internal::stacker::maybe_grow(
78                #stacksafe_crate::get_minimum_stack_size(),
79                #stacksafe_crate::get_stack_allocation_size(),
80                #stacksafe_crate::internal::with_protected(move || #ret { #block })
81            )
82        }
83    };
84
85    *item_fn.block = syn::parse(wrapped_block.into()).unwrap();
86    item_fn.into_token_stream().into()
87}
stacksafe_macro/lib.rs

stacksafe_macro/
lib.rs
