translation__fun_to_clos.ml

(** {0 Typed closure conversion on nameless terms}

    Translation from {!Lang.Fun} to {!Lang.Clos}.

    This translation converts functions into closures, separating the code of
    functions from the data implicitly captured from the surrounding environment.
    It’s a modified version of the algorithm described in {{:https://doi.org/10.1145/1291201.1291212}
    “A type-preserving closure conversion in haskell”}, using de Bruijn levels
    to simplify the handling of variable bindings.

    The implementation is somewhat fiddly due to the decision to use de Bruijn
    indices in the source and target languages. {!Translate.Fun_a_to_clos_a} uses
    alpha renamed terms, which simplifies this somewhat.

    Note that we don’t do any uncurrying during this translation, which results
    in more closures being constructed than might otherwise be necessary.
*)

module Fun = Lang.Fun
module Clos = Lang.Clos


(** {1 Values} *)

(** Values used when substituting source terms for closure converted target
    terms. We only need to worry about variables and projections on variables,
    as those are the only values that will be stored in the environment during
    compilation.

    Variables are represented using de Bruijn levels, allowing us to freely
    weaken the target environment (i.e. add new bindings) without needing to
    worry about shifting de Bruijn indices.
*)
type vtm =
  | Var of int
  | Tuple_proj of int * int

(** Quote a value back into a target environment. *)
let quote size' : vtm -> Clos.tm =
  function
  | Var level -> Var (size' - level - 1)
  | Tuple_proj (level, label) -> Tuple_proj (Var (size' - level - 1), label)


(** {1 Environment} *)

(** Compilation environment that maps source variables to closure converted
    values and their types. *)
type env = (vtm * Clos.ty) option list


(** {1 Translation} *)

(** Translation to closure converted types *)
let rec translate_ty : Fun.ty -> Clos.ty =
  function
  | Bool_type -> Bool_type
  | Int_type -> Int_type
  | Fun_type (param_ty, body_ty) ->
      Clos_type (translate_ty param_ty, translate_ty body_ty)

(** Translation to closure converted terms.

    Two size parameters are supplied:

    - [size]: the size of the source environment, used for computing free variables
    - [size']: the size of the target environment, used for quoting values into
      closure converted terms
*)
let rec translate env size size' : Fun.tm -> Clos.tm =
  function
  | Var index ->
      let vtm, _ = Option.get (List.nth env index) in
      quote size' vtm

  | Let (name, def_ty, def, body) ->
      let def_ty = translate_ty def_ty in
      let def = translate env size size' def in

      (* Compile the body of the let expression, adding a binding to the source
        and target environments respectively. *)
      let body_env = Some (Var size', def_ty) :: env in
      let body = translate body_env (size + 1) (size' + 1) body in

      Let (name, def_ty, def, body)

  | Bool_lit b -> Bool_lit b
  | Int_lit i -> Int_lit i

  | Prim_app (prim, args) ->
      let args = List.map (translate env size size') args in
      Prim_app (prim, args)

  | Fun_lit (name, param_ty, body) ->
      let param_ty = translate_ty param_ty in

      (* There are only two variables bound in the environment of the compiled
         code literal: one for the environment parameter, and another for the
         original parameter of the function. *)
      let env_level = 0 in
      let param_level = 1 in

      (* Create a mask over the environment that records the free variables used
         in the body of the function. *)
      let body_fvs = Fun.fvs (size + 1) body in

      (* Prunes the current environment with the bitmask, returning:

        - an environmnent that substitutes source variables that were used in
          the body of the function to projections off the environment parameter
        - a list of terms (in reverse order) to be used when constructing the environment
        - a list of types (in reverse order) to be used in the type of the environment
      *)
      let rec make_env env index =
        match env with
        | [] -> [], [], []
        | binding :: env ->
            let env', tms, tys = make_env env (index + 1) in
            (* Check if the current binding was used in the body *)
            if body_fvs.(size - index - 1) then
              (* This binding was used in the body of the function, so add
                it to the environment tuple and map any occurrences in the body
                of the closure to projections off the environment parameter. *)
              let vtm, ty = Option.get binding in
              Some (Tuple_proj (env_level, List.length tms), ty) :: env',
              quote size' vtm :: tms,
              ty :: tys
            else
              (* This binding was not used in the body of the function, so
                encountering it in the body of the closure is a bug. *)
              None :: env', tms, tys
      in

      (* Construct the environments *)
      let proj_env, env_tms, env_tys = make_env env 0 in

      (* Translate the body of the function *)
      let body_env = Some (Var param_level, param_ty) :: proj_env in
      let body = translate body_env (size + 1) 2 body in

      Clos_lit
        (Code_lit (Tuple_type (List.rev env_tys), (name, param_ty), body),
          Tuple_lit (List.rev env_tms))

  | Fun_app (head, arg) ->
      let head = translate env size size' head in
      let arg = translate env size size' arg in
      Clos_app (head, arg)