Add @res.hoistedFunction support for flat JS export

Author: cknitt

compiler/core/js_cmj_format.ml

@@ -37,6 +37,8 @@ type effect_ = string option
 
 let single_na = Single Lam_arity.na
 
+let module_field_path_key segments = String.concat "$" segments
+
 type keyed_cmj_value = {
   name: string;
   arity: arity;
@@ -47,12 +49,17 @@ type keyed_cmj_values = keyed_cmj_value array
 
 type t = {
   values: keyed_cmj_values;
+  (* Full source paths that can be read through flat JS exports.  The key uses
+     [module_field_path_key], so source path [A.B.make] is stored as
+     [A$B$make].  The same string is also the emitted export name. *)
+  hoisted_values: string array;
   pure: bool;
   package_spec: Js_packages_info.t;
   case: Ext_js_file_kind.case;
 }
 
-let make ~(values : cmj_value Map_string.t) ~effect_ ~package_spec ~case : t =
+let make ~(values : cmj_value Map_string.t) ~hoisted_values ~effect_
+    ~package_spec ~case : t =
   {
     values =
       Map_string.to_sorted_array_with_f values (fun k v ->
@@ -61,6 +68,7 @@ let make ~(values : cmj_value Map_string.t) ~effect_ ~package_spec ~case : t =
             arity = v.arity;
             persistent_closed_lambda = v.persistent_closed_lambda;
           });
+    hoisted_values = Array.of_list (Set_string.elements hoisted_values);
     pure = effect_ = None;
     package_spec;
     case;
@@ -158,6 +166,10 @@ let query_by_name (cmj_table : t) name : keyed_cmj_value =
   let values = cmj_table.values in
   binary_search values name
 
+let has_hoisted_value (cmj_table : t) name =
+  Array.exists (fun hoisted_name -> Ext_string.equal hoisted_name name)
+    cmj_table.hoisted_values
+
 type path = string
 
 type cmj_load_info = {

compiler/core/js_cmj_format.mli

@@ -62,22 +62,35 @@ type keyed_cmj_value = {
 
 type t = {
   values: keyed_cmj_value array;
+  (* Full source paths that can be read through flat JS exports.  The key uses
+     [module_field_path_key], so source path [A.B.make] is stored as
+     [A$B$make].  The same string is also the emitted export name. *)
+  hoisted_values: string array;
   pure: bool;
   package_spec: Js_packages_info.t;
   case: Ext_js_file_kind.case;
 }
 
 val make :
   values:cmj_value Map_string.t ->
+  hoisted_values:Set_string.t ->
   effect_:effect_ ->
   package_spec:Js_packages_info.t ->
   case:Ext_js_file_kind.case ->
   t
 
 val query_by_name : t -> string -> keyed_cmj_value
 
+val has_hoisted_value : t -> string -> bool
+
 val single_na : arity
 
+(* Encode a full module field path as the string key used in .cmj tables.
+   This is only needed when the whole path is itself a recorded value, such as
+   a hoisted [A.B.make] export.  Ordinary nested module metadata is still read
+   from the parent field, such as [A], and then interpreted by field index. *)
+val module_field_path_key : string list -> string
+
 val from_file : string -> t
 
 val from_file_with_digest : string -> t * Digest.t

compiler/core/lam_compile.ml

@@ -232,6 +232,40 @@ type initialization = J.block
 *)
 
 let compile output_prefix =
+  (* When compiling a read from another module, a nested source path like
+     Other.A.B.make reaches this point as nested module-field reads:
+
+       Pfield "make" (Pfield "B" (Pfield "A" (Lglobal_module Other)))
+
+     Normal compilation does not look up the full path.  It only queries the
+     first field, "A", and then emits the remaining fields as JS property
+     access: Other.A.B.make.  The "A" lookup may include Submodule arity data,
+     but it does not say whether A.B.make has a separate root-level export.
+
+     Hoisted functions need that extra question.  For them, query the separate
+     hoisted-values table with the full source path A.B.make.  If present, the
+     same key is the root-level JS export name, for example Other.A$B$make.
+     The table is only a marker set; normal export metadata still lives in the
+     regular .cmj values table. *)
+  let rec extract_field_path segments primitive args =
+    match (primitive, args) with
+    | ( Lam_primitive.Pfield (_, Fld_module {name}),
+        [Lam.Lprim {primitive; args; _}] ) ->
+      extract_field_path (name :: segments) primitive args
+    | ( Lam_primitive.Pfield (_, Fld_module {name}),
+        [Lam.Lglobal_module (id, dynamic_import)] ) ->
+      Some (id, dynamic_import, name :: segments)
+    | _ -> None
+  in
+  let hoisted_external_field_name primitive args =
+    match extract_field_path [] primitive args with
+    | Some (id, dynamic_import, (_ :: _ :: _ as segments)) ->
+      let name = Js_cmj_format.module_field_path_key segments in
+      if Lam_compile_env.has_hoisted_external_id ~dynamic_import id name then
+        Some (id, dynamic_import, name)
+      else None
+    | Some (_, _, ([] | [_])) | 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 =
@@ -1637,17 +1671,47 @@ let compile output_prefix =
              fn_code args))
   and compile_prim (prim_info : Lam.prim_info)
       (lambda_cxt : Lam_compile_context.t) =
