4 files changed, 60 insertions, 3 deletions

diff --git a/lld/ELF/Config.h b/lld/ELF/Config.h
index c860741e9bd..9963e15e698 100644
--- a/lld/ELF/Config.h
+++ b/lld/ELF/Config.h
@@ -171,6 +171,7 @@ struct Configuration {
   bool ZCopyreloc;
   bool ZExecstack;
   bool ZHazardplt;
+  bool ZKeepTextSectionPrefix;
   bool ZNodelete;
   bool ZNodlopen;
   bool ZNow;
diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index c15fcf04cc5..0d6602b767e 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -764,6 +764,8 @@ void LinkerDriver::readConfigs(opt::InputArgList &Args) {
   Config->ZCopyreloc = getZFlag(Args, "copyreloc", "nocopyreloc", true);
   Config->ZExecstack = getZFlag(Args, "execstack", "noexecstack", false);
   Config->ZHazardplt = hasZOption(Args, "hazardplt");
+  Config->ZKeepTextSectionPrefix = getZFlag(
+      Args, "keep-text-section-prefix", "nokeep-text-section-prefix", false);
   Config->ZNodelete = hasZOption(Args, "nodelete");
   Config->ZNodlopen = hasZOption(Args, "nodlopen");
   Config->ZNow = getZFlag(Args, "now", "lazy", false);
diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
index bb5b39b4785..3c76d63911c 100644
--- a/lld/ELF/Writer.cpp
+++ b/lld/ELF/Writer.cpp
@@ -88,6 +88,10 @@ private:
 };
 } // anonymous namespace
 
+static bool isSectionPrefix(StringRef Prefix, StringRef Name) {
+  return Name.startswith(Prefix) || Name == Prefix.drop_back();
+}
+
 StringRef elf::getOutputSectionName(InputSectionBase *S) {
   if (Config->Relocatable)
     return S->Name;
@@ -104,13 +108,25 @@ StringRef elf::getOutputSectionName(InputSectionBase *S) {
     }
   }
 
+  // This check is for -z keep-text-section-prefix.  This option separates text
+  // sections with prefix ".text.hot", ".text.unlikely", ".text.startup" or
+  // ".text.exit".
+  // When enabled, this allows identifying the hot code region (.text.hot) in
+  // the final binary which can be selectively mapped to huge pages or mlocked,
+  // for instance.
+  if (Config->ZKeepTextSectionPrefix)
+    for (StringRef V :
+         {".text.hot.", ".text.unlikely.", ".text.startup.", ".text.exit."}) {
+      if (isSectionPrefix(V, S->Name))
+        return V.drop_back();
+    }
+
   for (StringRef V :
        {".text.", ".rodata.", ".data.rel.ro.", ".data.", ".bss.rel.ro.",
         ".bss.", ".init_array.", ".fini_array.", ".ctors.", ".dtors.", ".tbss.",
         ".gcc_except_table.", ".tdata.", ".ARM.exidx.", ".ARM.extab."}) {
-    StringRef Prefix = V.drop_back();
-    if (S->Name.startswith(V) || S->Name == Prefix)
-      return Prefix;
+    if (isSectionPrefix(V, S->Name))
+      return V.drop_back();
   }
 
   // CommonSection is identified as "COMMON" in linker scripts.
diff --git a/lld/test/ELF/text-section-prefix.s b/lld/test/ELF/text-section-prefix.s
new file mode 100644
index 00000000000..578a504e7e5
--- /dev/null
+++ b/lld/test/ELF/text-section-prefix.s
@@ -0,0 +1,38 @@
+# RUN: llvm-mc -filetype=obj %s -o %t
+# RUN: ld.lld -z keep-text-section-prefix %t -o %t2
+# RUN: llvm-readelf -l %t2 | FileCheck %s
+# RUN: ld.lld %t -o %t3
+# RUN: llvm-readelf -l %t3 | FileCheck --check-prefix=CHECKNO %s
+# RUN: ld.lld -z nokeep-text-section-prefix %t -o %t4
+# RUN: llvm-readelf -l %t4 | FileCheck --check-prefix=CHECKNO %s
+
+# CHECK: .text
+# CHECK: .text.hot
+# CHECK: .text.startup
+# CHECK: .text.exit
+# CHECK: .text.unlikely
+# CHECKNO: .text
+# CHECKNO-NOT: .text.hot
+
+_start:
+  ret
+
+.section .text.f,"ax"
+f:
+  nop
+
+.section .text.hot.f_hot,"ax"
+f_hot:
+  nop
+
+.section .text.startup.f_startup,"ax"
+f_startup:
+  nop
+
+.section .text.exit.f_exit,"ax"
+f_exit:
+  nop
+
+.section .text.unlikely.f_unlikely,"ax"
+f_unlikely:
+  nop