Skip to main content

charon_lib/transform/simplify_output/
remove_unused_locals.rs

1//! Remove the locals (which are not used for the input arguments) which are
2//! never used in the function bodies.  This is useful to remove the locals with
3//! type `Never`. We actually check that there are no such local variables
4//! remaining afterwards.
5use derive_generic_visitor::Visitor;
6use std::mem;
7use std::ops::ControlFlow::Continue;
8
9use crate::ast::*;
10use crate::ids::IndexVec;
11use crate::transform::TransformCtx;
12use crate::transform::ctx::{TransformPass, UllbcPass};
13
14#[derive(Visitor)]
15struct LocalsUsageVisitor {
16    used_locals: IndexVec<LocalId, bool>,
17}
18
19impl VisitBody for LocalsUsageVisitor {
20    fn enter_local_id(&mut self, lid: &LocalId) {
21        self.used_locals[*lid] = true;
22    }
23    fn visit_local(&mut self, local: &Local) -> ControlFlow<Self::Break> {
24        // Skip the local id.
25        self.visit(&local.drop_flag_for)
26    }
27    fn visit_llbc_statement(&mut self, st: &llbc_ast::Statement) -> ControlFlow<Self::Break> {
28        match &st.kind {
29            llbc_ast::StatementKind::StorageDead(_) | llbc_ast::StatementKind::StorageLive(_) => {
30                // These statements don't count as a variable use.
31                Continue(())
32            }
33            _ => self.visit_inner(st),
34        }
35    }
36    fn visit_ullbc_statement(&mut self, st: &ullbc_ast::Statement) -> ControlFlow<Self::Break> {
37        match &st.kind {
38            ullbc_ast::StatementKind::StorageDead(_) | ullbc_ast::StatementKind::StorageLive(_) => {
39                // These statements don't count as a variable use.
40                Continue(())
41            }
42            _ => self.visit_inner(st),
43        }
44    }
45}
46
47#[derive(Visitor)]
48struct LocalsRenumberVisitor {
49    ids_map: IndexVec<LocalId, Option<LocalId>>,
50}
51
52impl VisitBodyMut for LocalsRenumberVisitor {
53    fn enter_local_id(&mut self, lid: &mut LocalId) {
54        *lid = self.ids_map[*lid].unwrap();
55    }
56    fn enter_llbc_statement(&mut self, st: &mut llbc_ast::Statement) {
57        match st.kind {
58            llbc_ast::StatementKind::StorageDead(lid)
59            | llbc_ast::StatementKind::StorageLive(lid)
60                if self.ids_map[lid].is_none() =>
61            {
62                st.kind = llbc_ast::StatementKind::Nop;
63            }
64            _ => {}
65        }
66    }
67    fn enter_ullbc_statement(&mut self, st: &mut ullbc_ast::Statement) {
68        match st.kind {
69            ullbc_ast::StatementKind::StorageDead(lid)
70            | ullbc_ast::StatementKind::StorageLive(lid)
71                if self.ids_map[lid].is_none() =>
72            {
73                st.kind = ullbc_ast::StatementKind::Nop;
74            }
75            _ => {}
76        }
77    }
78}
79
80fn remove_unused_locals<Body: BodyVisitable>(body: &mut GExprBody<Body>) {
81    // Compute the set of used locals.
82    // We always register the return variable and the input arguments.
83    let mut visitor = LocalsUsageVisitor {
84        used_locals: body
85            .locals
86            .locals
87            .map_ref(|local| body.locals.is_return_or_arg(local.index)),
88    };
89    body.drive_body(&mut visitor);
90    let used_locals = visitor.used_locals;
91    trace!("used_locals: {:?}", used_locals);
92
93    // Keep only the variables that are used (storage statements don't count) and update their
94    // indices to be contiguous.
95    let mut ids_map: IndexVec<LocalId, Option<LocalId>> = body.locals.locals.map_ref(|_| None);
96    for local in mem::take(&mut body.locals.locals) {
97        if used_locals[local.index] {
98            let old_id = local.index;
99            // The pushed local has the wrong index, which will be fixed by the renumbering pass.
100            let new_id = body.locals.locals.push(local);
101            ids_map[old_id] = Some(new_id);
102        }
103    }
104    trace!("ids_maps: {:?}", ids_map);
105
106    // Update all `LocalId`s.
107    let mut visitor = LocalsRenumberVisitor { ids_map };
108    body.drive_body_mut(&mut visitor);
109}
110
111pub struct Transform;
112impl UllbcPass for Transform {
113    fn transform_body(&self, _ctx: &mut TransformCtx, body: &mut ullbc_ast::ExprBody) {
114        remove_unused_locals(body)
115    }
116}
117impl TransformPass for Transform {
118    fn transform_ctx(&self, ctx: &mut TransformCtx) {
119        ctx.for_each_fun_decl(|_ctx, fun| match &mut fun.body {
120            Body::Unstructured(body) => remove_unused_locals(body),
121            Body::Structured(body) => remove_unused_locals(body),
122            _ => {}
123        });
124    }
125}