[ORC][ORC_RT] Handle ELF .init_array with non-default priority

ELF-based platforms currently support defining multiple static
initializer table sections with differing priorities, for example
.init_array.0 or .init_array.100; the default .init_array corresponds
to a priority of 65535. When building a shared library or executable,
the system linker normally sorts these sections and combines them into
a single .init_array section. This change adds the capability to
recognize ELF static initializers with priorities other than the
default, and to properly sort them by priority, to Orc and the Orc
runtime.

Reviewed By: lhames

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

6 files changed, 232 insertions, 32 deletions

diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index 7800bcd..6f502b2 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -15,6 +15,7 @@
 #include "error.h"
 #include "wrapper_function_utils.h"
 
+#include <algorithm>
 #include <map>
 #include <mutex>
 #include <sstream>
@@ -133,12 +134,6 @@ private:
 
   static ELFNixPlatformRuntimeState *MOPS;
 
-  using InitSectionHandler =
-      Error (*)(const std::vector<ExecutorAddrRange> &Sections,
-                const ELFNixJITDylibInitializers &MOJDIs);
-  const std::vector<std::pair<const char *, InitSectionHandler>> InitSections =
-      {{".init_array", runInitArray}};
-
   void *PlatformJDDSOHandle;
 
   // FIXME: Move to thread-state.
@@ -378,21 +373,29 @@ Expected<void *> ELFNixPlatformRuntimeState::dlopenInitialize(string_view Path,
   return JDS->Header;
 }
 
+long getPriority(const std::string &name) {
+  auto pos = name.find_last_not_of("0123456789");
+  if (pos == name.size() - 1)
+    return 65535;
+  else
+    return std::strtol(name.c_str() + pos + 1, nullptr, 10);
+}
+
 Error ELFNixPlatformRuntimeState::initializeJITDylib(
     ELFNixJITDylibInitializers &MOJDIs) {
 
   auto &JDS = getOrCreateJITDylibState(MOJDIs);
   ++JDS.RefCount;
 
-  for (auto &KV : InitSections) {
-    const auto &Name = KV.first;
-    const auto &Handler = KV.second;
-    auto I = MOJDIs.InitSections.find(Name);
-    if (I != MOJDIs.InitSections.end()) {
-      if (auto Err = Handler(I->second, MOJDIs))
-        return Err;
-    }
-  }
+  using SectionList = std::vector<ExecutorAddrRange>;
+  std::sort(MOJDIs.InitSections.begin(), MOJDIs.InitSections.end(),
+            [](const std::pair<std::string, SectionList> &LHS,
+               const std::pair<std::string, SectionList> &RHS) -> bool {
+              return getPriority(LHS.first) < getPriority(RHS.first);
+            });
+  for (auto &Entry : MOJDIs.InitSections)
+    if (auto Err = runInitArray(Entry.second, MOJDIs))
+      return Err;
 
   return Error::success();
 }
diff --git a/compiler-rt/lib/orc/elfnix_platform.h b/compiler-rt/lib/orc/elfnix_platform.h
index 12b9591..e0ee959 100644
--- a/compiler-rt/lib/orc/elfnix_platform.h
+++ b/compiler-rt/lib/orc/elfnix_platform.h
@@ -47,7 +47,7 @@ struct ELFNixJITDylibInitializers {
   std::string Name;
   ExecutorAddr DSOHandleAddress;
 
-  std::unordered_map<std::string, SectionList> InitSections;
+  std::vector<std::pair<std::string, SectionList>> InitSections;
 };
 
 class ELFNixJITDylibDeinitializers {};
