[PECOFF] Support COMDAT associative sections.

[ruiu]

Summary

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."

Implementing such feature is easy. We can add a reference from a
target atom to an original atom, so that if the target is linked,
the original atom is also linked. If not linked, both will be
dead-stripped. So they are treated as a group.

I added a new reference type, kindAssociate. It does nothing except
preventing referenced atoms from being dead-stripped.

No change to the Resolver is needed.

[atanasyan]

LGTM with minor question / note.

lib/ReaderWriter/PECOFF/ReaderCOFF.cpp

174 ↗

(On Diff #9905)

Is it really necessary to use the std::set here i.e. is it possible that a COFF file contains more than one "record" states that sections say "A" and "B" are associated?

[rnk]

Is it OK to rely on dead stripping to remove associative comdat sections when the associated data is discarded? I thought it was possible for the user to disable dead stripping, in which case the result will be incorrect, and data will be initialized multiple times.

[ruiu]

Closed by commit rL210240 (authored by @ruiu).

lld/trunk/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 prevents a referenced atom from being dead-stripped.
+    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;

lld/trunk/lib/ReaderWriter/PECOFF/Atoms.h

public:
   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 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;
   ContentType _contentType;
   ContentPermissions _permissions;
   uint64_t _ordinal;
   Alignment _alignment;
   std::vector<std::unique_ptr<COFFReference> > _references;

lld/trunk/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. 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();
 }
 
 /// 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) {

lld/trunk/lib/ReaderWriter/Reader.cpp

Registry::parseFile(std::unique_ptr<MemoryBuffer> &mb,
   return llvm::make_error_code(llvm::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,

lld/trunk/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
+...

lld/trunk/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
+...

lld/trunk/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'