Add a pass to inline __original_main() into main() (#2461)

clang/llvm introduce __original_main as a workaround for
the fact that main may have different signatures. A downside
to that is that users get it in stack traces, which is confusing.

In -O2 and above we normally inline __original_main anyhow,
but as this is for debugging, non-optimized builds matter too,
so add a pass for this.

The implementation is trivial, just call doInling. However we
must check some corner cases first.

Bonus minor fixes to FindAllPointers, which unnecessarily
created an object to get the class Id (which is not valid
for all classes), and that it didn't take the input by
reference properly, which meant we couldn't get the
pointer to the function body's toplevel.

4 files changed, 44 insertions, 3 deletions

diff --git a/src/ir/find_all.h b/src/ir/find_all.h
index 1abaab772..bd3d265e2 100644
--- a/src/ir/find_all.h
+++ b/src/ir/find_all.h
@@ -58,9 +58,11 @@ struct PointerFinder
 template<typename T> struct FindAllPointers {
   std::vector<Expression**> list;
 
-  FindAllPointers(Expression* ast) {
+  // Note that a pointer may be to the function->body itself, so we must
+  // take \ast by reference.
+  FindAllPointers(Expression*& ast) {
     PointerFinder finder;
-    finder.id = T()._id;
+    finder.id = (Expression::Id)T::SpecificId;
     finder.list = &list;
     finder.walk(ast);
   }
diff --git a/src/passes/Inlining.cpp b/src/passes/Inlining.cpp
index 9ba01c4c1..53a57cfeb 100644
--- a/src/passes/Inlining.cpp
+++ b/src/passes/Inlining.cpp
@@ -217,7 +217,7 @@ struct Updater : public PostWalker<Updater> {
 // Core inlining logic. Modifies the outside function (adding locals as
 // needed), and returns the inlined code.
 static Expression*
-doInlining(Module* module, Function* into, InliningAction& action) {
+doInlining(Module* module, Function* into, const InliningAction& action) {
   Function* from = action.contents;
   auto* call = (*action.callSite)->cast<Call>();
   // Works for return_call, too
@@ -415,4 +415,40 @@ Pass* createInliningOptimizingPass() {
   return ret;
 }
 
+static const char* MAIN = "main";
+static const char* ORIGINAL_MAIN = "__original_main";
+
+// Inline __original_main into main, if they exist. This works around the odd
+// thing that clang/llvm currently do, where __original_main contains the user's
+// actual main (this is done as a workaround for main having two different
+// possible signatures).
+struct InlineMainPass : public Pass {
+  void run(PassRunner* runner, Module* module) override {
+    auto* main = module->getFunctionOrNull(MAIN);
+    auto* originalMain = module->getFunctionOrNull(ORIGINAL_MAIN);
+    if (!main || main->imported() || !originalMain ||
+        originalMain->imported()) {
+      return;
+    }
+    FindAllPointers<Call> calls(main->body);
+    Expression** callSite = nullptr;
+    for (auto* call : calls.list) {
+      if ((*call)->cast<Call>()->target == ORIGINAL_MAIN) {
+        if (callSite) {
+          // More than one call site.
+          return;
+        }
+        callSite = call;
+      }
+    }
+    if (!callSite) {
+      // No call at all.
+      return;
+    }
+    doInlining(module, main, InliningAction(callSite, originalMain));
+  }
+};
+
+Pass* createInlineMainPass() { return new InlineMainPass(); }
+
 } // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 9afea38e6..906fb85e2 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -129,6 +129,8 @@ void PassRegistry::registerPasses() {
     "func-metrics", "reports function metrics", createFunctionMetricsPass);
   registerPass(
     "generate-stack-ir", "generate Stack IR", createGenerateStackIRPass);
+  registerPass(
+    "inline-main", "inline __original_main into main", createInlineMainPass);
   registerPass("inlining",
                "inline functions (you probably want inlining-optimizing)",
                createInliningPass);
diff --git a/src/passes/passes.h b/src/passes/passes.h
index cc33e4300..98c7d374d 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -45,6 +45,7 @@ Pass* createFullPrinterPass();
 Pass* createFunctionMetricsPass();
 Pass* createGenerateStackIRPass();
 Pass* createI64ToI32LoweringPass();
+Pass* createInlineMainPass();
 Pass* createInliningPass();
 Pass* createInliningOptimizingPass();
 Pass* createLegalizeJSInterfacePass();