Merge pull request #3621 from z0w0/jit-crates

Add support for crate loading to JIT

Author: brson

src/rustc/back/link.rs

@@ -74,17 +74,47 @@ mod jit {
             m: ModuleRef,
             opt: c_int,
             stacks: bool) unsafe {
-        let ptr = llvm::LLVMRustJIT(rusti::morestack_addr(),
-                                    pm, m, opt, stacks);
+        let manager = llvm::LLVMRustPrepareJIT(rusti::morestack_addr());
+
+        // We need to tell JIT where to resolve all linked
+        // symbols from. The equivalent of -lstd, -lcore, etc.
+        // By default the JIT will resolve symbols from the std and
+        // core linked into rustc. We don't want that,
+        // incase the user wants to use an older std library.
+
+        let cstore = sess.cstore;
+        for cstore::get_used_crate_files(cstore).each |cratepath| {
+            let path = cratepath.to_str();
+
+            debug!("linking: %s", path);
+
+            let _: () = str::as_c_str(
+                path,
+                |buf_t| {
+                    if !llvm::LLVMRustLoadCrate(manager, buf_t) {
+                        llvm_err(sess, ~"Could not link");
+                    }
+                    debug!("linked: %s", path);
+                });
+        }
+
+        // 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.
 
-        if ptr::is_null(ptr) {
+        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: ptr,
+                code: entry,
                 env: ptr::null()
             };
-            let func: fn(~[~str]) = cast::transmute(move closure);
+            let func: fn(++argv: ~[~str]) = cast::transmute(move closure);
 
             func(~[sess.opts.binary]);
         }
@@ -193,30 +223,9 @@ mod write {
                 // JIT execution takes ownership of the module,
                 // so don't dispose and return.
 
-                // We need to tell LLVM where to resolve all linked
-                // symbols from. The equivalent of -lstd, -lcore, etc.
-                // By default the JIT will resolve symbols from the std and
-                // core linked into rustc. We don't want that,
-                // incase the user wants to use an older std library.
-                /*let cstore = sess.cstore;
-                for cstore::get_used_crate_files(cstore).each |cratepath| {
-                    debug!{"linking: %s", cratepath};
-
-                    let _: () = str::as_c_str(
-                        cratepath,
-                        |buf_t| {
-                            if !llvm::LLVMRustLoadLibrary(buf_t) {
-                                llvm_err(sess, ~"Could not link");
-                            }
-                            debug!{"linked: %s", cratepath};
-                        });
-                }*/
-
                 jit::exec(sess, pm.llpm, llmod, CodeGenOptLevel, true);
 
-                if sess.time_llvm_passes() {
-                    llvm::LLVMRustPrintPassTimings();
-                }
+                if sess.time_llvm_passes() { llvm::LLVMRustPrintPassTimings(); }
                 return;
             }
 

src/rustc/lib/llvm.rs

@@ -990,15 +990,19 @@ extern mod llvm {
         call. */
     fn LLVMRustGetLastError() -> *c_char;
 
-    /** Load a shared library to resolve symbols against. */
-    fn LLVMRustLoadLibrary(Filename: *c_char) -> bool;
-
-    /** Create and execute the JIT engine. */
-    fn LLVMRustJIT(__morestack: *(),
-                   PM: PassManagerRef,
-                   M: ModuleRef,
-                   OptLevel: c_int,
-                   EnableSegmentedStacks: bool) -> *();
+    /** Prepare the JIT. Returns a memory manager that can load crates. */
+    fn LLVMRustPrepareJIT(__morestack: *()) -> *();
+
+    /** Load a crate into the memory manager. */
+    fn LLVMRustLoadCrate(MM: *(),
+                         Filename: *c_char) -> bool;
+
+    /** Execute the JIT engine. */
+    fn LLVMRustExecuteJIT(MM: *(),
+                          PM: PassManagerRef,
+                          M: ModuleRef,
+                          OptLevel: c_int,
+                          EnableSegmentedStacks: bool) -> *();
 
     /** Parses the bitcode in the given memory buffer. */
     fn LLVMRustParseBitcode(MemBuf: MemoryBufferRef) -> ModuleRef;

src/rustllvm/RustWrapper.cpp

@@ -20,6 +20,7 @@
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Assembly/Parser.h"
 #include "llvm/Assembly/PrintModulePass.h"
 #include "llvm/Support/FormattedStream.h"
@@ -42,7 +43,6 @@
 #include "llvm-c/Core.h"
 #include "llvm-c/BitReader.h"
 #include "llvm-c/Object.h"
-#include <cstdlib>
 
 // Used by RustMCJITMemoryManager::getPointerToNamedFunction()
 // to get around glibc issues. See the function for more information.
@@ -53,6 +53,7 @@
 #endif
 
 using namespace llvm;
+using namespace llvm::sys;
 
 static const char *LLVMRustError;
 
@@ -100,18 +101,6 @@ void LLVMRustInitializeTargets() {
   LLVMInitializeX86AsmParser();
 }
 
-extern "C" bool
-LLVMRustLoadLibrary(const char* file) {
-  std::string err;
-
-  if(llvm::sys::DynamicLibrary::LoadLibraryPermanently(file, &err)) {
-    LLVMRustError = err.c_str();
-    return false;
-  }
-
-  return true;
-}
-
 // Custom memory manager for MCJITting. It needs special features
 // that the generic JIT memory manager doesn't entail. Based on
 // code from LLI, change where needed for Rust.
@@ -121,10 +110,13 @@ class RustMCJITMemoryManager : public JITMemoryManager {
   SmallVector<sys::MemoryBlock, 16> AllocatedCodeMem;
   SmallVector<sys::MemoryBlock, 16> FreeCodeMem;
   void* __morestack;
+  DenseSet<DynamicLibrary*> crates;
 
   RustMCJITMemoryManager(void* sym) : __morestack(sym) { }
   ~RustMCJITMemoryManager();
 
+  bool loadCrate(const char*, std::string*);
+
   virtual uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment,
                                        unsigned SectionID);
 
@@ -197,6 +189,19 @@ class RustMCJITMemoryManager : public JITMemoryManager {
   }
 };
 
