From dbe7572eb5ceab049e7899bf3513d2ccc8cfc456 Mon Sep 17 00:00:00 2001
From: Danila Fedorin <danila.fedorin@gmail.com>
Date: Wed, 12 Jun 2019 04:11:28 -0700
Subject: [PATCH] Store arity of functions in llvm context to generate correct
 calls.

---
 src/gmachine.cpp |  2 +-
 src/llvm.cpp     | 10 ++++++++--
 src/llvm.hpp     |  7 +++++--
 src/main.cpp     |  4 ++--
 src/parser.cpp   |  5 +++--
 5 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/gmachine.cpp b/src/gmachine.cpp
index f61e643..b989c44 100644
--- a/src/gmachine.cpp
+++ b/src/gmachine.cpp
@@ -109,7 +109,7 @@ namespace lily {
     void instruction_push_global::gen_llvm(llvm_context& ctx) {
         llvm::Value* stack = ctx.get_current_function()->arg_begin();
         llvm::Value* new_node = builder.CreateCall(malloc_node_global_func,
-                { get_int8_constant(2), ctx.get_supercombinator(name) }, "temp");
+                { get_int8_constant(ctx.get_supercombinator_arity(name)), ctx.get_supercombinator_function(name) }, "temp");
         // TODO get arity
         builder.CreateCall(stack_push_func, { stack, new_node });
     }
diff --git a/src/llvm.cpp b/src/llvm.cpp
index c672d80..1d06d0e 100644
--- a/src/llvm.cpp
+++ b/src/llvm.cpp
@@ -217,12 +217,17 @@ namespace lily {
         }
     }
 
-    llvm::Function* llvm_context::get_supercombinator(const std::string& name) {
+    llvm::Function* llvm_context::get_supercombinator_function(const std::string& name) {
         if(!supercombinators.count(name)) throw error("unknown supercombinator");
         return supercombinators.find(name)->second;
     }
 
-    void llvm_context::add_supercombinator(const std::string& name) {
+    int llvm_context::get_supercombinator_arity(const std::string& name) {
+        if(!arities.count(name)) throw error("unknown supercombinator");
+        return arities.find(name)->second;
+    }
+
+    void llvm_context::add_supercombinator(const std::string& name, int arity) {
         if(supercombinators.count(name)) throw error("re-creating supercombinator");
         llvm::Function* new_function = llvm::Function::Create(
                 supercomb_function_type,
@@ -231,6 +236,7 @@ namespace lily {
                 &module);
         new_function->arg_begin()->setName("stack");
         supercombinators[name] = new_function;
+        arities[name] = arity;
     }
 
     llvm::Function* llvm_context::get_current_function() {
diff --git a/src/llvm.hpp b/src/llvm.hpp
index 50bb06c..de47efb 100644
--- a/src/llvm.hpp
+++ b/src/llvm.hpp
@@ -53,10 +53,13 @@ namespace lily {
     class llvm_context {
         private:
             std::map<std::string, llvm::Function*> supercombinators;
+            std::map<std::string, int> arities;
             llvm::Function* current_function;
         public:
-            llvm::Function* get_supercombinator(const std::string& name);
-            void add_supercombinator(const std::string& name);
+            void add_supercombinator(const std::string& name, int arity);
+            llvm::Function* get_supercombinator_function(const std::string& name);
+            int get_supercombinator_arity(const std::string& name);
+
             llvm::Function* get_current_function();
             void set_current_function(llvm::Function* f);
     };
diff --git a/src/main.cpp b/src/main.cpp
index cf95faa..e4efcf6 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -8,8 +8,8 @@ int main() {
     try {
 
         lily::program_ptr prog = lily::parse(
-                "defn magic x y = { let z = { x + y } in  { 326 + z + z } }\n"
-                "defn main = { magic 1 2 }"
+                "defn magic x = { let z = { x * 2 } in  { 326 + z + z } }\n"
+                "defn main = { magic 1 + magic 0 }"
                 );
         prog->gen_llvm();
     } catch(lily::error& e) {
diff --git a/src/parser.cpp b/src/parser.cpp
index 968af03..f60c254 100644
--- a/src/parser.cpp
+++ b/src/parser.cpp
@@ -342,12 +342,13 @@ namespace lily {
         compile(mgr, gcode);
 
         for(auto& pair : gcode) {
-            ctx.add_supercombinator(pair.first);
+            int arity = functions.count(pair.first) ? functions[pair.first].params.size() : 2;
+            ctx.add_supercombinator(pair.first, arity);
         }
 
         for(auto& pair : gcode) {
             std::vector<instruction*>& comb_gcode = pair.second;
-            llvm::Function* current_function = ctx.get_supercombinator(pair.first);
+            llvm::Function* current_function = ctx.get_supercombinator_function(pair.first);
             ctx.set_current_function(current_function);
             llvm::BasicBlock* new_block =
                 llvm::BasicBlock::Create(context, "entry", current_function);