COFF: Implement string tail merging.

[pcc]

Summary

In COFF, duplicate string literals are merged by placing them in a
comdat whose leader symbol name contains a specific prefix followed
by the hash and partial contents of the string literal. This gives
us an easy way to identify sections containing string literals in
the linker: check for leader symbol names with the given prefix.

Any sections that are identified in this way as containing string
literals may be tail merged. We do so using the StringTableBuilder
class, which is also used to tail merge string literals in the ELF
linker. Tail merging is enabled only if ICF is enabled, as this
provides a signal as to whether the user cares about binary size.

Depends on D44501

[pcc]

• Check contents of .rdata

[pcc]

• Align MergeChunks correctly

[majnemer]

Do you have a test to make sure that a wide string literal doesn't point into a misaligned offset into a narrow string literal?

Also, does the standard allow aliasing string literals of different underlying types together?

[pcc]

Do you have a test to make sure that a wide string literal doesn't point into a misaligned offset into a narrow string literal?

I will add one.

Also, does the standard allow aliasing string literals of different underlying types together?

That's an interesting question. The standardese is:

Whether all string literals are distinct (that is, are stored in nonoverlapping objects) and whether successive evaluations of a string-literal yield the same or a different object is unspecified.

This doesn't make any distinction between different kinds of string literal, so I think the answer is "yes". In any event, this seems to be the intent of the standard.

[pcc]

• Add a test that makes sure that we don't merge string literals with different alignments

[ruiu]

Overall looking very nice.

lld/COFF/Chunks.cpp

590–592 ↗

