From 32c74511914a77a29729b6bcacfe3cd2a0810cfd Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Mon, 17 Jun 2024 12:29:15 -0400
Subject: [PATCH] xo-jit: + codegen for if-expressions

---
 include/xo/jit/MachPipeline.hpp |  5 +-
 src/jit/MachPipeline.cpp        | 82 +++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/include/xo/jit/MachPipeline.hpp b/include/xo/jit/MachPipeline.hpp
index 07fd6050..ca59dfb3 100644
--- a/include/xo/jit/MachPipeline.hpp
+++ b/include/xo/jit/MachPipeline.hpp
@@ -18,7 +18,7 @@
 #include "xo/expression/Apply.hpp"
 #include "xo/expression/Lambda.hpp"
 #include "xo/expression/Variable.hpp"
-
+#include "xo/expression/IfExpr.hpp"
 
 /* stuff from kaleidoscope.cpp */
 #include "llvm/ADT/APFloat.h"
@@ -81,6 +81,7 @@ namespace xo {
             llvm::Value * codegen_apply(ref::brw<xo::ast::Apply> expr);
             llvm::Function * codegen_lambda(ref::brw<xo::ast::Lambda> expr);
             llvm::Value * codegen_variable(ref::brw<xo::ast::Variable> var);
+            llvm::Value * codegen_ifexpr(ref::brw<xo::ast::IfExpr> ifexpr);
 
             llvm::Value * codegen(ref::brw<Expression> expr);
 
@@ -117,7 +118,7 @@ namespace xo {
 
             // ----- this part adapted from kaleidoscope.cpp -----
 
-            /** everything bleow represents a pipeline
+            /** everything below represents a pipeline
              *  that takes expressions, and turns them into llvm IR.
              *
              *  llvm IR can be added to running JIT by calling
diff --git a/src/jit/MachPipeline.cpp b/src/jit/MachPipeline.cpp
index 1619d121..18c7245e 100644
--- a/src/jit/MachPipeline.cpp
+++ b/src/jit/MachPipeline.cpp
@@ -10,6 +10,7 @@ namespace xo {
     using xo::ast::Lambda;
     using xo::ast::Variable;
     using xo::ast::Apply;
+    using xo::ast::IfExpr;
     using xo::reflect::TypeDescr;
     using std::cerr;
     using std::endl;
@@ -370,6 +371,85 @@ namespace xo {
             return ix->second;
         } /*codegen_variable*/
 
+        llvm::Value *
+        MachPipeline::codegen_ifexpr(ref::brw<IfExpr> expr)
+        {
+            llvm::Value * test_ir = this->codegen(expr->test());
+
+            /** need test result in a variable
+             **/
+            llvm::Value * test_with_cmp_ir
+                = llvm_ir_builder_->CreateFCmpONE(test_ir,
+                                                  llvm::ConstantFP::get(llvm_cx_->llvm_cx_ref(),
+                                                                        llvm::APFloat(0.0)),
+                                                  "iftest");
+
+            llvm::Function * parent_fn = llvm_ir_builder_->GetInsertBlock()->getParent();
+
+            /* when_true_bb, when_false_bb, merge_bb:
+             * initially-empty basic-blocks for {when_true, when_false, merged} codegen
+             */
+
+            /* when_true branch inserted at (current) end of function */
+            llvm::BasicBlock * when_true_bb
+                = llvm::BasicBlock::Create(llvm_cx_->llvm_cx_ref(),
+                                           "when_true",
+                                           parent_fn);
+            llvm::BasicBlock * when_false_bb
+                = llvm::BasicBlock::Create(llvm_cx_->llvm_cx_ref(),
+                                           "when_false");
+
+            llvm::BasicBlock * merge_bb
+                = llvm::BasicBlock::Create(llvm_cx_->llvm_cx_ref(),
+                                           "merge");
+
+            /* IR to direct control flow to one of {when_true_bb, when_false_bb},
+             * depending on result of test_with_cmp_ir
+             */
+            llvm_ir_builder_->CreateCondBr(test_with_cmp_ir,
+                                           when_true_bb,
+                                           when_false_bb);
+
+            /* populate when_true_bb */
+            llvm_ir_builder_->SetInsertPoint(when_true_bb);
+
+            llvm::Value * when_true_ir = this->codegen(expr->when_true());
+
+            if (!when_true_ir)
+                return nullptr;
+
+            /* at end of when-true sequence, jump to merge suffix */
+            llvm_ir_builder_->CreateBr(merge_bb);
+            /* note: codegen for expr->when_true() may have altered builder's "current block" */
+            when_true_bb = llvm_ir_builder_->GetInsertBlock();
+
+            /* populate when_false_bb */
+            parent_fn->insert(parent_fn->end(), when_false_bb);
+            llvm_ir_builder_->SetInsertPoint(when_false_bb);
+
+            llvm::Value * when_false_ir = this->codegen(expr->when_false());
+            if (!when_false_ir)
+                return nullptr;
+
+            /* at end of when-false sequence, jump to merge suffix */
+            llvm_ir_builder_->CreateBr(merge_bb);
+            /* note: codegen for expr->when_false() may have altered builder's "current block" */
+            when_false_bb = llvm_ir_builder_->GetInsertBlock();
+
+            /* merged suffix sequence */
+            parent_fn->insert(parent_fn->end(), merge_bb);
+            llvm_ir_builder_->SetInsertPoint(merge_bb);
+
+            llvm::PHINode * phi_node
+                = llvm_ir_builder_->CreatePHI(llvm::Type::getDoubleTy(llvm_cx_->llvm_cx_ref()),
+                                              2 /*#of branches being merged (?)*/,
+                                              "iftmp");
+            phi_node->addIncoming(when_true_ir, when_true_bb);
+            phi_node->addIncoming(when_false_ir, when_false_bb);
+
+            return phi_node;
+        }
+
         llvm::Value *
         MachPipeline::codegen(ref::brw<Expression> expr)
         {
@@ -384,6 +464,8 @@ namespace xo {
                 return this->codegen_lambda(Lambda::from(expr));
             case exprtype::variable:
                 return this->codegen_variable(Variable::from(expr));
+            case exprtype::if_expr:
+                return this->codegen_ifexpr(IfExpr::from(expr));
             case exprtype::invalid:
             case exprtype::n_expr:
                 return nullptr;