From 34c28c1bbca08858966a90614a5b739efbb1f1c8 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Mon, 28 Jun 2021 16:41:35 +0200
Subject: [PATCH] Include `self` in usage search for modules in their
 definition source

---
 crates/ide/src/highlight_related.rs |  7 ++-
 crates/ide/src/references.rs        |  1 +
 crates/ide/src/rename.rs            | 10 ++++-
 crates/ide_db/src/search.rs         | 66 ++++++++++++++++++++++++++++-
 4 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/crates/ide/src/highlight_related.rs b/crates/ide/src/highlight_related.rs
index 2d27fb45e34..537d02e12b7 100644
--- a/crates/ide/src/highlight_related.rs
+++ b/crates/ide/src/highlight_related.rs
@@ -59,7 +59,11 @@ fn highlight_references(
     FilePosition { offset, file_id }: FilePosition,
 ) -> Option<Vec<HighlightedRange>> {
     let def = references::find_def(sema, syntax, offset)?;
-    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(file_id))).all();
+    let usages = def
+        .usages(sema)
+        .set_scope(Some(SearchScope::single_file(file_id)))
+        .include_self_refs()
+        .all();
 
     let declaration = match def {
         Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
@@ -315,6 +319,7 @@ use self$0;
 mod foo;
 //- /foo.rs
 use self$0;
+ // ^^^^
 "#,
         );
     }
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index 2d3a0f59812..aa03de4e417 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -707,6 +707,7 @@ use self$0;
             expect![[r#"
                 foo Module FileId(0) 0..8 4..7
 
+                FileId(1) 4..8
             "#]],
         );
     }
diff --git a/crates/ide/src/rename.rs b/crates/ide/src/rename.rs
index 4f5f2c935c2..7342ef4fcf4 100644
--- a/crates/ide/src/rename.rs
+++ b/crates/ide/src/rename.rs
@@ -934,12 +934,18 @@ mod outer { mod fo$0o; }
         check(
             "baz",
             r#"
-mod $0foo { pub fn bar() {} }
+mod $0foo {
+    pub use self::bar as qux;
+    pub fn bar() {}
+}
 
 fn main() { foo::bar(); }
 "#,
             r#"
-mod baz { pub fn bar() {} }
+mod baz {
+    pub use self::bar as qux;
+    pub fn bar() {}
+}
 
 fn main() { baz::bar(); }
 "#,
diff --git a/crates/ide_db/src/search.rs b/crates/ide_db/src/search.rs
index baed5e88082..7be76dca107 100644
--- a/crates/ide_db/src/search.rs
+++ b/crates/ide_db/src/search.rs
@@ -305,7 +305,13 @@ impl Definition {
     }
 
     pub fn usages<'a>(self, sema: &'a Semantics<RootDatabase>) -> FindUsages<'a> {
-        FindUsages { def: self, sema, scope: None, include_self_kw_refs: None }
+        FindUsages {
+            def: self,
+            sema,
+            scope: None,
+            include_self_kw_refs: None,
+            search_self_mod: false,
+        }
     }
 }
 
@@ -314,12 +320,14 @@ pub struct FindUsages<'a> {
     sema: &'a Semantics<'a, RootDatabase>,
     scope: Option<SearchScope>,
     include_self_kw_refs: Option<hir::Type>,
+    search_self_mod: bool,
 }
 
 impl<'a> FindUsages<'a> {
-    /// Enable searching for `Self` when the definition is a type.
+    /// Enable searching for `Self` when the definition is a type or `self` for modules.
     pub fn include_self_refs(mut self) -> FindUsages<'a> {
         self.include_self_kw_refs = def_to_ty(self.sema, &self.def);
+        self.search_self_mod = true;
         self
     }
 
@@ -416,6 +424,41 @@ impl<'a> FindUsages<'a> {
                 }
             }
         }
+
+        // search for module `self` references in our module's definition source
+        match self.def {
+            Definition::ModuleDef(hir::ModuleDef::Module(module)) if self.search_self_mod => {
+                let src = module.definition_source(sema.db);
+                let file_id = src.file_id.original_file(sema.db);
+                let (file_id, search_range) = match src.value {
+                    ModuleSource::Module(m) => (file_id, Some(m.syntax().text_range())),
+                    ModuleSource::BlockExpr(b) => (file_id, Some(b.syntax().text_range())),
+                    ModuleSource::SourceFile(_) => (file_id, None),
+                };
+
+                let text = sema.db.file_text(file_id);
+                let search_range =
+                    search_range.unwrap_or_else(|| TextRange::up_to(TextSize::of(text.as_str())));
+
+                let tree = Lazy::new(|| sema.parse(file_id).syntax().clone());
+
+                for (idx, _) in text.match_indices("self") {
+                    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_module_name_ref(&name_ref, sink) {
+                            return;
+                        }
+                    }
+                }
+            }
+            _ => {}
+        }
     }
 
     fn found_self_ty_name_ref(
@@ -440,6 +483,25 @@ impl<'a> FindUsages<'a> {
         }
     }
 
+    fn found_self_module_name_ref(
+        &self,
+        name_ref: &ast::NameRef,
+        sink: &mut dyn FnMut(FileId, FileReference) -> bool,
+    ) -> bool {
+        match NameRefClass::classify(self.sema, name_ref) {
+            Some(NameRefClass::Definition(def @ Definition::ModuleDef(_))) if def == 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,
+                };
+                sink(file_id, reference)
+            }
+            _ => false,
+        }
+    }
+
     fn found_lifetime(
         &self,
         lifetime: &ast::Lifetime,