feat: set type variable

mtshiba · mtshiba · commit bee2e70e1f83 · 2024-08-19T21:29:38.000+09:00
diff --git a/crates/erg_compiler/context/compare.rs b/crates/erg_compiler/context/compare.rs
@@ -418,9 +418,26 @@ impl Context {
             // ?T(<: Int) :> ?U(:> Int)
             // ?T(<: Nat) !:> ?U(:> Int) (if the upper bound of LHS is smaller than the lower bound of RHS, LHS cannot not be a supertype)
             // ?T(<: Nat) :> ?U(<: Int) (?U can be smaller than ?T)
+            // ?T(:> ?U) :> ?U
+            // ?U :> ?T(<: ?U)
+            // ?T(: {Int, Str}) :> ?U(<: Int)
             (FreeVar(lfv), FreeVar(rfv)) => match (lfv.get_subsup(), rfv.get_subsup()) {
                 (Some((_, l_sup)), Some((r_sub, _))) => self.supertype_of(&l_sup, &r_sub),
+                (Some((l_sub, _)), None) if &l_sub == rhs => true,
+                (None, Some((_, r_sup))) if lhs == &r_sup => true,
                 _ => {
+                    let lfvt = lfv.get_type();
+                    // lfv: T: {Int, Str}, rhs: Int
+                    if let Some(tys) = lfvt.as_ref().and_then(|t| t.refinement_values()) {
+                        for tp in tys {
+                            let Ok(ty) = self.convert_tp_into_type(tp.clone()) else {
+                                continue;
+                            };
+                            if self.supertype_of(&ty, rhs) {
+                                return true;
+                            }
+                        }
+                    }
                     if lfv.is_linked() {
                         self.supertype_of(lfv.unsafe_crack(), rhs)
                     } else if rfv.is_linked() {
@@ -504,6 +521,17 @@ impl Context {
                 if let Some((_sub, sup)) = lfv.get_subsup() {
                     lfv.do_avoiding_recursion_with(rhs, || self.supertype_of(&sup, rhs))
                 } else if let Some(lfvt) = lfv.get_type() {
+                    // lfv: T: {Int, Str}, rhs: Int
+                    if let Some(tys) = lfvt.refinement_values() {
+                        for tp in tys {
+                            let Ok(ty) = self.convert_tp_into_type(tp.clone()) else {
+                                continue;
+                            };
+                            if self.supertype_of(&ty, rhs) {
+                                return true;
+                            }
+                        }
+                    }
                     // e.g. lfv: ?L(: Int) is unreachable
                     // but
                     // ?L(: List(Type, 3)) :> List(Int, 3)
diff --git a/crates/erg_compiler/context/generalize.rs b/crates/erg_compiler/context/generalize.rs
@@ -857,6 +857,24 @@ impl<'c, 'q, 'l, L: Locational> Dereferencer<'c, 'q, 'l, L> {
                     Ok(Type::FreeVar(fv))
                 }
             }
+            FreeVar(fv) if fv.get_type().is_some() => {
+                let ty = fv.get_type().unwrap();
+                if self.level <= fv.level().unwrap() {
+                    // T: {Int, Str} => Int or Str
+                    if let Some(tys) = ty.refinement_values() {
+                        let mut union = Never;
+                        for tp in tys {
+                            if let Ok(ty) = self.ctx.convert_tp_into_type(tp.clone()) {
+                                union = self.ctx.union(&union, &ty);
+                            }
+                        }
+                        return Ok(union);
+                    }
+                    Ok(Type::FreeVar(fv))
+                } else {
+                    Ok(Type::FreeVar(fv))
+                }
+            }
             FreeVar(fv) if fv.is_unbound() => {
                 if self.level == 0 {
                     match &*fv.crack_constraint() {
diff --git a/crates/erg_compiler/context/instantiate_spec.rs b/crates/erg_compiler/context/instantiate_spec.rs
@@ -183,7 +183,7 @@ impl Context {
                     }
                     _ => unreachable!(),
                 };
-                if constr.get_sub_sup().is_none() {
+                if constr.get_type().is_some_and(|t| !t.is_meta_type()) {
                     let tp = TyParam::named_free_var(lhs.inspect().clone(), self.level, constr);
                     tv_cache.push_or_init_typaram(lhs, &tp, self)?;
                 } else {
diff --git a/crates/erg_compiler/ty/mod.rs b/crates/erg_compiler/ty/mod.rs
@@ -1619,9 +1619,17 @@ impl CanbeFree for Type {
     }
 
     fn destructive_update_constraint(&self, new_constraint: Constraint, in_instantiation: bool) {
-        if let Some(fv) = self.as_free() {
-            fv.update_constraint(new_constraint, in_instantiation);
+        let Some(fv) = self.as_free() else {
+            return;
+        };
+        // self: T
+        // new_constraint: (:> T, <: U) => <: U
+        if new_constraint.get_sub_sup().is_some_and(|(sub, sup)| {
+            sub.contains_tvar_in_constraint(fv) || sup.contains_tvar_in_constraint(fv)
+        }) {
+            return;
         }
+        fv.update_constraint(new_constraint, in_instantiation);
     }
 }
 
@@ -2731,6 +2739,28 @@ impl Type {
         }
     }
 
+    pub fn contains_tvar_in_constraint(&self, target: &FreeTyVar) -> bool {
+        match self {
+            Self::FreeVar(fv) if fv.is_linked() => fv.crack().contains_tvar_in_constraint(target),
+            Self::FreeVar(fv) if fv.constraint_is_typeof() => {
+                ref_addr_eq!(fv.forced_as_ref(), target.forced_as_ref())
+            }
+            Self::FreeVar(fv) => {
+                ref_addr_eq!(fv.forced_as_ref(), target.forced_as_ref())
+                    || fv
+                        .get_subsup()
+                        .map(|(sub, sup)| {
+                            fv.do_avoiding_recursion(|| {
+                                sub.contains_tvar_in_constraint(target)
+                                    || sup.contains_tvar_in_constraint(target)
+                            })
+                        })
+                        .unwrap_or(false)
+            }
+            _ => false,
+        }
+    }
+
     pub fn contains_type(&self, target: &Type) -> bool {
         if self == target {
             // This operation can also be performed for recursive types
@@ -3095,6 +3125,14 @@ impl Type {
         }
     }
 
+    pub fn get_meta_type(&self) -> Option<Type> {
+        match self {
+            Self::FreeVar(fv) if fv.is_linked() => fv.crack().get_meta_type(),
+            Self::FreeVar(fv) if fv.is_unbound() => fv.get_type(),
+            _ => None,
+        }
+    }
+
     pub const fn is_free_var(&self) -> bool {
         matches!(self, Self::FreeVar(_))
     }
@@ -3523,6 +3561,15 @@ impl Type {
         !self.has_unbound_var()
     }
 
+    pub fn is_meta_type(&self) -> bool {
+        match self {
+            Self::FreeVar(fv) if fv.is_linked() => fv.crack().is_meta_type(),
+            Self::Refinement(refine) => refine.t.is_meta_type(),
+            Self::ClassType | Self::TraitType | Self::Type => true,
+            _ => false,
+        }
+    }
+
     pub fn typarams_len(&self) -> Option<usize> {
         match self {
             Self::FreeVar(fv) if fv.is_linked() => fv.crack().typarams_len(),
diff --git a/crates/erg_parser/ast.rs b/crates/erg_parser/ast.rs
@@ -552,6 +552,22 @@ impl Args {
             })
         }
     }
+
+    #[to_owned(cloned)]
+    pub fn get_nth(&self, nth: usize) -> Option<&Expr> {
+        self.pos_args.get(nth).map(|a| &a.expr)
+    }
+
+    #[to_owned(cloned)]
+    pub fn get_with_key(&self, key: &str) -> Option<&Expr> {
+        self.kw_args.iter().find_map(|a| {
+            if &a.keyword.content[..] == key {
+                Some(&a.expr)
+            } else {
+                None
+            }
+        })
+    }
 }
 
 #[pyclass]
diff --git a/doc/EN/syntax/type/15_quantified.md b/doc/EN/syntax/type/15_quantified.md
@@ -86,6 +86,15 @@ f|X, Y, Z| x: X, y: Y, z: Z =
     x + y + z + x
 ```
 
+```python
+# T is a subclass of Int or Str
+f|T: {Int, Str}|(x: T, _: T): T = x
+
+_: Int = f 1, 2
+_: Str = f "a", "b"
+_ = f None, None # ERROR
+```
+
 Unlike many languages with generics, all declared type variables must be used either in the temporary argument list (the `x: X, y: Y, z: Z` part) or in the arguments of other type variables.
 This is a requirement from Erg's language design that all type variables are inferrable from real arguments.
 So information that cannot be inferred, such as the return type, is passed from real arguments; Erg allows types to be passed from real arguments.
diff --git a/doc/JA/syntax/type/15_quantified.md b/doc/JA/syntax/type/15_quantified.md
@@ -89,6 +89,15 @@ f|X, Y, Z| x: X, y: Y, z: Z  =
     x + y + z + x
 ```
 
+```python
+# TはIntまたはStrのサブクラス
+f|T: {Int, Str}|(x: T, _: T): T = x
+
+_: Int = f 1, 2
+_: Str = f "a", "b"
+_ = f None, None # ERROR
+```
+
 ジェネリクスを持つ多くの言語と違い、宣言した型変数はすべて、仮引数リスト内(`x: X, y: Y, z: Z`の部分)か他の型変数の引数内かで使用されていなければなりません。
 これは、型変数はすべて実引数から推論可能であるというErgの言語設計からの要求です。
 なので、戻り値の型など推論ができない情報は、実引数から渡します。Ergは型を実引数から渡すことができるのです。
diff --git a/tests/should_err/set_type.er b/tests/should_err/set_type.er
@@ -0,0 +1,4 @@
+f|T: {Int, Str}|(x: T, _: T): T = x
+
+_ = f 1.1, 1.2 # ERR
+_ = f "a", None # ERR
diff --git a/tests/should_ok/set_type.er b/tests/should_ok/set_type.er
@@ -0,0 +1,4 @@
+f|T: {Int, Str}|(x: T, _: T): T = x
+
+_ = f 1, 2
+_ = f "a", "b"
diff --git a/tests/test.rs b/tests/test.rs
@@ -397,6 +397,11 @@ fn exec_self_type() -> Result<(), ()> {
     expect_success("tests/should_ok/self_type.er", 0)
 }
 
+#[test]
+fn exec_set_type() -> Result<(), ()> {
+    expect_success("tests/should_ok/set_type.er", 0)
+}
+
 #[test]
 fn exec_slice() -> Result<(), ()> {
     expect_success("tests/should_ok/slice.er", 0)
@@ -638,6 +643,11 @@ fn exec_set() -> Result<(), ()> {
     expect_failure("examples/set.er", 3, 1)
 }
 
+#[test]
+fn exec_set_type_err() -> Result<(), ()> {
+    expect_failure("tests/should_err/set_type.er", 0, 3)
+}
+
 #[test]
 fn exec_side_effect() -> Result<(), ()> {
     expect_failure("examples/side_effect.er", 5, 4)

Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,7 @@ impl Context {`
`183`	`183`	`}`
`184`	`184`	`_ => unreachable!(),`
`185`	`185`	`};`
`186`		`- if constr.get_sub_sup().is_none() {`
	`186`	`+ if constr.get_type().is_some_and(\|t\| !t.is_meta_type()) {`
`187`	`187`	`let tp = TyParam::named_free_var(lhs.inspect().clone(), self.level, constr);`
`188`	`188`	`tv_cache.push_or_init_typaram(lhs, &tp, self)?;`
`189`	`189`	`} else {`