TS backend phase 0 - Typed lambda

compiler/bsc/rescript_compiler_main.ml

@@ -397,6 +397,12 @@ let command_line_flags : (string * Bsc_args.spec * string) array =
     ("-dtypedtree", set Clflags.dump_typedtree, "*internal* debug typedtree");
     ("-dparsetree", set Clflags.dump_parsetree, "*internal* debug parsetree");
     ("-drawlambda", set Clflags.dump_rawlambda, "*internal* debug raw lambda");
+    ( "-dlamtypes",
+      set Clflags.dump_lamtypes,
+      "*internal* dump Lam IR type annotations" );
+    ( "-emit-typedefs",
+      set Clflags.emit_typedefs,
+      "*internal* emit .d.ts declarations" );
     ("-dsource", set Clflags.dump_source, "*internal* print source");
     ( "-reprint-source",
       string_call reprint_source_file,

compiler/core/js_implementation.ml

@@ -133,7 +133,7 @@ let after_parsing_impl ppf outputprefix (ast : Parsetree.structure) =
       Lam_compile_env.reset ();
       let env = Res_compmisc.initial_env ~modulename () in
       Env.set_unit_name modulename;
-      let typedtree, coercion, _, _ =
+      let typedtree, coercion, finalenv, _ =
         Typemod.type_implementation_more
           ?check_exists:(if !Js_config.force_cmi then None else Some ())
           !Location.input_name outputprefix modulename env ast
@@ -143,12 +143,13 @@ let after_parsing_impl ppf outputprefix (ast : Parsetree.structure) =
         Printtyped.implementation_with_coercion typedtree_coercion;
       (if !Js_config.cmi_only then Warnings.check_fatal ()
        else
-         let lambda, exports =
-           Translmod.transl_implementation modulename typedtree_coercion
+         let env, lambda, exports =
+           Translmod.transl_implementation modulename finalenv
+             typedtree_coercion
          in
          let js_program =
            print_if_pipe ppf Clflags.dump_rawlambda Printlambda.lambda lambda
-           |> Lam_compile_main.compile outputprefix exports
+           |> Lam_compile_main.compile ~typedtree outputprefix exports
          in
          if not !Js_config.cmj_only then
            Lam_compile_main.lambda_as_module js_program outputprefix);

compiler/core/lam.ml

@@ -31,7 +31,7 @@ type ap_info = {
   ap_status: apply_status;
 }
 
-module Types = struct
+module Lam_types = struct
   type lambda_switch = {
     sw_consts_full: bool;
     (* TODO: refine its representation *)
@@ -47,6 +47,7 @@ module Types = struct
     params: ident list;
     body: t;
     attr: Lambda.function_attribute;
+    ty: Types.type_expr option;
   }
 
   (*
@@ -86,6 +87,7 @@ module Types = struct
     ap_args: t list;
     ap_info: ap_info;
     ap_transformed_jsx: bool;
+    ap_result_type: Types.type_expr option;
   }
 
   and t =
@@ -94,7 +96,7 @@ module Types = struct
     | Lconst of Lam_constant.t
     | Lapply of apply
     | Lfunction of lfunction
-    | Llet of Lam_compat.let_kind * ident * t * t
+    | Llet of Lam_compat.let_kind * ident * Types.type_expr option * t * t
     | Lletrec of (ident * t) list * t
     | Lprim of prim_info
     | Lswitch of t * lambda_switch
@@ -115,7 +117,7 @@ module Types = struct
 end
 
 module X = struct
-  type lambda_switch = Types.lambda_switch = {
+  type lambda_switch = Lam_types.lambda_switch = {
     sw_consts_full: bool;
     sw_consts: (int * t) list;
     sw_blocks_full: bool;
@@ -124,33 +126,35 @@ module X = struct
     sw_names: Ast_untagged_variants.switch_names option;
   }
 
-  and prim_info = Types.prim_info = {
+  and prim_info = Lam_types.prim_info = {
     primitive: Lam_primitive.t;
     args: t list;
     loc: Location.t;
   }
 
-  and apply = Types.apply = {
+  and apply = Lam_types.apply = {
     ap_func: t;
     ap_args: t list;
     ap_info: ap_info;
     ap_transformed_jsx: bool;
+    ap_result_type: Types.type_expr option;
   }
 
-  and lfunction = Types.lfunction = {
+  and lfunction = Lam_types.lfunction = {
     arity: int;
     params: ident list;
     body: t;
     attr: Lambda.function_attribute;
+    ty: Types.type_expr option;
   }
 
-  and t = Types.t =
+  and t = Lam_types.t =
     | Lvar of ident
     | Lglobal_module of ident * bool
     | Lconst of Lam_constant.t
     | Lapply of apply
     | Lfunction of lfunction
-    | Llet of Lam_compat.let_kind * ident * t * t
+    | Llet of Lam_compat.let_kind * ident * Types.type_expr option * t * t
     | Lletrec of (ident * t) list * t
     | Lprim of prim_info
     | Lswitch of t * lambda_switch
@@ -170,24 +174,24 @@ module X = struct
   (* | Lsend of Lam_compat.meth_kind * t * t * t list * Location.t *)
 end
 
-include Types
+include Lam_types
 
 (** apply [f] to direct successor which has type [Lam.t] *)
 
 let inner_map (l : t) (f : t -> X.t) : X.t =
   match l with
   | Lvar (_ : ident) | Lconst (_ : Lam_constant.t) -> ((* Obj.magic *) l : X.t)
-  | Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx} ->
+  | Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx; ap_result_type} ->
     let ap_func = f ap_func in
     let ap_args = Ext_list.map ap_args f in
-    Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx}
-  | Lfunction {body; arity; params; attr} ->
+    Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx; ap_result_type}
+  | Lfunction {body; arity; params; attr; ty} ->
     let body = f body in
-    Lfunction {body; arity; params; attr}
-  | Llet (str, id, arg, body) ->
+    Lfunction {body; arity; params; attr; ty}
+  | Llet (str, id, ty, arg, body) ->
     let arg = f arg in
     let body = f body in
-    Llet (str, id, arg, body)
+    Llet (str, id, ty, arg, body)
   | Lletrec (decl, body) ->
     let body = f body in
     let decl = Ext_list.map_snd decl f in
@@ -307,7 +311,8 @@ let rec is_eta_conversion_exn params inner_args outer_args : t list =
   | _, _, _ -> raise_notrace Not_simple_form
 
 (** FIXME: more robust inlining check later, we should inline it before we add stub code*)
-let rec apply ?(ap_transformed_jsx = false) fn args (ap_info : ap_info) : t =
+let rec apply ?(ap_transformed_jsx = false)
+    ~(ap_result_type : Types.type_expr option) fn args (ap_info : ap_info) : t =
   match fn with
   | Lfunction
       {
@@ -328,7 +333,14 @@ let rec apply ?(ap_transformed_jsx = false) fn args (ap_info : ap_info) : t =
       Lprim
         {primitive = wrap; args = [Lprim {primitive_call with args; loc}]; loc}
     | exception Not_simple_form ->
-      Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx})
+      Lapply
+        {
+          ap_func = fn;
+          ap_args = args;
+          ap_info;
+          ap_transformed_jsx;
+          ap_result_type;
+        })
   | Lfunction
       {
         params;
@@ -337,7 +349,14 @@ let rec apply ?(ap_transformed_jsx = false) fn args (ap_info : ap_info) : t =
     match is_eta_conversion_exn params inner_args args with
     | args -> Lprim {primitive_call with args; loc = ap_info.ap_loc}
     | exception _ ->
-      Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx})
+      Lapply
+        {
+          ap_func = fn;
+          ap_args = args;
+          ap_info;
+          ap_transformed_jsx;
+          ap_result_type;
+        })
   | Lfunction
       {
         params;
@@ -350,17 +369,37 @@ let rec apply ?(ap_transformed_jsx = false) fn args (ap_info : ap_info) : t =
     | args ->
       Lsequence (Lprim {primitive_call with args; loc = ap_info.ap_loc}, const)
     | exception _ ->
-      Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx}
+      Lapply
+        {
+          ap_func = fn;
+          ap_args = args;
+          ap_info;
+          ap_transformed_jsx;
+          ap_result_type;
+        }
       (* | Lfunction {params;body} when Ext_list.same_length params args ->
           Ext_list.fold_right2 (fun p arg acc ->
             Llet(Strict,p,arg,acc)
           ) params args body *)
       (* TODO: more rigirous analysis on [let_kind] *))
