1 files changed, 90 insertions, 30 deletions

diff --git a/crates/ra_assists/src/handlers/fix_visibility.rs b/crates/ra_assists/src/handlers/fix_visibility.rs
index 19d4dac5e..1aefa79cc 100644
--- a/crates/ra_assists/src/handlers/fix_visibility.rs
+++ b/crates/ra_assists/src/handlers/fix_visibility.rs
@@ -2,7 +2,8 @@ use hir::{db::HirDatabase, HasSource, HasVisibility, PathResolution};
 use ra_db::FileId;
 use ra_syntax::{ast, AstNode, TextRange, TextSize};
 
-use crate::{utils::vis_offset, AssistContext, AssistId, Assists};
+use crate::{utils::vis_offset, AssistContext, AssistId, AssistKind, Assists};
+use ast::VisibilityOwner;
 
 // FIXME: this really should be a fix for diagnostic, rather than an assist.
 
@@ -41,14 +42,15 @@ fn add_vis_to_referenced_module_def(acc: &mut Assists, ctx: &AssistContext) -> O
     };
 
     let current_module = ctx.sema.scope(&path.syntax()).module()?;
-    let target_module = def.module(ctx.db)?;
+    let target_module = def.module(ctx.db())?;
 
-    let vis = target_module.visibility_of(ctx.db, &def)?;
-    if vis.is_visible_from(ctx.db, current_module.into()) {
+    let vis = target_module.visibility_of(ctx.db(), &def)?;
+    if vis.is_visible_from(ctx.db(), current_module.into()) {
         return None;
     };
 
-    let (offset, target, target_file, target_name) = target_data_for_def(ctx.db, def)?;
+    let (offset, current_visibility, target, target_file, target_name) =
+        target_data_for_def(ctx.db(), def)?;
 
     let missing_visibility =
         if current_module.krate() == target_module.krate() { "pub(crate)" } else { "pub" };
@@ -58,54 +60,78 @@ fn add_vis_to_referenced_module_def(acc: &mut Assists, ctx: &AssistContext) -> O
         Some(name) => format!("Change visibility of {} to {}", name, missing_visibility),
     };
 
