2 files changed, 27 insertions, 28 deletions

diff --git a/lld/COFF/Chunks.cpp b/lld/COFF/Chunks.cpp
index cce68467f1b..c8b2262e330 100644
--- a/lld/COFF/Chunks.cpp
+++ b/lld/COFF/Chunks.cpp
@@ -30,9 +30,21 @@ SectionChunk::SectionChunk(ObjectFile *F, const coff_section *H, uint32_t SI)
     : File(F), Header(H), SectionIndex(SI) {
   // Initialize SectionName.
   File->getCOFFObj()->getSectionName(Header, SectionName);
+
   // Bit [20:24] contains section alignment.
   unsigned Shift = ((Header->Characteristics & 0xF00000) >> 20) - 1;
   Align = uint32_t(1) << Shift;
+
+  // When a new chunk is created, we don't if if it's going to make it
+  // to the final output. Initially all sections are unmarked in terms
+  // of garbage collection. The writer will call markLive() to mark
+  // all reachable section chunks.
+  Live = false;
+
+  // COMDAT sections are not GC root. Non-text sections are not
+  // subject of garbage collection (thus they are root).
+  if (!isCOMDAT() && !(Header->Characteristics & IMAGE_SCN_CNT_CODE))
+    Root = true;
 }
 
 void SectionChunk::writeTo(uint8_t *Buf) {
@@ -50,22 +62,7 @@ void SectionChunk::writeTo(uint8_t *Buf) {
   }
 }
 
-// Returns true if this chunk should be considered as a GC root.
-bool SectionChunk::isRoot() {
-  // COMDAT sections are live only when they are referenced by something else.
-  if (isCOMDAT())
-    return false;
-
-  // Associative sections are live if their parent COMDATs are live,
-  // and vice versa, so they are not considered live by themselves.
-  if (IsAssocChild)
-    return false;
-
-  // Only code is subject of dead-stripping.
-  return !(Header->Characteristics & IMAGE_SCN_CNT_CODE);
-}
-
-void SectionChunk::markLive() {
+void SectionChunk::mark() {
   if (Live)
     return;
   Live = true;
@@ -84,8 +81,10 @@ void SectionChunk::markLive() {
 }
 
 void SectionChunk::addAssociative(SectionChunk *Child) {
-  Child->IsAssocChild = true;
   AssocChildren.push_back(Child);
+  // Associative sections are live if their parent COMDATs are live,
+  // and vice versa, so they are not considered live by themselves.
+  Child->Root = false;
 }
 
 static void add16(uint8_t *P, int16_t V) { write16le(P, read16le(P) + V); }
diff --git a/lld/COFF/Chunks.h b/lld/COFF/Chunks.h
index 6275942a813..10ad33298c2 100644
--- a/lld/COFF/Chunks.h
+++ b/lld/COFF/Chunks.h
@@ -83,9 +83,9 @@ public:
   virtual bool isCOMDAT() const { return false; }
 
   // Used by the garbage collector.
-  virtual bool isRoot() { return false; }
-  virtual bool isLive() { return true; }
-  virtual void markLive() {}
+  bool isRoot() { return Root; }
+  bool isLive() { return Live; }
+  void markLive() { if (!Live) mark(); }
 
   // An output section has pointers to chunks in the section, and each
   // chunk has a back pointer to an output section.
@@ -99,11 +99,16 @@ protected:
   // The offset from beginning of the output file. The writer sets a value.
   uint64_t FileOff = 0;
 
+  // The output section for this chunk.
+  OutputSection *Out = nullptr;
+
   // The alignment of this chunk. The writer uses the value.
   uint32_t Align = 1;
 
-  // The output section for this chunk.
-  OutputSection *Out = nullptr;
+  // Used by the garbage collector.
+  virtual void mark() {}
+  bool Live = true;
+  bool Root = false;
 };
 
 // A chunk corresponding a section of an input file.
@@ -119,15 +124,12 @@ public:
   void printDiscardedMessage() override;
   bool isCOMDAT() const override;
 
-  bool isRoot() override;
-  void markLive() override;
-  bool isLive() override { return Live; }
-
   // Adds COMDAT associative sections to this COMDAT section. A chunk
   // and its children are treated as a group by the garbage collector.
   void addAssociative(SectionChunk *Child);
 
 private:
+  void mark() override;
   SectionRef getSectionRef();
   void applyReloc(uint8_t *Buf, const coff_relocation *Rel);
 
@@ -137,9 +139,7 @@ private:
   const coff_section *Header;
   uint32_t SectionIndex;
   StringRef SectionName;
-  bool Live = false;
   std::vector<Chunk *> AssocChildren;
-  bool IsAssocChild = false;
 };
 
 // A chunk for common symbols. Common chunks don't have actual data.