From 27ede539c034619a9b30208dd23eb8b93b15ca8f Mon Sep 17 00:00:00 2001
From: Shunsuke Shibayama <sbym1346@gmail.com>
Date: Wed, 14 Feb 2024 13:44:45 +0900
Subject: [PATCH] fix: refinement type assert cast bug

---
 crates/erg_common/triple.rs            |  8 ++++++
 crates/erg_compiler/context/compare.rs | 37 +++++++++++++++++++++-----
 crates/erg_compiler/context/inquire.rs |  1 +
 crates/erg_compiler/tests/infer.er     |  8 ++++++
 crates/erg_compiler/tests/test.rs      |  3 +++
 5 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/crates/erg_common/triple.rs b/crates/erg_common/triple.rs
index 91547caf2..a743d43ab 100644
--- a/crates/erg_common/triple.rs
+++ b/crates/erg_common/triple.rs
@@ -141,6 +141,14 @@ impl<T> Triple<T, T> {
             Triple::Ok(a) | Triple::Err(a) => Some(a),
         }
     }
+
+    pub fn merge_or(self, default: T) -> T {
+        match self {
+            Triple::None => default,
+            Triple::Ok(ok) => ok,
+            Triple::Err(err) => err,
+        }
+    }
 }
 
 impl<T, E: std::error::Error> Triple<T, E> {
diff --git a/crates/erg_compiler/context/compare.rs b/crates/erg_compiler/context/compare.rs
index 4fb64adb5..652afedc4 100644
--- a/crates/erg_compiler/context/compare.rs
+++ b/crates/erg_compiler/context/compare.rs
@@ -1282,6 +1282,7 @@ impl Context {
     /// union(Array(Int, 2), Array(Str, 3)) == Array(Int, 2) or Array(Int, 3)
     /// union({ .a = Int }, { .a = Str }) == { .a = Int or Str }
     /// union({ .a = Int }, { .a = Int; .b = Int }) == { .a = Int }
+    /// union((A and B) or C) == (A or C) and (B or C)
     /// ```
     pub(crate) fn union(&self, lhs: &Type, rhs: &Type) -> Type {
         if lhs == rhs {
@@ -1345,6 +1346,16 @@ impl Context {
                 _ => self.simple_union(lhs, rhs),
             },
             (other, or @ Or(_, _)) | (or @ Or(_, _), other) => self.union_add(or, other),
+            // (A and B) or C ==> (A or C) and (B or C)
+            (and_t @ And(_, _), other) | (other, and_t @ And(_, _)) => {
+                let ands = and_t.ands();
+                let mut t = Type::Obj;
+                for branch in ands.iter() {
+                    let union = self.union(branch, other);
+                    t = and(t, union);
+                }
+                t
+            }
             (t, Type::Never) | (Type::Never, t) => t.clone(),
             // Array({1, 2}, 2), Array({3, 4}, 2) ==> Array({1, 2, 3, 4}, 2)
             (
@@ -1497,12 +1508,6 @@ impl Context {
                 self.intersection(&fv.crack(), other)
             }
             (Refinement(l), Refinement(r)) => Type::Refinement(self.intersection_refinement(l, r)),
-            (other, Refinement(refine)) | (Refinement(refine), other) => {
-                let other = other.clone().into_refinement();
-                let intersec = self.intersection_refinement(&other, refine);
-                self.try_squash_refinement(intersec)
-                    .unwrap_or_else(Type::Refinement)
-            }
             (Structural(l), Structural(r)) => self.intersection(l, r).structuralize(),
             (Guard(l), Guard(r)) => {
                 if l.namespace == r.namespace && l.target == r.target {
@@ -1527,6 +1532,26 @@ impl Context {
             (other, and @ And(_, _)) | (and @ And(_, _), other) => {
                 self.intersection_add(and, other)
             }
+            // (A or B) and C == (A and C) or (B and C)
+            (or_t @ Or(_, _), other) | (other, or_t @ Or(_, _)) => {
+                let ors = or_t.ors();
+                let mut t = Type::Never;
+                for branch in ors.iter() {
+                    let isec = self.intersection(branch, other);
+                    if branch.is_unbound_var() {
+                        t = or(t, isec);
+                    } else {
+                        t = self.union(&t, &isec);
+                    }
+                }
+                t
+            }
+            (other, Refinement(refine)) | (Refinement(refine), other) => {
+                let other = other.clone().into_refinement();
+                let intersec = self.intersection_refinement(&other, refine);
+                self.try_squash_refinement(intersec)
+                    .unwrap_or_else(Type::Refinement)
+            }
             // overloading
             (l, r) if l.is_subr() && r.is_subr() => and(lhs.clone(), rhs.clone()),
             _ => self.simple_intersection(lhs, rhs),
diff --git a/crates/erg_compiler/context/inquire.rs b/crates/erg_compiler/context/inquire.rs
index d85c3cb9d..300ae91a6 100644
--- a/crates/erg_compiler/context/inquire.rs
+++ b/crates/erg_compiler/context/inquire.rs
@@ -3657,6 +3657,7 @@ impl Context {
     /// ```erg
     /// recover_typarams(Int, Nat) == Nat
     /// recover_typarams(Array!(Int, _), Array(Nat, 2)) == Array!(Nat, 2)
+    /// recover_typarams(Str or NoneType, {"a", "b"}) == {"a", "b"}
     /// ```
     /// ```erg
     /// # REVIEW: should be?
diff --git a/crates/erg_compiler/tests/infer.er b/crates/erg_compiler/tests/infer.er
index 9b5c148f0..2243ed302 100644
--- a/crates/erg_compiler/tests/infer.er
+++ b/crates/erg_compiler/tests/infer.er
@@ -30,3 +30,11 @@ c_new x, y = C.new x, y
 C = Class Int
 C.
     new x, y = Self x + y
+
+val!() =
+    for! [{ "a": "b" }], (pkg as {Str: Str}) =>
+        x = pkg.get("a", "c")
+        assert x in {"b"}
+        val!::return x
+    "d"
+val = val!()
diff --git a/crates/erg_compiler/tests/test.rs b/crates/erg_compiler/tests/test.rs
index 80bbec857..89bcdd598 100644
--- a/crates/erg_compiler/tests/test.rs
+++ b/crates/erg_compiler/tests/test.rs
@@ -87,6 +87,9 @@ fn _test_infer_types() -> Result<(), ()> {
     let c_new_t = func2(add_r, r, c.clone()).quantify();
     module.context.assert_var_type("c_new", &c_new_t)?;
     module.context.assert_attr_type(&c, "new", &c_new_t)?;
+    module
+        .context
+        .assert_var_type("val", &v_enum(set! { "b".into(), "d".into() }))?;
     Ok(())
 }