Rewrite nested JSX component paths to direct hoisted exports

Author: fhammerschmidt

compiler/core/js_of_lam_block.ml

@@ -43,7 +43,7 @@ let field (field_info : Lam_compat.field_dbg_info) e (i : int32) =
   | Fld_cons -> E.cons_access e i
   | Fld_record_inline {name} -> E.inline_record_access e name i
   | Fld_record {name} -> E.record_access e name i
-  | Fld_module {name} -> E.module_access e name i
+  | Fld_module {name; jsx_component = _} -> E.module_access e name i
 
 let field_by_exp e i = E.array_index e i
 

compiler/core/lam.ml

@@ -555,7 +555,8 @@ let prim ~primitive:(prim : Lam_primitive.t) ~args loc : t =
         | ( f :: fields,
             Lprim
               {
-                primitive = Pfield (pos, Fld_module {name = f1});
+                primitive =
+                  Pfield (pos, Fld_module {name = f1; jsx_component = false});
                 args = [(Lglobal_module (v1, _) | Lvar v1)];
               }
             :: args ) ->
@@ -566,7 +567,8 @@ let prim ~primitive:(prim : Lam_primitive.t) ~args loc : t =
       | ( field1 :: rest,
           Lprim
             {
-              primitive = Pfield (pos, Fld_module {name = f1});
+              primitive =
+                Pfield (pos, Fld_module {name = f1; jsx_component = false});
               args = [((Lglobal_module (v1, _) | Lvar v1) as lam)];
             }
           :: args1 ) ->

compiler/core/lam_analysis.ml

@@ -123,7 +123,12 @@ let rec no_side_effects (lam : Lam.t) : bool =
   (* | Lsend _ -> false  *)
   | Lapply
       {
-        ap_func = Lprim {primitive = Pfield (_, Fld_module {name = "from_fun"})};
+        ap_func =
+          Lprim
+            {
+              primitive =
+                Pfield (_, Fld_module {name = "from_fun"; jsx_component = _});
+            };
         ap_args = [arg];
       } ->
     no_side_effects arg

compiler/core/lam_arity_analysis.ml

@@ -42,7 +42,7 @@ let rec get_arity (meta : Lam_stats.t) (lam : Lam.t) : Lam_arity.t =
   | Llet (_, _, _, l) -> get_arity meta l
   | Lprim
       {
-        primitive = Pfield (_, Fld_module {name});
+        primitive = Pfield (_, Fld_module {name; jsx_component = _});
         args = [Lglobal_module (id, dynamic_import)];
         _;
       } -> (
@@ -58,7 +58,7 @@ let rec get_arity (meta : Lam_stats.t) (lam : Lam.t) : Lam_arity.t =
           [
             Lprim
               {
-                primitive = Pfield (_, Fld_module {name});
+                primitive = Pfield (_, Fld_module {name; jsx_component = _});
                 args = [Lglobal_module (id, dynamic_import)];
               };
           ];

compiler/core/lam_compat.ml

@@ -64,7 +64,7 @@ type let_kind = Lambda.let_kind = Strict | Alias | StrictOpt | Variable
 
 type field_dbg_info = Lambda.field_dbg_info =
   | Fld_record of {name: string; mutable_flag: Asttypes.mutable_flag}
-  | Fld_module of {name: string}
+  | Fld_module of {name: string; jsx_component: bool}
   | Fld_record_inline of {name: string}
   | Fld_record_extension of {name: string}
   | Fld_tuple

compiler/core/lam_compat.mli

@@ -28,7 +28,7 @@ type let_kind = Lambda.let_kind = Strict | Alias | StrictOpt | Variable
 
 type field_dbg_info = Lambda.field_dbg_info =
   | Fld_record of {name: string; mutable_flag: Asttypes.mutable_flag}
-  | Fld_module of {name: string}
+  | Fld_module of {name: string; jsx_component: bool}
   | Fld_record_inline of {name: string}
   | Fld_record_extension of {name: string}
   | Fld_tuple

compiler/core/lam_compile.ml

@@ -232,6 +232,64 @@ type initialization = J.block
 *)
 
 let compile output_prefix =
