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::RawStatement::StorageDead(loc) | llbc_ast::RawStatement::StorageLive(loc) => {
37                self.unmentioned_locals.remove(&loc);
38            }
39            _ => {}
40        }
41        Continue(())
42    }
43
44    fn visit_ullbc_statement(&mut self, st: &ullbc_ast::Statement) -> ControlFlow<Self::Break> {
45        match st.content {
46            ullbc_ast::RawStatement::StorageDead(loc)
47            | ullbc_ast::RawStatement::StorageLive(loc) => {
48                self.unmentioned_locals.remove(&loc);
49            }
50            _ => {}
51        }
52        Continue(())
53    }
54}
55
56pub struct Transform;
57impl TransformPass for Transform {
58    fn transform_ctx(&self, ctx: &mut TransformCtx) {
59        ctx.for_each_fun_decl(|_ctx, fun| {
60            let Ok(body) = &mut fun.body else {
61                return;
62            };
63
64            let mut storage_visitor = match body {
65                Body::Structured(body) => StorageVisitor::new(&body.locals),
66                Body::Unstructured(body) => StorageVisitor::new(&body.locals),
67            };
68            let _ = storage_visitor.visit(body);
69
70            // Insert StorageLive instructions for the always initialised locals.
71            match body {
72                Body::Structured(body) => {
73                    let first_span = body.body.statements.first().unwrap().span;
74                    let new_statements = storage_visitor.unmentioned_locals.iter().map(|local| {
75                        llbc_ast::Statement::new(
76                            first_span,
77                            llbc_ast::RawStatement::StorageLive(*local),
78                        )
79                    });
80                    body.body.statements.splice(0..0, new_statements);
81                }
82                Body::Unstructured(body) => {
83                    let first_block = body.body.get_mut(BlockId::ZERO).unwrap();
84                    let first_span = if let Some(fst) = first_block.statements.first() {
85                        fst.span
86                    } else {
87                        first_block.terminator.span
88                    };
89                    let new_statements = storage_visitor.unmentioned_locals.iter().map(|local| {
90                        ullbc_ast::Statement::new(
91                            first_span,
92                            ullbc_ast::RawStatement::StorageLive(*local),
93                        )
94                    });
95                    first_block.statements.splice(0..0, new_statements);
96                }
97            }
98        });
99    }
100}
charon_lib/transform/insert_storage_lives.rs

charon_lib/transform/
insert_storage_lives.rs
