[PECOFF] Support COMDAT associative sections.

[ruiu]

Summary

This is a rework of r210240 to do it in the right way. Previously
it depended on dead-stripping pass to remove unneeded atoms, so
if the pass is not enabled it wouldn't work. This patch does not
depends on dead-stripping pass.

Background:
COFF supports a feature similar to ELF's section groups. This
patch implements it.

In ELF, section groups are identified by their names, and they are
treated somewhat differently from regular symbols. In COFF, the
feature is realized in a more straightforward way. A section can
have an annotation saying "if Nth section is linked, link this
section too." The annotation type is SELECT_ASSOCIATE.

Design:
A new method associate() is added to DefinedAtom. The new method
is supposed to return a list of atoms with SELECT_ASSOCIATE. The
Resolver adds live atom's associates() return values to the live
atom list. By default, assocaites() returns an empty list, so the
default behavior is the same as before.

Caveats:
The reason why I added a new member function associate() to
DefinedAtom, rather than adding a new reference type is mainly
performance. In order to collect all references of some type,
we have to scan all references of all atoms, which is more
costly than calling a member function on all atoms.

[shankarke]

Could we use references to track the associated atoms, rather than adding a separate member to the defined atom ?

[ruiu]

As I wrote in the description of this patch, I thought about that and actually tried
implementing, and concluded that reference is not suitable for that. It was slow
and unnecessarily complicated.

[atanasyan]

Is the assocaites() list serialized to YAML representation if I run lld with the --output-filetype=yaml option?

[ruiu]

Good point, it doesn't. Will update the patch to serialize it from/to YAML.

[ruiu]

• Serialize associates() list from/to YAML

[• shankar.easwaran]

You have to still get native reader/writer to emit this in native format.

Please check with Nick/Bigcheese before committing this change.

[ruiu]

It turned out that dumping associate() contents to Native format
is too much, because it's not easy to dump variable length objects
to Native format. We would have had to define a new chunk for
associated atoms.

So I came back to the original design to use References to represent
"associate" relationships. I believe I found a way to implement it
in an efficient way. PTAL.

[Bigcheese]

lgtm

[ruiu]

Closed by commit rL211106 (authored by @ruiu).

include/lld/Core/Reference.h

enum {
     // kindLayoutAfter is treated as a bidirected edge by the dead-stripping
     // pass.
     kindLayoutAfter = 2,
     // kindLayoutBefore is currently used only by PECOFF port, and will
     // be removed soon. To enforce layout, use kindLayoutAfter instead.
     kindLayoutBefore = 3,
     // kindGroupChild is treated as a bidirected edge too.
     kindGroupChild = 4,
+    kindAssociate = 5,
   };
 
   // A value to be added to the value of a target
   typedef int64_t Addend;
 
   /// If the reference is a fixup in the Atom, then this returns the
   /// byte offset into the Atom's content to do the fix up.
   virtual uint64_t offsetInAtom() const = 0;

include/lld/Core/Resolver.h

private:
     atom_collection_vector<AbsoluteAtom>        _absoluteAtoms;
   };
 
   LinkingContext &_context;
   SymbolTable _symbolTable;
   std::vector<const Atom *>     _atoms;
   std::set<const Atom *>        _deadStripRoots;
   llvm::DenseSet<const Atom *>  _liveAtoms;
+  llvm::DenseSet<const Atom *>  _deadAtoms;
   std::unique_ptr<MergedFile>   _result;
   llvm::DenseMap<const Atom *, llvm::DenseSet<const Atom *>> _reverseRef;
 };
 
 } // namespace lld
 
 #endif // LLD_CORE_RESOLVER_H

lib/Core/Resolver.cpp

 
 // switch all references to undefined or coalesced away atoms
 // to the new defined atom
 void Resolver::updateReferences() {
   ScopedTask task(getDefaultDomain(), "updateReferences");
   for (const Atom *atom : _atoms) {
     if (const DefinedAtom *defAtom = dyn_cast<DefinedAtom>(atom)) {
       for (const Reference *ref : *defAtom) {
+        // A reference of type kindAssociate should't be updated.
+        // Instead, an atom having such reference will be removed
+        // if the target atom is coalesced away, so that they will
+        // go away as a group.
+        if (ref->kindNamespace() == lld::Reference::KindNamespace::all &&
+            ref->kindValue() == lld::Reference::kindAssociate) {
+          if (_symbolTable.isCoalescedAway(atom))
+            _deadAtoms.insert(ref->target());
+          continue;
+        }
         const Atom *newTarget = _symbolTable.replacement(ref->target());
         const_cast<Reference *>(ref)->setTarget(newTarget);
       }
     }
   }
 }
 
 // For dead code stripping, recursively mark atoms "live"
