shortcut push-pops; don't env-copy functions

4 files changed, 43 insertions, 8 deletions

diff --git a/ast.cpp b/ast.cpp
index 917d88f..095c373 100644
--- a/ast.cpp
+++ b/ast.cpp
@@ -60,8 +60,12 @@ PushAST::PushAST(const std::string& n): BaseAST(n) {}
 llvm::Value *PushAST::codegen(LLVMState& llvmState) const
 {
     if (auto [var, native] = Var::lookup(name, 1); var) {
-        auto index = llvm::ConstantInt::get(llvmState.inttype, llvmState.envidx++);
-        Var::addLocal(name, index);
+        if (!native) {
+            auto index = llvm::ConstantInt::get(llvmState.inttype, llvmState.envidx++);
+            Var::addLocal(name, index);
+        } else {
+            Var::addLocal(name, Var {var, native});
+        }
     }
 
     if (auto [var, native] = Var::lookupLocal(name); var) {
@@ -182,9 +186,11 @@ void ThunkAST::endGen(LLVMState& llvmState)
 
     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);
+            if (auto [c, nat] = Var::lookup(n, 1); c) {
+                if (!nat) {
+                    auto src = loadEnv(llvmState, c);
+                    storeEnv(llvmState, v.value, src);
+                }
             }
         }
     }
diff --git a/main.cpp b/main.cpp
index 8e65ec8..1880112 100644
--- a/main.cpp
+++ b/main.cpp
@@ -146,7 +146,9 @@ std::list<ThunkAST>::iterator buildThunk(std::list<ThunkAST>::iterator it)
     curEnv = it->env;
     llvmState.lastSp = it->lastSp;
 
-    for (auto& a : it->ast) {
+    for (auto ait = it->ast.begin(); ait != it->ast.end(); ++ait) {
+        auto& a = *ait;
+
         if (a->name.starts_with("__t")) {
             it->lastSp = llvmState.lastSp;
             next = buildThunk(next);
@@ -155,8 +157,17 @@ std::list<ThunkAST>::iterator buildThunk(std::list<ThunkAST>::iterator it)
             llvmState.lastSp = it->lastSp;
         }
 
-        if (a->codegen(llvmState) == nullptr) {
-            return scope.end();
+        auto next = ait;
+        ++next;
+        if (next != it->ast.end() && dynamic_cast<PushAST *>(ait->get()) && dynamic_cast<PopAST *>(next->get())) {
+            // Value (e.g. thunk) is immediately pushed then popped into a variable.
+            // Shortcut and simply rename the value:
+            Var::rename(a->name, (*next)->name);
+            it->ast.erase(next);
+        } else {
+            // Otherwise, just do the regular codegen.
+            if (a->codegen(llvmState) == nullptr)
+                return scope.end();
         }
     }
 
diff --git a/var.cpp b/var.cpp
index c907d7a..f4d64fd 100644
--- a/var.cpp
+++ b/var.cpp
@@ -33,6 +33,23 @@ Var Var::lookup(const std::string& name, int skip)
     return {};
 }
 
+void Var::rename(const std::string& old, const std::string& neww, int skip)
+{
+    for (auto sc = vars.rbegin(); sc != vars.rend(); ++sc) {
+        if (skip > 0) {
+            --skip;
+            continue;
+        }
+        if (sc->contains(old)) {
+            auto it = sc->find(old);
+            (*sc)[neww] = it->second;
+            //it->second.value->setName(neww);
+            sc->erase(it);
+            break;
+        }
+    }
+}
+
 Var Var::lookupLocal(const std::string& name)
 {
     return vars.back().contains(name) ? vars.back()[name] : Var();
diff --git a/var.hpp b/var.hpp
index 1ffb8b8..3a5665c 100644
--- a/var.hpp
+++ b/var.hpp
@@ -36,6 +36,7 @@ struct Var {
 
     static Var lookup(const std::string& name, int skip = 0);
     static Var lookupLocal(const std::string& name);
+    static void rename(const std::string& old, const std::string& neww, int skip = 0);
     static void pushScope();
     static void popScope();
     static Var& addGlobal(const std::string& name, Var var);