global env; alloca locals

kinda messy

10 files changed, 212 insertions, 95 deletions

diff --git a/ast.cpp b/ast.cpp
index a528cc3..c7a9309 100644
--- a/ast.cpp
+++ b/ast.cpp
@@ -20,7 +20,24 @@
 #include <charconv>
 #include <iostream>
 
+extern std::list<ThunkAST> scope;
+
 int ThunkAST::tcount = 0;
+int ThunkAST::envidx = 0;
+
+static inline auto loadEnv(LLVMState& llvmState, llvm::Value *index)
+{
+    auto ptrty = llvmState.inttype->getPointerTo();
+    auto gep = llvmState.builder.CreateGEP(ptrty, scope.back().env, {index});
+    return llvmState.builder.CreateLoad(ptrty, gep);
+}
+
+static inline auto storeEnv(LLVMState& llvmState, llvm::Value *index, llvm::Value *val)
+{
+    auto ptrty = llvmState.inttype->getPointerTo();
+    auto var = llvmState.builder.CreateGEP(ptrty, scope.back().env, {index});
+    return llvmState.builder.CreateStore(val, var, false);
+}
 
 NumberAST::NumberAST(const std::string& n): BaseAST(n) {}
 
@@ -34,7 +51,7 @@ llvm::Value *NumberAST::codegen(LLVMState& llvmState) const
         return nullptr;
     } else {
         auto val = llvmState.createInt(value);
-        return llvmState.builder.CreateStore(val, llvmState.createPush());
+        return llvmState.createPush(val);
     }
 }
 
