Allow clients to specify search order of DynamicLibraries.

Summary: Different JITs and other clients of LLVM may have different needs in how symbol resolution should occur.

Reviewers: v.g.vassilev, lhames, karies

Reviewed By: v.g.vassilev

Subscribers: pcanal, llvm-commits

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

llvm-svn: 307849

5 files changed, 63 insertions, 11 deletions

diff --git a/llvm/include/llvm/Support/DynamicLibrary.h b/llvm/include/llvm/Support/DynamicLibrary.h
index a8874a10d46..469d5dfad06 100644
--- a/llvm/include/llvm/Support/DynamicLibrary.h
+++ b/llvm/include/llvm/Support/DynamicLibrary.h
@@ -88,6 +88,22 @@ namespace sys {
       return !getPermanentLibrary(Filename, ErrMsg).isValid();
     }
 
+    enum SearchOrdering {
+      /// SO_Linker - Search as a call to dlsym(dlopen(NULL)) would when
+      /// DynamicLibrary::getPermanentLibrary(NULL) has been called or
+      /// search the list of explcitly loaded symbols if not.
+      SO_Linker,
+      /// SO_LoadedFirst - Search all loaded libraries, then as SO_Linker would.
+      SO_LoadedFirst,
+      /// SO_LoadedLast - Search as SO_Linker would, then loaded libraries.
+      /// Only useful to search if libraries with RTLD_LOCAL have been added.
+      SO_LoadedLast,
+      /// SO_LoadOrder - Or this in to search libraries in the ordered loaded.
+      /// The default bahaviour is to search loaded libraries in reverse.
+      SO_LoadOrder = 4
+    };
+    static SearchOrdering SearchOrder; // = SO_Linker
+
     /// This function will search through all previously loaded dynamic
     /// libraries for the symbol \p symbolName. If it is found, the address of
     /// that symbol is returned. If not, null is returned. Note that this will
diff --git a/llvm/lib/Support/DynamicLibrary.cpp b/llvm/lib/Support/DynamicLibrary.cpp
index 9398789cea8..d8422115eae 100644
--- a/llvm/lib/Support/DynamicLibrary.cpp
+++ b/llvm/lib/Support/DynamicLibrary.cpp
@@ -14,6 +14,7 @@
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm-c/Support.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Config/config.h"
 #include "llvm/Support/ManagedStatic.h"
@@ -73,19 +74,37 @@ public:
     return true;
   }
 
