mirror of
https://github.com/CatalaLang/catala.git
synced 2024-09-20 00:41:05 +03:00
323 lines
13 KiB
OCaml
323 lines
13 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 Ast
|
|
open Utils
|
|
|
|
(** TODO: This version is not yet debugged and ought to be specialized when
|
|
Lcalc has more structure. *)
|
|
|
|
type closure = { name : Var.t; expr : expr Pos.marked Bindlib.box }
|
|
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 rec closure_conversion_expr (ctx : ctx) (e : expr Pos.marked) :
|
|
expr Pos.marked Bindlib.box * VarSet.t =
|
|
match Pos.unmark e with
|
|
| EVar v ->
|
|
( Bindlib.box_apply
|
|
(fun new_v -> (new_v, Pos.get_position v))
|
|
(Bindlib.box_var (Pos.unmark v)),
|
|
VarSet.diff (VarSet.singleton (Pos.unmark v)) ctx.globally_bound_vars )
|
|
| ETuple (args, s) ->
|
|
let new_args, free_vars =
|
|
List.fold_left
|
|
(fun (new_args, free_vars) arg ->
|
|
let new_arg, new_free_vars = closure_conversion_expr ctx arg in
|
|
(new_arg :: new_args, VarSet.union new_free_vars free_vars))
|
|
([], VarSet.empty) args
|
|
in
|
|
( Bindlib.box_apply
|
|
(fun new_args -> (ETuple (List.rev new_args, s), Pos.get_position e))
|
|
(Bindlib.box_list new_args),
|
|
free_vars )
|
|
| ETupleAccess (e1, n, s, typs) ->
|
|
let new_e1, free_vars = closure_conversion_expr ctx e1 in
|
|
( Bindlib.box_apply
|
|
(fun new_e1 ->
|
|
(ETupleAccess (new_e1, n, s, typs), Pos.get_position e))
|
|
new_e1,
|
|
free_vars )
|
|
| EInj (e1, n, e_name, typs) ->
|
|
let new_e1, free_vars = closure_conversion_expr ctx e1 in
|
|
( Bindlib.box_apply
|
|
(fun new_e1 -> (EInj (new_e1, n, e_name, typs), Pos.get_position e))
|
|
new_e1,
|
|
free_vars )
|
|
| EMatch (e1, arms, e_name) ->
|
|
let new_e1, free_vars = closure_conversion_expr ctx e1 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 new_arms, free_vars =
|
|
List.fold_right
|
|
(fun arm (new_arms, free_vars) ->
|
|
match Pos.unmark arm with
|
|
| EAbs ((binder, binder_pos), typs) ->
|
|
let vars, body = Bindlib.unmbind binder in
|
|
let new_body, new_free_vars =
|
|
closure_conversion_expr ctx body
|
|
in
|
|
let new_binder = Bindlib.bind_mvar vars new_body in
|
|
( Bindlib.box_apply
|
|
(fun new_binder ->
|
|
( EAbs ((new_binder, binder_pos), typs),
|
|
Pos.get_position arm ))
|
|
new_binder
|
|
:: new_arms,
|
|
VarSet.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), Pos.get_position 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 = closure_conversion_expr ctx arg in
|
|
(new_arg :: new_args, VarSet.union free_vars new_free_vars))
|
|
args ([], VarSet.empty)
|
|
in
|
|
( Bindlib.box_apply
|
|
(fun new_args -> (EArray new_args, Pos.get_position e))
|
|
(Bindlib.box_list new_args),
|
|
free_vars )
|
|
| ELit l -> (Bindlib.box (ELit l, Pos.get_position e), VarSet.empty)
|
|
| EApp ((EAbs ((binder, binder_pos), typs_abs), e1_pos), args) ->
|
|
(* let-binding, we should not close these *)
|
|
let vars, body = Bindlib.unmbind binder in
|
|
let new_body, free_vars = closure_conversion_expr ctx 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 = closure_conversion_expr ctx arg in
|
|
(new_arg :: new_args, VarSet.union free_vars new_free_vars))
|
|
args ([], free_vars)
|
|
in
|
|
( Bindlib.box_apply2
|
|
(fun new_binder new_args ->
|
|
( EApp
|
|
((EAbs ((new_binder, binder_pos), typs_abs), e1_pos), new_args),
|
|
Pos.get_position e ))
|
|
new_binder
|
|
(Bindlib.box_list new_args),
|
|
free_vars )
|
|
| EAbs ((binder, binder_pos), typs) ->
|
|
(* λ x.t *)
|
|
(* Converting the closure. *)
|
|
let vars, body = Bindlib.unmbind binder in
|
|
(* t *)
|
|
let new_body, body_vars = closure_conversion_expr ctx body in
|
|
(* [[t]] *)
|
|
let extra_vars =
|
|
VarSet.diff body_vars (VarSet.of_list (Array.to_list vars))
|
|
in
|
|
let extra_vars_list = VarSet.elements extra_vars in
|
|
(* x1, ..., xn *)
|
|
let code_var = Var.make (ctx.name_context, binder_pos) in
|
|
(* code *)
|
|
let inner_c_var = Var.make ("env", binder_pos) in
|
|
let new_closure_body =
|
|
make_multiple_let_in
|
|
(Array.of_list extra_vars_list)
|
|
(List.init (List.length extra_vars_list) (fun _ ->
|
|
(Dcalc.Ast.TAny, binder_pos)))
|
|
(List.mapi
|
|
(fun i _ ->
|
|
Bindlib.box_apply
|
|
(fun inner_c_var ->
|
|
( ETupleAccess
|
|
( (inner_c_var, binder_pos),
|
|
i + 1,
|
|
None,
|
|
List.init
|
|
(List.length extra_vars_list + 1)
|
|
(fun _ -> (Dcalc.Ast.TAny, binder_pos)) ),
|
|
binder_pos ))
|
|
(Bindlib.box_var inner_c_var))
|
|
extra_vars_list)
|
|
new_body binder_pos
|
|
in
|
|
let new_closure =
|
|
make_abs
|
|
(Array.concat [ Array.make 1 inner_c_var; vars ])
|
|
new_closure_body binder_pos
|
|
((Dcalc.Ast.TAny, binder_pos) :: typs)
|
|
(Pos.get_position e)
|
|
in
|
|
( make_let_in code_var
|
|
(Dcalc.Ast.TAny, Pos.get_position e)
|
|
new_closure
|
|
(Bindlib.box_apply2
|
|
(fun code_var extra_vars ->
|
|
( ETuple
|
|
( (code_var, binder_pos)
|
|
:: List.map
|
|
(fun extra_var -> (extra_var, binder_pos))
|
|
extra_vars,
|
|
None ),
|
|
Pos.get_position e ))
|
|
(Bindlib.box_var code_var)
|
|
(Bindlib.box_list
|
|
(List.map
|
|
(fun extra_var -> Bindlib.box_var extra_var)
|
|
extra_vars_list)))
|
|
(Pos.get_position e),
|
|
extra_vars )
|
|
| EApp ((EOp op, pos_op), args) ->
|
|
(* This corresponds to an operator call, which we don't want to
|
|
transform*)
|
|
let new_args, free_vars =
|
|
List.fold_right
|
|
(fun arg (new_args, free_vars) ->
|
|
let new_arg, new_free_vars = closure_conversion_expr ctx arg in
|
|
(new_arg :: new_args, VarSet.union free_vars new_free_vars))
|
|
args ([], VarSet.empty)
|
|
in
|
|
( Bindlib.box_apply
|
|
(fun new_e2 -> (EApp ((EOp op, pos_op), new_e2), Pos.get_position e))
|
|
(Bindlib.box_list new_args),
|
|
free_vars )
|
|
| EApp ((EVar (v, _), v_pos), args) when VarSet.mem v ctx.globally_bound_vars
|
|
->
|
|
(* This corresponds to a scope call, which we don't want to transform*)
|
|
let new_args, free_vars =
|
|
List.fold_right
|
|
(fun arg (new_args, free_vars) ->
|
|
let new_arg, new_free_vars = closure_conversion_expr ctx arg in
|
|
(new_arg :: new_args, VarSet.union free_vars new_free_vars))
|
|
args ([], VarSet.empty)
|
|
in
|
|
( Bindlib.box_apply2
|
|
(fun new_v new_e2 ->
|
|
(EApp ((new_v, v_pos), new_e2), Pos.get_position e))
|
|
(Bindlib.box_var v)
|
|
(Bindlib.box_list new_args),
|
|
free_vars )
|
|
| EApp (e1, args) ->
|
|
let new_e1, free_vars = closure_conversion_expr ctx e1 in
|
|
let env_var = Var.make ("env", Pos.get_position e1) in
|
|
let code_var = Var.make ("code", Pos.get_position e1) in
|
|
let new_args, free_vars =
|
|
List.fold_right
|
|
(fun arg (new_args, free_vars) ->
|
|
let new_arg, new_free_vars = closure_conversion_expr ctx arg in
|
|
(new_arg :: new_args, VarSet.union free_vars new_free_vars))
|
|
args ([], free_vars)
|
|
in
|
|
let call_expr =
|
|
make_let_in code_var
|
|
(Dcalc.Ast.TAny, Pos.get_position e)
|
|
(Bindlib.box_apply
|
|
(fun env_var ->
|
|
( ETupleAccess
|
|
((env_var, Pos.get_position e1), 0, None, [ (*TODO: fill?*) ]),
|
|
Pos.get_position e ))
|
|
(Bindlib.box_var env_var))
|
|
(Bindlib.box_apply3
|
|
(fun code_var env_var new_args ->
|
|
( EApp
|
|
( (code_var, Pos.get_position e1),
|
|
(env_var, Pos.get_position e1) :: new_args ),
|
|
Pos.get_position e ))
|
|
(Bindlib.box_var code_var) (Bindlib.box_var env_var)
|
|
(Bindlib.box_list new_args))
|
|
(Pos.get_position e)
|
|
in
|
|
( make_let_in env_var
|
|
(Dcalc.Ast.TAny, Pos.get_position e)
|
|
new_e1 call_expr (Pos.get_position e),
|
|
free_vars )
|
|
| EAssert e1 ->
|
|
let new_e1, free_vars = closure_conversion_expr ctx e1 in
|
|
( Bindlib.box_apply
|
|
(fun new_e1 -> (EAssert new_e1, Pos.get_position e))
|
|
new_e1,
|
|
free_vars )
|
|
| EOp op -> (Bindlib.box (EOp op, Pos.get_position e), VarSet.empty)
|
|
| EIfThenElse (e1, e2, e3) ->
|
|
let new_e1, free_vars1 = closure_conversion_expr ctx e1 in
|
|
let new_e2, free_vars2 = closure_conversion_expr ctx e2 in
|
|
let new_e3, free_vars3 = closure_conversion_expr ctx e3 in
|
|
( Bindlib.box_apply3
|
|
(fun new_e1 new_e2 new_e3 ->
|
|
(EIfThenElse (new_e1, new_e2, new_e3), Pos.get_position e))
|
|
new_e1 new_e2 new_e3,
|
|
VarSet.union (VarSet.union free_vars1 free_vars2) free_vars3 )
|
|
| ERaise except ->
|
|
(Bindlib.box (ERaise except, Pos.get_position e), VarSet.empty)
|
|
| ECatch (e1, except, e2) ->
|
|
let new_e1, free_vars1 = closure_conversion_expr ctx e1 in
|
|
let new_e2, free_vars2 = closure_conversion_expr ctx e2 in
|
|
( Bindlib.box_apply2
|
|
(fun new_e1 new_e2 ->
|
|
(ECatch (new_e1, except, new_e2), Pos.get_position e))
|
|
new_e1 new_e2,
|
|
VarSet.union free_vars1 free_vars2 )
|
|
|
|
let closure_conversion (p : program) : program Bindlib.box * closure list =
|
|
let all_scope_variables =
|
|
List.fold_left
|
|
(fun acc scope -> VarSet.add scope.scope_body_var acc)
|
|
VarSet.empty p.scopes
|
|
in
|
|
let new_scopes, closures =
|
|
List.fold_left
|
|
(fun ((acc_new_scopes, acc_closures) :
|
|
scope_body Bindlib.box list * closure list) (scope : scope_body) ->
|
|
match Pos.unmark scope.scope_body_expr with
|
|
| EAbs ((binder, binder_pos), typs) ->
|
|
(* We do not hoist the outer-most EAbs which is the scope function
|
|
itself *)
|
|
let vars, body = Bindlib.unmbind binder in
|
|
let ctx =
|
|
{
|
|
name_context =
|
|
Pos.unmark
|
|
(Dcalc.Ast.ScopeName.get_info scope.scope_body_name);
|
|
globally_bound_vars =
|
|
VarSet.union all_scope_variables
|
|
(VarSet.of_list [ handle_default; handle_default_opt ]);
|
|
}
|
|
in
|
|
let new_body_expr, _ = closure_conversion_expr ctx body in
|
|
let new_binder = Bindlib.bind_mvar vars new_body_expr in
|
|
( Bindlib.box_apply
|
|
(fun new_binder ->
|
|
{
|
|
scope with
|
|
scope_body_expr =
|
|
( EAbs ((new_binder, binder_pos), typs),
|
|
Pos.get_position scope.scope_body_expr );
|
|
})
|
|
new_binder
|
|
:: acc_new_scopes,
|
|
[] @ acc_closures )
|
|
| _ -> failwith "should not happen")
|
|
([], []) p.scopes
|
|
in
|
|
( Bindlib.box_apply
|
|
(fun new_scopes -> { p with scopes = List.rev new_scopes })
|
|
(Bindlib.box_list new_scopes),
|
|
closures )
|