Merge pull request #19158 from PoignardAzur/expaned_pub_glob_imports

Implement expand_glob_reexport assist

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)?;
+    if !is_visible_from(ctx, &target_module, current_module) {
-    let imported_defs = find_imported_defs(ctx, star)?;
+        return None;
+    }
 
     let target = parent.either(|n| n.syntax().clone(), |n| n.syntax().clone());
     acc.add(
@@ -62,9 +65,108 @@ pub(crate) fn expand_glob_import(acc: &mut Assists, ctx: &AssistContext<'_>) ->
         "Expand glob import",
         target.text_range(),
         |builder| {
+            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(ctx, refs_in_target, imported_defs);
+    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(),
@@ -89,8 +191,21 @@ pub(crate) fn expand_glob_import(acc: &mut Assists, ctx: &AssistContext<'_>) ->
         }
         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 {
@@ -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) => {
+            ScopeDef::ModuleDef(def) => Some(Ref {
-                Some(Ref { visible_name: name, def: Definition::from(def) })
+                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> {
+    must_be_pub: bool,
-    if !is_expandable_visible_from(ctx, &expandable, visible_from) {
+) -> Refs {
-        return None;
-    }
-
     match expandable {
         Expandable::Module(module) => {
             let module_scope = module.scope(ctx.db(), Some(visible_from));
-            let refs =
+            let refs = module_scope
-                module_scope.into_iter().filter_map(|(n, d)| Ref::from_scope_def(n, d)).collect();
+                .into_iter()
-            Some(Refs(refs))
+                .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(
+fn is_visible_from(ctx: &AssistContext<'_>, expandable: &Expandable, from: Module) -> bool {
-    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,22 +364,11 @@ fn is_expandable_visible_from(
 // use foo::*$0;
 // use baz::Baz;
 // ↑ ---------------
-fn find_imported_defs(ctx: &AssistContext<'_>, star: SyntaxToken) -> Option<Vec<Definition>> {
+fn find_imported_defs(ctx: &AssistContext<'_>, use_item: Use) -> 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
+            use_item.syntax().siblings(dir.to_owned()).filter(|n| ast::Use::can_cast(n.kind()))
-                    .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)? {
@@ -279,17 +386,12 @@ fn find_imported_defs(ctx: &AssistContext<'_>, star: SyntaxToken) -> Option<Vec<
             ) => Some(def),
             _ => None,
         })
-            .collect(),
+        .collect()
-    )
 }
 
-fn find_names_to_import(
+fn find_names_to_import(refs_in_target: Refs, imported_defs: Vec<Definition>) -> Vec<Name> {
-    ctx: &AssistContext<'_>,
+    let final_refs = refs_in_target.filter_out_by_defs(imported_defs);
-    refs_in_target: Refs,
+    final_refs.0.iter().map(|r| r.visible_name.clone()).collect()
-    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()
 }
 
 #[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;
+",
+        );
+    }
 }

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,

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(