From cf726fe4a80b8bb7954ea5fc240c087de47ed637 Mon Sep 17 00:00:00 2001
From: Shunsuke Shibayama <sbym1346@gmail.com>
Date: Fri, 29 Dec 2023 12:09:59 +0900
Subject: [PATCH] fix: array type bug

---
 crates/erg_compiler/lower.rs   | 9 ++++++++-
 tests/should_err/collection.er | 5 +++++
 tests/should_ok/collection.er  | 5 +++++
 tests/test.rs                  | 2 +-
 4 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/crates/erg_compiler/lower.rs b/crates/erg_compiler/lower.rs
index 1c37a040..cfb06f14 100644
--- a/crates/erg_compiler/lower.rs
+++ b/crates/erg_compiler/lower.rs
@@ -413,7 +413,14 @@ impl<A: ASTBuildable> GenericASTLowerer<A> {
                             {
                                 return Err(self.elem_err(&l, &r, elem));
                             } // else(OK): e.g. [1, "a": Str or Int]
-                        } else {
+                        }
+                        // OK: ?T(:> {"a"}) or ?U(:> {"b"}) or {"c", "d"} => {"a", "b", "c", "d"} <: Str
+                        else if self
+                            .module
+                            .context
+                            .coerce(union_.clone(), &())
+                            .map_or(true, |coerced| coerced.union_pair().is_some())
+                        {
                             return Err(self.elem_err(&l, &r, elem));
                         }
                     }
diff --git a/tests/should_err/collection.er b/tests/should_err/collection.er
index a658d16a..e1c1fc00 100644
--- a/tests/should_err/collection.er
+++ b/tests/should_err/collection.er
@@ -5,3 +5,8 @@ print! arr[2][0] # ERR
 dict = {"a": {"a": 1}, "b": {"b": 2}}
 print! dict["c"]["a"] # ERR
 print! dict["a"]["c"] # ERR
+
+ab = if True:
+    do "a"
+    do "b"
+_ = [1, 2, ab] # NG
diff --git a/tests/should_ok/collection.er b/tests/should_ok/collection.er
index adf3ba8c..f01a97bc 100644
--- a/tests/should_ok/collection.er
+++ b/tests/should_ok/collection.er
@@ -5,3 +5,8 @@ print! x
 dic = {"a": {"a": 1}, "b": {"b": 2}}
 y = dic["a"]["a"]
 print! y
+
+ab = if True: # x: {"a"} or {"b"}
+    do "a"
+    do "b"
+_ = ["c", "d", ab] # OK
diff --git a/tests/test.rs b/tests/test.rs
index 42956017..66740bde 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -482,7 +482,7 @@ fn exec_class_attr_err() -> Result<(), ()> {
 
 #[test]
 fn exec_collection_err() -> Result<(), ()> {
-    expect_failure("tests/should_err/collection.er", 0, 4)
+    expect_failure("tests/should_err/collection.er", 0, 5)
 }
 
 #[test]