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
open Catala_utils
open Definitions
let map_exprs_in_lets :
?typ:(typ -> typ) ->
f:('expr1 -> 'expr2 boxed) ->
varf:('expr1 Var.t -> 'expr2 Var.t) ->
'expr1 scope_body_expr ->
'expr2 scope_body_expr Bindlib.box =
fun ?(typ = Fun.id) ~f ~varf scope_body_expr ->
let f e = Expr.Box.lift (f e) in
BoundList.map ~last:f
~f:(fun v scope_let ->
( varf v,
Bindlib.box_apply
(fun scope_let_expr ->
{
scope_let with
scope_let_expr;
scope_let_typ = typ scope_let.scope_let_typ;
})
(f scope_let.scope_let_expr) ))
scope_body_expr
let map_exprs ?(typ = Fun.id) ~f ~varf scopes =
let f v = function
| ScopeDef (name, body) ->
let scope_input_var, scope_lets = Bindlib.unbind body.scope_body_expr in
let new_body_expr = map_exprs_in_lets ~typ ~f ~varf scope_lets in
let new_body_expr =
Bindlib.bind_var (varf scope_input_var) new_body_expr
in
( varf v,
Bindlib.box_apply
(fun scope_body_expr ->
ScopeDef (name, { body with scope_body_expr }))
new_body_expr )
| Topdef (name, ty, expr) ->
( varf v,
Bindlib.box_apply
(fun e -> Topdef (name, typ ty, e))
(Expr.Box.lift (f expr)) )
in
BoundList.map ~f ~last:Bindlib.box scopes
let fold_exprs ~f ~init scopes =
let f acc def _ =
match def with
| Topdef (_, typ, e) -> f acc e typ
| ScopeDef (_, scope) ->
let _, body = Bindlib.unbind scope.scope_body_expr in
let acc, last =
BoundList.fold_left body ~init:acc ~f:(fun acc sl _ ->
f acc sl.scope_let_expr sl.scope_let_typ)
in
f acc last (TStruct scope.scope_body_output_struct, Expr.pos last)
in
fst @@ BoundList.fold_left ~f ~init scopes
let typ body =
let pos = Mark.get (StructName.get_info body.scope_body_input_struct) in
let input_typ = Mark.add pos (TStruct body.scope_body_input_struct) in
let result_typ = Mark.add pos (TStruct body.scope_body_output_struct) in
Mark.add pos (TArrow ([input_typ], result_typ))
let get_body_mark scope_body =
let m0 =
match Bindlib.unbind scope_body.scope_body_expr with
| _, Last (_, m) | _, Cons ({ scope_let_expr = _, m; _ }, _) -> m
in
Expr.with_ty m0 (typ scope_body)
let unfold_body_expr (_ctx : decl_ctx) (scope_let : 'e scope_body_expr) =
BoundList.fold_right scope_let ~init:Expr.rebox ~f:(fun sl var acc ->
Expr.make_let_in var sl.scope_let_typ
(Expr.rebox sl.scope_let_expr)
acc sl.scope_let_pos)
let input_type ty io =
match io, ty with
| (Runtime.Reentrant, iopos), (TArrow (args, ret), tpos) ->
TArrow (args, (TDefault ret, iopos)), tpos
| (Runtime.Reentrant, iopos), (ty, tpos) -> TDefault (ty, tpos), iopos
| _, ty -> ty
let to_expr (ctx : decl_ctx) (body : 'e scope_body) : 'e boxed =
let var, body_expr = Bindlib.unbind body.scope_body_expr in
let body_expr = unfold_body_expr ctx body_expr in
let pos = Expr.pos body_expr in
Expr.make_abs [| var |] body_expr
[TStruct body.scope_body_input_struct, pos]
pos
let unfold (ctx : decl_ctx) (s : 'e code_item_list) (main_scope : ScopeName.t) :
'e boxed =
BoundList.fold_lr s ~top:None
~down:(fun v item main ->
match main, item with
| None, ScopeDef (name, body) when ScopeName.equal name main_scope ->
Some (Expr.make_var v (get_body_mark body))
| r, _ -> r)
~bottom:(fun () -> function Some v -> v | None -> raise Not_found)
~up:(fun var item next ->
let e, typ =
match item with
| ScopeDef (_, body) -> to_expr ctx body, typ body
| Topdef (_, typ, expr) -> Expr.rebox expr, typ
in
Expr.make_let_in var typ e next (Expr.pos e))
let free_vars_body_expr scope_lets =
BoundList.fold_right scope_lets ~init:Expr.free_vars ~f:(fun sl v acc ->
Var.Set.union (Var.Set.remove v acc) (Expr.free_vars sl.scope_let_expr))
let free_vars_item = function
| ScopeDef (_, { scope_body_expr; _ }) ->
let v, body = Bindlib.unbind scope_body_expr in
Var.Set.remove v (free_vars_body_expr body)
| Topdef (_, _, expr) -> Expr.free_vars expr
let free_vars scopes =
BoundList.fold_right scopes
~init:(fun () -> Var.Set.empty)
~f:(fun item v acc ->
Var.Set.union (Var.Set.remove v acc) (free_vars_item item))