[WebAssemby] Re-apply r320041: "Support main functions with alternate signatures."

This includes a fix so that it doesn't transform declarations, and it
puts the functionality under control of a command-line option which is off
by default to avoid breaking existing setups.

llvm-svn: 320196

3 files changed, 87 insertions, 1 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
index f1d2c00d163..084576b2b17 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
@@ -36,6 +36,11 @@ using namespace llvm;
 
 #define DEBUG_TYPE "wasm-fix-function-bitcasts"
 
+static cl::opt<bool> TemporaryWorkarounds(
+  "wasm-temporary-workarounds",
+  cl::desc("Apply certain temporary workarounds"),
+  cl::init(true), cl::Hidden);
+
 namespace {
 class FixFunctionBitcasts final : public ModulePass {
   StringRef getPassName() const override {
@@ -142,11 +147,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 (!TemporaryWorkarounds && !F.isDeclaration() && 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 +216,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/llvm/test/CodeGen/WebAssembly/main-declaration.ll b/llvm/test/CodeGen/WebAssembly/main-declaration.ll
new file mode 100644
index 00000000000..4e337850b81
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/main-declaration.ll
@@ -0,0 +1,19 @@
+; RUN: llc < %s -asm-verbose=false -wasm-temporary-workarounds=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"
+
+declare void @main()
+
+define void @foo() {
+  call void @main()
+  ret void
+}
+
+; CHECK-NOT:   __original_main
+; CHECK-LABEL: foo:
+; CHECK-NEXT:    call main@FUNCTION
+; CHECK-NEXT:    end_function
+; CHECK-NOT:   __original_main
diff --git a/llvm/test/CodeGen/WebAssembly/main.ll b/llvm/test/CodeGen/WebAssembly/main.ll
new file mode 100644
index 00000000000..c77db8467d8
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/main.ll
@@ -0,0 +1,18 @@
+; RUN: llc < %s -asm-verbose=false -wasm-temporary-workarounds=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