anodized_core/instrument/
loops.rs

1#[cfg(test)]
2#[path = "loops_tests.rs"]
3mod loops_tests;
4
5use proc_macro2::Span;
6use syn::{
7    Block, Error, Expr, ExprClosure, ExprForLoop, ExprWhile, Ident, ItemFn, Result, Stmt,
8    parse_quote,
9    visit_mut::{self, VisitMut},
10};
11
12use crate::{
13    LoopSpec,
14    instrument::{Mode, find_spec_attr},
15};
16
17impl Mode {
18    pub fn instrument_loops_in_fn_body(&self, body: &mut Block) -> Result<()> {
19        let mut visitor = LoopSpecVisitor::new(self);
20        visitor.visit_block_mut(body);
21        visitor.finish()
22    }
23
24    pub fn instrument_expr_while(&self, spec: LoopSpec, expr_while: &mut ExprWhile) {
25        self.instrument_loop_body(spec, &mut expr_while.body.stmts);
26    }
27
28    pub fn instrument_expr_for_loop(&self, spec: LoopSpec, expr_for_loop: &mut ExprForLoop) {
29        self.instrument_loop_body(spec, &mut expr_for_loop.body.stmts);
30    }
31
32    fn instrument_loop_body(&self, spec: LoopSpec, stmts: &mut Vec<Stmt>) {
33        if let Self::EmbedSpecs = self {
34            let maintains_block = Self::build_precondition_fn_body(&[], &spec.maintains);
35            stmts.insert(
36                0,
37                parse_quote! {
38                    let __anodized_loop_maintains = || -> bool #maintains_block;
39                },
40            );
41
42            let mut variant_stmts: Vec<Stmt> = Vec::new();
43            let mut variant_names: Vec<Ident> = Vec::new();
44            if let Some(loop_variant) = &spec.decreases {
45                let i = variant_names.len();
46                let name = Ident::new(&format!("__anodized_value_{}", i + 1), Span::mixed_site());
47                let expr = &loop_variant.expr;
48                variant_stmts.push(parse_quote! { let #name = (|| #expr)(); });
49                variant_names.push(name);
50            }
51            let variant_expr: Option<Expr> = if !variant_names.is_empty() {
52                Some(parse_quote! { (#(#variant_names),*) })
53            } else {
54                None
55            };
56
57            stmts.insert(
58                1,
59                parse_quote! {
60                    let __anodized_loop_decreases = || {
61                        #(#variant_stmts)*
62                        #variant_expr
63                    };
64                },
65            );
66        }
67    }
68}
69
70struct LoopSpecVisitor<'a> {
71    config: &'a Mode,
72    errors: Option<Error>,
73}
74
75impl<'a> LoopSpecVisitor<'a> {
76    fn new(config: &'a Mode) -> Self {
77        Self {
78            config,
79            errors: None,
80        }
81    }
82
83    fn finish(self) -> Result<()> {
84        match self.errors {
85            Some(error) => Err(error),
86            None => Ok(()),
87        }
88    }
89
90    fn add_error(&mut self, error: Error) {
91        match &mut self.errors {
92            Some(existing) => existing.combine(error),
93            None => self.errors = Some(error),
94        }
95    }
96}
97
98impl VisitMut for LoopSpecVisitor<'_> {
99    fn visit_expr_while_mut(&mut self, expr_while: &mut ExprWhile) {
100        let attrs = std::mem::take(&mut expr_while.attrs);
101        let (spec_attr, other_attrs) = match find_spec_attr(attrs) {
102            Ok(result) => result,
103            Err(error) => {
104                self.add_error(error);
105                return;
106            }
107        };
108        expr_while.attrs = other_attrs;
109
110        visit_mut::visit_expr_while_mut(self, expr_while);
111
112        let Some(spec_attr) = spec_attr else {
113            return;
114        };
115
116        match spec_attr.parse_args::<LoopSpec>() {
117            Ok(spec) => self
118                .config
119                .instrument_loop_body(spec, &mut expr_while.body.stmts),
120            Err(error) => self.add_error(error),
121        }
122    }
123
124    fn visit_expr_for_loop_mut(&mut self, expr_for_loop: &mut ExprForLoop) {
125        let attrs = std::mem::take(&mut expr_for_loop.attrs);
126        let (spec_attr, other_attrs) = match find_spec_attr(attrs) {
127            Ok(result) => result,
128            Err(error) => {
129                self.add_error(error);
130                return;
131            }
132        };
133        expr_for_loop.attrs = other_attrs;
134
135        visit_mut::visit_expr_for_loop_mut(self, expr_for_loop);
136
137        let Some(spec_attr) = spec_attr else {
138            return;
139        };
140
141        match spec_attr.parse_args::<LoopSpec>() {
142            Ok(spec) => self.config.instrument_expr_for_loop(spec, expr_for_loop),
143            Err(error) => self.add_error(error),
144        }
145    }
146
147    // Nested closure scopes are independently analyzed by the outer function macro expansion.
148    fn visit_expr_closure_mut(&mut self, _expr_closure: &mut ExprClosure) {}
149
150    // Nested `fn` items are independently analyzed by the outer function macro expansion.
151    fn visit_item_fn_mut(&mut self, _item_fn: &mut ItemFn) {}
152}
anodized_core/instrument/loops.rs

anodized_core/instrument/
loops.rs
