From feeee2039b754c3490bf22598bb503cb8e734461 Mon Sep 17 00:00:00 2001
From: Danila Fedorin <danila.fedorin@gmail.com>
Date: Sat, 30 May 2020 16:40:27 -0700
Subject: [PATCH] Check for free variables in the environment before
 generalizing.

---
 12/type.cpp     |  9 +++++++++
 12/type.hpp     |  1 +
 12/type_env.cpp | 25 ++++++++++++++++++++++---
 12/type_env.hpp |  4 ++++
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/12/type.cpp b/12/type.cpp
index 05e85c6..237c9fb 100644
--- a/12/type.cpp
+++ b/12/type.cpp
@@ -185,3 +185,12 @@ void type_mgr::find_free(const type_ptr& t, std::set<std::string>& into) const {
         for(auto& arg : app->arguments) find_free(arg, into);
     }
 }
+
+void type_mgr::find_free(const type_scheme_ptr& t, std::set<std::string>& into) const {
+    std::set<std::string> monotype_free;
+    find_free(t->monotype, monotype_free);
+    for(auto& not_free : t->forall) {
+        monotype_free.erase(not_free);
+    }
+    into.insert(monotype_free.begin(), monotype_free.end());
+}
diff --git a/12/type.hpp b/12/type.hpp
index 4583f45..abf2b55 100644
--- a/12/type.hpp
+++ b/12/type.hpp
@@ -92,4 +92,5 @@ struct type_mgr {
     type_ptr resolve(type_ptr t, type_var*& var) const;
     void bind(const std::string& s, type_ptr t);
     void find_free(const type_ptr& t, std::set<std::string>& into) const;
+    void find_free(const type_scheme_ptr& t, std::set<std::string>& into) const;
 };
diff --git a/12/type_env.cpp b/12/type_env.cpp
index 3e654be..26a5a84 100644
--- a/12/type_env.cpp
+++ b/12/type_env.cpp
@@ -1,6 +1,22 @@
 #include "type_env.hpp"
 #include "type.hpp"
 
+void type_env::find_free(const type_mgr& mgr, std::set<std::string>& into) const {
+    if(parent != nullptr) parent->find_free(mgr, into);
+    for(auto& binding : names) {
+        mgr.find_free(binding.second, into);
+    }
+}
+
+void type_env::find_free_except(const type_mgr& mgr, const std::string& avoid,
+        std::set<std::string>& into) const {
+    if(parent != nullptr) parent->find_free(mgr, into);
+    for(auto& binding : names) {
+        if(binding.first == avoid) continue;
+        mgr.find_free(binding.second, into);
+    }
+}
+
 type_scheme_ptr type_env::lookup(const std::string& name) const {
     auto it = names.find(name);
     if(it != names.end()) return it->second;
@@ -33,9 +49,12 @@ void type_env::generalize(const std::string& name, type_mgr& mgr) {
     if(names_it == names.end()) throw 0;
     if(names_it->second->forall.size() > 0) throw 0;
 
-    std::set<std::string> free_variables;
-    mgr.find_free(names_it->second->monotype, free_variables);
-    for(auto& free : free_variables) {
+    std::set<std::string> free_in_type;
+    std::set<std::string> free_in_env;
+    mgr.find_free(names_it->second->monotype, free_in_type);
+    find_free_except(mgr, name, free_in_env);
+    for(auto& free : free_in_type) {
+        if(free_in_env.find(free) != free_in_env.end()) continue;
         names_it->second->forall.push_back(free);
     }
 }
diff --git a/12/type_env.hpp b/12/type_env.hpp
index ca91ef7..9b529d2 100644
--- a/12/type_env.hpp
+++ b/12/type_env.hpp
@@ -1,6 +1,7 @@
 #pragma once
 #include <map>
 #include <string>
+#include <set>
 #include "type.hpp"
 
 struct type_env;
@@ -14,6 +15,9 @@ struct type_env {
     type_env(type_env_ptr p) : parent(std::move(p)) {}
     type_env() : type_env(nullptr) {}
 
+    void find_free(const type_mgr& mgr, std::set<std::string>& into) const;
+    void find_free_except(const type_mgr& mgr, const std::string& avoid,
+            std::set<std::string>& into) const;
     type_scheme_ptr lookup(const std::string& name) const;
     type_ptr lookup_type(const std::string& name) const;
     void bind(const std::string& name, type_ptr t);