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