(On Diff #138489)

I believe you can use toStringRef(C->getContents()).

601 ↗

(On Diff #138489)

I'd flip the condition and continue early.

lld/COFF/Chunks.h

230–231 ↗

(On Diff #138489)

I'd expand this comment a little bit to say that this is an lld-specific feature (not implemented in MSVC) to minimize the output size by finding string literals sharing tail parts and merge them. Maybe I'd explain how it work a bit too.

232 ↗

(On Diff #138489)

Do all members have to be public? Even that's the case, I feel comfortable if you make it not a struct but a class with no private member.

lld/COFF/ICF.cpp

230 ↗

(On Diff #138489)

auto -> MergeSection

lld/COFF/Writer.cpp

429 ↗

(On Diff #138489)

auto -> MergeSection

[pcc]

• Address review comments

lld/COFF/ICF.cpp

230 ↗

(On Diff #138489)

It would actually be std::pair<const uint32_t, MergeChunk *>, which seems a little verbose, so I left it the way it is.

lld/COFF/Writer.cpp

429 ↗

(On Diff #138489)

Likewise

[ruiu]

LGTM

This feature is awesome. I strongly believe that Microsoft should do the same thing in the next version of MSVC. :)

lld/COFF/ICF.cpp

230 ↗

(On Diff #138489)

But what about SC? Isn't it SectionChunk?

[pcc]

Use SectionChunk

[pcc]

I will commit when D44501 is approved :)

lld/COFF/ICF.cpp

230 ↗

(On Diff #138489)

Oh, right.

COFF/Chunks.h

 #define LLD_COFF_CHUNKS_H
 
 #include "Config.h"
 #include "InputFiles.h"
 #include "lld/Common/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
+#include "llvm/MC/StringTableBuilder.h"
 #include "llvm/Object/COFF.h"
 #include <utility>
 #include <vector>
 
 namespace lld {
 namespace coff {
 
 using llvm::COFF::ImportDirectoryTableEntry;
public:
   virtual size_t getSize() const = 0;
 
   // Write this chunk to a mmap'ed file, assuming Buf is pointing to
   // beginning of the file. Because this function may use RVA values
   // of other chunks for relocations, you need to set them properly
   // before calling this function.
   virtual void writeTo(uint8_t *Buf) const {}
 
+  // Called by the writer after an RVA is assigned, but before calling
+  // getSize().
+  virtual void finalizeContents() {}
+
   // The writer sets and uses the addresses.
   uint64_t getRVA() const { return RVA; }
   void setRVA(uint64_t V) { RVA = V; }
 
   // Returns true if this has non-zero data. BSS chunks return
   // false. If false is returned, the space occupied by this chunk
   // will be filled with zeros.
   virtual bool hasData() const { return true; }
private:
   // Used by the garbage collector.
   bool Live;
 
   // Used for ICF (Identical COMDAT Folding)
   void replace(SectionChunk *Other);
   uint32_t Class[2] = {0, 0};
 };
 
+// This class is used to implement an lld-specific feature (not implemented in
+// MSVC) that minimizes the output size by finding string literals sharing tail
+// parts and merging them.
+//
+// If string tail merging is enabled and a section is identified as containing a
+// string literal, it is added to a MergeChunk with an appropriate alignment.
+// The MergeChunk then tail merges the strings using the StringTableBuilder
+// class and assigns RVAs and section offsets to each of the member chunks based
+// on the offsets assigned by the StringTableBuilder.
+class MergeChunk : public Chunk {
+public:
+  MergeChunk(uint32_t Alignment);
+  static void addSection(SectionChunk *C);
+  void finalizeContents() override;
+
+  uint32_t getPermissions() const override;
+  StringRef getSectionName() const override { return ".rdata"; }
+  size_t getSize() const override;
+  void writeTo(uint8_t *Buf) const override;
+
+  static std::map<uint32_t, MergeChunk *> Instances;
+  std::vector<SectionChunk *> Sections;
+
+private:
+  llvm::StringTableBuilder Builder;
+};
+
 // A chunk for common symbols. Common chunks don't have actual data.
 class CommonChunk : public Chunk {
 public:
   CommonChunk(const COFFSymbolRef Sym);
   size_t getSize() const override { return Sym.getValue(); }
   bool hasData() const override { return false; }
   uint32_t getPermissions() const override;
   StringRef getSectionName() const override { return ".bss"; }

COFF/Chunks.cpp

uint8_t Baserel::getDefaultType() {
   case I386:
   case ARMNT:
     return IMAGE_REL_BASED_HIGHLOW;
   default:
     llvm_unreachable("unknown machine type");
   }
 }
 
+std::map<uint32_t, MergeChunk *> MergeChunk::Instances;
+
+MergeChunk::MergeChunk(uint32_t Alignment)
+    : Builder(StringTableBuilder::RAW, Alignment) {
+  this->Alignment = Alignment;
+}
+
+void MergeChunk::addSection(SectionChunk *C) {
+  auto *&MC = Instances[C->Alignment];
+  if (!MC)
+    MC = make<MergeChunk>(C->Alignment);
+  MC->Sections.push_back(C);
+}
+
+void MergeChunk::finalizeContents() {
+  for (SectionChunk *C : Sections)
+    if (C->isLive())
+      Builder.add(toStringRef(C->getContents()));
+  Builder.finalize();
+
+  for (SectionChunk *C : Sections) {
+    if (!C->isLive())
+      continue;
+    size_t Off = Builder.getOffset(toStringRef(C->getContents()));
+    C->setOutputSection(Out);
+    C->setRVA(RVA + Off);
+    C->OutputSectionOff = OutputSectionOff + Off;
+  }
+}
+
+uint32_t MergeChunk::getPermissions() const {
+  return IMAGE_SCN_MEM_READ | IMAGE_SCN_CNT_INITIALIZED_DATA;
+}
+
+size_t MergeChunk::getSize() const {
+  return Builder.getSize();
+}
+
+void MergeChunk::writeTo(uint8_t *Buf) const {
+  Builder.write(Buf + OutputSectionOff);
+}
+
 } // namespace coff
 } // namespace lld

COFF/ICF.cpp

for (Chunk *C : Vec) {
     if (auto *SC = dyn_cast<SectionChunk>(C)) {
       if (isEligible(SC))
         Chunks.push_back(SC);
       else
         SC->Class[0] = NextId++;
     }
   }
 
+  // Make sure that ICF doesn't merge sections that are being handled by string
+  // tail merging.
+  for (auto &P : MergeChunk::Instances)
+    for (SectionChunk *SC : P.second->Sections)
+      SC->Class[0] = NextId++;
+
   // Initially, we use hash values to partition sections.
   for_each(parallel::par, Chunks.begin(), Chunks.end(), [&](SectionChunk *SC) {
     // Set MSB to 1 to avoid collisions with non-hash classs.
     SC->Class[0] = getHash(SC) | (1 << 31);
   });
 
   // From now on, sections in Chunks are ordered so that sections in
   // the same group are consecutive in the vector.

COFF/InputFiles.h

public:
   llvm::pdb::DbiModuleDescriptorBuilder *ModuleDBI = nullptr;
 
 private:
   void initializeChunks();
   void initializeSymbols();
 
   SectionChunk *
   readSection(uint32_t SectionNumber,
-              const llvm::object::coff_aux_section_definition *Def);
+              const llvm::object::coff_aux_section_definition *Def,
+              StringRef LeaderName);
 
   void readAssociativeDefinition(
       COFFSymbolRef COFFSym,
       const llvm::object::coff_aux_section_definition *Def);
 
   llvm::Optional<Symbol *>
   createDefined(COFFSymbolRef Sym,
                 std::vector<const llvm::object::coff_aux_section_definition *>

COFF/InputFiles.cpp

void ObjFile::initializeChunks() {
   for (uint32_t I = 1; I < NumSections + 1; ++I) {
     const coff_section *Sec;
     if (auto EC = COFFObj->getSection(I, Sec))
       fatal("getSection failed: #" + Twine(I) + ": " + EC.message());
 
     if (Sec->Characteristics & IMAGE_SCN_LNK_COMDAT)
       SparseChunks[I] = PendingComdat;
     else
-      SparseChunks[I] = readSection(I, nullptr);
+      SparseChunks[I] = readSection(I, nullptr, "");
   }
 }
 
 SectionChunk *ObjFile::readSection(uint32_t SectionNumber,
-                                   const coff_aux_section_definition *Def) {
+                                   const coff_aux_section_definition *Def,
+                                   StringRef LeaderName) {
   const coff_section *Sec;
   StringRef Name;
   if (auto EC = COFFObj->getSection(SectionNumber, Sec))
     fatal("getSection failed: #" + Twine(SectionNumber) + ": " + EC.message());
   if (auto EC = COFFObj->getSectionName(Sec, Name))
     fatal("getSectionName failed: #" + Twine(SectionNumber) + ": " +
           EC.message());
 
SectionChunk *ObjFile::readSection(uint32_t SectionNumber,
   if (C->isCodeView())
     DebugChunks.push_back(C);
   else if (Config->GuardCF != GuardCFLevel::Off && Name == ".gfids$y")
     GuardFidChunks.push_back(C);
   else if (Config->GuardCF != GuardCFLevel::Off && Name == ".gljmp$y")
     GuardLJmpChunks.push_back(C);
   else if (Name == ".sxdata")
     SXDataChunks.push_back(C);
+  else if (Config->DoICF && Sec->NumberOfRelocations == 0 && Name == ".rdata" &&
+           LeaderName.startswith("??_C@"))
+    // COFF sections that look like string literal sections (i.e. no
+    // relocations, in .rdata, leader symbol name matches the MSVC name mangling
+    // for string literals) are subject to string tail merging.
+    MergeChunk::addSection(C);
   else
     Chunks.push_back(C);
 
   return C;
 }
 
 void ObjFile::readAssociativeDefinition(
     COFFSymbolRef Sym, const coff_aux_section_definition *Def) {
   SectionChunk *Parent = SparseChunks[Def->getNumber(Sym.isBigObj())];
 
   // If the parent is pending, it probably means that its section definition
   // appears after us in the symbol table. Leave the associated section as
   // pending; we will handle it during the second pass in initializeSymbols().
   if (Parent == PendingComdat)
     return;
 
   // Check whether the parent is prevailing. If it is, so are we, and we read
   // the section; otherwise mark it as discarded.
   int32_t SectionNumber = Sym.getSectionNumber();
   if (Parent) {
-    SparseChunks[SectionNumber] = readSection(SectionNumber, Def);
+    SparseChunks[SectionNumber] = readSection(SectionNumber, Def, "");
     if (SparseChunks[SectionNumber])
       Parent->addAssociative(SparseChunks[SectionNumber]);
   } else {
     SparseChunks[SectionNumber] = nullptr;
   }
 }
 
 Symbol *ObjFile::createRegular(COFFSymbolRef Sym) {
if (Sym.isExternal()) {
       std::tie(Leader, Prevailing) =
           Symtab->addComdat(this, Name, Sym.getGeneric());
     } else {
       Leader = make<DefinedRegular>(this, /*Name*/ "", false,
                                     /*IsExternal*/ false, Sym.getGeneric());
       Prevailing = true;
     }
     if (Prevailing) {
-      SectionChunk *C = readSection(SectionNumber, Def);
+      SectionChunk *C = readSection(SectionNumber, Def, Name);
       SparseChunks[SectionNumber] = C;
       C->Sym = cast<DefinedRegular>(Leader);
       cast<DefinedRegular>(Leader)->Data = &C->Repl;
     } else {
       SparseChunks[SectionNumber] = nullptr;
     }
     return Leader;
   }

COFF/Writer.cpp

for (Chunk *C : Chunks) {
       Sec->addPermissions(C->getPermissions());
     }
   }
 }
 
 void Writer::createMiscChunks() {
   OutputSection *RData = createSection(".rdata");
 
+  for (auto &P : MergeChunk::Instances)
+    RData->addChunk(P.second);
+
   // Create thunks for locally-dllimported symbols.
   if (!Symtab->LocalImportChunks.empty()) {
     for (Chunk *C : Symtab->LocalImportChunks)
       RData->addChunk(C);
   }
 
   // Create Debug Information Chunks
   if (Config->Debug) {
for (OutputSection *Sec : OutputSections) {
     if (Sec->Name == ".reloc")
       addBaserels(Sec);
     uint64_t RawSize = 0, VirtualSize = 0;
     Sec->Header.VirtualAddress = RVA;
     for (Chunk *C : Sec->getChunks()) {
       VirtualSize = alignTo(VirtualSize, C->Alignment);
       C->setRVA(RVA + VirtualSize);
       C->OutputSectionOff = VirtualSize;
+      C->finalizeContents();
       VirtualSize += C->getSize();
       if (C->hasData())
         RawSize = alignTo(VirtualSize, SectorSize);
     }
     if (VirtualSize > UINT32_MAX)
       error("section larger than 4 GiB: " + Sec->Name);
     Sec->Header.VirtualSize = VirtualSize;
     Sec->Header.SizeOfRawData = RawSize;

test/COFF/string-tail-merge.s

+# REQUIRES: x86
+# RUN: llvm-mc -triple=x86_64-windows-msvc -filetype=obj -o %t.obj %s
+# RUN: lld-link %t.obj /out:%t.exe /entry:main /subsystem:console
+# RUN: llvm-objdump -s %t.exe | FileCheck %s
+
+# CHECK: Contents of section .rdata:
+# CHECK-NEXT:  140002000 68656c6c 6f20776f 726c6400 6fa26ca4  hello world.o.l.
+# CHECK-NEXT:  140002010 0068656c 6c6f2077 6f726c64 00006865  .hello world..he
+# CHECK-NEXT:  140002020 6c6c6f20 776f726c 64006800 65006c00  llo world.h.e.l.
+# CHECK-NEXT:  140002030 6c006f00 20007700 6f007200 6c006400  l.o. .w.o.r.l.d.
+# CHECK-NEXT:  140002040 0000                                 ..
+
+# CHECK: Contents of section .text:
+.globl main
+main:
+# CHECK-NEXT: 140003000 11200040 01000000 17200040 01000000
+.8byte "??_C@_0M@LACCCNMM@hello?5world?$AA@"
+.8byte "??_C@_05MCBCHHEJ@world?$AA@"
+# CHECK-NEXT: 140003010 2a200040 01000000 36200040 01000000
+.8byte "??_C@_1BI@HHJHKLLN@?$AAh?$AAe?$AAl?$AAl?$AAo?$AA?5?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@"
+.8byte "??_C@_1M@NBBDDHIO@?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@"
+# CHECK-NEXT: 140003020 00200040 01000000 0c200040 01000000
+.8byte "??_D@not_a_string_literal"
+.8byte "??_C@string_literal_with_relocs"
+# CHECK-NEXT: 140003030 00100040 01000000 1e200040 01000000
+.8byte "??_C@string_literal_in_wrong_section"
+.8byte "??_C@overaligned_string_literal"
+
+.section .rdata,"dr",discard,"??_C@_0M@LACCCNMM@hello?5world?$AA@"
+.globl "??_C@_0M@LACCCNMM@hello?5world?$AA@"
+"??_C@_0M@LACCCNMM@hello?5world?$AA@":
+.asciz "hello world"
+
+.section .rdata,"dr",discard,"??_C@_05MCBCHHEJ@world?$AA@"
+.globl "??_C@_05MCBCHHEJ@world?$AA@"
+"??_C@_05MCBCHHEJ@world?$AA@":
+.asciz "world"
+
+.section .rdata,"dr",discard,"??_C@_1BI@HHJHKLLN@?$AAh?$AAe?$AAl?$AAl?$AAo?$AA?5?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@"
+.globl "??_C@_1BI@HHJHKLLN@?$AAh?$AAe?$AAl?$AAl?$AAo?$AA?5?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@"
+.p2align 1
+"??_C@_1BI@HHJHKLLN@?$AAh?$AAe?$AAl?$AAl?$AAo?$AA?5?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@":
+.short 104
+.short 101
+.short 108
+.short 108
+.short 111
+.short 32
+.short 119
+.short 111
+.short 114
+.short 108
+.short 100
+.short 0
+
+.section .rdata,"dr",discard,"??_C@_1M@NBBDDHIO@?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@"
+.globl "??_C@_1M@NBBDDHIO@?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@"
+.p2align 1
+"??_C@_1M@NBBDDHIO@?$AAw?$AAo?$AAr?$AAl?$AAd?$AA?$AA@":
+.short 119
+.short 111
+.short 114
+.short 108
+.short 100
+.short 0
+
+.section .data,"drw",discard,"??_C@string_literal_in_wrong_section"
+.globl "??_C@string_literal_in_wrong_section"
+"??_C@string_literal_in_wrong_section":
+.asciz "hello world"
+
+.section .rdata,"dr",discard,"??_D@not_a_string_literal"
+.globl "??_D@not_a_string_literal"
+"??_D@not_a_string_literal":
+.asciz "hello world"
+
+.section .rdata,"dr",discard,"??_C@string_literal_with_relocs"
+.globl "??_C@string_literal_with_relocs"
+"??_C@string_literal_with_relocs":
+.4byte main + 111 + (114 << 8) + (108 << 16) + (100 << 24) # main + "orld"
+.byte 0
+
+.section .rdata,"dr",discard,"??_C@overaligned_string_literal"
+.globl "??_C@overaligned_string_literal"
+.p2align 1
+"??_C@overaligned_string_literal":
+.asciz "hello world"