4 files changed, 19 insertions, 13 deletions
diff --git a/lld/ELF/LinkerScript.cpp b/lld/ELF/LinkerScript.cpp
index 161909abf00..c303f0524ad 100644
--- a/lld/ELF/LinkerScript.cpp
+++ b/lld/ELF/LinkerScript.cpp
@@ -440,9 +440,6 @@ void LinkerScript::fabricateDefaultCommands() {
   // For each OutputSection that needs a VA fabricate an OutputSectionCommand
   // with an InputSectionDescription describing the InputSections
   for (OutputSection *Sec : *OutputSections) {
-    if (!(Sec->Flags & SHF_ALLOC))
-      continue;
-
     auto *OSCmd = make<OutputSectionCommand>(Sec->Name);
     OSCmd->Sec = Sec;
     SecToCommand[Sec] = OSCmd;
diff --git a/lld/ELF/MapFile.cpp b/lld/ELF/MapFile.cpp
index c3f4f5a4848..7b82eceba02 100644
--- a/lld/ELF/MapFile.cpp
+++ b/lld/ELF/MapFile.cpp
@@ -21,6 +21,7 @@
 
 #include "MapFile.h"
 #include "InputFiles.h"
+#include "LinkerScript.h"
 #include "OutputSections.h"
 #include "Strings.h"
 #include "SymbolTable.h"
@@ -99,7 +100,7 @@ getSymbolStrings(ArrayRef<DefinedRegular *> Syms) {
 }
 
 template <class ELFT>
-void elf::writeMapFile(ArrayRef<OutputSection *> OutputSections) {
+void elf::writeMapFile(llvm::ArrayRef<BaseCommand *> Script) {
   if (Config->MapFile.empty())
     return;
 
@@ -122,7 +123,11 @@ void elf::writeMapFile(ArrayRef<OutputSection *> OutputSections) {
      << " Align Out     In      Symbol\n";
 
   // Print out file contents.
-  for (OutputSection *OSec : OutputSections) {
+  for (BaseCommand *Base : Script) {
+    auto *Cmd = dyn_cast<OutputSectionCommand>(Base);
+    if (!Cmd)
+      continue;
+    OutputSection *OSec = Cmd->Sec;
     writeHeader<ELFT>(OS, OSec->Addr, OSec->Size, OSec->Alignment);
     OS << OSec->Name << '\n';
 
@@ -137,7 +142,7 @@ void elf::writeMapFile(ArrayRef<OutputSection *> OutputSections) {
   }
 }
 
-template void elf::writeMapFile<ELF32LE>(ArrayRef<OutputSection *>);
-template void elf::writeMapFile<ELF32BE>(ArrayRef<OutputSection *>);
-template void elf::writeMapFile<ELF64LE>(ArrayRef<OutputSection *>);
-template void elf::writeMapFile<ELF64BE>(ArrayRef<OutputSection *>);
+template void elf::writeMapFile<ELF32LE>(ArrayRef<BaseCommand *>);
+template void elf::writeMapFile<ELF32BE>(ArrayRef<BaseCommand *>);
+template void elf::writeMapFile<ELF64LE>(ArrayRef<BaseCommand *>);
+template void elf::writeMapFile<ELF64BE>(ArrayRef<BaseCommand *>);
diff --git a/lld/ELF/MapFile.h b/lld/ELF/MapFile.h
index 89d4d2d4119..f50ef00061f 100644
--- a/lld/ELF/MapFile.h
+++ b/lld/ELF/MapFile.h
@@ -14,9 +14,8 @@
 
 namespace lld {
 namespace elf {
-class OutputSection;
-template <class ELFT>
-void writeMapFile(llvm::ArrayRef<OutputSection *> OutputSections);
+struct BaseCommand;
+template <class ELFT> void writeMapFile(llvm::ArrayRef<BaseCommand *> Script);
 }
 }
 
diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
index bc9eee37772..7a21d2b9f13 100644
--- a/lld/ELF/Writer.cpp
+++ b/lld/ELF/Writer.cpp
@@ -288,8 +288,13 @@ template <class ELFT> void Writer<ELFT>::run() {
   if (ErrorCount)
     return;
 
+  // Clear the OutputSections to make sure it is not used anymore. Any
+  // code from this point on should be using the linker script
+  // commands.
+  OutputSections.clear();
+
   // Handle -Map option.
-  writeMapFile<ELFT>(OutputSections);
+  writeMapFile<ELFT>(Script->Opt.Commands);
   if (ErrorCount)
     return;