+    let compile_primitive_default primitive args loc =
+      let args_block, args_expr =
+        if args = [] then ([], [])
+        else
+          let new_cxt = {lambda_cxt with continuation = NeedValue Not_tail} in
+          Ext_list.split_map 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 args_block in
+      let exp =
+        (* TODO: all can be done in [compile_primitive] *)
+        Lam_compile_primitive.translate output_prefix loc lambda_cxt primitive
+          args_expr
+      in
+      Js_output.output_of_block_and_expression lambda_cxt.continuation args_code
+        exp
+    in
     match prim_info with
-    | {
-     primitive = Pfield (_, fld_info);
-     args = [Lglobal_module (id, dynamic_import)];
-     _;
-    } -> (
-      (* should be before Lglobal_global *)
-      match fld_info with
-      | Fld_module {name = field} ->
-        compile_external_field ~dynamic_import lambda_cxt id field
-      | _ -> assert false)
+    | {primitive = Pfield (_, Fld_module _); _} -> (
+      match hoisted_external_field_name prim_info.primitive prim_info.args with
+      | Some (id, dynamic_import, hoisted_name) ->
+        Js_output.output_of_expression lambda_cxt.continuation
+          ~no_effects:no_effects_const
+          (E.ml_var_dot ~dynamic_import id hoisted_name)
+      | None -> (
+        match prim_info with
+        | {
+         primitive = Pfield (_, fld_info);
+         args = [Lglobal_module (id, dynamic_import)];
+         _;
+        } -> (
+          (* should be before Lglobal_global *)
+          match fld_info with
+          | Fld_module {name = field} ->
+            compile_external_field ~dynamic_import lambda_cxt id field
+          | _ -> assert false)
+        | _ ->
+          compile_primitive_default prim_info.primitive prim_info.args
+            prim_info.loc))
     | {primitive = Praise; args = [e]; _} -> (
       match
         compile_lambda {lambda_cxt with continuation = NeedValue Not_tail} e
@@ -1816,24 +1880,7 @@ let compile output_prefix =
         Location.raise_errorf ~loc
           "Invalid argument: unsupported argument to dynamic import. If you \
            believe this should be supported, please open an issue.")
-    | {primitive; args; loc} ->
-      let args_block, args_expr =
-        if args = [] then ([], [])
-        else
-          let new_cxt = {lambda_cxt with continuation = NeedValue Not_tail} in
-          Ext_list.split_map 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 args_block in
-      let exp =
-        (* TODO: all can be done in [compile_primitive] *)
-        Lam_compile_primitive.translate output_prefix loc lambda_cxt primitive
-          args_expr
-      in
-      Js_output.output_of_block_and_expression lambda_cxt.continuation args_code
-        exp
+    | {primitive; args; loc} -> compile_primitive_default primitive args loc
   and compile_lambda (lambda_cxt : Lam_compile_context.t) (cur_lam : Lam.t) :
       Js_output.t =
     match cur_lam with

