Support inferring `Self` type in enum definitions

Signed-off-by: ice1000 <[email protected]>

7 files changed, 54 insertions, 9 deletions

diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs
index e3a7e8e3c..cd501d1f6 100644
--- a/crates/ra_hir/src/code_model.rs
+++ b/crates/ra_hir/src/code_model.rs
@@ -24,7 +24,7 @@ use crate::{
         U8, USIZE,
     },
     nameres::{CrateModuleId, ImportId, ModuleScope, Namespace},
-    resolve::{Resolver, TypeNs},
+    resolve::{Resolver, Scope, TypeNs},
     traits::TraitData,
     ty::{
         primitive::{FloatBitness, FloatTy, IntBitness, IntTy, Signedness},
@@ -465,7 +465,7 @@ impl Enum {
         // ...and add generic params, if present
         let p = self.generic_params(db);
         let r = if !p.params.is_empty() { r.push_generic_params_scope(p) } else { r };
-        r
+        r.push_scope(Scope::AdtScope(From::from(self)))
     }
 }
 
diff --git a/crates/ra_hir/src/resolve.rs b/crates/ra_hir/src/resolve.rs
index 39f8e1d8a..c46f7e157 100644
--- a/crates/ra_hir/src/resolve.rs
+++ b/crates/ra_hir/src/resolve.rs
@@ -43,8 +43,10 @@ pub(crate) enum Scope {
     ModuleScope(ModuleItemMap),
     /// Brings the generic parameters of an item into scope
     GenericParams(Arc<GenericParams>),
-    /// Brings `Self` into scope
+    /// Brings `Self` in `impl` block into scope
     ImplBlockScope(ImplBlock),
+    /// Brings `Self` in enum definition into scope
+    AdtScope(Adt),
     /// Local bindings
     ExprScope(ExprScope),
 }
@@ -54,6 +56,7 @@ pub enum TypeNs {
     SelfType(ImplBlock),
     GenericParam(u32),
     Adt(Adt),
+    AdtSelfType(Adt),
     EnumVariant(EnumVariant),
     TypeAlias(TypeAlias),
     BuiltinType(BuiltinType),
@@ -151,6 +154,12 @@ impl Resolver {
                         return Some((TypeNs::SelfType(*impl_), idx));
                     }
                 }
