3 files changed, 48 insertions, 11 deletions

diff --git a/lld/lib/ReaderWriter/PECOFF/Atoms.h b/lld/lib/ReaderWriter/PECOFF/Atoms.h
index c4ae1271c0b..7cd74951c0c 100644
--- a/lld/lib/ReaderWriter/PECOFF/Atoms.h
+++ b/lld/lib/ReaderWriter/PECOFF/Atoms.h
@@ -167,6 +167,13 @@ public:
   uint64_t ordinal() const override { return _ordinal; }
   Alignment alignment() const override { return _alignment; }
 
+  void associate(COFFDefinedFileAtom *other) {
+    auto *r = new COFFReference(other, 0, Reference::kindAssociate,
+				Reference::KindNamespace::all,
+				Reference::KindArch::all);
+    addReference(std::unique_ptr<COFFReference>(r));
+  }
+
 private:
   StringRef _sectionName;
   Scope _scope;
diff --git a/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp b/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
index 47520972bac..069c710ac3e 100644
--- a/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
@@ -172,6 +172,9 @@ private:
   // A map to get whether the section allows its contents to be merged or not.
   std::map<const coff_section *, DefinedAtom::Merge> _merge;
 
+  // COMDAT associative sections
+  std::map<const coff_section *, std::set<const coff_section *>> _association;
+
   // A sorted map to find an atom from a section and an offset within
   // the section.
   std::map<const coff_section *,
@@ -540,6 +543,7 @@ error_code FileCOFF::cacheSectionAttributes() {
   // section. It feels to me that it's unnecessarily complicated, but this is
   // how COFF works.
   for (auto i : _auxSymbol) {
+    // Read a section from the file
     const coff_symbol *sym = i.first;
     if (sym->SectionNumber == llvm::COFF::IMAGE_SYM_ABSOLUTE ||
         sym->SectionNumber == llvm::COFF::IMAGE_SYM_UNDEFINED)
@@ -548,19 +552,22 @@ error_code FileCOFF::cacheSectionAttributes() {
     const coff_section *sec;
     if (error_code ec = _obj->getSection(sym->SectionNumber, sec))
       return ec;
+    const coff_aux_section_definition *aux =
+      reinterpret_cast<const coff_aux_section_definition *>(i.second);
 
-    if (_merge.count(sec))
-      continue;
-    if (!(sec->Characteristics & llvm::COFF::IMAGE_SCN_LNK_COMDAT))
-      continue;
-
-    _comdatSections.insert(sec);
+    if (sec->Characteristics & llvm::COFF::IMAGE_SCN_LNK_COMDAT) {
+      // Read aux symbol data.
+      _comdatSections.insert(sec);
+      _merge[sec] = getMerge(aux);
+    }
 
-    if (sym->NumberOfAuxSymbols == 0)
-      return llvm::object::object_error::parse_failed;
-    const coff_aux_section_definition *aux =
-        reinterpret_cast<const coff_aux_section_definition *>(i.second);
-    _merge[sec] = getMerge(aux);
+    // Handle associative sections.
+    if (aux->Selection == llvm::COFF::IMAGE_COMDAT_SELECT_ASSOCIATIVE) {
+      const coff_section *parent;
+      if (error_code ec = _obj->getSection(aux->Number, parent))
+	return ec;
+      _association[parent].insert(sec);
+    }
   }
 
   // The sections that does not have auxiliary symbol are regular sections, in
@@ -694,6 +701,28 @@ error_code FileCOFF::AtomizeDefinedSymbols(
       definedAtoms.push_back(atom);
     }
   }
+
+  // A COMDAT section with SELECT_ASSOCIATIVE attribute refer to other
+  // section. They need to be linked to a binary or dead stripped as a group.
+  // Here We add a kindAssociate edge between them, so that if a target section
+  // is linked, the other is also linked.
+  //
+  // In a typical use case, parent is a comdat BSS section, and a child is a
+  // static initializer for the data.
+  for (auto i : _association) {
+    const coff_section *parent = i.first;
+    const std::set<const coff_section *> &childSections = i.second;
+    assert(_sectionAtoms[parent].size() > 0);
+
+    COFFDefinedFileAtom *p = _sectionAtoms[parent][0];
+    for (const coff_section *sec : childSections) {
+      if (_sectionAtoms.count(sec)) {
+	assert(_sectionAtoms[sec].size() > 0);
+	p->associate(_sectionAtoms[sec][0]);
+      }
+    }
+  }
+
   return error_code();
 }
 
diff --git a/lld/lib/ReaderWriter/Reader.cpp b/lld/lib/ReaderWriter/Reader.cpp
index ce81c69186e..644acd7fe96 100644
--- a/lld/lib/ReaderWriter/Reader.cpp
+++ b/lld/lib/ReaderWriter/Reader.cpp
@@ -52,6 +52,7 @@ static const Registry::KindStrings kindStrings[] = {
     {Reference::kindLayoutAfter, "layout-after"},
     {Reference::kindLayoutBefore, "layout-before"},
     {Reference::kindGroupChild, "group-child"},
+    {Reference::kindAssociate, "associate"},
     LLD_KIND_STRING_END};
 
 Registry::Registry() {