[WebAssembly] Re-enable main-function signature rewriting

Re-enable the code to rewrite main-function signatures into
"int main(int argc, char *argv[])", but limited to only handling
the case of "int main(void)", so that it doesn't silently strip
an argument in the "int main(int argc, char *argv[], char *envp[])"
case.

This allows main to be called by C startup code, since WebAssembly
requires caller and callee signatures to match, so it can't rely
on passing main a different number of arguments than it expects.

Differential Revision: https://reviews.llvm.org/D57323

llvm-svn: 352479

8 files changed, 59 insertions, 32 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
index a822ac1d9b1..e2ef4e0a81c 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp
@@ -35,11 +35,6 @@ 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 {
@@ -226,6 +221,17 @@ static Function *CreateWrapper(Function *F, FunctionType *Ty) {
   return Wrapper;
 }
 
+// Test whether a main function with type FuncTy should be rewritten to have
+// type MainTy.
+bool ShouldFixMainFunction(FunctionType *FuncTy, FunctionType *MainTy) {
+  // Only fix the main function if it's the standard zero-arg form. That way,
+  // the standard cases will work as expected, and users will see signature
+  // mismatches from the linker for non-standard cases.
+  return FuncTy->getReturnType() == MainTy->getReturnType() &&
+         FuncTy->getNumParams() == 0 &&
+         !FuncTy->isVarArg();
+}
+
 bool FixFunctionBitcasts::runOnModule(Module &M) {
   LLVM_DEBUG(dbgs() << "********** Fix Function Bitcasts **********\n");
 
@@ -242,14 +248,14 @@ bool FixFunctionBitcasts::runOnModule(Module &M) {
     // "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") {
+    if (F.getName() == "main") {
       Main = &F;
       LLVMContext &C = M.getContext();
       Type *MainArgTys[] = {Type::getInt32Ty(C),
                             PointerType::get(Type::getInt8PtrTy(C), 0)};
       FunctionType *MainTy = FunctionType::get(Type::getInt32Ty(C), MainArgTys,
                                                /*isVarArg=*/false);
-      if (F.getFunctionType() != MainTy) {
+      if (ShouldFixMainFunction(F.getFunctionType(), MainTy)) {
         LLVM_DEBUG(dbgs() << "Found `main` function with incorrect type: "
                           << *F.getFunctionType() << "\n");
         Value *Args[] = {UndefValue::get(MainArgTys[0]),
@@ -297,12 +303,18 @@ bool FixFunctionBitcasts::runOnModule(Module &M) {
     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;
+    if (Main->isDeclaration()) {
+      // The wrapper is not needed in this case as we don't need to export
+      // it to anyone else.
+      MainWrapper->eraseFromParent();
+    } else {
+      // Otherwise give the wrapper the same linkage as the original main
+      // function, so that it can be called from the same places.
+      MainWrapper->setName("main");
+      MainWrapper->setLinkage(Main->getLinkage());
+      MainWrapper->setVisibility(Main->getVisibility());
+    }
   }
 
   return true;
diff --git a/llvm/test/CodeGen/WebAssembly/call.ll b/llvm/test/CodeGen/WebAssembly/call.ll
index db666a6c366..77f17c850ed 100644
--- a/llvm/test/CodeGen/WebAssembly/call.ll
+++ b/llvm/test/CodeGen/WebAssembly/call.ll
@@ -1,5 +1,5 @@
-; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-keep-registers -wasm-temporary-workarounds=false -mattr=+sign-ext,+simd128 | FileCheck %s
-; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-keep-registers -fast-isel -fast-isel-abort=1 -wasm-temporary-workarounds=false -mattr=+sign-ext,+simd128 | FileCheck %s
+; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-keep-registers -mattr=+sign-ext,+simd128 | FileCheck %s
+; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-keep-registers -fast-isel -fast-isel-abort=1 -mattr=+sign-ext,+simd128 | FileCheck %s
 
 ; Test that basic call operations assemble as expected.
 
diff --git a/llvm/test/CodeGen/WebAssembly/function-bitcasts-varargs.ll b/llvm/test/CodeGen/WebAssembly/function-bitcasts-varargs.ll
index 515c5703d86..b542276e068 100644
--- a/llvm/test/CodeGen/WebAssembly/function-bitcasts-varargs.ll
+++ b/llvm/test/CodeGen/WebAssembly/function-bitcasts-varargs.ll
@@ -1,4 +1,4 @@
-; RUN: llc < %s -asm-verbose=false -wasm-temporary-workarounds=false -wasm-keep-registers | FileCheck %s
+; RUN: llc < %s -asm-verbose=false -wasm-keep-registers | FileCheck %s
 
 ; Test that function pointer casts casting away varargs are replaced with
 ; wrappers.
diff --git a/llvm/test/CodeGen/WebAssembly/function-bitcasts.ll b/llvm/test/CodeGen/WebAssembly/function-bitcasts.ll
index a779cbe4142..813e8420ae5 100644
--- a/llvm/test/CodeGen/WebAssembly/function-bitcasts.ll
+++ b/llvm/test/CodeGen/WebAssembly/function-bitcasts.ll
@@ -1,4 +1,4 @@
-; RUN: llc < %s -asm-verbose=false -wasm-disable-explicit-locals -wasm-keep-registers -enable-emscripten-cxx-exceptions -wasm-temporary-workarounds=false | FileCheck %s
+; RUN: llc < %s -asm-verbose=false -wasm-disable-explicit-locals -wasm-keep-registers -enable-emscripten-cxx-exceptions | FileCheck %s
 
 ; Test that function pointer casts are replaced with wrappers.
 
diff --git a/llvm/test/CodeGen/WebAssembly/main-declaration.ll b/llvm/test/CodeGen/WebAssembly/main-declaration.ll
index f9d68db2bae..544f5588c50 100644
--- a/llvm/test/CodeGen/WebAssembly/main-declaration.ll
+++ b/llvm/test/CodeGen/WebAssembly/main-declaration.ll
@@ -1,20 +1,18 @@
-; RUN: llc < %s -asm-verbose=false -wasm-temporary-workarounds=false | FileCheck %s
+; 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"
 
-declare void @main()
+declare i32 @main()
 
-define void @foo() {
-  call void @main()
-  ret void
+define i32 @foo() {
+  %t = call i32 @main()
+  ret i32 %t
 }
 
-; CHECK-NOT:   __original_main
 ; CHECK-LABEL: foo:
-; CHECK-NEXT:    .functype foo () -> ()
-; CHECK-NEXT:    call main@FUNCTION
+; CHECK-NEXT:    .functype foo () -> (i32)
+; CHECK-NEXT:    call __original_main@FUNCTION
 ; CHECK-NEXT:    end_function
-; CHECK-NOT:   __original_main
diff --git a/llvm/test/CodeGen/WebAssembly/main-no-args.ll b/llvm/test/CodeGen/WebAssembly/main-no-args.ll
index 0bc46717d97..97023e26945 100644
--- a/llvm/test/CodeGen/WebAssembly/main-no-args.ll
+++ b/llvm/test/CodeGen/WebAssembly/main-no-args.ll
@@ -1,18 +1,19 @@
-; RUN: llc < %s -asm-verbose=false -wasm-temporary-workarounds=false | FileCheck %s
+; 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"
 
-define void @main() {
-  ret void
+define i32 @main() {
+  ret i32 0
 }
 
-; CHECK-LABEL: .L__original_main:
-; CHECK-NEXT: .functype .L__original_main () -> ()
+; CHECK-LABEL: __original_main:
+; CHECK-NEXT: .functype __original_main () -> (i32)
+; CHECK-NEXT: i32.const 0
 ; CHECK-NEXT: end_function
 
 ; CHECK-LABEL: main:
 ; CHECK-NEXT: .functype main (i32, i32) -> (i32)
-; CHECK:      call .L__original_main@FUNCTION
+; CHECK:      call __original_main@FUNCTION
diff --git a/llvm/test/CodeGen/WebAssembly/main-three-args.ll b/llvm/test/CodeGen/WebAssembly/main-three-args.ll
new file mode 100644
index 00000000000..77b3e5b8c30
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/main-three-args.ll
@@ -0,0 +1,16 @@
+; RUN: llc < %s -asm-verbose=false | FileCheck %s
+
+; Test that main function with a non-standard third argument is
+; not wrapped.
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+define i32 @main(i32 %a, i8** %b, i8** %c) {
+  ret i32 0
+}
+
+; CHECK-LABEL: main:
+; CHECK-NEXT: .functype main (i32, i32, i32) -> (i32)
+
+; CHECK-NOT: __original_main:
diff --git a/llvm/test/CodeGen/WebAssembly/main-with-args.ll b/llvm/test/CodeGen/WebAssembly/main-with-args.ll
index d4a11ef14d4..205cb133f8c 100644
--- a/llvm/test/CodeGen/WebAssembly/main-with-args.ll
+++ b/llvm/test/CodeGen/WebAssembly/main-with-args.ll
@@ -1,4 +1,4 @@
-; RUN: llc < %s -asm-verbose=false -wasm-temporary-workarounds=false | FileCheck %s
+; RUN: llc < %s -asm-verbose=false | FileCheck %s
 
 ; Test that main function with expected signature is not wrapped