diff --git a/compiler-rt/test/orc/TestCases/FreeBSD/x86-64/priority-static-initializer.S b/compiler-rt/test/orc/TestCases/FreeBSD/x86-64/priority-static-initializer.S
new file mode 100644
index 0000000..13a52f3
--- /dev/null
+++ b/compiler-rt/test/orc/TestCases/FreeBSD/x86-64/priority-static-initializer.S
@@ -0,0 +1,99 @@
+// Test that ELF static initializers with different constructor priorities work
+// and are executed in the proper order.
+//
+// RUN: %clang -c -o %t %s
+// RUN: %llvm_jitlink %t | FileCheck %s
+
+// CHECK: constructor 100
+// CHECK-NEXT: constructor 200
+// CHECK-NEXT: constructor 65535
+// CHECK-NEXT: main
+// CHECK-NEXT: destructor
+
+	.text
+
+	.globl	destructor
+	.p2align	4, 0x90
+	.type	destructor,@function
+destructor:
+.Ldestructor$local:
+
+	leaq	.L.str.d(%rip), %rdi
+	jmp	puts@PLT
+
+	.globl	main
+	.p2align	4, 0x90
+	.type	main,@function
+main:
+.Lmain$local:
+
+	pushq	%rax
+	leaq	.L.str(%rip), %rdi
+	callq	puts@PLT
+	xorl	%eax, %eax
+	popq	%rcx
+	retq
+
+	.p2align	4, 0x90
+	.type	constructor.65535,@function
+constructor.65535:
+
+	pushq	%rax
+	leaq	.L.str.65535(%rip), %rdi
+	callq	puts@PLT
+	leaq	.Ldestructor$local(%rip), %rdi
+	leaq	__dso_handle(%rip), %rdx
+	xorl	%esi, %esi
+	popq	%rax
+	jmp	__cxa_atexit@PLT
+
+	.p2align	4, 0x90
+	.type	constructor.200,@function
+constructor.200:
+
+	leaq	.L.str.200(%rip), %rdi
+	jmp	puts@PLT
+
+	.p2align	4, 0x90
+	.type	constructor.100,@function
+constructor.100:
+
+	leaq	.L.str.100(%rip), %rdi
+	jmp	puts@PLT
+
+	.hidden	__dso_handle
+	.type	.L.str,@object
+	.section	.rodata.str1.1,"aMS",@progbits,1
+.L.str:
+	.asciz	"main"
+	.size	.L.str, 5
+
+	.type	.L.str.100,@object
+.L.str.100:
+	.asciz	"constructor 100"
+	.size	.L.str.100, 16
+
+	.type	.L.str.200,@object
+.L.str.200:
+	.asciz	"constructor 200"
+	.size	.L.str.200, 16
+
+	.type	.L.str.65535,@object
+.L.str.65535:
+	.asciz	"constructor 65535"
+	.size	.L.str.65535, 18
+
+	.type	.L.str.d,@object
+.L.str.d:
+	.asciz	"destructor"
+	.size	.L.str.d, 11
+
+	.section	.init_array.100,"aw",@init_array
+	.p2align	3
+	.quad	constructor.100
+	.section	.init_array.200,"aw",@init_array
+	.p2align	3
+	.quad	constructor.200
+	.section	.init_array,"aw",@init_array
+	.p2align	3
+	.quad	constructor.65535
diff --git a/compiler-rt/test/orc/TestCases/Linux/x86-64/priority-static-initializer.S b/compiler-rt/test/orc/TestCases/Linux/x86-64/priority-static-initializer.S
new file mode 100644
index 0000000..13a52f3
--- /dev/null
+++ b/compiler-rt/test/orc/TestCases/Linux/x86-64/priority-static-initializer.S
@@ -0,0 +1,99 @@
+// Test that ELF static initializers with different constructor priorities work
+// and are executed in the proper order.
+//
+// RUN: %clang -c -o %t %s
+// RUN: %llvm_jitlink %t | FileCheck %s
+
+// CHECK: constructor 100
+// CHECK-NEXT: constructor 200
+// CHECK-NEXT: constructor 65535
+// CHECK-NEXT: main
+// CHECK-NEXT: destructor
+
+	.text
+
+	.globl	destructor
+	.p2align	4, 0x90
+	.type	destructor,@function
+destructor:
+.Ldestructor$local:
+
+	leaq	.L.str.d(%rip), %rdi
+	jmp	puts@PLT
+
+	.globl	main
+	.p2align	4, 0x90
+	.type	main,@function
+main:
+.Lmain$local:
+
+	pushq	%rax
+	leaq	.L.str(%rip), %rdi
+	callq	puts@PLT
+	xorl	%eax, %eax
+	popq	%rcx
+	retq
+
+	.p2align	4, 0x90
+	.type	constructor.65535,@function
+constructor.65535:
+
+	pushq	%rax
+	leaq	.L.str.65535(%rip), %rdi
+	callq	puts@PLT
+	leaq	.Ldestructor$local(%rip), %rdi
+	leaq	__dso_handle(%rip), %rdx
+	xorl	%esi, %esi
+	popq	%rax
+	jmp	__cxa_atexit@PLT
+
+	.p2align	4, 0x90
+	.type	constructor.200,@function
+constructor.200:
+
+	leaq	.L.str.200(%rip), %rdi
+	jmp	puts@PLT
+
+	.p2align	4, 0x90
+	.type	constructor.100,@function
+constructor.100:
+
+	leaq	.L.str.100(%rip), %rdi
+	jmp	puts@PLT
+
+	.hidden	__dso_handle
+	.type	.L.str,@object
+	.section	.rodata.str1.1,"aMS",@progbits,1
+.L.str:
+	.asciz	"main"
+	.size	.L.str, 5
+
+	.type	.L.str.100,@object
+.L.str.100:
+	.asciz	"constructor 100"
+	.size	.L.str.100, 16
+
+	.type	.L.str.200,@object
+.L.str.200:
+	.asciz	"constructor 200"
+	.size	.L.str.200, 16
+
+	.type	.L.str.65535,@object
+.L.str.65535:
+	.asciz	"constructor 65535"
+	.size	.L.str.65535, 18
+
+	.type	.L.str.d,@object
+.L.str.d:
+	.asciz	"destructor"
+	.size	.L.str.d, 11
+
+	.section	.init_array.100,"aw",@init_array
+	.p2align	3
+	.quad	constructor.100
+	.section	.init_array.200,"aw",@init_array
+	.p2align	3
+	.quad	constructor.200
+	.section	.init_array,"aw",@init_array
+	.p2align	3
+	.quad	constructor.65535
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index 730e193..e476c54 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -101,8 +101,6 @@ StringRef InitArrayFuncSectionName = ".init_array";
 StringRef ThreadBSSSectionName = ".tbss";
 StringRef ThreadDataSectionName = ".tdata";
 
