anodized_core/instrument/function/
mod.rs

1#[cfg(test)]
2mod tests;
3
4use crate::{Spec, instrument::Backend};
5
6use proc_macro2::Span;
7use quote::{ToTokens, quote};
8use syn::{Block, Ident, ItemFn, parse::Result, parse_quote};
9
10impl Backend {
11    pub fn instrument_fn(self, spec: Spec, mut func: ItemFn) -> syn::Result<ItemFn> {
12        let is_async = func.sig.asyncness.is_some();
13
14        // Extract the return type from the function signature
15        let return_type = match &func.sig.output {
16            syn::ReturnType::Default => syn::parse_quote!(()),
17            syn::ReturnType::Type(_, ty) => ty.as_ref().clone(),
18        };
19
20        // Generate the new, instrumented function body.
21        let new_body = self.instrument_fn_body(&spec, &func.block, is_async, &return_type)?;
22
23        // Replace the old function body with the new one.
24        *func.block = new_body;
25
26        Ok(func)
27    }
28
29    fn instrument_fn_body(
30        self,
31        spec: &Spec,
32        original_body: &Block,
33        is_async: bool,
34        return_type: &syn::Type,
35    ) -> Result<Block> {
36        let build_check = self.build_check;
37
38        // The identifier for the return value binding.
39        let output_ident = Ident::new("__anodized_output", Span::mixed_site());
40
41        // --- Generate Precondition Checks ---
42        let precondition_checks = spec
43            .requires
44            .iter()
45            .map(|condition| {
46                let expr = condition.expr.to_token_stream();
47                build_check(
48                    condition.cfg.as_ref(),
49                    &expr,
50                    "Precondition failed: {}",
51                    &expr,
52                )
53            })
54            .chain(spec.maintains.iter().map(|condition| {
55                let expr = condition.expr.to_token_stream();
56                build_check(
57                    condition.cfg.as_ref(),
58                    &expr,
59                    "Pre-invariant failed: {}",
60                    &expr,
61                )
62            }));
63
64        // --- Generate Combined Body and Capture Statement ---
65        // Capture values and execute body in a single tuple assignment
66        // This ensures captured values aren't accessible to the body itself
67
68        // Chain capture aliases with output binding
69        let aliases = spec
70            .captures
71            .iter()
72            .map(|cb| &cb.alias)
73            .chain(std::iter::once(&output_ident));
74
75        // Chain capture expressions with body expression
76        let capture_exprs = spec.captures.iter().map(|cb| {
77            let expr = &cb.expr;
78            quote! { #expr }
79        });
80
81        // Chain underscore types with return type for tuple type annotation
82        let types = spec
83            .captures
84            .iter()
85            .map(|_| quote! { _ })
86            .chain(std::iter::once(quote! { #return_type }));
87
88        let body_expr = if is_async {
89            quote! { (async || #original_body)().await }
90        } else {
91            quote! { (|| #original_body)() }
92        };
93
94        let exprs = capture_exprs.chain(std::iter::once(body_expr));
95
96        // Build tuple assignment with type annotation on the tuple
97        let body_and_captures = quote! {
98            let (#(#aliases),*): (#(#types),*) = (#(#exprs),*);
99        };
100
101        // --- Generate Postcondition Checks ---
102        let postcondition_checks = spec
103            .maintains
104            .iter()
105            .map(|condition| {
106                let expr = condition.expr.to_token_stream();
107                build_check(
108                    condition.cfg.as_ref(),
109                    &expr,
110                    "Post-invariant failed: {}",
111                    &expr,
112                )
113            })
114            .chain(spec.ensures.iter().map(|postcondition| {
115                let closure = annotate_postcondition_closure_argument(
116                    postcondition.closure.clone(),
117                    return_type.clone(),
118                );
119
120                let expr = quote! { (#closure)(&#output_ident) };
121                build_check(
122                    postcondition.cfg.as_ref(),
123                    &expr,
124                    "Postcondition failed: {}",
125                    &postcondition.closure.to_token_stream(),
126                )
127            }));
128
129        Ok(parse_quote! {
130            {
131                #(#precondition_checks)*
132                #body_and_captures
133                #(#postcondition_checks)*
134                #output_ident
135            }
136        })
137    }
138}
139
140fn annotate_postcondition_closure_argument(
141    mut closure: syn::ExprClosure,
142    return_type: syn::Type,
143) -> syn::ExprClosure {
144    // Add type annotation: convert |param| to |param: &ReturnType|.
145    if let Some(first_input) = closure.inputs.first_mut() {
146        // Wrap the pattern with a type annotation
147        let pattern = first_input.clone();
148        *first_input = syn::Pat::Type(syn::PatType {
149            attrs: vec![],
150            pat: Box::new(pattern),
151            colon_token: Default::default(),
152            ty: Box::new(syn::Type::Reference(syn::TypeReference {
153                and_token: Default::default(),
154                lifetime: None,
155                mutability: None,
156                elem: Box::new(return_type),
157            })),
158        });
159    }
160    closure
161}
anodized_core/instrument/function/mod.rs

anodized_core/instrument/function/
mod.rs
