From 0b544a2515ad725c996e46870cd941f2e5d787d9 Mon Sep 17 00:00:00 2001
From: Danila Fedorin <danila.fedorin@gmail.com>
Date: Sun, 9 Jun 2019 19:57:58 -0700
Subject: [PATCH] Fix function return type checking.

---
 src/main.cpp   | 1 +
 src/parser.cpp | 7 ++++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/main.cpp b/src/main.cpp
index f2a6908..a5d2385 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -8,6 +8,7 @@ int main() {
                 "data Bool = { True, False }\n"
                 "data Color = { Red, Black }\n"
                 "data IntList = { Nil, Cons(Int, IntList) }\n"
+                "defn other x y = { 3 }\n"
                 "defn add x y = { x + add x x }");
     } catch(lily::error& e) {
         std::cout << e.message << std::endl;
diff --git a/src/parser.cpp b/src/parser.cpp
index 0116f04..f0b999d 100644
--- a/src/parser.cpp
+++ b/src/parser.cpp
@@ -248,6 +248,7 @@ namespace lily {
         // Each function has an environment, which will be used
         // as base for type checking.
         std::map<std::string, std::shared_ptr<type_env>> function_envs;
+        std::map<std::string, type*> function_output_types;
         auto base_env = std::make_shared<type_env>();
 
         // First step is to collect all function types.
@@ -261,6 +262,7 @@ namespace lily {
             // Create the return type parameter
             type* return_type = type_mgr.create_type<type_parameter>();
             type* current_type = return_type;
+            function_output_types[pair.first] = return_type;
 
             // Create type parameters for every variable
             // We also want to place them in a local environment
@@ -279,7 +281,10 @@ namespace lily {
 
         // Now that we have collected the functions, check their bodies.
         for(auto& pair : functions) {
-            pair.second.body->check(type_mgr, function_envs[pair.first]);
+            type* body_type =
+                pair.second.body->check(type_mgr, function_envs[pair.first]);
+            if(!function_output_types[pair.first]->unify_with(body_type))
+                throw error("unable to unify function type");
         }
     }
 }