charon_lib/transform/finish_translation/
insert_storage_lives.rs

1//! Add missing StorageLives -- in MIR, some locals are considered "always" initialised, and have
2//! no StorageLive and StorageDead instructions associated; this always includes the arguments
3//! and the return value, but also sometimes includes other locals. We make sure these additional
4//! locals get initialised at the start of the function if they're used anywhere.
5use derive_generic_visitor::Visitor;
6
7use crate::ast::*;
8use crate::transform::TransformCtx;
9use crate::transform::ctx::TransformPass;
10use crate::ullbc_ast::BlockId;
11
12#[derive(Visitor)]
13struct StorageVisitor {
14    local_status: Vector<LocalId, LocalStatus>,
15}
16
17enum LocalStatus {
18    Unused,
19    UsedAndNoExplicitStorage,
20    UsedAndHasExplicitStorage,
21}
22
23impl StorageVisitor {
24    fn new(locals: &Locals) -> Self {
25        let local_status = locals.locals.map_ref(|local| {
26            if locals.is_return_or_arg(local.index) {
27                // The return and argument places count as having a `StorageLive` already.
28                LocalStatus::UsedAndHasExplicitStorage
29            } else {
30                LocalStatus::Unused
31            }
32        });
33        Self { local_status }
34    }
35}
36
37impl VisitBody for StorageVisitor {
38    fn visit_locals(&mut self, _: &Locals) -> ::std::ops::ControlFlow<Self::Break> {
39        // Don't look inside the local declarations otherwise we'll think they're all used.
40        ControlFlow::Continue(())
41    }
42    fn enter_local_id(&mut self, lid: &LocalId) {
43        let status = &mut self.local_status[*lid];
44        if let LocalStatus::Unused = *status {
45            *status = LocalStatus::UsedAndNoExplicitStorage
46        }
47    }
48    fn enter_llbc_statement(&mut self, st: &llbc_ast::Statement) {
49        match &st.kind {
50            llbc_ast::StatementKind::StorageDead(lid)
51            | llbc_ast::StatementKind::StorageLive(lid) => {
52                self.local_status[*lid] = LocalStatus::UsedAndHasExplicitStorage
53            }
54            _ => {}
55        }
56    }
57    fn enter_ullbc_statement(&mut self, st: &ullbc_ast::Statement) {
58        match &st.kind {
59            ullbc_ast::StatementKind::StorageDead(lid)
60            | ullbc_ast::StatementKind::StorageLive(lid) => {
61                self.local_status[*lid] = LocalStatus::UsedAndHasExplicitStorage
62            }
63            _ => {}
64        }
65    }
66}
67
68pub struct Transform;
69impl TransformPass for Transform {
70    fn transform_ctx(&self, ctx: &mut TransformCtx) {
71        ctx.for_each_fun_decl(|_ctx, fun| {
72            let Ok(body) = &mut fun.body else {
73                return;
74            };
75
76            let mut storage_visitor = StorageVisitor::new(body.locals());
77            let _ = storage_visitor.visit(body);
78
79            // Insert StorageLive instructions for the always initialised locals.
80            let locals_with_missing_storage = storage_visitor
81                .local_status
82                .iter_indexed()
83                .filter(|(_, status)| matches!(status, LocalStatus::UsedAndNoExplicitStorage))
84                .map(|(local, _)| local);
85            match body {
86                Body::Structured(body) => {
87                    let first_span = body.body.statements.first().unwrap().span;
88                    let new_statements = locals_with_missing_storage.map(|local| {
89                        llbc_ast::Statement::new(
90                            first_span,
91                            llbc_ast::StatementKind::StorageLive(local),
92                        )
93                    });
94                    body.body.statements.splice(0..0, new_statements);
95                }
96                Body::Unstructured(body) => {
97                    let first_block = body.body.get_mut(BlockId::ZERO).unwrap();
98                    let first_span = if let Some(fst) = first_block.statements.first() {
99                        fst.span
100                    } else {
101                        first_block.terminator.span
102                    };
103                    let new_statements = locals_with_missing_storage.map(|local| {
104                        ullbc_ast::Statement::new(
105                            first_span,
106                            ullbc_ast::StatementKind::StorageLive(local),
107                        )
108                    });
109                    first_block.statements.splice(0..0, new_statements);
110                }
111            }
112        });
113    }
114}
charon_lib/transform/finish_translation/insert_storage_lives.rs

charon_lib/transform/finish_translation/
insert_storage_lives.rs
