Library autoloader

alexander-penev · alexander-penev · commit 7bda1c338cbb · 2024-07-10T17:27:57.000Z
diff --git a/include/clang/Interpreter/CppInterOp.h b/include/clang/Interpreter/CppInterOp.h
@@ -663,6 +663,14 @@ namespace Cpp {
                                    unsigned complete_line = 1U,
                                    unsigned complete_column = 1U);
 
+  /// Set libraries autoload.
+  ///\param[in] autoload - true = libraries autoload is on.
+  CPPINTEROP_API void SetLibrariesAutoload(bool autoload = true);
+
+  /// Get libraries autoload.
+  ///\returns LibraryAutoLoad state (true = libraries autoload is on).
+  CPPINTEROP_API bool GetLibrariesAutoload();
+
 } // end namespace Cpp
 
 #endif // CPPINTEROP_CPPINTEROP_H
diff --git a/lib/Interpreter/CppInterOp.cpp b/lib/Interpreter/CppInterOp.cpp
@@ -27,6 +27,7 @@
 #include "clang/Sema/TemplateDeduction.h"
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_os_ostream.h"
@@ -3267,4 +3268,175 @@ namespace Cpp {
                          complete_column);
   }
 
+#define DEBUG_TYPE "autoload"
+
+  static inline std::string DemangleNameForDlsym(const std::string& name) {
+    std::string nameForDlsym = name;
+
+#if defined(R__MACOSX) || defined(R__WIN32)
+    // The JIT gives us a mangled name which has an additional leading underscore
+    // on macOS and Windows, for instance __ZN8TRandom34RndmEv. However, dlsym
+    // requires us to remove it.
+    // FIXME: get this information from the DataLayout via getGlobalPrefix()!
+    if (nameForDlsym[0] == '_')
+      nameForDlsym.erase(0, 1);
+#endif //R__MACOSX
+
+    return nameForDlsym;
+  }
+
+  class AutoLoadLibrarySearchGenerator;
+  static AutoLoadLibrarySearchGenerator *ALLSG = nullptr;
+
+  class AutoLoadLibrarySearchGenerator : public llvm::orc::DefinitionGenerator {
+  public:
+    bool Enabled = false;
+
+    // Lazy materialization unit class helper
+    class AutoloadLibraryMU : public llvm::orc::MaterializationUnit {
+      std::string lib;
+      llvm::orc::SymbolNameVector syms;
+    public:
+      AutoloadLibraryMU(const std::string &Library, const llvm::orc::SymbolNameVector &Symbols)
+        : MaterializationUnit({getSymbolFlagsMap(Symbols), nullptr}), lib(Library), syms(Symbols) {}
+
+      StringRef getName() const override {
+        return "<Symbols from Autoloaded Library>";
+      }
+
+      void materialize(std::unique_ptr<llvm::orc::MaterializationResponsibility> R) override {
+        if (!ALLSG || !ALLSG->Enabled) {
+          R->failMaterialization();
+          return;
+        }
+
+        LLVM_DEBUG(dbgs() << "Materialize " << lib << " syms=" << syms);
+
+        auto& I = getInterp();
+        auto DLM = I.getDynamicLibraryManager();
+
+        llvm::orc::SymbolMap loadedSymbols;
+        llvm::orc::SymbolNameSet failedSymbols;
+        bool loadedLibrary = false;
+
+        for (auto symbol : syms) {
+          std::string symbolStr = (*symbol).str();
+          std::string nameForDlsym = DemangleNameForDlsym(symbolStr);
+
+          // Check if the symbol is available without loading the library.
+          void *addr = llvm::sys::DynamicLibrary::SearchForAddressOfSymbol(nameForDlsym);
+
+          if (!addr && !loadedLibrary) {
+            // Try to load the library which should provide the symbol definition.
+            if (DLM->loadLibrary(lib, false) != DynamicLibraryManager::LoadLibResult::kLoadLibSuccess) {
+              LLVM_DEBUG(dbgs() << "MU: Failed to load library " << lib);
+              string err = "MU: Failed to load library! " + lib;
+              perror(err.c_str());
+            } else {
+              LLVM_DEBUG(dbgs() << "MU: Autoload library " << lib);
+            }
+
+            // Only try loading the library once.
+            loadedLibrary = true;
+
+            addr = llvm::sys::DynamicLibrary::SearchForAddressOfSymbol(nameForDlsym);
+          }
+
+          if (addr) {
+            loadedSymbols[symbol] =
+              llvm::orc::ExecutorSymbolDef(llvm::orc::ExecutorAddr::fromPtr(addr), JITSymbolFlags::Exported);
+          } else {
+            // Collect all failing symbols, delegate their responsibility and then
+            // fail their materialization. R->defineNonExistent() sounds like it
+            // should do that, but it's not implemented?!
+            failedSymbols.insert(symbol);
+          }
+        }
+
+        if (!failedSymbols.empty()) {
+          auto failingMR = R->delegate(failedSymbols);
+          if (failingMR) {
+            (*failingMR)->failMaterialization();
+          }
+        }
+
+        if (!loadedSymbols.empty()) {
+          llvm::cantFail(R->notifyResolved(loadedSymbols));
+          llvm::cantFail(R->notifyEmitted());
+        }
+      }
+
+      void discard(const llvm::orc::JITDylib &JD, const llvm::orc::SymbolStringPtr &Name) override {}
+
+      private:
+        static llvm::orc::SymbolFlagsMap getSymbolFlagsMap(const llvm::orc::SymbolNameVector &Symbols) {
+          llvm::orc::SymbolFlagsMap map;
+          for (auto symbolName : Symbols)
+            map[symbolName] = llvm::JITSymbolFlags::Exported;
+          return map;
+        }
+    };
+
+    llvm::Error tryToGenerate(llvm::orc::LookupState &LS, llvm::orc::LookupKind K, llvm::orc::JITDylib &JD,
+                              llvm::orc::JITDylibLookupFlags JDLookupFlags, const llvm::orc::SymbolLookupSet &Symbols) override {
+      if (Enabled) {
+        LLVM_DEBUG(dbgs() << "tryToGenerate");
+
+        auto& I = getInterp();
+        auto DLM = I.getDynamicLibraryManager();
+
+        std::unordered_map<std::string, llvm::orc::SymbolNameVector> found;
+        llvm::orc::SymbolMap NewSymbols;
+        for (auto &KV : Symbols) {
+          auto &Name = KV.first;
+          if ((*Name).empty())
+            continue;
+
+          auto lib = DLM->searchLibrariesForSymbol(*Name, /*searchSystem=*/true); // false?
+          if (lib.empty())
+            continue;
+
+          found[lib].push_back(Name);
+
+          // Workaround: This getAddressOfGlobal call make first symbol search
+          // to work, immediatelly after library auto load. This approach do not
+          // use MU
+          //DLM->loadLibrary(lib, true);
+          //I.getAddressOfGlobal(*Name);
+        }
+
+        for (auto &&KV : found) {
+          auto MU = std::make_unique<AutoloadLibraryMU>(KV.first, std::move(KV.second));
+          if (auto Err = JD.define(MU))
+            return Err;
+        }
+      }
+
+      return llvm::Error::success();
+    }
+  };
+
+  void SetLibrariesAutoload(bool autoload /* = true */) {
+    auto& I = getInterp();
+    llvm::orc::LLJIT& EE = *compat::getExecutionEngine(I);
+#if CLANG_VERSION_MAJOR < 17
+    llvm::orc::JITDylib& DyLib = EE.getMainJITDylib();
+#else
+    llvm::orc::JITDylib& DyLib = *EE.getProcessSymbolsJITDylib().get();
+#endif // CLANG_VERSION_MAJOR
+
+    if (!ALLSG)
+      ALLSG = &DyLib.addGenerator(std::make_unique<AutoLoadLibrarySearchGenerator>());
+    ALLSG->Enabled = autoload;
+
+    LLVM_DEBUG(dbgs() << "Autoload=" << (ALLSG && ALLSG->Enabled ? "ON" : "OFF"));
+  }
+
+  bool GetLibrariesAutoload() {
+    LLVM_DEBUG(dbgs() << "Autoload is " << (ALLSG && ALLSG->Enabled ? "ON" : "OFF"));
+    return ALLSG && ALLSG->Enabled;
+  }
+
+#undef DEBUG_TYPE
+
   } // end namespace Cpp
