start experimenting with extracting and rewriting embeds

Author: zth

jscomp/bsc/rescript_compiler_main.ml

@@ -264,6 +264,9 @@ let buckle_script_flags : (string * Bsc_args.spec * string) array =
     "-I", string_list_add  Clflags.include_dirs ,
     "*internal* <dir>  Add <dir> to the list of include directories" ;
 
+    "-embed", string_list_add  Js_config.embeds ,
+    "TODO: Explain." ;
+
     "-w", string_call (Warnings.parse_options false),
     "<list>  Enable or disable warnings according to <list>:\n\
      +<spec>   enable warnings in <spec>\n\

jscomp/common/js_config.ml

@@ -31,6 +31,8 @@ type jsx_mode = Classic | Automatic
 let no_version_header = ref false
 
 let directives = ref []
+
+let embeds = ref []
 let cross_module_inline = ref false
 let diagnose = ref false
 

jscomp/common/js_config.mli

@@ -32,6 +32,9 @@ type jsx_mode = Classic | Automatic
 val no_version_header : bool ref
 (** set/get header *)
 
+val embeds : string list ref
+(** embeds *)
+
 val directives : string list ref
 (** directives printed verbatims just after the version header *)
 

jscomp/core/js_embeds.ml

@@ -0,0 +1,39 @@
+let write_embeds ~extension_points ~output ast =
+  let content = ref [] in
+  let append item = content := item :: !content in
+  let extension (iterator : Ast_iterator.iterator) (ext : Parsetree.extension) =
+    (match ext with
+    | ( {txt},
+        PStr
+          [
+            {
+              pstr_desc =
+                Pstr_eval
+                  ( {
+                      pexp_loc;
+                      pexp_desc = Pexp_constant (Pconst_string (contents, _));
+                    },
+                    _ );
+            };
+          ] )
+      when extension_points |> List.mem txt ->
+      append (pexp_loc, txt, contents)
+    | _ -> ());
+    Ast_iterator.default_iterator.extension iterator ext
+  in
+  let iterator = {Ast_iterator.default_iterator with extension} in
+  iterator.structure iterator ast;
+  match !content with
+  | [] -> ()
+  | content ->
+    let text =
+      content
+      |> List.map (fun (loc, extensionName, contents) ->
+             Printf.sprintf "<<- item begin ->>\n%s\n%s\n%i:%i-%i:%i"
+               extensionName contents loc.Location.loc_start.pos_lnum
+               loc.loc_start.pos_cnum loc.loc_end.pos_lnum loc.loc_end.pos_cnum)
+      |> List.rev |> String.concat "\n\n"
+    in
+    let oc = open_out_bin output in
+    output_string oc text;
+    close_out oc

jscomp/core/js_implementation.ml

@@ -127,6 +127,15 @@ let no_export (rest : Parsetree.structure) : Parsetree.structure =
         ]
   | _ -> rest
 
