From 8ca733934faa936c8651f4f1373d3ebc3b68b573 Mon Sep 17 00:00:00 2001
From: programsnail <purpurebacterium@yandex.ru>
Date: Sun, 21 Apr 2024 18:31:44 +0300
Subject: [PATCH] type checker arrow type fix and type check test

---
 include/types.hpp  |  2 +-
 src/main.cpp       | 54 ++++++++++++++++++++++++++++++++++++++++------
 src/type_check.cpp |  9 +++++++-
 3 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/include/types.hpp b/include/types.hpp
index a6807e3..e5ade3d 100644
--- a/include/types.hpp
+++ b/include/types.hpp
@@ -231,7 +231,7 @@ struct Storage {
     }
   }
 
-private:
+// private: // TODO: temporary, to beautify type checker output
   size_t first_unused_generic_id = 0;
 
   vector<Type> types;
diff --git a/src/main.cpp b/src/main.cpp
index 79ee770..b236dc9 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -19,11 +19,19 @@ auto make_program_2(bool uniq) {
   using namespace nodes;
   return make_expr<Let>(
       Arg("f"),
-      lambda2(Arg("x", uniq ? types::Mode(types::Mode::Uniq::UNIQUE) : types::Mode()), Arg("y"),
+      lambda2(Arg("x", uniq ? types::Mode(types::Mode::Uniq::UNIQUE)
+                            : types::Mode()),
+              Arg("y"),
               operator_call("+", make_expr<Var>("x"), make_expr<Var>("y"))),
       make_expr<Var>("f"));
 }
 
+auto make_program_3() {
+  using namespace nodes;
+  return make_expr<Let>(Arg("f"), lambda1(Arg("x"), make_expr<Var>("x")),
+                        make_expr<Var>("f"));
+}
+
 void add_builtin_functions_types(type_check::State &state, bool uniq) {
   auto sum_type = state.type_storage.add(types::make_operator(
       state.type_storage.get_int_type(
@@ -47,7 +55,7 @@ void print_error(const std::string &general_message,
             << std::endl;
 }
 
-void run_example(const auto& make_program, bool arg_uniq, bool sum_uniq) {
+void run_example(const auto &make_program, bool arg_uniq, bool sum_uniq) {
   const auto program = make_program(arg_uniq);
 
   std::cout << "\x1b[1;34mPROGRAM:\x1b[0m \x1b[1;90m" << *program
@@ -83,13 +91,47 @@ void run_example(const auto& make_program, bool arg_uniq, bool sum_uniq) {
   }
 
   std::cout << "\x1b[1;92mPROGRAM IS CORRECT\x1b[0m\n";
-  return;
+}
+
+void run_example_2() {
+  const auto program = make_program_3();
+
+  std::cout << "\n\x1b[1;34m--- TYPE CHECK ---\x1b[0m\n";
+  std::cout << "\x1b[1;34mPROGRAM:\x1b[0m \x1b[1;90m" << *program
+            << "\x1b[0m\n\n";
+
+  try {
+    type_check::State state;
+
+    std::cout << "expression type is " << type_check::check_expr(program, state).get().type.index()
+              << "\n";
+
+    for (auto &type : state.type_storage.types) {
+      std::cout << type.type.index();
+      if (auto *arrow_type = get_if<types::ArrowType>(&type.type);
+          arrow_type != nullptr) {
+        std::cout << "[-";
+        for (const auto &arg : arrow_type->types) {
+          std::cout << arg.get().type.index() << "-";
+        }
+        std::cout << "]";
+      }
+      std::cout << ' ';
+    }
+  } catch (utils::Error error) {
+    print_error("\x1b[1;31mTYPE CHECK ERROR:\x1b[0m", error);
+    return;
+  }
+
+  std::cout << "\n\n\x1b[1;34m--- END ---\x1b[0m\n";
 }
 
 int main() {
   for (size_t n = 0; n < 8; ++n) {
-    std::cout << "\n\x1b[1;34m--- START TEST ---\x1b[0m\n";
-    run_example(n / 4 == 0 ? &make_program_1 : &make_program_2, n % 2 == 1, (n / 2) % 2 ==1);
-    std::cout << "\n\x1b[1;34m--- END TEST ---\x1b[0m\n";
+    std::cout << "\n\x1b[1;34m--- TEST ---\x1b[0m\n";
+    run_example(n / 4 == 0 ? &make_program_1 : &make_program_2, n % 2 == 1,
+                (n / 2) % 2 == 1);
+    std::cout << "\n\x1b[1;34m--- END ---\x1b[0m\n";
   }
+  run_example_2();
 }
diff --git a/src/type_check.cpp b/src/type_check.cpp
index 938e99f..81e0d92 100644
--- a/src/type_check.cpp
+++ b/src/type_check.cpp
@@ -36,13 +36,20 @@ types::TypeID check_let(nodes::Let &expr, State &state) {
 types::TypeID check_lambda(nodes::Lambda &expr, State &state) {
   Context context(state.manager);
 
+  ArrowType lambda_arrow_type;
+
+    lambda_arrow_type.types.reserve(expr.args.size() + 1);
   for (auto &arg : expr.args) {
     types::TypeID new_type = state.type_storage.introduce_new_generic(arg.name, arg.mode_hint);
     arg.type = new_type;
+    lambda_arrow_type.types.push_back(new_type);
     state.manager.add_var(arg.name, new_type);
   }
 
-  types::TypeID lambda_type = check_expr(expr.expr, state);
+  types::TypeID ret_type = check_expr(expr.expr, state);
+  lambda_arrow_type.types.push_back(ret_type);
+
+  types::TypeID lambda_type = state.type_storage.add(lambda_arrow_type);
   return (expr.type = lambda_type).value();
 }