+  let root_module_name (id : Ident.t) =
+    match String.index_opt id.name '$' with
+    | Some index -> String.sub id.name 0 index
+    | None -> id.name
+  in
+  let rec extract_nested_external_component_segments segments
+      ((lam : Lam.t), (make_dynamic_import : bool option ref)) :
+      (Ident.t * bool * string list) option =
+    match lam with
+    | Lprim
+        {
+          primitive = Pfield (_, Fld_module {name; jsx_component = _});
+          args = [arg];
+          _;
+        } ->
+      extract_nested_external_component_segments (name :: segments)
+        (arg, make_dynamic_import)
+    | Lvar id ->
+      make_dynamic_import := Some false;
+      Some (id, false, List.rev segments)
+    | Lglobal_module (id, dynamic_import) ->
+      make_dynamic_import := Some dynamic_import;
+      Some (id, dynamic_import, List.rev segments)
+    | _ -> None
+  in
+  let extract_nested_external_component_field (lam : Lam.t) :
+      (Ident.t * bool * string) option =
+    match lam with
+    | Lprim
+        {
+          primitive = Pfield (_, Fld_module {name = "make"; jsx_component = _});
+          args = [arg];
+          _;
+        } -> (
+      let dynamic_import = ref None in
+      match
+        extract_nested_external_component_segments [] (arg, dynamic_import)
+      with
+      | Some (id, dynamic_import, segments) -> (
+        let segments =
+          match segments with
+          | head :: rest
+            when head = id.name
+                 || head = root_module_name id
+                 || Ext_string.starts_with head (root_module_name id ^ "$") ->
+            rest
+          | _ -> segments
+        in
+        match segments with
+        | [] -> None
+        | _ ->
+          Some
+            ( id,
+              dynamic_import,
+              String.concat "$" (root_module_name id :: segments) ))
+      | None -> None)
+    | _ -> None
+  in
   let rec compile_external_field (* Like [List.empty]*)
       ?(dynamic_import = false) (lamba_cxt : Lam_compile_context.t)
       (id : Ident.t) name : Js_output.t =
@@ -300,6 +358,17 @@ let compile output_prefix =
                 (Ext_list.append block args_code, b :: args)
               | _ -> assert false)
       in
+      let args =
+        if appinfo.ap_transformed_jsx then
+          match (appinfo.ap_args, args) with
+          | jsx_tag :: _, _ :: rest_args -> (
+            match extract_nested_external_component_field jsx_tag with
+            | Some (id, dynamic_import, hidden_name) ->
+              E.ml_var_dot ~dynamic_import id hidden_name :: rest_args
+            | None -> args)
+          | _ -> args
+        else args
+      in
 
       let fn = E.ml_var_dot ~dynamic_import module_id ident_info.name in
       let expression =
@@ -1524,6 +1593,17 @@ let compile output_prefix =
               (Ext_list.append block args_code, b :: fn_code)
             | {value = None} -> assert false)
       in
