charon_lib/transform/simplify_output/inline_promoted_consts.rs
1use std::{collections::HashMap, mem};
2
3use crate::transform::{TransformCtx, ctx::UllbcPass};
4use crate::{ids::Generator, ullbc_ast::*};
5
6pub struct Transform;
7impl UllbcPass for Transform {
8 fn transform_ctx(&self, ctx: &mut TransformCtx) {
9 // Currently the only anon consts that are not already evaluated are promoted consts. If
10 // that changes, we'll have to restrict this pass to the consts that can be inlined into a
11 // body.
12
13 // Map each anon const id to its initializer, and remove both from `translated`.
14 let anon_consts: HashMap<GlobalDeclId, ExprBody> = ctx
15 .translated
16 .global_decls
17 .extract(|gdecl| matches!(gdecl.global_kind, GlobalKind::AnonConst))
18 .filter_map(|(id, gdecl)| {
19 let fdecl = ctx.translated.fun_decls.remove(gdecl.init)?;
20 let body = fdecl.body.ok()?;
21 let body = body.to_unstructured()?;
22 Some((id, body))
23 })
24 .collect();
25
26 ctx.for_each_fun_decl(|_ctx, decl| {
27 if let Ok(outer_body) = &mut decl.body {
28 let outer_body = outer_body.as_unstructured_mut().unwrap();
29 for block_id in outer_body.body.all_indices() {
30 // Subtle: This generator must be managed to correctly track the indices that will
31 // be generated when pushing onto `outer_body.body`.
32 let mut bid_generator =
33 Generator::new_with_init_value(outer_body.body.next_id());
34 let start_new_bodies = bid_generator.next_id();
35 let Some(block) = outer_body.body.get_mut(block_id) else {
36 continue;
37 };
38 let mut new_blocks = vec![];
39 block.dyn_visit_in_body_mut(|op: &mut Operand| {
40 if let Operand::Const(c) = op
41 && let ConstantExprKind::Global(gref) = &mut c.kind
42 && let Some(inner_body) = anon_consts.get(&gref.id)
43 {
44 // We inline the required body by shifting its local ids and block ids
45 // and adding its blocks to the outer body. The inner body's return
46 // local becomes a normal local that we can read from. We redirect some
47 // gotos so that the inner body is executed before the current block.
48 let mut inner_body = inner_body.clone().substitute(&gref.generics);
49
50 // The init function of a global assumes the return place is live;
51 // this is not the case once we inline it
52 inner_body.body[0].statements.insert(
53 0,
54 Statement::new(
55 Span::dummy(),
56 StatementKind::StorageLive(LocalId::ZERO),
57 ),
58 );
59
60 let return_local = outer_body.locals.locals.next_id();
61 inner_body.dyn_visit_in_body_mut(|l: &mut LocalId| {
62 *l += return_local;
63 });
64
65 let start_block = bid_generator.next_id();
66 bid_generator.advance(inner_body.body.elem_count());
67 let end_block = bid_generator.next_id();
68 inner_body.dyn_visit_in_body_mut(|b: &mut BlockId| {
69 *b += start_block;
70 });
71 // Make all returns point to `end_block`. This block doesn't exist yet,
72 // it will either be the start block of another inner body, or the
73 // current outer block that we'll push at the end.
74 inner_body.body.dyn_visit_in_body_mut(|t: &mut Terminator| {
75 if let TerminatorKind::Return = t.kind {
76 t.kind = TerminatorKind::Goto { target: end_block };
77 }
78 });
79
80 outer_body
81 .locals
82 .locals
83 .extend(inner_body.locals.locals.into_iter());
84 new_blocks.extend(inner_body.body);
85 *op = Operand::Move(outer_body.locals.place_for_var(return_local));
86 }
87 });
88 if !new_blocks.is_empty() {
89 // Instead of the current block, start evaluating the new bodies.
90 let block = mem::replace(
91 block,
92 BlockData::new_goto(Span::dummy(), start_new_bodies),
93 );
94 // Add the new blocks. They've been set up so that each new inner body
95 // returns to what follows it in the sequence. Hence the last added body
96 // points to the not-yet-existing block at `start_new_bodies`
97 outer_body.body.extend(new_blocks.into_iter());
98 // Push the current block to be executed after the newly-added ones.
99 outer_body.body.push(block);
100 }
101 }
102 }
103 });
104 }
105}