Have JIT execution take ownership of the LLVMContextRef

Also stop leaking the ExecutionEngine created for jit code by forcibly disposing
of it after the JIT code has finished executing

Author: alexcrichton

src/librustc/back/link.rs

@@ -93,7 +93,7 @@ pub mod jit {
     use back::link::llvm_err;
     use driver::session::Session;
     use lib::llvm::llvm;
-    use lib::llvm::{ModuleRef, PassManagerRef};
+    use lib::llvm::{ModuleRef, PassManagerRef, ContextRef};
     use metadata::cstore;
 
     use core::libc::c_int;
@@ -113,6 +113,7 @@ pub mod jit {
 
     pub fn exec(sess: Session,
                 pm: PassManagerRef,
+                c: ContextRef,
                 m: ModuleRef,
                 opt: c_int,
                 stacks: bool) {
@@ -141,26 +142,43 @@ pub mod jit {
                     });
             }
 
-            // The execute function will return a void pointer
-            // to the _rust_main function. We can do closure
-            // magic here to turn it straight into a callable rust
-            // closure. It will also cleanup the memory manager
-            // for us.
-
-            let entry = llvm::LLVMRustExecuteJIT(manager,
-                                                 pm, m, opt, stacks);
-
-            if ptr::is_null(entry) {
-                llvm_err(sess, ~"Could not JIT");
-            } else {
-                let closure = Closure {
-                    code: entry,
-                    env: ptr::null()
-                };
-                let func: &fn() = cast::transmute(closure);
+            // We custom-build a JIT execution engine via some rust wrappers
+            // first. This wrappers takes ownership of the module passed in.
+            let ee = llvm::LLVMRustBuildJIT(manager, pm, m, opt, stacks);
+            if ptr::is_null(ee) {
+                llvm::LLVMContextDispose(c);
+                llvm_err(sess, ~"Could not create the JIT");
+            }
 
-                func();
+            // Next, we need to get a handle on the _rust_main function by
+            // looking up it's corresponding ValueRef and then requesting that
+            // the execution engine compiles the function.
+            let fun = do str::as_c_str("_rust_main") |entry| {
+                llvm::LLVMGetNamedFunction(m, entry)
+            };
+            if fun.is_null() {
+                llvm::LLVMDisposeExecutionEngine(ee);
+                llvm::LLVMContextDispose(c);
+                llvm_err(sess, ~"Could not find _rust_main in the JIT");
             }
+
+            // Finally, once we have the pointer to the code, we can do some
+            // closure magic here to turn it straight into a callable rust
+            // closure
+            let code = llvm::LLVMGetPointerToGlobal(ee, fun);
+            assert!(!code.is_null());
+            let closure = Closure {
+                code: code,
+                env: ptr::null()
+            };
+            let func: &fn() = cast::transmute(closure);
+            func();
+
+            // Sadly, there currently is no interface to re-use this execution
+            // engine, so it's disposed of here along with the context to
+            // prevent leaks.
+            llvm::LLVMDisposeExecutionEngine(ee);
+            llvm::LLVMContextDispose(c);
         }
     }
 }
