Lifetime reference search

3 files changed, 107 insertions, 18 deletions

diff --git a/crates/ide_db/src/defs.rs b/crates/ide_db/src/defs.rs
index d4a774261..f2d1e4c39 100644
--- a/crates/ide_db/src/defs.rs
+++ b/crates/ide_db/src/defs.rs
@@ -6,12 +6,12 @@
 // FIXME: this badly needs rename/rewrite (matklad, 2020-02-06).
 
 use hir::{
-    db::HirDatabase, Crate, Field, HasVisibility, ImplDef, Local, MacroDef, Module, ModuleDef,
-    Name, PathResolution, Semantics, TypeParam, Visibility,
+    db::HirDatabase, Crate, Field, HasVisibility, ImplDef, LifetimeParam, Local, MacroDef, Module,
+    ModuleDef, Name, PathResolution, Semantics, TypeParam, Visibility,
 };
 use syntax::{
     ast::{self, AstNode},
-    match_ast, SyntaxNode,
+    match_ast, SyntaxKind, SyntaxNode,
 };
 
 use crate::RootDatabase;
@@ -25,6 +25,8 @@ pub enum Definition {
     SelfType(ImplDef),
     Local(Local),
     TypeParam(TypeParam),
+    LifetimeParam(LifetimeParam),
+    // FIXME: Label
 }
 
 impl Definition {
@@ -36,6 +38,7 @@ impl Definition {
             Definition::SelfType(it) => Some(it.module(db)),
             Definition::Local(it) => Some(it.module(db)),
             Definition::TypeParam(it) => Some(it.module(db)),
+            Definition::LifetimeParam(it) => Some(it.module(db)),
         }
     }
 
@@ -47,6 +50,7 @@ impl Definition {
             Definition::SelfType(_) => None,
             Definition::Local(_) => None,
             Definition::TypeParam(_) => None,
+            Definition::LifetimeParam(_) => None,
         }
     }
 
@@ -72,6 +76,7 @@ impl Definition {
             Definition::SelfType(_) => return None,
             Definition::Local(it) => it.name(db)?,
             Definition::TypeParam(it) => it.name(db),
+            Definition::LifetimeParam(it) => it.name(db),
         };
         Some(name)
     }
@@ -229,6 +234,25 @@ impl NameClass {
             }
         }
     }
