Correctly support SelfType when searching for usages

6 files changed, 210 insertions, 76 deletions

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index ac23e385e..c9ef4b420 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2071,6 +2071,10 @@ impl Type {
         Some(adt.into())
     }
 
+    pub fn as_builtin(&self) -> Option<BuiltinType> {
+        self.ty.as_builtin().map(|inner| BuiltinType { inner })
+    }
+
     pub fn as_dyn_trait(&self) -> Option<Trait> {
         self.ty.dyn_trait().map(Into::into)
     }
diff --git a/crates/hir_ty/src/chalk_ext.rs b/crates/hir_ty/src/chalk_ext.rs
index 8c4542956..5232a7d80 100644
--- a/crates/hir_ty/src/chalk_ext.rs
+++ b/crates/hir_ty/src/chalk_ext.rs
@@ -1,8 +1,10 @@
 //! Various extensions traits for Chalk types.
 
-use chalk_ir::Mutability;
+use chalk_ir::{FloatTy, IntTy, Mutability, Scalar, UintTy};
 use hir_def::{
-    type_ref::Rawness, AssocContainerId, FunctionId, GenericDefId, HasModule, Lookup, TraitId,
+    builtin_type::{BuiltinFloat, BuiltinInt, BuiltinType, BuiltinUint},
+    type_ref::Rawness,
+    AssocContainerId, FunctionId, GenericDefId, HasModule, Lookup, TraitId,
 };
 
 use crate::{
@@ -18,6 +20,7 @@ pub trait TyExt {
     fn is_unknown(&self) -> bool;
 
     fn as_adt(&self) -> Option<(hir_def::AdtId, &Substitution)>;
+    fn as_builtin(&self) -> Option<BuiltinType>;
     fn as_tuple(&self) -> Option<&Substitution>;
     fn as_fn_def(&self, db: &dyn HirDatabase) -> Option<FunctionId>;
     fn as_reference(&self) -> Option<(&Ty, Lifetime, Mutability)>;
@@ -59,6 +62,35 @@ impl TyExt for Ty {
         }
     }
 
+    fn as_builtin(&self) -> Option<BuiltinType> {
+        match self.kind(&Interner) {
+            TyKind::Str => Some(BuiltinType::Str),
+            TyKind::Scalar(Scalar::Bool) => Some(BuiltinType::Bool),
+            TyKind::Scalar(Scalar::Char) => Some(BuiltinType::Char),
+            TyKind::Scalar(Scalar::Float(fty)) => Some(BuiltinType::Float(match fty {
+                FloatTy::F64 => BuiltinFloat::F64,
+                FloatTy::F32 => BuiltinFloat::F32,
+            })),
+            TyKind::Scalar(Scalar::Int(ity)) => Some(BuiltinType::Int(match ity {
+                IntTy::Isize => BuiltinInt::Isize,
+                IntTy::I8 => BuiltinInt::I8,
+                IntTy::I16 => BuiltinInt::I16,
+                IntTy::I32 => BuiltinInt::I32,
+                IntTy::I64 => BuiltinInt::I64,
+                IntTy::I128 => BuiltinInt::I128,
+            })),
+            TyKind::Scalar(Scalar::Uint(ity)) => Some(BuiltinType::Uint(match ity {
+                UintTy::Usize => BuiltinUint::Usize,
+                UintTy::U8 => BuiltinUint::U8,
+                UintTy::U16 => BuiltinUint::U16,
+                UintTy::U32 => BuiltinUint::U32,
+                UintTy::U64 => BuiltinUint::U64,
+                UintTy::U128 => BuiltinUint::U128,
+            })),
+            _ => None,
+        }
+    }
+
     fn as_tuple(&self) -> Option<&Substitution> {
         match self.kind(&Interner) {
             TyKind::Tuple(_, substs) => Some(substs),
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index 11ca7ec6b..1551cd2a8 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -65,7 +65,7 @@ pub(crate) fn find_all_refs(
             (find_def(&sema, &syntax, position)?, false)
         };
 
-    let mut usages = def.usages(sema).set_scope(search_scope).all();
+    let mut usages = def.usages(sema).set_scope(search_scope).include_self_refs().all();
     if is_literal_search {
         // filter for constructor-literals
         let refs = usages.references.values_mut();
@@ -1166,19 +1166,68 @@ fn foo<const FOO$0: usize>() -> usize {
     fn test_find_self_ty_in_trait_def() {
         check(
             r#"
-trait Foo {
+trait Foo where Self: {
     fn f() -> Self$0;
 }
 "#,
             expect![[r#"
                 Self TypeParam FileId(0) 6..9 6..9
 
-                FileId(0) 26..30
+                FileId(0) 16..20
+                FileId(0) 38..42
             "#]],
         );
+        //         check(
+        //             r#"
+        // trait Foo$0 where Self: {
+        //     fn f() -> Self;
+        // }
+        // "#,
+        //             expect![[r#"
+        //                 Foo Trait FileId(0) 0..45 6..9
+
+        //                 FileId(0) 16..20
+        //                 FileId(0) 38..42
+        //             "#]],
+        //         );
     }
 
     #[test]
+    fn test_self_ty() {
+        check(
+            r#"
+        struct $0Foo;
+
+        impl Foo where Self: {
+            fn f() -> Self;
+        }
+        "#,
+            expect![[r#"
+                        Foo Struct FileId(0) 0..11 7..10
+
+                        FileId(0) 18..21
+                        FileId(0) 28..32
+                        FileId(0) 50..54
+                    "#]],
+        );
+        check(
+            r#"
+struct Foo;
+
+impl Foo where Self: {
+    fn f() -> Self$0;
+}
+"#,
+            expect![[r#"
+                impl Impl FileId(0) 13..57 18..21
+
+                FileId(0) 18..21
+                FileId(0) 28..32
+                FileId(0) 50..54
+            "#]],
+        );
+    }
+    #[test]
     fn test_self_variant_with_payload() {
         check(
             r#"
diff --git a/crates/ide/src/references/rename.rs b/crates/ide/src/references/rename.rs
index 175e7a31d..2bf953305 100644
--- a/crates/ide/src/references/rename.rs
+++ b/crates/ide/src/references/rename.rs
@@ -1888,4 +1888,21 @@ impl Foo {
             "error: Cannot rename `Self`",
         );
     }
+
+    #[test]
+    fn test_rename_ignores_self_ty() {
+        check(
+            "Fo0",
+            r#"
+struct $0Foo;
+
+impl Foo where Self: {}
+"#,
+            r#"
+struct Fo0;
+
+impl Fo0 where Self: {}
+"#,
+        );
+    }
 }
diff --git a/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs b/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs
index b5b5ada5e..70949ca35 100644
--- a/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs
+++ b/crates/ide_assists/src/handlers/convert_tuple_struct_to_named_struct.rs
@@ -107,7 +107,7 @@ fn edit_struct_references(
     names: &[ast::Name],
 ) {
     let strukt_def = Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Struct(strukt)));
-    let usages = strukt_def.usages(&ctx.sema).include_self_kw_refs(true).all();
+    let usages = strukt_def.usages(&ctx.sema).include_self_refs().all();
 
     let edit_node = |edit: &mut AssistBuilder, node: SyntaxNode| -> Option<()> {
         match_ast! {
diff --git a/crates/ide_db/src/search.rs b/crates/ide_db/src/search.rs
index 8f899ea56..891b8cf0a 100644
--- a/crates/ide_db/src/search.rs
+++ b/crates/ide_db/src/search.rs
@@ -233,6 +233,13 @@ impl Definition {
             };
         }
 
+        if let Definition::SelfType(impl_) = self {
+            return match impl_.source(db).map(|src| src.value.syntax().text_range()) {
+                Some(range) => SearchScope::file_range(FileRange { file_id, range }),
+                None => SearchScope::single_file(file_id),
+            };
+        }
+
         if let Definition::GenericParam(hir::GenericParam::LifetimeParam(param)) = self {
             let range = match param.parent(db) {
                 hir::GenericDef::Function(it) => {
@@ -297,7 +304,7 @@ impl Definition {
     }
 
     pub fn usages<'a>(&'a self, sema: &'a Semantics<RootDatabase>) -> FindUsages<'a> {
-        FindUsages { def: self, sema, scope: None, include_self_kw_refs: false }
+        FindUsages { def: self, sema, scope: None, include_self_kw_refs: None }
     }
 }
 
@@ -305,12 +312,13 @@ pub struct FindUsages<'a> {
     def: &'a Definition,
     sema: &'a Semantics<'a, RootDatabase>,
     scope: Option<SearchScope>,
-    include_self_kw_refs: bool,
+    include_self_kw_refs: Option<hir::Type>,
 }
 
 impl<'a> FindUsages<'a> {
-    pub fn include_self_kw_refs(mut self, include: bool) -> FindUsages<'a> {
-        self.include_self_kw_refs = include;
+    /// Enable searching for `Self` when the definition is a type.
+    pub fn include_self_refs(mut self) -> FindUsages<'a> {
+        self.include_self_kw_refs = def_to_ty(self.sema.db, self.def);
         self
     }
 
@@ -354,13 +362,18 @@ impl<'a> FindUsages<'a> {
             }
         };
 
-        let name = match self.def.name(sema.db) {
-            Some(it) => it.to_string(),
+        let name = self.def.name(sema.db).or_else(|| {
+            self.include_self_kw_refs.as_ref().and_then(|ty| {
+                ty.as_adt()
+                    .map(|adt| adt.name(self.sema.db))
+                    .or_else(|| ty.as_builtin().map(|builtin| builtin.name()))
+            })
+        });
+        let name = match name {
+            Some(name) => name.to_string(),
             None => return,
         };
-
-        let pat = name.as_str();
-        let search_for_self = self.include_self_kw_refs;
+        let name = name.as_str();
 
         for (file_id, search_range) in search_scope {
             let text = sema.db.file_text(file_id);
@@ -369,51 +382,63 @@ impl<'a> FindUsages<'a> {
 
             let tree = Lazy::new(|| sema.parse(file_id).syntax().clone());
 
-            let mut handle_match = |idx: usize| -> bool {
+            for (idx, _) in text.match_indices(name) {
                 let offset: TextSize = idx.try_into().unwrap();
                 if !search_range.contains_inclusive(offset) {
-                    return false;
+                    continue;
                 }
 
                 if let Some(name) = sema.find_node_at_offset_with_descend(&tree, offset) {
-                    match name {
-                        ast::NameLike::NameRef(name_ref) => {
-                            if self.found_name_ref(&name_ref, sink) {
-                                return true;
-                            }
-                        }
-                        ast::NameLike::Name(name) => {
-                            if self.found_name(&name, sink) {
-                                return true;
-                            }
-                        }
-                        ast::NameLike::Lifetime(lifetime) => {
-                            if self.found_lifetime(&lifetime, sink) {
-                                return true;
-                            }
-                        }
+                    if match name {
+                        ast::NameLike::NameRef(name_ref) => self.found_name_ref(&name_ref, sink),
+                        ast::NameLike::Name(name) => self.found_name(&name, sink),
+                        ast::NameLike::Lifetime(lifetime) => self.found_lifetime(&lifetime, sink),
+                    } {
+                        return;
                     }
                 }
-
-                return false;
-            };
-
-            for (idx, _) in text.match_indices(pat) {
-                if handle_match(idx) {
-                    return;
-                }
             }
-
-            if search_for_self {
+            if let Some(self_ty) = &self.include_self_kw_refs {
                 for (idx, _) in text.match_indices("Self") {
-                    if handle_match(idx) {
-                        return;
+                    let offset: TextSize = idx.try_into().unwrap();
+                    if !search_range.contains_inclusive(offset) {
+                        continue;
+                    }
+
+                    if let Some(ast::NameLike::NameRef(name_ref)) =
+                        sema.find_node_at_offset_with_descend(&tree, offset)
+                    {
+                        if self.found_self_ty_name_ref(&self_ty, &name_ref, sink) {
+                            return;
+                        }
                     }
                 }
             }
         }
     }
 
+    fn found_self_ty_name_ref(
+        &self,
+        self_ty: &hir::Type,
+        name_ref: &ast::NameRef,
+        sink: &mut dyn FnMut(FileId, FileReference) -> bool,
+    ) -> bool {
+        match NameRefClass::classify(self.sema, &name_ref) {
+            Some(NameRefClass::Definition(Definition::SelfType(impl_)))
+                if impl_.self_ty(self.sema.db) == *self_ty =>
+            {
+                let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
+                let reference = FileReference {
+                    range,
+                    name: ast::NameLike::NameRef(name_ref.clone()),
+                    access: None,
+                };
+                sink(file_id, reference)
+            }
+            _ => false,
+        }
+    }
+
     fn found_lifetime(
         &self,
         lifetime: &ast::Lifetime,
@@ -429,7 +454,7 @@ impl<'a> FindUsages<'a> {
                 };
                 sink(file_id, reference)
             }
-            _ => false, // not a usage
+            _ => false,
         }
     }
 
@@ -448,42 +473,35 @@ impl<'a> FindUsages<'a> {
                 };
                 sink(file_id, reference)
             }
-            Some(NameRefClass::Definition(Definition::SelfType(impl_))) => {
-                let ty = impl_.self_ty(self.sema.db);
-
-                if let Some(adt) = ty.as_adt() {
-                    if &Definition::ModuleDef(ModuleDef::Adt(adt)) == self.def {
-                        let FileRange { file_id, range } =
-                            self.sema.original_range(name_ref.syntax());
-                        let reference = FileReference {
-                            range,
-                            name: ast::NameLike::NameRef(name_ref.clone()),
-                            access: None,
-                        };
-                        return sink(file_id, reference);
-                    }
+            Some(NameRefClass::Definition(def)) if self.include_self_kw_refs.is_some() => {
+                if self.include_self_kw_refs == def_to_ty(self.sema.db, &def) {
+                    let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
+                    let reference = FileReference {
+                        range,
+                        name: ast::NameLike::NameRef(name_ref.clone()),
+                        access: reference_access(&def, &name_ref),
+                    };
+                    sink(file_id, reference)
+                } else {
+                    false
                 }
-
-                false
             }
             Some(NameRefClass::FieldShorthand { local_ref: local, field_ref: field }) => {
                 let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
-                let reference = match self.def {
-                    Definition::Field(_) if &field == self.def => FileReference {
-                        range,
-                        name: ast::NameLike::NameRef(name_ref.clone()),
-                        access: reference_access(&field, &name_ref),
-                    },
-                    Definition::Local(l) if &local == l => FileReference {
-                        range,
-                        name: ast::NameLike::NameRef(name_ref.clone()),
-                        access: reference_access(&Definition::Local(local), &name_ref),
-                    },
-                    _ => return false, // not a usage
+                let access = match self.def {
+                    Definition::Field(_) if &field == self.def => {
+                        reference_access(&field, &name_ref)
+                    }
+                    Definition::Local(l) if &local == l => {
+                        reference_access(&Definition::Local(local), &name_ref)
+                    }
+                    _ => return false,
                 };
+                let reference =
+                    FileReference { range, name: ast::NameLike::NameRef(name_ref.clone()), access };
                 sink(file_id, reference)
             }
-            _ => false, // not a usage
+            _ => false,
         }
     }
 
@@ -513,11 +531,25 @@ impl<'a> FindUsages<'a> {
                     FileReference { range, name: ast::NameLike::Name(name.clone()), access: None };
                 sink(file_id, reference)
             }
-            _ => false, // not a usage
+            _ => false,
         }
     }
 }
 
+fn def_to_ty(db: &RootDatabase, def: &Definition) -> Option<hir::Type> {
+    match def {
+        Definition::ModuleDef(def) => match def {
+            ModuleDef::Adt(adt) => Some(adt.ty(db)),
+            ModuleDef::TypeAlias(it) => Some(it.ty(db)),
+            ModuleDef::BuiltinType(_it) => None, // FIXME somehow acquire some module to construct the builtin type
+            ModuleDef::Trait(_it) => None,       // FIXME turn trait into its self-type
+            _ => None,
+        },
+        Definition::SelfType(it) => Some(it.self_ty(db)),
+        _ => None,
+    }
+}
+
 fn reference_access(def: &Definition, name_ref: &ast::NameRef) -> Option<ReferenceAccess> {
     // Only Locals and Fields have accesses for now.
     if !matches!(def, Definition::Local(_) | Definition::Field(_)) {