diff --git a/codon/compiler/jit.cpp b/codon/compiler/jit.cpp index 38f197b1..4f3b01e7 100644 --- a/codon/compiler/jit.cpp +++ b/codon/compiler/jit.cpp @@ -58,23 +58,14 @@ llvm::Error JIT::init() { return llvm::Error::success(); } -llvm::Expected JIT::run(const ir::Func *input, - const std::vector &newGlobals) { +llvm::Expected JIT::run(const ir::Func *input) { auto *module = compiler->getModule(); auto *pm = compiler->getPassManager(); auto *llvisitor = compiler->getLLVMVisitor(); pm->run(module); const std::string name = ir::LLVMVisitor::getNameForFunction(input); - llvisitor->registerGlobal(input); - for (auto *var : newGlobals) { - llvisitor->registerGlobal(var); - } - for (auto *var : newGlobals) { - if (auto *func = ir::cast(var)) - func->accept(*llvisitor); - } - input->accept(*llvisitor); + llvisitor->processNewGlobals(module); auto pair = llvisitor->takeModule(); if (auto err = engine->addModule({std::move(pair.first), std::move(pair.second)})) @@ -138,11 +129,7 @@ llvm::Expected JIT::exec(const std::string &code) { auto *cache = compiler->getCache(); auto typechecked = ast::TypecheckVisitor::apply(cache, simplified); - std::vector globalNames; - for (auto &g : cache->globals) { - if (!g.second) - globalNames.push_back(g.first); - } + // add newly realized functions std::vector v; std::vector frs; @@ -155,16 +142,7 @@ llvm::Expected JIT::exec(const std::string &code) { ast::TranslateVisitor::apply(cache, std::make_shared(v, false)); cache->jitCell++; - std::vector globalVars; - for (auto &g : globalNames) { - seqassert(cache->globals[g], "JIT global {} not set", g); - globalVars.push_back(cache->globals[g]); - } - for (auto &i : frs) { - seqassert(*i, "JIT fn not set"); - globalVars.push_back(*i); - } - return run(func, globalVars); + return run(func); } catch (const exc::ParserException &e) { *cache = bCache; *(cache->imports[MAIN_IMPORT].ctx) = bSimplify; diff --git a/codon/compiler/jit.h b/codon/compiler/jit.h index b781299d..ab3e42a7 100644 --- a/codon/compiler/jit.h +++ b/codon/compiler/jit.h @@ -28,8 +28,7 @@ public: Engine *getEngine() const { return engine.get(); } llvm::Error init(); - llvm::Expected run(const ir::Func *input, - const std::vector &newGlobals = {}); + llvm::Expected run(const ir::Func *input); llvm::Expected exec(const std::string &code); }; diff --git a/codon/sir/llvm/llvisitor.cpp b/codon/sir/llvm/llvisitor.cpp index 1b877e9b..d108a810 100644 --- a/codon/sir/llvm/llvisitor.cpp +++ b/codon/sir/llvm/llvisitor.cpp @@ -113,6 +113,25 @@ void LLVMVisitor::registerGlobal(const Var *var) { } } +void LLVMVisitor::processNewGlobals(Module *module) { + std::vector newFuncs; + for (auto *var : *module) { + if (!var->isGlobal()) + continue; + auto id = var->getId(); + auto *func = cast(var); + bool isNewFunc = (func && funcs.find(id) == funcs.end()); + if (isNewFunc || (!func && vars.find(id) == vars.end())) + registerGlobal(var); + if (isNewFunc) + newFuncs.push_back(func); + } + + for (auto *func : newFuncs) { + func->accept(*this); + } +} + llvm::Value *LLVMVisitor::getVar(const Var *var) { auto it = vars.find(var->getId()); if (db.jit && var->isGlobal()) { diff --git a/codon/sir/llvm/llvisitor.h b/codon/sir/llvm/llvisitor.h index 2a29e7c4..4f28252b 100644 --- a/codon/sir/llvm/llvisitor.h +++ b/codon/sir/llvm/llvisitor.h @@ -275,6 +275,11 @@ public: /// @param var the global variable (or function) to register void registerGlobal(const Var *var); + /// Processes new globals that were not previously + /// compiled. Used in JIT mode. + /// @param module the IR module + void processNewGlobals(Module *module); + /// Returns the default LLVM linkage type for the module. /// @return LLVM linkage type llvm::GlobalValue::LinkageTypes getDefaultLinkage();