From 728cc0075e5dfdb454ebe6418b63bdffae71ec14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Storsj=C3=B6?= <martin@martin.st>
Date: Mon, 24 May 2021 15:06:56 +0300
Subject: [PATCH] [LLD] [COFF] Fix autoexport from LTO objects with comdat
 symbols

Make sure that comdat symbols also have a non-null dummy
SectionChunk associated.

This requires moving around an existing FIXME regarding comdats in
LTO.

Differential Revision: https://reviews.llvm.org/D103012
---
 COFF/InputFiles.cpp         | 40 ++++++++++++++++++++-------------
 test/COFF/export-all-lto.ll | 14 ++++++++++--
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/COFF/InputFiles.cpp b/COFF/InputFiles.cpp
index 9900903bfe7cc..f4abeb4acfb4a 100644
--- a/COFF/InputFiles.cpp
+++ b/COFF/InputFiles.cpp
@@ -503,14 +503,9 @@ void ObjFile::handleComdatSelection(
   SectionChunk *leaderChunk = nullptr;
   COMDATType leaderSelection = IMAGE_COMDAT_SELECT_ANY;
 
-  if (leader->data) {
-    leaderChunk = leader->getChunk();
-    leaderSelection = leaderChunk->selection;
-  } else {
-    // FIXME: comdats from LTO files don't know their selection; treat them
-    // as "any".
-    selection = leaderSelection;
-  }
+  assert(leader->data && "Comdat leader without SectionChunk?");
+  leaderChunk = leader->getChunk();
+  leaderSelection = leaderChunk->selection;
 
   if ((selection == IMAGE_COMDAT_SELECT_ANY &&
        leaderSelection == IMAGE_COMDAT_SELECT_LARGEST) ||
@@ -1051,10 +1046,22 @@ class FakeSection {
   coff_section section;
 };
 
+// Convenience class for initializing a SectionChunk with specific flags.
+class FakeSectionChunk {
+public:
+  FakeSectionChunk(const coff_section *section) : chunk(nullptr, section) {
+    // FIXME: comdats from LTO files don't know their selection; treat them
+    // as "any".
+    chunk.selection = IMAGE_COMDAT_SELECT_ANY;
+  }
+
+  SectionChunk chunk;
+};
+
 FakeSection ltoTextSection(IMAGE_SCN_MEM_EXECUTE);
 FakeSection ltoDataSection(IMAGE_SCN_CNT_INITIALIZED_DATA);
-SectionChunk ltoTextSectionChunk(nullptr, &ltoTextSection.section);
-SectionChunk ltoDataSectionChunk(nullptr, &ltoDataSection.section);
+FakeSectionChunk ltoTextSectionChunk(&ltoTextSection.section);
+FakeSectionChunk ltoDataSectionChunk(&ltoDataSection.section);
 } // namespace
 
 void BitcodeFile::parse() {
@@ -1069,9 +1076,9 @@ void BitcodeFile::parse() {
     Symbol *sym;
     SectionChunk *fakeSC = nullptr;
     if (objSym.isExecutable())
-      fakeSC = &ltoTextSectionChunk;
+      fakeSC = &ltoTextSectionChunk.chunk;
     else
-      fakeSC = &ltoDataSectionChunk;
+      fakeSC = &ltoDataSectionChunk.chunk;
     if (objSym.isUndefined()) {
       sym = symtab->addUndefined(symName, this, false);
     } else if (objSym.isCommon()) {
@@ -1083,12 +1090,15 @@ void BitcodeFile::parse() {
       Symbol *alias = symtab->addUndefined(saver.save(fallback));
       checkAndSetWeakAlias(symtab, this, sym, alias);
     } else if (comdatIndex != -1) {
-      if (symName == obj->getComdatTable()[comdatIndex])
+      if (symName == obj->getComdatTable()[comdatIndex]) {
         sym = comdat[comdatIndex].first;
-      else if (comdat[comdatIndex].second)
+        if (cast<DefinedRegular>(sym)->data == nullptr)
+          cast<DefinedRegular>(sym)->data = &fakeSC->repl;
+      } else if (comdat[comdatIndex].second) {
         sym = symtab->addRegular(this, symName, nullptr, fakeSC);
-      else
+      } else {
         sym = symtab->addUndefined(symName, this, false);
+      }
     } else {
       sym = symtab->addRegular(this, symName, nullptr, fakeSC);
     }
diff --git a/test/COFF/export-all-lto.ll b/test/COFF/export-all-lto.ll
index 9d68809a77e07..13620e11d9b93 100644
--- a/test/COFF/export-all-lto.ll
+++ b/test/COFF/export-all-lto.ll
@@ -6,18 +6,28 @@
 ; RUN: llvm-readobj --coff-exports %t.dll | grep Name: | FileCheck %s
 ; RUN: cat %t.def | FileCheck --check-prefix=IMPLIB %s
 
+; CHECK: Name: MyComdatFunc
 ; CHECK: Name: MyExtData
 ; CHECK: Name: MyLibFunc
 
-; IMPLIB: MyExtData @1 DATA
-; IMPLIB: MyLibFunc @2{{$}}
+; IMPLIB: MyComdatFunc @1{{$}}
+; IMPLIB: MyExtData @2 DATA
+; IMPLIB: MyLibFunc @3{{$}}
 
 target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-w64-windows-gnu"
 
 @MyExtData = dso_local global i32 42, align 4
 
+$MyComdatFunc = comdat any
+
 define dso_local void @MyLibFunc() {
+entry:
+  call void @MyComdatFunc()
+  ret void
+}
+
+define linkonce_odr void @MyComdatFunc() comdat {
 entry:
   ret void
 }