[J2Cl] Make J2clOpts more effective with transitive deps in constant intialization (#6571)

Constants that need to be hoisted sometimes are initialized by calling
getters of other constants that need to be hoisted. These getters are
non-trivial, e.g.

  (func $getConst1_<once>_@X (result (ref null $A))
    (block (result (ref null $A))
      (if (i32.eqz (ref.is_null (global.get $$const1@X)))
        (then
          (return (global.get $$const1@X))
        )
      )
      (global.set $$const1@X (struct.new $A (i32.const 2)))
      (global.get $$const1@X)
    )

  (func $getConst2_<once>_@X (result (ref null $A))
    (block (result (ref null $A))
      (if (i32.eqz (ref.is_null (global.get $$const2@X)))
        (then
          (return (global.get $$const2@X))
        )
      )
      (global.set $$const2@X  .... expression with (call $getConst1_<once>_@X)  ....))
      (global.get $$const2@X)
    )

and can only be simplified after the constants they initialize are hoisted. After
the constant is hoisted the getter can be inlined and constants that depend on
it for their initialization can now be hoisted.

Before this pass, inlining would happen after the pass was run by a subsequent
run of the inliner (likely as part of -O3), requiring as many runs of this pass,
interleaved with the inliner, as the depth in the call sequence.

By having a simpler inliner run as part of the loop in this pass, the pass becomes
more effective and more independent of the call depths.

1 files changed, 129 insertions, 43 deletions

diff --git a/src/passes/J2CLOpts.cpp b/src/passes/J2CLOpts.cpp
index 97129d2cc..cc73e873e 100644
--- a/src/passes/J2CLOpts.cpp
+++ b/src/passes/J2CLOpts.cpp
@@ -30,9 +30,52 @@ namespace wasm {
 
 namespace {
 
-bool isOnceFunction(Function* f) { return f->name.hasSubstring("_<once>_"); }
-
 using AssignmentCountMap = std::unordered_map<Name, Index>;
+using TrivialFunctionMap = std::unordered_map<Name, Expression*>;
+
+bool isOnceFunction(Name name) { return name.hasSubstring("_<once>_"); }
+bool isOnceFunction(Function* func) { return isOnceFunction(func->name); }
+
+// Returns the function body if it is a trivial function, null otherwise.
+Expression* getTrivialFunctionBody(Function* func) {
+  auto* body = func->body;
+
+  // Only consider trivial the following instructions which can be safely
+  // inlined and note that their size is at most 2.
+  if (body->is<Nop>() || body->is<GlobalGet>() || body->is<Const>() ||
+      // Call with no arguments.
+      (body->is<Call>() && body->dynCast<Call>()->operands.size() == 0) ||
+      // Simple global.set with a constant.
+      (body->is<GlobalSet>() &&
+       body->dynCast<GlobalSet>()->value->is<Const>())) {
+    return body;
+  }
+  return nullptr;
+}
+
+// Adds the function to the map if it is trivial.
+void maybeCollectTrivialFunction(Function* func,
+                                 TrivialFunctionMap& trivialFunctionMap) {
+  auto* body = getTrivialFunctionBody(func);
+  if (body == nullptr) {
+    return;
+  }
+
+  trivialFunctionMap[func->name] = body;
+}
+
+// Cleans up a once function that has been modified in the hopes it
+// becomes trivial.
+void cleanupFunction(Module* module, Function* func) {
+  PassRunner runner(module);
+  runner.add("precompute");
+  runner.add("vacuum");
+  // Run after vacuum to remove the extra returns that vacuum might
+  // leave when reducing a block that ends up with just one instruction.
+  runner.add("remove-unused-brs");
+  runner.setIsNested(true);
+  runner.runOnFunction(func);
+}
 
 // A visitor to count the number of GlobalSets of each global so we can later
 // identify the number of assignments of the global.
@@ -47,9 +90,9 @@ public:
     if (isInitialValue(curr->init)) {
       return;
     }
-    // J2CL normally doesn't set non-default initial value however just in
-    // case if other passes in bineryen does something we should back off
-    // by recording this as an assignment.
+    // J2CL normally doesn't set non-default initial values, however, just in
+    // case other passes in binaryen do something and set a value to the global
+    // we should back off by recording this as an assignment.
     recordGlobalAssignment(curr->name);
   }
   void visitGlobalSet(GlobalSet* curr) { recordGlobalAssignment(curr->name); }
@@ -81,8 +124,10 @@ private:
 // TODO: parallelize
 class ConstantHoister : public WalkerPass<PostWalker<ConstantHoister>> {
 public:
-  ConstantHoister(AssignmentCountMap& assignmentCounts)
-    : assignmentCounts(assignmentCounts) {}
+  ConstantHoister(AssignmentCountMap& assignmentCounts,
+                  TrivialFunctionMap& trivialFunctionMap)
+    : assignmentCounts(assignmentCounts),
+      trivialFunctionMap(trivialFunctionMap) {}
 
   int optimized = 0;
 
@@ -101,15 +146,8 @@ public:
     }
 
     if (optimized != optimizedBefore) {
-      // We introduced "nop" instruction. Run the vacuum to cleanup.
-      // TODO: maybe we should not introduce "nop" in the first place and try
-      // removing instead.
-      PassRunner runner(getModule());
-      runner.add("precompute");
-      runner.add("remove-unused-brs");
-      runner.add("vacuum");
-      runner.setIsNested(true);
-      runner.runOnFunction(curr);
+      cleanupFunction(getModule(), curr);
+      maybeCollectTrivialFunction(curr, trivialFunctionMap);
     }
   }
 
@@ -154,54 +192,105 @@ private:
   }
 
   AssignmentCountMap& assignmentCounts;
+  TrivialFunctionMap& trivialFunctionMap;
 };
 
