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_macro2::Span;
21use proc_macro2::TokenStream;
22use quote::ToTokens;
23use quote::quote;
24use syn::ItemFn;
25use syn::Path;
26use syn::ReturnType;
27use syn::Type;
28use syn::parse_quote;
29use syn::spanned::Spanned;
30
31#[proc_macro_attribute]
32pub fn stacksafe(
33    args: proc_macro::TokenStream,
34    item: proc_macro::TokenStream,
35) -> proc_macro::TokenStream {
36    let args = TokenStream::from(args);
37    let item = TokenStream::from(item);
38    match stacksafe_impl(args, item) {
39        Ok(tokens) => tokens.into(),
40        Err(err) => err.into_compile_error().into(),
41    }
42}
43
44fn stacksafe_impl(args: TokenStream, item: TokenStream) -> syn::Result<TokenStream> {
45    let mut crate_path: Option<Path> = None;
46    let arg_parser = syn::meta::parser(|meta| {
47        if meta.path.is_ident("crate") {
48            crate_path = Some(meta.value()?.parse()?);
49            Ok(())
50        } else {
51            Err(meta.error(format!(
52                "unknown attribute parameter `{}`",
53                meta.path
54                    .get_ident()
55                    .map_or("unknown".to_string(), |i| i.to_string())
56            )))
57        }
58    });
59    syn::parse::Parser::parse2(arg_parser, args)?;
60
61    let item_fn = match syn::parse2::<ItemFn>(item.clone()) {
62        Ok(item) => item,
63        Err(_) => {
64            return Err(syn::Error::new(
65                Span::call_site(),
66                "#[stacksafe] can only be applied to functions",
67            ));
68        }
69    };
70
71    if item_fn.sig.asyncness.is_some() {
72        return Err(syn::Error::new(
73            item_fn.sig.asyncness.span(),
74            "#[stacksafe] does not support async functions",
75        ));
76    }
77
78    let mut item_fn = item_fn;
79    let ret = match &item_fn.sig.output {
80        // impl trait is not supported in closure return type, override with
81        // default, which is inferring.
82        ReturnType::Type(_, ty) if matches!(**ty, Type::ImplTrait(_)) => ReturnType::Default,
83        _ => item_fn.sig.output.clone(),
84    };
85
86    let stacksafe_crate = crate_path.unwrap_or_else(|| parse_quote!(::stacksafe));
87    let block = &item_fn.block;
88    let wrapped_block = quote! {
89        {
90            #stacksafe_crate::internal::stacker::maybe_grow(
91                #stacksafe_crate::get_minimum_stack_size(),
92                #stacksafe_crate::get_stack_allocation_size(),
93                #stacksafe_crate::internal::with_protected(move || #ret { #block })
94            )
95        }
96    };
97
98    *item_fn.block = syn::parse(wrapped_block.into())?;
99    Ok(item_fn.into_token_stream())
100}
stacksafe_macro/lib.rs

stacksafe_macro/
lib.rs
