6 files changed, 127 insertions, 48 deletions

diff --git a/crates/ra_ide_api/src/lib.rs b/crates/ra_ide_api/src/lib.rs
index 1746b58ae..57a490fa7 100644
--- a/crates/ra_ide_api/src/lib.rs
+++ b/crates/ra_ide_api/src/lib.rs
@@ -56,6 +56,7 @@ pub use crate::{
     completion::{CompletionItem, CompletionItemKind, InsertTextFormat},
     runnables::{Runnable, RunnableKind},
     navigation_target::NavigationTarget,
+    references::ReferenceSearchResult,
 };
 pub use ra_ide_api_light::{
     Fold, FoldKind, HighlightedRange, Severity, StructureNode, LocalEdit,
@@ -319,7 +320,10 @@ impl Analysis {
     }
 
     /// Finds all usages of the reference at point.
-    pub fn find_all_refs(&self, position: FilePosition) -> Cancelable<Vec<(FileId, TextRange)>> {
+    pub fn find_all_refs(
+        &self,
+        position: FilePosition,
+    ) -> Cancelable<Option<ReferenceSearchResult>> {
         self.with_db(|db| references::find_all_refs(db, position))
     }
 
diff --git a/crates/ra_ide_api/src/navigation_target.rs b/crates/ra_ide_api/src/navigation_target.rs
index 004921863..fd001179a 100644
--- a/crates/ra_ide_api/src/navigation_target.rs
+++ b/crates/ra_ide_api/src/navigation_target.rs
@@ -23,6 +23,12 @@ pub struct NavigationTarget {
 }
 
 impl NavigationTarget {
+    /// When `focus_range` is specified, returns it. otherwise
+    /// returns `full_range`
+    pub fn range(&self) -> TextRange {
+        self.focus_range.unwrap_or(self.full_range)
+    }
+
     pub fn name(&self) -> &SmolStr {
         &self.name
     }
@@ -43,14 +49,18 @@ impl NavigationTarget {
         self.full_range
     }
 
-    /// A "most interesting" range withing the `range_full`.
+    /// A "most interesting" range withing the `full_range`.
     ///
-    /// Typically, `range` is the whole syntax node, including doc comments, and
-    /// `focus_range` is the range of the identifier.
+    /// Typically, `full_range` is the whole syntax node,
+    /// including doc comments, and `focus_range` is the range of the identifier.
     pub fn focus_range(&self) -> Option<TextRange> {
         self.focus_range
     }
 
+    pub(crate) fn from_bind_pat(file_id: FileId, pat: &ast::BindPat) -> NavigationTarget {
+        NavigationTarget::from_named(file_id, pat)
+    }
+
     pub(crate) fn from_symbol(symbol: FileSymbol) -> NavigationTarget {
         NavigationTarget {
             file_id: symbol.file_id,
diff --git a/crates/ra_ide_api/src/references.rs b/crates/ra_ide_api/src/references.rs
index e7ebf9f6e..b7784e577 100644
--- a/crates/ra_ide_api/src/references.rs
+++ b/crates/ra_ide_api/src/references.rs
@@ -1,42 +1,77 @@
 use relative_path::{RelativePath, RelativePathBuf};
 use hir::{ModuleSource, source_binder};
-use ra_db::{FileId, SourceDatabase};
+use ra_db::{SourceDatabase};
 use ra_syntax::{
-    AstNode, SyntaxNode, TextRange, SourceFile,
-    ast::{self, NameOwner},
+    AstNode, SyntaxNode, SourceFile,
+    ast,
     algo::find_node_at_offset,
 };
 
 use crate::{
     db::RootDatabase,
     FilePosition,
+    FileRange,
+    FileId,
+    NavigationTarget,
     FileSystemEdit,
     SourceChange,
     SourceFileEdit,
+    TextRange,
 };
 
-pub(crate) fn find_all_refs(db: &RootDatabase, position: FilePosition) -> Vec<(FileId, TextRange)> {
+#[derive(Debug, Clone)]
+pub struct ReferenceSearchResult {
+    declaration: NavigationTarget,
+    references: Vec<FileRange>,
+}
+
+impl ReferenceSearchResult {
+    pub fn declaration(&self) -> &NavigationTarget {
+        &self.declaration
+    }
+
+    pub fn references(&self) -> &[FileRange] {
+        &self.references
+    }
+
+    /// Total number of references
+    /// At least 1 since all valid references should
+    /// Have a declaration
+    pub fn len(&self) -> usize {
+        self.references.len() + 1
+    }
+}
+
+// allow turning ReferenceSearchResult into an iterator
+// over FileRanges
+impl IntoIterator for ReferenceSearchResult {
+    type Item = FileRange;
+    type IntoIter = std::vec::IntoIter<FileRange>;
+
+    fn into_iter(mut self) -> Self::IntoIter {
+        let mut v = Vec::with_capacity(self.len());
+        v.push(FileRange { file_id: self.declaration.file_id(), range: self.declaration.range() });
+        v.append(&mut self.references);
+        v.into_iter()
+    }
+}
+
+pub(crate) fn find_all_refs(
+    db: &RootDatabase,
+    position: FilePosition,
+) -> Option<ReferenceSearchResult> {
     let file = db.parse(position.file_id);
-    // Find the binding associated with the offset
-    let (binding, descr) = match find_binding(db, &file, position) {
-        None => return Vec::new(),
-        Some(it) => it,
-    };
+    let (binding, descr) = find_binding(db, &file, position)?;
+    let declaration = NavigationTarget::from_bind_pat(position.file_id, binding);
 
-    let mut ret = binding
-        .name()
+    let references = descr
+        .scopes(db)
+        .find_all_refs(binding)
         .into_iter()
-        .map(|name| (position.file_id, name.syntax().range()))
+        .map(move |ref_desc| FileRange { file_id: position.file_id, range: ref_desc.range })
         .collect::<Vec<_>>();
-    ret.extend(
-        descr
-            .scopes(db)
-            .find_all_refs(binding)
-            .into_iter()
-            .map(|ref_desc| (position.file_id, ref_desc.range)),
-    );
 
-    return ret;
+    return Some(ReferenceSearchResult { declaration, references });
 
     fn find_binding<'a>(
         db: &RootDatabase,
@@ -88,6 +123,21 @@ fn find_name_and_module_at_offset(
     None
 }
 
+fn source_edit_from_fileid_range(
+    file_id: FileId,
+    range: TextRange,
+    new_name: &str,
+) -> SourceFileEdit {
+    SourceFileEdit {
+        file_id,
+        edit: {
+            let mut builder = ra_text_edit::TextEditBuilder::default();
+            builder.replace(range, new_name.into());
+            builder.finish()
+        },
+    }
+}
+
 fn rename_mod(
     db: &RootDatabase,
     ast_name: &ast::Name,
@@ -150,17 +200,13 @@ fn rename_reference(
     position: FilePosition,
     new_name: &str,
 ) -> Option<SourceChange> {
-    let edit = find_all_refs(db, position)
-        .iter()
-        .map(|(file_id, text_range)| SourceFileEdit {
-            file_id: *file_id,
-            edit: {
-                let mut builder = ra_text_edit::TextEditBuilder::default();
-                builder.replace(*text_range, new_name.into());
-                builder.finish()
-            },
-        })
+    let refs = find_all_refs(db, position)?;
+
+    let edit = refs
+        .into_iter()
+        .map(|range| source_edit_from_fileid_range(range.file_id, range.range, new_name))
         .collect::<Vec<_>>();
+
     if edit.is_empty() {
         return None;
     }
diff --git a/crates/ra_ide_api/tests/test/main.rs b/crates/ra_ide_api/tests/test/main.rs
index 0526f7584..a83fbe07b 100644
--- a/crates/ra_ide_api/tests/test/main.rs
+++ b/crates/ra_ide_api/tests/test/main.rs
@@ -1,7 +1,8 @@
 use insta::assert_debug_snapshot_matches;
 use ra_ide_api::{
     mock_analysis::{single_file, single_file_with_position, MockAnalysis},
-    AnalysisChange, CrateGraph, Edition::Edition2018, FileId, Query, NavigationTarget
+    AnalysisChange, CrateGraph, Edition::Edition2018, Query, NavigationTarget,
+    ReferenceSearchResult,
 };
 use ra_syntax::{TextRange, SmolStr};
 
@@ -44,9 +45,9 @@ fn test_resolve_crate_root() {
     assert_eq!(host.analysis().crate_for(mod_file).unwrap(), vec![crate_id]);
 }
 
-fn get_all_refs(text: &str) -> Vec<(FileId, TextRange)> {
+fn get_all_refs(text: &str) -> ReferenceSearchResult {
     let (analysis, position) = single_file_with_position(text);
-    analysis.find_all_refs(position).unwrap()
+    analysis.find_all_refs(position).unwrap().unwrap()
 }
 
 fn get_symbols_matching(text: &str, query: &str) -> Vec<NavigationTarget> {
diff --git a/crates/ra_lsp_server/src/conv.rs b/crates/ra_lsp_server/src/conv.rs
index 20077a48a..c3192a1e5 100644
--- a/crates/ra_lsp_server/src/conv.rs
+++ b/crates/ra_lsp_server/src/conv.rs
@@ -333,7 +333,7 @@ impl TryConvWith for &NavigationTarget {
     type Output = Location;
     fn try_conv_with(self, world: &ServerWorld) -> Result<Location> {
         let line_index = world.analysis().file_line_index(self.file_id());
-        let range = self.focus_range().unwrap_or(self.full_range());
+        let range = self.range();
         to_location(self.file_id(), range, &world, &line_index)
     }
 }
diff --git a/crates/ra_lsp_server/src/main_loop/handlers.rs b/crates/ra_lsp_server/src/main_loop/handlers.rs
index 09d896c40..9208ee473 100644
--- a/crates/ra_lsp_server/src/main_loop/handlers.rs
+++ b/crates/ra_lsp_server/src/main_loop/handlers.rs
@@ -456,14 +456,16 @@ pub fn handle_prepare_rename(
 
     // We support renaming references like handle_rename does.
     // In the future we may want to reject the renaming of things like keywords here too.
-    let refs = world.analysis().find_all_refs(position)?;
-    let r = match refs.first() {
-        Some(r) => r,
+    let refs = match world.analysis().find_all_refs(position)? {
         None => return Ok(None),
+        Some(refs) => refs,
     };
+
+    // Refs should always have a declaration
+    let r = refs.declaration();
     let file_id = params.text_document.try_conv_with(&world)?;
     let line_index = world.analysis().file_line_index(file_id);
-    let loc = to_location(r.0, r.1, &world, &line_index)?;
+    let loc = to_location(r.file_id(), r.range(), &world, &line_index)?;
 
     Ok(Some(PrepareRenameResponse::Range(loc.range)))
 }
@@ -501,11 +503,24 @@ pub fn handle_references(
     let line_index = world.analysis().file_line_index(file_id);
     let offset = params.position.conv_with(&line_index);
 
-    let refs = world.analysis().find_all_refs(FilePosition { file_id, offset })?;
+    let refs = match world.analysis().find_all_refs(FilePosition { file_id, offset })? {
+        None => return Ok(None),
+        Some(refs) => refs,
+    };
 
-    Ok(Some(
-        refs.into_iter().filter_map(|r| to_location(r.0, r.1, &world, &line_index).ok()).collect(),
-    ))
+    let locations = if params.context.include_declaration {
+        refs.into_iter()
+            .filter_map(|r| to_location(r.file_id, r.range, &world, &line_index).ok())
+            .collect()
+    } else {
+        // Only iterate over the references if include_declaration was false
+        refs.references()
+            .iter()
+            .filter_map(|r| to_location(r.file_id, r.range, &world, &line_index).ok())
+            .collect()
+    };
+
+    Ok(Some(locations))
 }
 
 pub fn handle_formatting(
@@ -712,11 +727,14 @@ pub fn handle_document_highlight(
     let file_id = params.text_document.try_conv_with(&world)?;
     let line_index = world.analysis().file_line_index(file_id);
 
-    let refs = world.analysis().find_all_refs(params.try_conv_with(&world)?)?;
+    let refs = match world.analysis().find_all_refs(params.try_conv_with(&world)?)? {
+        None => return Ok(None),
+        Some(refs) => refs,
+    };
 
     Ok(Some(
         refs.into_iter()
-            .map(|r| DocumentHighlight { range: r.1.conv_with(&line_index), kind: None })
+            .map(|r| DocumentHighlight { range: r.range.conv_with(&line_index), kind: None })
             .collect(),
     ))
 }