OnceReduction: Consider calls to "once" functions in "once" functions

src/passes/OnceReduction.cpp

@@ -197,9 +197,10 @@ struct Scanner : public WalkerPass<PostWalker<Scanner>> {
   OptInfo& optInfo;
 };
 
-// Information in a basic block. We track relevant expressions, which are calls
-// calls to "once" functions, and writes to "once" globals.
+// Information in a basic block.
 struct BlockInfo {
+  // We track relevant expressions, which are call to "once" functions, and
+  // writes to "once" globals.
   std::vector<Expression*> exprs;
 };
 
@@ -312,18 +313,17 @@ struct Optimizer
             optimizeOnce(set->name);
           }
         } else if (auto* call = expr->dynCast<Call>()) {
-          if (optInfo.onceFuncs.at(call->target).is()) {
+          auto target = call->target;
+          if (optInfo.onceFuncs.at(target).is()) {
             // The global used by the "once" func is written.
             assert(call->operands.empty());
-            optimizeOnce(optInfo.onceFuncs.at(call->target));
-            continue;
+            optimizeOnce(optInfo.onceFuncs.at(target));
+            // Fall through to the code below to apply other things we know
+            // about the caller.
           }
 
-          // This is not a call to a "once" func. However, we may have inferred
-          // that it definitely sets some "once" globals before it returns, and
-          // we can use that information.
-          for (auto globalName :
-               optInfo.onceGlobalsSetInFuncs.at(call->target)) {
+          // Note as written all globals the called function is known to write.
+          for (auto globalName : optInfo.onceGlobalsSetInFuncs.at(target)) {
             onceGlobalsWritten.insert(globalName);
           }
         } else {
@@ -339,6 +339,80 @@ struct Optimizer
     // TODO: Aside from the entry block, we could intersect all the exit blocks.
     optInfo.newOnceGlobalsSetInFuncs[func->name] =
       std::move(onceGlobalsWrittenVec[0]);
+
+    if (optInfo.onceFuncs.at(func->name).is()) {
+      // This is a "once" function, that is, it has this shape:
+      //
+      //  function foo() {
+      //    if (!foo$once) return;
+      //    foo$once = 1;
+      //    CODE
+      //  }
+      //
+      // If in addition CODE = bar() (a call to another function, and nothing
+      // else) then we can potentially optimize further here. We know that this
+      // function has either been called before, in which case it exits, or it
+      // has not, in which case it sets the global and calls bar. We can
+      // therefore assume that bar has been entered before foo exits, because
+      // there are only two cases:
+      //
+      //  1. bar is on the stack, that is, bar called foo. Then bar has already
+      //     been entered even before we are reached.
+      //  2. bar is not on the stack. In this case, if foo exits immediately
+      //     then we executed the call to bar before; and if not, then we enter
+      //     bar from here.
+      //
+      // Note that we cannot assume bar will have *fully executed* before we do;
+      // only that it was entered. That is because of the first case, where bar
+      // is on the stack. Imagine that bar is also a "once" function, then this
+      // happens:
+      //
+      //  a. bar is called.
+      //  b. bar sets its global to 1.
+      //  c. bar calls foo
+      //  d. foo calls bar, which immediately exits (since bar's global is 1).
+      //
+      // When foo exits after step (d) we are in a situation where foo has
+      // exited and bar has only been entered; it also early-exited, but it did
+      // not execute its body. The distinction between whether it was only
+      // entered or whether it executed fully can matter in situations like
+      // this:
+      //
+      //  function A() {
+      //    if (!A$once) return;
+      //    A$once = 1;
+      //    B(); // this calls A and C
+      //    C(); // can this be removed?
+      //  }
+      //
+      // Imagine that we called B which calls A. Inside A, the call to B will
+      // exit immediately (as shown in steps a-d above), and back in A we will
+      // continue to call C. Naively, we might think that we have a call to B
+      // here, and B calls C, so we can assume C has been called, and can remove
+      // the call to C from A, but that is false: B has been entered, and then
+      // early-exited, but it has not had a chance to fully execute yet, so its
+      // call to C has not happened. If we removed the call to C here, we could
+      // end up with a noticeable change to runtime execution, as then C would
+      // end up called from B after A returns to its caller.
+      //
+      // More cases are optimizable here than what we handle, which is just a
+      // "once" function that calls another and does nothing else at all, but
+      // this is a common case in Java-style code, and anything more would add
+      // significant complexity, as the paragraphs above show.
+      auto& list = func->body->cast<Block>()->list;
+      if (list.size() == 3) {
+        if (auto* call = list[2]->dynCast<Call>()) {
+          if (optInfo.onceFuncs.at(call->target).is()) {
+            // This other "once" global will be entered before we exit, so we
+            // can assume it will be set. Due to the concerns above we cannot
+            // assume anything about that function fully executing, only that it
+            // was entered (and if it is entered, the global will be set, at
+            // least).
+            optInfo.newOnceGlobalsSetInFuncs[func->name].insert(call->target);
+          }
+        }
+      }
+    }
   }
 
 private: