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

charon_lib/transform/finish_translation/
insert_storage_lives.rs