bool Resolver::checkUndefines() {
   }
   return false;
 }
 
 // remove from _atoms all coaleseced away atoms
 void Resolver::removeCoalescedAwayAtoms() {
   ScopedTask task(getDefaultDomain(), "removeCoalescedAwayAtoms");
   _atoms.erase(std::remove_if(_atoms.begin(), _atoms.end(), [&](const Atom *a) {
-                 return _symbolTable.isCoalescedAway(a);
+                 return _symbolTable.isCoalescedAway(a) || _deadAtoms.count(a);
                }),
                _atoms.end());
 }
 
 bool Resolver::resolve() {
   if (!resolveUndefines())
     return false;
   updateReferences();

lib/ReaderWriter/PECOFF/Atoms.h

COFFDefinedFileAtom(const File &file, StringRef name, StringRef sectionName,
         _scope(scope), _contentType(contentType), _permissions(perms),
         _ordinal(ordinal), _alignment(0) {}
 
   static bool classof(const COFFBaseDefinedAtom *atom) {
     return atom->getKind() == Kind::File;
   }
 
   void setAlignment(Alignment val) { _alignment = val; }
-
   SectionChoice sectionChoice() const override { return sectionCustomRequired; }
   StringRef customSectionName() const override { return _sectionName; }
   Scope scope() const override { return _scope; }
   ContentType contentType() const override { return _contentType; }
   ContentPermissions permissions() const override { return _permissions; }
   uint64_t ordinal() const override { return _ordinal; }
   Alignment alignment() const override { return _alignment; }
 
+  void addAssociate(const DefinedAtom *other) {
+    auto *ref = new COFFReference(other, 0, lld::Reference::kindAssociate,
+                                  Reference::KindNamespace::all,
+                                  Reference::KindArch::all);
+    addReference(std::unique_ptr<COFFReference>(ref));
+  }
+
 private:
   StringRef _sectionName;
   Scope _scope;
   ContentType _contentType;
   ContentPermissions _permissions;
   uint64_t _ordinal;
   Alignment _alignment;
   std::vector<std::unique_ptr<COFFReference> > _references;

lib/ReaderWriter/PECOFF/ReaderCOFF.cpp

private:
   _sectionAtoms;
 
   // A set of COMDAT sections.
   std::set<const coff_section *> _comdatSections;
 
   // 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 *,
            std::map<uint32_t, std::vector<COFFDefinedAtom *>>>
   _definedAtomLocations;
 
   uint64_t _ordinal;
 };

 // Cache the COMDAT attributes, which indicate whether the symbols in the
 // section can be merged or not.
 error_code FileCOFF::cacheSectionAttributes() {
   // The COMDAT section attribute is not an attribute of coff_section, but is
   // stored in the auxiliary symbol for the first symbol referring a COMDAT
   // 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)
       continue;
 
     const coff_section *sec;
     if (error_code ec = _obj->getSection(sym->SectionNumber, sec))
       return ec;
-
-    if (_merge.count(sec))
-      continue;
-    if (!(sec->Characteristics & llvm::COFF::IMAGE_SCN_LNK_COMDAT))
-      continue;
-
-    _comdatSections.insert(sec);
-
-    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);
+
+    if (sec->Characteristics & llvm::COFF::IMAGE_SCN_LNK_COMDAT) {
+      // Read aux symbol data.
+      _comdatSections.insert(sec);
       _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
   // which symbols are not allowed to be merged.
   for (const auto &section : _obj->sections()) {
     const coff_section *sec = _obj->getCOFFSection(section);
     if (!_merge.count(sec))
       _merge[sec] = DefinedAtom::mergeNo;
   }
   return error_code();
for (auto &i : definedSymbols) {
     // Connect atoms with layout-before/layout-after edges.
     connectAtomsWithLayoutEdge(atoms);
 
     for (COFFDefinedFileAtom *atom : atoms) {
       _sectionAtoms[section].push_back(atom);
       definedAtoms.push_back(atom);
     }
   }
+
+  // A COMDAT section with SELECT_ASSOCIATIVE attribute refer to other
+  // section. If the referred section is linked to a binary, the
+  // referring section needs to be linked too. A typical use case of
+  // this attribute is a static initializer; a parent is a comdat BSS
+  // section, and a child is a static initializer code for the data.
+  //
+  // We add referring section contents to the referred section's
+  // associate list, so that Resolver takes care of them.
+  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->addAssociate(_sectionAtoms[sec][0]);
+      }
+    }
+  }
+
   return error_code();
 }
 
 /// Find the atom that is at \p targetAddress in \p section.
 error_code FileCOFF::findAtomAt(const coff_section *section,
                                 uint32_t targetAddress,
                                 COFFDefinedFileAtom *&result,
                                 uint32_t &offsetInAtom) {

lib/ReaderWriter/Reader.cpp

Registry::parseFile(std::unique_ptr<MemoryBuffer> &mb,
   return llvm::make_error_code(std::errc::executable_format_error);
 }
 
 static const Registry::KindStrings kindStrings[] = {
     {Reference::kindInGroup, "in-group"},
     {Reference::kindLayoutAfter, "layout-after"},
     {Reference::kindLayoutBefore, "layout-before"},
     {Reference::kindGroupChild, "group-child"},
+    {Reference::kindAssociate, "associate"},
     LLD_KIND_STRING_END};
 
 Registry::Registry() {
   addKindTable(Reference::KindNamespace::all, Reference::KindArch::all,
                kindStrings);
 }
 
 bool Registry::handleTaggedDoc(llvm::yaml::IO &io,

test/core/associates.objtxt

+# RUN: lld -core %s | FileCheck %s
+
+---
+defined-atoms:
+  - name:  f1
+    merge: as-weak
+    scope: global
+    references:
+      - kind:            associate
+        target:          f2
+  - name:  f2
+---
+defined-atoms:
+  - name:  f1
+    merge: as-weak
+    scope: global
+    references:
+      - kind:            associate
+        target:          f2
+  - name:  f2
+...
+
+# CHECK: defined-atoms:
+# CHECK:   - name:  f1
+# CHECK:     scope: global
+# CHECK:     references:
+# CHECK:       - kind:            associate
+# CHECK:         target:          f2
+# CHECK:   - name:  f2
+# CHECK-NOT: - name:  f2

test/pecoff/Inputs/associative1.obj.yaml

+---
+header:
+  Machine:         IMAGE_FILE_MACHINE_I386
+  Characteristics: []
+sections:
+  - Name:            .data
+    Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ]
+    Alignment:       4
+    SectionData:     00000000
+  - Name:            '.CRT$XCU'
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       4
+    SectionData:     77777777
+symbols:
+  - Name:            .data
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition:
+      Length:          4
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          0
+      Selection:       IMAGE_COMDAT_SELECT_ANY
+  - Name:            _var
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            '.CRT$XCU'
+    Value:           0
+    SectionNumber:   2
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition:
+      Length:          4
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          1
+      Selection:       IMAGE_COMDAT_SELECT_ASSOCIATIVE
+  - Name:            _init
+    Value:           0
+    SectionNumber:   2
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...

test/pecoff/Inputs/associative3.obj.yaml

+---
+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:     0000000000000000
+    Relocations:
+      - VirtualAddress:  4
+        SymbolName:      _var
+        Type:            IMAGE_REL_I386_DIR32
+symbols:
+  - Name:            .text
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+  - Name:            _main
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            _var
+    Value:           0
+    SectionNumber:   0
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...

test/pecoff/associative.test

+# RUN: yaml2obj %p/Inputs/associative1.obj.yaml > %t1.obj
+# RUN: yaml2obj %p/Inputs/associative1.obj.yaml > %t2.obj
+# RUN: yaml2obj %p/Inputs/associative3.obj.yaml > %t3.obj
+#
+# RUN: lld -flavor link /machine:x86 /subsystem:console /entry:main \
+# RUN:   /out:%t.exe -- %t1.obj %t2.obj %t3.obj
+# RUN: obj2yaml %t.exe | FileCheck %s
+
+CHECK: - Name: .CRT
+CHECK:   SectionData: '77777777'