Find live toplevel bindings in Js_shake with a worklist

CHANGELOG.md

@@ -52,6 +52,7 @@
 - Make compilation slightly faster and stop rewatch from copying files twice for modules with an interface. https://github.com/rescript-lang/rescript/pull/8769
 - Speed up inlining decisions for large functions. https://github.com/rescript-lang/rescript/pull/8771
 - Speed up the rewriting of JS globals shadowed by local bindings in nested functions. https://github.com/rescript-lang/rescript/pull/8772
+- Speed up dead code elimination of toplevel bindings in large modules. https://github.com/rescript-lang/rescript/pull/8773
 - Represent explicit expression braces as `Pexp_braces` in parsetree v1 and format `else` branches consistently with `if` branches. https://github.com/rescript-lang/rescript/pull/8678
 - Omit redundant braces around multi-statement switch case bodies when formatting. https://github.com/rescript-lang/rescript/pull/8677
 - Avoid running `rescript-schema-ppx` and `sury-ppx` on source files without an `@schema` annotation. https://github.com/rescript-lang/rescript/pull/8662

compiler/core/js_shake.ml

@@ -22,57 +22,53 @@
  * along with this program; if not, write to the Free Software
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA. *)
 
-(** we also need make it complete 
-*)
-let get_initial_exports count_non_variable_declaration_statement
-    (export_set : Set_ident.t) (block : J.block) =
-  let result =
-    Ext_list.fold_left block export_set (fun acc st ->
-        match st.statement_desc with
-        | Variable {ident; value; _} -> (
-          if Set_ident.mem acc ident then
-            match value with
-            | None -> acc
-            | Some x ->
-              (* If not a function, we have to calcuate again and again
-                  TODO: add hashtbl for a cache
-              *)
-              Set_ident.(union (Js_analyzer.free_variables_of_expression x) acc)
-          else
-            match value with
-            | None -> acc
-            | Some x ->
-              if Js_analyzer.no_side_effect_expression x then acc
-              else
-                Set_ident.(
-                  union
-                    (Js_analyzer.free_variables_of_expression x)
-                    (add acc ident)))
-        | _ ->
-          (* recalcuate again and again ... *)
-          if
-            Js_analyzer.no_side_effect_statement st
-            || not count_non_variable_declaration_statement
-          then acc
-          else
-            Set_ident.(union (Js_analyzer.free_variables_of_statement st) acc))
+(** The identifiers that the toplevel [block] needs: the exports, everything a
+    statement with side effects refers to, and transitively the free variables
+    of the initializers of needed declarations. Each initializer's free
+    variables are computed once, and a worklist follows the dependencies. *)
+let live_idents (export_set : Set_ident.t) (block : J.block) : Set_ident.t =
+  (* free variables of the initializer of each toplevel declaration *)
+  let deps = Hash_ident.create 64 in
+  let live = ref Set_ident.empty in
+  let worklist = ref [] in
+  let mark id =
+    if not (Set_ident.mem !live id) then (
+      live := Set_ident.add !live id;
+      worklist := id :: !worklist)
+  in
+  Ext_list.iter block (fun (st : J.statement) ->
+      match st.statement_desc with
+      | Variable {ident; value = Some x; _} ->
+        let fv = Js_analyzer.free_variables_of_expression x in
+        let fv =
+          match Hash_ident.find_opt deps ident with
+          | None -> fv
+          | Some other -> Set_ident.union fv other
+        in
+        Hash_ident.replace deps ident fv;
+        if not (Js_analyzer.no_side_effect_expression x) then mark ident
+      | Variable {value = None; _} -> ()
+      | _ ->
+        if not (Js_analyzer.no_side_effect_statement st) then
+          Set_ident.iter (Js_analyzer.free_variables_of_statement st) mark);
+  Set_ident.iter export_set mark;
+  let rec drain () =
+    match !worklist with
+    | [] -> ()
+    | id :: rest ->
+      worklist := rest;
+      (match Hash_ident.find_opt deps id with
+      | Some fv -> Set_ident.iter fv mark
+      | None -> ());
+      drain ()
   in
-  (result, Set_ident.(diff result export_set))
+  drain ();
+  !live
 
 let shake_program (program : J.program) =
   let shake_block block export_set =
     let block = List.rev @@ Js_analyzer.rev_toplevel_flatten block in
