charon_lib/transform/simplify_output/inline_anon_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 should_run(&self, options: &crate::options::TranslateOptions) -> bool {
9 !options.raw_consts
10 }
11
12 fn transform_ctx(&self, ctx: &mut TransformCtx) {
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.to_unstructured()?;
21 Some((id, body))
22 })
23 .collect();
24
25 ctx.for_each_fun_decl(|_ctx, decl| {
26 if let Some(outer_body) = decl.body.as_unstructured_mut() {
27 for block_id in outer_body.body.all_indices() {
28 // Subtle: This generator must be managed to correctly track the indices that will
29 // be generated when pushing onto `outer_body.body`.
30 let mut bid_generator: Generator<BlockId> =
31 Generator::new_with_init_value(outer_body.body.next_idx());
32 let start_new_bodies = bid_generator.next_id();
33 let Some(block) = outer_body.body.get_mut(block_id) else {
34 continue;
35 };
36 let mut new_blocks = vec![];
37 block.dyn_visit_in_body_mut(|op: &mut Operand| {
38 if let Operand::Const(c) = op
39 && let ConstantExprKind::Global(gref) = &mut c.kind
40 && let Some(inner_body) = anon_consts.get(&gref.id)
41 {
42 // We inline the required body by shifting its local ids and block ids
43 // and adding its blocks to the outer body. The inner body's return
44 // local becomes a normal local that we can read from. We redirect some
45 // gotos so that the inner body is executed before the current block.
46 let mut inner_body = inner_body.clone();
47 let inner_bound = inner_body.bound_body_regions;
48
49 // Shift all the body regions in the inner body BEFORE substitution,
50 // so that we only shift the inner body's own regions.
51 inner_body.dyn_visit_mut(|r: &mut Region| {
52 if let Region::Body(v) = r {
53 *v += outer_body.bound_body_regions;
54 }
55 });
56 outer_body.bound_body_regions += inner_bound;
57
58 // Now substitute generics. This may inject outer-body Region::Body
59 // IDs, which is correct since they don't need shifting.
60 let mut inner_body = inner_body.substitute(&gref.generics);
61
62 // The init function of a global assumes the return place is live;
63 // this is not the case once we inline it
64 inner_body.body[0].statements.insert(
65 0,
66 Statement::new(
67 Span::dummy(),
68 StatementKind::StorageLive(LocalId::ZERO),
69 ),
70 );
71
72 let return_local = outer_body.locals.locals.next_idx();
73 inner_body.dyn_visit_in_body_mut(|l: &mut LocalId| {
74 *l += return_local;
75 });
76
77 let start_block = bid_generator.next_id();
78 bid_generator.advance(inner_body.body.len());
79 let end_block = bid_generator.next_id();
80 inner_body.dyn_visit_in_body_mut(|b: &mut BlockId| {
81 *b += start_block;
82 });
83 // Make all returns point to `end_block`. This block doesn't exist yet,
84 // it will either be the start block of another inner body, or the
85 // current outer block that we'll push at the end.
86 inner_body.body.dyn_visit_in_body_mut(|t: &mut Terminator| {
87 if let TerminatorKind::Return = t.kind {
88 t.kind = TerminatorKind::Goto { target: end_block };
89 }
90 });
91
92 outer_body
93 .locals
94 .locals
95 .extend(inner_body.locals.locals.into_iter());
96 new_blocks.extend(inner_body.body);
97 *op = Operand::Move(outer_body.locals.place_for_var(return_local));
98 }
99 });
100 if !new_blocks.is_empty() {
101 // Instead of the current block, start evaluating the new bodies.
102 let block = mem::replace(
103 block,
104 BlockData::new_goto(Span::dummy(), start_new_bodies),
105 );
106 // Add the new blocks. They've been set up so that each new inner body
107 // returns to what follows it in the sequence. Hence the last added body
108 // points to the not-yet-existing block at `start_new_bodies`
109 outer_body.body.extend(new_blocks.into_iter());
110 // Push the current block to be executed after the newly-added ones.
111 outer_body.body.push(block);
112 }
113 }
114 }
115 });
116 }
117}