From 913914f0f83bfc0ac1a04fe15bf98cb8ed5e118e Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Wed, 26 Jan 2022 10:23:56 -0800
Subject: [PATCH] [ELF] Simplify writing the Elf_Chdr header. NFC

And avoiding changing `size` in `writeTo`.
---
 lld/ELF/OutputSections.cpp | 18 +++++++-----------
 lld/ELF/OutputSections.h   |  2 +-
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/lld/ELF/OutputSections.cpp b/lld/ELF/OutputSections.cpp
index c4c1123..c73d6e4 100644
--- a/lld/ELF/OutputSections.cpp
+++ b/lld/ELF/OutputSections.cpp
@@ -330,13 +330,6 @@ template <class ELFT> void OutputSection::maybeCompress() {
 
   llvm::TimeTraceScope timeScope("Compress debug sections");
 
-  // Create a section header.
-  zDebugHeader.resize(sizeof(Elf_Chdr));
-  auto *hdr = reinterpret_cast<Elf_Chdr *>(zDebugHeader.data());
-  hdr->ch_type = ELFCOMPRESS_ZLIB;
-  hdr->ch_size = size;
-  hdr->ch_addralign = alignment;
-
   // Write uncompressed data to a temporary zero-initialized buffer.
   auto buf = std::make_unique<uint8_t[]>(size);
   writeTo<ELFT>(buf.get());
@@ -369,6 +362,7 @@ template <class ELFT> void OutputSection::maybeCompress() {
 
   // Update section size and combine Alder-32 checksums.
   uint32_t checksum = 1;       // Initial Adler-32 value
+  compressed.uncompressedSize = size;
   size = sizeof(Elf_Chdr) + 2; // Elf_Chdir and zlib header
   for (size_t i = 0; i != numShards; ++i) {
     size += shardsOut[i].size();
@@ -405,9 +399,11 @@ template <class ELFT> void OutputSection::writeTo(uint8_t *buf) {
   // we've already compressed section contents. If that's the case,
   // just write it down.
   if (compressed.shards) {
-    memcpy(buf, zDebugHeader.data(), zDebugHeader.size());
-    buf += zDebugHeader.size();
-    size -= zDebugHeader.size();
+    auto *chdr = reinterpret_cast<typename ELFT::Chdr *>(buf);
+    chdr->ch_type = ELFCOMPRESS_ZLIB;
+    chdr->ch_size = compressed.uncompressedSize;
+    chdr->ch_addralign = alignment;
+    buf += sizeof(*chdr);
 
     // Compute shard offsets.
     auto offsets = std::make_unique<size_t[]>(compressed.numShards);
@@ -422,7 +418,7 @@ template <class ELFT> void OutputSection::writeTo(uint8_t *buf) {
              compressed.shards[i].size());
     });
 
-    write32be(buf + size - 4, compressed.checksum);
+    write32be(buf + (size - sizeof(*chdr) - 4), compressed.checksum);
     return;
   }
 
diff --git a/lld/ELF/OutputSections.h b/lld/ELF/OutputSections.h
index 957e676..ad656e4 100644
--- a/lld/ELF/OutputSections.h
+++ b/lld/ELF/OutputSections.h
@@ -29,6 +29,7 @@ struct CompressedData {
   std::unique_ptr<SmallVector<uint8_t, 0>[]> shards;
   uint32_t numShards = 0;
   uint32_t checksum = 0;
+  uint64_t uncompressedSize;
 };
 
 // This represents a section in an output file.
@@ -118,7 +119,6 @@ public:
 
 private:
   // Used for implementation of --compress-debug-sections option.
-  SmallVector<uint8_t, 0> zDebugHeader;
   CompressedData compressed;
 
   std::array<uint8_t, 4> getFiller();
-- 
2.7.4