Add support for better pretty printing JS output using flow pretty printer

api/dune

@@ -7,6 +7,7 @@
     (package js_of_ocaml-lwt)
     (package js_of_ocaml-tyxml)
     (package js_of_ocaml-toplevel)
+    (package js_of_ocaml-flow_pretty_printer)
     (package js_of_ocaml-ocamlbuild))
   (action (run
     %{dep:../tools/docgen/docgen.exe}

api/index

@@ -7,6 +7,7 @@ Graphics_js
 Deriving_Json
 Js_of_ocaml_tyxml
 Js_of_ocaml_toplevel
+Js_of_ocaml_flow_pretty_printer
 }
 
 {1 Indexes}

buildTest.sh

@@ -139,3 +139,7 @@ export OCAMLRUNPARAM=b && time refmterr js_of_ocaml --noinline --keep-unit-names
 # Tests recursive functions
 echo "./rehack_tests/output/recursion/Recursion.cma.js"
 export OCAMLRUNPARAM=b && time refmterr js_of_ocaml --enable vardecl --keep-unit-names --enable excwrap --enable wrapped-exceptions --disable simplify_ifdecl --disable shortvar --pretty --custom-header "file:./rehack_tests/templates/common-js-module-header.js" --backend js "${1}"/default/rehack_tests/recursion/Recursion.cma -o ./rehack_tests/output/recursion/Recursion.cma.js/
+
+# Tests mutually recursive functions
+echo "./rehack_tests/output/flow_pretty_printing/MutualRecursion.cma.js"
+export OCAMLRUNPARAM=b && time refmterr js_of_ocaml --enable vardecl --keep-unit-names --enable excwrap --enable wrapped-exceptions --disable simplify_ifdecl --disable shortvar --pretty --flow-pretty-js --custom-header "file:./rehack_tests/templates/common-js-module-header.js" --backend js "${1}"/default/rehack_tests/flow_pretty_printing/MutualRecursion.cma -o ./rehack_tests/output/flow_pretty_printing/MutualRecursion.cma.js/

compiler/commonArg.ml

@@ -67,6 +67,11 @@ let prettiestJs =
   let doc = "Whether or not the pretty printing of Js should be extra pretty." in
   Arg.(value & flag & info ["prettiest-js"] ~doc)
 
+let flowPrettyJs =
+  (*TODO: change docs*)
+  let doc = "Whether or not the pretty printing of Js should be extra pretty." in
+  Arg.(value & flag & info ["flow-pretty-js"] ~doc)
+
 let debuginfo =
   let doc = "Output debug information." in
   Arg.(value & flag & info ["debuginfo"; "debug-info"] ~doc)
@@ -127,6 +132,7 @@ let t =
            disable
            pretty
            prettiestJs
+           flowPrettyJs
            debuginfo
            noinline
            quiet
@@ -138,6 +144,7 @@ let t =
            ->
         let enable = if pretty then "pretty" :: enable else enable in
         let enable = if prettiestJs then "prettiest-js" :: enable else enable in
+        let enable = if flowPrettyJs then "flow-pretty-js" :: enable else enable in
         let enable = if debuginfo then "debuginfo" :: enable else enable in
         let disable = if noinline then "inline" :: disable else disable in
         let disable_if_pretty name disable =
@@ -158,6 +165,7 @@ let t =
     $ disable
     $ pretty
     $ prettiestJs
+    $ flowPrettyJs
     $ debuginfo
     $ noinline
     $ is_quiet

compiler/js_of_ocaml.ml

@@ -158,6 +158,13 @@ let gen_file file f =
   in
   try
     let ch = open_out_bin (Fp.toString f_tmp) in
+    let ch = Custom_channel.ActualChannel.of_out_channel ch in
+    let onBeforeWrite = match Config.Flag.flowPrettyJs () with
+      | true -> Some (Backend.Current.extra_pretty_print ())
+      | false -> None
+    in
+    let ch = Custom_channel.Channel.make ch ?onBeforeWrite () in
+    let close_out = Custom_channel.Channel.close in
     (try f ch
      with e ->
        close_out ch;
@@ -187,7 +194,7 @@ let ensure_file fp =
   then ()
   else
     gen_file fp (fun chan ->
-        let fmt = Pretty_print.to_out_channel chan in
+        let fmt = Pretty_print.to_custom_channel chan in
         Pretty_print.string fmt "true";
         Pretty_print.newline fmt)
 
@@ -372,7 +379,14 @@ let f
             ; env_instr () ]
         in
         let code = Code.prepend one.code instr in
-        let fmt = Pretty_print.to_out_channel stdout in
+
+        let ch = Custom_channel.ActualChannel.of_stdout stdout in
+        let onBeforeWrite = match Config.Flag.flowPrettyJs () with
+          | true -> Some (Backend.Current.extra_pretty_print ())
+          | false -> None
+        in
+        let ch = Custom_channel.Channel.make ch ?onBeforeWrite () in
+        let fmt = Pretty_print.to_custom_channel ch in
         RehpDriver.f
           ~standalone
           ?projectRoot:(get_project_root())
@@ -383,7 +397,8 @@ let f
           ?source_map
           fmt
           one.debug
-          code
+          code;
+        Custom_channel.Channel.close ch;
     | `Name file ->
         if file_needs_update use_hashing hashes_comment file
         then (
@@ -397,7 +412,7 @@ let f
                 List.concat [fs_instr1; pseudo_fs_init_instr (); env_instr ()]
               in
               let code = Code.prepend one.code instr in
-              let fmt = Pretty_print.to_out_channel chan in
+              let fmt = Pretty_print.to_custom_channel chan in
               RehpDriver.f
                 ~file
                 ?projectRoot:(get_project_root())
@@ -416,7 +431,7 @@ let f
               gen_file file (fun chan ->
                   let instr = fs_instr2 in
                   let code = Code.prepend Code.empty instr in
-                  let pfs_fmt = Pretty_print.to_out_channel chan in
+                  let pfs_fmt = Pretty_print.to_custom_channel chan in
                   RehpDriver.f
                     ~file
                     ?projectRoot:(get_project_root())

compiler/lib/backend.re

@@ -195,6 +195,7 @@ module type Backend_implementation = {
       )
     ) =>
     unit;
+  let extra_pretty_print: (unit, string) => string;
 };
 
 type t = (module Backend_implementation);
@@ -282,4 +283,9 @@ module Current: Backend_implementation = {
     | None => raise(Invalid_argument("Runtime setter not set"))
     | Some((module CurrentBackend)) => CurrentBackend.runtime_module_var()
     };
+  let extra_pretty_print = () =>
+    switch (current^) {
+    | None => raise(Invalid_argument("Extra pretty printer not set"))
+    | Some((module CurrentBackend)) => CurrentBackend.extra_pretty_print()
+    };
 };

compiler/lib/backend_js.re

@@ -318,6 +318,11 @@ let output =
   };
 };
 
+let extra_pretty_print = ((), partially_formatted_output) =>
+  Js_of_ocaml_flow_pretty_printer.Util.pretty_print(
+    partially_formatted_output,
+  );
+
 /**
  * Primitive.exists only returns true if it was provided by the linker or
  * something that actually registers it with its arity/kind etc typically

compiler/lib/backend_php.re

@@ -341,6 +341,8 @@ let output =
   };
 };
 
+let extra_pretty_print = ((), partially_formatted_output) => partially_formatted_output;
+
 /**
  * In PHP or other backends, we aren't doing any intelligent detection of
  * which external bindings are linked in. So instead we keep a manual list of

compiler/lib/config.ml

@@ -46,6 +46,8 @@ module Flag = struct
 
   let prettiestJs = o ~name:"prettiest-js" ~default:false
 
+  let flowPrettyJs = o ~name:"flow-pretty-js" ~default:false
+
   let stable_var = o ~name:"stable_var" ~default:false
 
   let debuginfo = o ~name:"debuginfo" ~default:false

compiler/lib/config.mli

@@ -46,6 +46,8 @@ module Flag : sig
 
   val prettiestJs : unit -> bool
 
+  val flowPrettyJs : unit -> bool
+
   val stable_var : unit -> bool
 
   val debuginfo : unit -> bool

compiler/lib/custom_channel.re

@@ -0,0 +1,58 @@
+module ActualChannel = {
+  type t =
+    | OutChannel(out_channel)
+    | Stdout(out_channel)
+    | Buffer(Buffer.t);
+
+  let of_out_channel = oc => OutChannel(oc);
+  let of_stdout = so => Stdout(so);
+  let of_buffer = b => Buffer(b);
+};
+
+module type CustomChannel = {
+  type t;
+  let make: (ActualChannel.t, ~onBeforeWrite: string => string=?, unit) => t;
+  let add_substring: (t, string, int, int) => unit;
+  let close: t => unit;
+};
+
+module Channel: CustomChannel = {
+  type t = {
+    channel: ActualChannel.t,
+    buffer: Buffer.t,
+    onBeforeWrite: string => string,
+  };
+  let make = (ch, ~onBeforeWrite=?, ()) => {
+    channel: ch,
+    buffer: Buffer.create(100),
+    onBeforeWrite:
+      switch (onBeforeWrite) {
+      | Some(fn) => fn
+      | None => (x => x)
+      },
+  };
+  let add_substring = (t, s, i, l) =>
+    Buffer.add_substring(t.buffer, s, i, l);
+  let close = t =>
+    t.buffer
+    |> Buffer.contents
+    |> t.onBeforeWrite
+    |> (
+      switch (t.channel) {
+      | ActualChannel.OutChannel(oc) => output_string(oc)
+      | ActualChannel.Stdout(so) => output_string(so)
+      | ActualChannel.Buffer(b) => Buffer.add_string(b)
+      }
+    )
+    |> (
+      () => {
+        // Cleanup
+        Buffer.clear(t.buffer);
+        switch (t.channel) {
+        | ActualChannel.OutChannel(oc) => close_out(oc)
+        | ActualChannel.Stdout(_)
+        | ActualChannel.Buffer(_) => ()
+        };
+      }
+    );
+};

compiler/lib/dune

@@ -2,7 +2,7 @@
   (name js_of_ocaml_compiler)
   (public_name js_of_ocaml-compiler)
   (synopsis "Js_of_ocaml compiler library")
-  (libraries compiler-libs.common compiler-libs.bytecomp bytes
+  (libraries compiler-libs.common compiler-libs.bytecomp bytes js_of_ocaml-flow_pretty_printer
    (select source_map_io.ml from
     (yojson -> source_map_yojson.ml.in)
     (       -> source_map_unsupported.ml.in))

compiler/lib/pretty_print.ml

@@ -182,7 +182,6 @@ let start_group st n = if not st.compact then push st (Start_group n)
 let end_group st = if not st.compact then push st End_group
 
 (*
-
 let render l =
   let st = { indent = 0; box_indent = 0; prev_indents = [];
              limit = 78; cur = 0; l = []; n = 0; w = 0;
@@ -191,15 +190,12 @@ let render l =
   List.iter (fun e -> push st e) l;
   push st End_group;
   output_newline st
-
 let rec tree n =
   if n = 0 then [Text "Leaf"] else
   [Start_group 10; Text "Node.... ("] @ tree (n - 1) @
   [Text ","; Break (" ", 0)] @ tree (n - 1) @ [Text ")"; End_group]
-
 let _ =
 for i = 1 to 10 do render (tree i) done
-
 *)
 
 let total t = t.total
@@ -235,6 +231,14 @@ let to_buffer b =
     compact = false; pending_space = None; last_char = None; needed_space = None;
     output = fun s i l -> Buffer.add_substring b s i l }
 
+let to_custom_channel ch = 
+  { indent = 0; box_indent = 0; prev_indents = [];
+    limit = 78; cur = 0; l = []; n = 0; w = 0;
+    col = 0; line = 0; total = 0;
+    compact = false; pending_space = None; last_char = None; needed_space = None;
+    output = Custom_channel.Channel.add_substring ch
+  }
+
 let set_compact st v = st.compact <- v
 
 let set_needed_space_function st f = st.needed_space <- Some f

compiler/lib/pretty_print.mli

@@ -40,6 +40,7 @@ val newline : t -> unit
 
 val to_out_channel : out_channel -> t
 val to_buffer : Buffer.t -> t
+val to_custom_channel : Custom_channel.Channel.t -> t
 val pos : t -> pos
 val total : t -> int