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::TokenStream;
21use quote::ToTokens;
22use quote::quote;
23use syn::Item;
24use syn::Path;
25use syn::ReturnType;
26use syn::Type;
27use syn::parse_quote;
28use syn::spanned::Spanned;
29
30#[proc_macro_attribute]
31pub fn stacksafe(
32    args: proc_macro::TokenStream,
33    item: proc_macro::TokenStream,
34) -> proc_macro::TokenStream {
35    let args = TokenStream::from(args);
36    let item = TokenStream::from(item);
37    match stacksafe_impl(args, item) {
38        Ok(tokens) => tokens.into(),
39        Err(err) => err.into_compile_error().into(),
40    }
41}
42
43fn stacksafe_impl(args: TokenStream, item: TokenStream) -> syn::Result<TokenStream> {
44    let mut crate_path: Option<Path> = None;
45    let arg_parser = syn::meta::parser(|meta| {
46        if meta.path.is_ident("crate") {
47            if crate_path.is_some() {
48                return Err(meta.error("duplicate attribute parameter `crate`"));
49            }
50            crate_path = Some(meta.value()?.parse()?);
51            Ok(())
52        } else {
53            Err(meta.error(format!(
54                "unknown attribute parameter `{}`",
55                meta.path.to_token_stream()
56            )))
57        }
58    });
59    syn::parse::Parser::parse2(arg_parser, args)?;
60
61    let mut item_fn = match syn::parse2::<Item>(item)? {
62        Item::Fn(item_fn) => item_fn,
63        item => {
64            return Err(syn::Error::new_spanned(
65                item,
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    if item_fn.sig.constness.is_some() {
79        return Err(syn::Error::new(
80            item_fn.sig.constness.span(),
81            "#[stacksafe] does not support const functions",
82        ));
83    }
84
85    let ret = match &item_fn.sig.output {
86        // Closures cannot use `impl Trait` return types, so omit the return
87        // type and let the compiler infer it.
88        ReturnType::Type(_, ty) if matches!(**ty, Type::ImplTrait(_)) => None,
89        ret => Some(ret),
90    };
91
92    let stacksafe_crate = crate_path.unwrap_or_else(|| parse_quote!(::stacksafe));
93    let block = &item_fn.block;
94    let wrapped_block = quote! {
95        {
96            #stacksafe_crate::internal::stacker::maybe_grow(
97                #stacksafe_crate::get_minimum_stack_size(),
98                #stacksafe_crate::get_stack_allocation_size(),
99                #stacksafe_crate::internal::with_protected(move || #ret { #block })
100            )
101        }
102    };
103
104    *item_fn.block = syn::parse2(wrapped_block)?;
105    Ok(item_fn.into_token_stream())
106}
stacksafe_macro/lib.rs

stacksafe_macro/
lib.rs
