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,48 @@ 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() then we can optimize further here: this
+      // function does nothing but set the global and forward the call onwards,
+      // which is simple to reason about. To see why, consider that it is very
+      // different from this situation:
+      //
+      //  function foo() {
+      //    if (foo$once) return;
+      //    foo$once = 1;
+      //    bar();
+      //    baz();
+      //  }
+      //
+      // Now there are two functions there, but we cannot assume that both have
+      // been called by the time that foo exits: imagine that we called bar
+      // which calls foo, that is, bar is on the stack when we get to foo; then
+      // if bar is an "only" function as well, it will early-exit (since it set
+      // its global), and any setting of baz in its body has not yet been
+      // reached.
+      //
+      // In comparison, when there is just one call at the end, then we know
+      // that if it is "once" function then its global will be set - we can't
+      // tell about anything in its body, but at least the global is known
+      // (since the global is set right at the top of such a function).
+      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()) {
+            optInfo.newOnceGlobalsSetInFuncs[func->name].insert(call->target);
+          }
+        }
+      }
+    }
   }
 
 private: