3 files changed, 89 insertions, 4 deletions

diff --git a/lld/lib/ReaderWriter/ELF/TargetLayout.cpp b/lld/lib/ReaderWriter/ELF/TargetLayout.cpp
index 87cbcd5cab5..c0dae207bd6 100644
--- a/lld/lib/ReaderWriter/ELF/TargetLayout.cpp
+++ b/lld/lib/ReaderWriter/ELF/TargetLayout.cpp
@@ -320,6 +320,7 @@ template <class ELFT> void TargetLayout<ELFT>::createOutputSections() {
     } else {
       outputSection = new (_allocator.Allocate<OutputSection<ELFT>>())
           OutputSection<ELFT>(section->outputSectionName());
+      checkOutputSectionSegment(outputSection);
       _outputSections.push_back(outputSection);
       outputSectionInsert.first->second = outputSection;
     }
@@ -327,6 +328,17 @@ template <class ELFT> void TargetLayout<ELFT>::createOutputSections() {
   }
 }
 
+// Check that output section has proper segment set
+template <class ELFT>
+void TargetLayout<ELFT>::checkOutputSectionSegment(
+    const OutputSection<ELFT> *sec) {
+  std::vector<const script::PHDR *> phdrs;
+  if (_linkerScriptSema.getPHDRsForOutputSection(sec->name(), phdrs)) {
+    llvm::report_fatal_error(
+        "Linker script has wrong segments set for output sections");
+  }
+}
+
 template <class ELFT>
 uint64_t
 TargetLayout<ELFT>::getLookupSectionFlags(const OutputSection<ELFT> *os) const {
diff --git a/lld/lib/ReaderWriter/ELF/TargetLayout.h b/lld/lib/ReaderWriter/ELF/TargetLayout.h
index 7162cd5fe08..f35bf031146 100644
--- a/lld/lib/ReaderWriter/ELF/TargetLayout.h
+++ b/lld/lib/ReaderWriter/ELF/TargetLayout.h
@@ -209,6 +209,9 @@ public:
   // Output sections with the same name into a OutputSection
   void createOutputSections();
 
+  // Check that output section has proper segment set
+  void checkOutputSectionSegment(const OutputSection<ELFT> *sec);
+
   /// \brief Sort the sections by their order as defined by the layout,
   /// preparing all sections to be assigned to a segment.
   virtual void sortInputSections();
diff --git a/lld/lib/ReaderWriter/LinkerScript.cpp b/lld/lib/ReaderWriter/LinkerScript.cpp
index 6d82e29d95a..79a367d15bf 100644
--- a/lld/lib/ReaderWriter/LinkerScript.cpp
+++ b/lld/lib/ReaderWriter/LinkerScript.cpp
@@ -953,6 +953,9 @@ void PHDR::dump(raw_ostream &os) const {
   os << ";\n";
 }
 
+static PHDR none("NONE", 0, false, false, NULL, 0);
+const PHDR *PHDR::NONE = &none;
+
 void PHDRS::dump(raw_ostream &os) const {
   os << "PHDRS\n{\n";
   for (auto &&phdr : _phdrs) {
@@ -2353,10 +2356,7 @@ Extern *Parser::parseExtern() {
 }
 
 // Sema member functions
-Sema::Sema()
-    : _scripts(), _layoutCommands(), _memberToLayoutOrder(),
-      _memberNameWildcards(), _cacheSectionOrder(), _cacheExpressionOrder(),
-      _deliveredExprs(), _symbolTable() {}
+Sema::Sema() : _parsedPHDRS(false) {}
 
 void Sema::perform() {
   for (auto &parser : _scripts)
@@ -2465,6 +2465,18 @@ uint64_t Sema::getLinkerScriptExprValue(StringRef name) const {
   return it->second;
 }
 
+std::error_code
+Sema::getPHDRsForOutputSection(StringRef name,
+                               std::vector<const PHDR *> &phdrs) const {
+  // Cache results if not done yet.
+  if (auto ec = const_cast<Sema *>(this)->buildSectionToPHDR())
+    return ec;
+
+  auto vec = _sectionToPHDR.lookup(name);
+  std::copy(std::begin(vec), std::end(vec), std::back_inserter(phdrs));
+  return std::error_code();
+}
+
 void Sema::dump() const {
   raw_ostream &os = llvm::outs();
   os << "Linker script semantics dump\n";
@@ -2701,6 +2713,64 @@ bool Sema::localCompare(int order, const SectionKey &lhs,
   return false;
 }
 
+std::error_code Sema::buildSectionToPHDR() {
+  if (_parsedPHDRS)
+    return std::error_code();
+  _parsedPHDRS = true;
+
+  // No scripts - nothing to do.
+  if (_scripts.empty() || _layoutCommands.empty())
+    return std::error_code();
+
+  // Collect all header declarations.
+  llvm::StringMap<const PHDR *> phdrs;
+  for (auto &parser : _scripts) {
+    for (auto *cmd : parser->get()->_commands) {
+      if (auto *ph = dyn_cast<PHDRS>(cmd)) {
+        for (auto *p : *ph)
+          phdrs[p->name()] = p;
+      }
+    }
+  }
+  const bool noPhdrs = phdrs.empty();
+
+  // Add NONE header to the map provided there's no user-defined
+  // header with the same name.
+  if (!_sectionToPHDR.count(PHDR::NONE->name()))
+    phdrs[PHDR::NONE->name()] = PHDR::NONE;
+
+  // Match output sections to available headers.
+  llvm::SmallVector<const PHDR *, 2> phdrsCur, phdrsLast { PHDR::NONE };
+  for (const Command *cmd : _layoutCommands) {
+    auto osd = dyn_cast<OutputSectionDescription>(cmd);
+    if (!osd || osd->isDiscarded())
+      continue;
+
+    phdrsCur.clear();
+    for (StringRef name : osd->PHDRs()) {
+      auto it = phdrs.find(name);
+      if (it == phdrs.end()) {
+        return LinkerScriptReaderError::unknown_phdr_ids;
+      }
+      phdrsCur.push_back(it->second);
+    }
+
+    // If no headers and no errors - insert empty headers set.
+    // If the current set of headers is empty, then use the last non-empty
+    // set. Otherwise mark the current set to be the last non-empty set for
+    // successors.
+    if (noPhdrs)
+      phdrsCur.clear();
+    else if (phdrsCur.empty())
+      phdrsCur = phdrsLast;
+    else
+      phdrsLast = phdrsCur;
+
+    _sectionToPHDR[osd->name()] = phdrsCur;
+  }
+  return std::error_code();
+}
+
 static bool hasWildcard(StringRef name) {
   for (auto ch : name)
     if (ch == '*' || ch == '?' || ch == '[' || ch == '\\')