Update new global detection and processing in JIT

codon/compiler/jit.cpp

@@ -58,23 +58,14 @@ llvm::Error JIT::init() {
   return llvm::Error::success();
 }
 
-llvm::Expected<std::string> JIT::run(const ir::Func *input,
-                                     const std::vector<ir::Var *> &newGlobals) {
+llvm::Expected<std::string> 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<ir::Func>(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<std::string> JIT::exec(const std::string &code) {
 
     auto *cache = compiler->getCache();
     auto typechecked = ast::TypecheckVisitor::apply(cache, simplified);
-    std::vector<std::string> globalNames;
-    for (auto &g : cache->globals) {
-      if (!g.second)
-        globalNames.push_back(g.first);
-    }
+
     // add newly realized functions
     std::vector<ast::StmtPtr> v;
     std::vector<ir::Func **> frs;
@@ -155,16 +142,7 @@ llvm::Expected<std::string> JIT::exec(const std::string &code) {
         ast::TranslateVisitor::apply(cache, std::make_shared<ast::SuiteStmt>(v, false));
     cache->jitCell++;
 
-    std::vector<ir::Var *> 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;

codon/compiler/jit.h

@@ -28,8 +28,7 @@ public:
   Engine *getEngine() const { return engine.get(); }
 
   llvm::Error init();
-  llvm::Expected<std::string> run(const ir::Func *input,
-                                  const std::vector<ir::Var *> &newGlobals = {});
+  llvm::Expected<std::string> run(const ir::Func *input);
   llvm::Expected<std::string> exec(const std::string &code);
 };
 

codon/sir/llvm/llvisitor.cpp

@@ -113,6 +113,25 @@ void LLVMVisitor::registerGlobal(const Var *var) {
   }
 }
 
+void LLVMVisitor::processNewGlobals(Module *module) {
+  std::vector<Func *> newFuncs;
+  for (auto *var : *module) {
+    if (!var->isGlobal())
+      continue;
+    auto id = var->getId();
+    auto *func = cast<Func>(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()) {

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();