[ELF/LinkerScript] Support ONLY_IF_{RO, RW} directive.

Differential Revision:   https://reviews.llvm.org/D22660

3 files changed, 69 insertions, 3 deletions

diff --git a/lld/ELF/LinkerScript.cpp b/lld/ELF/LinkerScript.cpp
index 1f6fbbcfca8..33ccf715fe6 100644
--- a/lld/ELF/LinkerScript.cpp
+++ b/lld/ELF/LinkerScript.cpp
@@ -262,15 +262,26 @@ std::vector<OutputSectionBase<ELFT> *>
 LinkerScript<ELFT>::createSections(OutputSectionFactory<ELFT> &Factory) {
   typedef const std::unique_ptr<ObjectFile<ELFT>> ObjectFile;
   std::vector<OutputSectionBase<ELFT> *> Result;
+  DenseSet<OutputSectionBase<ELFT> *> Removed;
 
   // Add input section to output section. If there is no output section yet,
   // then create it and add to output section list.
-  auto AddInputSec = [&](InputSectionBase<ELFT> *C, StringRef Name) {
+  auto AddInputSec = [&](InputSectionBase<ELFT> *C, StringRef Name,
+                         ConstraintKind Constraint) {
     OutputSectionBase<ELFT> *Sec;
     bool IsNew;
     std::tie(Sec, IsNew) = Factory.create(C, Name);
     if (IsNew)
       Result.push_back(Sec);
+    if ((!(C->getSectionHdr()->sh_flags & SHF_WRITE)) &&
+        Constraint == ReadWrite) {
+      Removed.insert(Sec);
+      return;
+    }
+    if ((C->getSectionHdr()->sh_flags & SHF_WRITE) && Constraint == ReadOnly) {
+      Removed.insert(Sec);
+      return;
+    }
     Sec->addSection(C);
   };
 
@@ -296,7 +307,7 @@ LinkerScript<ELFT>::createSections(OutputSectionFactory<ELFT> &Factory) {
             if (OutCmd->Name == "/DISCARD/")
               S->Live = false;
             else
-              AddInputSec(S, OutCmd->Name);
+              AddInputSec(S, OutCmd->Name, OutCmd->Constraint);
           }
         }
       }
@@ -308,11 +319,17 @@ LinkerScript<ELFT>::createSections(OutputSectionFactory<ELFT> &Factory) {
     for (InputSectionBase<ELFT> *S : F->getSections()) {
       if (!isDiscarded(S)) {
         if (!S->OutSec)
-          AddInputSec(S, getOutputSectionName(S));
+          AddInputSec(S, getOutputSectionName(S), NoConstraint);
       } else
         reportDiscarded(S, F);
     }
 
+  // Remove from the output all the sections which did not met the constraints.
+  Result.erase(std::remove_if(Result.begin(), Result.end(),
+                              [&](OutputSectionBase<ELFT> *Sec) {
+                                return Removed.count(Sec);
+                              }),
+               Result.end());
   return Result;
 }
 
@@ -793,6 +810,12 @@ void ScriptParser::readOutputSectionDescription(StringRef OutSec) {
   OutputSectionCommand *Cmd = new OutputSectionCommand(OutSec);
   Opt.Commands.emplace_back(Cmd);
   expect(":");
+
+  // Parse constraints.
+  if (skip("ONLY_IF_RO"))
+    Cmd->Constraint = ReadOnly;
+  if (skip("ONLY_IF_RW"))
+    Cmd->Constraint = ReadWrite;
   expect("{");
 
   while (!Error && !skip("}")) {
diff --git a/lld/ELF/LinkerScript.h b/lld/ELF/LinkerScript.h
index 0af38b4a401..2277cd04cf0 100644
--- a/lld/ELF/LinkerScript.h
+++ b/lld/ELF/LinkerScript.h
@@ -53,6 +53,13 @@ struct SymbolAssignment : BaseCommand {
   std::vector<StringRef> Expr;
 };
 
+// Linker scripts allow additional constraints to be put on ouput sections.
+// An output section will only be created if all of its input sections are
+// read-only
+// or all of its input sections are read-write by using the keyword ONLY_IF_RO
+// and ONLY_IF_RW respectively.
+enum ConstraintKind { NoConstraint, ReadOnly, ReadWrite };
+
 struct OutputSectionCommand : BaseCommand {
   OutputSectionCommand(StringRef Name)
       : BaseCommand(OutputSectionKind), Name(Name) {}
@@ -61,6 +68,7 @@ struct OutputSectionCommand : BaseCommand {
   std::vector<std::unique_ptr<BaseCommand>> Commands;
   std::vector<StringRef> Phdrs;
   std::vector<uint8_t> Filler;
+  ConstraintKind Constraint = NoConstraint;
 };
 
 struct InputSectionDescription : BaseCommand {
diff --git a/lld/test/ELF/linkerscript-sections-constraint.s b/lld/test/ELF/linkerscript-sections-constraint.s
new file mode 100644
index 00000000000..84711e7bff8
--- /dev/null
+++ b/lld/test/ELF/linkerscript-sections-constraint.s
@@ -0,0 +1,35 @@
+# REQUIRES: x86
+# RUN: llvm-mc -filetype=obj -triple=x86_64-unknown-linux %s -o %t
+# RUN: echo "SECTIONS { \
+# RUN:  .writable : ONLY_IF_RW { *(.writable) } \
+# RUN:  .readable : ONLY_IF_RO { *(.readable) }}" > %t.script
+# RUN: ld.lld -o %t1 --script %t.script %t
+# RUN: llvm-objdump -section-headers %t1 | \
+# RUN:   FileCheck -check-prefix=BASE %s
+# BASE: Sections:
+# BASE-NEXT: Idx Name          Size      Address          Type
+# BASE-NEXT:   0               00000000 0000000000000000
+# BASE-NEXT:   1 .writable     00000004 0000000000000190 DATA
+# BASE-NEXT:   2 .readable     00000004 0000000000000194 DATA
+
+# RUN: echo "SECTIONS { \
+# RUN:  .writable : ONLY_IF_RO { *(.writable) } \
+# RUN:  .readable : ONLY_IF_RW { *(.readable) }}" > %t.script
+# RUN: ld.lld -o %t1 --script %t.script %t
+# RUN: llvm-objdump -section-headers %t1 | \
+# RUN:   FileCheck -check-prefix=NOSECTIONS %s
+# NOSECTIONS: Sections:
+# NOSECTIONS-NOT: .writable
+# NOSECTIONS-NOT: .readable
+
+.global _start
+_start:
+  nop
+
+.section .writable, "aw"
+writable:
+ .long 1
+
+.section .readable, "a"
+readable:
+ .long 2