TS backend phase 0 - Typed lambda

This is a split from #8118 and explains the most basic part of the entire work.

Therefore, this should provide sufficient explanation to the LLM agents
and should not conflict with other backend modifications (e.g., sourcemap) yet.

- Before: Gentype; a separated pipeline that parses and reanalyzes serialized type information.
- After: Attach the full type information to the Lambda IR and pass it directly to the codegen backend.

Although the compiler's intermediate layer still does not use type information,
passing through it simplifying the type-related code generation as
it can directly access already computed type information without any extra parsing.

Next phases:

- 1. Implementing dts file emission:
  Provide it as an additional option. Stabilizes the dts format independently of the existing JS IR and deprecates Gentype.
- 2. Rewriting codegen backend entirely:
  New, fully-typed, well-structured codegen IR that unifies all the fragmented codegen logic.
  It should covers JS/TS/d.ts in a single path (Need research. It may not be possible or complicated)
  - Implementation only: JavaScript output
  - Type only: d.ts output
  - Implementation + Type: TypeScript output

Author: cometkim

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,8 +143,9 @@ 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

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_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)

compiler/core/lam_bounded_vars.ml

@@ -75,11 +75,11 @@ let rewrite (map : _ Hash_ident.t) (lam : Lam.t) : Lam.t =
   and aux (lam : Lam.t) : Lam.t =
     match lam with
     | Lvar v -> Hash_ident.find_default map v lam
-    | Llet (str, v, l1, l2) ->
+    | Llet (str, v, ty, l1, l2) ->
       let v = rebind v in
       let l1 = aux l1 in
       let l2 = aux l2 in
-      Lam.let_ str v l1 l2
+      Lam.let_ str v ty l1 l2
     | Lletrec (bindings, body) ->
       (*order matters see GPR #405*)
       let vars = Ext_list.map bindings (fun (k, _) -> rebind k) in
@@ -88,10 +88,10 @@ let rewrite (map : _ Hash_ident.t) (lam : Lam.t) : Lam.t =
       in
       let body = aux body in
       Lam.letrec bindings body
-    | Lfunction {arity; params; body; attr} ->
+    | Lfunction {arity; params; body; attr; ty} ->
       let params = Ext_list.map params rebind in
       let body = aux body in
-      Lam.function_ ~arity ~params ~body ~attr
+      Lam.function_ ~arity ~params ~body ~attr ~ty
     | Lstaticcatch (l1, (i, xs), l2) ->
       let l1 = aux l1 in
       let xs = Ext_list.map xs rebind in
@@ -118,10 +118,10 @@ let rewrite (map : _ Hash_ident.t) (lam : Lam.t) : Lam.t =
       (* here it makes sure that global vars are not rebound *)
       Lam.prim ~primitive ~args:(Ext_list.map args aux) loc
     | Lglobal_module _ -> lam
-    | Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx} ->
+    | Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx; ap_result_type} ->
       let fn = aux ap_func in
       let args = Ext_list.map ap_args aux in
-      Lam.apply ~ap_transformed_jsx fn args ap_info
+      Lam.apply ~ap_transformed_jsx ~ap_result_type fn args ap_info
     | Lswitch
         ( l,
           {