From f1a4cd3edf6d7fcd4358ec8a415f641e5845e8a0 Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Mon, 28 Oct 2024 10:41:22 -0300
Subject: [PATCH] fix: (LSP) check visibility of module that re-exports item,
 if any

---
 .../requests/code_action/import_or_qualify.rs |  1 +
 .../src/requests/completion/auto_import.rs    |  1 +
 tooling/lsp/src/requests/completion/tests.rs  | 23 +++++++++++++++++++
 tooling/lsp/src/visibility.rs                 | 14 +++++++----
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/tooling/lsp/src/requests/code_action/import_or_qualify.rs b/tooling/lsp/src/requests/code_action/import_or_qualify.rs
index 2e051890544..ffc83b05a5b 100644
--- a/tooling/lsp/src/requests/code_action/import_or_qualify.rs
+++ b/tooling/lsp/src/requests/code_action/import_or_qualify.rs
@@ -46,6 +46,7 @@ impl<'a> CodeActionFinder<'a> {
                     *module_def_id,
                     self.module_id,
                     *visibility,
+                    *defining_module,
                     self.interner,
                     self.def_maps,
                 ) {
diff --git a/tooling/lsp/src/requests/completion/auto_import.rs b/tooling/lsp/src/requests/completion/auto_import.rs
index 3b12d941c98..0e80b284f32 100644
--- a/tooling/lsp/src/requests/completion/auto_import.rs
+++ b/tooling/lsp/src/requests/completion/auto_import.rs
@@ -38,6 +38,7 @@ impl<'a> NodeFinder<'a> {
                     *module_def_id,
                     self.module_id,
                     *visibility,
+                    *defining_module,
                     self.interner,
                     self.def_maps,
                 ) {
diff --git a/tooling/lsp/src/requests/completion/tests.rs b/tooling/lsp/src/requests/completion/tests.rs
index b399088d05b..adc9b8ba140 100644
--- a/tooling/lsp/src/requests/completion/tests.rs
+++ b/tooling/lsp/src/requests/completion/tests.rs
@@ -1642,6 +1642,29 @@ fn main() {
         assert!(items.is_empty());
     }
 
+    #[test]
+    async fn checks_visibility_of_module_that_exports_item_if_any() {
+        let src = r#"
+            mod foo {
+                mod bar {
+                    pub fn hello_world() {}
+                }
+
+                pub use bar::hello_world;
+            }
+
+            fn main() {
+                hello_w>|<
+            }
+        "#;
+        let mut items = get_completions(src).await;
+        assert_eq!(items.len(), 1);
+
+        let item = items.remove(0);
+        assert_eq!(item.label, "hello_world()");
+        assert_eq!(item.label_details.unwrap().detail.unwrap(), "(use foo::hello_world)");
+    }
+
     #[test]
     async fn test_auto_import_suggests_modules_too() {
         let src = r#"mod foo {
diff --git a/tooling/lsp/src/visibility.rs b/tooling/lsp/src/visibility.rs
index f5b807055d5..6a21525f069 100644
--- a/tooling/lsp/src/visibility.rs
+++ b/tooling/lsp/src/visibility.rs
@@ -37,10 +37,13 @@ pub(super) fn item_in_module_is_visible(
 
 /// Returns true if the given ModuleDefId is visible from the current module, given its visibility.
 /// This will in turn check if the ModuleDefId parent modules are visible from the current module.
+/// If `defining_module` is Some, it will be considered as the parent of the item to check
+/// (this is the case when the item is re-exported with `pub use` or similar).
 pub(super) fn module_def_id_is_visible(
     module_def_id: ModuleDefId,
     current_module_id: ModuleId,
     mut visibility: ItemVisibility,
+    mut defining_module: Option<ModuleId>,
     interner: &NodeInterner,
     def_maps: &BTreeMap<CrateId, CrateDefMap>,
 ) -> bool {
@@ -49,7 +52,7 @@ pub(super) fn module_def_id_is_visible(
     let mut target_module_id = if let ModuleDefId::ModuleId(module_id) = module_def_id {
         Some(module_id)
     } else {
-        get_parent_module(interner, module_def_id)
+        std::mem::take(&mut defining_module).or_else(|| get_parent_module(interner, module_def_id))
     };
 
     // Then check if it's visible, and upwards
@@ -58,10 +61,11 @@ pub(super) fn module_def_id_is_visible(
             return false;
         }
 
-        let module_data = &def_maps[&module_id.krate].modules()[module_id.local_id.0];
-        let parent_local_id = module_data.parent;
-        target_module_id =
-            parent_local_id.map(|local_id| ModuleId { krate: module_id.krate, local_id });
+        target_module_id = std::mem::take(&mut defining_module).or_else(|| {
+            let module_data = &def_maps[&module_id.krate].modules()[module_id.local_id.0];
+            let parent_local_id = module_data.parent;
+            parent_local_id.map(|local_id| ModuleId { krate: module_id.krate, local_id })
+        });
 
         // This is a bit strange, but the visibility is always that of the item inside another module,
         // so the visibility we update here is for the next loop check.