+                Scope::AdtScope(adt) => {
+                    if first_name == &SELF_TYPE {
+                        let idx = if path.segments.len() == 1 { None } else { Some(1) };
+                        return Some((TypeNs::AdtSelfType(*adt), idx));
+                    }
+                }
                 Scope::ModuleScope(m) => {
                     let (module_def, idx) = m.crate_def_map.resolve_path(db, m.module_id, path);
                     let res = match module_def.take_types()? {
@@ -200,7 +209,10 @@ impl Resolver {
         let skip_to_mod = path.kind != PathKind::Plain && !path.is_self();
         for scope in self.scopes.iter().rev() {
             match scope {
-                Scope::ExprScope(_) | Scope::GenericParams(_) | Scope::ImplBlockScope(_)
+                Scope::AdtScope(_)
+                | Scope::ExprScope(_)
+                | Scope::GenericParams(_)
+                | Scope::ImplBlockScope(_)
                     if skip_to_mod =>
                 {
                     continue
@@ -233,7 +245,13 @@ impl Resolver {
                         return Some(ResolveValueResult::Partial(ty, 1));
                     }
                 }
-                Scope::ImplBlockScope(_) => continue,
+                Scope::AdtScope(adt) if n_segments > 1 => {
+                    if first_name == &SELF_TYPE {
+                        let ty = TypeNs::AdtSelfType(*adt);
+                        return Some(ResolveValueResult::Partial(ty, 1));
+                    }
+                }
+                Scope::ImplBlockScope(_) | Scope::AdtScope(_) => continue,
 
                 Scope::ModuleScope(m) => {
                     let (module_def, idx) = m.crate_def_map.resolve_path(db, m.module_id, path);
@@ -389,7 +407,8 @@ pub enum ScopeDef {
     ModuleDef(ModuleDef),
     MacroDef(MacroDef),
     GenericParam(u32),
-    SelfType(ImplBlock),
+    ImplSelfType(ImplBlock),
+    AdtSelfType(Adt),
     LocalBinding(PatId),
     Unknown,
 }
@@ -437,7 +456,10 @@ impl Scope {
                 }
             }
             Scope::ImplBlockScope(i) => {
-                f(SELF_TYPE, ScopeDef::SelfType(*i));
+                f(SELF_TYPE, ScopeDef::ImplSelfType(*i));
+            }
+            Scope::AdtScope(i) => {
+                f(SELF_TYPE, ScopeDef::AdtSelfType(*i));
             }
             Scope::ExprScope(e) => {
                 e.expr_scopes.entries(e.scope_id).iter().for_each(|e| {
diff --git a/crates/ra_hir/src/source_binder.rs b/crates/ra_hir/src/source_binder.rs
index 088335e66..a907d6a9f 100644
--- a/crates/ra_hir/src/source_binder.rs
+++ b/crates/ra_hir/src/source_binder.rs
@@ -216,7 +216,7 @@ impl SourceAnalyzer {
         let types = self.resolver.resolve_path_in_type_ns_fully(db, &path).map(|ty| match ty {
             TypeNs::SelfType(it) => PathResolution::SelfType(it),
             TypeNs::GenericParam(it) => PathResolution::GenericParam(it),
-            TypeNs::Adt(it) => PathResolution::Def(it.into()),
+            TypeNs::AdtSelfType(it) | TypeNs::Adt(it) => PathResolution::Def(it.into()),
             TypeNs::EnumVariant(it) => PathResolution::Def(it.into()),
             TypeNs::TypeAlias(it) => PathResolution::Def(it.into()),
             TypeNs::BuiltinType(it) => PathResolution::Def(it.into()),
diff --git a/crates/ra_hir/src/ty/infer.rs b/crates/ra_hir/src/ty/infer.rs
index ca9aefc42..9e5d42509 100644
--- a/crates/ra_hir/src/ty/infer.rs
+++ b/crates/ra_hir/src/ty/infer.rs
@@ -559,6 +559,9 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
             match resolver.resolve_path_in_type_ns_fully(self.db, &path) {
                 Some(TypeNs::Adt(Adt::Struct(it))) => it.into(),
                 Some(TypeNs::Adt(Adt::Union(it))) => it.into(),
+                Some(TypeNs::AdtSelfType(Adt::Struct(it))) => it.into(),
+                Some(TypeNs::AdtSelfType(Adt::Union(it))) => it.into(),
+                Some(TypeNs::AdtSelfType(Adt::Enum(it))) => it.into(),
                 Some(TypeNs::EnumVariant(it)) => it.into(),
                 Some(TypeNs::TypeAlias(it)) => it.into(),
 
diff --git a/crates/ra_hir/src/ty/lower.rs b/crates/ra_hir/src/ty/lower.rs
index 4b67c82e7..a604c02e2 100644
--- a/crates/ra_hir/src/ty/lower.rs
+++ b/crates/ra_hir/src/ty/lower.rs
@@ -175,6 +175,7 @@ impl Ty {
                 Ty::Param { idx, name }
             }
             TypeNs::SelfType(impl_block) => impl_block.target_ty(db),
+            TypeNs::AdtSelfType(adt) => adt.ty(db),
 
             TypeNs::Adt(it) => Ty::from_hir_path_inner(db, resolver, resolved_segment, it.into()),
             TypeNs::BuiltinType(it) => {
diff --git a/crates/ra_hir/src/ty/tests.rs b/crates/ra_hir/src/ty/tests.rs
index 25dad81eb..03b30adcd 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -135,6 +135,25 @@ mod boxed {
 }
 
 #[test]
+fn infer_adt_self() {
+    let (db, pos) = MockDatabase::with_position(
+        r#"
+//- /main.rs
+enum Nat { Succ(Self), Demo(Nat), Zero }
+
+fn test() {
+    let foo: Nat = Nat::Zero;
+    if let Nat::Succ(x) = foo {
+        x<|>
+    }
+}
+
+"#,
+    );
+    assert_eq!("Nat", type_at_pos(&db, pos));
+}
+
+#[test]
 fn infer_try() {
     let (mut db, pos) = MockDatabase::with_position(
         r#"
diff --git a/crates/ra_ide_api/src/completion/presentation.rs b/crates/ra_ide_api/src/completion/presentation.rs
index 48028a2f9..eb480a775 100644
--- a/crates/ra_ide_api/src/completion/presentation.rs
+++ b/crates/ra_ide_api/src/completion/presentation.rs
@@ -72,7 +72,7 @@ impl Completions {
             }
             ScopeDef::GenericParam(..) => (CompletionItemKind::TypeParam, None),
             ScopeDef::LocalBinding(..) => (CompletionItemKind::Binding, None),
-            ScopeDef::SelfType(..) => (
+            ScopeDef::AdtSelfType(..) | ScopeDef::ImplSelfType(..) => (
                 CompletionItemKind::TypeParam, // (does this need its own kind?)
                 None,
             ),