5 files changed, 135 insertions, 35 deletions

diff --git a/crates/ra_ide_api/src/lib.rs b/crates/ra_ide_api/src/lib.rs
index 19669a7f0..0832229fd 100644
--- a/crates/ra_ide_api/src/lib.rs
+++ b/crates/ra_ide_api/src/lib.rs
@@ -71,7 +71,7 @@ pub use crate::{
     inlay_hints::{InlayHint, InlayKind},
     line_index::{LineCol, LineIndex},
     line_index_utils::translate_offset_with_edit,
-    references::ReferenceSearchResult,
+    references::{ReferenceSearchResult, SearchScope},
     runnables::{Runnable, RunnableKind},
     syntax_highlighting::HighlightedRange,
 };
@@ -481,8 +481,9 @@ impl Analysis {
     pub fn find_all_refs(
         &self,
         position: FilePosition,
+        search_scope: Option<SearchScope>,
     ) -> Cancelable<Option<ReferenceSearchResult>> {
-        self.with_db(|db| references::find_all_refs(db, position).map(|it| it.info))
+        self.with_db(|db| references::find_all_refs(db, position, search_scope).map(|it| it.info))
     }
 
     /// Returns a short text describing element at position.
diff --git a/crates/ra_ide_api/src/references.rs b/crates/ra_ide_api/src/references.rs
index 8200bd1ef..b5b1c9a16 100644
--- a/crates/ra_ide_api/src/references.rs
+++ b/crates/ra_ide_api/src/references.rs
@@ -27,6 +27,8 @@ pub(crate) use self::{
     rename::rename,
 };
 
+pub use self::search_scope::SearchScope;
+
 #[derive(Debug, Clone)]
 pub struct ReferenceSearchResult {
     declaration: NavigationTarget,
@@ -67,6 +69,7 @@ impl IntoIterator for ReferenceSearchResult {
 pub(crate) fn find_all_refs(
     db: &RootDatabase,
     position: FilePosition,
+    search_scope: Option<SearchScope>,
 ) -> Option<RangeInfo<ReferenceSearchResult>> {
     let parse = db.parse(position.file_id);
     let syntax = parse.tree().syntax().clone();
@@ -86,7 +89,15 @@ pub(crate) fn find_all_refs(
         NameKind::GenericParam(_) => return None,
     };
 
-    let references = process_definition(db, def, name);
+    let search_scope = {
+        let base = def.search_scope(db);
+        match search_scope {
+            None => base,
+            Some(scope) => base.intersection(&scope),
+        }
+    };
+
+    let references = process_definition(db, def, name, search_scope);
 
     Some(RangeInfo::new(range, ReferenceSearchResult { declaration, references }))
 }
@@ -107,11 +118,15 @@ fn find_name<'a>(
     Some(RangeInfo::new(range, (name_ref.text().to_string(), def)))
 }
 
-fn process_definition(db: &RootDatabase, def: NameDefinition, name: String) -> Vec<FileRange> {
+fn process_definition(
+    db: &RootDatabase,
+    def: NameDefinition,
+    name: String,
+    scope: SearchScope,
+) -> Vec<FileRange> {
     let _p = profile("process_definition");
 
     let pat = name.as_str();
-    let scope = def.search_scope(db);
     let mut refs = vec![];
 
     for (file_id, search_range) in scope {
@@ -144,8 +159,8 @@ fn process_definition(db: &RootDatabase, def: NameDefinition, name: String) -> V
 #[cfg(test)]
 mod tests {
     use crate::{
-        mock_analysis::analysis_and_position, mock_analysis::single_file_with_position,
-        ReferenceSearchResult,
+        mock_analysis::{analysis_and_position, single_file_with_position, MockAnalysis},
+        ReferenceSearchResult, SearchScope,
     };
 
     #[test]
@@ -270,7 +285,7 @@ mod tests {
         "#;
 
         let (analysis, pos) = analysis_and_position(code);
-        let refs = analysis.find_all_refs(pos).unwrap().unwrap();
+        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
         assert_eq!(refs.len(), 3);
     }
 
@@ -296,7 +311,7 @@ mod tests {
         "#;
 
         let (analysis, pos) = analysis_and_position(code);
-        let refs = analysis.find_all_refs(pos).unwrap().unwrap();
+        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
         assert_eq!(refs.len(), 2);
     }
 
@@ -321,12 +336,40 @@ mod tests {
         "#;
 
         let (analysis, pos) = analysis_and_position(code);
-        let refs = analysis.find_all_refs(pos).unwrap().unwrap();
+        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
+        assert_eq!(refs.len(), 3);
+    }
+
+    #[test]
+    fn test_find_all_refs_with_scope() {
+        let code = r#"
+            //- /lib.rs
+            mod foo;
+            mod bar;
+
+            pub fn quux<|>() {}
+
+            //- /foo.rs
+            fn f() { super::quux(); }
+
+            //- /bar.rs
+            fn f() { super::quux(); }
+        "#;
+
+        let (mock, pos) = MockAnalysis::with_files_and_position(code);
+        let bar = mock.id_of("/bar.rs");
+        let analysis = mock.analysis();
+
+        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
         assert_eq!(refs.len(), 3);
+
+        let refs =
+            analysis.find_all_refs(pos, Some(SearchScope::single_file(bar))).unwrap().unwrap();
+        assert_eq!(refs.len(), 2);
     }
 
     fn get_all_refs(text: &str) -> ReferenceSearchResult {
         let (analysis, position) = single_file_with_position(text);
-        analysis.find_all_refs(position).unwrap().unwrap()
+        analysis.find_all_refs(position, None).unwrap().unwrap()
     }
 }
diff --git a/crates/ra_ide_api/src/references/rename.rs b/crates/ra_ide_api/src/references/rename.rs
index 0e2e088e0..ee6e73e1b 100644
--- a/crates/ra_ide_api/src/references/rename.rs
+++ b/crates/ra_ide_api/src/references/rename.rs
@@ -110,7 +110,7 @@ fn rename_reference(
     position: FilePosition,
     new_name: &str,
 ) -> Option<RangeInfo<SourceChange>> {
-    let RangeInfo { range, info: refs } = find_all_refs(db, position)?;
+    let RangeInfo { range, info: refs } = find_all_refs(db, position, None)?;
 
     let edit = refs
         .into_iter()
@@ -255,13 +255,13 @@ mod tests {
         "#;
 
         let (analysis, pos) = analysis_and_position(code);
-        let refs = analysis.find_all_refs(pos).unwrap().unwrap();
+        let refs = analysis.find_all_refs(pos, None).unwrap().unwrap();
         assert_eq!(refs.len(), 3);
     }
 
     fn get_all_refs(text: &str) -> ReferenceSearchResult {
         let (analysis, position) = single_file_with_position(text);
-        analysis.find_all_refs(position).unwrap().unwrap()
+        analysis.find_all_refs(position, None).unwrap().unwrap()
     }
 
     #[test]
diff --git a/crates/ra_ide_api/src/references/search_scope.rs b/crates/ra_ide_api/src/references/search_scope.rs
index 21e667c83..b6eb248b7 100644
--- a/crates/ra_ide_api/src/references/search_scope.rs
+++ b/crates/ra_ide_api/src/references/search_scope.rs
@@ -2,33 +2,84 @@
 //! For `pub(crate)` things it's a crate, for `pub` things it's a crate and dependant crates.
 //! In some cases, the location of the references is known to within a `TextRange`,
 //! e.g. for things like local variables.
+use std::mem;
 
 use hir::{DefWithBody, HasSource, ModuleSource};
 use ra_db::{FileId, SourceDatabase, SourceDatabaseExt};
 use ra_prof::profile;
 use ra_syntax::{AstNode, TextRange};
-use rustc_hash::FxHashSet;
+use rustc_hash::FxHashMap;
 
 use crate::db::RootDatabase;
 
 use super::{NameDefinition, NameKind};
 
+pub struct SearchScope {
+    entries: FxHashMap<FileId, Option<TextRange>>,
+}
+
+impl SearchScope {
+    fn new(entries: FxHashMap<FileId, Option<TextRange>>) -> SearchScope {
+        SearchScope { entries }
+    }
+    pub fn single_file(file: FileId) -> SearchScope {
+        SearchScope::new(std::iter::once((file, None)).collect())
+    }
+    pub(crate) fn intersection(&self, other: &SearchScope) -> SearchScope {
+        let (mut small, mut large) = (&self.entries, &other.entries);
+        if small.len() > large.len() {
+            mem::swap(&mut small, &mut large)
+        }
+
+        let res = small
+            .iter()
+            .filter_map(|(file_id, r1)| {
+                let r2 = large.get(file_id)?;
+                let r = intersect_ranges(*r1, *r2)?;
+                Some((*file_id, r))
+            })
+            .collect();
+        return SearchScope::new(res);
+
+        fn intersect_ranges(
+            r1: Option<TextRange>,
+            r2: Option<TextRange>,
+        ) -> Option<Option<TextRange>> {
+            match (r1, r2) {
+                (None, r) | (r, None) => Some(r),
+                (Some(r1), Some(r2)) => {
+                    let r = r1.intersection(&r2)?;
+                    Some(Some(r))
+                }
+            }
+        }
+    }
+}
+
+impl IntoIterator for SearchScope {
+    type Item = (FileId, Option<TextRange>);
+    type IntoIter = std::collections::hash_map::IntoIter<FileId, Option<TextRange>>;
+    fn into_iter(self) -> Self::IntoIter {
+        self.entries.into_iter()
+    }
+}
+
 impl NameDefinition {
-    pub(crate) fn search_scope(&self, db: &RootDatabase) -> FxHashSet<(FileId, Option<TextRange>)> {
+    pub(crate) fn search_scope(&self, db: &RootDatabase) -> SearchScope {
         let _p = profile("search_scope");
 
         let module_src = self.container.definition_source(db);
         let file_id = module_src.file_id.original_file(db);
 
         if let NameKind::Pat((def, _)) = self.kind {
-            let mut res = FxHashSet::default();
+            let mut res = FxHashMap::default();
             let range = match def {
                 DefWithBody::Function(f) => f.source(db).ast.syntax().text_range(),
                 DefWithBody::Const(c) => c.source(db).ast.syntax().text_range(),
                 DefWithBody::Static(s) => s.source(db).ast.syntax().text_range(),
             };
-            res.insert((file_id, Some(range)));
-            return res;
+            res.insert(file_id, Some(range));
+            return SearchScope::new(res);
         }
 
         let vis =
@@ -36,37 +87,37 @@ impl NameDefinition {
 
         if vis.as_str() == "pub(super)" {
             if let Some(parent_module) = self.container.parent(db) {
-                let mut files = FxHashSet::default();
+                let mut res = FxHashMap::default();
                 let parent_src = parent_module.definition_source(db);
                 let file_id = parent_src.file_id.original_file(db);
 
                 match parent_src.ast {
                     ModuleSource::Module(m) => {
                         let range = Some(m.syntax().text_range());
-                        files.insert((file_id, range));
+                        res.insert(file_id, range);
                     }
                     ModuleSource::SourceFile(_) => {
-                        files.insert((file_id, None));
-                        files.extend(parent_module.children(db).map(|m| {
+                        res.insert(file_id, None);
+                        res.extend(parent_module.children(db).map(|m| {
                             let src = m.definition_source(db);
                             (src.file_id.original_file(db), None)
                         }));
                     }
                 }
-                return files;
+                return SearchScope::new(res);
             }
         }
 
         if vis.as_str() != "" {
             let source_root_id = db.file_source_root(file_id);
             let source_root = db.source_root(source_root_id);
-            let mut files =
-                source_root.walk().map(|id| (id.into(), None)).collect::<FxHashSet<_>>();
+            let mut res =
+                source_root.walk().map(|id| (id.into(), None)).collect::<FxHashMap<_, _>>();
 
             // FIXME: add "pub(in path)"
 
             if vis.as_str() == "pub(crate)" {
-                return files;
+                return SearchScope::new(res);
             }
             if vis.as_str() == "pub" {
                 let krate = self.container.krate(db).unwrap();
@@ -77,19 +128,19 @@ impl NameDefinition {
                         let root_file = crate_graph.crate_root(crate_id);
                         let source_root_id = db.file_source_root(root_file);
                         let source_root = db.source_root(source_root_id);
-                        files.extend(source_root.walk().map(|id| (id.into(), None)));
+                        res.extend(source_root.walk().map(|id| (id.into(), None)));
                     }
                 }
-                return files;
+                return SearchScope::new(res);
             }
         }
 
-        let mut res = FxHashSet::default();
+        let mut res = FxHashMap::default();
         let range = match module_src.ast {
             ModuleSource::Module(m) => Some(m.syntax().text_range()),
             ModuleSource::SourceFile(_) => None,
         };
-        res.insert((file_id, range));
-        res
+        res.insert(file_id, range);
+        SearchScope::new(res)
     }
 }
diff --git a/crates/ra_lsp_server/src/main_loop/handlers.rs b/crates/ra_lsp_server/src/main_loop/handlers.rs
index e65f075a6..a29971d10 100644
--- a/crates/ra_lsp_server/src/main_loop/handlers.rs
+++ b/crates/ra_lsp_server/src/main_loop/handlers.rs
@@ -9,7 +9,9 @@ use lsp_types::{
     Hover, HoverContents, Location, MarkupContent, MarkupKind, Position, PrepareRenameResponse,
     Range, RenameParams, SymbolInformation, TextDocumentIdentifier, TextEdit, WorkspaceEdit,
 };
-use ra_ide_api::{AssistId, FileId, FilePosition, FileRange, Query, Runnable, RunnableKind};
+use ra_ide_api::{
+    AssistId, FileId, FilePosition, FileRange, Query, Runnable, RunnableKind, SearchScope,
+};
 use ra_prof::profile;
 use ra_syntax::{AstNode, SyntaxKind, TextRange, TextUnit};
 use rustc_hash::FxHashMap;
@@ -485,7 +487,7 @@ pub fn handle_references(
     let _p = profile("handle_references");
     let position = params.text_document_position.try_conv_with(&world)?;
 
-    let refs = match world.analysis().find_all_refs(position)? {
+    let refs = match world.analysis().find_all_refs(position, None)? {
         None => return Ok(None),
         Some(refs) => refs,
     };
@@ -748,7 +750,10 @@ 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 = match world.analysis().find_all_refs(params.try_conv_with(&world)?)? {
+    let refs = match world
+        .analysis()
+        .find_all_refs(params.try_conv_with(&world)?, Some(SearchScope::single_file(file_id)))?
+    {
         None => return Ok(None),
         Some(refs) => refs,
     };