-  | Llet (kind, id, e, (Lfunction _ as fn)) ->
-    Llet (kind, id, e, apply fn args ap_info ~ap_transformed_jsx)
+  | Llet (kind, id, ty, e, (Lfunction _ as fn)) ->
+    Llet
+      ( kind,
+        id,
+        ty,
+        e,
+        apply fn args ap_info ~ap_transformed_jsx ~ap_result_type )
   (* | Llet (kind0, id0, e0, Llet (kind,id, e, (Lfunction _ as fn))) ->
      Llet(kind0,id0,e0,Llet (kind, id, e, apply fn args loc status)) *)
-  | _ -> Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx}
+  | _ ->
+    Lapply
+      {
+        ap_func = fn;
+        ap_args = args;
+        ap_info;
+        ap_transformed_jsx;
+        ap_result_type;
+      }
 
 let rec eq_approx (l1 : t) (l2 : t) =
   match l1 with
@@ -419,7 +458,7 @@ let rec eq_approx (l1 : t) (l2 : t) =
            (fun ((k : string), v) (k2, v2) -> k = k2 && eq_approx v v2)
     | _ -> false)
   | Lfunction _
-  | Llet (_, _, _, _)
+  | Llet (_, _, _, _, _)
   | Lletrec _ | Lswitch _ | Lstaticcatch _ | Ltrywith _
   | Lfor (_, _, _, _, _)
   | Lfor_of (_, _, _)
