mirror of
https://github.com/CatalaLang/catala.git
synced 2024-11-08 07:51:43 +03:00
666 lines
26 KiB
OCaml
666 lines
26 KiB
OCaml
(* This file is part of the Catala compiler, a specification language for tax
|
|
and social benefits computation rules. Copyright (C) 2022 Inria, contributor:
|
|
Denis Merigoux <denis.merigoux@inria.fr>
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License"); you may not
|
|
use this file except in compliance with the License. You may obtain a copy of
|
|
the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
License for the specific language governing permissions and limitations under
|
|
the License. *)
|
|
|
|
open Catala_utils
|
|
open Shared_ast
|
|
open Ast
|
|
module D = Dcalc.Ast
|
|
|
|
type 'm ctx = {
|
|
name_context : string;
|
|
globally_bound_vars : ('m expr, typ) Var.Map.t;
|
|
}
|
|
|
|
let tys_as_tanys tys = List.map (fun x -> Mark.map (fun _ -> TAny) x) tys
|
|
|
|
(** {1 Transforming closures}*)
|
|
|
|
(** 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=10
|
|
(environment-passing closure conversion). *)
|
|
let rec transform_closures_expr :
|
|
type m. m ctx -> m expr -> m expr Var.Set.t * m expr boxed =
|
|
fun ctx e ->
|
|
let m = Mark.get e in
|
|
match Mark.remove e with
|
|
| EStruct _ | EStructAccess _ | ETuple _ | ETupleAccess _ | EInj _ | EArray _
|
|
| ELit _ | EExternal _ | EAssert _ | EIfThenElse _ | ERaise _ | ECatch _ ->
|
|
Expr.map_gather ~acc:Var.Set.empty ~join:Var.Set.union
|
|
~f:(transform_closures_expr ctx)
|
|
e
|
|
| EVar v -> (
|
|
match Var.Map.find_opt v ctx.globally_bound_vars with
|
|
| None -> Var.Set.singleton v, (Bindlib.box_var v, m)
|
|
| Some (TArrow (targs, tret), _) ->
|
|
(* Here we eta-expand the argument to make sure function pointers are
|
|
correctly casted as closures *)
|
|
let args = Array.init (List.length targs) (fun _ -> Var.make "eta_arg") in
|
|
let arg_vars =
|
|
List.map2
|
|
(fun v ty -> Expr.evar v (Expr.with_ty m ty))
|
|
(Array.to_list args) targs
|
|
in
|
|
let e =
|
|
Expr.eabs
|
|
(Expr.bind args
|
|
(Expr.eapp ~f:(Expr.rebox e) ~args:arg_vars ~tys:targs
|
|
(Expr.with_ty m tret)))
|
|
targs m
|
|
in
|
|
let boxed =
|
|
let ctx =
|
|
(* We hide the type of the toplevel definition so that the function
|
|
doesn't loop *)
|
|
{
|
|
ctx with
|
|
globally_bound_vars =
|
|
Var.Map.add v (TAny, Pos.no_pos) ctx.globally_bound_vars;
|
|
}
|
|
in
|
|
Bindlib.box_apply (transform_closures_expr ctx) (Expr.Box.lift e)
|
|
in
|
|
Bindlib.unbox boxed
|
|
| Some _ -> Var.Set.empty, (Bindlib.box_var v, m))
|
|
| EMatch { e; cases; name } ->
|
|
let free_vars, new_e = (transform_closures_expr ctx) e in
|
|
(* We do not close the clotures inside the arms of the match expression,
|
|
since they get a special treatment at compilation to Scalc. *)
|
|
let free_vars, new_cases =
|
|
EnumConstructor.Map.fold
|
|
(fun cons e1 (free_vars, new_cases) ->
|
|
match Mark.remove e1 with
|
|
| EAbs { binder; tys } ->
|
|
let vars, body = Bindlib.unmbind binder in
|
|
let new_free_vars, new_body = (transform_closures_expr ctx) body in
|
|
let new_free_vars =
|
|
Array.fold_left
|
|
(fun acc v -> Var.Set.remove v acc)
|
|
new_free_vars vars
|
|
in
|
|
let new_binder = Expr.bind vars new_body in
|
|
( Var.Set.union free_vars
|
|
(Var.Set.diff new_free_vars
|
|
(Var.Set.of_list (Array.to_list vars))),
|
|
EnumConstructor.Map.add cons
|
|
(Expr.eabs new_binder tys (Mark.get e1))
|
|
new_cases )
|
|
| _ -> failwith "should not happen")
|
|
cases
|
|
(free_vars, EnumConstructor.Map.empty)
|
|
in
|
|
free_vars, Expr.ematch ~e:new_e ~name ~cases:new_cases m
|
|
| EApp { f = EAbs { binder; tys }, e1_pos; args; _ } ->
|
|
(* let-binding, we should not close these *)
|
|
let vars, body = Bindlib.unmbind binder in
|
|
let free_vars, new_body = (transform_closures_expr ctx) body in
|
|
let free_vars =
|
|
Array.fold_left (fun acc v -> Var.Set.remove v acc) free_vars vars
|
|
in
|
|
let new_binder = Expr.bind vars new_body in
|
|
let free_vars, new_args =
|
|
List.fold_right
|
|
(fun arg (free_vars, new_args) ->
|
|
let new_free_vars, new_arg = (transform_closures_expr ctx) arg in
|
|
Var.Set.union free_vars new_free_vars, new_arg :: new_args)
|
|
args (free_vars, [])
|
|
in
|
|
( free_vars,
|
|
Expr.eapp
|
|
~f:(Expr.eabs new_binder (tys_as_tanys tys) e1_pos)
|
|
~args:new_args ~tys m )
|
|
| EAbs { binder; tys } ->
|
|
(* λ x.t *)
|
|
let binder_mark = Expr.with_ty m (TAny, Expr.mark_pos m) in
|
|
let binder_pos = Expr.mark_pos binder_mark in
|
|
(* Converting the closure. *)
|
|
let vars, body = Bindlib.unmbind binder in
|
|
(* t *)
|
|
let body_vars, new_body = (transform_closures_expr ctx) body in
|
|
(* [[t]] *)
|
|
let extra_vars =
|
|
Var.Set.diff body_vars (Var.Set.of_list (Array.to_list vars))
|
|
in
|
|
let extra_vars_list = Var.Set.elements extra_vars in
|
|
(* x1, ..., xn *)
|
|
let code_var = Var.make ctx.name_context in
|
|
(* code *)
|
|
let closure_env_arg_var = Var.make "env" in
|
|
let closure_env_var = Var.make "env" in
|
|
let any_ty = TAny, binder_pos in
|
|
(* let env = from_closure_env env in let arg0 = env.0 in ... *)
|
|
let new_closure_body =
|
|
Expr.make_let_in closure_env_var any_ty
|
|
(Expr.eappop ~op:Operator.FromClosureEnv
|
|
~tys:[TClosureEnv, binder_pos]
|
|
~args:[Expr.evar closure_env_arg_var binder_mark]
|
|
binder_mark)
|
|
(Expr.make_multiple_let_in
|
|
(Array.of_list extra_vars_list)
|
|
(List.map (fun _ -> any_ty) extra_vars_list)
|
|
(List.mapi
|
|
(fun i _ ->
|
|
Expr.make_tupleaccess
|
|
(Expr.evar closure_env_var binder_mark)
|
|
i
|
|
(List.length extra_vars_list)
|
|
binder_pos)
|
|
extra_vars_list)
|
|
new_body binder_pos)
|
|
binder_pos
|
|
in
|
|
(* fun env arg0 ... -> new_closure_body *)
|
|
let new_closure =
|
|
Expr.make_abs
|
|
(Array.concat [Array.make 1 closure_env_arg_var; vars])
|
|
new_closure_body
|
|
((TClosureEnv, binder_pos) :: tys)
|
|
(Expr.pos e)
|
|
in
|
|
( extra_vars,
|
|
Expr.make_let_in code_var
|
|
(TAny, Expr.pos e)
|
|
new_closure
|
|
(Expr.make_tuple
|
|
((Bindlib.box_var code_var, binder_mark)
|
|
:: [
|
|
Expr.eappop ~op:Operator.ToClosureEnv
|
|
~tys:[TAny, Expr.pos e]
|
|
~args:
|
|
[
|
|
(if extra_vars_list = [] then Expr.elit LUnit binder_mark
|
|
else
|
|
Expr.etuple
|
|
(List.map
|
|
(fun extra_var ->
|
|
Bindlib.box_var extra_var, binder_mark)
|
|
extra_vars_list)
|
|
m);
|
|
]
|
|
(Mark.get e);
|
|
])
|
|
m)
|
|
(Expr.pos e) )
|
|
| EAppOp
|
|
{
|
|
op = (HandleDefaultOpt | Fold | Map | Filter | Reduce) as op;
|
|
tys;
|
|
args;
|
|
} ->
|
|
(* Special case for some operators: its arguments shall remain thunks (which
|
|
are closures) because if you want to extract it as a function you need
|
|
these closures to preserve evaluation order, but backends that don't
|
|
support closures will simply extract these operators in a inlined way and
|
|
skip the thunks. *)
|
|
let free_vars, new_args =
|
|
List.fold_right
|
|
(fun (arg : (lcalc, m) gexpr) (free_vars, new_args) ->
|
|
let m_arg = Mark.get arg in
|
|
match Mark.remove arg with
|
|
| EAbs { binder; tys } ->
|
|
let vars, arg = Bindlib.unmbind binder in
|
|
let new_free_vars, new_arg = (transform_closures_expr ctx) arg in
|
|
let new_arg =
|
|
Expr.make_abs vars new_arg tys (Expr.mark_pos m_arg)
|
|
in
|
|
Var.Set.union free_vars new_free_vars, new_arg :: new_args
|
|
| _ ->
|
|
let new_free_vars, new_arg = transform_closures_expr ctx arg in
|
|
Var.Set.union free_vars new_free_vars, new_arg :: new_args)
|
|
args (Var.Set.empty, [])
|
|
in
|
|
free_vars, Expr.eappop ~op ~tys ~args:new_args (Mark.get e)
|
|
| EAppOp _ ->
|
|
(* This corresponds to an operator call, which we don't want to transform *)
|
|
Expr.map_gather ~acc:Var.Set.empty ~join:Var.Set.union
|
|
~f:(transform_closures_expr ctx)
|
|
e
|
|
| EApp { f = EVar v, f_m; args; tys }
|
|
when Var.Map.mem v ctx.globally_bound_vars ->
|
|
(* This corresponds to a scope or toplevel function call, which we don't
|
|
want to transform *)
|
|
let free_vars, new_args =
|
|
List.fold_right
|
|
(fun arg (free_vars, new_args) ->
|
|
let new_free_vars, new_arg = (transform_closures_expr ctx) arg in
|
|
Var.Set.union free_vars new_free_vars, new_arg :: new_args)
|
|
args (Var.Set.empty, [])
|
|
in
|
|
free_vars, Expr.eapp ~f:(Expr.evar v f_m) ~args:new_args ~tys m
|
|
| EApp { f = e1; args; tys } ->
|
|
let free_vars, new_e1 = (transform_closures_expr ctx) e1 in
|
|
let code_env_var = Var.make "code_and_env" in
|
|
let code_env_expr =
|
|
let pos = Expr.pos e1 in
|
|
Expr.evar code_env_var
|
|
(Expr.with_ty (Mark.get e1)
|
|
( TTuple
|
|
[
|
|
( TArrow ((TClosureEnv, pos) :: tys, (TAny, Expr.pos e)),
|
|
Expr.pos e );
|
|
TClosureEnv, pos;
|
|
],
|
|
pos ))
|
|
in
|
|
let env_var = Var.make "env" in
|
|
let code_var = Var.make "code" in
|
|
let free_vars, new_args =
|
|
List.fold_right
|
|
(fun arg (free_vars, new_args) ->
|
|
let new_free_vars, new_arg = (transform_closures_expr ctx) arg in
|
|
Var.Set.union free_vars new_free_vars, new_arg :: new_args)
|
|
args (free_vars, [])
|
|
in
|
|
let call_expr =
|
|
let m1 = Mark.get e1 in
|
|
let pos = Expr.mark_pos m in
|
|
let env_arg_ty = TClosureEnv, Expr.pos e1 in
|
|
let fun_ty = TArrow (env_arg_ty :: tys, (TAny, Expr.pos e)), Expr.pos e in
|
|
Expr.make_multiple_let_in [| code_var; env_var |] [fun_ty; env_arg_ty]
|
|
[
|
|
Expr.make_tupleaccess code_env_expr 0 2 pos;
|
|
Expr.make_tupleaccess code_env_expr 1 2 pos;
|
|
]
|
|
(Expr.eapp
|
|
~f:(Bindlib.box_var code_var, m1)
|
|
~args:((Bindlib.box_var env_var, m1) :: new_args)
|
|
~tys:(env_arg_ty :: tys) m)
|
|
(Expr.pos e)
|
|
in
|
|
( free_vars,
|
|
Expr.make_let_in code_env_var
|
|
(TAny, Expr.pos e)
|
|
new_e1 call_expr (Expr.pos e) )
|
|
| _ -> .
|
|
|
|
let transform_closures_scope_let ctx scope_body_expr =
|
|
BoundList.map
|
|
~f:(fun var_next scope_let ->
|
|
let _free_vars, new_scope_let_expr =
|
|
(transform_closures_expr
|
|
{ ctx with name_context = Bindlib.name_of var_next })
|
|
scope_let.scope_let_expr
|
|
in
|
|
( var_next,
|
|
Bindlib.box_apply
|
|
(fun scope_let_expr ->
|
|
{
|
|
scope_let with
|
|
scope_let_expr;
|
|
scope_let_typ = Mark.copy scope_let.scope_let_typ TAny;
|
|
})
|
|
(Expr.Box.lift new_scope_let_expr) ))
|
|
~last:(fun res ->
|
|
let _free_vars, new_scope_let_expr = (transform_closures_expr ctx) res in
|
|
(* INVARIANT here: the result expr of a scope is simply a struct
|
|
containing all output variables so nothing should be converted here, so
|
|
no need to take into account free variables. *)
|
|
Expr.Box.lift new_scope_let_expr)
|
|
scope_body_expr
|
|
|
|
let transform_closures_program (p : 'm program) : 'm program Bindlib.box =
|
|
let (), new_code_items =
|
|
BoundList.fold_map
|
|
~f:(fun toplevel_vars var code_item ->
|
|
match code_item with
|
|
| ScopeDef (name, body) ->
|
|
let scope_input_var, scope_body_expr =
|
|
Bindlib.unbind body.scope_body_expr
|
|
in
|
|
let ctx =
|
|
{
|
|
name_context = Mark.remove (ScopeName.get_info name);
|
|
globally_bound_vars = toplevel_vars;
|
|
}
|
|
in
|
|
let new_scope_lets =
|
|
transform_closures_scope_let ctx scope_body_expr
|
|
in
|
|
let new_scope_body_expr =
|
|
Bindlib.bind_var scope_input_var new_scope_lets
|
|
in
|
|
let ty =
|
|
let pos = Mark.get (ScopeName.get_info name) in
|
|
( TArrow
|
|
( [TStruct body.scope_body_input_struct, pos],
|
|
(TStruct body.scope_body_output_struct, pos) ),
|
|
pos )
|
|
in
|
|
( Var.Map.add var ty toplevel_vars,
|
|
var,
|
|
Bindlib.box_apply
|
|
(fun scope_body_expr ->
|
|
ScopeDef (name, { body with scope_body_expr }))
|
|
new_scope_body_expr )
|
|
| Topdef (name, ty, (EAbs { binder; tys }, m)) ->
|
|
let v, expr = Bindlib.unmbind binder in
|
|
let ctx =
|
|
{
|
|
name_context = Mark.remove (TopdefName.get_info name);
|
|
globally_bound_vars = toplevel_vars;
|
|
}
|
|
in
|
|
let _free_vars, new_expr = transform_closures_expr ctx expr in
|
|
let new_binder = Expr.bind v new_expr in
|
|
( Var.Map.add var ty toplevel_vars,
|
|
var,
|
|
Bindlib.box_apply
|
|
(fun e -> Topdef (name, ty, e))
|
|
(Expr.Box.lift (Expr.eabs new_binder tys m)) )
|
|
| Topdef (name, ty, expr) ->
|
|
let ctx =
|
|
{
|
|
name_context = Mark.remove (TopdefName.get_info name);
|
|
globally_bound_vars = toplevel_vars;
|
|
}
|
|
in
|
|
let _free_vars, new_expr = transform_closures_expr ctx expr in
|
|
( Var.Map.add var ty toplevel_vars,
|
|
var,
|
|
Bindlib.box_apply
|
|
(fun e -> Topdef (name, (TAny, Mark.get ty), e))
|
|
(Expr.Box.lift new_expr) ))
|
|
~last:(fun _ () -> (), Bindlib.box ())
|
|
~init:Var.Map.empty p.code_items
|
|
in
|
|
(* Now we need to further tweak [decl_ctx] because some of the user-defined
|
|
types can have closures in them and these closured might have changed type.
|
|
So we reset them to [TAny] and leave the typechecker to figure it out. This
|
|
will not yield any type unification conflicts because of the special type
|
|
[TClosureEnv]. Indeed, consider the following closure: [let f = if ... then
|
|
fun v -> x + v else fun v -> v]. To be typed correctly once converted, this
|
|
closure needs an existential type, this is what [TClosureEnv] is for. This
|
|
kind of situation is difficult to produce using the Catala surface
|
|
language: it can only happen if you store a closure which is the output of
|
|
a scope inside a user-defined data structure, and if you do it in two
|
|
different places in the code with two closures that don't have the same
|
|
capture footprint. See
|
|
[tests/tests_func/good/scope_call_func_struct_closure.catala_en]. *)
|
|
let new_decl_ctx =
|
|
let rec replace_fun_typs t =
|
|
match Mark.remove t with
|
|
| TArrow (t1, t2) ->
|
|
( TTuple
|
|
[
|
|
( TArrow
|
|
( (TClosureEnv, Pos.no_pos) :: List.map replace_fun_typs t1,
|
|
replace_fun_typs t2 ),
|
|
Pos.no_pos );
|
|
TClosureEnv, Pos.no_pos;
|
|
],
|
|
Mark.get t )
|
|
| TDefault t' -> TDefault (replace_fun_typs t'), Mark.get t
|
|
| TOption t' -> TOption (replace_fun_typs t'), Mark.get t
|
|
| TAny | TClosureEnv | TLit _ | TEnum _ | TStruct _ -> t
|
|
| TArray ts -> TArray (replace_fun_typs ts), Mark.get t
|
|
| TTuple ts -> TTuple (List.map replace_fun_typs ts), Mark.get t
|
|
in
|
|
{
|
|
p.decl_ctx with
|
|
ctx_structs =
|
|
StructName.Map.map
|
|
(StructField.Map.map replace_fun_typs)
|
|
p.decl_ctx.ctx_structs;
|
|
ctx_enums =
|
|
EnumName.Map.map
|
|
(EnumConstructor.Map.map replace_fun_typs)
|
|
p.decl_ctx.ctx_enums;
|
|
(* Toplevel definitions may not contain scope calls or take functions as
|
|
arguments at the moment, which ensures that their interfaces aren't
|
|
changed by the conversion *)
|
|
}
|
|
in
|
|
Bindlib.box_apply
|
|
(fun new_code_items ->
|
|
{
|
|
code_items = new_code_items;
|
|
decl_ctx = new_decl_ctx;
|
|
module_name = p.module_name;
|
|
lang = p.lang;
|
|
})
|
|
new_code_items
|
|
|
|
(** {1 Hoisting closures}*)
|
|
|
|
type 'm hoisted_closure = {
|
|
name : 'm expr Var.t;
|
|
ty : typ;
|
|
closure : (lcalc, 'm) boxed_gexpr (* Starts with [EAbs]. *);
|
|
}
|
|
|
|
let rec hoist_closures_expr :
|
|
type m. string -> m expr -> m hoisted_closure list * m expr boxed =
|
|
fun name_context e ->
|
|
let m = Mark.get e in
|
|
match Mark.remove e with
|
|
| EMatch { e; cases; name } ->
|
|
let collected_closures, new_e = (hoist_closures_expr name_context) e in
|
|
(* We do not close the closures inside the arms of the match expression,
|
|
since they get a special treatment at compilation to Scalc. *)
|
|
let collected_closures, new_cases =
|
|
EnumConstructor.Map.fold
|
|
(fun cons e1 (collected_closures, new_cases) ->
|
|
match Mark.remove e1 with
|
|
| EAbs { binder; tys } ->
|
|
let vars, body = Bindlib.unmbind binder in
|
|
let new_collected_closures, new_body =
|
|
(hoist_closures_expr name_context) body
|
|
in
|
|
let new_binder = Expr.bind vars new_body in
|
|
( collected_closures @ new_collected_closures,
|
|
EnumConstructor.Map.add cons
|
|
(Expr.eabs new_binder tys (Mark.get e1))
|
|
new_cases )
|
|
| _ -> failwith "should not happen")
|
|
cases
|
|
(collected_closures, EnumConstructor.Map.empty)
|
|
in
|
|
collected_closures, Expr.ematch ~e:new_e ~name ~cases:new_cases m
|
|
| EApp { f = EAbs { binder; tys }, e1_pos; args; _ } ->
|
|
(* let-binding, we should not close these *)
|
|
let vars, body = Bindlib.unmbind binder in
|
|
let collected_closures, new_body =
|
|
(hoist_closures_expr name_context) body
|
|
in
|
|
let new_binder = Expr.bind vars new_body in
|
|
let collected_closures, new_args =
|
|
List.fold_right
|
|
(fun arg (collected_closures, new_args) ->
|
|
let new_collected_closures, new_arg =
|
|
(hoist_closures_expr name_context) arg
|
|
in
|
|
collected_closures @ new_collected_closures, new_arg :: new_args)
|
|
args (collected_closures, [])
|
|
in
|
|
( collected_closures,
|
|
Expr.eapp
|
|
~f:(Expr.eabs new_binder (tys_as_tanys tys) e1_pos)
|
|
~args:new_args ~tys m )
|
|
| EAppOp
|
|
{
|
|
op = (HandleDefaultOpt | Fold | Map | Filter | Reduce) as op;
|
|
tys;
|
|
args;
|
|
} ->
|
|
(* Special case for some operators: its arguments closures thunks because if
|
|
you want to extract it as a function you need these closures to preserve
|
|
evaluation order, but backends that don't support closures will simply
|
|
extract these operators in a inlined way and skip the thunks. *)
|
|
let collected_closures, new_args =
|
|
List.fold_right
|
|
(fun (arg : (lcalc, m) gexpr) (collected_closures, new_args) ->
|
|
let m_arg = Mark.get arg in
|
|
match Mark.remove arg with
|
|
| EAbs { binder; tys } ->
|
|
let vars, arg = Bindlib.unmbind binder in
|
|
let new_collected_closures, new_arg =
|
|
(hoist_closures_expr name_context) arg
|
|
in
|
|
let new_arg =
|
|
Expr.make_abs vars new_arg tys (Expr.mark_pos m_arg)
|
|
in
|
|
new_collected_closures @ collected_closures, new_arg :: new_args
|
|
| _ ->
|
|
let new_collected_closures, new_arg =
|
|
hoist_closures_expr name_context arg
|
|
in
|
|
new_collected_closures @ collected_closures, new_arg :: new_args)
|
|
args ([], [])
|
|
in
|
|
collected_closures, Expr.eappop ~op ~args:new_args ~tys (Mark.get e)
|
|
| EAbs { tys; _ } ->
|
|
(* this is the closure we want to hoist*)
|
|
let closure_var = Var.make ("closure_" ^ name_context) in
|
|
(* TODO: This will end up as a toplevel name. However for now we assume
|
|
toplevel names are unique, but this breaks this assertions and can lead
|
|
to name wrangling in the backends. We need to have a better system for
|
|
name disambiguation when for instance printing to Dcalc/Lcalc/Scalc but
|
|
also OCaml, Python, etc. *)
|
|
( [
|
|
{
|
|
name = closure_var;
|
|
ty = TArrow (tys, (TAny, Expr.mark_pos m)), Expr.mark_pos m;
|
|
closure = Expr.rebox e;
|
|
};
|
|
],
|
|
Expr.make_var closure_var m )
|
|
| EApp _ | EStruct _ | EStructAccess _ | ETuple _ | ETupleAccess _ | EInj _
|
|
| EArray _ | ELit _ | EAssert _ | EAppOp _ | EIfThenElse _ | ERaise _
|
|
| ECatch _ | EVar _ ->
|
|
Expr.map_gather ~acc:[] ~join:( @ ) ~f:(hoist_closures_expr name_context) e
|
|
| EExternal _ -> failwith "unimplemented"
|
|
| _ -> .
|
|
|
|
let hoist_closures_scope_let name_context scope_body_expr =
|
|
BoundList.fold_right
|
|
~f:(fun scope_let var_next (hoisted_closures, next_scope_lets) ->
|
|
let new_hoisted_closures, new_scope_let_expr =
|
|
(hoist_closures_expr (Bindlib.name_of var_next))
|
|
scope_let.scope_let_expr
|
|
in
|
|
( new_hoisted_closures @ hoisted_closures,
|
|
Bindlib.box_apply2
|
|
(fun scope_let_next scope_let_expr ->
|
|
Cons ({ scope_let with scope_let_expr }, scope_let_next))
|
|
(Bindlib.bind_var var_next next_scope_lets)
|
|
(Expr.Box.lift new_scope_let_expr) ))
|
|
~init:(fun res ->
|
|
let hoisted_closures, new_scope_let_expr =
|
|
(hoist_closures_expr name_context) res
|
|
in
|
|
(* INVARIANT here: the result expr of a scope is simply a struct
|
|
containing all output variables so nothing should be converted here, so
|
|
no need to take into account free variables. *)
|
|
( hoisted_closures,
|
|
Bindlib.box_apply
|
|
(fun res -> Last res)
|
|
(Expr.Box.lift new_scope_let_expr) ))
|
|
scope_body_expr
|
|
|
|
let rec hoist_closures_code_item_list
|
|
(code_items : (lcalc, 'm) gexpr code_item_list) :
|
|
(lcalc, 'm) gexpr code_item_list Bindlib.box =
|
|
match code_items with
|
|
| Last () -> Bindlib.box (Last ())
|
|
| Cons (code_item, next_code_items) ->
|
|
let code_item_var, next_code_items = Bindlib.unbind next_code_items in
|
|
let hoisted_closures, new_code_item =
|
|
match code_item with
|
|
| ScopeDef (name, body) ->
|
|
let scope_input_var, scope_body_expr =
|
|
Bindlib.unbind body.scope_body_expr
|
|
in
|
|
let new_hoisted_closures, new_scope_lets =
|
|
hoist_closures_scope_let
|
|
(fst (ScopeName.get_info name))
|
|
scope_body_expr
|
|
in
|
|
let new_scope_body_expr =
|
|
Bindlib.bind_var scope_input_var new_scope_lets
|
|
in
|
|
( new_hoisted_closures,
|
|
Bindlib.box_apply
|
|
(fun scope_body_expr ->
|
|
ScopeDef (name, { body with scope_body_expr }))
|
|
new_scope_body_expr )
|
|
| Topdef (name, ty, (EAbs { binder; tys }, m)) ->
|
|
let v, expr = Bindlib.unmbind binder in
|
|
let new_hoisted_closures, new_expr =
|
|
hoist_closures_expr (Mark.remove (TopdefName.get_info name)) expr
|
|
in
|
|
let new_binder = Expr.bind v new_expr in
|
|
( new_hoisted_closures,
|
|
Bindlib.box_apply
|
|
(fun e -> Topdef (name, ty, e))
|
|
(Expr.Box.lift (Expr.eabs new_binder tys m)) )
|
|
| Topdef (name, ty, expr) ->
|
|
let new_hoisted_closures, new_expr =
|
|
hoist_closures_expr (Mark.remove (TopdefName.get_info name)) expr
|
|
in
|
|
( new_hoisted_closures,
|
|
Bindlib.box_apply
|
|
(fun e -> Topdef (name, (TAny, Mark.get ty), e))
|
|
(Expr.Box.lift new_expr) )
|
|
in
|
|
let next_code_items = hoist_closures_code_item_list next_code_items in
|
|
let next_code_items =
|
|
Bindlib.box_apply2
|
|
(fun next_code_items new_code_item ->
|
|
Cons (new_code_item, next_code_items))
|
|
(Bindlib.bind_var code_item_var next_code_items)
|
|
new_code_item
|
|
in
|
|
let next_code_items =
|
|
List.fold_left
|
|
(fun (next_code_items : (lcalc, 'm) gexpr code_item_list Bindlib.box)
|
|
(hoisted_closure : 'm hoisted_closure) ->
|
|
let next_code_items =
|
|
Bindlib.bind_var hoisted_closure.name next_code_items
|
|
in
|
|
let closure, closure_mark = hoisted_closure.closure in
|
|
Bindlib.box_apply2
|
|
(fun next_code_items closure ->
|
|
Cons
|
|
( Topdef
|
|
( TopdefName.fresh []
|
|
( Bindlib.name_of hoisted_closure.name,
|
|
Expr.mark_pos closure_mark ),
|
|
hoisted_closure.ty,
|
|
(closure, closure_mark) ),
|
|
next_code_items ))
|
|
next_code_items closure)
|
|
next_code_items hoisted_closures
|
|
in
|
|
next_code_items
|
|
|
|
let hoist_closures_program (p : 'm program) : 'm program Bindlib.box =
|
|
let new_code_items = hoist_closures_code_item_list p.code_items in
|
|
(*TODO: we need to insert the hoisted closures just before the scopes they
|
|
belong to, because some of them call sub-scopes and putting them all at the
|
|
beginning breaks dependency ordering. *)
|
|
Bindlib.box_apply
|
|
(fun new_code_items -> { p with code_items = new_code_items })
|
|
new_code_items
|
|
|
|
(** {1 Closure conversion}*)
|
|
|
|
let closure_conversion (p : 'm program) : untyped program =
|
|
let new_p = transform_closures_program p in
|
|
let new_p = hoist_closures_program (Bindlib.unbox new_p) in
|
|
(* FIXME: either fix the types of the marks, or remove the types annotations
|
|
during the main processing (rather than requiring a new traversal) *)
|
|
Program.untype (Bindlib.unbox new_p)
|