Merge pull request #19158 from PoignardAzur/expaned_pub_glob_imports

Implement expand_glob_reexport assist
2025-02-17 09:46:29 +00:00 · 2025-02-17 09:46:29 +00:00 · c7a2986c20
commit c7a2986c20
parent d7fb8f15c7 824b18248c
3 changed files with 297 additions and 92 deletions
--- a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/expand_glob_import.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/expand_glob_import.rs
@ -3,10 +3,11 @@ use hir::{AssocItem, Enum, HasVisibility, Module, ModuleDef, Name, PathResolutio
 use ide_db::{
    defs::{Definition, NameRefClass},
    search::SearchScope,
+    source_change::SourceChangeBuilder,
 };
 use stdx::never;
 use syntax::{
-    ast::{self, make},
+    ast::{self, make, Use, UseTree, VisibilityKind},
    ted, AstNode, Direction, SyntaxNode, SyntaxToken, T,
 };

@ -43,6 +44,7 @@ use crate::{
 pub(crate) fn expand_glob_import(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let star = ctx.find_token_syntax_at_offset(T![*])?;
    let use_tree = star.parent().and_then(ast::UseTree::cast)?;
+    let use_item = star.parent_ancestors().find_map(ast::Use::cast)?;
    let (parent, mod_path) = find_parent_and_path(&star)?;
    let target_module = match ctx.sema.resolve_path(&mod_path)? {
        PathResolution::Def(ModuleDef::Module(it)) => Expandable::Module(it),
@ -53,8 +55,9 @@ pub(crate) fn expand_glob_import(acc: &mut Assists, ctx: &AssistContext<'_>) ->
    let current_scope = ctx.sema.scope(&star.parent()?)?;
    let current_module = current_scope.module();

-    let refs_in_target = find_refs_in_mod(ctx, target_module, current_module)?;
-    let imported_defs = find_imported_defs(ctx, star)?;
+    if !is_visible_from(ctx, &target_module, current_module) {
+        return None;
+    }

    let target = parent.either(|n| n.syntax().clone(), |n| n.syntax().clone());
    acc.add(
@ -62,37 +65,149 @@ pub(crate) fn expand_glob_import(acc: &mut Assists, ctx: &AssistContext<'_>) ->
        "Expand glob import",
        target.text_range(),
        |builder| {
-            let use_tree = builder.make_mut(use_tree);
-
-            let names_to_import = find_names_to_import(ctx, refs_in_target, imported_defs);
-            let expanded = make::use_tree_list(names_to_import.iter().map(|n| {
-                let path = make::ext::ident_path(
-                    &n.display(ctx.db(), current_module.krate().edition(ctx.db())).to_string(),
-                );
-                make::use_tree(path, None, None, false)
-            }))
-            .clone_for_update();
-
-            match use_tree.star_token() {
-                Some(star) => {
-                    let needs_braces = use_tree.path().is_some() && names_to_import.len() != 1;
-                    if needs_braces {
-                        ted::replace(star, expanded.syntax())
-                    } else {
-                        let without_braces = expanded
-                            .syntax()
-                            .children_with_tokens()
-                            .filter(|child| !matches!(child.kind(), T!['{'] | T!['}']))
-                            .collect();
-                        ted::replace_with_many(star, without_braces)
-                    }
-                }
-                None => never!(),
-            }
+            build_expanded_import(
+                ctx,
+                builder,
+                use_tree,
+                use_item,
+                target_module,
+                current_module,
+                false,
+            )
        },
    )
 }

+// Assist: expand_glob_reexport
+//
+// Expands non-private glob imports.
+//
+// ```
+// mod foo {
+//     pub struct Bar;
+//     pub struct Baz;
+// }
+//
+// pub use foo::*$0;
+// ```
+// ->
+// ```
+// mod foo {
+//     pub struct Bar;
+//     pub struct Baz;
+// }
+//
+// pub use foo::{Bar, Baz};
+// ```
+pub(crate) fn expand_glob_reexport(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let star = ctx.find_token_syntax_at_offset(T![*])?;
+    let use_tree = star.parent().and_then(ast::UseTree::cast)?;
+    let use_item = star.parent_ancestors().find_map(ast::Use::cast)?;
+    let (parent, mod_path) = find_parent_and_path(&star)?;
+    let target_module = match ctx.sema.resolve_path(&mod_path)? {
+        PathResolution::Def(ModuleDef::Module(it)) => Expandable::Module(it),
+        PathResolution::Def(ModuleDef::Adt(hir::Adt::Enum(e))) => Expandable::Enum(e),
+        _ => return None,
+    };
+
+    let current_scope = ctx.sema.scope(&star.parent()?)?;
+    let current_module = current_scope.module();
+
+    if let VisibilityKind::PubSelf = get_export_visibility_kind(&use_item) {
+        return None;
+    }
+    if !is_visible_from(ctx, &target_module, current_module) {
+        return None;
+    }
+
+    let target = parent.either(|n| n.syntax().clone(), |n| n.syntax().clone());
+    acc.add(
+        AssistId("expand_glob_reexport", AssistKind::RefactorRewrite),
+        "Expand glob reexport",
+        target.text_range(),
+        |builder| {
+            build_expanded_import(
+                ctx,
+                builder,
+                use_tree,
+                use_item,
+                target_module,
+                current_module,
+                true,
+            )
+        },
+    )
+}
+
+fn build_expanded_import(
+    ctx: &AssistContext<'_>,
+    builder: &mut SourceChangeBuilder,
+    use_tree: UseTree,
+    use_item: Use,
+    target_module: Expandable,
+    current_module: Module,
+    reexport_public_items: bool,
+) {
+    let (must_be_pub, visible_from) = if !reexport_public_items {
+        (false, current_module)
+    } else {
+        match get_export_visibility_kind(&use_item) {
+            VisibilityKind::Pub => (true, current_module.krate().root_module()),
+            VisibilityKind::PubCrate => (false, current_module.krate().root_module()),
+            _ => (false, current_module),
+        }
+    };
+
+    let refs_in_target = find_refs_in_mod(ctx, target_module, visible_from, must_be_pub);
+    let imported_defs = find_imported_defs(ctx, use_item);
+
+    let filtered_defs =
+        if reexport_public_items { refs_in_target } else { refs_in_target.used_refs(ctx) };
+
+    let use_tree = builder.make_mut(use_tree);
+
+    let names_to_import = find_names_to_import(filtered_defs, imported_defs);
+    let expanded = make::use_tree_list(names_to_import.iter().map(|n| {
+        let path = make::ext::ident_path(
+            &n.display(ctx.db(), current_module.krate().edition(ctx.db())).to_string(),
+        );
+        make::use_tree(path, None, None, false)
+    }))
+    .clone_for_update();
+
+    match use_tree.star_token() {
+        Some(star) => {
+            let needs_braces = use_tree.path().is_some() && names_to_import.len() != 1;
+            if needs_braces {
+                ted::replace(star, expanded.syntax())
+            } else {
+                let without_braces = expanded
+                    .syntax()
+                    .children_with_tokens()
+                    .filter(|child| !matches!(child.kind(), T!['{'] | T!['}']))
+                    .collect();
+                ted::replace_with_many(star, without_braces)
+            }
+        }
+        None => never!(),
+    }
+}
+
+fn get_export_visibility_kind(use_item: &Use) -> VisibilityKind {
+    use syntax::ast::HasVisibility as _;
+    match use_item.visibility() {
+        Some(vis) => match vis.kind() {
+            VisibilityKind::PubCrate => VisibilityKind::PubCrate,
+            VisibilityKind::Pub => VisibilityKind::Pub,
+            VisibilityKind::PubSelf => VisibilityKind::PubSelf,
+            // We don't handle pub(in ...) and pub(super) yet
+            VisibilityKind::In(_) => VisibilityKind::PubSelf,
+            VisibilityKind::PubSuper => VisibilityKind::PubSelf,
+        },
+        None => VisibilityKind::PubSelf,
+    }
+}
+
 enum Expandable {
    Module(Module),
    Enum(Enum),
@ -130,14 +245,17 @@ struct Ref {
    // could be alias
    visible_name: Name,
    def: Definition,
+    is_pub: bool,
 }

 impl Ref {
-    fn from_scope_def(name: Name, scope_def: ScopeDef) -> Option<Self> {
+    fn from_scope_def(ctx: &AssistContext<'_>, name: Name, scope_def: ScopeDef) -> Option<Self> {
        match scope_def {
-            ScopeDef::ModuleDef(def) => {
-                Some(Ref { visible_name: name, def: Definition::from(def) })
-            }
+            ScopeDef::ModuleDef(def) => Some(Ref {
+                visible_name: name,
+                def: Definition::from(def),
+                is_pub: matches!(def.visibility(ctx.db()), hir::Visibility::Public),
+            }),
            _ => None,
        }
    }
@ -180,32 +298,32 @@ fn find_refs_in_mod(
    ctx: &AssistContext<'_>,
    expandable: Expandable,
    visible_from: Module,
-) -> Option<Refs> {
-    if !is_expandable_visible_from(ctx, &expandable, visible_from) {
-        return None;
-    }
-
+    must_be_pub: bool,
+) -> Refs {
    match expandable {
        Expandable::Module(module) => {
            let module_scope = module.scope(ctx.db(), Some(visible_from));
-            let refs =
-                module_scope.into_iter().filter_map(|(n, d)| Ref::from_scope_def(n, d)).collect();
-            Some(Refs(refs))
+            let refs = module_scope
+                .into_iter()
+                .filter_map(|(n, d)| Ref::from_scope_def(ctx, n, d))
+                .filter(|r| !must_be_pub || r.is_pub)
+                .collect();
+            Refs(refs)
        }
-        Expandable::Enum(enm) => Some(Refs(
+        Expandable::Enum(enm) => Refs(
            enm.variants(ctx.db())
                .into_iter()
-                .map(|v| Ref { visible_name: v.name(ctx.db()), def: Definition::Variant(v) })
+                .map(|v| Ref {
+                    visible_name: v.name(ctx.db()),
+                    def: Definition::Variant(v),
+                    is_pub: true,
+                })
                .collect(),
-        )),
+        ),
    }
 }

-fn is_expandable_visible_from(
-    ctx: &AssistContext<'_>,
-    expandable: &Expandable,
-    from: Module,
-) -> bool {
+fn is_visible_from(ctx: &AssistContext<'_>, expandable: &Expandable, from: Module) -> bool {
    fn is_mod_visible_from(ctx: &AssistContext<'_>, module: Module, from: Module) -> bool {
        match module.parent(ctx.db()) {
            Some(parent) => {
@ -246,50 +364,34 @@ fn is_expandable_visible_from(
 // use foo::*$0;
 // use baz::Baz;
 // ↑ ---------------
-fn find_imported_defs(ctx: &AssistContext<'_>, star: SyntaxToken) -> Option<Vec<Definition>> {
-    let parent_use_item_syntax = star.parent_ancestors().find_map(|n| {
-        if ast::Use::can_cast(n.kind()) {
-            Some(n)
-        } else {
-            None
-        }
-    })?;
-
-    Some(
-        [Direction::Prev, Direction::Next]
-            .into_iter()
-            .flat_map(|dir| {
-                parent_use_item_syntax
-                    .siblings(dir.to_owned())
-                    .filter(|n| ast::Use::can_cast(n.kind()))
-            })
-            .flat_map(|n| n.descendants().filter_map(ast::NameRef::cast))
-            .filter_map(|r| match NameRefClass::classify(&ctx.sema, &r)? {
-                NameRefClass::Definition(
-                    def @ (Definition::Macro(_)
-                    | Definition::Module(_)
-                    | Definition::Function(_)
-                    | Definition::Adt(_)
-                    | Definition::Variant(_)
-                    | Definition::Const(_)
-                    | Definition::Static(_)
-                    | Definition::Trait(_)
-                    | Definition::TypeAlias(_)),
-                    _,
-                ) => Some(def),
-                _ => None,
-            })
-            .collect(),
-    )
+fn find_imported_defs(ctx: &AssistContext<'_>, use_item: Use) -> Vec<Definition> {
+    [Direction::Prev, Direction::Next]
+        .into_iter()
+        .flat_map(|dir| {
+            use_item.syntax().siblings(dir.to_owned()).filter(|n| ast::Use::can_cast(n.kind()))
+        })
+        .flat_map(|n| n.descendants().filter_map(ast::NameRef::cast))
+        .filter_map(|r| match NameRefClass::classify(&ctx.sema, &r)? {
+            NameRefClass::Definition(
+                def @ (Definition::Macro(_)
+                | Definition::Module(_)
+                | Definition::Function(_)
+                | Definition::Adt(_)
+                | Definition::Variant(_)
+                | Definition::Const(_)
+                | Definition::Static(_)
+                | Definition::Trait(_)
+                | Definition::TypeAlias(_)),
+                _,
+            ) => Some(def),
+            _ => None,
+        })
+        .collect()
 }

-fn find_names_to_import(
-    ctx: &AssistContext<'_>,
-    refs_in_target: Refs,
-    imported_defs: Vec<Definition>,
-) -> Vec<Name> {
-    let used_refs = refs_in_target.used_refs(ctx).filter_out_by_defs(imported_defs);
-    used_refs.0.iter().map(|r| r.visible_name.clone()).collect()
+fn find_names_to_import(refs_in_target: Refs, imported_defs: Vec<Definition>) -> Vec<Name> {
+    let final_refs = refs_in_target.filter_out_by_defs(imported_defs);
+    final_refs.0.iter().map(|r| r.visible_name.clone()).collect()
 }

 #[cfg(test)]
@ -1036,4 +1138,83 @@ mod abc {
 }"#,
        )
    }
+
+    #[test]
+    fn expanding_glob_reexport() {
+        check_assist(
+            expand_glob_reexport,
+            r"
+mod foo {
+    pub struct Bar;
+    pub struct Baz;
+    struct Qux;
+
+    pub fn f() {}
+
+    pub(crate) fn g() {}
+    pub(self) fn h() {}
+}
+
+pub use foo::*$0;
+",
+            r"
+mod foo {
+    pub struct Bar;
+    pub struct Baz;
+    struct Qux;
+
+    pub fn f() {}
+
+    pub(crate) fn g() {}
+    pub(self) fn h() {}
+}
+
+pub use foo::{Bar, Baz, f};
+",
+        )
+    }
+
+    #[test]
+    fn expanding_recursive_glob_reexport() {
+        check_assist(
+            expand_glob_reexport,
+            r"
+mod foo {
+    pub use bar::*;
+    mod bar {
+        pub struct Bar;
+        pub struct Baz;
+    }
+}
+
+pub use foo::*$0;
+",
+            r"
+mod foo {
+    pub use bar::*;
+    mod bar {
+        pub struct Bar;
+        pub struct Baz;
+    }
+}
+
+pub use foo::{Bar, Baz};
+",
+        )
+    }
+
+    #[test]
+    fn expanding_reexport_is_not_applicable_for_private_import() {
+        check_assist_not_applicable(
+            expand_glob_reexport,
+            r"
+mod foo {
+    pub struct Bar;
+    pub struct Baz;
+}
+
+use foo::*$0;
+",
+        );
+    }
 }
--- a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
@ -270,6 +270,7 @@ mod handlers {
            destructure_tuple_binding::destructure_tuple_binding,
            destructure_struct_binding::destructure_struct_binding,
            expand_glob_import::expand_glob_import,
+            expand_glob_import::expand_glob_reexport,
            explicit_enum_discriminant::explicit_enum_discriminant,
            extract_expressions_from_format_string::extract_expressions_from_format_string,
            extract_struct_from_enum_variant::extract_struct_from_enum_variant,
--- a/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
@ -909,6 +909,29 @@ fn qux(bar: Bar, baz: Baz) {}
    )
 }

+#[test]
+fn doctest_expand_glob_reexport() {
+    check_doc_test(
+        "expand_glob_reexport",
+        r#####"
+mod foo {
+    pub struct Bar;
+    pub struct Baz;
+}
+
+pub use foo::*$0;
+"#####,
+        r#####"
+mod foo {
+    pub struct Bar;
+    pub struct Baz;
+}
+
+pub use foo::{Bar, Baz};
+"#####,
+    )
+}
+
 #[test]
 fn doctest_explicit_enum_discriminant() {
    check_doc_test(