JIT imports

codon/compiler/jit.cpp

@@ -35,6 +35,7 @@ llvm::Error JIT::init() {
 
   auto transformed = ast::SimplifyVisitor::apply(
       cache, std::make_shared<ast::SuiteStmt>(), JIT_FILENAME, {});
+
   auto typechecked = ast::TypecheckVisitor::apply(cache, std::move(transformed));
   ast::TranslateVisitor::apply(cache, std::move(typechecked));
   cache->isJit = true; // we still need main(), so set isJit after it has been set
@@ -90,6 +91,9 @@ llvm::Error JIT::run(const ir::Func *input, const std::vector<ir::Var *> &newGlo
 std::pair<ir::Func *, std::vector<ir::Var *>>
 JIT::transformSimplified(const ast::StmtPtr &simplified) {
   auto *cache = compiler->getCache();
+
+  // LOG("-- {}", simplified->toString(1));
+
   auto typechecked = ast::TypecheckVisitor::apply(cache, simplified);
   std::vector<std::string> globalNames;
   for (auto &g : cache->globals) {

codon/parser/common.cpp

@@ -250,10 +250,12 @@ std::shared_ptr<ImportFile> getImportFile(const std::string &argv0,
                                           bool forceStdlib, const std::string &module0,
                                           const std::vector<std::string> &plugins) {
   std::vector<fs::path> paths;
-  auto parentRelativeTo = fs::path(relativeTo).parent_path();
-  if (!forceStdlib) {
-    addPath(paths, (parentRelativeTo / what).replace_extension("codon"));
-    addPath(paths, parentRelativeTo / what / "__init__.codon");
+  if (what != "<jit>") {
+    auto parentRelativeTo = fs::path(relativeTo).parent_path();
+    if (!forceStdlib) {
+      addPath(paths, (parentRelativeTo / what).replace_extension("codon"));
+      addPath(paths, parentRelativeTo / what / "__init__.codon");
+    }
   }
   for (auto &p : getStdLibPaths(argv0, plugins)) {
     addPath(paths, (p / what).replace_extension("codon"));

codon/parser/visitors/simplify/simplify.cpp

@@ -88,7 +88,7 @@ SimplifyVisitor::apply(Cache *cache, const StmtPtr &node, const std::string &fil
     }
     // Reserve the following static identifiers.
     for (auto name : {"staticlen", "compile_error", "isinstance", "hasattr", "type",
-                      "TypeVar", "Callable"})
+                      "TypeVar", "Callable", "argv"})
       stdlib->generateCanonicalName(name);
 
     // This code must be placed in a preamble (these are not POD types but are

codon/parser/visitors/simplify/simplify_stmt.cpp

@@ -1043,7 +1043,8 @@ StmtPtr SimplifyVisitor::transformAssignment(const ExprPtr &lhs, const ExprPtr &
     // ctx->moduleName != MODULE_MAIN;
     // ⚠️ TODO: should we make __main__ top-level variables NOT global by default?
     // Problem: a = [1]; def foo(): a.append(2) won't work anymore as in Python.
-    if (global && !isStatic && !(r && r->isType()))
+    if (global && !isStatic && !(r && r->isType()) &&
+        !in(ctx->cache->globals, canonical))
       ctx->cache->globals[canonical] = nullptr;
     // Handle type aliases as well!
     ctx->add(r && r->isType() ? SimplifyItem::Type : SimplifyItem::Var, e->value,
@@ -1337,7 +1338,8 @@ void SimplifyVisitor::transformNewImport(const ImportFile &file) {
     // loaded)
     preamble->globals.push_back(N<AssignStmt>(
         N<IdExpr>(importDoneVar = importVar + "_done"), N<BoolExpr>(false)));
-    ctx->cache->globals[importDoneVar] = nullptr;
+    if (!in(ctx->cache->globals, importDoneVar))
+      ctx->cache->globals[importDoneVar] = nullptr;
     std::vector<StmtPtr> stmts;
     stmts.push_back(nullptr); // placeholder to be filled later!
     // We need to wrap all imported top-level statements (not signatures! they have