anodized_core/instrument/fns/
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, Pat, PatIdent, 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 = Pat::Ident(PatIdent {
40            attrs: vec![],
41            by_ref: None,
42            mutability: None,
43            ident: Ident::new("__anodized_output", Span::mixed_site()),
44            subpat: None,
45        });
46
47        // --- Generate Precondition Checks ---
48        let precondition_checks = spec
49            .requires
50            .iter()
51            .map(|condition| {
52                let closure = condition.closure.to_token_stream();
53                let expr = quote! { (#closure)() };
54                let repr = condition.closure.body.to_token_stream();
55                build_check(
56                    condition.cfg.as_ref(),
57                    &expr,
58                    "Precondition failed: {}",
59                    &repr,
60                )
61            })
62            .chain(spec.maintains.iter().map(|condition| {
63                let closure = condition.closure.to_token_stream();
64                let expr = quote! { (#closure)() };
65                let repr = condition.closure.body.to_token_stream();
66                build_check(
67                    condition.cfg.as_ref(),
68                    &expr,
69                    "Pre-invariant failed: {}",
70                    &repr,
71                )
72            }));
73
74        // --- Generate Combined Body and Capture Statement ---
75        // Capture values and execute body in a single tuple assignment
76        // This ensures captured values aren't accessible to the body itself
77
78        // Chain capture aliases with output binding
79        let aliases = spec
80            .captures
81            .iter()
82            .map(|cb| &cb.pat)
83            .chain(std::iter::once(&output_ident));
84
85        // Chain capture expressions with body expression
86        let capture_exprs = spec.captures.iter().map(|cb| {
87            let expr = &cb.expr;
88            // Evaluate expression in a closure to prevent early return.
89            quote! { (|| #expr)() }
90        });
91
92        // Chain underscore types with return type for tuple type annotation
93        let types = spec
94            .captures
95            .iter()
96            .map(|_| quote! { _ })
97            .chain(std::iter::once(quote! { #return_type }));
98
99        let body_expr = if is_async {
100            quote! { (async || #original_body)().await }
101        } else {
102            quote! { (|| #original_body)() }
103        };
104
105        let exprs = capture_exprs.chain(std::iter::once(body_expr));
106
107        // Build tuple assignment with type annotation on the tuple
108        let body_and_captures = quote! {
109            let (#(#aliases),*): (#(#types),*) = (#(#exprs),*);
110        };
111
112        // --- Generate Postcondition Checks ---
113        let postcondition_checks = spec
114            .maintains
115            .iter()
116            .map(|condition| {
117                let closure = condition.closure.to_token_stream();
118                let expr = quote! { (#closure)() };
119                let repr = condition.closure.body.to_token_stream();
120                build_check(
121                    condition.cfg.as_ref(),
122                    &expr,
123                    "Post-invariant failed: {}",
124                    &repr,
125                )
126            })
127            .chain(spec.ensures.iter().map(|postcondition| {
128                let closure = annotate_postcondition_closure_argument(
129                    postcondition.closure.clone(),
130                    return_type.clone(),
131                );
132
133                let expr = quote! { (#closure)(&#output_ident) };
134                build_check(
135                    postcondition.cfg.as_ref(),
136                    &expr,
137                    "Postcondition failed: {}",
138                    &postcondition.closure.to_token_stream(),
139                )
140            }));
141
142        Ok(parse_quote! {
143            {
144                #(#precondition_checks)*
145                #body_and_captures
146                #(#postcondition_checks)*
147                #output_ident
148            }
149        })
150    }
151}
152
153fn annotate_postcondition_closure_argument(
154    mut closure: syn::ExprClosure,
155    return_type: syn::Type,
156) -> syn::ExprClosure {
157    // Add type annotation: convert |param| to |param: &ReturnType|.
158    if let Some(first_input) = closure.inputs.first_mut() {
159        // Wrap the pattern with a type annotation
160        let pattern = first_input.clone();
161        *first_input = syn::Pat::Type(syn::PatType {
162            attrs: vec![],
163            pat: Box::new(pattern),
164            colon_token: Default::default(),
165            ty: Box::new(syn::Type::Reference(syn::TypeReference {
166                and_token: Default::default(),
167                lifetime: None,
168                mutability: None,
169                elem: Box::new(return_type),
170            })),
171        });
172    }
173    closure
174}
anodized_core/instrument/fns/mod.rs

anodized_core/instrument/fns/
mod.rs