@@ -475,10 +514,10 @@ let global_module ?(dynamic_import = false) id =
   Lglobal_module (id, dynamic_import)
 let const ct : t = Lconst ct
 
-let function_ ~attr ~arity ~params ~body : t =
-  Lfunction {arity; params; body; attr}
+let function_ ~attr ~arity ~params ~body ~ty : t =
+  Lfunction {arity; params; body; attr; ty}
 
-let let_ kind id e body : t = Llet (kind, id, e, body)
+let let_ kind id ty e body : t = Llet (kind, id, ty, e, body)
 let letrec bindings body : t = Lletrec (bindings, body)
 let while_ a b : t = Lwhile (a, b)
 let try_ body id handler : t = Ltrywith (body, id, handler)

compiler/core/lam.mli

@@ -46,13 +46,15 @@ and apply = private {
   ap_args: t list;
   ap_info: ap_info;
   ap_transformed_jsx: bool;
+  ap_result_type: Types.type_expr option;
 }
 
 and lfunction = {
   arity: int;
   params: ident list;
   body: t;
   attr: Lambda.function_attribute;
+  ty: Types.type_expr option;
 }
 
 and prim_info = private {
@@ -67,7 +69,7 @@ and t = private
   | Lconst of Lam_constant.t
   | Lapply of apply
   | Lfunction of lfunction
-  | Llet of Lam_compat.let_kind * ident * t * t
+  | Llet of Lam_compat.let_kind * ident * Types.type_expr option * t * t
   | Lletrec of (ident * t) list * t
   | Lprim of prim_info
   | Lswitch of t * lambda_switch
@@ -113,16 +115,23 @@ val global_module : ?dynamic_import:bool -> ident -> t
 
 val const : Lam_constant.t -> t
 
-val apply : ?ap_transformed_jsx:bool -> t -> t list -> ap_info -> t
+val apply :
+  ?ap_transformed_jsx:bool ->
+  ap_result_type:Types.type_expr option ->
+  t ->
+  t list ->
+  ap_info ->
+  t
 
 val function_ :
   attr:Lambda.function_attribute ->
   arity:int ->
   params:ident list ->
   body:t ->
+  ty:Types.type_expr option ->
   t
 
-val let_ : Lam_compat.let_kind -> ident -> t -> t -> t
+val let_ : Lam_compat.let_kind -> ident -> Types.type_expr option -> t -> t -> t
 
 val letrec : (ident * t) list -> t -> t
 

compiler/core/lam_analysis.ml

@@ -102,7 +102,7 @@ let rec no_side_effects (lam : Lam.t) : bool =
     (* byte swap *)
     | Parraysets | Parraysetu | Poffsetref _ | Praise | Psetfield _ ->
       false)
