Align set_visibility with the rest of the set_ edit-in-place methods

Author: DropDemBits

crates/ide-assists/src/handlers/fix_visibility.rs

@@ -79,7 +79,7 @@ fn add_vis_to_referenced_module_def(acc: &mut Assists, ctx: &AssistContext<'_>)
         edit.edit_file(target_file);
 
         let vis_owner = edit.make_mut(vis_owner);
-        vis_owner.set_visibility(missing_visibility.clone_for_update());
+        vis_owner.set_visibility(Some(missing_visibility.clone_for_update()));
 
         if let Some((cap, vis)) = ctx.config.snippet_cap.zip(vis_owner.visibility()) {
             edit.add_tabstop_before(cap, vis);
@@ -131,7 +131,7 @@ fn add_vis_to_referenced_record_field(acc: &mut Assists, ctx: &AssistContext<'_>
         edit.edit_file(target_file);
 
         let vis_owner = edit.make_mut(vis_owner);
-        vis_owner.set_visibility(missing_visibility.clone_for_update());
+        vis_owner.set_visibility(Some(missing_visibility.clone_for_update()));
 
         if let Some((cap, vis)) = ctx.config.snippet_cap.zip(vis_owner.visibility()) {
             edit.add_tabstop_before(cap, vis);

crates/syntax/src/ast/edit_in_place.rs

@@ -1007,20 +1007,24 @@ impl ast::IdentPat {
 }
 
 pub trait HasVisibilityEdit: ast::HasVisibility {
-    fn set_visibility(&self, visibility: ast::Visibility) {
-        match self.visibility() {
-            Some(current_visibility) => {
-                ted::replace(current_visibility.syntax(), visibility.syntax())
-            }
-            None => {
-                let vis_before = self
-                    .syntax()
-                    .children_with_tokens()
-                    .find(|it| !matches!(it.kind(), WHITESPACE | COMMENT | ATTR))
-                    .unwrap_or_else(|| self.syntax().first_child_or_token().unwrap());
-
-                ted::insert(ted::Position::before(vis_before), visibility.syntax());
+    fn set_visibility(&self, visibility: Option<ast::Visibility>) {
+        if let Some(visibility) = visibility {
+            match self.visibility() {
+                Some(current_visibility) => {
+                    ted::replace(current_visibility.syntax(), visibility.syntax())
+                }
+                None => {
+                    let vis_before = self
+                        .syntax()
+                        .children_with_tokens()
+                        .find(|it| !matches!(it.kind(), WHITESPACE | COMMENT | ATTR))
+                        .unwrap_or_else(|| self.syntax().first_child_or_token().unwrap());
+
+                    ted::insert(ted::Position::before(vis_before), visibility.syntax());
+                }
             }
+        } else if let Some(visibility) = self.visibility() {
+            ted::remove(visibility.syntax());
         }
     }
 }