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::TransformPass;
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 TransformPass for Transform {
71    fn transform_ctx(&self, ctx: &mut TransformCtx) {
72        ctx.for_each_fun_decl(|_ctx, fun| {
73            let body = &mut fun.body;
74            if !body.has_contents() {
75                return;
76            }
77
78            let mut storage_visitor = StorageVisitor::new(body.locals());
79            let _ = storage_visitor.visit(body);
80
81            // Insert StorageLive instructions for the always initialised locals.
82            let locals_with_missing_storage = storage_visitor
83                .local_status
84                .iter_indexed()
85                .filter(|(_, status)| matches!(status, LocalStatus::UsedAndNoExplicitStorage))
86                .map(|(local, _)| local);
87            match body {
88                Body::Structured(body) => {
89                    let first_span = body.body.statements.first().unwrap().span;
90                    let new_statements = locals_with_missing_storage.map(|local| {
91                        llbc_ast::Statement::new(
92                            first_span,
93                            llbc_ast::StatementKind::StorageLive(local),
94                        )
95                    });
96                    body.body.statements.splice(0..0, new_statements);
97                }
98                Body::Unstructured(body) => {
99                    let first_block = body.body.get_mut(BlockId::ZERO).unwrap();
100                    let first_span = if let Some(fst) = first_block.statements.first() {
101                        fst.span
102                    } else {
103                        first_block.terminator.span
104                    };
105                    let new_statements = locals_with_missing_storage.map(|local| {
106                        ullbc_ast::Statement::new(
107                            first_span,
108                            ullbc_ast::StatementKind::StorageLive(local),
109                        )
110                    });
111                    first_block.statements.splice(0..0, new_statements);
112                }
113                _ => unreachable!(),
114            }
115        });
116    }
117}
charon_lib/transform/finish_translation/insert_storage_lives.rs

charon_lib/transform/finish_translation/
insert_storage_lives.rs
