5 files changed, 27 insertions, 40 deletions

diff --git a/lld/ELF/MapFile.cpp b/lld/ELF/MapFile.cpp
index dcc829315e6..9f3b1eee6b0 100644
--- a/lld/ELF/MapFile.cpp
+++ b/lld/ELF/MapFile.cpp
@@ -137,17 +137,11 @@ void elf::writeMapFile() {
     OS << OSec->Name << '\n';
 
     // Dump symbols for each input section.
-    for (BaseCommand *Base : OSec->SectionCommands) {
-      auto *ISD = dyn_cast<InputSectionDescription>(Base);
-      if (!ISD)
-        continue;
-      for (InputSection *IS : ISD->Sections) {
-        writeHeader(OS, OSec->Addr + IS->OutSecOff, IS->getSize(),
-                    IS->Alignment);
-        OS << indent(1) << toString(IS) << '\n';
-        for (Symbol *Sym : SectionSyms[IS])
-          OS << SymStr[Sym] << '\n';
-      }
+    for (InputSection *IS : getInputSections(OSec)) {
+      writeHeader(OS, OSec->Addr + IS->OutSecOff, IS->getSize(), IS->Alignment);
+      OS << indent(1) << toString(IS) << '\n';
+      for (Symbol *Sym : SectionSyms[IS])
+        OS << SymStr[Sym] << '\n';
     }
   }
 }
diff --git a/lld/ELF/OutputSections.cpp b/lld/ELF/OutputSections.cpp
index 59fc5bd4bd3..6a893f03e85 100644
--- a/lld/ELF/OutputSections.cpp
+++ b/lld/ELF/OutputSections.cpp
@@ -231,12 +231,7 @@ template <class ELFT> void OutputSection::writeTo(uint8_t *Buf) {
   }
 
   // Write leading padding.
-  std::vector<InputSection *> Sections;
-  for (BaseCommand *Cmd : SectionCommands)
-    if (auto *ISD = dyn_cast<InputSectionDescription>(Cmd))
-      for (InputSection *IS : ISD->Sections)
-        if (IS->Live)
-          Sections.push_back(IS);
+  std::vector<InputSection *> Sections = getInputSections(this);
   uint32_t Filler = getFiller();
   if (Filler)
     fill(Buf, Sections.empty() ? Size : Sections[0]->OutSecOff, Filler);
@@ -281,17 +276,13 @@ static void finalizeShtGroup(OutputSection *OS,
 }
 
 template <class ELFT> void OutputSection::finalize() {
-  InputSection *First = nullptr;
-  for (BaseCommand *Base : SectionCommands) {
-    if (auto *ISD = dyn_cast<InputSectionDescription>(Base)) {
-      if (ISD->Sections.empty())
-        continue;
-      if (First == nullptr)
-        First = ISD->Sections.front();
-    }
-    if (isa<ByteCommand>(Base) && Type == SHT_NOBITS)
-      Type = SHT_PROGBITS;
-  }
+  if (Type == SHT_NOBITS)
+    for (BaseCommand *Base : SectionCommands)
+      if (isa<ByteCommand>(Base))
+        Type = SHT_PROGBITS;
+
+  std::vector<InputSection *> V = getInputSections(this);
+  InputSection *First = V.empty() ? nullptr : V[0];
 
   if (Flags & SHF_LINK_ORDER) {
     // We must preserve the link order dependency of sections with the
@@ -394,6 +385,14 @@ int elf::getPriority(StringRef S) {
   return V;
 }
 
+std::vector<InputSection *> elf::getInputSections(OutputSection *OS) {
+  std::vector<InputSection *> Ret;
+  for (BaseCommand *Base : OS->SectionCommands)
+    if (auto *ISD = dyn_cast<InputSectionDescription>(Base))
+      Ret.insert(Ret.end(), ISD->Sections.begin(), ISD->Sections.end());
+  return Ret;
+}
+
 // Sorts input sections by section name suffixes, so that .foo.N comes
 // before .foo.M if N < M. Used to sort .{init,fini}_array.N sections.
 // We want to keep the original order if the priorities are the same
diff --git a/lld/ELF/OutputSections.h b/lld/ELF/OutputSections.h
index 2e2f7d3c9a3..35a873488f2 100644
--- a/lld/ELF/OutputSections.h
+++ b/lld/ELF/OutputSections.h
@@ -120,6 +120,8 @@ private:
 
 int getPriority(StringRef S);
 
+std::vector<InputSection *> getInputSections(OutputSection* OS);
+
 // All output sections that are handled by the linker specially are
 // globally accessible. Writer initializes them, so don't use them
 // until Writer is initialized.
diff --git a/lld/ELF/SyntheticSections.cpp b/lld/ELF/SyntheticSections.cpp
index f699e415f7e..ebf5d7f25e5 100644
--- a/lld/ELF/SyntheticSections.cpp
+++ b/lld/ELF/SyntheticSections.cpp
@@ -2579,12 +2579,9 @@ void ARMExidxSentinelSection::writeTo(uint8_t *Buf) {
 
 // The sentinel has to be removed if there are no other .ARM.exidx entries.
 bool ARMExidxSentinelSection::empty() const {
-  OutputSection *OS = getParent();
-  for (auto *B : OS->SectionCommands)
-    if (auto *ISD = dyn_cast<InputSectionDescription>(B))
-      for (auto *S : ISD->Sections)
-        if (!isa<ARMExidxSentinelSection>(S))
-          return false;
+  for (InputSection *IS : getInputSections(getParent()))
+    if (!isa<ARMExidxSentinelSection>(IS))
+      return false;
   return true;
 }
 
diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
index b58e1bf1ab7..c208a5a6216 100644
--- a/lld/ELF/Writer.cpp
+++ b/lld/ELF/Writer.cpp
@@ -1392,12 +1392,7 @@ static void removeUnusedSyntheticSections() {
 
     // If there are no other alive sections or commands left in the output
     // section description, we remove it from the output.
-    bool IsEmpty = llvm::all_of(OS->SectionCommands, [](BaseCommand *B) {
-      if (auto *ISD = dyn_cast<InputSectionDescription>(B))
-        return ISD->Sections.empty();
-      return false;
-    });
-    if (IsEmpty)
+    if (getInputSections(OS).empty())
       OS->Live = false;
   }
 }