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

charon_lib/ast/
ullbc_ast_utils.rs