+let write_embeds outputprefix (ast : Parsetree.structure) =
+  if !Clflags.only_parse = false && !Js_config.binary_ast then (
+    match !Js_config.embeds with
+    | [] -> ()
+    | embeds -> Js_embeds.write_embeds ~extension_points:embeds 
+      ~output:(outputprefix ^ Literals.suffix_embeds) 
+      ast);
+    ast
+
 let after_parsing_impl ppf outputprefix (ast : Parsetree.structure) =
   if !Clflags.only_parse = false then (
     Js_config.all_module_aliases :=
@@ -180,6 +189,7 @@ let implementation ~parser ppf ?outputprefix fname =
   in
   Res_compmisc.init_path ();
   parser fname
+  |> write_embeds outputprefix
   |> Cmd_ppx_apply.apply_rewriters ~restore:false ~tool_name:Js_config.tool_name
        Ml
   |> Ppx_entry.rewrite_implementation

jscomp/ext/literals.ml

@@ -115,6 +115,8 @@ let suffix_cmti = ".cmti"
 
 let suffix_ast = ".ast"
 
+let suffix_embeds = ".embeds"
+
 let suffix_iast = ".iast"
 
 let suffix_d = ".d"

jscomp/frontend/bs_builtin_ppx.ml

@@ -377,6 +377,11 @@ let signature_item_mapper (self : mapper) (sigi : Parsetree.signature_item) :
 
 let structure_item_mapper (self : mapper) (str : Parsetree.structure_item) :
     Parsetree.structure_item =
+  let str =
+    match !Js_config.embeds with
+    | [] -> str
+    | _ -> Bs_embed_lang.structure_item str
+  in
   match str.pstr_desc with
   | Pstr_type (rf, tdcls) (* [ {ptype_attributes} as tdcl ] *) ->
     Ast_tdcls.handle_tdcls_in_stru self str rf tdcls

jscomp/frontend/bs_embed_lang.ml

@@ -0,0 +1,142 @@
+let should_transform name = !Js_config.embeds |> List.mem name
+
+let extract_extension str =
+  match String.split_on_char '.' str with
+  | ["generated"; tag] -> Some (tag, None)
+  | ["generated"; tag; fn_name] -> Some (tag, Some fn_name)
+  | [tag] -> Some (tag, None)
+  | [tag; fn_name] -> Some (tag, Some fn_name)
+  | _ -> None
+
+let transformed_count = Hashtbl.create 10
+
+let increment_transformed_count (ext_name : string) =
+  match Hashtbl.find_opt transformed_count ext_name with
+  | None -> Hashtbl.add transformed_count ext_name 1
+  | Some count -> Hashtbl.replace transformed_count ext_name (count + 1)
+
+let get_transformed_count ext_name =
+  match Hashtbl.find_opt transformed_count ext_name with
+  | None -> 0
+  | Some count -> count
+
+type transformMode = LetBinding | ModuleBinding
+
+let make_lident ?fn_name ~extension_name ~transform_mode filename =
+  Longident.parse
+    (Printf.sprintf "%s__%s.M%i%s"
+       (if String.ends_with filename ~suffix:".res" then
+          Filename.(chop_suffix (basename filename) ".res")
+        else Filename.(chop_suffix (basename filename) ".resi"))
+       extension_name
+       (get_transformed_count extension_name)
+       (match (transform_mode, fn_name) with
+       | LetBinding, Some fn_name -> "." ^ fn_name
+       | LetBinding, None -> ".default"
+       | ModuleBinding, _ -> ""))
+
+let transform_expr expr =
+  match expr.Parsetree.pexp_desc with
+  | Pexp_extension
+      ( {txt = ext_name},
+        PStr
+          [
+            {
+              pstr_desc =
+                Pstr_eval ({pexp_desc = Pexp_constant (Pconst_string (_, _))}, _);
+            };
+          ] )
+    when should_transform ext_name -> (
+    match extract_extension ext_name with
+    | None -> expr
+    | Some (extension_name, fn_name) ->
+      increment_transformed_count extension_name;
+      let loc = expr.pexp_loc in
+      let filename = loc.loc_start.pos_fname in
+      let lid =
+        make_lident ?fn_name ~extension_name ~transform_mode:LetBinding filename
+      in
+      Ast_helper.Exp.ident ~loc {txt = lid; loc})
+  | _ -> expr
+
+let structure_item structure_item =
+  match structure_item.Parsetree.pstr_desc with
+  | Pstr_value
+      ( recFlag,
+        [
+          ({
+             pvb_expr =
+               {pexp_desc = Pexp_extension ({txt = ext_name}, _)} as expr;
+           } as valueBinding);
+        ] )
+    when should_transform ext_name -> (
+    match extract_extension ext_name with
+    | None -> structure_item
+    | Some _ ->
+      {
+        structure_item with
+        pstr_desc =
+          Pstr_value
+            (recFlag, [{valueBinding with pvb_expr = transform_expr expr}]);
+      })
+  | Pstr_include
+      ({
+         pincl_mod =
+           {pmod_desc = Pmod_extension ({txt = ext_name; loc}, _)} as pmod;
+       } as pincl)
+    when ext_name |> should_transform -> (
+    match extract_extension ext_name with
+    | None -> structure_item
+    | Some (extension_name, _fn_name) ->
+      increment_transformed_count extension_name;
+      {
+        structure_item with
+        pstr_desc =
+          Pstr_include
+            {
+              pincl with
+              pincl_mod =
+                {
+                  pmod with
+                  pmod_desc =
+                    Pmod_ident
+                      {
+                        txt =
+                          make_lident loc.loc_start.pos_fname ~extension_name
+                            ~transform_mode:ModuleBinding;
+                        loc;
+                      };
+                };
+            };
+      })
+  | Pstr_module
+      ({
+         pmb_expr =
+           {pmod_desc = Pmod_extension ({txt = ext_name; loc}, _)} as pmod;
+       } as pmb)
+    when ext_name |> should_transform -> (
+    match extract_extension ext_name with
+    | None -> structure_item
+    | Some (extension_name, _fn_name) ->
+      increment_transformed_count extension_name;
+      {
+        structure_item with
+        pstr_desc =
+          Pstr_module
+            {
+              pmb with
+              pmb_expr =
+                {
+                  pmod with
+                  pmod_desc =
+                    Pmod_ident
+                      {
+                        txt =
+                          make_lident loc.loc_start.pos_fname ~extension_name
+                            ~transform_mode:ModuleBinding;
+                        loc;
+                      };
+                };
+            };
+      })
+  | _ -> structure_item