JIT: isolate JIT runtime symbols via explicit symbol mapping

codon/cir/llvm/llvm.h

@@ -31,6 +31,7 @@
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/MCJIT.h"
+#include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h"
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
@@ -49,6 +50,7 @@
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/ExecutionEngine/Orc/MachOPlatform.h"
+#include "llvm/ExecutionEngine/Orc/Mangling.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/Shared/AllocationActions.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"

codon/compiler/engine.cpp

@@ -2,13 +2,23 @@
 
 #include "engine.h"
 
+#include <dlfcn.h>
+
 #include "codon/cir/llvm/optimize.h"
 #include "codon/compiler/memory_manager.h"
 
 namespace codon {
 namespace jit {
 
-Engine::Engine() : jit(), debug(nullptr) {
+namespace {
+
+llvm::Error makeRuntimeLoadError(const std::string &message) {
+  return llvm::make_error<llvm::StringError>(message, llvm::inconvertibleErrorCode());
+}
+
+} // namespace
+
+Engine::Engine() : jit(), debug(nullptr), globalPrefix('\0') {
   auto eb = llvm::EngineBuilder();
   eb.setMArch(llvm::codegen::getMArch());
   eb.setMCPU(llvm::codegen::getCPUStr());
@@ -39,10 +49,13 @@ Engine::Engine() : jit(), debug(nullptr) {
   builder.setJITTargetMachineBuilder(
       llvm::orc::JITTargetMachineBuilder(target->getTargetTriple()));
   jit = llvm::cantFail(builder.create());
+  globalPrefix = layout.getGlobalPrefix();
 
-  jit->getMainJITDylib().addGenerator(
-      llvm::cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
-          layout.getGlobalPrefix())));
+  jit->getMainJITDylib().addGenerator(llvm::cantFail(
+      llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(globalPrefix)));
+  jit->getMainJITDylib().addGenerator(llvm::cantFail(
+      llvm::orc::DynamicLibrarySearchGenerator::Load(kLibgcc_sName, globalPrefix)));
+  tryAddDynamicLibrarySearchGenerator(kLibstdcxxName);
 
   jit->getIRTransformLayer().setTransform(
       [&](llvm::orc::ThreadSafeModule module,
@@ -54,6 +67,12 @@ Engine::Engine() : jit(), debug(nullptr) {
       });
 }
 
+Engine::~Engine() {
+  jit.reset();
+  for (auto *handle : runtimeHandles)
+    dlclose(handle);
+}
+
 llvm::Error Engine::addModule(llvm::orc::ThreadSafeModule module,
                               llvm::orc::ResourceTrackerSP rt) {
   if (!rt)
@@ -66,5 +85,68 @@ llvm::Expected<llvm::orc::ExecutorAddr> Engine::lookup(llvm::StringRef name) {
   return jit->lookup(name);
 }
 
+llvm::Error Engine::registerSymbols(
+    llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)> symbolMap) {
+  auto &mainJitDylib = jit->getMainJITDylib();
+  return mainJitDylib.define(
+      llvm::orc::absoluteSymbols(symbolMap(llvm::orc::MangleAndInterner(
+          mainJitDylib.getExecutionSession(), jit->getDataLayout()))));
+}
+
+void Engine::tryAddDynamicLibrarySearchGenerator(const char *path) {
+  auto gen = llvm::orc::DynamicLibrarySearchGenerator::Load(path, globalPrefix);
+  if (!gen) {
+    llvm::consumeError(gen.takeError());
+    return;
+  }
+  jit->getMainJITDylib().addGenerator(std::move(*gen));
+}
+
+llvm::Error Engine::addRuntimeSymbolMap(const std::string &path) {
+  void *handle = dlopen(path.c_str(), RTLD_NOW | RTLD_LOCAL);
+  if (!handle) {
+    const char *err = dlerror();
+    return makeRuntimeLoadError("cannot load codon runtime '" + path +
+                                "': " + (err ? err : "unknown error"));
+  }
+
+  llvm::sys::DynamicLibrary lib(handle);
+  void *initSym = lib.getAddressOfSymbol(kRuntimeInitFnName);
+
+  if (!initSym) {
+    const char *err = dlerror();
+    dlclose(handle);
+    return makeRuntimeLoadError("cannot find " + std::string(kRuntimeInitFnName) +
+                                " in '" + path + "': " + (err ? err : "unknown error"));
+  }
+
+  RuntimeSymbolMap runtimeSymbols;
+  auto initFn = reinterpret_cast<RuntimeInitFunc>(initSym);
+  initFn(
+      [](void *ctx, const char *name, void *address) {
+        if (!ctx || !name || !address)
+          return;
+        auto *symbols = static_cast<RuntimeSymbolMap *>(ctx);
+        (*symbols)[name] = address;
+      },
+      &runtimeSymbols);
+
+  // Build a runtime symbol map from the exported symbols and register them.
+  auto runtimeSymbolMap = [&](llvm::orc::MangleAndInterner interner) {
+    auto symbolMap = llvm::orc::SymbolMap();
+    for (auto &[name, address] : runtimeSymbols)
+      symbolMap[interner(name)] = {llvm::orc::ExecutorAddr::fromPtr(address),
+                                   llvm::JITSymbolFlags::Exported};
+    return symbolMap;
+  };
+  if (auto err = registerSymbols(runtimeSymbolMap)) {
+    dlclose(handle);
+    return err;
+  }
+
+  runtimeHandles.push_back(handle);
+  return llvm::Error::success();
+}
+
 } // namespace jit
 } // namespace codon

codon/compiler/engine.h

@@ -2,7 +2,9 @@
 
 #pragma once
 
+#include <map>
 #include <memory>
+#include <string>
 #include <vector>
 
 #include "codon/cir/llvm/llvm.h"
@@ -13,11 +15,29 @@ namespace jit {
 
 class Engine {
 private:
+  constexpr static const char *kLibgcc_sName = "libgcc_s.so.1";
+  constexpr static const char *kLibstdcxxName = "libstdc++.so.6";
+  constexpr static const char *kRuntimeInitFnName = "__codon_jit_runtime_init";
+
+  using RuntimeSymbolMap = std::map<std::string, void *>;
+  using RuntimeAddSymbolFunc = void (*)(void *, const char *, void *);
+  using RuntimeInitFunc = void (*)(RuntimeAddSymbolFunc, void *);
+
   std::unique_ptr<llvm::orc::LLJIT> jit;
   DebugPlugin *debug;
+  char globalPrefix;
+  std::vector<void *> runtimeHandles;
+
+  /// Register symbols with this Engine.
+  llvm::Error registerSymbols(
+      llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)> symbolMap);
+
+  /// Best-effort dynamic library search generator registration.
+  void tryAddDynamicLibrarySearchGenerator(const char *path);
 
 public:
   Engine();
+  ~Engine();
 
   const llvm::DataLayout &getDataLayout() const { return jit->getDataLayout(); }
 
@@ -29,6 +49,11 @@ class Engine {
                         llvm::orc::ResourceTrackerSP rt = nullptr);
 
   llvm::Expected<llvm::orc::ExecutorAddr> lookup(llvm::StringRef name);
+
+  /// Load the Codon runtime locally and register its JIT symbol map with ORC.
+  /// @param path Path to the Codon runtime library file (.so/.dll/.dylib)
+  /// @return llvm::Error::success() on success, error code on failure
+  llvm::Error addRuntimeSymbolMap(const std::string &path);
 };
 
 } // namespace jit