-    acc.add(AssistId("fix_visibility"), assist_label, target, |builder| {
+    acc.add(AssistId("fix_visibility", AssistKind::QuickFix), assist_label, target, |builder| {
         builder.edit_file(target_file);
         match ctx.config.snippet_cap {
-            Some(cap) => builder.insert_snippet(cap, offset, format!("$0{} ", missing_visibility)),
-            None => builder.insert(offset, format!("{} ", missing_visibility)),
+            Some(cap) => match current_visibility {
+                Some(current_visibility) => builder.replace_snippet(
+                    cap,
+                    current_visibility.syntax().text_range(),
+                    format!("$0{}", missing_visibility),
+                ),
+                None => builder.insert_snippet(cap, offset, format!("$0{} ", missing_visibility)),
+            },
+            None => match current_visibility {
+                Some(current_visibility) => {
+                    builder.replace(current_visibility.syntax().text_range(), missing_visibility)
+                }
+                None => builder.insert(offset, format!("{} ", missing_visibility)),
+            },
         }
     })
 }
 
 fn add_vis_to_referenced_record_field(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let record_field: ast::RecordField = ctx.find_node_at_offset()?;
+    let record_field: ast::RecordExprField = ctx.find_node_at_offset()?;
     let (record_field_def, _) = ctx.sema.resolve_record_field(&record_field)?;
 
     let current_module = ctx.sema.scope(record_field.syntax()).module()?;
-    let visibility = record_field_def.visibility(ctx.db);
-    if visibility.is_visible_from(ctx.db, current_module.into()) {
+    let visibility = record_field_def.visibility(ctx.db());
+    if visibility.is_visible_from(ctx.db(), current_module.into()) {
         return None;
     }
 
-    let parent = record_field_def.parent_def(ctx.db);
-    let parent_name = parent.name(ctx.db);
-    let target_module = parent.module(ctx.db);
+    let parent = record_field_def.parent_def(ctx.db());
+    let parent_name = parent.name(ctx.db());
+    let target_module = parent.module(ctx.db());
 
-    let in_file_source = record_field_def.source(ctx.db);
-    let (offset, target) = match in_file_source.value {
+    let in_file_source = record_field_def.source(ctx.db());
+    let (offset, current_visibility, target) = match in_file_source.value {
         hir::FieldSource::Named(it) => {
             let s = it.syntax();
-            (vis_offset(s), s.text_range())
+            (vis_offset(s), it.visibility(), s.text_range())
         }
         hir::FieldSource::Pos(it) => {
             let s = it.syntax();
-            (vis_offset(s), s.text_range())
+            (vis_offset(s), it.visibility(), s.text_range())
         }
     };
 
     let missing_visibility =
         if current_module.krate() == target_module.krate() { "pub(crate)" } else { "pub" };
-    let target_file = in_file_source.file_id.original_file(ctx.db);
+    let target_file = in_file_source.file_id.original_file(ctx.db());
 
-    let target_name = record_field_def.name(ctx.db);
+    let target_name = record_field_def.name(ctx.db());
     let assist_label =
         format!("Change visibility of {}.{} to {}", parent_name, target_name, missing_visibility);
 
-    acc.add(AssistId("fix_visibility"), assist_label, target, |builder| {
+    acc.add(AssistId("fix_visibility", AssistKind::QuickFix), assist_label, target, |builder| {
         builder.edit_file(target_file);
         match ctx.config.snippet_cap {
-            Some(cap) => builder.insert_snippet(cap, offset, format!("$0{} ", missing_visibility)),
-            None => builder.insert(offset, format!("{} ", missing_visibility)),
+            Some(cap) => match current_visibility {
+                Some(current_visibility) => builder.replace_snippet(
+                    cap,
+                    dbg!(current_visibility.syntax()).text_range(),
+                    format!("$0{}", missing_visibility),
+                ),
+                None => builder.insert_snippet(cap, offset, format!("$0{} ", missing_visibility)),
+            },
+            None => match current_visibility {
+                Some(current_visibility) => {
+                    builder.replace(current_visibility.syntax().text_range(), missing_visibility)
+                }
+                None => builder.insert(offset, format!("{} ", missing_visibility)),
+            },
         }
     })
 }
@@ -113,24 +139,30 @@ fn add_vis_to_referenced_record_field(acc: &mut Assists, ctx: &AssistContext) ->
 fn target_data_for_def(
     db: &dyn HirDatabase,
     def: hir::ModuleDef,
-) -> Option<(TextSize, TextRange, FileId, Option<hir::Name>)> {
+) -> Option<(TextSize, Option<ast::Visibility>, TextRange, FileId, Option<hir::Name>)> {
     fn offset_target_and_file_id<S, Ast>(
         db: &dyn HirDatabase,
         x: S,
-    ) -> (TextSize, TextRange, FileId)
+    ) -> (TextSize, Option<ast::Visibility>, TextRange, FileId)
     where
         S: HasSource<Ast = Ast>,
-        Ast: AstNode,
+        Ast: AstNode + ast::VisibilityOwner,
     {
         let source = x.source(db);
         let in_file_syntax = source.syntax();
         let file_id = in_file_syntax.file_id;
         let syntax = in_file_syntax.value;
-        (vis_offset(syntax), syntax.text_range(), file_id.original_file(db.upcast()))
+        let current_visibility = source.value.visibility();
+        (
+            vis_offset(syntax),
+            current_visibility,
+            syntax.text_range(),
+            file_id.original_file(db.upcast()),
+        )
     }
 
     let target_name;
-    let (offset, target, target_file) = match def {
+    let (offset, current_visibility, target, target_file) = match def {
         hir::ModuleDef::Function(f) => {
             target_name = Some(f.name(db));
             offset_target_and_file_id(db, f)
@@ -164,13 +196,13 @@ fn target_data_for_def(
             let in_file_source = m.declaration_source(db)?;
             let file_id = in_file_source.file_id.original_file(db.upcast());
             let syntax = in_file_source.value.syntax();
-            (vis_offset(syntax), syntax.text_range(), file_id)
+            (vis_offset(syntax), in_file_source.value.visibility(), syntax.text_range(), file_id)
         }
         // Enum variants can't be private, we can't modify builtin types
         hir::ModuleDef::EnumVariant(_) | hir::ModuleDef::BuiltinType(_) => return None,
     };
 
-    Some((offset, target, target_file, target_name))
+    Some((offset, current_visibility, target, target_file, target_name))
 }
 
 #[cfg(test)]
@@ -523,6 +555,34 @@ struct Bar;
     }
 
     #[test]
+    fn replaces_pub_crate_with_pub() {
+        check_assist(
+            fix_visibility,
+            r"
+//- /main.rs crate:a deps:foo
+foo::Bar<|>
+//- /lib.rs crate:foo
+pub(crate) struct Bar;
+",
+            r"$0pub struct Bar;
+",
+        );
+        check_assist(
+            fix_visibility,
+            r"
+//- /main.rs crate:a deps:foo
+fn main() {
+    foo::Foo { <|>bar: () };
+}
+//- /lib.rs crate:foo
+pub struct Foo { pub(crate) bar: () }
+",
+            r"pub struct Foo { $0pub bar: () }
+",
+        );
+    }
+
+    #[test]
     #[ignore]
     // FIXME handle reexports properly
     fn fix_visibility_of_reexport() {