[PECOFF] Fix entry point functions selection

On Windows there are four "main" functions -- main, wmain, WinMain,
or wWinMain. Their parameter types are diffferent. The standard
library provides four different entry functions (i.e.
{w,}{WinMain,main}CRTStartup) for them. You need to use the right
entry routine for your "main" function.

If you give an /entry option, the specified name is used
unconditionally.

Otherwise, the linker needs to select the right one based on
user-supplied entry point function. This can be done after the
linker reads all the input files.

This patch moves the code to determine the entry point function
from the driver to a virtual input file. It also implements the
correct logic for the entry point function selection.

llvm-svn: 213713

11 files changed, 178 insertions, 121 deletions

diff --git a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
index 002edb6477f..224ff072dff 100644
--- a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
+++ b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
@@ -113,10 +113,8 @@ public:
   StringRef decorateSymbol(StringRef name) const;
   StringRef undecorateSymbol(StringRef name) const;
 
-  void setEntrySymbolName(StringRef name) {
-    if (!name.empty())
-      LinkingContext::setEntrySymbolName(decorateSymbol(name));
-  }
+  void setEntrySymbolName(StringRef name) { _entry = name; }
+  StringRef getEntrySymbolName() const { return _entry; }
 
   void setHasEntry(bool val) { _hasEntry = val; }
   bool hasEntry() const { return _hasEntry; }
@@ -308,6 +306,8 @@ private:
   std::recursive_mutex _mutex;
   mutable std::mutex _allocMutex;
 
+  std::string _entry;
+
   // False if /noentry option is given.
   bool _hasEntry;
 
diff --git a/lld/lib/Driver/WinLinkDriver.cpp b/lld/lib/Driver/WinLinkDriver.cpp
index d226aa75f5e..2fa2adb6e40 100644
--- a/lld/lib/Driver/WinLinkDriver.cpp
+++ b/lld/lib/Driver/WinLinkDriver.cpp
@@ -638,22 +638,6 @@ static void processLibEnv(PECOFFLinkingContext &ctx) {
       ctx.appendInputSearchPath(ctx.allocate(path));
 }
 
