change visibility can change pub to pub(crate)

1 files changed, 30 insertions, 2 deletions

diff --git a/crates/ra_editor/src/assists/change_visibility.rs b/crates/ra_editor/src/assists/change_visibility.rs
index ac75f635e..6c8466394 100644
--- a/crates/ra_editor/src/assists/change_visibility.rs
+++ b/crates/ra_editor/src/assists/change_visibility.rs
@@ -7,10 +7,19 @@ use ra_syntax::{
 use crate::assists::{AssistCtx, Assist};
 
 pub fn change_visibility(ctx: AssistCtx) -> Option<Assist> {
-    let offset = if let Some(keyword) = ctx.leaf_at_offset().find(|leaf| match leaf.kind() {
+    if let Some(vis) = ctx.node_at_offset::<ast::Visibility>() {
+        return change_vis(ctx, vis);
+    }
+    add_vis(ctx)
+}
+
+fn add_vis(ctx: AssistCtx) -> Option<Assist> {
+    let item_keyword = ctx.leaf_at_offset().find(|leaf| match leaf.kind() {
         FN_KW | MOD_KW | STRUCT_KW | ENUM_KW | TRAIT_KW => true,
         _ => false,
-    }) {
+    });
+
+    let offset = if let Some(keyword) = item_keyword {
         let parent = keyword.parent()?;
         let def_kws = vec![FN_DEF, MODULE, STRUCT_DEF, ENUM_DEF, TRAIT_DEF];
         // Parent is not a definition, can't add visibility
@@ -37,6 +46,16 @@ pub fn change_visibility(ctx: AssistCtx) -> Option<Assist> {
     })
 }
 
+fn change_vis(ctx: AssistCtx, vis: ast::Visibility) -> Option<Assist> {
+    if vis.syntax().text() != "pub" {
+        return None;
+    }
+    ctx.build("chage to pub(crate)", |edit| {
+        edit.replace(vis.syntax().range(), "pub(crate)");
+        edit.set_cursor(vis.syntax().range().start());
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -85,4 +104,13 @@ mod tests {
             "struct S { <|>pub(crate) field: u32 }",
         )
     }
+
+    #[test]
+    fn change_visibility_pub_to_pub_crate() {
+        check_assist(
+            change_visibility,
+            "<|>pub fn foo() {}",
+            "<|>pub(crate) fn foo() {}",
+        )
+    }
 }