COFF: /include'd symbols must be preserved.

Not only entry point symbol but also symbols specified by /include
option must be preserved, as they will never be dead-stripped.

http://reviews.llvm.org/D10220

llvm-svn: 239005

6 files changed, 47 insertions, 16 deletions

diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h
index afe5670..66ffa65 100644
--- a/lld/COFF/Config.h
+++ b/lld/COFF/Config.h
@@ -27,7 +27,10 @@ public:
   llvm::COFF::MachineTypes MachineType = llvm::COFF::IMAGE_FILE_MACHINE_AMD64;
   bool Verbose = false;
   WindowsSubsystem Subsystem = llvm::COFF::IMAGE_SUBSYSTEM_UNKNOWN;
-  std::string EntryName;
+  StringRef EntryName;
+
+  // Symbols in this set are considered as live by the garbage collector.
+  std::set<StringRef> GCRoots;
 
   std::set<StringRef> NoDefaultLibs;
   bool NoDefaultLibAll = false;
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 8e972a6..8a91865 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -312,8 +312,11 @@ bool LinkerDriver::link(int Argc, const char *Argv[]) {
 
   // Add undefined symbols given via the command line.
   // (/include is equivalent to Unix linker's -u option.)
-  for (auto *Arg : Args->filtered(OPT_incl))
-    Symtab.addUndefined(Arg->getValue());
+  for (auto *Arg : Args->filtered(OPT_incl)) {
+    StringRef Sym = Arg->getValue();
+    Symtab.addUndefined(Sym);
+    Config->GCRoots.insert(Sym);
+  }
 
   // Parse all input files and put all symbols to the symbol table.
   // The symbol table will take care of name resolution.
@@ -362,11 +365,14 @@ bool LinkerDriver::link(int Argc, const char *Argv[]) {
     }
     Config->EntryName = EntryOrErr.get();
   }
+  Config->GCRoots.insert(Config->EntryName);
 
   // Make sure we have resolved all symbols.
   if (Symtab.reportRemainingUndefines())
     return false;
 
+  // Do LTO by compiling bitcode input files to a native COFF file
+  // then link that file.
   if (auto EC = Symtab.addCombinedLTOObject()) {
     llvm::errs() << EC.message() << "\n";
     return false;
diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
index 015d707..35042e5 100644
--- a/lld/COFF/SymbolTable.cpp
+++ b/lld/COFF/SymbolTable.cpp
@@ -227,20 +227,17 @@ std::error_code SymbolTable::addCombinedLTOObject() {
     return std::error_code();
 
   llvm::LTOCodeGenerator CG;
-  std::set<DefinedBitcode *> PreservedBitcodeSymbols;
 
   // All symbols referenced by non-bitcode objects must be preserved.
   for (std::unique_ptr<ObjectFile> &File : ObjectFiles)
     for (SymbolBody *Body : File->getSymbols())
       if (auto *S = dyn_cast<DefinedBitcode>(Body->getReplacement()))
-        PreservedBitcodeSymbols.insert(S);
+        CG.addMustPreserveSymbol(S->getName());
 
-  // Likewise for the linker-generated reference to the entry point.
-  if (auto *S = dyn_cast<DefinedBitcode>(Symtab[Config->EntryName]->Body))
-    PreservedBitcodeSymbols.insert(S);
-
-  for (DefinedBitcode *S : PreservedBitcodeSymbols)
-    CG.addMustPreserveSymbol(S->getName());
+  // Likewise for other symbols that must be preserved.
+  for (StringRef Name : Config->GCRoots)
+    if (isa<DefinedBitcode>(Symtab[Name]->Body))
+      CG.addMustPreserveSymbol(Name);
 
   CG.setModule(BitcodeFiles[0]->releaseModule());
   for (unsigned I = 1, E = BitcodeFiles.size(); I != E; ++I)
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 51ed9e1..6fd5197 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -87,8 +87,8 @@ void OutputSection::writeHeader(uint8_t *Buf) {
 }
 
 void Writer::markLive() {
-  Entry = cast<Defined>(Symtab->find(Config->EntryName));
-  Entry->markLive();
+  for (StringRef Name : Config->GCRoots)
+    cast<Defined>(Symtab->find(Name))->markLive();
   for (Chunk *C : Symtab->getChunks())
     if (C->isRoot())
       C->markLive();
@@ -291,6 +291,7 @@ void Writer::writeHeader() {
   PE->Subsystem = Config->Subsystem;
   PE->SizeOfImage = SizeOfImage;
   PE->SizeOfHeaders = SizeOfHeaders;
+  Defined *Entry = cast<Defined>(Symtab->find(Config->EntryName));
   PE->AddressOfEntryPoint = Entry->getRVA();
   PE->SizeOfStackReserve = Config->StackReserve;
   PE->SizeOfStackCommit = Config->StackCommit;
diff --git a/lld/COFF/Writer.h b/lld/COFF/Writer.h
index 73de9b5..4705435 100644
--- a/lld/COFF/Writer.h
+++ b/lld/COFF/Writer.h
@@ -98,7 +98,6 @@ private:
   uint32_t ImportDirectoryTableSize = 0;
   uint32_t ImportAddressTableSize = 0;
 
-  Defined *Entry;
   uint64_t FileSize;
   uint64_t SizeOfImage;
   uint64_t SizeOfHeaders;
diff --git a/lld/test/COFF/lto.ll b/lld/test/COFF/lto.ll
index 5af4da3..f84ca07 100644
--- a/lld/test/COFF/lto.ll
+++ b/lld/test/COFF/lto.ll
@@ -8,10 +8,10 @@
 ; RUN: rm -f %T/foo.lib
 ; RUN: llvm-ar cru %T/foo.lib %T/foo.obj
 
-; RUN: lld -flavor link2 /out:%T/main.exe /entry:main /subsystem:console %T/main.lto.obj %T/foo.lto.obj
+; RUN: lld -flavor link2 /out:%T/main.exe /entry:main /include:f2 /subsystem:console %T/main.lto.obj %T/foo.lto.obj
 ; RUN: llvm-readobj -file-headers %T/main.exe | FileCheck -check-prefix=HEADERS-11 %s
 ; RUN: llvm-objdump -d %T/main.exe | FileCheck -check-prefix=TEXT-11 %s
-; RUN: lld -flavor link2 /out:%T/main.exe /entry:main /subsystem:console %T/main.lto.obj %T/foo.lto.lib
+; RUN: lld -flavor link2 /out:%T/main.exe /entry:main /include:f2 /subsystem:console %T/main.lto.obj %T/foo.lto.lib
 ; RUN: llvm-readobj -file-headers %T/main.exe | FileCheck -check-prefix=HEADERS-11 %s
 ; RUN: llvm-objdump -d %T/main.exe | FileCheck -check-prefix=TEXT-11 %s
 
@@ -34,6 +34,21 @@
 ; TEXT-11-NEXT: .text:
 ; TEXT-11-NEXT: xorl	%eax, %eax
 ; TEXT-11-NEXT: retq
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: int3
+; TEXT-11-NEXT: movl	$2, %eax
+; TEXT-11-NEXT: retq
 
 ; HEADERS-01: AddressOfEntryPoint: 0x1000
 ; TEXT-01: Disassembly of section .text:
@@ -79,3 +94,13 @@ define i32 @main() {
 }
 
 declare void @foo()
+
+$f1 = comdat any
+define i32 @f1() comdat($f1) {
+  ret i32 1
+}
+
+$f2 = comdat any
+define i32 @f2() comdat($f2) {
+  ret i32 2
+}