[PECOFF] Fix entry point address.

Because of a bug, the entry point address in the PE/COFF header
was not correct.

llvm-svn: 213802

2 files changed, 6 insertions, 4 deletions

diff --git a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
index 9d5c37937f2..57aae835c25 100644
--- a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
+++ b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
@@ -303,6 +303,8 @@ private:
       StringRef entrySym = _ctx->allocate(_ctx->decorateSymbol(getEntry()));
       _undefinedAtoms._atoms.push_back(
           new (_alloc) SimpleUndefinedAtom(*this, entrySym));
+      _ctx->setHasEntry(true);
+      _ctx->setEntrySymbolName(entrySym);
       if (_ctx->deadStrip())
         _ctx->addDeadStripRoot(entrySym);
     }
diff --git a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
index 73fa4f7195e..22dd2f01ca9 100644
--- a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
@@ -968,14 +968,14 @@ void PECOFFWriter::build(const File &linkedFile) {
       // Find the virtual address of the entry point symbol if any.  PECOFF spec
       // says that entry point for dll images is optional, in which case it must
       // be set to 0.
-      if (_ctx.entrySymbolName().empty() && _ctx.isDll()) {
-        peHeader->setAddressOfEntryPoint(0);
-      } else {
+      if (_ctx.hasEntry()) {
         uint64_t entryPointAddress =
             dyn_cast<AtomChunk>(section)
-                ->getAtomVirtualAddress(_ctx.entrySymbolName());
+                ->getAtomVirtualAddress(_ctx.getEntrySymbolName());
         if (entryPointAddress != 0)
           peHeader->setAddressOfEntryPoint(entryPointAddress);
+      } else {
+        peHeader->setAddressOfEntryPoint(0);
       }
     }
     if (section->getSectionName() == ".data")