charon_lib/ast/
ullbc_ast_utils.rs

1//! Implementations for [crate::ullbc_ast]
2use crate::meta::Span;
3use crate::ullbc_ast::*;
4use std::mem;
5
6impl SwitchTargets {
7    pub fn get_targets(&self) -> Vec<BlockId> {
8        match self {
9            SwitchTargets::If(then_tgt, else_tgt) => {
10                vec![*then_tgt, *else_tgt]
11            }
12            SwitchTargets::SwitchInt(_, targets, otherwise) => {
13                let mut all_targets = vec![];
14                for (_, target) in targets {
15                    all_targets.push(*target);
16                }
17                all_targets.push(*otherwise);
18                all_targets
19            }
20        }
21    }
22}
23
24impl Statement {
25    pub fn new(span: Span, content: RawStatement) -> Self {
26        Statement {
27            span,
28            content,
29            comments_before: vec![],
30        }
31    }
32}
33
34impl Terminator {
35    pub fn new(span: Span, content: RawTerminator) -> Self {
36        Terminator {
37            span,
38            content,
39            comments_before: vec![],
40        }
41    }
42}
43
44impl BlockData {
45    pub fn targets(&self) -> Vec<BlockId> {
46        match &self.terminator.content {
47            RawTerminator::Goto { target } => {
48                vec![*target]
49            }
50            RawTerminator::Switch { targets, .. } => targets.get_targets(),
51            RawTerminator::Abort(..) | RawTerminator::Return => {
52                vec![]
53            }
54        }
55    }
56
57    /// See [body_transform_operands]
58    pub fn transform_operands<F: FnMut(&Span, &mut Vec<Statement>, &mut Operand)>(
59        &mut self,
60        mut f: F,
61    ) {
62        // Explore the operands in the statements
63        for mut st in mem::take(&mut self.statements) {
64            st.content
65                .dyn_visit_in_body_mut(|op: &mut Operand| f(&st.span, &mut self.statements, op));
66            // Add the statement to the vector of statements
67            self.statements.push(st)
68        }
69
70        // Explore the terminator
71        self.terminator
72            .content
73            .dyn_visit_in_body_mut(|op: &mut Operand| {
74                f(&self.terminator.span, &mut self.statements, op)
75            });
76    }
77
78    /// Apply a transformer to all the statements.
79    ///
80    /// The transformer should:
81    /// - mutate the current statement in place
82    /// - return the sequence of statements to introduce before the current statement
83    pub fn transform<F: FnMut(&mut Statement) -> Vec<Statement>>(&mut self, mut f: F) {
84        self.transform_sequences(|slice| {
85            let new_statements = f(&mut slice[0]);
86            if new_statements.is_empty() {
87                vec![]
88            } else {
89                vec![(0, new_statements)]
90            }
91        });
92    }
93
94    /// Apply a transformer to all the statements.
95    ///
96    /// The transformer should:
97    /// - mutate the current statements in place
98    /// - return a list of `(i, statements)` where `statements` will be inserted before index `i`.
99    pub fn transform_sequences<F>(&mut self, mut f: F)
100    where
101        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
102    {
103        let mut to_insert = vec![];
104        let mut final_len = self.statements.len();
105        for i in (0..self.statements.len()).rev() {
106            let new_to_insert = f(&mut self.statements[i..]);
107            to_insert.extend(new_to_insert.into_iter().map(|(j, stmts)| {
108                final_len += stmts.len();
109                (i + j, stmts)
110            }));
111        }
112        if !to_insert.is_empty() {
113            to_insert.sort_by_key(|(i, _)| *i);
114            // Make it so the first element is always at the end so we can pop it.
115            to_insert.reverse();
116            // Construct the merged list of statements.
117            let old_statements = mem::replace(&mut self.statements, Vec::with_capacity(final_len));
118            for (i, stmt) in old_statements.into_iter().enumerate() {
119                while let Some((j, _)) = to_insert.last()
120                    && *j == i
121                {
122                    let (_, mut stmts) = to_insert.pop().unwrap();
123                    self.statements.append(&mut stmts);
124                }
125                self.statements.push(stmt);
126            }
127        }
128    }
129}
130
131impl ExprBody {
132    pub fn transform_sequences<F>(&mut self, mut f: F)
133    where
134        F: FnMut(&mut Locals, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
135    {
136        for block in &mut self.body {
137            block.transform_sequences(|seq| f(&mut self.locals, seq));
138        }
139    }
140
141    /// Apply a function to all the statements, in a bottom-up manner.
142    pub fn visit_statements<F: FnMut(&mut Statement)>(&mut self, mut f: F) {
143        for block in self.body.iter_mut().rev() {
144            for st in block.statements.iter_mut().rev() {
145                f(st);
146            }
147        }
148    }
149}