charon_lib/transform/
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.
5use std::collections::BTreeSet;
6use std::ops::ControlFlow::{self, Continue};
7
8use derive_generic_visitor::Visitor;
9
10use crate::ast::*;
11use crate::transform::TransformCtx;
12use crate::ullbc_ast::BlockId;
13
14use super::ctx::TransformPass;
15
16#[derive(Visitor)]
17struct StorageVisitor {
18    unmentioned_locals: BTreeSet<LocalId>,
19}
20
21impl StorageVisitor {
22    fn new(locals: &Locals) -> Self {
23        let mut unmentioned_locals = BTreeSet::new();
24        for local in locals.locals.iter() {
25            if local.index > locals.arg_count {
26                unmentioned_locals.insert(local.index);
27            }
28        }
29        Self { unmentioned_locals }
30    }
31}
32
33impl VisitAst for StorageVisitor {
34    fn visit_llbc_statement(&mut self, st: &llbc_ast::Statement) -> ControlFlow<Self::Break> {
35        match st.content {
36            llbc_ast::StatementKind::StorageDead(loc)
37            | llbc_ast::StatementKind::StorageLive(loc) => {
38                self.unmentioned_locals.remove(&loc);
39            }
40            _ => {}
41        }
42        Continue(())
43    }
44
45    fn visit_ullbc_statement(&mut self, st: &ullbc_ast::Statement) -> ControlFlow<Self::Break> {
46        match st.content {
47            ullbc_ast::StatementKind::StorageDead(loc)
48            | ullbc_ast::StatementKind::StorageLive(loc) => {
49                self.unmentioned_locals.remove(&loc);
50            }
51            _ => {}
52        }
53        Continue(())
54    }
55}
56
57pub struct Transform;
58impl TransformPass for Transform {
59    fn transform_ctx(&self, ctx: &mut TransformCtx) {
60        ctx.for_each_fun_decl(|_ctx, fun| {
61            let Ok(body) = &mut fun.body else {
62                return;
63            };
64
65            let mut storage_visitor = match body {
66                Body::Structured(body) => StorageVisitor::new(&body.locals),
67                Body::Unstructured(body) => StorageVisitor::new(&body.locals),
68            };
69            let _ = storage_visitor.visit(body);
70
71            // Insert StorageLive instructions for the always initialised locals.
72            match body {
73                Body::Structured(body) => {
74                    let first_span = body.body.statements.first().unwrap().span;
75                    let new_statements = storage_visitor.unmentioned_locals.iter().map(|local| {
76                        llbc_ast::Statement::new(
77                            first_span,
78                            llbc_ast::StatementKind::StorageLive(*local),
79                        )
80                    });
81                    body.body.statements.splice(0..0, new_statements);
82                }
83                Body::Unstructured(body) => {
84                    let first_block = body.body.get_mut(BlockId::ZERO).unwrap();
85                    let first_span = if let Some(fst) = first_block.statements.first() {
86                        fst.span
87                    } else {
88                        first_block.terminator.span
89                    };
90                    let new_statements = storage_visitor.unmentioned_locals.iter().map(|local| {
91                        ullbc_ast::Statement::new(
92                            first_span,
93                            ullbc_ast::StatementKind::StorageLive(*local),
94                        )
95                    });
96                    first_block.statements.splice(0..0, new_statements);
97                }
98            }
99        });
100    }
101}
charon_lib/transform/insert_storage_lives.rs

charon_lib/transform/
insert_storage_lives.rs