diff --git a/lib/Interpreter/DynamicLibraryManager.cpp b/lib/Interpreter/DynamicLibraryManager.cpp
@@ -242,7 +242,7 @@ namespace Cpp {
     // get canonical path name and check if already loaded
     const std::string Path = platform::NormalizePath(foundDyLib);
     if (Path.empty()) {
-      LLVM_DEBUG(dbgs() << "cling::DynamicLibraryManager::lookupLibMaybeAddExt(): "
+      LLVM_DEBUG(dbgs() << "DynamicLibraryManager::lookupLibMaybeAddExt(): "
                         << "error getting real (canonical) path of library " << foundDyLib << '\n');
       return foundDyLib;
     }
@@ -391,8 +391,7 @@ namespace Cpp {
       return;
 
     DyLibHandle dyLibHandle = nullptr;
-    for (DyLibs::const_iterator I = m_DyLibs.begin(), E = m_DyLibs.end();
-         I != E; ++I) {
+    for (DyLibs::const_iterator I = m_DyLibs.begin(), E = m_DyLibs.end(); I != E; ++I) {
       if (I->second == canonicalLoadedLib) {
         dyLibHandle = I->first;
         break;
@@ -404,7 +403,7 @@ namespace Cpp {
     std::string errMsg;
     platform::DLClose(dyLibHandle, &errMsg);
     if (!errMsg.empty()) {
-      LLVM_DEBUG(dbgs() << "cling::DynamicLibraryManager::unloadLibrary(): "
+      LLVM_DEBUG(dbgs() << "DynamicLibraryManager::unloadLibrary(): "
                         << errMsg << '\n');
     }
 
diff --git a/lib/Interpreter/DynamicLibraryManagerSymbol.cpp b/lib/Interpreter/DynamicLibraryManagerSymbol.cpp
@@ -1153,8 +1153,8 @@ namespace Cpp {
   std::string Dyld::searchLibrariesForSymbol(StringRef mangledName,
                                              bool searchSystem/* = true*/) {
 #define DEBUG_TYPE "Dyld:searchLibrariesForSymbol:"
-    assert(!llvm::sys::DynamicLibrary::SearchForAddressOfSymbol(mangledName.str()) &&
-           "Library already loaded, please use dlsym!");
+//    assert(!llvm::sys::DynamicLibrary::SearchForAddressOfSymbol(mangledName.str()) &&
+//           "Library already loaded, please use dlsym!");
     assert(!mangledName.empty());
 
     using namespace llvm::sys::path;
diff --git a/unittests/CppInterOp/CMakeLists.txt b/unittests/CppInterOp/CMakeLists.txt
@@ -36,3 +36,7 @@ set_output_directory(DynamicLibraryManagerTests BINARY_DIR ${CMAKE_BINARY_DIR}/u
 add_dependencies(DynamicLibraryManagerTests TestSharedLib)
 #export_executable_symbols_for_plugins(TestSharedLib)
 add_subdirectory(TestSharedLib)
+# TODO: Remove when libraryunload work correctly
+add_dependencies(DynamicLibraryManagerTests TestSharedLib1)
+#export_executable_symbols_for_plugins(TestSharedLib1)
+add_subdirectory(TestSharedLib1)
diff --git a/unittests/CppInterOp/DynamicLibraryManagerTest.cpp b/unittests/CppInterOp/DynamicLibraryManagerTest.cpp
@@ -5,6 +5,9 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 
+
+// Helper functions
+
 // This function isn't referenced outside its translation unit, but it
 // can't use the "static" keyword because its address is used for
 // GetMainExecutable (since some platforms don't support taking the
@@ -17,9 +20,20 @@ std::string GetExecutablePath(const char* Argv0) {
   return llvm::sys::fs::getMainExecutable(Argv0, MainAddr);
 }
 
+static inline std::string MangleNameForDlsym(const std::string& name) {
+  std::string nameForDlsym = name;
+#if defined(R__MACOSX) || defined(R__WIN32)
+  if (nameForDlsym[0] != '_')
+    nameForDlsym.insert (0, 1, '_');
+#endif //R__MACOSX
+  return nameForDlsym;
+}
+
+// Tests
+
 TEST(DynamicLibraryManagerTest, Sanity) {
   EXPECT_TRUE(Cpp::CreateInterpreter());
-  EXPECT_FALSE(Cpp::GetFunctionAddress("ret_zero"));
+  EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_zero").c_str()));
 
   std::string BinaryPath = GetExecutablePath(/*Argv0=*/nullptr);
   llvm::StringRef Dir = llvm::sys::path::parent_path(BinaryPath);
@@ -28,13 +42,9 @@ TEST(DynamicLibraryManagerTest, Sanity) {
   // FIXME: dlsym on mach-o takes the C-level name, however, the macho-o format
   // adds an additional underscore (_) prefix to the lowered names. Figure out
   // how to harmonize that API.
-#ifdef __APPLE__
-  std::string PathToTestSharedLib =
-      Cpp::SearchLibrariesForSymbol("_ret_zero", /*system_search=*/false);
-#else
+
   std::string PathToTestSharedLib =
-      Cpp::SearchLibrariesForSymbol("ret_zero", /*system_search=*/false);
-#endif // __APPLE__
+      Cpp::SearchLibrariesForSymbol(MangleNameForDlsym("ret_zero").c_str(), /*system_search=*/false);
 
   EXPECT_STRNE("", PathToTestSharedLib.c_str())
       << "Cannot find: '" << PathToTestSharedLib << "' in '" << Dir.str()
@@ -44,13 +54,76 @@ TEST(DynamicLibraryManagerTest, Sanity) {
   // Force ExecutionEngine to be created.
   Cpp::Process("");
   // FIXME: Conda returns false to run this code on osx.
-#ifndef __APPLE__
-  EXPECT_TRUE(Cpp::GetFunctionAddress("ret_zero"));
-#endif //__APPLE__
+  EXPECT_TRUE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_zero").c_str()));
 
   Cpp::UnloadLibrary("TestSharedLib");
   // We have no reliable way to check if it was unloaded because posix does not
   // require the library to be actually unloaded but just the handle to be
   // invalidated...
-  // EXPECT_FALSE(Cpp::GetFunctionAddress("ret_zero"));
+  // EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_zero").c_str()));
+}
+
+TEST(DynamicLibraryManagerTest, LibrariesAutoload) {
+  EXPECT_TRUE(Cpp::CreateInterpreter());
+
+  // Autoload by default is OFF. Symbol search must fail.
+  EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str()));
+
+  // Libraries Search path by default is set to main executable directory.
+  std::string BinaryPath = GetExecutablePath(/*Argv0=*/nullptr);
+  llvm::StringRef Dir = llvm::sys::path::parent_path(BinaryPath);
+  Cpp::AddSearchPath(Dir.str().c_str());
+
+  // Find library with "rec_one" symbol defined and exported
+  //
+  // FIXME: dlsym on mach-o takes the C-level name, however, the macho-o format
+  // adds an additional underscore (_) prefix to the lowered names. Figure out
+  // how to harmonize that API. For now we use out minimal implementation of
+  // helper function.
+  std::string PathToTestSharedLib1 =
+      Cpp::SearchLibrariesForSymbol(MangleNameForDlsym("ret_one").c_str(), /*system_search=*/false);
+
+  // If result is "" then we cannot find this library.
+  EXPECT_STRNE("", PathToTestSharedLib1.c_str())
+      << "Cannot find: '" << PathToTestSharedLib1 << "' in '" << Dir.str() << "'";
+
+  // Force ExecutionEngine to be created.
+  Cpp::Process("");
+
+  // Check default Autoload is OFF
+  EXPECT_FALSE(Cpp::GetLibrariesAutoload());
+  // Find symbol must fail (when auotload=off)
+  EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str()));
+
+  // Autoload turn ON
+  Cpp::SetLibrariesAutoload(true);
+  // Check autorum status (must be turned ON)
+  EXPECT_TRUE(Cpp::GetLibrariesAutoload());
+
+  // FIXME: Conda returns false to run this code on osx.
+  EXPECT_TRUE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str()));
+  // Check for some symbols (exists and not exists)
+  EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_not_exist").c_str()));
+  EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_not_exist1").c_str()));
+  EXPECT_TRUE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_two").c_str()));
+
+  //Cpp::UnloadLibrary("TestSharedLib1");
+  //// We have no reliable way to check if it was unloaded because posix does not
+  //// require the library to be actually unloaded but just the handle to be
+  //// invalidated...
+  //EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str()));
+
+  // Autoload turn OFF
+  Cpp::SetLibrariesAutoload(false);
+  // Check autorum status (must be turned OFF)
+  EXPECT_FALSE(Cpp::GetLibrariesAutoload());
+  //EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str())); // if unload works
+  //EXPECT_TRUE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str()));
+
+  // Autoload turn ON again
+  Cpp::SetLibrariesAutoload(true);
+  // Check autorum status (must be turned ON)
+  EXPECT_TRUE(Cpp::GetLibrariesAutoload());
+  //EXPECT_FALSE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str())); // if unload works
+  //EXPECT_TRUE(Cpp::GetFunctionAddress(MangleNameForDlsym("ret_one").c_str()));
 }
diff --git a/unittests/CppInterOp/TestSharedLib1/CMakeLists.txt b/unittests/CppInterOp/TestSharedLib1/CMakeLists.txt
@@ -0,0 +1,12 @@
+# TODO: Remove when library unload work correctly
+add_llvm_library(TestSharedLib1
+  SHARED
+  DISABLE_LLVM_LINK_LLVM_DYLIB
+  BUILDTREE_ONLY
+  TestSharedLib1.cpp)
+# Put TestSharedLib1 next to the unit test executable.
+set_output_directory(TestSharedLib1
+  BINARY_DIR ${CMAKE_BINARY_DIR}/unittests/bin/$<CONFIG>/
+  LIBRARY_DIR ${CMAKE_BINARY_DIR}/unittests/bin/$<CONFIG>/
+  )
+set_target_properties(TestSharedLib1 PROPERTIES FOLDER "Tests")
diff --git a/unittests/CppInterOp/TestSharedLib1/TestSharedLib1.cpp b/unittests/CppInterOp/TestSharedLib1/TestSharedLib1.cpp
@@ -0,0 +1,5 @@
+#include "TestSharedLib1.h"
+
+int ret_one() { return 1; }
+
+int ret_two() { return 2; }
diff --git a/unittests/CppInterOp/TestSharedLib1/TestSharedLib1.h b/unittests/CppInterOp/TestSharedLib1/TestSharedLib1.h