compiler/core/lam_compile_env.ml

@@ -84,20 +84,26 @@ let add_js_module ?import_attributes
     id
   | Some old_key -> old_key.id
 
+let cmj_table_of_module_id ~dynamic_import (module_id : Ident.t) =
+  let oid = Lam_module_ident.of_ml ~dynamic_import module_id in
+  match Lam_module_ident.Hash.find_opt cached_tbl oid with
+  | None ->
+    let cmj_load_info = !Js_cmj_load.load_unit module_id.name in
+    oid +> Ml cmj_load_info;
+    cmj_load_info.cmj_table
+  | Some (Ml {cmj_table}) -> cmj_table
+  | Some External -> assert false
+
 let query_external_id_info ?(dynamic_import = false) (module_id : Ident.t)
     (name : string) : ident_info =
-  let oid = Lam_module_ident.of_ml ~dynamic_import module_id in
-  let cmj_table =
-    match Lam_module_ident.Hash.find_opt cached_tbl oid with
-    | None ->
-      let cmj_load_info = !Js_cmj_load.load_unit module_id.name in
-      oid +> Ml cmj_load_info;
-      cmj_load_info.cmj_table
-    | Some (Ml {cmj_table}) -> cmj_table
-    | Some External -> assert false
-  in
+  let cmj_table = cmj_table_of_module_id ~dynamic_import module_id in
   Js_cmj_format.query_by_name cmj_table name
 
+let has_hoisted_external_id ?(dynamic_import = false) (module_id : Ident.t)
+    (name : string) : bool =
+  let cmj_table = cmj_table_of_module_id ~dynamic_import module_id in
+  Js_cmj_format.has_hoisted_value cmj_table name
+
 let get_package_path_from_cmj (id : Lam_module_ident.t) :
     string * Js_packages_info.t * Ext_js_file_kind.case =
   let cmj_load_info =

compiler/core/lam_compile_env.mli

@@ -71,6 +71,8 @@ val query_external_id_info :
    will raise if not found
 *)
 
+val has_hoisted_external_id : ?dynamic_import:bool -> Ident.t -> string -> bool
+
 val is_pure_module : Lam_module_ident.t -> bool
 
 val get_package_path_from_cmj :

compiler/core/lam_compile_main.ml

@@ -115,6 +115,78 @@ let no_side_effects (rest : Lam_group.t list) : string option =
           Some ""
         else None (* TODO :*))
 
