[ELF] Add --compress-section to compress matched non-SHF_ALLOC sections

--compress-sections <section-glib>=[none|zlib|zstd] is similar to
--compress-debug-sections but applies to broader sections without the
SHF_ALLOC flag. lld will report an error if a SHF_ALLOC section is
matched. An interesting use case is to compress `.strtab`/`.symtab`,
which consume a significant portion of the file size (15.1% for a
release build of Clang).

An older revision is available at https://reviews.llvm.org/D154641 .
This patch focuses on non-allocated sections for safety. Moving
`maybeCompress` as D154641 does not handle STT_SECTION symbols for
`-r --compress-debug-sections=zlib` (see `relocatable-section-symbol.s`
from #66804).

Since different output sections may use different compression
algorithms, we need CompressedData::type to generalize
config->compressDebugSections.

GNU ld feature request: https://sourceware.org/bugzilla/show_bug.cgi?id=27452

Link: https://discourse.llvm.org/t/rfc-compress-arbitrary-sections-with-ld-lld-compress-sections/71674

Pull Request: https://github.com/llvm/llvm-project/pull/84855

11 files changed, 259 insertions, 19 deletions

diff --git a/lld/ELF/Config.h b/lld/ELF/Config.h
index 691ebfc..9ae01eb 100644
--- a/lld/ELF/Config.h
+++ b/lld/ELF/Config.h
@@ -222,7 +222,9 @@ struct Config {
   CGProfileSortKind callGraphProfileSort;
   bool checkSections;
   bool checkDynamicRelocs;
-  llvm::DebugCompressionType compressDebugSections;
+  std::optional<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
index de4b2e3..2439d14 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -1224,9 +1224,10 @@ static void readConfigs(opt::InputArgList &args) {
   config->checkSections =
       args.hasFlag(OPT_check_sections, OPT_no_check_sections, true);
   config->chroot = args.getLastArgValue(OPT_chroot);
-  config->compressDebugSections = getCompressionType(
-      args.getLastArgValue(OPT_compress_debug_sections, "none"),
-      "--compress-debug-sections");
+  if (auto *arg = args.getLastArg(OPT_compress_debug_sections)) {
+    config->compressDebugSections =
+        getCompressionType(arg->getValue(), "--compress-debug-sections");
+  }
   config->cref = args.hasArg(OPT_cref);
   config->optimizeBBJumps =
       args.hasFlag(OPT_optimize_bb_jumps, OPT_no_optimize_bb_jumps, false);
@@ -1516,6 +1517,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=[none|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/Options.td b/lld/ELF/Options.td
index c10a73e..3819b86 100644
--- a/lld/ELF/Options.td
+++ b/lld/ELF/Options.td
@@ -67,6 +67,10 @@ defm compress_debug_sections:
   Eq<"compress-debug-sections", "Compress DWARF debug sections">,
   MetaVarName<"[none,zlib,zstd]">;
 
+defm compress_sections: EEq<"compress-sections",
+  "Compress non-SHF_ALLOC output sections matching <section-glob>">,
+  MetaVarName<"<section-glob>=[none|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
index ee93741..55e6a14 100644
--- a/lld/ELF/OutputSections.cpp
+++ b/lld/ELF/OutputSections.cpp
@@ -326,17 +326,30 @@ static SmallVector<uint8_t, 0> deflateShard(ArrayRef<uint8_t> in, int level,
 }
 #endif
 
-// Compress section contents if this section contains debug info.
+// Compress certain non-SHF_ALLOC sections:
+//
+// * (if --compress-debug-sections is specified) non-empty .debug_* sections
+// * (if --compress-sections is specified) matched sections
 template <class ELFT> void OutputSection::maybeCompress() {
   using Elf_Chdr = typename ELFT::Chdr;
   (void)sizeof(Elf_Chdr);
 
-  // Compress only DWARF debug sections.
-  if (config->compressDebugSections == DebugCompressionType::None ||
-      (flags & SHF_ALLOC) || !name.starts_with(".debug_") || size == 0)
+  DebugCompressionType ctype = DebugCompressionType::None;
+  for (auto &[glob, t] : config->compressSections)
+    if (glob.match(name))
+      ctype = t;
+  if (!(flags & SHF_ALLOC) && config->compressDebugSections &&
+      name.starts_with(".debug_") && size)
+    ctype = *config->compressDebugSections;
+  if (ctype == DebugCompressionType::None)
+    return;
+  if (flags & SHF_ALLOC) {
+    errorOrWarn("--compress-sections: section '" + name +
+                "' with the SHF_ALLOC flag cannot be compressed");
     return;
+  }
 
-  llvm::TimeTraceScope timeScope("Compress debug sections");
+  llvm::TimeTraceScope timeScope("Compress sections");
   compressed.uncompressedSize = size;
   auto buf = std::make_unique<uint8_t[]>(size);
   // Write uncompressed data to a temporary zero-initialized buffer.
@@ -344,14 +357,21 @@ template <class ELFT> void OutputSection::maybeCompress() {
     parallel::TaskGroup tg;
     writeTo<ELFT>(buf.get(), tg);
   }
+  // The generic ABI specifies "The sh_size and sh_addralign fields of the
+  // section header for a compressed section reflect the requirements of the
+  // compressed section." However, 1-byte alignment has been wildly accepted
+  // and utilized for a long time. Removing alignment padding is particularly
+  // useful when there are many compressed output sections.
+  addralign = 1;
 
 #if LLVM_ENABLE_ZSTD
   // 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 (ctype == 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));
@@ -424,6 +444,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;
@@ -450,20 +471,18 @@ void OutputSection::writeTo(uint8_t *buf, parallel::TaskGroup &tg) {
   if (type == SHT_NOBITS)
     return;
 
-  // If --compress-debug-section is specified and if this is a debug section,
-  // we've already compressed section contents. If that's the case,
-  // just write it down.
+  // If the section is compressed due to
+  // --compress-debug-section/--compress-sections, the content is already known.
   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
index c793147..421a018 100644
--- 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 @@ public:
   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/docs/ReleaseNotes.rst b/lld/docs/ReleaseNotes.rst
index 6f60efd..97ed060 100644
--- a/lld/docs/ReleaseNotes.rst
+++ b/lld/docs/ReleaseNotes.rst
@@ -26,6 +26,10 @@ Non-comprehensive list of changes in this release
 ELF Improvements
 ----------------
 
+* ``--compress-sections <section-glib>=[none|zlib|zstd]`` is added to compress
+  matched output sections without the ``SHF_ALLOC`` flag.
+  (`#84855 <https://github.com/llvm/llvm-project/pull/84855>`_)
+
 Breaking changes
 ----------------
 
diff --git a/lld/docs/ld.lld.1 b/lld/docs/ld.lld.1
index e4d39e4..e759776 100644
--- a/lld/docs/ld.lld.1
+++ b/lld/docs/ld.lld.1
@@ -164,6 +164,10 @@ to set the compression level to 6.
 The compression level is 5.
 .El
 .Pp
+.It Fl -compress-sections Ns = Ns Ar section-glob=[none|zlib|zstd]
+Compress output sections that match the glob and do not have the SHF_ALLOC flag.
+This is like a generalized
+.Cm --compress-debug-sections.
 .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
index 0978038..1b46aea 100644
--- 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-special.s b/lld/test/ELF/compress-sections-special.s
new file mode 100644
index 0000000..80c61fe
--- /dev/null
+++ b/lld/test/ELF/compress-sections-special.s
@@ -0,0 +1,31 @@
+# REQUIRES: x86, zlib
+
+# RUN: rm -rf %t && mkdir %t && cd %t
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %s -o a.o
+# RUN: ld.lld -pie a.o --compress-sections .strtab=zlib --compress-sections .symtab=zlib -o out
+# RUN: llvm-readelf -Ss -x .strtab out 2>&1 | FileCheck %s
+
+# CHECK:      nonalloc0  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00     0 0  1
+# CHECK:      .symtab    SYMTAB   0000000000000000 [[#%x,]] [[#%x,]] 18  C 12 3  1
+# CHECK-NEXT: .shstrtab  STRTAB   0000000000000000 [[#%x,]] [[#%x,]] 00     0 0  1
+# CHECK-NEXT: .strtab    STRTAB   0000000000000000 [[#%x,]] [[#%x,]] 00  C  0 0  1
+
+## TODO Add compressed SHT_STRTAB/SHT_SYMTAB support to llvm-readelf
+# CHECK:      warning: {{.*}}: unable to get the string table for the SHT_SYMTAB section: SHT_STRTAB string table section
+
+# CHECK:      Hex dump of section '.strtab':
+# CHECK-NEXT: 01000000 00000000 1a000000 00000000
+# CHECK-NEXT: 01000000 00000000 {{.*}}
+
+# RUN: not ld.lld -shared a.o --compress-sections .dynstr=zlib 2>&1 | FileCheck %s --check-prefix=ERR-ALLOC
+# ERR-ALLOC: error: --compress-sections: section '.dynstr' with the SHF_ALLOC flag cannot be compressed
+
+.globl _start, g0, g1
+_start:
+l0:
+g0:
+g1:
+
+.section nonalloc0,""
+.quad .text+1
+.quad .text+2
diff --git a/lld/test/ELF/compress-sections.s b/lld/test/ELF/compress-sections.s
new file mode 100644
index 0000000..59b5408
--- /dev/null
+++ b/lld/test/ELF/compress-sections.s
@@ -0,0 +1,91 @@
+# REQUIRES: x86, zlib, zstd
+
+# RUN: rm -rf %t && mkdir %t && cd %t
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %s -o a.o
+# RUN: ld.lld -pie a.o -o out --compress-sections '*0=zlib' --compress-sections '*0=none' --compress-sections 'nomatch=none'
+# RUN: llvm-readelf -SrsX out | FileCheck %s --check-prefix=CHECK1
+
+# CHECK1:      Name       Type          Address     Off      Size     ES Flg Lk Inf Al
+# CHECK1:      foo0       PROGBITS [[#%x,FOO0:]]    [[#%x,]] [[#%x,]] 00 A    0   0  8
+# CHECK1-NEXT: foo1       PROGBITS [[#%x,FOO1:]]    [[#%x,]] [[#%x,]] 00 A    0   0  8
+# CHECK1-NEXT: .text      PROGBITS [[#%x,TEXT:]]    [[#%x,]] [[#%x,]] 00 AX   0   0  4
+# CHECK1:      nonalloc0  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00      0   0  8
+# CHECK1-NEXT: nonalloc1  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00      0   0  8
+# CHECK1-NEXT: .debug_str PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 01 MS   0   0  1
+
+# CHECK1: 0000000000000010  0 NOTYPE  LOCAL  DEFAULT   [[#]] (nonalloc0) sym0
+# CHECK1: 0000000000000008  0 NOTYPE  LOCAL  DEFAULT   [[#]] (nonalloc1) sym1
+
+# RUN: ld.lld -pie a.o --compress-sections '*c0=zlib' --compress-sections .debug_str=zstd -o out2
+# RUN: llvm-readelf -SrsX -x nonalloc0 -x .debug_str out2 | FileCheck %s --check-prefix=CHECK2
+
+# CHECK2:      Name       Type          Address     Off      Size     ES Flg Lk Inf Al
+# CHECK2:      foo0       PROGBITS [[#%x,FOO0:]]    [[#%x,]] [[#%x,]] 00 A    0   0  8
+# CHECK2-NEXT: foo1       PROGBITS [[#%x,FOO1:]]    [[#%x,]] [[#%x,]] 00 A    0   0  8
+# CHECK2-NEXT: .text      PROGBITS [[#%x,TEXT:]]    [[#%x,]] [[#%x,]] 00 AX   0   0  4
+# CHECK2:      nonalloc0  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00 C    0   0  1
+# CHECK2-NEXT: nonalloc1  PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00      0   0  8
+# CHECK2-NEXT: .debug_str PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 01 MSC  0   0  1
+
+# CHECK2: 0000000000000010  0 NOTYPE  LOCAL  DEFAULT   [[#]] (nonalloc0) sym0
+# CHECK2: 0000000000000008  0 NOTYPE  LOCAL  DEFAULT   [[#]] (nonalloc1) sym1
+
+# 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 {{.*}}
+
+## --compress-debug-sections=none takes precedence.
+# RUN: ld.lld a.o --compress-debug-sections=none --compress-sections .debug_str=zstd -o out3
+# RUN: llvm-readelf -S out3 | FileCheck %s --check-prefix=CHECK3
+
+# CHECK3:      .debug_str PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 01 MS   0   0  1
+
+# RUN: not ld.lld a.o --compress-sections '*0=zlib' 2>&1 | \
+# RUN:   FileCheck %s --check-prefix=ERR-ALLOC --implicit-check-not=error:
+# ERR-ALLOC: error: --compress-sections: section 'foo0' with the SHF_ALLOC flag cannot be compressed
+
+# RUN: not ld.lld --compress-sections=foo a.o 2>&1 | \
+# RUN:   FileCheck %s --check-prefix=ERR1 --implicit-check-not=error:
+# ERR1:      error: --compress-sections: parse error, not 'section-glob=[none|zlib|zstd]'
+
+# RUN: not ld.lld --compress-sections 'a[=zlib' a.o 2>&1 | \
+# RUN:   FileCheck %s --check-prefix=ERR2 --implicit-check-not=error:
+# ERR2:      error: --compress-sections: invalid glob pattern, unmatched '['
+
+# RUN: not ld.lld a.o --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=[none|zlib|zstd]'
+
+.globl _start
+_start:
+  ret
+
+.section foo0,"a"
+.balign 8
+.quad .text-.
+.quad .text-.
+.section foo1,"a"
+.balign 8
+.quad .text-.
+.quad .text-.
+.section nonalloc0,""
+.balign 8
+.quad .text+1
+.quad .text+2
+sym0:
+.section nonalloc1,""
+.balign 8
+.quad 42
+sym1:
+
+.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
new file mode 100644
index 0000000..9b4574a
--- /dev/null
+++ b/lld/test/ELF/linkerscript/compress-sections.s
@@ -0,0 +1,62 @@
+# 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: ld.lld -T a.lds a.o --compress-sections nonalloc=zlib --compress-sections str=zlib -o out
+# RUN: llvm-readelf -SsXz -p str out | FileCheck %s
+
+# CHECK:      Name     Type            Address   Off      Size     ES Flg  Lk Inf Al
+# CHECK:      nonalloc PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 00   C   0   0  1
+# CHECK-NEXT: str      PROGBITS 0000000000000000 [[#%x,]] [[#%x,]] 01 MSC   0   0  1
+
+# CHECK:      0000000000000000  0 NOTYPE  GLOBAL DEFAULT [[#]] (nonalloc) nonalloc_start
+# CHECK:      0000000000000023  0 NOTYPE  GLOBAL DEFAULT [[#]] (nonalloc) nonalloc_end
+# CHECK:      String dump of section 'str':
+# CHECK-NEXT: [     0] AAA
+# CHECK-NEXT: [     4] BBB
+
+## TODO The uncompressed size of 'nonalloc' is dependent on linker script
+## commands, which is not handled. We should report an error.
+# RUN: ld.lld -T b.lds a.o --compress-sections nonalloc=zlib
+
+#--- a.s
+.globl _start
+_start:
+  ret
+
+.section nonalloc0,""
+.balign 8
+.quad .text
+.quad .text
+.section nonalloc1,""
+.balign 8
+.quad 42
+
+.section str,"MS",@progbits,1
+  .asciz "AAA"
+  .asciz "BBB"
+
+#--- a.lds
+SECTIONS {
+  .text : { *(.text) }
+  c = SIZEOF(.text);
+  b = c+1;
+  a = b+1;
+  nonalloc : {
+    nonalloc_start = .;
+## In general, using data commands is error-prone. This case is correct, though.
+    *(nonalloc*) QUAD(SIZEOF(.text))
+    . += a;
+    nonalloc_end = .;
+  }
+  str : { *(str) }
+}
+
+#--- b.lds
+SECTIONS {
+  nonalloc : { *(nonalloc*) . += a; }
+  .text : { *(.text) }
+  a = b+1;
+  b = c+1;
+  c = SIZEOF(.text);
+}