From c0bf885c477823290b712fe2e0a043a58f624338 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ibrahim=20Numanagic=CC=81?= <ibrahimpasa@gmail.com>
Date: Sat, 23 Oct 2021 16:45:48 -0700
Subject: [PATCH] Parser JIT support [wip]

---
 codon/parser/parser.cpp                           |  6 +-----
 codon/parser/visitors/simplify/simplify_stmt.cpp  |  2 +-
 codon/parser/visitors/translate/translate.cpp     | 14 ++++++++++++--
 codon/parser/visitors/translate/translate_ctx.cpp |  6 +-----
 codon/parser/visitors/translate/translate_ctx.h   |  3 +--
 5 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/codon/parser/parser.cpp b/codon/parser/parser.cpp
index b101b6fc..b3a76d0b 100644
--- a/codon/parser/parser.cpp
+++ b/codon/parser/parser.cpp
@@ -172,18 +172,14 @@ void generateDocstr(const std::string &argv0) {
 int jitLoop(const std::string &argv0) {
   fmt::print("Loading Codon JIT...");
   auto cache = std::make_shared<ast::Cache>(argv0);
-  cache->isJit = true;
-
   string fileName = "<jit>";
   // Initialize JIT (load stdlib by parsing an empty AST node)
-
-  // ast::SimplifyVisitor::apply(cache, move(codeStmt), abs, defines, (isTest > 1));
-
   auto transformed =
       ast::SimplifyVisitor::apply(cache, make_shared<ast::SuiteStmt>(), fileName, {});
   auto typechecked = ast::TypecheckVisitor::apply(cache, move(transformed));
   ast::TranslateVisitor::apply(cache, move(typechecked));
 
+  cache->isJit = true; // we still need main(), so set isJit after it has been set
   auto jit = jit::JIT(cache->module);
   cache->module->setSrcInfo({fileName, 0, 0, 0});
   jit.init();
diff --git a/codon/parser/visitors/simplify/simplify_stmt.cpp b/codon/parser/visitors/simplify/simplify_stmt.cpp
index d1ec10cb..535c0b9d 100644
--- a/codon/parser/visitors/simplify/simplify_stmt.cpp
+++ b/codon/parser/visitors/simplify/simplify_stmt.cpp
@@ -1043,7 +1043,7 @@ 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)
+    if (global && !isStatic && !(r && r->isType()))
       ctx->cache->globals[canonical] = nullptr;
     // Handle type aliases as well!
     ctx->add(r && r->isType() ? SimplifyItem::Type : SimplifyItem::Var, e->value,
diff --git a/codon/parser/visitors/translate/translate.cpp b/codon/parser/visitors/translate/translate.cpp
index f447c2a5..7cd40a6d 100644
--- a/codon/parser/visitors/translate/translate.cpp
+++ b/codon/parser/visitors/translate/translate.cpp
@@ -24,8 +24,13 @@ TranslateVisitor::TranslateVisitor(std::shared_ptr<TranslateContext> ctx)
 ir::Func *TranslateVisitor::apply(std::shared_ptr<Cache> cache, StmtPtr stmts) {
   ir::BodiedFunc *main;
   if (cache->isJit) {
-    main = cache->module->Nr<ir::BodiedFunc>(format("_jit_{}", cache->jitCell));
+    auto fnName = format("_jit_{}", cache->jitCell);
+    main = cache->module->Nr<ir::BodiedFunc>(fnName);
     main->setSrcInfo({"<jit>", 0, 0, 0});
+    main->setGlobal();
+    auto irType = cache->module->unsafeGetFuncType(
+        fnName, cache->classes["void"].realizations["void"]->ir, {}, false);
+    main->realize(irType, {});
   } else {
     main = cast<ir::BodiedFunc>(cache->module->getMainFunc());
     char buf[PATH_MAX + 1];
@@ -36,7 +41,11 @@ ir::Func *TranslateVisitor::apply(std::shared_ptr<Cache> cache, StmtPtr stmts) {
   auto block = cache->module->Nr<ir::SeriesFlow>("body");
   main->setBody(block);
 
-  cache->codegenCtx = std::make_shared<TranslateContext>(cache, block, main);
+  if (!cache->codegenCtx)
+    cache->codegenCtx = std::make_shared<TranslateContext>(cache);
+  cache->codegenCtx->bases = {main};
+  cache->codegenCtx->series = {block};
+
   TranslateVisitor(cache->codegenCtx).transform(stmts);
   return main;
 }
@@ -233,6 +242,7 @@ void TranslateVisitor::visit(AssignStmt *stmt) {
   auto var = stmt->lhs->getId()->value;
   if (!stmt->rhs && var == VAR_ARGV) {
     ctx->add(TranslateItem::Var, var, ctx->getModule()->getArgVar());
+    ctx->cache->globals[var] = ctx->getModule()->getArgVar();
   } else if (!stmt->rhs || !stmt->rhs->isType()) {
     auto *newVar =
         make<ir::Var>(stmt, getType((stmt->rhs ? stmt->rhs : stmt->lhs)->getType()),
diff --git a/codon/parser/visitors/translate/translate_ctx.cpp b/codon/parser/visitors/translate/translate_ctx.cpp
index 9506aef0..e3b58749 100644
--- a/codon/parser/visitors/translate/translate_ctx.cpp
+++ b/codon/parser/visitors/translate/translate_ctx.cpp
@@ -11,13 +11,9 @@
 namespace codon {
 namespace ast {
 
-TranslateContext::TranslateContext(std::shared_ptr<Cache> cache,
-                                   codon::ir::SeriesFlow *series,
-                                   codon::ir::BodiedFunc *base)
+TranslateContext::TranslateContext(std::shared_ptr<Cache> cache)
     : Context<TranslateItem>(""), cache(std::move(cache)) {
   stack.push_front(std::vector<std::string>());
-  bases.push_back(base);
-  addSeries(series);
 }
 
 std::shared_ptr<TranslateItem> TranslateContext::find(const std::string &name) const {
diff --git a/codon/parser/visitors/translate/translate_ctx.h b/codon/parser/visitors/translate/translate_ctx.h
index 6f641d50..c82172a2 100644
--- a/codon/parser/visitors/translate/translate_ctx.h
+++ b/codon/parser/visitors/translate/translate_ctx.h
@@ -52,8 +52,7 @@ struct TranslateContext : public Context<TranslateItem> {
   std::vector<codon::ir::SeriesFlow *> series;
 
 public:
-  TranslateContext(std::shared_ptr<Cache> cache, codon::ir::SeriesFlow *series,
-                   codon::ir::BodiedFunc *base);
+  TranslateContext(std::shared_ptr<Cache> cache);
 
   using Context<TranslateItem>::add;
   /// Convenience method for adding an object to the context.