Improve Js_assign logic

CHANGES.md

@@ -28,7 +28,7 @@
 * Compiler: improve debug/sourcemap location of closures (#1947)
 * Compiler: improve tailcall optimization (#1943)
 * Compiler: improve deadcode optimization (#1963, #1962, #1967)
-* Compiler: improve coloring optimization (#1971, #1984, #1989)
+* Compiler: improve coloring optimization (#1971, #1984, #1986, #1989)
 * Compiler: faster constant sharing (#1988)
 * Compiler: more efficient code generation from bytecode (#1972)
 * Runtime: use Dataview to convert between floats and bit representation

compiler/bin-jsoo_minify/jsoo_minify.ml

@@ -45,7 +45,6 @@ let f { Cmd_arg.common; output_file; use_stdin; files } =
   let gen pp =
     let pretty = Config.Flag.pretty () in
     Pretty_print.set_compact pp (not pretty);
-    Code.Var.set_pretty pretty;
     let error_of_pi pi =
       match pi with
       | { Parse_info.name = Some src; line; col; _ }
@@ -67,17 +66,6 @@ let f { Cmd_arg.common; output_file; use_stdin; files } =
         try p @ Parse_js.parse lex with Parse_js.Parsing_error pi -> error_of_pi pi
       else p
     in
-    let free = new Js_traverse.free in
-    let (_ : Javascript.program) = free#program p in
-    let toplevel_def_and_use =
-      let state = free#state in
-      Javascript.IdentSet.union state.def_var state.use
-    in
-    Javascript.IdentSet.iter
-      (function
-        | V _ -> ()
-        | S { name = Utf8_string.Utf8 x; _ } -> Var_printer.add_reserved x)
-      toplevel_def_and_use;
     let true_ () = true in
     let open Config in
     let passes : ((unit -> bool) * (unit -> Js_traverse.mapper)) list =

compiler/bin-wasm_of_ocaml/compile.ml

@@ -74,10 +74,7 @@ let opt_with action x f =
   | None -> f None
   | Some x -> action x (fun y -> f (Some y))
 
-let output_gen output_file f =
-  Code.Var.set_pretty true;
-  Code.Var.set_stable (Config.Flag.stable_var ());
-  Filename.gen_file output_file f
+let output_gen output_file f = Filename.gen_file output_file f
 
 let with_runtime_files ~runtime_wasm_files f =
   let inputs =

compiler/lib-wasm/wasm_output.ml

@@ -1182,6 +1182,4 @@ let f ch fields =
 
     let string = output_string
   end) in
-  Code.Var.set_pretty true;
-  Code.Var.set_stable (Config.Flag.stable_var ());
   O.output_module ch fields

compiler/lib/code.ml

@@ -50,8 +50,6 @@ module Var : sig
 
   val of_idx : int -> t
 
-  val to_string : ?origin:t -> t -> string
-
   val fresh : unit -> t
 
   val fresh_n : string -> t
@@ -62,18 +60,14 @@ module Var : sig
 
   val compare : t -> t -> int
 
-  val name : t -> string -> unit
+  val set_name : t -> string -> unit
 
   val get_name : t -> string option
 
   val propagate_name : t -> t -> unit
 
   val reset : unit -> unit
 
-  val set_pretty : bool -> unit
-
-  val set_stable : bool -> unit
-
   module Set : Set.S with type elt = t
 
   module Map : Map.S with type key = t
@@ -130,32 +124,106 @@ end = struct
 
   let printer = Var_printer.create Var_printer.Alphabet.javascript
 
+  module Name = struct
+    let names = Int.Hashtbl.create 100
+
+    let reset () = Int.Hashtbl.clear names
+
+    let reserved = String.Hashtbl.create 100
+
+    let () = StringSet.iter (fun s -> String.Hashtbl.add reserved s ()) Reserved.keyword
+
+    let is_reserved s = String.Hashtbl.mem reserved s
+
+    let merge n1 n2 =
+      match n1, n2 with
+      | "", n2 -> n2
+      | n1, "" -> n1
+      | n1, n2 ->
+          if generated_name n1
+          then n2
+          else if generated_name n2
+          then n1
+          else if String.length n1 > String.length n2
+          then n1
+          else n2
+
+    let set_raw v nm = Int.Hashtbl.replace names v nm
+
+    let propagate v v' =
+      try
+        let name = Int.Hashtbl.find names v in
+        match Int.Hashtbl.find names v' with
+        | exception Not_found -> set_raw v' name
+        | name' -> set_raw v' (merge name name')
+      with Not_found -> ()
+
+    let set v nm_orig =
+      let len = String.length nm_orig in
+      if len > 0
+      then (
+        let buf = Buffer.create (String.length nm_orig) in
+        let idx = ref 0 in
+        while !idx < len && not (Char.is_letter nm_orig.[!idx]) do
+          incr idx
+        done;
+        let pending = ref false in
+        if !idx >= len
+        then (
+          pending := true;
+          idx := 0);
+        for i = !idx to len - 1 do
+          if Char.is_letter nm_orig.[i] || Char.is_digit nm_orig.[i]
+          then (
+            if !pending then Buffer.add_char buf '_';
+            Buffer.add_char buf nm_orig.[i];
+            pending := false)
+          else pending := true
+        done;
+        let str = Buffer.contents buf in
+        let str =
+          match str, nm_orig with
+          | "", ">>=" -> "symbol_bind"
+          | "", ">>|" -> "symbol_map"
+          | "", "^" -> "symbol_concat"
+          | "", _ -> "symbol"
+          | str, _ -> if is_reserved str then str ^ "$" else str
+        in
+        (* protect against large names *)
+        let max_len = 30 in
+        let str =
+          if String.length str > max_len then String.sub str ~pos:0 ~len:max_len else str
+        in
+        set_raw v str)
+
+    let get v = try Some (Int.Hashtbl.find names v) with Not_found -> None
+  end
+
   let last_var = ref 0
 
   let reset () =
     last_var := 0;
-    Var_printer.reset printer
-
-  let to_string ?origin i = Var_printer.to_string printer ?origin i
+    Var_printer.reset printer;
+    Name.reset ()
 
   let print f x =
     Format.fprintf
       f
       "v%d%s"
       x
-      (match Var_printer.get_name printer x with
+      (match Name.get x with
       | None -> ""
       | Some nm -> "{" ^ nm ^ "}")
 
-  let name i nm = Var_printer.name printer i nm
+  let set_name i nm = Name.set i nm
 
   let fresh () =
     incr last_var;
     !last_var
 
   let fresh_n nm =
     incr last_var;
-    name !last_var nm;
+    set_name !last_var nm;
     !last_var
 
   let count () = !last_var + 1
@@ -164,13 +232,9 @@ end = struct
 
   let of_idx v = v
 
-  let get_name i = Var_printer.get_name printer i
-
-  let propagate_name i j = Var_printer.propagate_name printer i j
-
-  let set_pretty b = Var_printer.set_pretty printer b
+  let get_name i = Name.get i
 
-  let set_stable b = Var_printer.set_stable printer b
+  let propagate_name i j = Name.propagate i j
 
   let fork o =
     let n = fresh () in

compiler/lib/code.mli

@@ -48,8 +48,6 @@ module Var : sig
 
   val of_idx : int -> t
 
-  val to_string : ?origin:t -> t -> string
-
   val fresh : unit -> t
 
   val fresh_n : string -> t
@@ -62,16 +60,12 @@ module Var : sig
 
   val get_name : t -> string option
 
-  val name : t -> string -> unit
+  val set_name : t -> string -> unit
 
   val propagate_name : t -> t -> unit
 
   val reset : unit -> unit
 
-  val set_pretty : bool -> unit
-
-  val set_stable : bool -> unit
-
   module Set : Set.S with type elt = t
 
   module Map : Map.S with type key = t

compiler/lib/driver.ml

@@ -421,8 +421,6 @@ let check_js js =
   let missing = StringSet.inter free all_external in
   let missing = StringSet.diff missing Reserved.provided in
   let other = StringSet.diff free missing in
-  let res = Var_printer.get_reserved () in
-  let other = StringSet.diff other res in
   if not (StringSet.is_empty missing) then report_missing_primitives missing;
   let probably_prov = StringSet.inter other Reserved.provided in
   let other = StringSet.diff other probably_prov in
@@ -449,8 +447,6 @@ let name_variables js =
       js)
     else js
   in
-  let o = new Js_traverse.fast_freevar (fun s -> Var_printer.add_reserved s) in
-  o#program js;
   let js = Js_assign.program js in
   if times () then Format.eprintf "  coloring: %a@." Timer.print t;
   js
@@ -618,9 +614,7 @@ let simplify_js js =
 
 let configure formatter =
   let pretty = Config.Flag.pretty () in
-  Pretty_print.set_compact formatter (not pretty);
-  Code.Var.set_pretty (pretty && not (Config.Flag.shortvar ()));
-  Code.Var.set_stable (Config.Flag.stable_var ())
+  Pretty_print.set_compact formatter (not pretty)
 
 let link_and_pack ?(standalone = true) ?(wrap_with_fun = `Iife) ?(link = `No) p =
   let export_runtime =

compiler/lib/effects.ml

@@ -775,8 +775,7 @@ let cps_transform ~live_vars ~flow_info ~cps_needed p =
     Code.fold_closures_innermost_first
       p
       (fun name_opt params (start, args) _cloc ({ Code.blocks; free_pc; _ } as p) ->
-        Option.iter name_opt ~f:(fun v ->
-            debug_print "@[<v>cname = %s@,@]" @@ Var.to_string v);
+        Option.iter name_opt ~f:(fun v -> debug_print "@[<v>cname = %a@,@]" Var.print v);
         (* We speculatively add a block at the beginning of the
            function. In case of tail-recursion optimization, the
            function implementing the loop body may have to be placed
@@ -1128,7 +1127,7 @@ let f ~flow_info ~live_vars p =
       then (
         debug_print "@]";
         debug_print "@[<v>cps_needed (after lifting) = @[<hov 2>";
-        Var.Set.iter (fun v -> debug_print "%s,@ " (Var.to_string v)) cps_needed;
+        Var.Set.iter (fun v -> debug_print "%a,@ " Var.print v) cps_needed;
         debug_print "@]@,@]";
         debug_print "@[<v>After lambda lifting...@,";
         Code.Print.program Format.err_formatter (fun _ _ -> "") p;

compiler/lib/generate_closure.ml

@@ -169,8 +169,11 @@ module Trampoline = struct
         then (
           Format.eprintf "Detect cycles of size (%d).\n%!" (List.length all);
           Format.eprintf
-            "%s\n%!"
-            (String.concat ~sep:", " (List.map all ~f:(fun x -> Var.to_string x))));
+            "%a\n%!"
+            (Format.pp_print_list
+               ~pp_sep:(fun fmt () -> Format.pp_print_string fmt ", ")
+               Var.print)
+            all);
         let tailcall_max_depth = Config.Param.tailcall_max_depth () in
         let all =
           List.map all ~f:(fun id ->
@@ -183,7 +186,7 @@ module Trampoline = struct
             ~init:(blocks, free_pc, [])
             ~f:(fun (blocks, free_pc, closures) (counter, ci) ->
               if debug_tc ()
-              then Format.eprintf "Rewriting for %s\n%!" (Var.to_string ci.f_name);
+              then Format.eprintf "Rewriting for %a\n%!" Var.print ci.f_name;
               let new_f = Code.Var.fork ci.f_name in
               let new_args = List.map ci.args ~f:Code.Var.fork in
               let wrapper_pc = free_pc in