charon_lib/ast/
ullbc_ast_utils.rs

1//! Implementations for [crate::ullbc_ast]
2use smallvec::{SmallVec, smallvec};
3
4use crate::ids::IndexVec;
5use crate::meta::Span;
6use crate::ullbc_ast::*;
7use std::collections::HashMap;
8use std::mem;
9use std::ops::{Index, IndexMut};
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
12pub struct StmtLoc {
13    pub block: BlockId,
14    pub statement: usize,
15}
16
17impl StmtLoc {
18    pub fn new(block: BlockId, statement: usize) -> Self {
19        StmtLoc { block, statement }
20    }
21
22    pub fn block_start(block: BlockId) -> Self {
23        StmtLoc {
24            block,
25            statement: 0,
26        }
27    }
28
29    pub fn after(self) -> Self {
30        StmtLoc {
31            block: self.block,
32            statement: self.statement + 1,
33        }
34    }
35}
36
37impl SwitchTargets {
38    pub fn targets(&self) -> SmallVec<[BlockId; 2]> {
39        match self {
40            SwitchTargets::If(then_tgt, else_tgt) => {
41                smallvec![*then_tgt, *else_tgt]
42            }
43            SwitchTargets::SwitchInt(_, targets, otherwise) => targets
44                .iter()
45                .map(|(_, t)| t)
46                .chain([otherwise])
47                .copied()
48                .collect(),
49        }
50    }
51    pub fn targets_mut(&mut self) -> SmallVec<[&mut BlockId; 2]> {
52        match self {
53            SwitchTargets::If(then_tgt, else_tgt) => {
54                smallvec![then_tgt, else_tgt]
55            }
56            SwitchTargets::SwitchInt(_, targets, otherwise) => targets
57                .iter_mut()
58                .map(|(_, t)| t)
59                .chain([otherwise])
60                .collect(),
61        }
62    }
63}
64
65impl Statement {
66    pub fn new(span: Span, kind: StatementKind) -> Self {
67        Statement {
68            span,
69            kind,
70            comments_before: vec![],
71        }
72    }
73}
74
75impl Terminator {
76    pub fn new(span: Span, kind: TerminatorKind) -> Self {
77        Terminator {
78            span,
79            kind,
80            comments_before: vec![],
81        }
82    }
83    pub fn goto(span: Span, target: BlockId) -> Self {
84        Self::new(span, TerminatorKind::Goto { target })
85    }
86    /// Whether this terminator is an unconditional error (panic).
87    pub fn is_error(&self) -> bool {
88        use TerminatorKind::*;
89        match &self.kind {
90            Abort(..) => true,
91            Goto { .. }
92            | Switch { .. }
93            | InlineAsm { .. }
94            | Return
95            | Call { .. }
96            | Drop { .. }
97            | UnwindResume
98            | Assert { .. } => false,
99        }
100    }
101
102    pub fn into_block(self) -> BlockData {
103        BlockData {
104            statements: vec![],
105            terminator: self,
106        }
107    }
108
109    pub fn targets(&self) -> SmallVec<[BlockId; 2]> {
110        match &self.kind {
111            TerminatorKind::Goto { target } => {
112                smallvec![*target]
113            }
114            TerminatorKind::Switch { targets, .. } => targets.targets(),
115            TerminatorKind::InlineAsm {
116                targets, on_unwind, ..
117            } => targets.iter().copied().chain([*on_unwind]).collect(),
118            TerminatorKind::Call {
119                target, on_unwind, ..
120            }
121            | TerminatorKind::Drop {
122                target, on_unwind, ..
123            }
124            | TerminatorKind::Assert {
125                target, on_unwind, ..
126            } => smallvec![*target, *on_unwind],
127            TerminatorKind::Abort(..) | TerminatorKind::Return | TerminatorKind::UnwindResume => {
128                smallvec![]
129            }
130        }
131    }
132    pub fn targets_mut(&mut self) -> SmallVec<[&mut BlockId; 2]> {
133        match &mut self.kind {
134            TerminatorKind::Goto { target } => {
135                smallvec![target]
136            }
137            TerminatorKind::Switch { targets, .. } => targets.targets_mut(),
138            TerminatorKind::InlineAsm {
139                targets, on_unwind, ..
140            } => targets.iter_mut().chain([on_unwind]).collect(),
141            TerminatorKind::Call {
142                target, on_unwind, ..
143            }
144            | TerminatorKind::Drop {
145                target, on_unwind, ..
146            }
147            | TerminatorKind::Assert {
148                target, on_unwind, ..
149            } => smallvec![target, on_unwind],
150            TerminatorKind::Abort(..) | TerminatorKind::Return | TerminatorKind::UnwindResume => {
151                smallvec![]
152            }
153        }
154    }
155
156    pub fn targets_ignoring_unwind(&self) -> SmallVec<[BlockId; 2]> {
157        match &self.kind {
158            TerminatorKind::Goto { target } => {
159                smallvec![*target]
160            }
161            TerminatorKind::Switch { targets, .. } => targets.targets(),
162            TerminatorKind::InlineAsm { targets, .. } => targets.iter().copied().collect(),
163            TerminatorKind::Call { target, .. }
164            | TerminatorKind::Drop { target, .. }
165            | TerminatorKind::Assert { target, .. } => {
166                smallvec![*target]
167            }
168            TerminatorKind::Abort(..) | TerminatorKind::Return | TerminatorKind::UnwindResume => {
169                smallvec![]
170            }
171        }
172    }
173}
174
175impl BlockData {
176    /// Build a block that's just a goto terminator.
177    pub fn new_goto(span: Span, target: BlockId) -> Self {
178        BlockData {
179            statements: vec![],
180            terminator: Terminator::goto(span, target),
181        }
182    }
183    pub fn as_goto(&self) -> Option<BlockId> {
184        if let TerminatorKind::Goto { target } = self.terminator.kind {
185            Some(target)
186        } else {
187            None
188        }
189    }
190    pub fn as_trivial_goto(&self) -> Option<BlockId> {
191        self.as_goto().filter(|_| {
192            self.statements
193                .iter()
194                .all(|st| matches!(st.kind, StatementKind::Nop))
195        })
196    }
197
198    pub fn as_abort(&self) -> Option<AbortKind> {
199        if self.statements.iter().all(|st| {
200            matches!(
201                st.kind,
202                StatementKind::Nop | StatementKind::StorageLive(_) | StatementKind::StorageDead(_)
203            )
204        }) && let TerminatorKind::Abort(abort) = &self.terminator.kind
205        {
206            Some(abort.clone())
207        } else {
208            None
209        }
210    }
211
212    /// Build a block that's UB to reach.
213    pub fn new_unreachable() -> Self {
214        Terminator::new(
215            Span::dummy(),
216            TerminatorKind::Abort(AbortKind::UndefinedBehavior),
217        )
218        .into_block()
219    }
220
221    pub fn targets(&self) -> SmallVec<[BlockId; 2]> {
222        self.terminator.targets()
223    }
224    pub fn targets_ignoring_unwind(&self) -> SmallVec<[BlockId; 2]> {
225        self.terminator.targets_ignoring_unwind()
226    }
227
228    /// Apply a transformer to all the statements.
229    ///
230    /// The transformer should:
231    /// - mutate the current statement in place
232    /// - return the sequence of statements to introduce before the current statement
233    pub fn transform<F: FnMut(&mut Statement) -> Vec<Statement>>(&mut self, mut f: F) {
234        self.transform_sequences_fwd(|slice| {
235            let new_statements = f(&mut slice[0]);
236            if new_statements.is_empty() {
237                vec![]
238            } else {
239                vec![(0, new_statements)]
240            }
241        });
242    }
243
244    /// Helper, see `transform_sequences_fwd` and `transform_sequences_bwd`.
245    fn transform_sequences<F>(&mut self, mut f: F, forward: bool)
246    where
247        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
248    {
249        let mut to_insert = vec![];
250        let mut final_len = self.statements.len();
251        if forward {
252            for i in 0..self.statements.len() {
253                let new_to_insert = f(&mut self.statements[i..]);
254                to_insert.extend(new_to_insert.into_iter().map(|(j, stmts)| {
255                    final_len += stmts.len();
256                    (i + j, stmts)
257                }));
258            }
259        } else {
260            for i in (0..self.statements.len()).rev() {
261                let new_to_insert = f(&mut self.statements[i..]);
262                to_insert.extend(new_to_insert.into_iter().map(|(j, stmts)| {
263                    final_len += stmts.len();
264                    (i + j, stmts)
265                }));
266            }
267        }
268        if !to_insert.is_empty() {
269            to_insert.sort_by_key(|(i, _)| *i);
270            // Make it so the first element is always at the end so we can pop it.
271            to_insert.reverse();
272            // Construct the merged list of statements.
273            let old_statements = mem::replace(&mut self.statements, Vec::with_capacity(final_len));
274            for (i, stmt) in old_statements.into_iter().enumerate() {
275                while let Some((j, _)) = to_insert.last()
276                    && *j == i
277                {
278                    let (_, mut stmts) = to_insert.pop().unwrap();
279                    self.statements.append(&mut stmts);
280                }
281                self.statements.push(stmt);
282            }
283        }
284    }
285
286    /// Apply a transformer to all the statements.
287    ///
288    /// The transformer should:
289    /// - mutate the current statements in place
290    /// - return a list of `(i, statements)` where `statements` will be inserted before index `i`.
291    pub fn transform_sequences_fwd<F>(&mut self, f: F)
292    where
293        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
294    {
295        self.transform_sequences(f, true);
296    }
297
298    /// Apply a transformer to all the statements.
299    ///
300    /// The transformer should:
301    /// - mutate the current statements in place
302    /// - return a list of `(i, statements)` where `statements` will be inserted before index `i`.
303    pub fn transform_sequences_bwd<F>(&mut self, f: F)
304    where
305        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
306    {
307        self.transform_sequences(f, false);
308    }
309}
310
311impl ExprBody {
312    /// Returns a map from blocks in this body to their abort kind, if they correspond to an
313    /// abort block (ie. a block with only bookkeeping statements and a
314    /// [TerminatorKind::Abort] terminator).
315    pub fn as_abort_map(&self) -> HashMap<BlockId, AbortKind> {
316        self.body
317            .iter_enumerated()
318            .filter_map(|(bid, block)| block.as_abort().map(|abort| (bid, abort)))
319            .collect()
320    }
321
322    pub fn transform_sequences_fwd<F>(&mut self, mut f: F)
323    where
324        F: FnMut(BlockId, &mut Locals, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
325    {
326        for (id, block) in &mut self.body.iter_mut_enumerated() {
327            block.transform_sequences_fwd(|seq| f(id, &mut self.locals, seq));
328        }
329    }
330
331    pub fn transform_sequences_bwd<F>(&mut self, mut f: F)
332    where
333        F: FnMut(&mut Locals, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
334    {
335        for block in &mut self.body {
336            block.transform_sequences_bwd(|seq| f(&mut self.locals, seq));
337        }
338    }
339
340    /// Apply a function to all the statements, in a bottom-up manner.
341    pub fn visit_statements<F: FnMut(&mut Statement)>(&mut self, mut f: F) {
342        for block in self.body.iter_mut().rev() {
343            for st in block.statements.iter_mut().rev() {
344                f(st);
345            }
346        }
347    }
348}
349
350impl Index<StmtLoc> for ExprBody {
351    type Output = Statement;
352    fn index(&self, loc: StmtLoc) -> &Self::Output {
353        &self.body[loc.block].statements[loc.statement]
354    }
355}
356
357impl IndexMut<StmtLoc> for ExprBody {
358    fn index_mut(&mut self, loc: StmtLoc) -> &mut Self::Output {
359        &mut self.body[loc.block].statements[loc.statement]
360    }
361}
362
363/// Helper to construct a small ullbc body.
364pub struct BodyBuilder {
365    /// The span to use for everything.
366    pub span: Span,
367    /// Body under construction.
368    pub body: ExprBody,
369    /// Block onto which we're adding statements. Its terminator is always `Return`.
370    pub current_block: BlockId,
371    /// Block to unwind to; created on demand.
372    pub unwind_block: Option<BlockId>,
373}
374
375fn mk_block(span: Span, term: TerminatorKind) -> BlockData {
376    BlockData {
377        statements: vec![],
378        terminator: Terminator::new(span, term),
379    }
380}
381
382impl BodyBuilder {
383    pub fn new(span: Span, arg_count: usize) -> Self {
384        let mut body: ExprBody = GExprBody {
385            span,
386            locals: Locals::new(arg_count),
387            bound_body_regions: 0,
388            body: IndexVec::new(),
389            comments: vec![],
390        };
391        let current_block = body.body.push(BlockData {
392            statements: Default::default(),
393            terminator: Terminator::new(span, TerminatorKind::Return),
394        });
395        Self {
396            span,
397            body,
398            current_block,
399            unwind_block: None,
400        }
401    }
402
403    /// Finalize the builder by returning the built body.
404    pub fn build(mut self) -> ExprBody {
405        // Replace erased regions with fresh ones.
406        let mut freshener: IndexMap<RegionId, ()> = IndexMap::new();
407        self.body.dyn_visit_mut(|r: &mut Region| {
408            if r.is_erased() || r.is_body() {
409                *r = Region::Body(freshener.push(()));
410            }
411        });
412        self.body.bound_body_regions = freshener.slot_count();
413        // Return the built body.
414        self.body
415    }
416
417    /// Create a new local. Adds a `StorageLive` statement if the local is not one of the special
418    /// ones (return or function argument).
419    pub fn new_var(&mut self, name: Option<String>, ty: Ty) -> Place {
420        let place = self.body.locals.new_var(name, ty);
421        let local_id = place.as_local().unwrap();
422        if !self.body.locals.is_return_or_arg(local_id) {
423            self.push_statement(StatementKind::StorageLive(local_id));
424        }
425        place
426    }
427
428    /// Helper.
429    fn current_block(&mut self) -> &mut BlockData {
430        &mut self.body.body[self.current_block]
431    }
432
433    pub fn push_statement(&mut self, kind: StatementKind) {
434        let st = Statement::new(self.span, kind);
435        self.current_block().statements.push(st);
436    }
437
438    fn unwind_block(&mut self) -> BlockId {
439        *self.unwind_block.get_or_insert_with(|| {
440            self.body
441                .body
442                .push(mk_block(self.span, TerminatorKind::UnwindResume))
443        })
444    }
445
446    pub fn call(&mut self, call: Call) {
447        let next_block = self
448            .body
449            .body
450            .push(mk_block(self.span, TerminatorKind::Return));
451        let term = TerminatorKind::Call {
452            target: next_block,
453            call,
454            on_unwind: self.unwind_block(),
455        };
456        self.current_block().terminator.kind = term;
457        self.current_block = next_block;
458    }
459
460    pub fn insert_drop(&mut self, place: Place, fn_ptr: FnPtr) {
461        let next_block = self
462            .body
463            .body
464            .push(mk_block(self.span, TerminatorKind::Return));
465        let term = TerminatorKind::Drop {
466            kind: DropKind::Precise,
467            place,
468            fn_ptr,
469            target: next_block,
470            on_unwind: self.unwind_block(),
471        };
472        self.current_block().terminator.kind = term;
473        self.current_block = next_block;
474    }
475}
charon_lib/ast/ullbc_ast_utils.rs

charon_lib/ast/
ullbc_ast_utils.rs
