From 1518624938da81e8ebc11f8f56f7e7fafd774add Mon Sep 17 00:00:00 2001
From: Dan Gohman <dan433584@gmail.com>
Date: Thu, 7 Dec 2017 13:49:27 +0000
Subject: [PATCH] [WebAssemby] Support main functions with alternate
 signatures.

WebAssembly requires caller and callee signatures to match, so the usual
C runtime trick of calling main and having it just work regardless of
whether main is defined as '()' or '(int argc, char *argv[])' doesn't
work. Extend the FixFunctionBitcasts pass to rewrite main to use the
latter form.

llvm-svn: 320041
---
 .../WebAssemblyFixFunctionBitcasts.cpp        | 46 ++++++++++++++++++-
 test/CodeGen/WebAssembly/main.ll              | 18 ++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 test/CodeGen/WebAssembly/main.ll

diff --git a/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp b/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
index f1d2c00d163..4888bf3a747 100644
--- a/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
+++ b/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
@@ -142,11 +142,41 @@ static Function *CreateWrapper(Function *F, FunctionType *Ty) {
 }
 
 bool FixFunctionBitcasts::runOnModule(Module &M) {
+  Function *Main = nullptr;
+  CallInst *CallMain = nullptr;
   SmallVector<std::pair<Use *, Function *>, 0> Uses;
   SmallPtrSet<Constant *, 2> ConstantBCs;
 
   // Collect all the places that need wrappers.
-  for (Function &F : M) FindUses(&F, F, Uses, ConstantBCs);
+  for (Function &F : M) {
+    FindUses(&F, F, Uses, ConstantBCs);
+
+    // If we have a "main" function, and its type isn't
+    // "int main(int argc, char *argv[])", create an artificial call with it
+    // bitcasted to that type so that we generate a wrapper for it, so that
+    // the C runtime can call it.
+    if (F.getName() == "main") {
+      Main = &F;
+      LLVMContext &C = M.getContext();
+      Type *MainArgTys[] = {
+        PointerType::get(Type::getInt8PtrTy(C), 0),
+        Type::getInt32Ty(C)
+      };
+      FunctionType *MainTy = FunctionType::get(Type::getInt32Ty(C), MainArgTys,
+                                               /*isVarArg=*/false);
+      if (F.getFunctionType() != MainTy) {
+        Value *Args[] = {
+          UndefValue::get(MainArgTys[0]),
+          UndefValue::get(MainArgTys[1])
+        };
+        Value *Casted = ConstantExpr::getBitCast(Main,
+                                                 PointerType::get(MainTy, 0));
+        CallMain = CallInst::Create(Casted, Args, "call_main");
+        Use *UseMain = &CallMain->getOperandUse(2);
+        Uses.push_back(std::make_pair(UseMain, &F));
+      }
+    }
+  }
 
   DenseMap<std::pair<Function *, FunctionType *>, Function *> Wrappers;
 
@@ -181,5 +211,19 @@ bool FixFunctionBitcasts::runOnModule(Module &M) {
       U->set(Wrapper);
   }
 
+  // If we created a wrapper for main, rename the wrapper so that it's the
+  // one that gets called from startup.
+  if (CallMain) {
+    Main->setName("__original_main");
+    Function *MainWrapper =
+        cast<Function>(CallMain->getCalledValue()->stripPointerCasts());
+    MainWrapper->setName("main");
+    MainWrapper->setLinkage(Main->getLinkage());
+    MainWrapper->setVisibility(Main->getVisibility());
+    Main->setLinkage(Function::PrivateLinkage);
+    Main->setVisibility(Function::DefaultVisibility);
+    delete CallMain;
+  }
+
   return true;
 }
diff --git a/test/CodeGen/WebAssembly/main.ll b/test/CodeGen/WebAssembly/main.ll
new file mode 100644
index 00000000000..314dde28307
--- /dev/null
+++ b/test/CodeGen/WebAssembly/main.ll
@@ -0,0 +1,18 @@
+; RUN: llc < %s -asm-verbose=false | FileCheck %s
+
+; Test main functions with alternate signatures.
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown-wasm"
+
+define void @main() {
+  ret void
+}
+
+; CHECK-LABEL: .L__original_main:
+; CHECK-NEXT: end_function
+
+; CHECK-LABEL: main:
+; CHECK-NEXT: .param i32, i32
+; CHECK-NEXT: .result i32
+; CHECK:      call .L__original_main@FUNCTION