-StringRef InitSectionNames[] = {InitArrayFuncSectionName};
-
 } // end anonymous namespace
 
 namespace llvm {
@@ -274,10 +272,9 @@ ELFNixPlatform::standardRuntimeUtilityAliases() {
 }
 
 bool ELFNixPlatform::isInitializerSection(StringRef SecName) {
-  for (auto &Name : InitSectionNames) {
-    if (Name.equals(SecName))
-      return true;
-  }
+  if (SecName.consume_front(InitArrayFuncSectionName) &&
+      (SecName.empty() || SecName[0] == '.'))
+    return true;
   return false;
 }
 
@@ -781,16 +778,15 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::preserveInitSections(
     jitlink::LinkGraph &G, MaterializationResponsibility &MR) {
 
   JITLinkSymbolSet InitSectionSymbols;
-  for (auto &InitSectionName : InitSectionNames) {
+  for (auto &InitSection : G.sections()) {
     // Skip non-init sections.
-    auto *InitSection = G.findSectionByName(InitSectionName);
-    if (!InitSection)
+    if (!isInitializerSection(InitSection.getName()))
       continue;
 
     // Make a pass over live symbols in the section: those blocks are already
     // preserved.
     DenseSet<jitlink::Block *> AlreadyLiveBlocks;
-    for (auto &Sym : InitSection->symbols()) {
+    for (auto &Sym : InitSection.symbols()) {
       auto &B = Sym->getBlock();
       if (Sym->isLive() && Sym->getOffset() == 0 &&
           Sym->getSize() == B.getSize() && !AlreadyLiveBlocks.count(&B)) {
@@ -800,7 +796,7 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::preserveInitSections(
     }
 
     // Add anonymous symbols to preserve any not-already-preserved blocks.
-    for (auto *B : InitSection->blocks())
+    for (auto *B : InitSection.blocks())
       if (!AlreadyLiveBlocks.count(B))
         InitSectionSymbols.insert(
             &G.addAnonymousSymbol(*B, 0, B->getSize(), false, true));
@@ -821,9 +817,9 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
 
   LLVM_DEBUG({ dbgs() << "ELFNixPlatform::registerInitSections\n"; });
 
-  for (auto InitSectionName : InitSectionNames) {
-    if (auto *Sec = G.findSectionByName(InitSectionName)) {
-      InitSections.push_back(Sec);
+  for (auto &Sec : G.sections()) {
+    if (isInitializerSection(Sec.getName())) {
+      InitSections.push_back(&Sec);
     }
   }
 
diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 4461b6c..5ddb35c 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -78,9 +78,12 @@ private:
   }
 
   static bool hasELFInitSection(LinkGraph &G) {
-    for (auto &Sec : G.sections())
-      if (Sec.getName() == ".init_array")
+    for (auto &Sec : G.sections()) {
+      auto SecName = Sec.getName();
+      if (SecName.consume_front(".init_array") &&
+          (SecName.empty() || SecName[0] == '.'))
         return true;
+    }
     return false;
   }