Source file closure_conversion.ml
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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
open Ast
open Utils
module D = Dcalc.Ast
(** TODO: This version is not yet debugged and ought to be specialized when
Lcalc has more structure. *)
type ctx = { name_context : string; globally_bound_vars : VarSet.t }
(** Returns the expression with closed closures and the set of free variables
inside this new expression. Implementation guided by
http://gallium.inria.fr/~fpottier/mpri/cours04.pdf#page=9. *)
let closure_conversion_expr (type m) (ctx : ctx) (e : m marked_expr) :
m marked_expr Bindlib.box =
let module MVarSet = Set.Make (struct
type t = m var
let compare = Bindlib.compare_vars
end) in
let rec aux e =
match Marked.unmark e with
| EVar v ->
( Bindlib.box_apply
(fun new_v -> new_v, Marked.get_mark e)
(Bindlib.box_var v),
if VarSet.mem (Var.t v) ctx.globally_bound_vars then MVarSet.empty
else MVarSet.singleton v )
| ETuple (args, s) ->
let new_args, free_vars =
List.fold_left
(fun (new_args, free_vars) arg ->
let new_arg, new_free_vars = aux arg in
new_arg :: new_args, MVarSet.union new_free_vars free_vars)
([], MVarSet.empty) args
in
( Bindlib.box_apply
(fun new_args -> ETuple (List.rev new_args, s), Marked.get_mark e)
(Bindlib.box_list new_args),
free_vars )
| ETupleAccess (e1, n, s, typs) ->
let new_e1, free_vars = aux e1 in
( Bindlib.box_apply
(fun new_e1 -> ETupleAccess (new_e1, n, s, typs), Marked.get_mark e)
new_e1,
free_vars )
| EInj (e1, n, e_name, typs) ->
let new_e1, free_vars = aux e1 in
( Bindlib.box_apply
(fun new_e1 -> EInj (new_e1, n, e_name, typs), Marked.get_mark e)
new_e1,
free_vars )
| EMatch (e1, arms, e_name) ->
let new_e1, free_vars = aux e1 in
let new_arms, free_vars =
List.fold_right
(fun arm (new_arms, free_vars) ->
match Marked.unmark arm with
| EAbs (binder, typs) ->
let vars, body = Bindlib.unmbind binder in
let new_body, new_free_vars = aux body in
let new_binder = Bindlib.bind_mvar vars new_body in
( Bindlib.box_apply
(fun new_binder ->
EAbs (new_binder, typs), Marked.get_mark arm)
new_binder
:: new_arms,
MVarSet.union free_vars new_free_vars )
| _ -> failwith "should not happen")
arms ([], free_vars)
in
( Bindlib.box_apply2
(fun new_e1 new_arms ->
EMatch (new_e1, new_arms, e_name), Marked.get_mark e)
new_e1
(Bindlib.box_list new_arms),
free_vars )
| EArray args ->
let new_args, free_vars =
List.fold_right
(fun arg (new_args, free_vars) ->
let new_arg, new_free_vars = aux arg in
new_arg :: new_args, MVarSet.union free_vars new_free_vars)
args ([], MVarSet.empty)
in
( Bindlib.box_apply
(fun new_args -> EArray new_args, Marked.get_mark e)
(Bindlib.box_list new_args),
free_vars )
| ELit l -> Bindlib.box (ELit l, Marked.get_mark e), MVarSet.empty
| EApp ((EAbs (binder, typs_abs), e1_pos), args) ->
let vars, body = Bindlib.unmbind binder in
let new_body, free_vars = aux body in
let new_binder = Bindlib.bind_mvar vars new_body in
let new_args, free_vars =
List.fold_right
(fun arg (new_args, free_vars) ->
let new_arg, new_free_vars = aux arg in
new_arg :: new_args, MVarSet.union free_vars new_free_vars)
args ([], free_vars)
in
( Bindlib.box_apply2
(fun new_binder new_args ->
( EApp ((EAbs (new_binder, typs_abs), e1_pos), new_args),
Marked.get_mark e ))
new_binder
(Bindlib.box_list new_args),
free_vars )
| EAbs (binder, typs) ->
let binder_mark = Marked.get_mark e in
let binder_pos = D.mark_pos binder_mark in
let vars, body = Bindlib.unmbind binder in
let new_body, body_vars = aux body in
let =
MVarSet.diff body_vars (MVarSet.of_list (Array.to_list vars))
in
let = MVarSet.elements extra_vars in
let code_var = new_var ctx.name_context in
let inner_c_var = new_var "env" in
let any_ty = Dcalc.Ast.TAny, binder_pos in
let new_closure_body =
make_multiple_let_in
(Array.of_list extra_vars_list)
(List.map (fun _ -> any_ty) extra_vars_list)
(List.mapi
(fun i _ ->
Bindlib.box_apply
(fun inner_c_var ->
( ETupleAccess
( (inner_c_var, binder_mark),
i + 1,
None,
List.map (fun _ -> any_ty) extra_vars_list ),
binder_mark ))
(Bindlib.box_var inner_c_var))
extra_vars_list)
new_body (D.mark_pos binder_mark)
in
let new_closure =
make_abs
(Array.concat [Array.make 1 inner_c_var; vars])
new_closure_body
((Dcalc.Ast.TAny, binder_pos) :: typs)
(Marked.get_mark e)
in
( make_let_in code_var
(Dcalc.Ast.TAny, D.pos e)
new_closure
(Bindlib.box_apply2
(fun code_var ->
( ETuple
( (code_var, binder_mark)
:: List.map
(fun -> extra_var, binder_mark)
extra_vars,
None ),
Marked.get_mark e ))
(Bindlib.box_var code_var)
(Bindlib.box_list
(List.map
(fun -> Bindlib.box_var extra_var)
extra_vars_list)))
(D.pos e),
extra_vars )
| EApp ((EOp op, pos_op), args) ->
let new_args, free_vars =
List.fold_right
(fun arg (new_args, free_vars) ->
let new_arg, new_free_vars = aux arg in
new_arg :: new_args, MVarSet.union free_vars new_free_vars)
args ([], MVarSet.empty)
in
( Bindlib.box_apply
(fun new_e2 -> EApp ((EOp op, pos_op), new_e2), Marked.get_mark e)
(Bindlib.box_list new_args),
free_vars )
| EApp ((EVar v, v_pos), args)
when VarSet.mem (Var.t v) ctx.globally_bound_vars ->
let new_args, free_vars =
List.fold_right
(fun arg (new_args, free_vars) ->
let new_arg, new_free_vars = aux arg in
new_arg :: new_args, MVarSet.union free_vars new_free_vars)
args ([], MVarSet.empty)
in
( Bindlib.box_apply2
(fun new_v new_e2 -> EApp ((new_v, v_pos), new_e2), Marked.get_mark e)
(Bindlib.box_var v)
(Bindlib.box_list new_args),
free_vars )
| EApp (e1, args) ->
let new_e1, free_vars = aux e1 in
let env_var = new_var "env" in
let code_var = new_var "code" in
let new_args, free_vars =
List.fold_right
(fun arg (new_args, free_vars) ->
let new_arg, new_free_vars = aux arg in
new_arg :: new_args, MVarSet.union free_vars new_free_vars)
args ([], free_vars)
in
let call_expr =
make_let_in code_var
(Dcalc.Ast.TAny, D.pos e)
(Bindlib.box_apply
(fun env_var ->
( ETupleAccess
((env_var, Marked.get_mark e1), 0, None, [ ]),
Marked.get_mark e ))
(Bindlib.box_var env_var))
(Bindlib.box_apply3
(fun code_var env_var new_args ->
( EApp
( (code_var, Marked.get_mark e1),
(env_var, Marked.get_mark e1) :: new_args ),
Marked.get_mark e ))
(Bindlib.box_var code_var) (Bindlib.box_var env_var)
(Bindlib.box_list new_args))
(D.pos e)
in
( make_let_in env_var (Dcalc.Ast.TAny, D.pos e) new_e1 call_expr (D.pos e),
free_vars )
| EAssert e1 ->
let new_e1, free_vars = aux e1 in
( Bindlib.box_apply
(fun new_e1 -> EAssert new_e1, Marked.get_mark e)
new_e1,
free_vars )
| EOp op -> Bindlib.box (EOp op, Marked.get_mark e), MVarSet.empty
| EIfThenElse (e1, e2, e3) ->
let new_e1, free_vars1 = aux e1 in
let new_e2, free_vars2 = aux e2 in
let new_e3, free_vars3 = aux e3 in
( Bindlib.box_apply3
(fun new_e1 new_e2 new_e3 ->
EIfThenElse (new_e1, new_e2, new_e3), Marked.get_mark e)
new_e1 new_e2 new_e3,
MVarSet.union (MVarSet.union free_vars1 free_vars2) free_vars3 )
| ERaise except ->
Bindlib.box (ERaise except, Marked.get_mark e), MVarSet.empty
| ECatch (e1, except, e2) ->
let new_e1, free_vars1 = aux e1 in
let new_e2, free_vars2 = aux e2 in
( Bindlib.box_apply2
(fun new_e1 new_e2 ->
ECatch (new_e1, except, new_e2), Marked.get_mark e)
new_e1 new_e2,
MVarSet.union free_vars1 free_vars2 )
in
let e', _vars = aux e in
e'
let closure_conversion (p : 'm program) : 'm program Bindlib.box =
let new_scopes, _ =
D.fold_left_scope_defs
~f:(fun (acc_new_scopes, global_vars) scope scope_var ->
let scope_input_var, scope_body_expr =
Bindlib.unbind scope.scope_body.scope_body_expr
in
let global_vars = VarSet.add (Var.t scope_var) global_vars in
let ctx =
{
name_context =
Marked.unmark (Dcalc.Ast.ScopeName.get_info scope.scope_name);
globally_bound_vars = global_vars;
}
in
let new_scope_lets =
D.map_exprs_in_scope_lets
~f:(closure_conversion_expr ctx)
~varf:(fun v -> v)
scope_body_expr
in
let new_scope_body_expr =
Bindlib.bind_var scope_input_var new_scope_lets
in
( (fun next ->
acc_new_scopes
(Bindlib.box_apply2
(fun new_scope_body_expr next ->
D.ScopeDef
{
scope with
scope_body =
{
scope.scope_body with
scope_body_expr = new_scope_body_expr;
};
scope_next = next;
})
new_scope_body_expr
(Bindlib.bind_var scope_var next))),
global_vars ))
~init:(Fun.id, VarSet.of_list [handle_default; handle_default_opt])
p.scopes
in
Bindlib.box_apply
(fun new_scopes -> { p with scopes = new_scopes })
(new_scopes (Bindlib.box D.Nil))