fix Match case on type don't exaust on types #598

pyrefly/lib/binding/narrow.rs

@@ -320,6 +320,45 @@ impl NarrowOp {
             _ => *self = Self::Or(vec![self.clone(), other]),
         }
     }
+
+    pub fn rebase_subject(&self, subject: &NarrowingSubject) -> Self {
+        match self {
+            Self::Atomic(prop, op) => {
+                Self::Atomic(rebase_facet_subject(subject, prop.as_ref()), op.clone())
+            }
+            Self::And(ops) => Self::And(ops.map(|op| op.rebase_subject(subject))),
+            Self::Or(ops) => Self::Or(ops.map(|op| op.rebase_subject(subject))),
+        }
+    }
+}
+
+fn rebase_facet_subject(
+    subject: &NarrowingSubject,
+    prop: Option<&FacetSubject>,
+) -> Option<FacetSubject> {
+    match (subject, prop) {
+        (NarrowingSubject::Name(_), None) => None,
+        (NarrowingSubject::Name(_), Some(facet)) => Some(facet.clone()),
+        (NarrowingSubject::Facets(_, base_facet), None) => Some(base_facet.clone()),
+        (NarrowingSubject::Facets(_, base_facet), Some(facet)) => {
+            Some(merge_facet_subjects(base_facet, facet))
+        }
+    }
+}
+
+fn merge_facet_subjects(base: &FacetSubject, extra: &FacetSubject) -> FacetSubject {
+    let mut facets: Vec<_> = base.chain.facets().iter().cloned().collect();
+    facets.extend(extra.chain.facets().iter().cloned());
+    let chain = Vec1::try_from_vec(facets)
+        .expect("FacetChain is always non-empty when merging facet subjects");
+    let origin = match (base.origin, extra.origin) {
+        (FacetOrigin::GetMethod, _) | (_, FacetOrigin::GetMethod) => FacetOrigin::GetMethod,
+        _ => FacetOrigin::Direct,
+    };
+    FacetSubject {
+        chain: FacetChain::new(chain),
+        origin,
+    }
 }
 
 #[derive(Clone, Debug, Default)]
@@ -330,6 +369,19 @@ impl NarrowOps {
         Self(SmallMap::new())
     }
 
+    pub fn and_for_subject(&mut self, subject: &NarrowingSubject, op: NarrowOp, range: TextRange) {
+        let name = subject.name().clone();
+        match self.0.entry(name) {
+            Entry::Occupied(mut entry) => {
+                entry.get_mut().0.and(op);
+                entry.get_mut().1 = range;
+            }
+            Entry::Vacant(entry) => {
+                entry.insert((op, range));
+            }
+        }
+    }
+
     pub fn negate(&self) -> Self {
         Self(
             self.0

pyrefly/lib/binding/pattern.rs

@@ -79,13 +79,27 @@ impl<'a> BindingsBuilder<'a> {
             Pattern::MatchAs(p) => {
                 // If there's no name for this pattern, refine the variable being matched
                 // If there is a new name, refine that instead
+                let original_subject = match_subject.clone();
+                let alias_name = p.name.as_ref().map(|name| name.id.clone());
                 let mut subject = match_subject;
                 if let Some(name) = &p.name {
                     self.bind_definition(name, Binding::Forward(subject_idx), FlowStyle::Other);
                     subject = Some(NarrowingSubject::Name(name.id.clone()));
                 };
                 if let Some(pattern) = p.pattern {
-                    self.bind_pattern(subject, *pattern, subject_idx)
+                    let mut narrow_ops = self.bind_pattern(subject, *pattern, subject_idx);
+                    if let (Some(alias_name), Some(original_subject)) =
+                        (&alias_name, &original_subject)
+                        && alias_name != original_subject.name()
+                        && let Some((alias_op, range)) = narrow_ops.0.get(alias_name).cloned()
+                    {
+                        narrow_ops.and_for_subject(
+                            original_subject,
+                            alias_op.rebase_subject(original_subject),
+                            range,
+                        );
+                    }
+                    narrow_ops
                 } else {
                     NarrowOps::new()
                 }

pyrefly/lib/test/flow_branching.rs

@@ -401,7 +401,7 @@ def test(x: int):
         case 1:
             assert_type(x, Literal[1])
         case 2 as q:
-            assert_type(x, int)
+            assert_type(x, Literal[2])
             assert_type(q, Literal[2])
         case q:
             assert_type(x, int)
@@ -518,7 +518,7 @@ def fun(x: A | B | C) -> None:
             assert_type(x, B)
     match x:
         case B(3, "B") as y:
-            assert_type(x, A | B | C)
+            assert_type(x, B)
             assert_type(y, B)
     match x:
         case A(1, "a") | B(2, "b"):
@@ -635,7 +635,7 @@ def test(x: Foo | Bar) -> None:
             assert_type(x, Bar)
             assert_type(x.x, str)  # we want to narrow this to Literal["bar"]
         case Bar(a) as b:
-            assert_type(x, Foo | Bar)
+            assert_type(x, Bar)
             assert_type(b, Bar)
             assert_type(a, str)
             assert_type(b, Bar)

pyrefly/lib/test/pattern_match.rs

@@ -99,6 +99,24 @@ def f0(x: int | str):
 "#,
 );
 
+testcase!(
+    test_match_alias_narrows_subject,
+    r#"
+from typing import assert_never, assert_type
+
+def my_method(str_or_int: str | int) -> str:
+    match str_or_int:
+        case str() as str_data:
+            assert_type(str_or_int, str)
+            return str_data
+        case int() as int_data:
+            assert_type(str_or_int, int)
+            return str(int_data)
+        case _:
+            assert_never(str_or_int)
+"#,
+);
+
 testcase!(
     test_class_match_with_args_not_exhaustive,
     r#"