MaskRay · MaskRay · Jun 27, 2023
diff --git a/lld/ELF/Config.h b/lld/ELF/Config.h
@@ -212,6 +212,8 @@ struct Config {
   bool checkSections;
   bool checkDynamicRelocs;
   llvm::DebugCompressionType compressDebugSections;
+  llvm::SmallVector<std::pair<llvm::GlobPattern, llvm::DebugCompressionType>, 0>
+      compressSections;
   bool cref;
   llvm::SmallVector<std::pair<llvm::GlobPattern, uint64_t>, 0>
       deadRelocInNonAlloc;

diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
@@ -1431,6 +1431,23 @@ static void readConfigs(opt::InputArgList &args) {
     }
   }
 
+  for (opt::Arg *arg : args.filtered(OPT_compress_sections)) {
+    SmallVector<StringRef, 0> fields;
+    StringRef(arg->getValue()).split(fields, '=');
+    if (fields.size() != 2 || fields[1].empty()) {
+      error(arg->getSpelling() +
+            ": parse error, not 'section-glob=[zlib|zstd]'");
+      continue;
+    }
+    auto type = getCompressionType(fields[1], arg->getSpelling());
+    if (Expected<GlobPattern> pat = GlobPattern::create(fields[0])) {
+      config->compressSections.emplace_back(std::move(*pat), type);
+    } else {
+      error(arg->getSpelling() + ": " + toString(pat.takeError()));
+      continue;
+    }
+  }
+
   for (opt::Arg *arg : args.filtered(OPT_z)) {
     std::pair<StringRef, StringRef> option =
         StringRef(arg->getValue()).split('=');

diff --git a/lld/ELF/LinkerScript.cpp b/lld/ELF/LinkerScript.cpp
@@ -1033,6 +1033,8 @@ void LinkerScript::assignOffsets(OutputSection *sec) {
     if (sec == findFirstSection(l))
       l->lmaOffset = state->lmaOffset;
 
+  const uint64_t savedDot2 = dot;
+  const size_t savedSize = sec->size;
   // We can call this method multiple times during the creation of
   // thunks and want to start over calculation each time.
   sec->size = 0;
@@ -1074,6 +1076,15 @@ void LinkerScript::assignOffsets(OutputSection *sec) {
     }
   }
 
+  // See the comment in finalizeAddressDependentContent.
+  if (sec->compressed.shards) {
+    if (sec->size != sec->compressed.uncompressedSize)
+      fatal("uncompressed size of SHF_COMPRESSED section '" + sec->name +
+            "' is dependent on linker script commands");
+    sec->size = savedSize;
+    dot = savedDot2 + savedSize;
+  }
+
   // Non-SHF_ALLOC sections do not affect the addresses of other OutputSections
   // as they are not part of the process image.
   if (!(sec->flags & SHF_ALLOC)) {

diff --git a/lld/ELF/Options.td b/lld/ELF/Options.td
@@ -64,6 +64,9 @@ defm compress_debug_sections:
   Eq<"compress-debug-sections", "Compress DWARF debug sections">,
   MetaVarName<"[none,zlib,zstd]">;
 
+defm compress_sections: EEq<"compress-sections", "Compress output sections matching <section-glob>">,
+  MetaVarName<"<section-glob>=[zlib|zstd]">;
+
 defm defsym: Eq<"defsym", "Define a symbol alias">, MetaVarName<"<symbol>=<value>">;
 
 defm optimize_bb_jumps: BB<"optimize-bb-jumps",

diff --git a/lld/ELF/OutputSections.cpp b/lld/ELF/OutputSections.cpp
@@ -330,8 +330,14 @@ template <class ELFT> void OutputSection::maybeCompress() {
   (void)sizeof(Elf_Chdr);
 
   // Compress only DWARF debug sections.
-  if (config->compressDebugSections == DebugCompressionType::None ||
-      (flags & SHF_ALLOC) || !name.starts_with(".debug_") || size == 0)
+  DebugCompressionType type = DebugCompressionType::None;
+  for (auto &[glob, t] : config->compressSections)
+    if (glob.match(name))
+      type = t;
+  if (config->compressDebugSections != DebugCompressionType::None &&
+      !(flags & SHF_ALLOC) && name.starts_with(".debug_") && size)
+    type = config->compressDebugSections;
+  if (type == DebugCompressionType::None)
     return;
 
   llvm::TimeTraceScope timeScope("Compress debug sections");
@@ -347,9 +353,10 @@ template <class ELFT> void OutputSection::maybeCompress() {
   // Use ZSTD's streaming compression API which permits parallel workers working
   // on the stream. See http://facebook.github.io/zstd/zstd_manual.html
   // "Streaming compression - HowTo".
-  if (config->compressDebugSections == DebugCompressionType::Zstd) {
+  if (type == DebugCompressionType::Zstd) {
     // Allocate a buffer of half of the input size, and grow it by 1.5x if
     // insufficient.
+    compressed.type = ELFCOMPRESS_ZSTD;
     compressed.shards = std::make_unique<SmallVector<uint8_t, 0>[]>(1);
     SmallVector<uint8_t, 0> &out = compressed.shards[0];
     out.resize_for_overwrite(std::max<size_t>(size / 2, 32));
@@ -422,6 +429,7 @@ template <class ELFT> void OutputSection::maybeCompress() {
   }
   size += 4; // checksum
 
+  compressed.type = ELFCOMPRESS_ZLIB;
   compressed.shards = std::move(shardsOut);
   compressed.numShards = numShards;
   compressed.checksum = checksum;
@@ -453,15 +461,14 @@ void OutputSection::writeTo(uint8_t *buf, parallel::TaskGroup &tg) {
   // just write it down.
   if (compressed.shards) {
     auto *chdr = reinterpret_cast<typename ELFT::Chdr *>(buf);
+    chdr->ch_type = compressed.type;
     chdr->ch_size = compressed.uncompressedSize;
     chdr->ch_addralign = addralign;
     buf += sizeof(*chdr);
-    if (config->compressDebugSections == DebugCompressionType::Zstd) {
-      chdr->ch_type = ELFCOMPRESS_ZSTD;
+    if (compressed.type == ELFCOMPRESS_ZSTD) {
       memcpy(buf, compressed.shards[0].data(), compressed.shards[0].size());
       return;
     }
-    chdr->ch_type = ELFCOMPRESS_ZLIB;
 
     // Compute shard offsets.
     auto offsets = std::make_unique<size_t[]>(compressed.numShards);

diff --git a/lld/ELF/OutputSections.h b/lld/ELF/OutputSections.h
@@ -23,6 +23,7 @@ struct PhdrEntry;
 
 struct CompressedData {
   std::unique_ptr<SmallVector<uint8_t, 0>[]> shards;
+  uint32_t type = 0;
   uint32_t numShards = 0;
   uint32_t checksum = 0;
   uint64_t uncompressedSize;
@@ -116,12 +117,13 @@ class OutputSection final : public SectionBase {
   void sortInitFini();
   void sortCtorsDtors();
 
+  // Used for implementation of --compress-debug-sections and
+  // --compress-sections.
+  CompressedData compressed;
+
 private:
   SmallVector<InputSection *, 0> storage;
 
-  // Used for implementation of --compress-debug-sections option.
-  CompressedData compressed;
-
   std::array<uint8_t, 4> getFiller();
 };
 

diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
@@ -539,8 +539,6 @@ template <class ELFT> void Writer<ELFT>::run() {
 
   // If --compressed-debug-sections is specified, compress .debug_* sections.
   // Do it right now because it changes the size of output sections.
-  for (OutputSection *sec : outputSections)
-    sec->maybeCompress<ELFT>();
 
   if (script->hasSectionsCommand)
     script->allocateHeaders(mainPart->phdrs);
@@ -1620,6 +1618,14 @@ template <class ELFT> void Writer<ELFT>::finalizeAddressDependentContent() {
   if (config->emachine == EM_HEXAGON)
     hexagonTLSSymbolUpdate(outputSections);
 
+  // Compress SHF_COMPRESSED sections using assignAddresses computed content and
+  // sizes. If there are data commands with changed values, the compressed
+  // content will be invalid, but we do not detect the case. If the section has
+  // changes size due to future assignAddresses calls, we will report an error
+  // in assignOffsets.
+  for (OutputSection *sec : outputSections)
+    sec->maybeCompress<ELFT>();
+
   uint32_t pass = 0, assignPasses = 0;
   for (;;) {
     bool changed = target->needsThunks ? tc.createThunks(pass, outputSections)

diff --git a/lld/docs/ld.lld.1 b/lld/docs/ld.lld.1
@@ -148,6 +148,8 @@ to set the compression level to 6.
 The compression level is 5.
 .El
 .Pp
+.It Fl -compress-sections Ns = Ns Ar section-glob=[zlib|zstd]
+Compress output sections matching the glob with zlib or zstd.
 .It Fl -cref
 Output cross reference table. If
 .Fl Map

diff --git a/lld/test/ELF/compress-sections-err.s b/lld/test/ELF/compress-sections-err.s
@@ -5,8 +5,11 @@
 # RUN: ld.lld %t.o --compress-debug-sections=zlib --compress-debug-sections=none -o /dev/null 2>&1 | count 0
 # RUN: not ld.lld %t.o --compress-debug-sections=zlib -o /dev/null 2>&1 | \
 # RUN:   FileCheck %s --implicit-check-not=error:
+# RUN: not ld.lld %t.o --compress-sections=foo=zlib -o /dev/null 2>&1 | \
+# RUN:   FileCheck %s --check-prefix=CHECK2 --implicit-check-not=error:
 
 # CHECK: error: --compress-debug-sections: LLVM was not built with LLVM_ENABLE_ZLIB or did not find zlib at build time
+# CHECK2: error: --compress-sections: LLVM was not built with LLVM_ENABLE_ZLIB or did not find zlib at build time
 
 .globl _start
 _start:
diff --git a/lld/test/ELF/compress-sections.s b/lld/test/ELF/compress-sections.s
@@ -0,0 +1,96 @@
+# REQUIRES: x86, zlib, zstd
+
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %s -o %t.o
+# RUN: ld.lld -pie %t.o -o %t --compress-sections '*0=zlib' --compress-sections '*0=none'
+# RUN: llvm-readelf -Srs %t | FileCheck %s --check-prefix=CHECK1
+
+# CHECK1:      foo0       PROGBITS [[#%x,FOO0:]]    [[#%x,]] [[#%x,]] 00 A   0 0  1
+# CHECK1-NEXT: foo1       PROGBITS [[#%x,FOO1:]]    [[#%x,]] [[#%x,]] 00 A   0 0  1
+# CHECK1-NEXT: .text      PROGBITS [[#%x,TEXT:]]    [[#%x,]] [[#%x,]] 00 AX  0 0  4
+# CHECK1:      write0     PROGBITS [[#%x,WRITE0:]]  [[#%x,]] [[#%x,]] 00 WA  0 0  1
+# CHECK1-NEXT: nonalloc0  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00     0 0  1
+# CHECK1-NEXT: nonalloc1  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00     0 0  1
+# CHECK1-NEXT: .debug_str PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 01 MS  0 0  1
+
+# CHECK1:          Offset          {{.*}} Type              Symbol's Value  Symbol's Name + Addend
+# CHECK1-NEXT: {{0*}}[[#WRITE0]]   {{.*}} R_X86_64_RELATIVE                 [[#TEXT]]
+# CHECK1-NEXT: {{0*}}[[#WRITE0+8]] {{.*}} R_X86_64_RELATIVE                 [[#TEXT]]
+
+# CHECK1: [[#FOO0]]      0 NOTYPE  LOCAL  DEFAULT   [[#]] foo0_sym
+# CHECK1: [[#FOO1]]      0 NOTYPE  LOCAL  DEFAULT   [[#]] foo1_sym
+# CHECK1: [[#FOO0]]      0 NOTYPE  GLOBAL PROTECTED [[#]] __start_foo0
+# CHECK1: [[#FOO1]]      0 NOTYPE  GLOBAL PROTECTED [[#]] __stop_foo0
+
+# RUN: ld.lld -pie %t.o -o %t --compress-sections '*0=zlib' --compress-sections .debug_str=zstd
+# RUN: llvm-readelf -Srs -x foo0 -x write0 -x nonalloc0 -x .debug_str %t | FileCheck %s --check-prefix=CHECK2
+
+# CHECK2:      foo0       PROGBITS [[#%x,FOO0:]]    [[#%x,]] [[#%x,]] 00 AC  0 0  1
+# CHECK2-NEXT: foo1       PROGBITS [[#%x,FOO1:]]    [[#%x,]] [[#%x,]] 00 A   0 0  1
+# CHECK2-NEXT: .text      PROGBITS [[#%x,TEXT:]]    [[#%x,]] [[#%x,]] 00 AX  0 0  4
+# CHECK2:      write0     PROGBITS [[#%x,WRITE0:]]  [[#%x,]] [[#%x,]] 00 WAC 0 0  1
+# CHECK2-NEXT: nonalloc0  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00 C   0 0  1
+# CHECK2-NEXT: nonalloc1  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00     0 0  1
+# CHECK2-NEXT: .debug_str PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 01 MSC 0 0  1
+
+# CHECK2:          Offset          {{.*}} Type              Symbol's Value  Symbol's Name + Addend
+# CHECK2-NEXT: {{0*}}[[#WRITE0]]   {{.*}} R_X86_64_RELATIVE                 [[#TEXT]]
+# CHECK2-NEXT: {{0*}}[[#WRITE0+8]] {{.*}} R_X86_64_RELATIVE                 [[#TEXT]]
+
+# CHECK2:      Hex dump of section 'foo0':
+## zlib with ch_size=0x10
+# CHECK2-NEXT: 01000000 00000000 10000000 00000000
+# CHECK2-NEXT: 01000000 00000000 {{.*}}
+# CHECK2:      Hex dump of section 'write0':
+## zlib with ch_size=0x10
+# CHECK2-NEXT: 01000000 00000000 10000000 00000000
+# CHECK2-NEXT: 01000000 00000000 {{.*}}
+# CHECK2:      Hex dump of section 'nonalloc0':
+## zlib with ch_size=0x10
+# CHECK2-NEXT: 01000000 00000000 10000000 00000000
+# CHECK2-NEXT: 01000000 00000000 {{.*}}
+# CHECK2:      Hex dump of section '.debug_str':
+## zstd with ch_size=0x38
+# CHECK2-NEXT: 02000000 00000000 38000000 00000000
+# CHECK2-NEXT: 01000000 00000000 {{.*}}
+
+# RUN: not ld.lld --compress-sections=foo %t.o -o /dev/null 2>&1 | \
+# RUN:   FileCheck %s --check-prefix=ERR1 --implicit-check-not=error:
+# ERR1:      error: --compress-sections: parse error, not 'section-glob=[zlib|zstd]'
+
+# RUN: not ld.lld --compress-sections 'a[=zlib' %t.o -o /dev/null 2>&1 | \
+# RUN:   FileCheck %s --check-prefix=ERR2 --implicit-check-not=error:
+# ERR2:      error: --compress-sections: invalid glob pattern: a[
+
+# RUN: not ld.lld %t.o -o /dev/null --compress-sections='.debug*=zlib-gabi' --compress-sections='.debug*=' 2>&1 | \
+# RUN:   FileCheck -check-prefix=ERR3 %s
+# ERR3:      unknown --compress-sections value: zlib-gabi
+# ERR3-NEXT: --compress-sections: parse error, not 'section-glob=[zlib|zstd]'
+
+.globl _start
+_start:
+  leaq __start_foo0(%rip), %rax
+  leaq __stop_foo0(%rip), %rax
+  ret
+
+.section foo0,"a"
+foo0_sym:
+.quad .text-.
+.quad .text-.
+.section foo1,"a"
+foo1_sym:
+.quad .text-.
+.quad .text-.
+.section write0,"aw"
+.quad .text
+.quad .text
+.section nonalloc0,""
+.quad .text
+.quad .text
+.section nonalloc1,""
+.quad 42
+
+.section .debug_str,"MS",@progbits,1
+.Linfo_string0:
+  .asciz "AAAAAAAAAAAAAAAAAAAAAAAAAAA"
+.Linfo_string1:
+  .asciz "BBBBBBBBBBBBBBBBBBBBBBBBBBB"
diff --git a/lld/test/ELF/linkerscript/compress-sections.s b/lld/test/ELF/linkerscript/compress-sections.s
@@ -0,0 +1,53 @@
+# REQUIRES: x86, zlib
+
+# RUN: rm -rf %t && split-file %s %t && cd %t
+# RUN: llvm-mc -filetype=obj -triple=x86_64 a.s -o a.o
+# RUN: not ld.lld -T a.lds a.o --compress-sections 'foo=zlib' 2>&1 | FileCheck %s --check-prefix=ERR --implicit-check-not=error:
+
+# ERR: error: uncompressed size of SHF_COMPRESSED section 'foo' is dependent on linker script commands
+
+# RUN: ld.lld -T b.lds a.o --compress-sections 'foo=zlib' -o a
+# RUN: llvm-readelf -Ss a | FileCheck %s
+
+# CHECK: .text    PROGBITS [[#%x,]]        [[#%x,]] [[#%x,]] 00 AX 0 0  4
+# CHECK: foo      PROGBITS [[#%x,FOO:]]    [[#%x,]] [[#%x,]] 00 AC 0 0  1
+# CHECK: bar      PROGBITS [[#%x,BAR:]]    [[#%x,]] [[#%x,]] 00 A  0 0  1
+
+# CHECK: [[#FOO]]      0 NOTYPE  LOCAL  DEFAULT   [[#]] foo0_sym
+# CHECK: [[#FOO+8]]    0 NOTYPE  LOCAL  DEFAULT   [[#]] foo1_sym
+# CHECK: [[#FOO]]      0 NOTYPE  GLOBAL PROTECTED [[#]] __start_foo
+# CHECK: [[#BAR]]      0 NOTYPE  GLOBAL PROTECTED [[#]] __stop_foo
+
+#--- a.s
+.globl _start
+_start:
+  leaq __start_foo(%rip), %rax
+  leaq __stop_foo(%rip), %rax
+  ret
+
+.section foo0,"a"
+foo0_sym:
+.quad 42
+.section foo1,"a"
+foo1_sym:
+.quad 42
+.section bar,"a"
+.quad 42
+
+#--- a.lds
+SECTIONS {
+  foo : { *(foo*) . += a; }
+  .text : { *(.text) }
+  a = b+1;
+  b = c+1;
+  c = SIZEOF(.text);
+}
+
+#--- b.lds
+SECTIONS {
+  .text : { *(.text) }
+  c = SIZEOF(.text);
+  b = c+1;
+  a = b+1;
+  foo : { *(foo*) QUAD(SIZEOF(foo)) . += a; }
+}