1 files changed, 28 insertions, 15 deletions
diff --git a/ELF/ICF.cpp b/ELF/ICF.cpp
index 8b01d06b0248..dce76f79c9b3 100644
--- a/ELF/ICF.cpp
+++ b/ELF/ICF.cpp
@@ -74,6 +74,8 @@
 
 #include "ICF.h"
 #include "Config.h"
+#include "LinkerScript.h"
+#include "OutputSections.h"
 #include "SymbolTable.h"
 #include "Symbols.h"
 #include "SyntheticSections.h"
@@ -86,12 +88,12 @@
 #include <algorithm>
 #include <atomic>
 
-using namespace lld;
-using namespace lld::elf;
 using namespace llvm;
 using namespace llvm::ELF;
 using namespace llvm::object;
 
+namespace lld {
+namespace elf {
 namespace {
 template <class ELFT> class ICF {
 public:
@@ -304,10 +306,8 @@ bool ICF<ELFT>::equalsConstant(const InputSection *a, const InputSection *b) {
     return false;
 
   // If two sections have different output sections, we cannot merge them.
-  // FIXME: This doesn't do the right thing in the case where there is a linker
-  // script. We probably need to move output section assignment before ICF to
-  // get the correct behaviour here.
-  if (getOutputSectionName(a) != getOutputSectionName(b))
+  assert(a->getParent() && b->getParent());
+  if (a->getParent() != b->getParent())
     return false;
 
   if (a->areRelocsRela)
@@ -446,10 +446,11 @@ static void print(const Twine &s) {
 // The main function of ICF.
 template <class ELFT> void ICF<ELFT>::run() {
   // Collect sections to merge.
-  for (InputSectionBase *sec : inputSections)
-    if (auto *s = dyn_cast<InputSection>(sec))
-      if (isEligible(s))
-        sections.push_back(s);
+  for (InputSectionBase *sec : inputSections) {
+    auto *s = cast<InputSection>(sec);
+    if (isEligible(s))
+      sections.push_back(s);
+  }
 
   // Initially, we use hash values to partition sections.
   parallelForEach(sections, [&](InputSection *s) {
@@ -499,12 +500,24 @@ template <class ELFT> void ICF<ELFT>::run() {
         isec->markDead();
     }
   });
+
+  // InputSectionDescription::sections is populated by processSectionCommands().
+  // ICF may fold some input sections assigned to output sections. Remove them.
+  for (BaseCommand *base : script->sectionCommands)
+    if (auto *sec = dyn_cast<OutputSection>(base))
+      for (BaseCommand *sub_base : sec->sectionCommands)
+        if (auto *isd = dyn_cast<InputSectionDescription>(sub_base))
+          llvm::erase_if(isd->sections,
+                         [](InputSection *isec) { return !isec->isLive(); });
 }
 
 // ICF entry point function.
-template <class ELFT> void elf::doIcf() { ICF<ELFT>().run(); }
+template <class ELFT> void doIcf() { ICF<ELFT>().run(); }
+
+template void doIcf<ELF32LE>();
+template void doIcf<ELF32BE>();
+template void doIcf<ELF64LE>();
+template void doIcf<ELF64BE>();
 
-template void elf::doIcf<ELF32LE>();
-template void elf::doIcf<ELF32BE>();
-template void elf::doIcf<ELF64LE>();
-template void elf::doIcf<ELF64BE>();
+} // namespace elf
+} // namespace lld