anodized_core/instrument/loops/
mod.rs

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

anodized_core/instrument/loops/
mod.rs
