charon_lib/transform/
update_closure_signatures.rs

1//! # Micro-pass: the first local variable of closures is (a borrow to) the closure itself. This is
2//! not consistent with the closure signature, which represents the captured state as a tuple. This
3//! micro-pass updates this.
4use crate::ids::Vector;
5use crate::transform::TransformCtx;
6use crate::ullbc_ast::*;
7use derive_generic_visitor::{Continue, ControlFlow, Visitor};
8
9use super::ctx::UllbcPass;
10
11#[derive(Visitor)]
12struct InsertRegions<'a> {
13    regions: &'a mut Vector<RegionId, RegionVar>,
14    // The number of region groups we dived into (we don't count the regions
15    // at the declaration level). We use this for the DeBruijn indices.
16    depth: usize,
17}
18
19impl VisitAstMut for InsertRegions<'_> {
20    fn exit_region(&mut self, r: &mut Region) {
21        if r == &Region::Erased {
22            // Insert a fresh region
23            let index = self
24                .regions
25                .push_with(|index| RegionVar { index, name: None });
26            *r = Region::Var(DeBruijnVar::bound(DeBruijnId::new(self.depth), index));
27        }
28    }
29
30    fn visit_region_binder<T: AstVisitable>(
31        &mut self,
32        x: &mut RegionBinder<T>,
33    ) -> ControlFlow<Self::Break> {
34        self.depth += 1;
35        self.visit_inner(x)?;
36        self.depth -= 1;
37        Continue(())
38    }
39}
40
41fn transform_function(_ctx: &TransformCtx, def: &mut FunDecl) -> Result<(), Error> {
42    let FunSig {
43        closure_info,
44        inputs,
45        generics,
46        ..
47    } = &mut def.signature;
48    if let Some(info) = closure_info {
49        // Explore the state and introduce fresh regions for the erased regions we find.
50        let mut visitor = InsertRegions {
51            regions: &mut generics.regions,
52            depth: 0,
53        };
54        let _ = visitor.visit(&mut inputs[0]);
55
56        // Update the body.
57        // We change the type of the local variable of index 1, which is a reference to the closure
58        // itself, so that it has the type of (a reference to) the state of the closure.
59        //
60        // For instance, in the example below:
61        // ```text
62        // pub fn test_closure_capture(x: u32, y: u32) -> u32 {
63        //   let f = &|z| x + y + z;
64        //   (f)(0)
65        // }
66        // ```
67        //
68        // The first local variable in the body of the closure has type:
69        // `&'_ (fn(u32) -> u32)`
70        // We change it to:
71        // ```
72        // &'_ (&u32, &u32)
73        // ```
74        // We also update the corresponding field accesses in the body of
75        // the function.
76
77        if let Ok(body) = &mut def.body {
78            let body = body.as_unstructured_mut().unwrap();
79
80            // Update the type of the local 1 (which is the closure)
81            assert!(body.locals.locals.elem_count() > 1);
82            let state_var = &mut body.locals.locals[1];
83            state_var.ty = inputs[0].clone();
84            state_var.name = Some("state".to_string());
85
86            // Update the body, and in particular the accesses to the states
87            // TODO: translate to ADT field access
88            // TODO: handle directly during translation
89            let num_fields = info.state.elem_count();
90            body.body.dyn_visit_in_body_mut(|pe: &mut ProjectionElem| {
91                if let ProjectionElem::Field(pk @ FieldProjKind::ClosureState, _) = pe {
92                    *pk = FieldProjKind::Tuple(num_fields);
93                }
94            });
95        }
96
97        Ok(())
98    } else {
99        Ok(())
100    }
101}
102
103pub struct Transform;
104impl UllbcPass for Transform {
105    fn transform_function(&self, ctx: &mut TransformCtx, def: &mut FunDecl) {
106        // Ignore the errors, which should have been reported
107        let _ = transform_function(ctx, def);
108    }
109}
charon_lib/transform/update_closure_signatures.rs

charon_lib/transform/
update_closure_signatures.rs
