From 41563fd6128e94e1fe1ae4006e4e183dff1a645c Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Wed, 7 Apr 2021 11:31:50 +0200
Subject: [PATCH] Infer variants through type aliased enums

---
 crates/hir_ty/src/infer.rs        | 70 ++++++++++++++++---------------
 crates/hir_ty/src/tests/simple.rs | 33 +++++++++++++++
 2 files changed, 70 insertions(+), 33 deletions(-)

diff --git a/crates/hir_ty/src/infer.rs b/crates/hir_ty/src/infer.rs
index c63878e7a61..796f77ab0de 100644
--- a/crates/hir_ty/src/infer.rs
+++ b/crates/hir_ty/src/infer.rs
@@ -484,36 +484,13 @@ impl<'a> InferenceContext<'a> {
                 let generics = crate::utils::generics(self.db.upcast(), impl_id.into());
                 let substs = generics.type_params_subst(self.db);
                 let ty = self.db.impl_self_ty(impl_id).substitute(&Interner, &substs);
-                match unresolved {
-                    None => {
-                        let variant = ty_variant(&ty);
-                        (ty, variant)
-                    }
-                    Some(1) => {
-                        let segment = path.mod_path().segments().last().unwrap();
-                        // this could be an enum variant or associated type
-                        if let Some((AdtId::EnumId(enum_id), _)) = ty.as_adt() {
-                            let enum_data = self.db.enum_data(enum_id);
-                            if let Some(local_id) = enum_data.variant(segment) {
-                                let variant = EnumVariantId { parent: enum_id, local_id };
-                                return (ty, Some(variant.into()));
-                            }
-                        }
-                        // FIXME potentially resolve assoc type
-                        (self.err_ty(), None)
-                    }
-                    Some(_) => {
-                        // FIXME diagnostic
-                        (self.err_ty(), None)
-                    }
-                }
+                self.resolve_variant_on_alias(ty, unresolved, path)
             }
             TypeNs::TypeAliasId(it) => {
                 let ty = TyBuilder::def_ty(self.db, it.into())
                     .fill(std::iter::repeat_with(|| self.table.new_type_var()))
                     .build();
-                let variant = ty_variant(&ty);
-                forbid_unresolved_segments((ty, variant), unresolved)
+                self.resolve_variant_on_alias(ty, unresolved, path)
             }
             TypeNs::AdtSelfType(_) => {
                 // FIXME this could happen in array size expressions, once we're checking them
@@ -540,16 +517,43 @@ impl<'a> InferenceContext<'a> {
                 (TyKind::Error.intern(&Interner), None)
             }
         }
+    }
 
-        fn ty_variant(ty: &Ty) -> Option<VariantId> {
-            ty.as_adt().and_then(|(adt_id, _)| match adt_id {
-                AdtId::StructId(s) => Some(VariantId::StructId(s)),
-                AdtId::UnionId(u) => Some(VariantId::UnionId(u)),
-                AdtId::EnumId(_) => {
-                    // FIXME Error E0071, expected struct, variant or union type, found enum `Foo`
-                    None
+    fn resolve_variant_on_alias(
+        &mut self,
+        ty: Ty,
+        unresolved: Option<usize>,
+        path: &Path,
+    ) -> (Ty, Option<VariantId>) {
+        match unresolved {
+            None => {
+                let variant = ty.as_adt().and_then(|(adt_id, _)| match adt_id {
+                    AdtId::StructId(s) => Some(VariantId::StructId(s)),
+                    AdtId::UnionId(u) => Some(VariantId::UnionId(u)),
+                    AdtId::EnumId(_) => {
+                        // FIXME Error E0071, expected struct, variant or union type, found enum `Foo`
+                        None
+                    }
+                });
+                (ty, variant)
+            }
+            Some(1) => {
+                let segment = path.mod_path().segments().last().unwrap();
+                // this could be an enum variant or associated type
+                if let Some((AdtId::EnumId(enum_id), _)) = ty.as_adt() {
+                    let enum_data = self.db.enum_data(enum_id);
+                    if let Some(local_id) = enum_data.variant(segment) {
+                        let variant = EnumVariantId { parent: enum_id, local_id };
+                        return (ty, Some(variant.into()));
+                    }
                 }
-            })
+                // FIXME potentially resolve assoc type
+                (self.err_ty(), None)
+            }
+            Some(_) => {
+                // FIXME diagnostic
+                (self.err_ty(), None)
+            }
         }
     }
 
diff --git a/crates/hir_ty/src/tests/simple.rs b/crates/hir_ty/src/tests/simple.rs
index 361cd63026f..84c5c05fd30 100644
--- a/crates/hir_ty/src/tests/simple.rs
+++ b/crates/hir_ty/src/tests/simple.rs
@@ -2564,3 +2564,36 @@ fn f() {
     "#,
     )
 }
+
+#[test]
+fn infer_type_alias_variant() {
+    check_infer(
+        r#"
+type Qux = Foo;
+enum Foo {
+    Bar(i32),
+    Baz { baz: f32 }
+}
+
+fn f() {
+    match Foo::Bar(3) {
+        Qux::Bar(bar) => (),
+        Qux::Baz { baz } => (),
+    }
+}
+    "#,
+        expect![[r#"
+            72..166 '{     ...   } }': ()
+            78..164 'match ...     }': ()
+            84..92 'Foo::Bar': Bar(i32) -> Foo
+            84..95 'Foo::Bar(3)': Foo
+            93..94 '3': i32
+            106..119 'Qux::Bar(bar)': Foo
+            115..118 'bar': i32
+            123..125 '()': ()
+            135..151 'Qux::B... baz }': Foo
+            146..149 'baz': f32
+            155..157 '()': ()
+        "#]],
+    )
+}