-  | Llet (_, _, arg, body) -> no_side_effects arg && no_side_effects body
+  | Llet (_, _, _, arg, body) -> no_side_effects arg && no_side_effects body
   | Lswitch (_, _) -> false
   | Lstringswitch (_, _, _) -> false
   | Lstaticraise _ -> false
@@ -149,7 +149,7 @@ let rec size (lam : Lam.t) =
     match lam with
     | Lvar _ -> 1
     | Lconst c -> size_constant c
-    | Llet (_, _, l1, l2) -> 1 + size l1 + size l2
+    | Llet (_, _, _, l1, l2) -> 1 + size l1 + size l2
     | Lletrec _ -> really_big ()
     | Lprim
         {

compiler/core/lam_arity_analysis.ml

@@ -39,7 +39,7 @@ let rec get_arity (meta : Lam_stats.t) (lam : Lam.t) : Lam_arity.t =
   match lam with
   | Lvar v -> arity_of_var meta v
   | Lconst _ -> Lam_arity.non_function_arity_info
-  | Llet (_, _, _, l) -> get_arity meta l
+  | Llet (_, _, _, _, l) -> get_arity meta l
   | Lprim
       {
         primitive = Pfield (_, Fld_module {name});

compiler/core/lam_beta_reduce.ml

@@ -71,7 +71,7 @@ let propagate_beta_reduce (meta : Lam_stats.t) (params : Ident.t list)
         | Lprim {primitive = Psome | Psome_not_nest; args = [v]; _} ->
           Hash_ident.replace meta.ident_tbl param (Normal_optional v)
         | _ -> ());
-        Lam_util.refine_let ~kind:Strict param arg l)
+        Lam_util.refine_let ~kind:Strict ~ty:None param arg l)
 
 let propagate_beta_reduce_with_map (meta : Lam_stats.t)
     (map : Lam_var_stats.stats Map_ident.t) params body args =
@@ -112,11 +112,11 @@ let propagate_beta_reduce_with_map (meta : Lam_stats.t)
         | Lprim {primitive = Psome | Psome_not_nest; args = [v]} ->
           Hash_ident.replace meta.ident_tbl param (Normal_optional v)
         | _ -> ());
-        Lam_util.refine_let ~kind:Strict param arg l)
+        Lam_util.refine_let ~kind:Strict ~ty:None param arg l)
 
 let no_names_beta_reduce params body args =
   match Lam_beta_reduce_util.simple_beta_reduce params body args with
   | Some x -> x
   | None ->
     Ext_list.fold_left2 params args body (fun param arg l ->
-        Lam_util.refine_let ~kind:Strict param arg l)
+        Lam_util.refine_let ~kind:Strict ~ty:None param arg l)

compiler/core/lam_beta_reduce_util.ml

@@ -113,7 +113,8 @@ let simple_beta_reduce params body args =
         | _ -> f
       in
       let result =
-        Hash_ident.fold param_hash (Lam.apply f new_args ap_info)
+        Hash_ident.fold param_hash
+          (Lam.apply ~ap_result_type:None f new_args ap_info)
           (fun _param stat acc ->
             let {lambda; used} = stat in
             if not used then Lam.seq lambda acc else acc)