charon_lib/transform/simplify_output/
inline_selected_functions.rs

1use std::{collections::HashMap, mem};
2
3use crate::transform::CowBox;
4use crate::transform::{TransformCtx, ctx::UllbcPass};
5use crate::ullbc_ast::*;
6
7pub struct Transform {
8    to_inline: HashMap<FunDeclId, FunDecl>,
9}
10
11impl Transform {
12    pub fn new(ctx: &mut TransformCtx) -> CowBox<dyn UllbcPass> {
13        let panic_name = Name::from_path(names::EXPLICIT_PANIC_NAME);
14        let panic_terminator = TerminatorKind::Abort(AbortKind::Panic(Some(panic_name)));
15
16        // Collect and remove the functions that we want to inline.
17        let to_inline = ctx
18            .translated
19            .fun_decls
20            .extract(|_, decl| {
21                decl.body.as_unstructured().is_some_and(|body| {
22                    // If the whole body is only a call to this specific panic function.
23                    // FIXME: also check that the name of the function is `panic_cold_explicit`?
24                    let is_local_panic_fn = body.body.len() == 1 && {
25                        let block = &body.body[0];
26                        block.statements.is_empty() && block.terminator.kind == panic_terminator
27                    };
28                    // The `anon_consts_to_call` pass already transformed references to anon consts
29                    // into calls to their initializers so we only have to inline these.
30                    let is_anon_const_initializer = if let FunSource::GlobalInitializer(global) =
31                        &decl.src
32                        && let Some(gdecl) = ctx.translated.global_decls.get(global.id)
33                    {
34                        matches!(gdecl.global_kind, GlobalKind::AnonConst)
35                    } else {
36                        false
37                    };
38                    let is_vec_construction_fn = decl.item_meta.diagnostic_item.as_deref()
39                        == Some(names::BOX_ASSUME_INIT_INTO_VEC_UNSAFE);
40                    is_local_panic_fn
41                        || (is_anon_const_initializer && !ctx.options.raw_consts)
42                        || (is_vec_construction_fn && ctx.options.treat_box_as_builtin)
43                })
44            })
45            .collect();
46
47        CowBox::Owned(Box::new(Transform { to_inline }))
48    }
49}
50impl UllbcPass for Transform {
51    fn should_run(&self, _options: &crate::options::TranslateOptions) -> bool {
52        !self.to_inline.is_empty()
53    }
54    fn apply_preceding_passes(&mut self, ctx: &mut TransformCtx, passes: &[CowBox<dyn UllbcPass>]) {
55        for decl in self.to_inline.values_mut() {
56            for pass in passes {
57                pass.transform_item(ctx, ItemRefMut::Fun(decl));
58            }
59        }
60    }
61    fn transform_body(&self, _ctx: &mut TransformCtx, outer_body: &mut ullbc_ast::ExprBody) {
62        for block_id in outer_body.body.indices() {
63            let Some(block) = outer_body.body.get_mut(block_id) else {
64                continue;
65            };
66            let TerminatorKind::Call {
67                call: Call { func, args, dest },
68                target,
69                on_unwind,
70            } = &mut block.terminator.kind
71            else {
72                continue;
73            };
74            let target = *target;
75            let on_unwind = *on_unwind;
76            let dest_place = dest.clone();
77            let args = args.clone();
78            let FnOperand::Regular(fn_ptr) = &func else {
79                continue;
80            };
81            let FnPtrKind::Fun(FunId::Regular(fun_id)) = fn_ptr.kind.as_ref() else {
82                continue;
83            };
84            let Some(initializer) = self.to_inline.get(fun_id) else {
85                continue;
86            };
87            let span = initializer.item_meta.span;
88            let Some(inner_body) = initializer.body.as_unstructured() else {
89                continue;
90            };
91
92            // We inline the required body by shifting its local ids and block ids
93            // and adding its blocks to the outer body. The inner body's return
94            // local becomes a normal local that we can read from. We redirect some
95            // gotos so that the inner body is executed before the current block.
96            let mut inner_body = {
97                let mut inner_body = inner_body.clone();
98                let inner_bound = inner_body.bound_body_regions;
99
100                // Shift all the body regions in the inner body BEFORE substitution,
101                // so that we only shift the inner body's own regions.
102                inner_body.dyn_visit_mut(|r: &mut Region| {
103                    if let Region::Body(v) = r {
104                        *v += outer_body.bound_body_regions;
105                    }
106                });
107                outer_body.bound_body_regions += inner_bound;
108
109                // Now substitute generics. This may inject outer-body Region::Body
110                // IDs, which is correct since they don't need shifting.
111                inner_body.substitute(&fn_ptr.generics)
112            };
113
114            let return_local = outer_body.locals.locals.next_idx();
115            inner_body.dyn_visit_in_body_mut(|l: &mut LocalId| {
116                *l += return_local;
117            });
118            outer_body
119                .locals
120                .locals
121                .extend(mem::take(&mut inner_body.locals.locals));
122
123            // The inner body assumes the arg places are live; allocate them, and initialize the
124            // args.
125            inner_body.body[0].statements.splice(
126                0..0,
127                args.into_iter()
128                    .enumerate()
129                    .flat_map(|(i, arg)| {
130                        let arg_local = return_local + i + 1;
131                        let arg_place = outer_body.locals.place_for_var(arg_local);
132                        [
133                            StatementKind::StorageLive(arg_local),
134                            StatementKind::Assign(arg_place, Rvalue::Use(arg, WithRetag::Yes)),
135                        ]
136                    })
137                    .map(|kind| Statement::new(span, kind)),
138            );
139
140            let mut final_block = BlockData::new_goto(span, target);
141
142            // The inner body will write to `return_place`, but the outer body expects the value at
143            // `dest_place`.
144            let return_place = outer_body.locals.place_for_var(return_local);
145            final_block.statements.push(Statement::new(
146                span,
147                StatementKind::Assign(
148                    dest_place,
149                    Rvalue::Use(Operand::Move(return_place), WithRetag::Yes),
150                ),
151            ));
152            let final_block = outer_body.body.push(final_block);
153
154            // Shift all block ids in the inner body and point return/unwind to where they should.
155            let start_block = outer_body.body.next_idx();
156            inner_body.dyn_visit_in_body_mut(|b: &mut BlockId| {
157                *b += start_block;
158            });
159            inner_body
160                .body
161                .dyn_visit_in_body_mut(|t: &mut Terminator| match t.kind {
162                    TerminatorKind::Return => {
163                        t.kind = TerminatorKind::Goto {
164                            target: final_block,
165                        };
166                    }
167                    TerminatorKind::UnwindResume => {
168                        t.kind = TerminatorKind::Goto { target: on_unwind };
169                    }
170                    _ => (),
171                });
172            // At the end of the current block, start evaluating the inner body.
173            outer_body.body[block_id].terminator.kind = TerminatorKind::Goto {
174                target: start_block,
175            };
176            // Add the blocks for the inner body.
177            outer_body.body.extend(inner_body.body);
178        }
179    }
180}
charon_lib/transform/simplify_output/inline_selected_functions.rs

charon_lib/transform/simplify_output/
inline_selected_functions.rs