-    let loop block export_set : Set_ident.t =
-      let rec aux acc block =
-        let result, diff = get_initial_exports false acc block in
-        if Set_ident.is_empty diff then result else aux result block
-      in
-      let first_iteration, delta = get_initial_exports true export_set block in
-      if not @@ Set_ident.is_empty delta then aux first_iteration block
-      else first_iteration
-    in
-
-    let really_set = loop block export_set in
+    let really_set = live_idents export_set block in
     Ext_list.fold_right block [] (fun (st : J.statement) acc ->
         match st.statement_desc with
         | Variable {ident; value; _} -> (

tests/tests/src/shake_liveness_test.mjs

@@ -0,0 +1,54 @@
+// Generated by ReScript, PLEASE EDIT WITH CARE
+
+import * as Mocha from "mocha";
+import * as Test_utils from "./test_utils.mjs";
+
+let counter = {
+  contents: 0
+};
+
+function chain0(x) {
+  return ((x * 3 | 0) + (x * x | 0) | 0) + 1 | 0;
+}
+
+function chain1(x) {
+  return ((chain0(x) * 3 | 0) + chain0(x + 1 | 0) | 0) - (x << 1) | 0;
+}
+
+function chain2(x) {
+  return ((chain1(x) * 3 | 0) + chain1(x + 1 | 0) | 0) - (x << 1) | 0;
+}
+
+function chainEnd(x) {
+  return chain2(x) + chain2(x + 1 | 0) | 0;
+}
+
+function effectHelper(x) {
+  return (((x * x | 0) * 7 | 0) + (x * 3 | 0) | 0) + 5 | 0;
+}
+
+function bump(x) {
+  counter.contents = counter.contents + effectHelper(x) | 0;
+  return counter.contents;
+}
+
+bump(1);
+
+function statementHelper(x) {
+  return (((x * x | 0) * 5 | 0) + (x * 11 | 0) | 0) + 2 | 0;
+}
+
+counter.contents = counter.contents + statementHelper(2) | 0;
+
+Mocha.describe("File \"shake_liveness_test.res\", line 36, characters 9-16", () => {
+  Mocha.test("kept bindings work", () => {
+    Test_utils.eq("File \"shake_liveness_test.res\", line 39, characters 7-14", counter.contents, 59);
+    Test_utils.eq("File \"shake_liveness_test.res\", line 40, characters 7-14", chainEnd(1), 338);
+  });
+});
+
+export {
+  counter,
+  chainEnd,
+}
+/*  Not a pure module */

tests/tests/src/shake_liveness_test.res

@@ -0,0 +1,42 @@
+// Which toplevel bindings Js_shake keeps: exports, bindings with side effects,
+// bindings used by statements with side effects, and transitively whatever
+// they use. The generated JS shows that only the `dead*` bindings are removed.
+open Mocha
+open Test_utils
+
+let counter = ref(0)
+
+// A chain that runs backwards through the module: each binding is only needed
+// because a later one uses it, and only the last one is exported.
+%%private(let chain0 = (x: int) => x * 3 + x * x + 1)
+%%private(let chain1 = (x: int) => chain0(x) * 3 + chain0(x + 1) - 2 * x)
+%%private(let chain2 = (x: int) => chain1(x) * 3 + chain1(x + 1) - 2 * x)
+let chainEnd = (x: int) => chain2(x) + chain2(x + 1)
+
+// Unused, but its initializer has a side effect: it compiles to the statement
+// `bump(1)`, which keeps `bump` and, through it, `effectHelper`.
+%%private(let effectHelper = (x: int) => x * x * 7 + x * 3 + 5)
+%%private(
+  let bump = (x: int) => {
+    counter := counter.contents + effectHelper(x)
+    counter.contents
+  }
+)
+%%private(let effectful = bump(1))
+
+// Only used by a toplevel statement with a side effect.
+%%private(let statementHelper = (x: int) => x * x * 5 + x * 11 + 2)
+counter := counter.contents + statementHelper(2)
+
+// Unused chain without side effects: removed.
+%%private(let deadUsedByDead = (x: int) => x * x * 9 + x * 13 + 4)
+%%private(let dead0 = (x: int) => deadUsedByDead(x) * 3 + deadUsedByDead(x + 1))
+%%private(let dead1 = (x: int) => dead0(x) * 3 + dead0(x + 1) - x)
+
+describe(__LOC__, () => {
+  test("kept bindings work", () => {
+    // effectHelper(1) = 15, statementHelper(2) = 44
+    eq(__LOC__, counter.contents, 59)
+    eq(__LOC__, chainEnd(1), 338)
+  })
+})