jit: Separate JIT execution into two functions and load crates before main lookup

4 files changed, 112 insertions, 34 deletions

diff --git a/src/rustc/back/link.rs b/src/rustc/back/link.rs
index 4bbd51524c4..751f0369141 100644
--- a/src/rustc/back/link.rs
+++ b/src/rustc/back/link.rs
@@ -74,14 +74,44 @@ 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.
+
+        let entry = llvm::LLVMRustExecuteJIT(manager,
+                                      pm, m, opt, stacks);
 
-        if ptr::is_null(ptr) {
+        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);
diff --git a/src/rustc/lib/llvm.rs b/src/rustc/lib/llvm.rs
index 031ce219308..c16fe31f933 100644
--- a/src/rustc/lib/llvm.rs
+++ b/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;
diff --git a/src/rustllvm/RustWrapper.cpp b/src/rustllvm/RustWrapper.cpp
index 6ea433e6f1f..2def9c1c030 100644
--- a/src/rustllvm/RustWrapper.cpp
+++ b/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 @@ public:
   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 @@ public:
   }
 };
 
+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,19 @@ void *RustMCJITMemoryManager::getPointerToNamedFunction(const std::string &Name,
   if (Name == "__morestack") return &__morestack;
 
   const char *NameStr = Name.c_str();
+
+  // Look through loaded crates for symbols.
+
+  for (DenseSet<DynamicLibrary*>::iterator I = crates.begin(),
+       E = crates.end(); I != E; ++I) {
+    void *Ptr = (*I)->getAddressOfSymbol(NameStr);
+
+    if (Ptr) return Ptr;
+  }
+
+  // Fallback to using any symbols LLVM has loaded (generally
+  // from the main program).
+
   void *Ptr = sys::DynamicLibrary::SearchForAddressOfSymbol(NameStr);
   if (Ptr) return Ptr;
 
@@ -293,11 +311,34 @@ 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();
@@ -308,6 +349,9 @@ LLVMRustJIT(void* __morestack,
   Options.NoFramePointerElim = true;
   Options.EnableSegmentedStacks = EnableSegmentedStacks;
   PassManager *PM = unwrap<PassManager>(PMR);
+  RustMCJITMemoryManager* MM = (RustMCJITMemoryManager*) mem;
+
+  assert(MM);
 
   PM->add(createBasicAliasAnalysisPass());
   PM->add(createInstructionCombiningPass());
@@ -318,7 +362,6 @@ LLVMRustJIT(void* __morestack,
   PM->add(createPromoteMemoryToRegisterPass());
   PM->run(*unwrap(M));
 
-  RustMCJITMemoryManager* MM = new RustMCJITMemoryManager(__morestack);
   ExecutionEngine* EE = EngineBuilder(unwrap(M))
     .setTargetOptions(Options)
     .setJITMemoryManager(MM)
diff --git a/src/rustllvm/rustllvm.def.in b/src/rustllvm/rustllvm.def.in
index 1de1e3ba58f..36833e5175e 100644
--- a/src/rustllvm/rustllvm.def.in
+++ b/src/rustllvm/rustllvm.def.in
@@ -4,8 +4,9 @@ LLVMRustWriteOutputFile
 LLVMRustGetLastError
 LLVMRustConstSmallInt
 LLVMRustConstInt
-LLVMRustLoadLibrary
-LLVMRustJIT
+LLVMRustLoadCrate
+LLVMRustPrepareJIT
+LLVMRustExecuteJIT
 LLVMRustParseBitcode
 LLVMRustParseAssemblyFile
 LLVMRustPrintPassTimings