Compiler: minify labels

Author: hhugo

compiler/lib/generate.ml

@@ -1014,11 +1014,6 @@ let throw_statement ctx cx k loc =
         , loc )
       ]
 
-let next_label scope_stack =
-  match scope_stack with
-  | (_, (l, _, _)) :: _ -> J.Label.succ l
-  | [] -> J.Label.zero
-
 let rec translate_expr ctx queue loc x e level : _ * J.statement_list =
   match e with
   | Apply { f; args; exact } ->
@@ -1400,7 +1395,7 @@ and compile_block st queue (pc : Addr.t) scope_stack ~fall_through =
     | true ->
         if debug () then Format.eprintf "@[<hv 2>for(;;) {@,";
         let never_body, body =
-          let lab = next_label scope_stack in
+          let lab = J.Label.fresh () in
           let lab_used = ref false in
           let exit_branch_used = ref false in
           let scope_stack = (pc, (lab, lab_used, Loop)) :: scope_stack in
@@ -1463,7 +1458,7 @@ and compile_block_no_loop st queue (pc : Addr.t) ~fall_through scope_stack =
     match l with
     | [] -> compile_conditional st queue ~fall_through block.branch scope_stack
     | x :: xs -> (
-        let l = next_label scope_stack in
+        let l = J.Label.fresh () in
         let used = ref false in
         let scope_stack = (x, (l, used, Forward)) :: scope_stack in
         let _never_inner, inner = loop ~scope_stack ~fall_through:(Block x) xs in
@@ -1515,7 +1510,7 @@ and compile_decision_tree kind st scope_stack loc cx dtree ~fall_through =
         let all_never = ref true in
         let len = Array.length a in
         let last_index = len - 1 in
-        let lab = next_label scope_stack in
+        let lab = J.Label.fresh () in
         let label_used = ref false in
         let exit_branch_used = ref false in
         let scope_stack =

compiler/lib/javascript.ml

@@ -143,20 +143,10 @@ end
 
 module Label = struct
   type t =
-    | L of int
+    | L of Code.Var.t
     | S of Utf8_string.t
 
-  let printer = Var_printer.create Var_printer.Alphabet.javascript
-
-  let zero = L 0
-
-  let succ = function
-    | L t -> L (succ t)
-    | S _ -> assert false
-
-  let to_string = function
-    | L t -> Utf8_string.of_string_exn (Var_printer.to_string printer t)
-    | S s -> s
+  let fresh () = L (Code.Var.fresh ())
 
   let of_string s = S s
 end

compiler/lib/javascript.mli

@@ -51,13 +51,11 @@ module Num : sig
 end
 
 module Label : sig
-  type t
-
-  val zero : t
-
-  val succ : t -> t
+  type t =
+    | L of Code.Var.t
+    | S of Utf8_string.t
 
-  val to_string : t -> Utf8_string.t
+  val fresh : unit -> t
 
   val of_string : Utf8_string.t -> t
 end

compiler/lib/js_assign.ml

@@ -329,11 +329,48 @@ class traverse record_block =
       super#record_block b
   end
 
+class traverse_labels h =
+  object
+    inherit Js_traverse.iter as super
+
+    val ldepth = 0
+
+    method fun_decl (_k, _params, body, _loc) =
+      let m = {<ldepth = 0>} in
+      m#function_body body
+
+    method statement =
+      function
+      | Labelled_statement (L l, (s, _)) ->
+          let m = {<ldepth = ldepth + 1>} in
+          Hashtbl.add h l ldepth;
+          m#statement s
+      | s -> super#statement s
+  end
+
+class name ident label =
+  object (m)
+    inherit Js_traverse.subst ident as super
+
+    method statement =
+      function
+      | Labelled_statement (l, (s, loc)) ->
+          Labelled_statement (label l, (m#statement s, loc))
+      | Break_statement (Some l) -> Break_statement (Some (label l))
+      | Continue_statement (Some l) -> Continue_statement (Some (label l))
+      | s -> super#statement s
+  end
+
 let program' (module Strategy : Strategy) p =
   let nv = Var.count () in
   let state = Strategy.create nv in
+  let labels = Hashtbl.create 20 in
   let mapper = new traverse (Strategy.record_block state) in
   let p = mapper#program p in
+  let () =
+    let o = new traverse_labels labels in
+    o#program p
+  in
   mapper#record_block Normal;
   let free =
     IdentSet.filter
@@ -350,7 +387,7 @@ let program' (module Strategy : Strategy) p =
       | S _ -> ()
       | V x -> names.(Var.idx x) <- "")
     free;
-  let color = function
+  let ident = function
     | V v -> (
         let name = names.(Var.idx v) in
         match name, has_free_var with
@@ -359,7 +396,18 @@ let program' (module Strategy : Strategy) p =
         | _, (true | false) -> ident ~var:v (Utf8_string.of_string_exn name))
     | x -> x
   in
-  let p = (new Js_traverse.subst color)#program p in
+  let label_printer = Var_printer.create Var_printer.Alphabet.javascript in
+  let max_label_depth = Hashtbl.fold (fun _ d acc -> max d acc) labels 0 in
+  let lname_per_depth =
+    Array.init (max_label_depth + 1) ~f:(fun i -> Var_printer.to_string label_printer i)
+  in
+  let label = function
+    | Label.S _ as l -> l
+    | L v ->
+        let i = Hashtbl.find labels v in
+        S (Utf8_string.of_string_exn lname_per_depth.(i))
+  in
+  let p = (new name ident label)#program p in
   (if has_free_var
    then
      let () =

compiler/lib/js_output.ml

@@ -60,6 +60,10 @@ end) =
 struct
   open D
 
+  let nane_of_label = function
+    | Javascript.Label.L _ -> assert false
+    | Javascript.Label.S n -> n
+
   let debug_enabled = Config.Flag.debuginfo ()
 
   let output_debug_info f loc =
@@ -1260,15 +1264,15 @@ struct
         last_semi ()
     | Continue_statement (Some s) ->
         PP.string f "continue ";
-        let (Utf8 l) = Javascript.Label.to_string s in
+        let (Utf8 l) = nane_of_label s in
         PP.string f l;
         last_semi ()
     | Break_statement None ->
         PP.string f "break";
         last_semi ()
     | Break_statement (Some s) ->
         PP.string f "break ";
-        let (Utf8 l) = Javascript.Label.to_string s in
+        let (Utf8 l) = nane_of_label s in
         PP.string f l;
         last_semi ()
     | Return_statement e -> (
@@ -1309,7 +1313,7 @@ struct
             (* There MUST be a space between the return and its
                argument. A line return will not work *))
     | Labelled_statement (i, s) ->
-        let (Utf8 l) = Javascript.Label.to_string i in
+        let (Utf8 l) = nane_of_label i in
         PP.string f l;
         PP.string f ":";
         PP.space f;

compiler/lib/js_traverse.ml

@@ -976,6 +976,8 @@ class rename_variable =
 
     val decl = StringSet.empty
 
+    val labels = StringMap.empty
+
     method private update_state scope params iter_body =
       let declared_names = declared scope params iter_body in
       {<subst = StringSet.fold
@@ -1011,6 +1013,30 @@ class rename_variable =
 
     method statement s =
       match s with
+      | Labelled_statement (l, (s, loc)) ->
+          let l, m =
+            match l with
+            | L _ -> l, m
+            | S (Utf8 u) ->
+                let l = Label.fresh () in
+                let m = {<labels = StringMap.add u l labels>} in
+                l, m
+          in
+          Labelled_statement (l, (m#statement s, loc))
+      | Break_statement (Some l) -> (
+          match l with
+          | L _ -> s
+          | S (Utf8 l) -> (
+              match StringMap.find_opt l labels with
+              | None -> s
+              | Some l -> Break_statement (Some l)))
+      | Continue_statement (Some l) -> (
+          match l with
+          | L _ -> s
+          | S (Utf8 l) -> (
+              match StringMap.find_opt l labels with
+              | None -> s
+              | Some l -> Continue_statement (Some l)))
       | Function_declaration (id, (k, params, body, nid)) ->
           let ids = bound_idents_of_params params in
           let m' = m#update_state (Fun_block None) ids body in