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;
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_llbc_statement(&mut self, st: &llbc_ast::Statement) -> ControlFlow<Self::Break> {
24        match &st.kind {
25            llbc_ast::StatementKind::StorageDead(_) | llbc_ast::StatementKind::StorageLive(_) => {
26                // These statements don't count as a variable use.
27                Continue(())
28            }
29            _ => self.visit_inner(st),
30        }
31    }
32    fn visit_ullbc_statement(&mut self, st: &ullbc_ast::Statement) -> ControlFlow<Self::Break> {
33        match &st.kind {
34            ullbc_ast::StatementKind::StorageDead(_) | ullbc_ast::StatementKind::StorageLive(_) => {
35                // These statements don't count as a variable use.
36                Continue(())
37            }
38            _ => self.visit_inner(st),
39        }
40    }
41}
42
43#[derive(Visitor)]
44struct LocalsRenumberVisitor {
45    ids_map: IndexVec<LocalId, Option<LocalId>>,
46}
47
48impl VisitBodyMut for LocalsRenumberVisitor {
49    fn enter_local_id(&mut self, lid: &mut LocalId) {
50        *lid = self.ids_map[*lid].unwrap();
51    }
52    fn enter_llbc_statement(&mut self, st: &mut llbc_ast::Statement) {
53        match st.kind {
54            llbc_ast::StatementKind::StorageDead(lid)
55            | llbc_ast::StatementKind::StorageLive(lid)
56                if self.ids_map[lid].is_none() =>
57            {
58                st.kind = llbc_ast::StatementKind::Nop;
59            }
60            _ => {}
61        }
62    }
63    fn enter_ullbc_statement(&mut self, st: &mut ullbc_ast::Statement) {
64        match st.kind {
65            ullbc_ast::StatementKind::StorageDead(lid)
66            | ullbc_ast::StatementKind::StorageLive(lid)
67                if self.ids_map[lid].is_none() =>
68            {
69                st.kind = ullbc_ast::StatementKind::Nop;
70            }
71            _ => {}
72        }
73    }
74}
75
76fn remove_unused_locals<Body: BodyVisitable>(body: &mut GExprBody<Body>) {
77    // Compute the set of used locals.
78    // We always register the return variable and the input arguments.
79    let mut visitor = LocalsUsageVisitor {
80        used_locals: body
81            .locals
82            .locals
83            .map_ref(|local| body.locals.is_return_or_arg(local.index)),
84    };
85    let _ = body.body.drive_body(&mut visitor);
86    let used_locals = visitor.used_locals;
87    trace!("used_locals: {:?}", used_locals);
88
89    // Keep only the variables that are used (storage statements don't count) and update their
90    // indices to be contiguous.
91    let mut ids_map: IndexVec<LocalId, Option<LocalId>> = body.locals.locals.map_ref(|_| None);
92    for local in mem::take(&mut body.locals.locals) {
93        if used_locals[local.index] {
94            let old_id = local.index;
95            let new_id = body
96                .locals
97                .locals
98                .push_with(|index| Local { index, ..local });
99            ids_map[old_id] = Some(new_id);
100        }
101    }
102    trace!("ids_maps: {:?}", ids_map);
103
104    // Update all `LocalId`s.
105    let mut visitor = LocalsRenumberVisitor { ids_map };
106    let _ = body.body.drive_body_mut(&mut visitor);
107}
108
109pub struct Transform;
110impl TransformPass for Transform {
111    fn transform_ctx(&self, ctx: &mut TransformCtx) {
112        ctx.for_each_fun_decl(|_ctx, fun| match &mut fun.body {
113            Body::Unstructured(body) => remove_unused_locals(body),
114            Body::Structured(body) => remove_unused_locals(body),
115            _ => {}
116        });
117    }
118}