charon_lib/ast/
ullbc_ast_utils.rs

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