From 18915f4e492dacd0b7b1c59a024531ee26aa742f Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Tue, 13 May 2025 19:30:18 +0200
Subject: [PATCH] ORC: Add support for the default linking layer (JITLink).

---
 deps/LLVMExtra/include/LLVMExtra.h |  2 ++
 deps/LLVMExtra/lib/Core.cpp        | 13 +++++++++++++
 lib/16/libLLVM_extra.jl            |  4 ++++
 src/orc.jl                         | 16 ++++++++++++++++
 4 files changed, 35 insertions(+)

diff --git a/deps/LLVMExtra/include/LLVMExtra.h b/deps/LLVMExtra/include/LLVMExtra.h
index 42669bf0..52488149 100644
--- a/deps/LLVMExtra/include/LLVMExtra.h
+++ b/deps/LLVMExtra/include/LLVMExtra.h
@@ -4,6 +4,7 @@
 #include "llvm/Config/llvm-config.h"
 #include <llvm-c/Core.h>
 #include <llvm-c/Orc.h>
+#include <llvm-c/LLJIT.h>
 #include <llvm-c/Target.h>
 #include <llvm-c/Transforms/PassBuilder.h>
 #include <llvm-c/Types.h>
@@ -109,6 +110,7 @@ void LLVMOrcIRCompileLayerEmit(LLVMOrcIRCompileLayerRef IRLayer,
                                LLVMOrcMaterializationResponsibilityRef MR,
                                LLVMOrcThreadSafeModuleRef TSM);
 char *LLVMDumpJitDylibToString(LLVMOrcJITDylibRef JD);
+LLVMOrcObjectLayerRef LLVMOrcLLJITGetObjectLinkingLayer(LLVMOrcLLJITRef J);
 
 // Cloning functionality
 typedef enum {
diff --git a/deps/LLVMExtra/lib/Core.cpp b/deps/LLVMExtra/lib/Core.cpp
index 519e2ff4..a3995426 100644
--- a/deps/LLVMExtra/lib/Core.cpp
+++ b/deps/LLVMExtra/lib/Core.cpp
@@ -9,6 +9,7 @@
 #include <llvm/Analysis/TargetLibraryInfo.h>
 #include <llvm/Analysis/TargetTransformInfo.h>
 #include <llvm/CodeGen/Passes.h>
+#include <llvm/ExecutionEngine/Orc/LLJIT.h>
 #include <llvm/ExecutionEngine/Orc/IRCompileLayer.h>
 #include <llvm/IR/Attributes.h>
 #include <llvm/IR/DebugInfo.h>
@@ -384,6 +385,18 @@ char *LLVMDumpJitDylibToString(LLVMOrcJITDylibRef JD) {
   return strdup(str.c_str());
 }
 
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(orc::ObjectLayer, LLVMOrcObjectLayerRef)
+
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(orc::LLJIT, LLVMOrcLLJITRef)
+
+// expose JITLink
+LLVMOrcObjectLayerRef LLVMOrcLLJITGetObjectLinkingLayer(LLVMOrcLLJITRef J) {
+    if (!J)
+      return nullptr;
+    orc::LLJIT *LLJITInstance = unwrap(J);
+    return wrap(&LLJITInstance->getObjLinkingLayer());
+}
+
 
 //
 // Cloning functionality
diff --git a/lib/16/libLLVM_extra.jl b/lib/16/libLLVM_extra.jl
index 2c3565fd..81ba4b11 100644
--- a/lib/16/libLLVM_extra.jl
+++ b/lib/16/libLLVM_extra.jl
@@ -190,6 +190,10 @@ function LLVMDumpJitDylibToString(JD)
     ccall((:LLVMDumpJitDylibToString, libLLVMExtra), Cstring, (LLVMOrcJITDylibRef,), JD)
 end
 
+function LLVMOrcLLJITGetObjectLinkingLayer(J)
+    ccall((:LLVMOrcLLJITGetObjectLinkingLayer, libLLVMExtra), LLVMOrcObjectLayerRef, (LLVMOrcLLJITRef,), J)
+end
+
 @cenum LLVMCloneFunctionChangeType::UInt32 begin
     LLVMCloneFunctionChangeTypeLocalChangesOnly = 0
     LLVMCloneFunctionChangeTypeGlobalChanges = 1
diff --git a/src/orc.jl b/src/orc.jl
index b2fcd185..e1573a93 100644
--- a/src/orc.jl
+++ b/src/orc.jl
@@ -45,11 +45,27 @@ end
 end
 Base.unsafe_convert(::Type{API.LLVMOrcObjectLayerRef}, oll::ObjectLinkingLayer) = oll.ref
 
+# RTDyld
 function ObjectLinkingLayer(es::ExecutionSession)
     ref = API.LLVMOrcCreateRTDyldObjectLinkingLayerWithSectionMemoryManager(es)
     ObjectLinkingLayer(ref)
 end
 
+"""
+    JITLinkObjectLayer(lljit::LLJIT)
+
+Gets the default JITLink-based ObjectLinkingLayer from an LLJIT instance.
+The lifetime of this layer is tied to the LLJIT instance.
+"""
+function JITLinkObjectLayer(lljit::LLJIT) # Takes LLJIT instance
+    ref = API.LLVMOrcLLJITGetObjectLinkingLayer(lljit)
+    if ref == C_NULL
+        error("Failed to get ObjectLinkingLayer from LLJIT instance.")
+    end
+    # XXX: owned by the JIT, so can't dispose
+    return ObjectLinkingLayer(ref)
+end
+
 function dispose(oll::ObjectLinkingLayer)
     API.LLVMOrcDisposeObjectLayer(oll)
 end