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::ids::IndexVec;
9use crate::transform::TransformCtx;
10use crate::transform::ctx::UllbcPass;
11use crate::ullbc_ast::BlockId;
12
13#[derive(Visitor)]
14struct StorageVisitor {
15    local_status: IndexVec<LocalId, LocalStatus>,
16}
17
18enum LocalStatus {
19    Unused,
20    UsedAndNoExplicitStorage,
21    UsedAndHasExplicitStorage,
22}
23
24impl StorageVisitor {
25    fn new(locals: &Locals) -> Self {
26        let local_status = locals.locals.map_ref(|local| {
27            if locals.is_return_or_arg(local.index) {
28                // The return and argument places count as having a `StorageLive` already.
29                LocalStatus::UsedAndHasExplicitStorage
30            } else {
31                LocalStatus::Unused
32            }
33        });
34        Self { local_status }
35    }
36}
37
38impl VisitBody for StorageVisitor {
39    fn visit_locals(&mut self, _: &Locals) -> ::std::ops::ControlFlow<Self::Break> {
40        // Don't look inside the local declarations otherwise we'll think they're all used.
41        ControlFlow::Continue(())
42    }
43    fn enter_local_id(&mut self, lid: &LocalId) {
44        let status = &mut self.local_status[*lid];
45        if let LocalStatus::Unused = *status {
46            *status = LocalStatus::UsedAndNoExplicitStorage
47        }
48    }
49    fn enter_llbc_statement(&mut self, st: &llbc_ast::Statement) {
50        match &st.kind {
51            llbc_ast::StatementKind::StorageDead(lid)
52            | llbc_ast::StatementKind::StorageLive(lid) => {
53                self.local_status[*lid] = LocalStatus::UsedAndHasExplicitStorage
54            }
55            _ => {}
56        }
57    }
58    fn enter_ullbc_statement(&mut self, st: &ullbc_ast::Statement) {
59        match &st.kind {
60            ullbc_ast::StatementKind::StorageDead(lid)
61            | ullbc_ast::StatementKind::StorageLive(lid) => {
62                self.local_status[*lid] = LocalStatus::UsedAndHasExplicitStorage
63            }
64            _ => {}
65        }
66    }
67}
68
69pub struct Transform;
70impl UllbcPass for Transform {
71    fn transform_function(&self, _ctx: &mut TransformCtx, fun: &mut FunDecl) {
72        let Body::Unstructured(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_enumerated()
83            .filter(|(_, status)| matches!(status, LocalStatus::UsedAndNoExplicitStorage))
84            .map(|(local, _)| local);
85        let first_block = body.body.get_mut(BlockId::ZERO).unwrap();
86        let first_span = if let Some(fst) = first_block.statements.first() {
87            fst.span
88        } else {
89            first_block.terminator.span
90        };
91        let new_statements = locals_with_missing_storage.map(|local| {
92            ullbc_ast::Statement::new(first_span, ullbc_ast::StatementKind::StorageLive(local))
93        });
94        first_block.statements.splice(0..0, new_statements);
95    }
96}
charon_lib/transform/finish_translation/insert_storage_lives.rs

charon_lib/transform/finish_translation/
insert_storage_lives.rs