+(* Materialize JS-hoisted values as root-level aliases and exports.  The source
+   value still lives at its normal module path, but downstream tools can import
+   the flat name directly when the .cmj metadata marks it as hoisted. *)
+let js_hoisted_aliases (export_idents : Set_ident.t) (groups : Lam_group.t list)
+    =
+  let group_map =
+    Ext_list.fold_left groups Map_ident.empty (fun acc group ->
+        match group with
+        | Single (_, id, lam) -> Map_ident.add acc id lam
+        | Recursive bindings ->
+          Ext_list.fold_left bindings acc (fun acc (id, lam) ->
+              Map_ident.add acc id lam)
+        | Nop _ -> acc)
+  in
+  let rec access loc base fields =
+    match fields with
+    | [] -> base
+    | (pos, name) :: fields ->
+      access loc
+        (Lam.prim
+           ~primitive:(Lam_primitive.Pfield (pos, Lam_compat.Fld_module {name}))
+           ~args:[base] loc)
+        fields
+  in
+  let resolve = function
+    | Lam.Lvar id as lam -> (
+      match Map_ident.find_opt group_map id with
+      | Some resolved -> resolved
+      | None -> lam)
+    | lam -> lam
+  in
+  let rec fold_module_fields fields args index acc f =
+    match (fields, args) with
+    | [], [] -> acc
+    | field :: fields, arg :: args ->
+      fold_module_fields fields args (index + 1) (f index field arg acc) f
+    | _, _ -> invalid_arg "fold_module_fields"
+  in
+  let rec scan top_id path acc lam =
+    match resolve lam with
+    | Lam.Lprim
+        {
+          primitive =
+            Lam_primitive.Pmakeblock (_, Blk_module fields, _);
+          args;
+          loc;
+        } ->
+      fold_module_fields fields args 0 acc (fun pos field arg acc ->
+          let path = path @ [(pos, field)] in
+          let acc = scan top_id path acc arg in
+          match arg with
+          | Lam.Lvar id when Ident.js_hoisted id ->
+            let name =
+              Js_cmj_format.module_field_path_key
+                (top_id.Ident.name
+                :: Ext_list.map path (fun (_pos, name) -> name))
+            in
+            let alias_id = Ident.create name in
+            Ident.make_js_hoisted alias_id;
+            let alias =
+              access loc (Lam.var top_id) path
+            in
+            (Lam_group.Single (Alias, alias_id, alias), alias_id, alias) :: acc
+          | _ -> acc)
+    | _ -> acc
+  in
+  Ext_list.fold_left groups [] (fun acc group ->
+      match group with
+      | Single (_, id, lam) when Set_ident.mem export_idents id ->
+        scan id [] acc lam
+      | Single _ | Recursive _ | Nop _ -> acc)
+
 
 let _d  = fun  s lam -> 
 #ifndef RELEASE
@@ -227,6 +299,28 @@ let () =
 in
 #endif  
 let maybe_pure = no_side_effects groups in
+(* Add the generated alias groups before JS lowering so regular export
+   printing, tree shaking, and .cmj metadata all see the flat runtime value. *)
+let hoisted_aliases = js_hoisted_aliases meta.export_idents groups in
+let hoisted_groups, hoisted_exports, hoisted_export_map =
+  Ext_list.fold_left hoisted_aliases ([], [], Map_ident.empty)
+    (fun (groups, exports, export_map) (group, id, lam) ->
+      (group :: groups, id :: exports, Map_ident.add export_map id lam))
+in
+let groups = groups @ List.rev hoisted_groups in
+let meta =
+  {
+    meta with
+    exports = meta.exports @ List.rev hoisted_exports;
+    export_idents =
+      Ext_list.fold_left hoisted_exports meta.export_idents (fun acc id ->
+          Set_ident.add acc id);
+  }
+in
+let export_map =
+  Map_ident.fold hoisted_export_map coerced_input.export_map (fun id lam acc ->
+      Map_ident.add acc id lam)
+in
 #ifndef RELEASE
 let () = Ext_log.dwarn ~__POS__ "\n@[[TIME:]Pre-compile: %f@]@."  (Sys.time () *. 1000.) in      
 #endif  
@@ -277,18 +371,18 @@ js
           ) 
     in
     Warnings.check_fatal();
-    let effect_ = 
+    let effect_ =
       Lam_stats_export.get_dependent_module_effect
-        maybe_pure external_module_ids in 
-    let v : Js_cmj_format.t = 
-      Lam_stats_export.export_to_cmj 
-        meta  
-        effect_ 
-        coerced_input.export_map          
+        maybe_pure external_module_ids in
+    let v : Js_cmj_format.t =
+      Lam_stats_export.export_to_cmj
+        meta
+        effect_
+        export_map
         (if Ext_char.is_lower_case (Filename.basename output_prefix).[0] then Little else Upper)
     in
     (if not !Clflags.dont_write_files then
-       Js_cmj_format.to_file 
+       Js_cmj_format.to_file
          ~check_exists:(not !Js_config.force_cmj)
          (output_prefix ^ Literals.suffix_cmj) v);
     {J.program = program ; side_effect = effect_ ; modules = external_module_ids }