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, kind: StatementKind) -> Self {
26        Statement {
27            span,
28            kind,
29            comments_before: vec![],
30        }
31    }
32}
33
34impl Terminator {
35    pub fn new(span: Span, kind: TerminatorKind) -> Self {
36        Terminator {
37            span,
38            kind,
39            comments_before: vec![],
40        }
41    }
42    pub fn goto(span: Span, target: BlockId) -> Self {
43        Self::new(span, TerminatorKind::Goto { target })
44    }
45    /// Whether this terminator is an unconditional error (panic).
46    pub fn is_error(&self) -> bool {
47        use TerminatorKind::*;
48        match &self.kind {
49            Abort(..) => true,
50            Goto { .. } | Switch { .. } | Return | Call { .. } | Drop { .. } | UnwindResume => {
51                false
52            }
53        }
54    }
55
56    pub fn into_block(self) -> BlockData {
57        BlockData {
58            statements: vec![],
59            terminator: self,
60        }
61    }
62}
63
64impl BlockData {
65    /// Build a block that's just a goto terminator.
66    pub fn new_goto(span: Span, target: BlockId) -> Self {
67        BlockData {
68            statements: vec![],
69            terminator: Terminator::goto(span, target),
70        }
71    }
72
73    pub fn targets(&self) -> Vec<BlockId> {
74        match &self.terminator.kind {
75            TerminatorKind::Goto { target } => {
76                vec![*target]
77            }
78            TerminatorKind::Switch { targets, .. } => targets.get_targets(),
79            TerminatorKind::Call {
80                target, on_unwind, ..
81            }
82            | TerminatorKind::Drop {
83                target, on_unwind, ..
84            } => vec![*target, *on_unwind],
85            TerminatorKind::Abort(..) | TerminatorKind::Return | TerminatorKind::UnwindResume => {
86                vec![]
87            }
88        }
89    }
90
91    pub fn targets_ignoring_unwind(&self) -> Vec<BlockId> {
92        match &self.terminator.kind {
93            TerminatorKind::Goto { target } => {
94                vec![*target]
95            }
96            TerminatorKind::Switch { targets, .. } => targets.get_targets(),
97            TerminatorKind::Call { target, .. } | TerminatorKind::Drop { target, .. } => {
98                vec![*target]
99            }
100            TerminatorKind::Abort(..) | TerminatorKind::Return | TerminatorKind::UnwindResume => {
101                vec![]
102            }
103        }
104    }
105
106    /// TODO: Write new documentation
107    pub fn transform_operands<F: FnMut(&Span, &mut Vec<Statement>, &mut Operand)>(
108        &mut self,
109        mut f: F,
110    ) {
111        // Explore the operands in the statements
112        for mut st in mem::take(&mut self.statements) {
113            st.kind
114                .dyn_visit_in_body_mut(|op: &mut Operand| f(&st.span, &mut self.statements, op));
115            // Add the statement to the vector of statements
116            self.statements.push(st)
117        }
118
119        // Explore the terminator
120        self.terminator
121            .kind
122            .dyn_visit_in_body_mut(|op: &mut Operand| {
123                f(&self.terminator.span, &mut self.statements, op)
124            });
125    }
126
127    /// Apply a transformer to all the statements.
128    ///
129    /// The transformer should:
130    /// - mutate the current statement in place
131    /// - return the sequence of statements to introduce before the current statement
132    pub fn transform<F: FnMut(&mut Statement) -> Vec<Statement>>(&mut self, mut f: F) {
133        self.transform_sequences_bwd(|slice| {
134            let new_statements = f(&mut slice[0]);
135            if new_statements.is_empty() {
136                vec![]
137            } else {
138                vec![(0, new_statements)]
139            }
140        });
141    }
142
143    /// Helper, see `transform_sequences_fwd` and `transform_sequences_bwd`.
144    fn transform_sequences<F>(&mut self, mut f: F, forward: bool)
145    where
146        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
147    {
148        let mut to_insert = vec![];
149        let mut final_len = self.statements.len();
150        if forward {
151            for i in 0..self.statements.len() {
152                let new_to_insert = f(&mut self.statements[i..]);
153                to_insert.extend(new_to_insert.into_iter().map(|(j, stmts)| {
154                    final_len += stmts.len();
155                    (i + j, stmts)
156                }));
157            }
158        } else {
159            for i in (0..self.statements.len()).rev() {
160                let new_to_insert = f(&mut self.statements[i..]);
161                to_insert.extend(new_to_insert.into_iter().map(|(j, stmts)| {
162                    final_len += stmts.len();
163                    (i + j, stmts)
164                }));
165            }
166        }
167        if !to_insert.is_empty() {
168            to_insert.sort_by_key(|(i, _)| *i);
169            // Make it so the first element is always at the end so we can pop it.
170            to_insert.reverse();
171            // Construct the merged list of statements.
172            let old_statements = mem::replace(&mut self.statements, Vec::with_capacity(final_len));
173            for (i, stmt) in old_statements.into_iter().enumerate() {
174                while let Some((j, _)) = to_insert.last()
175                    && *j == i
176                {
177                    let (_, mut stmts) = to_insert.pop().unwrap();
178                    self.statements.append(&mut stmts);
179                }
180                self.statements.push(stmt);
181            }
182        }
183    }
184
185    /// Apply a transformer to all the statements.
186    ///
187    /// The transformer should:
188    /// - mutate the current statements in place
189    /// - return a list of `(i, statements)` where `statements` will be inserted before index `i`.
190    pub fn transform_sequences_fwd<F>(&mut self, f: F)
191    where
192        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
193    {
194        self.transform_sequences(f, true);
195    }
196
197    /// Apply a transformer to all the statements.
198    ///
199    /// The transformer should:
200    /// - mutate the current statements in place
201    /// - return a list of `(i, statements)` where `statements` will be inserted before index `i`.
202    pub fn transform_sequences_bwd<F>(&mut self, f: F)
203    where
204        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
205    {
206        self.transform_sequences(f, false);
207    }
208}
209
210impl ExprBody {
211    pub fn transform_sequences_fwd<F>(&mut self, mut f: F)
212    where
213        F: FnMut(&mut Locals, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
214    {
215        for block in &mut self.body {
216            block.transform_sequences_fwd(|seq| f(&mut self.locals, seq));
217        }
218    }
219
220    pub fn transform_sequences_bwd<F>(&mut self, mut f: F)
221    where
222        F: FnMut(&mut Locals, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
223    {
224        for block in &mut self.body {
225            block.transform_sequences_bwd(|seq| f(&mut self.locals, seq));
226        }
227    }
228
229    /// Apply a function to all the statements, in a bottom-up manner.
230    pub fn visit_statements<F: FnMut(&mut Statement)>(&mut self, mut f: F) {
231        for block in self.body.iter_mut().rev() {
232            for st in block.statements.iter_mut().rev() {
233                f(st);
234            }
235        }
236    }
237}
238
239/// Helper to construct a small ullbc body.
240pub struct BodyBuilder {
241    /// The span to use for everything.
242    pub span: Span,
243    /// Body under construction.
244    pub body: ExprBody,
245    /// Block onto which we're adding statements. Its terminator is always `Return`.
246    pub current_block: BlockId,
247    /// Block to unwind to; created on demand.
248    pub unwind_block: Option<BlockId>,
249}
250
251fn mk_block(span: Span, term: TerminatorKind) -> BlockData {
252    BlockData {
253        statements: vec![],
254        terminator: Terminator::new(span, term),
255    }
256}
257
258impl BodyBuilder {
259    pub fn new(span: Span, arg_count: usize) -> Self {
260        let mut body: ExprBody = GExprBody {
261            span,
262            locals: Locals::new(arg_count),
263            comments: vec![],
264            body: Vector::new(),
265        };
266        let current_block = body.body.push(BlockData {
267            statements: Default::default(),
268            terminator: Terminator::new(span, TerminatorKind::Return),
269        });
270        Self {
271            span,
272            body,
273            current_block,
274            unwind_block: None,
275        }
276    }
277
278    /// Finalize the builder by returning the built body.
279    pub fn build(self) -> ExprBody {
280        self.body
281    }
282
283    /// Create a new local. Adds a `StorageLive` statement if the local is not one of the special
284    /// ones (return or function argument).
285    pub fn new_var(&mut self, name: Option<String>, ty: Ty) -> Place {
286        let place = self.body.locals.new_var(name, ty);
287        let local_id = place.as_local().unwrap();
288        if !self.body.locals.is_return_or_arg(local_id) {
289            self.push_statement(StatementKind::StorageLive(local_id));
290        }
291        place
292    }
293
294    /// Helper.
295    fn current_block(&mut self) -> &mut BlockData {
296        &mut self.body.body[self.current_block]
297    }
298
299    pub fn push_statement(&mut self, kind: StatementKind) {
300        let st = Statement::new(self.span, kind);
301        self.current_block().statements.push(st);
302    }
303
304    fn unwind_block(&mut self) -> BlockId {
305        *self.unwind_block.get_or_insert_with(|| {
306            self.body
307                .body
308                .push(mk_block(self.span, TerminatorKind::UnwindResume))
309        })
310    }
311
312    pub fn call(&mut self, call: Call) {
313        let next_block = self
314            .body
315            .body
316            .push(mk_block(self.span, TerminatorKind::Return));
317        let term = TerminatorKind::Call {
318            target: next_block,
319            call,
320            on_unwind: self.unwind_block(),
321        };
322        self.current_block().terminator.kind = term;
323        self.current_block = next_block;
324    }
325
326    pub fn insert_drop(&mut self, place: Place, tref: TraitRef) {
327        let next_block = self
328            .body
329            .body
330            .push(mk_block(self.span, TerminatorKind::Return));
331        let term = TerminatorKind::Drop {
332            place: place,
333            tref: tref,
334            target: next_block,
335            on_unwind: self.unwind_block(),
336        };
337        self.current_block().terminator.kind = term;
338        self.current_block = next_block;
339    }
340}
charon_lib/ast/ullbc_ast_utils.rs

charon_lib/ast/
ullbc_ast_utils.rs
