From 69dfaa931ab30cad66095107e6c1e70d6155f107 Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Thu, 13 Jun 2024 16:21:19 -0400
Subject: [PATCH] xo-jit: + compile Apply expressions [wip]

---
 example/ex1/ex1.cpp    | 24 ++++++++++++++++
 include/xo/jit/Jit.hpp |  4 +++
 src/jit/Jit.cpp        | 63 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+)

diff --git a/example/ex1/ex1.cpp b/example/ex1/ex1.cpp
index 571b288f..bad239be 100644
--- a/example/ex1/ex1.cpp
+++ b/example/ex1/ex1.cpp
@@ -3,6 +3,7 @@
 #include "xo/jit/Jit.hpp"
 #include "xo/expression/Constant.hpp"
 #include "xo/expression/Primitive.hpp"
+#include "xo/expression/Apply.hpp"
 #include <iostream>
 
 int
@@ -10,6 +11,7 @@ main() {
     using xo::jit::Jit;
     using xo::ast::make_constant;
     using xo::ast::make_primitive;
+    using xo::ast::make_apply;
     using xo::xtag;
     using std::cerr;
     using std::endl;
@@ -51,6 +53,28 @@ main() {
                  << endl;
         }
     }
+
+    {
+        /* (sqrt 2) */
+
+        auto fn = make_primitive("sqrt", ::sqrt);
+        auto arg = make_constant(2.0);
+
+        auto call = make_apply(fn, arg);
+
+        auto llvm_ircode = jit->codegen(call);
+
+        if (llvm_ircode) {
+            /* note: llvm:errs() is 'raw stderr stream' */
+            cerr << "ex1 llvm_ircode:" << endl;
+            llvm_ircode->print(llvm::errs());
+            cerr << endl;
+        } else {
+            cerr << "ex1: code generation failed"
+                 << xtag("expr", call)
+                 << endl;
+        }
+    }
 }
 
 /** end ex1.cpp **/
diff --git a/include/xo/jit/Jit.hpp b/include/xo/jit/Jit.hpp
index f7adf0cc..7e32fb26 100644
--- a/include/xo/jit/Jit.hpp
+++ b/include/xo/jit/Jit.hpp
@@ -11,6 +11,7 @@
 #include "xo/expression/Expression.hpp"
 #include "xo/expression/ConstantInterface.hpp"
 #include "xo/expression/PrimitiveInterface.hpp"
+#include "xo/expression/Apply.hpp"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/BasicBlock.h"
@@ -50,6 +51,7 @@ namespace xo {
 
             llvm::Value * codegen_constant(ref::brw<xo::ast::ConstantInterface> expr);
             llvm::Function * codegen_primitive(ref::brw<xo::ast::PrimitiveInterface> expr);
+            llvm::Value * codegen_apply(ref::brw<xo::ast::Apply> expr);
 
             llvm::Value * codegen(ref::brw<Expression> expr);
 
@@ -64,6 +66,8 @@ namespace xo {
              *  each with its own LLVMContext
              **/
             std::unique_ptr<llvm::LLVMContext> llvm_cx_;
+            /** builder for intermediate-representation objects **/
+            std::unique_ptr<llvm::IRBuilder<>> llvm_ir_builder_;
             /** a module (aka library) being prepared by llvm.
              *  - function names are unique within a module.
              **/
diff --git a/src/jit/Jit.cpp b/src/jit/Jit.cpp
index b5d3241e..6bf39631 100644
--- a/src/jit/Jit.cpp
+++ b/src/jit/Jit.cpp
@@ -7,11 +7,15 @@ namespace xo {
     using xo::ast::Expression;
     using xo::ast::ConstantInterface;
     using xo::ast::PrimitiveInterface;
+    using xo::ast::Apply;
     using xo::reflect::TypeDescr;
+    using std::cerr;
+    using std::endl;
 
     namespace jit {
         Jit::Jit()
             : llvm_cx_{std::make_unique<llvm::LLVMContext>()},
+              llvm_ir_builder_{std::make_unique<llvm::IRBuilder<>>(*llvm_cx_)},
               llvm_module_{std::make_unique<llvm::Module>("xojit", *llvm_cx_)}
         {}
 
@@ -75,6 +79,60 @@ namespace xo {
             return fn;
         } /*codegen_primitive*/
 
+        llvm::Value *
+        Jit::codegen_apply(ref::brw<Apply> apply)
+        {
+            using std::cerr;
+            using std::endl;
+
+            /* editorial:
+             *
+             * to handle (computed functions) properly,
+             * we will need a runtime representation for a 'primitive function pointer'
+             *
+             * For now, finesse by only handling PrimitiveInterface in function-callee position
+             */
+            if (apply->fn()->extype() == exprtype::primitive) {
+                auto pm = PrimitiveInterface::from(apply->fn());
+                auto * fn = this->codegen_primitive(pm);
+
+#ifdef NOT_USING_DEBUG
+                cerr << "Jit::codegen_apply: fn:" << endl;
+                fn->print(llvm::errs());
+                cerr << endl;
+#endif
+
+                if (fn->arg_size() != apply->argv().size()) {
+                    cerr << "Jit::codegen_apply: error: callee f expecting n1 args where n2 supplied"
+                              << xtag("f", pm->name())
+                              << xtag("n1", pm->n_arg())
+                              << xtag("n2", apply->argv().size())
+                              << endl;
+                    return nullptr;
+                }
+
+                std::vector<llvm::Value *> args;
+                args.reserve(apply->argv().size());
+
+                for (const auto & arg_expr : apply->argv()) {
+                    auto * arg = this->codegen(arg_expr);
+
+#ifdef NOT_USING_DEBUG
+                    cerr << "Jit::codegen_apply: arg:" << endl;
+                    arg->print(llvm::errs());
+                    cerr << endl;
+#endif
+
+                    args.push_back(arg);
+                }
+
+                return llvm_ir_builder_->CreateCall(fn, args, "calltmp");
+            } else {
+                cerr << "Jit::codegen_apply: error: only allowing call to known primitives at present" << endl;
+                return nullptr;
+            }
+        } /*codegen_apply*/
+
         llvm::Value *
         Jit::codegen(ref::brw<Expression> expr)
         {
@@ -84,6 +142,7 @@ namespace xo {
             case exprtype::primitive:
                 return this->codegen_primitive(PrimitiveInterface::from(expr));
             case exprtype::apply:
+                return this->codegen_apply(Apply::from(expr));
                 break;
             case exprtype::invalid:
             case exprtype::n_expr:
@@ -91,6 +150,10 @@ namespace xo {
                 break;
             }
 
+            cerr << "Jit::codegen: error: no handler for expression of type T"
+                 << xtag("T", expr->extype())
+                 << endl;
+
             return nullptr;
         } /*codegen*/