-// Returns a default entry point symbol name depending on context image type and
-// subsystem. These default names are MS CRT compliant.
-static StringRef getDefaultEntrySymbolName(PECOFFLinkingContext &ctx) {
-  if (ctx.isDll()) {
-    if (ctx.getMachineType() == llvm::COFF::IMAGE_FILE_MACHINE_I386)
-      return "_DllMainCRTStartup@12";
-    return "_DllMainCRTStartup";
-  }
-  llvm::COFF::WindowsSubsystem subsystem = ctx.getSubsystem();
-  if (subsystem == llvm::COFF::WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_GUI)
-    return "WinMainCRTStartup";
-  if (subsystem == llvm::COFF::WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_CUI)
-    return "mainCRTStartup";
-  return "";
-}
-
 namespace {
 class DriverStringSaver : public llvm::cl::StringSaver {
 public:
@@ -1285,13 +1269,6 @@ bool WinLinkDriver::parse(int argc, const char *argv[],
     }
   }
 
-  // Use the default entry name if /entry option is not given.
-  if (ctx.entrySymbolName().empty() && !parsedArgs->getLastArg(OPT_noentry))
-    ctx.setEntrySymbolName(getDefaultEntrySymbolName(ctx));
-  StringRef entry = ctx.entrySymbolName();
-  if (!entry.empty())
-    ctx.addInitialUndefinedSymbol(entry);
-
   // Specify /noentry without /dll is an error.
   if (!ctx.hasEntry() && !parsedArgs->getLastArg(OPT_dll)) {
     diag << "/noentry must be specified with /dll\n";
diff --git a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
index 14576e2695f..9d5c37937f2 100644
--- a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
+++ b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
@@ -16,6 +16,8 @@
 
 #include <mutex>
 
+using llvm::COFF::WindowsSubsystem;
+
 namespace lld {
 namespace pecoff {
 
@@ -274,5 +276,104 @@ private:
   mutable llvm::BumpPtrAllocator _alloc;
 };
 
+// Windows has not only one but many entry point functions. The
+// appropriate one is automatically selected based on the subsystem
+// setting and the user-supplied entry point function.
+//
+// http://msdn.microsoft.com/en-us/library/f9t8842e.aspx
+class EntryPointFile : public SimpleFile {
+public:
+  EntryPointFile(const PECOFFLinkingContext &ctx,
+                 std::shared_ptr<ResolvableSymbols> syms)
+      : SimpleFile("<entry>"), _ctx(const_cast<PECOFFLinkingContext *>(&ctx)),
+        _syms(syms), _firstTime(true) {}
+
+  const atom_collection<UndefinedAtom> &undefined() const override {
+    return const_cast<EntryPointFile *>(this)->getUndefinedAtoms();
+  }
+
+private:
+  const atom_collection<UndefinedAtom> &getUndefinedAtoms() {
+    std::lock_guard<std::mutex> lock(_mutex);
+    if (!_firstTime)
+      return _undefinedAtoms;
+    _firstTime = false;
+
+    if (_ctx->hasEntry()) {
+      StringRef entrySym = _ctx->allocate(_ctx->decorateSymbol(getEntry()));
+      _undefinedAtoms._atoms.push_back(
+          new (_alloc) SimpleUndefinedAtom(*this, entrySym));
+      if (_ctx->deadStrip())
+        _ctx->addDeadStripRoot(entrySym);
+    }
+    return _undefinedAtoms;
+  }
+
+  // Returns the entry point function name. It also sets the inferred
+  // subsystem if it's unknown.
+  std::string getEntry() const {
+    StringRef opt = _ctx->getEntrySymbolName();
+    if (!opt.empty())
+      return opt;
+
+    const std::string wWinMainCRTStartup = "wWinMainCRTStartup";
+    const std::string WinMainCRTStartup = "WinMainCRTStartup";
+    const std::string wmainCRTStartup = "wmainCRTStartup";
+    const std::string mainCRTStartup = "mainCRTStartup";
+    auto windows = WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_GUI;
+    auto console = WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_CUI;
+
+    if (_ctx->isDll()) {
+      _ctx->setSubsystem(windows);
+      if (_ctx->getMachineType() == llvm::COFF::IMAGE_FILE_MACHINE_I386)
+        return "_DllMainCRTStartup@12";
+      return "_DllMainCRTStartup";
+    }
+
+    // Returns true if a given name exists in an input object file.
+    auto defined = [&](StringRef name) -> bool {
+      return _syms->defined().count(_ctx->decorateSymbol(name));
+    };
+
+    switch (_ctx->getSubsystem()) {
+    case WindowsSubsystem::IMAGE_SUBSYSTEM_UNKNOWN: {
+      if (defined("wWinMain")) {
+        _ctx->setSubsystem(windows);
+        return wWinMainCRTStartup;
+      }
+      if (defined("WinMain")) {
+        _ctx->setSubsystem(windows);
+        return WinMainCRTStartup;
+      }
+      if (defined("wmain")) {
+        _ctx->setSubsystem(console);
+        return wmainCRTStartup;
+      }
+      if (!defined("main"))
+        llvm::errs() << "Cannot infer subsystem; assuming /subsystem:console\n";
+      _ctx->setSubsystem(console);
+      return mainCRTStartup;
+    }
+    case WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_GUI:
+      if (defined("WinMain"))
+        return WinMainCRTStartup;
+      return wWinMainCRTStartup;
+    case WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_CUI:
+      if (defined("wmain"))
+        return wmainCRTStartup;
+      return mainCRTStartup;
+    default:
+      return mainCRTStartup;
+    }
+  }
+
+  PECOFFLinkingContext *_ctx;
+  atom_collection_vector<UndefinedAtom> _undefinedAtoms;
+  std::mutex _mutex;
+  std::shared_ptr<ResolvableSymbols> _syms;
+  llvm::BumpPtrAllocator _alloc;
+  bool _firstTime;
+};
+
 } // end namespace pecoff
 } // end namespace lld
diff --git a/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp b/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
index 0ac98693785..94d06775903 100644
--- a/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
@@ -13,7 +13,6 @@
 #include "IdataPass.h"
 #include "LinkerGeneratedSymbolFile.h"
 #include "LoadConfigPass.h"
-#include "SetSubsystemPass.h"
 
 #include "lld/Core/PassManager.h"
 #include "lld/Core/Simple.h"
@@ -128,6 +127,13 @@ bool PECOFFLinkingContext::createImplicitFiles(
   auto *renameFile = new pecoff::ExportedSymbolRenameFile(*this, syms);
   exportNode->appendInputFile(std::unique_ptr<File>(renameFile));
   getLibraryGroup()->addFile(std::move(exportNode));
+
+  // Create a file for the entry point function.
+  std::unique_ptr<SimpleFileNode> entryFileNode(new SimpleFileNode("<entry>"));
+  entryFileNode->appendInputFile(
+      std::unique_ptr<File>(new pecoff::EntryPointFile(*this, syms)));
+  getInputGraph().insertElementAt(std::move(entryFileNode),
+                                  InputGraph::Position::END);
   return true;
 }
 
@@ -296,7 +302,6 @@ std::string PECOFFLinkingContext::getOutputImportLibraryPath() const {
 }
 
 void PECOFFLinkingContext::addPasses(PassManager &pm) {
-  pm.add(std::unique_ptr<Pass>(new pecoff::SetSubsystemPass(*this)));
   pm.add(std::unique_ptr<Pass>(new pecoff::EdataPass(*this)));
   pm.add(std::unique_ptr<Pass>(new pecoff::IdataPass(*this)));
   pm.add(std::unique_ptr<Pass>(new LayoutPass(registry())));
diff --git a/lld/lib/ReaderWriter/PECOFF/SetSubsystemPass.h b/lld/lib/ReaderWriter/PECOFF/SetSubsystemPass.h
deleted file mode 100644
index bf5721c951c..00000000000
--- a/lld/lib/ReaderWriter/PECOFF/SetSubsystemPass.h
+++ /dev/null
@@ -1,58 +0,0 @@
-//===- lib/ReaderWriter/PECOFF/SetSubsystemPass.h -------------------------===//
-//
-//                             The LLVM Linker
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef LLD_READER_WRITER_PE_COFF_SET_SUBSYSTEM_PASS_H
-#define LLD_READER_WRITER_PE_COFF_SET_SUBSYSTEM_PASS_H
-
-#include "lld/ReaderWriter/PECOFFLinkingContext.h"
-
-using llvm::COFF::WindowsSubsystem;
-
-namespace lld {
-namespace pecoff {
-
-/// If "main" or "wmain" is defined, /subsystem:console is the default. If
-/// "WinMain" or "wWinMain" is defined, /subsystem:windows is the default.
-class SetSubsystemPass : public lld::Pass {
-public:
-  SetSubsystemPass(PECOFFLinkingContext &ctx) : _ctx(ctx) {}
-
-  void perform(std::unique_ptr<MutableFile> &file) override {
-    if (_ctx.getSubsystem() != WindowsSubsystem::IMAGE_SUBSYSTEM_UNKNOWN)
-      return;
-    StringRef main = _ctx.decorateSymbol("main");
-    StringRef wmain = _ctx.decorateSymbol("wmain");
-    StringRef winmain = _ctx.decorateSymbol("WinMain");
-    StringRef wwinmain = _ctx.decorateSymbol("wWinMain");
-    for (auto *atom : file->defined()) {
-      StringRef s = atom->name();
-      if (s == main || s == wmain) {
-        _ctx.setSubsystem(WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_CUI);
-        return;
-      }
-      if (s == winmain || s == wwinmain) {
-        _ctx.setSubsystem(WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_GUI);
-        return;
-      }
-    }
-    if (_ctx.isDll()) {
-      _ctx.setSubsystem(WindowsSubsystem::IMAGE_SUBSYSTEM_WINDOWS_GUI);
-      return;
-    }
-    llvm_unreachable("Failed to infer the subsystem.");
-  }
-
-private:
-  PECOFFLinkingContext &_ctx;
-};
-
-} // namespace pecoff
-} // namespace lld
-
-#endif
diff --git a/lld/test/pecoff/Inputs/entry.obj.yaml b/lld/test/pecoff/Inputs/entry.obj.yaml
new file mode 100644
index 00000000000..595ea18fc9f
--- /dev/null
+++ b/lld/test/pecoff/Inputs/entry.obj.yaml
@@ -0,0 +1,26 @@
+---
+header:
+  Machine:         IMAGE_FILE_MACHINE_I386
+  Characteristics: [  ]
+
+sections:
+  - Name:            .text
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       4
+    SectionData:     A100000000030500000000C3
+
+symbols:
+  - Name:            .text
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+
+  - Name:            _foo
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...
diff --git a/lld/test/pecoff/Inputs/subsystem.main.yaml b/lld/test/pecoff/Inputs/subsystem.main.yaml
index 01cb63bb4b0..25fbe1be014 100644
--- a/lld/test/pecoff/Inputs/subsystem.main.yaml
+++ b/lld/test/pecoff/Inputs/subsystem.main.yaml
@@ -26,4 +26,10 @@ symbols:
     SimpleType:      IMAGE_SYM_TYPE_NULL
     ComplexType:     IMAGE_SYM_DTYPE_NULL
     StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            _mainCRTStartup
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
 ...
diff --git a/lld/test/pecoff/Inputs/subsystem.winmain.yaml b/lld/test/pecoff/Inputs/subsystem.winmain.yaml
index b3b1fd62066..7f0381e2dc6 100644
--- a/lld/test/pecoff/Inputs/subsystem.winmain.yaml
+++ b/lld/test/pecoff/Inputs/subsystem.winmain.yaml
@@ -26,4 +26,10 @@ symbols:
     SimpleType:      IMAGE_SYM_TYPE_NULL
     ComplexType:     IMAGE_SYM_DTYPE_NULL
     StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            _WinMainCRTStartup
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
 ...
diff --git a/lld/test/pecoff/alternatename.test b/lld/test/pecoff/alternatename.test
index eb0dda2abeb..ff6617ebdd6 100644
--- a/lld/test/pecoff/alternatename.test
+++ b/lld/test/pecoff/alternatename.test
@@ -2,13 +2,13 @@
 # RUN: yaml2obj %p/Inputs/alternatename2.obj.yaml > %t2.obj
 # RUN: yaml2obj %p/Inputs/alternatename3.obj.yaml > %t3.obj
 #
-# RUN: lld -flavor link /out:%t1.exe /alternatename:_main=_foo -- %t1.obj
+# RUN: lld -flavor link /force /out:%t1.exe /alternatename:_main=_foo -- %t1.obj
 # RUN: llvm-objdump -d %t1.exe | FileCheck -check-prefix=CHECK1 %s
 #
-# RUN: lld -flavor link /out:%t2.exe /alternatename:_main=_foo  -- %t1.obj %t2.obj
+# RUN: lld -flavor link /force /out:%t2.exe /alternatename:_main=_foo  -- %t1.obj %t2.obj
 # RUN: llvm-objdump -d %t2.exe | FileCheck -check-prefix=CHECK2 %s
 #
-# RUN: lld -flavor link /out:%t3.exe -- %t3.obj
+# RUN: lld -flavor link /force /out:%t3.exe -- %t3.obj
 # RUN: llvm-objdump -d %t3.exe | FileCheck -check-prefix=CHECK3 %s
 
 CHECK1:      nop
diff --git a/lld/test/pecoff/entry.test b/lld/test/pecoff/entry.test
index c3af25c65e3..404cf431ea5 100644
--- a/lld/test/pecoff/entry.test
+++ b/lld/test/pecoff/entry.test
@@ -1,9 +1,27 @@
 # REQUIRES: asserts
-# Verify that entry atom will not be dead-stripped.
 
-# RUN: yaml2obj %p/Inputs/main.obj.yaml > %t.obj
-# RUN: lld -flavor link /mllvm:-debug-only=WriterPECOFF /out:%t1.exe \
-# RUN:   /subsystem:console /entry:main /force -- %t.obj >& %t1.log
-# RUN: FileCheck -check-prefix=CHECK %s < %t1.log
+# RUN: yaml2obj %p/Inputs/entry.obj.yaml > %t.obj
 
-CHECK: : _main
+# RUN: not lld -flavor link /out:%t.exe /alternatename:_main=_foo \
+# RUN:   -- %t.obj 2> %t.log
+# RUN: FileCheck -check-prefix=MAIN %s < %t.log
+
+MAIN: _mainCRTStartup
+
+# RUN: not lld -flavor link /out:%t.exe /alternatename:_wmain=_foo \
+# RUN:   -- %t.obj 2> %t.log
+# RUN: FileCheck -check-prefix=WMAIN %s < %t.log
+
+WMAIN: _wmainCRTStartup
+
+# RUN: not lld -flavor link /out:%t.exe /alternatename:_WinMain=_foo \
+# RUN:   -- %t.obj 2> %t.log
+# RUN: FileCheck -check-prefix=WINMAIN %s < %t.log
+
+WINMAIN: _WinMainCRTStartup
+
+# RUN: not lld -flavor link /out:%t.exe /alternatename:_wWinMain=_foo \
+# RUN:   -- %t.obj 2> %t.log
+# RUN: FileCheck -check-prefix=WWINMAIN %s < %t.log
+
+WWINMAIN: _wWinMainCRTStartup
diff --git a/lld/unittests/DriverTests/WinLinkDriverTest.cpp b/lld/unittests/DriverTests/WinLinkDriverTest.cpp
index 2246d82856e..2f74da790d4 100644
--- a/lld/unittests/DriverTests/WinLinkDriverTest.cpp
+++ b/lld/unittests/DriverTests/WinLinkDriverTest.cpp
@@ -38,7 +38,7 @@ TEST_F(WinLinkParserTest, Basic) {
   EXPECT_EQ(llvm::COFF::IMAGE_SUBSYSTEM_WINDOWS_CUI, _context.getSubsystem());
   EXPECT_EQ(llvm::COFF::IMAGE_FILE_MACHINE_I386, _context.getMachineType());
   EXPECT_EQ("a.exe", _context.outputPath());
-  EXPECT_EQ("_start", _context.entrySymbolName());
+  EXPECT_EQ("start", _context.getEntrySymbolName());
   EXPECT_EQ(4, inputFileCount());
   EXPECT_EQ("a.obj", inputFile(0));
   EXPECT_EQ("b.obj", inputFile(1));
@@ -684,27 +684,3 @@ TEST_F(WinLinkParserTest, DashDash) {
   EXPECT_EQ("b.obj", inputFile(1));
   EXPECT_EQ("-c.obj", inputFile(2));
 }
-
-//
-// Tests for entry symbol name.
-//
-
-TEST_F(WinLinkParserTest, DefEntryNameConsole) {
-  EXPECT_TRUE(parse("link.exe", "/subsystem:console", "a.obj", nullptr));
-  EXPECT_EQ("_mainCRTStartup", _context.entrySymbolName());
-}
-
-TEST_F(WinLinkParserTest, DefEntryNameWindows) {
-  EXPECT_TRUE(parse("link.exe", "/subsystem:windows", "a.obj", nullptr));
-  EXPECT_EQ("_WinMainCRTStartup", _context.entrySymbolName());
-}
-
-TEST_F(WinLinkParserTest, DefEntryNameDll32) {
-  EXPECT_TRUE(parse("link.exe", "/dll", "/machine:x86", "a.obj", nullptr));
-  EXPECT_EQ("__DllMainCRTStartup@12", _context.entrySymbolName());
-}
-
-TEST_F(WinLinkParserTest, DefEntryNameDll64) {
-  EXPECT_TRUE(parse("link.exe", "/dll", "/machine:x64", "a.obj", nullptr));
-  EXPECT_EQ("_DllMainCRTStartup", _context.entrySymbolName());
-}