charon_lib/transform/
update_closure_signatures.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
//! # Micro-pass: the first local variable of closures is (a borrow to) the
//! closure itself. This is not consistent with the closure signature,
//! which ignores this first variable. This micro-pass updates this.
use derive_visitor::{visitor_enter_fn_mut, DriveMut, VisitorMut};

use crate::ids::Vector;
use crate::transform::TransformCtx;
use crate::ullbc_ast::*;

use super::ctx::UllbcPass;

#[derive(VisitorMut)]
#[visitor(Region(exit), Ty(enter, exit))]
struct InsertRegions<'a> {
    regions: &'a mut Vector<RegionId, RegionVar>,
    // The number of region groups we dived into (we don't count the regions
    // at the declaration level). We use this for the DeBruijn indices.
    depth: usize,
}

impl<'a> InsertRegions<'a> {
    fn exit_region(&mut self, r: &mut Region) {
        if r == &Region::Erased {
            // Insert a fresh region
            let index = self
                .regions
                .push_with(|index| RegionVar { index, name: None });
            *r = Region::BVar(DeBruijnId::new(self.depth), index);
        }
    }

    fn enter_ty(&mut self, ty: &mut Ty) {
        if let TyKind::Arrow(..) = ty.kind() {
            self.depth += 1;
        }
    }

    fn exit_ty(&mut self, ty: &mut Ty) {
        if let TyKind::Arrow(..) = ty.kind() {
            self.depth -= 1;
        }
    }
}

fn transform_function(
    _ctx: &TransformCtx,
    def: &mut FunDecl,
    body: Option<&mut ExprBody>,
) -> Result<(), Error> {
    let FunSig {
        closure_info,
        inputs,
        generics,
        ..
    } = &mut def.signature;
    if let Some(info) = closure_info {
        // Update the signature.
        // We add as first parameter the state of the closure, that is
        // a borrow to a tuple (of borrows, usually).
        // Remark: the types used in the closure state may contain erased
        // regions. In particular, the regions coming from the parent
        // function are often erased. TODO:
        // However, we introduce fresh regions for the state (in particular
        // because it is easy to do so).

        // Group the types into a tuple
        let num_fields = info.state.len();
        let state = TyKind::Adt(
            TypeId::Tuple,
            GenericArgs::new_from_types(info.state.clone()),
        )
        .into_ty();
        // Depending on the kind of the closure, add a reference
        let mut state = match &info.kind {
            ClosureKind::FnOnce => state,
            ClosureKind::Fn | ClosureKind::FnMut => {
                // We introduce an erased region, that we replace later
                //let index = RegionId::new(generics.regions.len());
                //generics.regions.push_back(RegionVar { index, name: None });

                let mutability = if info.kind == ClosureKind::Fn {
                    RefKind::Shared
                } else {
                    RefKind::Mut
                };
                //let r = Region::BVar(DeBruijnId::new(0), index);
                TyKind::Ref(Region::Erased, state, mutability).into_ty()
            }
        };

        // Explore the state and introduce fresh regions for the erased
        // regions we find.
        let mut visitor = Ty::visit_inside(InsertRegions {
            regions: &mut generics.regions,
            depth: 0,
        });
        state.drive_mut(&mut visitor);

        // Update the inputs (slightly annoying to push to the front of
        // a vector...).
        let mut original_inputs = std::mem::take(inputs);
        let mut ninputs = vec![state.clone()];
        ninputs.append(&mut original_inputs);
        *inputs = ninputs;

        // Update the body.
        // We change the type of the local variable of index 1, which is
        // a reference to the closure itself, so that it has the type of
        // (a reference to) the state of the closure.
        //
        // For instance, in the example below:
        // ```text
        // pub fn test_closure_capture(x: u32, y: u32) -> u32 {
        //   let f = &|z| x + y + z;
        //   (f)(0)
        // }
        // ```
        //
        // The first local variable in the body of the closure has type:
        // `&'_ (fn(u32) -> u32)`
        // We change it to:
        // ```
        // &'_ (&u32, &u32)
        // ```
        // We also update the corresponding field accesses in the body of
        // the function.

        if let Some(body) = body {
            body.locals.arg_count += 1;

            // Update the type of the local 1 (which is the closure)
            assert!(body.locals.vars.len() > 1);
            let state_var = &mut body.locals.vars[1];
            state_var.ty = state;
            state_var.name = Some("state".to_string());

            // Update the body, and in particular the accesses to the states
            body.body
                .drive_mut(&mut visitor_enter_fn_mut(|pe: &mut ProjectionElem| {
                    if let ProjectionElem::Field(pk @ FieldProjKind::ClosureState, _) = pe {
                        *pk = FieldProjKind::Tuple(num_fields);
                    }
                }));
        }

        Ok(())
    } else {
        Ok(())
    }
}

pub struct Transform;
impl UllbcPass for Transform {
    fn transform_function(
        &self,
        ctx: &mut TransformCtx,
        def: &mut FunDecl,
        body: Result<&mut ExprBody, Opaque>,
    ) {
        // Ignore the errors, which should have been reported
        let _ = transform_function(ctx, def, body.ok());
    }
}