+
+    pub fn classify_lifetime(
+        sema: &Semantics<RootDatabase>,
+        lifetime: &ast::Lifetime,
+    ) -> Option<NameClass> {
+        let _p = profile::span("classify_lifetime").detail(|| lifetime.to_string());
+        let parent = lifetime.syntax().parent()?;
+
+        match_ast! {
+            match parent {
+                ast::LifetimeParam(it) => {
+                    let def = sema.to_def(&it)?;
+                    Some(NameClass::Definition(Definition::LifetimeParam(def)))
+                },
+                ast::Label(_it) => None,
+                _ => None,
+            }
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -338,6 +362,35 @@ impl NameRefClass {
         let resolved = sema.resolve_extern_crate(&extern_crate)?;
         Some(NameRefClass::ExternCrate(resolved))
     }
+
+    pub fn classify_lifetime(
+        sema: &Semantics<RootDatabase>,
+        lifetime: &ast::Lifetime,
+    ) -> Option<NameRefClass> {
+        let _p = profile::span("classify_lifetime_ref").detail(|| lifetime.to_string());
+        let parent = lifetime.syntax().parent()?;
+        match parent.kind() {
+            SyntaxKind::LIFETIME_ARG
+            | SyntaxKind::SELF_PARAM
+            | SyntaxKind::TYPE_BOUND
+            | SyntaxKind::WHERE_PRED
+            | SyntaxKind::REF_TYPE => sema
+                .resolve_lifetime_param(lifetime)
+                .map(Definition::LifetimeParam)
+                .map(NameRefClass::Definition),
+            // lifetime bounds, as in the 'b in 'a: 'b aren't wrapped in TypeBound nodes so we gotta check
+            // if our lifetime is in a LifetimeParam without being the constrained lifetime
+            _ if ast::LifetimeParam::cast(parent).and_then(|param| param.lifetime()).as_ref()
+                != Some(lifetime) =>
+            {
+                sema.resolve_lifetime_param(lifetime)
+                    .map(Definition::LifetimeParam)
+                    .map(NameRefClass::Definition)
+            }
+            SyntaxKind::BREAK_EXPR | SyntaxKind::CONTINUE_EXPR => None,
+            _ => None,
+        }
+    }
 }
 
 impl From<PathResolution> for Definition {
diff --git a/crates/ide_db/src/search.rs b/crates/ide_db/src/search.rs
index 3936c7390..5b3997bcf 100644
--- a/crates/ide_db/src/search.rs
+++ b/crates/ide_db/src/search.rs
@@ -33,6 +33,7 @@ pub enum ReferenceKind {
     RecordFieldExprOrPat,
     SelfKw,
     EnumLiteral,
+    Lifetime,
     Other,
 }
 
@@ -129,6 +130,25 @@ impl Definition {
             return SearchScope::new(res);
         }
 
+        if let Definition::LifetimeParam(param) = self {
+            let range = match param.parent(db) {
+                hir::GenericDef::Function(it) => it.source(db).value.syntax().text_range(),
+                hir::GenericDef::Adt(it) => match it {
+                    hir::Adt::Struct(it) => it.source(db).value.syntax().text_range(),
+                    hir::Adt::Union(it) => it.source(db).value.syntax().text_range(),
+                    hir::Adt::Enum(it) => it.source(db).value.syntax().text_range(),
+                },
+                hir::GenericDef::Trait(it) => it.source(db).value.syntax().text_range(),
+                hir::GenericDef::TypeAlias(it) => it.source(db).value.syntax().text_range(),
+                hir::GenericDef::ImplDef(it) => it.source(db).value.syntax().text_range(),
+                hir::GenericDef::EnumVariant(it) => it.source(db).value.syntax().text_range(),
+                hir::GenericDef::Const(it) => it.source(db).value.syntax().text_range(),
+            };
+            let mut res = FxHashMap::default();
+            res.insert(file_id, Some(range));
+            return SearchScope::new(res);
+        }
+
         let vis = self.visibility(db);
 
         if let Some(Visibility::Module(module)) = vis.and_then(|it| it.into()) {
@@ -255,25 +275,42 @@ impl<'a> FindUsages<'a> {
                     continue;
                 }
 
-                match sema.find_node_at_offset_with_descend(&tree, offset) {
-                    Some(name_ref) => {
-                        if self.found_name_ref(&name_ref, sink) {
-                            return;
-                        }
+                if let Some(name_ref) = sema.find_node_at_offset_with_descend(&tree, offset) {
+                    if self.found_name_ref(&name_ref, sink) {
+                        return;
+                    }
+                } else if let Some(name) = sema.find_node_at_offset_with_descend(&tree, offset) {
+                    if self.found_name(&name, sink) {
+                        return;
+                    }
+                } else if let Some(lifetime) = sema.find_node_at_offset_with_descend(&tree, offset)
+                {
+                    if self.found_lifetime(&lifetime, sink) {
+                        return;
                     }
-                    None => match sema.find_node_at_offset_with_descend(&tree, offset) {
-                        Some(name) => {
-                            if self.found_name(&name, sink) {
-                                return;
-                            }
-                        }
-                        None => {}
-                    },
                 }
             }
         }
     }
 
+    fn found_lifetime(
+        &self,
+        lifetime: &ast::Lifetime,
+        sink: &mut dyn FnMut(Reference) -> bool,
+    ) -> bool {
+        match NameRefClass::classify_lifetime(self.sema, lifetime) {
+            Some(NameRefClass::Definition(def)) if &def == self.def => {
+                let reference = Reference {
+                    file_range: self.sema.original_range(lifetime.syntax()),
+                    kind: ReferenceKind::Lifetime,
+                    access: None,
+                };
+                sink(reference)
+            }
+            _ => false, // not a usage
+        }
+    }
+
     fn found_name_ref(
         &self,
         name_ref: &ast::NameRef,
diff --git a/crates/ide_db/src/symbol_index.rs b/crates/ide_db/src/symbol_index.rs
index 121063aea..ca455fa03 100644
--- a/crates/ide_db/src/symbol_index.rs
+++ b/crates/ide_db/src/symbol_index.rs
@@ -209,8 +209,7 @@ pub fn crate_symbols(db: &RootDatabase, krate: CrateId, query: Query) -> Vec<Fil
     query.search(&buf)
 }
 
-pub fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
-    let name = name_ref.text();
+pub fn index_resolve(db: &RootDatabase, name: &SmolStr) -> Vec<FileSymbol> {
     let mut query = Query::new(name.to_string());
     query.exact();
     query.limit(4);