+      let args =
+        if appinfo.ap_transformed_jsx then
+          match (appinfo.ap_args, args) with
+          | jsx_tag :: _, _ :: rest_args -> (
+            match extract_nested_external_component_field jsx_tag with
+            | Some (id, dynamic_import, hidden_name) ->
+              E.ml_var_dot ~dynamic_import id hidden_name :: rest_args
+            | None -> args)
+          | _ -> args
+        else args
+      in
       match (ap_func, lambda_cxt.continuation) with
       | ( Lvar fn_id,
           ( EffectCall (Maybe_tail_is_return (Tail_with_name {label = Some ret}))
@@ -1583,6 +1663,48 @@ let compile output_prefix =
   and compile_prim (prim_info : Lam.prim_info)
       (lambda_cxt : Lam_compile_context.t) =
     match prim_info with
+    | {
+     primitive =
+       Pjs_call
+         {
+           prim_name = "jsx" | "jsxs" | "jsxKeyed" | "jsxsKeyed";
+           transformed_jsx = true;
+           _;
+         };
+     args = jsx_tag :: rest_args;
+     loc;
+    } ->
+      let new_cxt = {lambda_cxt with continuation = NeedValue Not_tail} in
+      let tag_block, tag_expr =
+        match extract_nested_external_component_field jsx_tag with
+        | Some (id, dynamic_import, hidden_name) -> (
+          match
+            Lam_compile_env.query_external_id_info ~dynamic_import id
+              (hidden_name ^ "$jsx")
+          with
+          | exception Not_found -> (
+            match compile_lambda new_cxt jsx_tag with
+            | {block; value = Some b} -> (block, b)
+            | {value = None} -> assert false)
+          | _ -> ([], E.ml_var_dot ~dynamic_import id hidden_name))
+        | None -> (
+          match compile_lambda new_cxt jsx_tag with
+          | {block; value = Some b} -> (block, b)
+          | {value = None} -> assert false)
+      in
+      let rest_blocks, rest_exprs =
+        Ext_list.split_map rest_args (fun x ->
+            match compile_lambda new_cxt x with
+            | {block; value = Some b} -> (block, b)
+            | {value = None} -> assert false)
+      in
+      let args_code : J.block = List.concat (tag_block :: rest_blocks) in
+      let exp =
+        Lam_compile_primitive.translate output_prefix loc lambda_cxt
+          prim_info.primitive (tag_expr :: rest_exprs)
+      in
+      Js_output.output_of_block_and_expression lambda_cxt.continuation args_code
+        exp
     | {
      primitive = Pfield (_, fld_info);
      args = [Lglobal_module (id, dynamic_import)];

compiler/core/lam_pass_remove_alias.ml

@@ -133,7 +133,8 @@ let simplify_alias (meta : Lam_stats.t) (lam : Lam.t) : Lam.t =
           ap_func =
             Lprim
               {
-                primitive = Pfield (_, Fld_module {name = fld_name});
+                primitive =
+                  Pfield (_, Fld_module {name = fld_name; jsx_component = _});
                 args = [Lglobal_module (ident, dynamic_import)];
                 _;
               } as l1;

compiler/core/lam_print.ml

@@ -323,7 +323,7 @@ let lambda ppf v =
       fprintf ppf ")@ %a)@]" lam body
     | Lprim
         {
-          primitive = Pfield (n, Fld_module {name = s});
+          primitive = Pfield (n, Fld_module {name = s; jsx_component = _});
           args = [Lglobal_module (id, dynamic_import)];
           _;
         } ->

compiler/ml/lambda.ml

@@ -113,7 +113,7 @@ let ref_tag_info : tag_info =
 
 type field_dbg_info =
   | Fld_record of {name: string; mutable_flag: Asttypes.mutable_flag}
-  | Fld_module of {name: string}
+  | Fld_module of {name: string; jsx_component: bool}
   | Fld_record_inline of {name: string}
   | Fld_record_extension of {name: string}
   | Fld_tuple
@@ -134,6 +134,8 @@ let fld_record_extension (lbl : label) =
   Fld_record_extension
     {name = Ext_list.find_def lbl.lbl_attributes find_name lbl.lbl_name}
 
+let fld_module ~name ~jsx_component = Fld_module {name; jsx_component}
+
 let ref_field_info : field_dbg_info =
   Fld_record {name = "contents"; mutable_flag = Mutable}
 
@@ -620,11 +622,28 @@ let rec transl_normal_path = function
     else Lvar id
   | Pdot (p, s, pos) ->
     Lprim
-      ( Pfield (pos, Fld_module {name = s}),
+      ( Pfield (pos, Fld_module {name = s; jsx_component = false}),
         [transl_normal_path p],
         Location.none )
   | Papply _ -> assert false
 
+let transl_jsx_path path =
+  let rec aux ~is_final = function
+    | Path.Pident id ->
+      if Ident.global id then Lprim (Pgetglobal id, [], Location.none)
+      else Lvar id
+    | Pdot (p, s, pos) ->
+      Lprim
+        ( Pfield
+            ( pos,
+              Fld_module
+                {name = s; jsx_component = is_final && String.equal s "make"} ),
+          [aux ~is_final:false p],
+          Location.none )
+    | Papply _ -> assert false
+  in
+  aux ~is_final:true path
+
 (* Translation of identifiers *)
 
 let transl_module_path ?(loc = Location.none) env path =
@@ -633,6 +652,9 @@ let transl_module_path ?(loc = Location.none) env path =
 let transl_value_path ?(loc = Location.none) env path =
   transl_normal_path (Env.normalize_path_prefix (Some loc) env path)
 
+let transl_jsx_value_path ?(loc = Location.none) env path =
+  transl_jsx_path (Env.normalize_path_prefix (Some loc) env path)
+
 let transl_extension_path = transl_value_path
 
 (* Apply a substitution to a lambda-term.