-  void *Lookup(const char *Symbol) {
-    // Process handle gets first try.
+  void *LibLookup(const char *Symbol, DynamicLibrary::SearchOrdering Order) {
+    if (Order & SO_LoadOrder) {
+      for (void *Handle : Handles) {
+        if (void *Ptr = DLSym(Handle, Symbol))
+          return Ptr;
+      }
+    } else {
+      for (void *Handle : llvm::reverse(Handles)) {
+        if (void *Ptr = DLSym(Handle, Symbol))
+          return Ptr;
+      }
+    }
+    return nullptr;
+  }
+
+  void *Lookup(const char *Symbol, DynamicLibrary::SearchOrdering Order) {
+    assert(!((Order & SO_LoadedFirst) && (Order & SO_LoadedLast)) &&
+           "Invalid Ordering");
+
+    if (!Process || (Order & SO_LoadedFirst)) {
+      if (void *Ptr = LibLookup(Symbol, Order))
+        return Ptr;
+    }
     if (Process) {
+      // Use OS facilities to search the current binary and all loaded libs.
       if (void *Ptr = DLSym(Process, Symbol))
         return Ptr;
-#ifndef NDEBUG
-      for (void *Handle : Handles)
-        assert(!DLSym(Handle, Symbol) && "Symbol exists in non process handle");
-#endif
-    } else {
-      // Iterate in reverse, so newer libraries/symbols override older.
-      for (auto &&I = Handles.rbegin(), E = Handles.rend(); I != E; ++I) {
-        if (void *Ptr = DLSym(*I, Symbol))
+
+      // Search any libs that might have been skipped because of RTLD_LOCAL.
+      if (Order & SO_LoadedLast) {
+        if (void *Ptr = LibLookup(Symbol, Order))
           return Ptr;
       }
     }
@@ -113,6 +132,8 @@ static llvm::ManagedStatic<llvm::sys::SmartMutex<true>> SymbolsMutex;
 #endif
 
 char DynamicLibrary::Invalid;
+DynamicLibrary::SearchOrdering DynamicLibrary::SearchOrder =
+    DynamicLibrary::SO_Linker;
 
 namespace llvm {
 void *SearchForAddressOfSpecialSymbol(const char *SymbolName) {
@@ -170,7 +191,7 @@ void *DynamicLibrary::SearchForAddressOfSymbol(const char *SymbolName) {
 
     // Now search the libraries.
     if (OpenedHandles.isConstructed()) {
-      if (void *Ptr = OpenedHandles->Lookup(SymbolName))
+      if (void *Ptr = OpenedHandles->Lookup(SymbolName, SearchOrder))
         return Ptr;
     }
   }
diff --git a/llvm/lib/Support/Unix/DynamicLibrary.inc b/llvm/lib/Support/Unix/DynamicLibrary.inc
index aad77f19c35..f05103ccd1e 100644
--- a/llvm/lib/Support/Unix/DynamicLibrary.inc
+++ b/llvm/lib/Support/Unix/DynamicLibrary.inc
@@ -20,6 +20,9 @@ DynamicLibrary::HandleSet::~HandleSet() {
     ::dlclose(Handle);
   if (Process)
     ::dlclose(Process);
+
+  // llvm_shutdown called, Return to default
+  DynamicLibrary::SearchOrder = DynamicLibrary::SO_Linker;
 }
 
 void *DynamicLibrary::HandleSet::DLOpen(const char *File, std::string *Err) {
diff --git a/llvm/lib/Support/Windows/DynamicLibrary.inc b/llvm/lib/Support/Windows/DynamicLibrary.inc
index caf1a0a658d..083ea902eeb 100644
--- a/llvm/lib/Support/Windows/DynamicLibrary.inc
+++ b/llvm/lib/Support/Windows/DynamicLibrary.inc
@@ -28,6 +28,8 @@ DynamicLibrary::HandleSet::~HandleSet() {
 
   // 'Process' should not be released on Windows.
   assert((!Process || Process==this) && "Bad Handle");
+  // llvm_shutdown called, Return to default
+  DynamicLibrary::SearchOrder = DynamicLibrary::SO_Linker;
 }
 
 void *DynamicLibrary::HandleSet::DLOpen(const char *File, std::string *Err) {
diff --git a/llvm/unittests/Support/DynamicLibrary/DynamicLibraryTest.cpp b/llvm/unittests/Support/DynamicLibrary/DynamicLibraryTest.cpp
index c54e1b7eed2..370e1c5ed5e 100644
--- a/llvm/unittests/Support/DynamicLibrary/DynamicLibraryTest.cpp
+++ b/llvm/unittests/Support/DynamicLibrary/DynamicLibraryTest.cpp
@@ -77,6 +77,7 @@ TEST(DynamicLibrary, Overload) {
     EXPECT_TRUE(DL.isValid());
     EXPECT_TRUE(Err.empty());
 
+    // Test overloading local symbols does not occur by default
     GS = FuncPtr<GetString>(DynamicLibrary::SearchForAddressOfSymbol("TestA"));
     EXPECT_TRUE(GS != nullptr && GS == &TestA);
     EXPECT_EQ(StdString(GS()), "ProcessCall");
@@ -85,6 +86,12 @@ TEST(DynamicLibrary, Overload) {
     EXPECT_TRUE(GS != nullptr && GS == &TestA);
     EXPECT_EQ(StdString(GS()), "ProcessCall");
 
+    // Test overloading by forcing library priority when searching for a symbol
+    DynamicLibrary::SearchOrder = DynamicLibrary::SO_LoadedFirst;
+    GS = FuncPtr<GetString>(DynamicLibrary::SearchForAddressOfSymbol("TestA"));
+    EXPECT_TRUE(GS != nullptr && GS != &TestA);
+    EXPECT_EQ(StdString(GS()), "LibCall");
+
     DynamicLibrary::AddSymbol("TestA", PtrFunc(&OverloadTestA));
     GS = FuncPtr<GetString>(DL.getAddressOfSymbol("TestA"));
     EXPECT_TRUE(GS != nullptr && GS != &OverloadTestA);
@@ -95,6 +102,9 @@ TEST(DynamicLibrary, Overload) {
   }
   EXPECT_TRUE(FuncPtr<GetString>(DynamicLibrary::SearchForAddressOfSymbol(
                   "TestA")) == nullptr);
+
+  // Check serach ordering is reset to default after call to llvm_shutdown
+  EXPECT_TRUE(DynamicLibrary::SearchOrder == DynamicLibrary::SO_Linker);
 }
 
 TEST(DynamicLibrary, Shutdown) {