From 0af0241d347cbe19a9a1adf99d43e5d02a402a35 Mon Sep 17 00:00:00 2001
From: Shunsuke Shibayama <sbym1346@gmail.com>
Date: Tue, 27 Dec 2022 20:15:21 +0900
Subject: [PATCH] Add `Context::resolve_ctx_vars`

---
 compiler/erg_compiler/context/mod.rs   |  1 +
 compiler/erg_compiler/context/tyvar.rs | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/compiler/erg_compiler/context/mod.rs b/compiler/erg_compiler/context/mod.rs
index 102523e4..0d56af43 100644
--- a/compiler/erg_compiler/context/mod.rs
+++ b/compiler/erg_compiler/context/mod.rs
@@ -957,6 +957,7 @@ impl Context {
         }
     }
 
+    /// unlike `pop`, `outer` must be `None`.
     pub fn pop_mod(&mut self) -> Option<Context> {
         if self.outer.is_some() {
             log!(err "not in the top-level context");
diff --git a/compiler/erg_compiler/context/tyvar.rs b/compiler/erg_compiler/context/tyvar.rs
index 2805333f..3f659b4d 100644
--- a/compiler/erg_compiler/context/tyvar.rs
+++ b/compiler/erg_compiler/context/tyvar.rs
@@ -696,6 +696,7 @@ impl Context {
                 errs.extend(es);
             }
         }
+        self.resolve_ctx_vars();
         if errs.is_empty() {
             Ok(hir)
         } else {
@@ -703,6 +704,29 @@ impl Context {
         }
     }
 
+    fn resolve_ctx_vars(&mut self) {
+        let mut locals = mem::take(&mut self.locals);
+        let mut params = mem::take(&mut self.params);
+        let mut methods_list = mem::take(&mut self.methods_list);
+        for (name, vi) in locals.iter_mut() {
+            if let Ok(t) = self.deref_tyvar(mem::take(&mut vi.t), Variance::Covariant, name.loc()) {
+                vi.t = t;
+            }
+        }
+        for (name, vi) in params.iter_mut() {
+            let loc = name.as_ref().map(|n| n.loc()).unwrap_or_default();
+            if let Ok(t) = self.deref_tyvar(mem::take(&mut vi.t), Variance::Covariant, loc) {
+                vi.t = t;
+            }
+        }
+        for (_, methods) in methods_list.iter_mut() {
+            methods.resolve_ctx_vars();
+        }
+        self.locals = locals;
+        self.params = params;
+        self.methods_list = methods_list;
+    }
+
     fn resolve_expr_t(&self, expr: &mut hir::Expr) -> TyCheckResult<()> {
         match expr {
             hir::Expr::Lit(_) => Ok(()),