-// A visitor that marks trivial once functions for removal.
-// TODO: parallelize
-class EnableInline : public WalkerPass<PostWalker<EnableInline>> {
+// Class to collect functions that are already trivial before the pass is run.
+// When this pass is run, other optimizations that preceded it might have left
+// the body of some of these functions trivial.
+// Since the loop in this pass will only inline the functions that are made
+// trivial by this pass, the functions that were already trivial before would
+// not be inlined if they were not collected by this visitor.
+class TrivialOnceFunctionCollector
+  : public WalkerPass<PostWalker<TrivialOnceFunctionCollector>> {
 public:
+  TrivialOnceFunctionCollector(TrivialFunctionMap& trivialFunctionMap)
+    : trivialFunctionMap(trivialFunctionMap) {}
+
   void visitFunction(Function* curr) {
     if (!isOnceFunction(curr)) {
       return;
     }
-    auto* body = curr->body;
-    if (Measurer::measure(body) <= 2) {
-      // Once function has become trivial: enable inlining so it could be
-      // removed.
-      //
-      // Note that the size of a trivial function is set to 2 to cover things
-      // like
-      //  - nop (e.g. a clinit that is emptied)
-      //  - call (e.g. a clinit just calling another one; note that clinits take
-      //  no parameters)
-      //  - return global.get abc (e.g. a string literal moved to global)
-      //  - global.set abc 1 (e.g. partially inlined clinit that is empty)
-      // while rejecting more complex scenario like condition checks. Optimizing
-      // complex functions could change the shape of "once" functions and make
-      // the ConstantHoister in this pass and OnceReducer ineffective.
-      curr->noFullInline = false;
-      curr->noPartialInline = false;
+    maybeCollectTrivialFunction(curr, trivialFunctionMap);
+  }
+
+private:
+  TrivialFunctionMap& trivialFunctionMap;
+};
+
+// A visitor that inlines trivial once functions.
+// TODO: parallelize
+class InlineTrivialOnceFunctions
+  : public WalkerPass<PostWalker<InlineTrivialOnceFunctions>> {
+public:
+  InlineTrivialOnceFunctions(TrivialFunctionMap& trivialFunctionMap)
+    : trivialFunctionMap(trivialFunctionMap) {}
+
+  void visitCall(Call* curr) {
+    if (curr->operands.size() != 0 || !isOnceFunction(curr->target)) {
+      return;
+    }
+
+    auto iter = trivialFunctionMap.find(curr->target);
+    if (iter == trivialFunctionMap.end()) {
+      return;
     }
+    auto* expr = iter->second;
+
+    // The call was to a trivial once function which consists of the expression
+    // in <expr>; replace the call with it.
+    Builder builder(*getModule());
+    auto* replacement = ExpressionManipulator::copy(expr, *getModule());
+    replaceCurrent(replacement);
+
+    lastModifiedFunction = getFunction();
+    inlined++;
   }
+
+  void visitFunction(Function* curr) {
+    // Since the traversal is in post-order, we only need to check if the
+    // current function is the function that was last inlined into.
+    // We also do not want to do any cleanup for a non-once function (we leave
+    // that for other passes, as it will not end up helping further work here).
+    if (lastModifiedFunction != curr || !isOnceFunction(curr)) {
+      return;
+    }
+
+    cleanupFunction(getModule(), curr);
+    maybeCollectTrivialFunction(curr, trivialFunctionMap);
+  }
+
+  int inlined = 0;
+
+private:
+  TrivialFunctionMap& trivialFunctionMap;
+  Function* lastModifiedFunction = nullptr;
 };
 
 struct J2CLOpts : public Pass {
   void hoistConstants(Module* module) {
     AssignmentCountMap assignmentCounts;
+    TrivialFunctionMap trivialFunctionMap;
 
     GlobalAssignmentCollector collector(assignmentCounts);
     collector.run(getPassRunner(), module);
 
-    // TODO surprisingly looping help some but not a lot. Maybe we are missing
-    // something that helps globals to be considered immutable.
+    TrivialOnceFunctionCollector trivialFunctionCollector(trivialFunctionMap);
+    trivialFunctionCollector.run(getPassRunner(), module);
+
     while (1) {
-      ConstantHoister hoister(assignmentCounts);
+      ConstantHoister hoister(assignmentCounts, trivialFunctionMap);
       hoister.run(getPassRunner(), module);
       int optimized = hoister.optimized;
+
+      InlineTrivialOnceFunctions inliner(trivialFunctionMap);
+      inliner.run(getPassRunner(), module);
+      int inlined = inliner.inlined;
+
 #ifdef J2CL_OPT_DEBUG
       std::cout << "Optimized " << optimized << " global fields\n";
 #endif
-      if (optimized == 0) {
+      if (optimized == 0 && inlined == 0) {
         break;
       }
     }
@@ -215,9 +304,6 @@ struct J2CLOpts : public Pass {
     // initialization.
     hoistConstants(module);
 
-    EnableInline enableInline;
-    enableInline.run(getPassRunner(), module);
-
     // We might have introduced new globals depending on other globals. Reorder
     // order them so they follow initialization order.
     // TODO: do this only if have introduced a new global.