+bool RustMCJITMemoryManager::loadCrate(const char* file, std::string* err) {
+  DynamicLibrary crate = DynamicLibrary::getPermanentLibrary(file,
+                                                             err);
+
+  if(crate.isValid()) {
+    crates.insert(&crate);
+
+    return true;
+  }
+
+  return false;
+}
+
 uint8_t *RustMCJITMemoryManager::allocateDataSection(uintptr_t Size,
                                                     unsigned Alignment,
                                                     unsigned SectionID) {
@@ -276,6 +281,9 @@ void *RustMCJITMemoryManager::getPointerToNamedFunction(const std::string &Name,
   if (Name == "__morestack") return &__morestack;
 
   const char *NameStr = Name.c_str();
+
+  // Look through loaded crates and main for symbols.
+
   void *Ptr = sys::DynamicLibrary::SearchForAddressOfSymbol(NameStr);
   if (Ptr) return Ptr;
 
@@ -293,21 +301,49 @@ RustMCJITMemoryManager::~RustMCJITMemoryManager() {
 }
 
 extern "C" void*
-LLVMRustJIT(void* __morestack,
-            LLVMPassManagerRef PMR,
-            LLVMModuleRef M,
-            CodeGenOpt::Level OptLevel,
-            bool EnableSegmentedStacks) {
+LLVMRustPrepareJIT(void* __morestack) {
+  // An execution engine will take ownership of this later
+  // and clean it up for us.
+
+  return (void*) new RustMCJITMemoryManager(__morestack);
+}
+
+extern "C" bool
+LLVMRustLoadCrate(void* mem, const char* crate) {
+  RustMCJITMemoryManager* manager = (RustMCJITMemoryManager*) mem;
+  std::string Err;
+
+  assert(manager);
+
+  if(!manager->loadCrate(crate, &Err)) {
+    LLVMRustError = Err.c_str();
+    return false;
+  }
+
+  return true;
+}
+
+extern "C" void*
+LLVMRustExecuteJIT(void* mem,
+                   LLVMPassManagerRef PMR,
+                   LLVMModuleRef M,
+                   CodeGenOpt::Level OptLevel,
+                   bool EnableSegmentedStacks) {
 
   InitializeNativeTarget();
   InitializeNativeTargetAsmPrinter();
+  InitializeNativeTargetAsmParser();
 
   std::string Err;
   TargetOptions Options;
+  Options.JITExceptionHandling = true;
   Options.JITEmitDebugInfo = true;
   Options.NoFramePointerElim = true;
   Options.EnableSegmentedStacks = EnableSegmentedStacks;
   PassManager *PM = unwrap<PassManager>(PMR);
+  RustMCJITMemoryManager* MM = (RustMCJITMemoryManager*) mem;
+
+  assert(MM);
 
   PM->add(createBasicAliasAnalysisPass());
   PM->add(createInstructionCombiningPass());
@@ -318,8 +354,8 @@ LLVMRustJIT(void* __morestack,
   PM->add(createPromoteMemoryToRegisterPass());
   PM->run(*unwrap(M));
 
-  RustMCJITMemoryManager* MM = new RustMCJITMemoryManager(__morestack);
   ExecutionEngine* EE = EngineBuilder(unwrap(M))
+    .setErrorStr(&Err)
     .setTargetOptions(Options)
     .setJITMemoryManager(MM)
     .setOptLevel(OptLevel)

src/rustllvm/rustllvm.def.in

@@ -4,8 +4,9 @@ LLVMRustWriteOutputFile
 LLVMRustGetLastError
 LLVMRustConstSmallInt
 LLVMRustConstInt
-LLVMRustLoadLibrary
-LLVMRustJIT
+LLVMRustLoadCrate
+LLVMRustPrepareJIT
+LLVMRustExecuteJIT
 LLVMRustParseBitcode
 LLVMRustParseAssemblyFile
 LLVMRustPrintPassTimings