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

charon_lib/ast/
ullbc_ast_utils.rs