@@ -42,12 +59,16 @@ PushAST::PushAST(const std::string& n): BaseAST(n) {}
 
 llvm::Value *PushAST::codegen(LLVMState& llvmState) const
 {
-    if (auto [var, thunk] = Var::lookup(name); var) {
-        auto dsget = llvmState.createPush();
-        if (!thunk)
-            var = llvmState.builder.CreateLoad(llvmState.inttype, var);
+    if (auto [var, native] = Var::lookup(name, 1); var) {
+        auto index = llvm::ConstantInt::get(llvmState.inttype, ThunkAST::envidx++);
+        Var::addLocal(name, index);
+    }
+
+    if (auto [var, native] = Var::lookupLocal(name); var) {
+        if (!native)
+            var = loadEnv(llvmState, var);
 
-        return llvmState.builder.CreateStore(var, dsget);
+        return llvmState.createPush(var);
     } else {
         std::cerr << "error: not defined: " << name << std::endl;
         return nullptr;
@@ -58,42 +79,77 @@ PopAST::PopAST(const std::string& n): BaseAST(n) {}
 
 llvm::Value *PopAST::codegen(LLVMState& llvmState) const
 {
+    Var v;
+
     if (name == "self") {
-        extern std::list<ThunkAST> scope;
-        auto func = scope.back().func;
-        auto var = llvmState.createVariable(name);
-        llvmState.builder.CreateStore(func, var, false);
-        return Var::addLocal(name, var).value;
+        v = {scope.back().func, true};
     } else {
-        auto gep = llvmState.createPop();
-        auto var = llvmState.createVariable(name);
-        auto load = llvmState.builder.CreateLoad(llvmState.inttype, gep);
-        llvmState.builder.CreateStore(load, var, false);
+        auto index = llvm::ConstantInt::get(llvmState.inttype, ThunkAST::envidx++);
+        auto pop = llvmState.createPop();
+        storeEnv(llvmState, index, pop);
 
-        return Var::addLocal(name, var).value;
+        v = index;
     }
+
+    return Var::addLocal(name, v).value;
 }
 
 CallAST::CallAST(const std::string& n): BaseAST(n) {}
 
 llvm::Value *CallAST::codegen(LLVMState& llvmState) const
 {
-    if (auto [var, call] = Var::lookup(name); var) {
-        if (call) {
-            return llvmState.builder.CreateCall(llvmState.ftype, var);
-        } else {
-            auto val = llvmState.builder.CreateLoad(llvmState.inttype, var);
-            auto cast = llvmState.builder.CreateIntToPtr(val,
-                llvmState.inttype->getPointerTo());
-            return llvmState.builder.CreateCall(llvmState.ftype, cast);
-        }
+    if (auto [var, native] = Var::lookup(name, 1); var && !native) {
+        auto index = llvm::ConstantInt::get(llvmState.inttype, ThunkAST::envidx++);
+        Var::addLocal(name, index);
+    }
+
+    llvm::Value *fn;
+    if (auto [var, native] = Var::lookup(name); var) {
+        if (!native)
+            var = loadEnv(llvmState, var);
+
+        fn = var;
     } else {
         std::cerr << "warning: anticipating external function: "
             << name << std::endl;
-        auto func = llvmState.createFunction(name);
-        Var::addGlobal(name, Var {func, true});
-        return llvmState.builder.CreateCall(llvmState.ftype, func);
+
+        fn = llvmState.createFunction(name);
+        Var::addGlobal(name, Var {fn, true});
+    }
+
+    auto ptrty = llvmState.inttype->getPointerTo();
+    llvm::Type *type;
+    llvm::Value *mem;
+
+    if (auto sz = Var::vars.back().size(); sz > 0) {
+        type = llvm::VectorType::get(llvmState.inttype, sz, false);
+        mem = llvmState.builder.CreateAlloca(type, nullptr);
+
+        int i = 0;
+        for (auto& [_, v] : Var::vars.back()) {
+            if (!v.native) {
+                auto index = llvm::ConstantInt::get(llvmState.inttype, i++);
+                auto m = llvmState.builder.CreateGEP(ptrty, mem, {index});
+                llvmState.builder.CreateStore(loadEnv(llvmState, v.value), m, false);
+            }
+        }
+    }
+
+    auto call = llvmState.builder.CreateCall(llvmState.ftype, fn, llvm::ArrayRef {scope.back().env});
+
+    if (auto sz = Var::vars.back().size(); sz > 0) {
+        int i = 0;
+        for (auto& [_, v] : Var::vars.back()) {
+            if (!v.native) {
+                auto index = llvm::ConstantInt::get(llvmState.inttype, i++);
+                auto m = llvmState.builder.CreateGEP(ptrty, mem, {index});
+                auto l = llvmState.builder.CreateLoad(ptrty, m);
+                storeEnv(llvmState, v.value, l);
+            }
+        }
     }
+
+    return call;
 }
 
 ThunkAST::ThunkAST(LLVMState& llvmState):
@@ -103,15 +159,30 @@ ThunkAST::ThunkAST(LLVMState& llvmState, std::string n): BaseAST(n)
 {
     parent = llvmState.builder.saveIP();
     func = llvmState.createFunction(name);
-    auto BB = llvmState.createEntry(func);
-    llvmState.builder.SetInsertPoint(BB);
+    entry = llvmState.createEntry(func);
+    body = llvm::BasicBlock::Create(llvmState.ctx, "body", func);
+    env = func->getArg(0);
+
+    llvmState.builder.SetInsertPoint(entry);
+    llvmState.builder.SetInsertPoint(body);
 }
 
 llvm::Value *ThunkAST::codegen(LLVMState& llvmState) const
 {
     llvmState.builder.CreateRetVoid();
-    llvmState.builder.restoreIP(parent);
+    llvmState.builder.SetInsertPoint(entry);
+
+    if (Var::vars.back().size() > 0) {
+        for (auto& [n, v] : Var::vars.back()) {
+            if (auto [c, _] = Var::lookup(n, 1); c) {
+                auto src = loadEnv(llvmState, c);
+                storeEnv(llvmState, v.value, src);
+            }
+        }
+    }
 
+    llvmState.builder.CreateBr(body);
+    llvmState.builder.restoreIP(parent);
     return func;
 }
 
diff --git a/ast.hpp b/ast.hpp
index 7d213ba..b7802c1 100644
--- a/ast.hpp
+++ b/ast.hpp
@@ -66,10 +66,12 @@ struct CallAST : public BaseAST
 struct ThunkAST : public BaseAST
 {
     static int tcount;
+    static int envidx;
 
-    std::list<std::unique_ptr<BaseAST>> body;
     llvm::IRBuilderBase::InsertPoint parent;
     llvm::Function *func;
+    llvm::BasicBlock *entry, *body;
+    llvm::Value *env;
 
     explicit ThunkAST(LLVMState& llvmState);
     explicit ThunkAST(LLVMState& llvmState, std::string n);
diff --git a/examples/fibonacci.fp b/examples/fibonacci.fp
new file mode 100644
index 0000000..493ca87
--- /dev/null
+++ b/examples/fibonacci.fp
@@ -0,0 +1,56 @@
+(
+  ; core utilities
+  ($x ^x ^x)          $dup
+  ($_)                $drop
+  ($x $y ^x ^y)       $swap
+  ($a $b $c ^b ^a ^c) $rot
+  (sub)               $-
+  (0 swap - -)        $+
+  (())                $nil
+  (nil eq)            $null?
+  ($x x)              $force
+  (10 emit)           $cr
+  (print)             $print
+  (4 +)               $cell+
+  ($a $d 2 alloc $p
+    ^a ^p poke
+    ^d ^p cell+ poke
+    ^p
+  )                   $cons
+  (peek)              $car
+  (cell+ peek)        $cdr
+
+  ; if-stmt
+  ($c $t $f c ^f ^t rot cswap $_ force) $if
+  ($f $t $c $fn ^f ^t ^c fn)     $endif
+
+  ; range
+  ($self $start $end
+    ^if (^start ^end eq)
+      ^nil
+      (^start ^end ^start 1 + self swap cons)
+    endif
+  ) $range
+
+  ; map [$fn $list -> $out-list]
+  ($self $fn $list
+    ^if (^list null?)
+      ^nil
+      (^list car fn ^list cdr ^fn self swap cons)
+    endif
+  ) $map
+
+  ; each [$fn $list]
+  ($fn (fn ^nil) map drop) $each
+
+  ; implementation
+  (0 1 ($self $a $b $n
+    ^if (^n 0 eq) (^b) (
+      ^n 1 - ^a ^b + ^b self
+    ) endif
+  ) force) $fibonacci
+
+  30 10 range
+  ^fibonacci map
+  ^print each
+)
diff --git a/llvm.cpp b/llvm.cpp
index f1ee27a..51d4706 100644
--- a/llvm.cpp
+++ b/llvm.cpp
@@ -22,8 +22,8 @@ LLVMState::LLVMState():
     modul("forsp", ctx),
     builder(ctx),
     inttype(llvm::Type::getInt32Ty(ctx)),
-    stacktype(llvm::VectorType::get(inttype, 12, false)),
-    ftype(llvm::FunctionType::get(llvm::Type::getVoidTy(ctx), {}, false)),
+    stacktype(llvm::VectorType::get(inttype, 16, false)),
+    ftype(llvm::FunctionType::get(llvm::Type::getVoidTy(ctx), llvm::ArrayRef<llvm::Type *> {inttype->getPointerTo()}, false)),
     one(llvm::ConstantInt::get(inttype, 1)),
     zero(llvm::ConstantInt::get(inttype, 0))
 {
@@ -34,13 +34,14 @@ LLVMState::LLVMState():
         llvm::GlobalValue::ExternalLinkage, zerovec, "stack");
 }
 
-llvm::Value *LLVMState::createPush()
+llvm::Value *LLVMState::createPush(llvm::Value *var)
 {
     auto dspval = builder.CreateLoad(inttype, llvmSp);
     auto inc = builder.CreateAdd(dspval, one);
     builder.CreateStore(inc, llvmSp, false);
 
-    return builder.CreateGEP(stacktype, llvmStack, {zero, dspval});
+    auto gep = builder.CreateGEP(stacktype, llvmStack, {zero, dspval});
+    return builder.CreateStore(var, gep);
 }
 
 llvm::Value *LLVMState::createPop()
@@ -49,13 +50,16 @@ llvm::Value *LLVMState::createPop()
     auto dec = builder.CreateSub(dspval, one);
     builder.CreateStore(dec, llvmSp, false);
 
-    return builder.CreateGEP(stacktype, llvmStack, {zero, dec});
+    auto gep = builder.CreateGEP(stacktype, llvmStack, {zero, dec});
+    return builder.CreateLoad(inttype, gep);
 }
 
 llvm::Function *LLVMState::createFunction(const std::string& name)
 {
-    return llvm::Function::Create(ftype, llvm::Function::ExternalLinkage,
+    auto func = llvm::Function::Create(ftype, llvm::Function::ExternalLinkage,
         name.c_str(), modul);
+    func->getArg(0)->setName("penv");
+    return func;
 }
 
 llvm::BasicBlock *LLVMState::createEntry(llvm::Function *func)
diff --git a/llvm.hpp b/llvm.hpp
index 0cd132f..09e2012 100644
--- a/llvm.hpp
+++ b/llvm.hpp
@@ -44,7 +44,7 @@ struct LLVMState
 
     LLVMState();
 
-    llvm::Value *createPush();
+    llvm::Value *createPush(llvm::Value *var);
     llvm::Value *createPop();
     llvm::Function *createFunction(const std::string& name);
     llvm::BasicBlock *createEntry(llvm::Function *func);
diff --git a/main.cpp b/main.cpp
index 19b76fe..5ae41bc 100644
--- a/main.cpp
+++ b/main.cpp
@@ -48,7 +48,17 @@ int main()
         }
     }
 
+    auto func = llvmState.createFunction("main");
+    auto entry = llvmState.createEntry(func);
+    auto envtype = llvm::VectorType::get(llvmState.inttype, ThunkAST::envidx, false);
+    auto [t0, _] = Var::lookup("__t0");
+    llvmState.builder.SetInsertPoint(entry);
+    auto env = llvmState.builder.CreateAlloca(envtype, nullptr);
+    llvmState.builder.CreateCall(llvmState.ftype, t0, llvm::ArrayRef<llvm::Value *> {env});
+    llvmState.builder.CreateRetVoid();
+
     llvmState.output();
+    std::cerr << "envidx: " << ThunkAST::envidx << std::endl;
     std::cout << std::endl;
 }
 
@@ -65,10 +75,7 @@ bool parseString(std::string_view sv)
 
             switch (tok) {
             case Token::ThunkOpen:
-                if (scope.empty())
-                    scope.emplace_back(llvmState, "main");
-                else
-                    scope.emplace_back(llvmState);
+                scope.emplace_back(llvmState);
                 Var::pushScope();
                 break;
             case Token::ThunkClose:
diff --git a/support.c b/support.c
index b400a05..e079de6 100644
--- a/support.c
+++ b/support.c
@@ -10,7 +10,7 @@ void emit()
     putchar(*(&stack + --sp));
 }
 
-void print()
+void _print()
 {
     printf("%d\n", *(&stack + --sp));
 }
diff --git a/test.fp b/test.fp
index 2ae975d..222d919 100644
--- a/test.fp
+++ b/test.fp
@@ -1,56 +1,25 @@
 (
   ; core utilities
-  ($x ^x ^x)          $dup
   ($_)                $drop
   ($x $y ^x ^y)       $swap
   ($a $b $c ^b ^a ^c) $rot
   (sub)               $-
   (0 swap - -)        $+
-  (())                $nil
-  (nil eq)            $null?
   ($x x)              $force
+  (32 emit)           $bl
   (10 emit)           $cr
-  (print)             $print
-  (4 +)               $cell+
-  ($a $d 2 alloc $p
-    ^a ^p poke
-    ^d ^p cell+ poke
-    ^p
-  )                   $cons
-  (peek)              $car
-  (cell+ peek)        $cdr
 
   ; if-stmt
   ($c $t $f c ^f ^t rot cswap $_ force) $if
   ($f $t $c $fn ^f ^t ^c fn)     $endif
 
   ; range
-  ($self $start $end
-    ^if (^start ^end eq)
-      ^nil
-      (^start ^end ^start 1 + self swap cons)
+  ($self $fn $start $end
+    ^if (^start ^end eq) ()
+      (^start fn ^end ^start 1 + ^fn self)
     endif
-  ) $range
+  ) $each
 
-  ; map [$fn $list -> $out-list]
-  ($self $fn $list
-    ^if (^list null?)
-      ^nil
-      (^list car fn ^list cdr ^fn self swap cons)
-    endif
-  ) $map
-
-  ; each [$fn $list]
-  ($fn (fn ^nil) map drop) $each
-
-  ; implementation
-  (0 1 ($self $a $b $n
-    ^if (^n 0 eq) (^b) (
-      ^n 1 - ^a ^b + ^b self
-    ) endif
-  ) force) $fibonacci
-
-  10 1 range
-  ^fibonacci map
-  ^print each
+  10 0 ($i 10 0 ($j ^i 48 + emit ^j 48 + emit bl) each) each
+  cr
 )
diff --git a/var.cpp b/var.cpp
index 6fe9104..c907d7a 100644
--- a/var.cpp
+++ b/var.cpp
@@ -17,15 +17,11 @@
  */
 #include "var.hpp"
 
-#include <list>
-#include <map>
-#include <string>
-
-static std::list<std::map<std::string, Var>> llvmVars;
+std::list<std::map<std::string, Var>> Var::vars;
 
 Var Var::lookup(const std::string& name, int skip)
 {
-    for (auto sc = llvmVars.rbegin(); sc != llvmVars.rend(); ++sc) {
+    for (auto sc = vars.rbegin(); sc != vars.rend(); ++sc) {
         if (skip > 0) {
             --skip;
             continue;
@@ -37,23 +33,28 @@ Var Var::lookup(const std::string& name, int skip)
     return {};
 }
 
+Var Var::lookupLocal(const std::string& name)
+{
+    return vars.back().contains(name) ? vars.back()[name] : Var();
+}
+
 void Var::pushScope()
 {
-    llvmVars.emplace_back();
+    vars.emplace_back();
 }
 
 void Var::popScope()
 {
-    llvmVars.pop_back();
+    vars.pop_back();
 }
 
 Var& Var::addGlobal(const std::string& name, Var var)
 {
-    return llvmVars.front().emplace(name, var).first->second;
+    return vars.front().emplace(name, var).first->second;
 }
 
 Var& Var::addLocal(const std::string& name, Var var)
 {
-    return llvmVars.back().emplace(name, var).first->second;
+    return vars.back().emplace(name, var).first->second;
 }
 
diff --git a/var.hpp b/var.hpp
index f61883f..1ffb8b8 100644
--- a/var.hpp
+++ b/var.hpp
@@ -21,14 +21,21 @@
 #include <llvm/IR/Type.h>
 #include <llvm/IR/DerivedTypes.h>
 
+#include <list>
+#include <map>
+#include <string>
+
 struct Var {
     llvm::Value *value;
-    bool callable;
+    bool native;
+
+    Var(llvm::Value *v = nullptr, bool n = false):
+        value(v), native(n) {}
 
-    Var(llvm::Value *v = nullptr, bool c = false):
-        value(v), callable(c) {}
+    static std::list<std::map<std::string, Var>> vars;
 
     static Var lookup(const std::string& name, int skip = 0);
+    static Var lookupLocal(const std::string& name);
     static void pushScope();
     static void popScope();
     static Var& addGlobal(const std::string& name, Var var);