@@ -174,7 +192,7 @@ pub mod write {
     use driver::session::Session;
     use driver::session;
     use lib::llvm::llvm;
-    use lib::llvm::{False, ModuleRef, mk_pass_manager, mk_target_data};
+    use lib::llvm::{False, ContextRef, ModuleRef, mk_pass_manager, mk_target_data};
     use lib;
 
     use core::libc::{c_int, c_uint};
@@ -190,6 +208,7 @@ pub mod write {
     }
 
     pub fn run_passes(sess: Session,
+                      llcx: ContextRef,
                       llmod: ModuleRef,
                       output_type: output_type,
                       output: &Path) {
@@ -287,7 +306,7 @@ pub mod write {
                     // JIT execution takes ownership of the module,
                     // so don't dispose and return.
 
-                    jit::exec(sess, pm.llpm, llmod, CodeGenOptLevel, true);
+                    jit::exec(sess, pm.llpm, llcx, llmod, CodeGenOptLevel, true);
 
                     if sess.time_llvm_passes() {
                         llvm::LLVMRustPrintPassTimings();
@@ -358,6 +377,7 @@ pub mod write {
                 // Clean up and return
 
                 llvm::LLVMDisposeModule(llmod);
+                llvm::LLVMContextDispose(llcx);
                 if sess.time_llvm_passes() {
                     llvm::LLVMRustPrintPassTimings();
                 }
@@ -377,6 +397,7 @@ pub mod write {
             }
 
             llvm::LLVMDisposeModule(llmod);
+            llvm::LLVMContextDispose(llcx);
             if sess.time_llvm_passes() { llvm::LLVMRustPrintPassTimings(); }
         }
     }

src/librustc/driver/driver.rs

@@ -236,7 +236,7 @@ pub fn compile_rest(sess: Session,
 
     let outputs = outputs.get();
 
-    let (llmod, link_meta) = {
+    let (llcx, llmod, link_meta) = {
 
         let ty_cx = ty::mk_ctxt(sess, def_map, ast_map, freevars,
                                 region_map, rp_set, lang_items);
@@ -318,14 +318,14 @@ pub fn compile_rest(sess: Session,
         let obj_filename = outputs.obj_filename.with_filetype("s");
 
         time(time_passes, ~"LLVM passes", ||
-            link::write::run_passes(sess, llmod, output_type,
-                            &obj_filename));
+            link::write::run_passes(sess, llcx, llmod, output_type,
+                                    &obj_filename));
 
         link::write::run_ndk(sess, &obj_filename, &outputs.obj_filename);
     } else {
         time(time_passes, ~"LLVM passes", ||
-            link::write::run_passes(sess, llmod, sess.opts.output_type,
-                                &outputs.obj_filename));
+            link::write::run_passes(sess, llcx, llmod, sess.opts.output_type,
+                                    &outputs.obj_filename));
     }
 
     let stop_after_codegen =

src/librustc/lib/llvm.rs

@@ -199,6 +199,8 @@ pub enum BasicBlock_opaque {}
 pub type BasicBlockRef = *BasicBlock_opaque;
 pub enum Builder_opaque {}
 pub type BuilderRef = *Builder_opaque;
+pub enum ExecutionEngine_opaque {}
+pub type ExecutionEngineRef = *ExecutionEngine_opaque;
 pub enum MemoryBuffer_opaque {}
 pub type MemoryBufferRef = *MemoryBuffer_opaque;
 pub enum PassManager_opaque {}
@@ -215,7 +217,7 @@ pub enum SectionIterator_opaque {}
 pub type SectionIteratorRef = *SectionIterator_opaque;
 
 pub mod llvm {
-    use super::{AtomicBinOp, AtomicOrdering, BasicBlockRef};
+    use super::{AtomicBinOp, AtomicOrdering, BasicBlockRef, ExecutionEngineRef};
     use super::{Bool, BuilderRef, ContextRef, MemoryBufferRef, ModuleRef};
     use super::{ObjectFileRef, Opcode, PassManagerRef, PassManagerBuilderRef};
     use super::{SectionIteratorRef, TargetDataRef, TypeKind, TypeRef, UseRef};
@@ -355,6 +357,10 @@ pub mod llvm {
         pub unsafe fn LLVMGetPointerAddressSpace(PointerTy: TypeRef)
                                               -> c_uint;
         #[fast_ffi]
+        pub unsafe fn LLVMGetPointerToGlobal(EE: ExecutionEngineRef,
+                                             V: ValueRef)
+                                              -> *();
+        #[fast_ffi]
         pub unsafe fn LLVMGetVectorSize(VectorTy: TypeRef) -> c_uint;
 
         /* Operations on other types */
@@ -995,6 +1001,8 @@ pub mod llvm {
                                                 Name: *c_char);
         #[fast_ffi]
         pub unsafe fn LLVMDisposeBuilder(Builder: BuilderRef);
+        #[fast_ffi]
+        pub unsafe fn LLVMDisposeExecutionEngine(EE: ExecutionEngineRef);
 
         /* Metadata */
         #[fast_ffi]
@@ -1785,11 +1793,11 @@ pub mod llvm {
 
         /** Execute the JIT engine. */
         #[fast_ffi]
-        pub unsafe fn LLVMRustExecuteJIT(MM: *(),
+        pub unsafe fn LLVMRustBuildJIT(MM: *(),
                               PM: PassManagerRef,
                               M: ModuleRef,
                               OptLevel: c_int,
-                              EnableSegmentedStacks: bool) -> *();
+                              EnableSegmentedStacks: bool) -> ExecutionEngineRef;
 
         /** Parses the bitcode in the given memory buffer. */
         #[fast_ffi]

src/librustc/middle/trans/base.rs

@@ -2984,7 +2984,7 @@ pub fn trans_crate(sess: session::Session,
                    tcx: ty::ctxt,
                    output: &Path,
                    emap2: resolve::ExportMap2,
-                   maps: astencode::Maps) -> (ModuleRef, LinkMeta) {
+                   maps: astencode::Maps) -> (ContextRef, ModuleRef, LinkMeta) {
 
     let symbol_hasher = @mut hash::default_state();
     let link_meta = link::build_link_meta(sess, crate, output, symbol_hasher);
@@ -3148,7 +3148,8 @@ pub fn trans_crate(sess: session::Session,
                 io::println(fmt!("%-7u %s", v, k));
             }
         }
-        return (llmod, link_meta);
+        unset_task_llcx();
+        return (llcx, llmod, link_meta);
     }
 }
 
@@ -3159,8 +3160,10 @@ pub fn task_llcx() -> ContextRef {
     *opt.expect("task-local LLVMContextRef wasn't ever set!")
 }
 
-fn set_task_llcx(c: ContextRef) {
-    unsafe {
-        local_data::local_data_set(task_local_llcx_key, @c);
-    }
+unsafe fn set_task_llcx(c: ContextRef) {
+    local_data::local_data_set(task_local_llcx_key, @c);
+}
+
+unsafe fn unset_task_llcx() {
+    local_data::local_data_pop(task_local_llcx_key);
 }

src/rustllvm/RustWrapper.cpp

@@ -50,6 +50,7 @@
 #include "llvm/Transforms/IPO.h"
 #include "llvm-c/Core.h"
 #include "llvm-c/BitReader.h"
+#include "llvm-c/ExecutionEngine.h"
 #include "llvm-c/Object.h"
 
 // Used by RustMCJITMemoryManager::getPointerToNamedFunction()
@@ -372,12 +373,12 @@ LLVMRustLoadCrate(void* mem, const char* crate) {
   return true;
 }
 
-extern "C" void*
-LLVMRustExecuteJIT(void* mem,
-                   LLVMPassManagerRef PMR,
-                   LLVMModuleRef M,
-                   CodeGenOpt::Level OptLevel,
-                   bool EnableSegmentedStacks) {
+extern "C" LLVMExecutionEngineRef
+LLVMRustBuildJIT(void* mem,
+                 LLVMPassManagerRef PMR,
+                 LLVMModuleRef M,
+                 CodeGenOpt::Level OptLevel,
+                 bool EnableSegmentedStacks) {
 
   InitializeNativeTarget();
   InitializeNativeTargetAsmPrinter();
@@ -414,21 +415,15 @@ LLVMRustExecuteJIT(void* mem,
 
   if(!EE || Err != "") {
     LLVMRustError = Err.c_str();
-    return 0;
+    // The EngineBuilder only takes ownership of these two structures if the
+    // create() call is successful, but here it wasn't successful.
+    LLVMDisposeModule(M);
+    delete MM;
+    return NULL;
   }
 
   MM->invalidateInstructionCache();
-  Function* func = EE->FindFunctionNamed("_rust_main");
-
-  if(!func || Err != "") {
-    LLVMRustError = Err.c_str();
-    return 0;
-  }
-
-  void* entry = EE->getPointerToFunction(func);
-  assert(entry);
-
-  return entry;
+  return wrap(EE);
 }
 
 extern "C" bool

src/rustllvm/rustllvm.def.in

@@ -6,13 +6,14 @@ LLVMRustConstSmallInt
 LLVMRustConstInt
 LLVMRustLoadCrate
 LLVMRustPrepareJIT
-LLVMRustExecuteJIT
+LLVMRustBuildJIT
 LLVMRustParseBitcode
 LLVMRustParseAssemblyFile
 LLVMRustPrintPassTimings
 LLVMRustStartMultithreading
 LLVMCreateObjectFile
 LLVMDisposeObjectFile
+LLVMDisposeExecutionEngine
 LLVMGetSections
 LLVMDisposeSectionIterator
 LLVMIsSectionIteratorAtEnd
@@ -356,6 +357,7 @@ LLVMGetParamParent
 LLVMGetParamTypes
 LLVMGetParams
 LLVMGetPointerAddressSpace
+LLVMGetPointerToGlobal
 LLVMGetPreviousBasicBlock
 LLVMGetPreviousFunction
 LLVMGetPreviousGlobal