From e671686a3a0c65208ae01f31679d448a3a32b9e6 Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Fri, 21 Jun 2024 17:00:06 -0400
Subject: [PATCH] xo-jit: refactor MachPipeline to use stack for lambda formals

---
 include/xo/jit/MachPipeline.hpp | 18 ++++++++++++-
 src/jit/MachPipeline.cpp        | 46 +++++++++++++++++++++++++++++++--
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/include/xo/jit/MachPipeline.hpp b/include/xo/jit/MachPipeline.hpp
index 967e3127..75244b93 100644
--- a/include/xo/jit/MachPipeline.hpp
+++ b/include/xo/jit/MachPipeline.hpp
@@ -54,6 +54,7 @@ namespace xo {
         class MachPipeline : public ref::Refcount {
         public:
             using Expression = xo::ast::Expression;
+            using TypeDescr = xo::reflect::TypeDescr;
             //using ConstantInterface = xo::ast::ConstantInterface;
 
         public:
@@ -111,6 +112,16 @@ namespace xo {
             /** iniitialize native builder (i.e. for platform we're running on) **/
             static void init_once();
 
+            /** codegen helper for a user-defined function (codegen_lambda()):
+             *  create stack slot on behalf of some formal parameter to a function,
+             *  so we can avoid SSA restriction on function body
+             *
+             *  @p var_type.  variable type
+             **/
+            llvm::AllocaInst * create_entry_block_alloca(llvm::Function * llvm_fn,
+                                                         const std::string & var_name,
+                                                         TypeDescr var_type);
+
             /** (re)create pipeline to turn expressions into llvm IR code **/
             void recreate_llvm_ir_pipeline();
 
@@ -155,8 +166,13 @@ namespace xo {
              *  corresponding llvm IR.
              *
              *  only supports one level atm (i.e. only top-level functions)
+             *
+             *  All values  live on the stack, so that we can evade single-assignment
+             *  restrictions.
+             *
+             *  rhs identifies logical stack location of a variable
              **/
-            std::map<std::string, llvm::Value*> nested_env_;
+            std::map<std::string, llvm::AllocaInst*> nested_env_; /* <-> kaleidoscope NamedValues */
 
         }; /*MachPipeline*/
 
diff --git a/src/jit/MachPipeline.cpp b/src/jit/MachPipeline.cpp
index d0abc7c6..c1c65b37 100644
--- a/src/jit/MachPipeline.cpp
+++ b/src/jit/MachPipeline.cpp
@@ -391,6 +391,25 @@ namespace xo {
             }
         } /*codegen_apply*/
 
+        /* in kaleidoscope7.cpp: CreateEntryBlockAlloca */
+        llvm::AllocaInst *
+        MachPipeline::create_entry_block_alloca(llvm::Function * llvm_fn,
+                                                const std::string & var_name,
+                                                TypeDescr var_type)
+        {
+            llvm::IRBuilder<> tmp_ir_builder(&llvm_fn->getEntryBlock(),
+                                             llvm_fn->getEntryBlock().begin());
+
+            llvm::Type * llvm_var_type = td_to_llvm_type(llvm_cx_.borrow(),
+                                                         var_type);
+            if (!llvm_var_type)
+                return nullptr;
+
+            return tmp_ir_builder.CreateAlloca(llvm_var_type,
+                                               nullptr,
+                                               var_name);
+        } /*create_entry_block_alloca*/
+
         llvm::Function *
         MachPipeline::codegen_lambda(ref::brw<Lambda> lambda)
         {
@@ -462,7 +481,25 @@ namespace xo {
                 for (auto & arg : fn->args()) {
                     log && log("nested environment", xtag("i", i), xtag("param", std::string(arg.getName())));
 
-                    nested_env_[std::string(arg.getName())] = &arg;
+                    /* stack location for arg[i] */
+                    llvm::AllocaInst * alloca
+                        = create_entry_block_alloca(fn,
+                                                    std::string(arg.getName()),
+                                                    lambda->fn_arg(i));
+
+                    if (!alloca)
+                        return nullptr;
+
+                    /* store on function entry
+                     *   see codegen_variable() for corresponding load
+                     */
+                    this->llvm_ir_builder_->CreateStore(&arg, alloca);
+
+                    /* remember stack location for reference + assignment
+                     * in lambda body.
+                     *
+                     */
+                    nested_env_[std::string(arg.getName())] = alloca;
                     ++i;
                 }
             }
@@ -500,7 +537,12 @@ namespace xo {
                 return nullptr;
             }
 
-            return ix->second;
+            llvm::AllocaInst * alloca = ix->second;
+
+            /* code to load value from stack */
+            return this->llvm_ir_builder_->CreateLoad(alloca->getAllocatedType(),
+                                                      alloca,
+                                                      var->name().c_str());
         } /*codegen_variable*/
 
         llvm::Value *