codon/compiler/jit.cpp

@@ -4,6 +4,9 @@
 
 #include <sstream>
 
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+
 #include "codon/parser/common.h"
 #include "codon/parser/peg/peg.h"
 #include "codon/parser/visitors/doc/doc.h"
@@ -374,11 +377,83 @@ JIT::JITResult JIT::executePython(const std::string &name,
   }
 }
 
+namespace {
+
+#ifdef __APPLE__
+constexpr const char *kCodonRTBaseName = "libcodonrt.dylib";
+#else
+constexpr const char *kCodonRTBaseName = "libcodonrt.so";
+#endif
+
+} // namespace
+
+// Locate the absolute path of the Codon runtime shared library
+// (libcodonrt.so / .dylib). Tries, in order:
+//   1. <library_path()>/../<RT>     (libcodonrt sits next to libcodonc)
+//   2. <library_path()>/<RT>
+//   3. $CODON_DIR/lib/codon/<RT>
+//   4. bare soname (let the dynamic linker resolve it)
+// Returns canonicalized path on success; bare soname on failure.
+std::string findCodonRuntime() {
+  std::vector<std::string> candidates;
+
+  const std::string libcodonc = codon::ast::library_path();
+  if (!libcodonc.empty()) {
+    auto dir = llvm::sys::path::parent_path(libcodonc);
+    if (!dir.empty()) {
+      llvm::SmallString<256> p1(dir);
+      llvm::sys::path::append(p1, kCodonRTBaseName);
+      candidates.emplace_back(p1.str());
+
+      llvm::SmallString<256> p2(dir);
+      llvm::sys::path::append(p2, "..", "lib", "codon", kCodonRTBaseName);
+      candidates.emplace_back(p2.str());
+    }
+  }
+
+  if (const char *codonDir = std::getenv("CODON_DIR")) {
+    llvm::SmallString<256> p(codonDir);
+    llvm::sys::path::append(p, "lib", "codon", kCodonRTBaseName);
+    candidates.emplace_back(p.str());
+  }
+
+  for (const auto &c : candidates) {
+    if (!llvm::sys::fs::exists(c))
+      continue;
+    // Canonicalize to prevent the dynamic linker from treating two different
+    // paths to the same physical file as distinct shared objects (which would
+    // otherwise cause the runtime / Boehm GC to be loaded twice).
+    llvm::SmallString<256> real;
+    if (!llvm::sys::fs::real_path(c, real))
+      return std::string(real.str());
+    return c;
+  }
+
+  // Last resort: bare soname; rely on RUNPATH / LD_LIBRARY_PATH / ldconfig.
+  return std::string(kCodonRTBaseName);
+}
+
 } // namespace jit
 } // namespace codon
 
 void *jit_init(char *name) {
   auto jit = new codon::jit::JIT(std::string(name));
+
+  // Register Codon runtime symbols with ORC using the runtime-provided symbol map.
+  // The runtime is loaded locally so its symbols do not leak into the
+  // process-global symbol table. Failure here is fatal: without the runtime the JIT
+  // cannot link any compiled code.
+  const std::string rt = codon::jit::findCodonRuntime();
+  if (auto err = jit->getEngine()->addRuntimeSymbolMap(rt)) {
+    auto info = llvm::toString(std::move(err));
+    llvm::report_fatal_error(
+        llvm::StringRef("cannot register codon runtime symbols from '" + rt +
+                        "': " + info),
+        /*gen_crash_diag=*/false);
+  }
+  if (std::getenv("CODON_JIT_DEBUG"))
+    fmt::print(stderr, "[codon::jit] loaded runtime: {}\n", rt);
+
   llvm::cantFail(jit->init());
   return jit;
 }