2 files changed, 164 insertions, 35 deletions

diff --git a/crates/ra_assists/src/handlers/fill_match_arms.rs b/crates/ra_assists/src/handlers/fill_match_arms.rs
index add82e5b1..5f279d25a 100644
--- a/crates/ra_assists/src/handlers/fill_match_arms.rs
+++ b/crates/ra_assists/src/handlers/fill_match_arms.rs
@@ -7,7 +7,7 @@ use itertools::Itertools;
 use ra_ide_db::RootDatabase;
 
 use crate::{Assist, AssistCtx, AssistId};
-use ra_syntax::ast::{self, edit::IndentLevel, make, AstNode, NameOwner};
+use ra_syntax::ast::{self, make, AstNode, NameOwner};
 
 use ast::{MatchArm, Pat};
 
@@ -97,10 +97,8 @@ pub(crate) fn fill_match_arms(ctx: AssistCtx) -> Option<Assist> {
     }
 
     ctx.add_assist(AssistId("fill_match_arms"), "Fill match arms", |edit| {
-        arms.extend(missing_arms);
-
-        let indent_level = IndentLevel::from_node(match_arm_list.syntax());
-        let new_arm_list = indent_level.increase_indent(make::match_arm_list(arms));
+        let new_arm_list =
+            match_arm_list.remove_placeholder().append_arms(missing_arms.into_iter());
 
         edit.target(match_expr.syntax().text_range());
         edit.set_cursor(expr.syntax().text_range().start());
@@ -655,4 +653,69 @@ mod tests {
             "#,
         );
     }
+
+    #[test]
+    fn fill_match_arms_preserves_comments() {
+        check_assist(
+            fill_match_arms,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            fn foo(a: A) {
+                match a {
+                    // foo bar baz<|>
+                    A::One => {}
+                    // This is where the rest should be
+                }
+            }
+            "#,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            fn foo(a: A) {
+                match <|>a {
+                    // foo bar baz
+                    A::One => {}
+                    // This is where the rest should be
+                    A::Two => {}
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn fill_match_arms_preserves_comments_empty() {
+        check_assist(
+            fill_match_arms,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            fn foo(a: A) {
+                match a {
+                    // foo bar baz<|>
+                }
+            }
+            "#,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            fn foo(a: A) {
+                match <|>a {
+                    // foo bar baz
+                    A::One => {}
+                    A::Two => {}
+                }
+            }
+            "#,
+        );
+    }
 }
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
index 2304e00cf..baf9a1b4b 100644
--- a/crates/ra_syntax/src/ast/edit.rs
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -46,12 +46,44 @@ impl ast::FnDef {
     }
 }
 
+fn make_multiline<N>(node: N) -> N
+where
+    N: AstNode + Clone,
+{
+    let l_curly = match node.syntax().children_with_tokens().find(|it| it.kind() == T!['{']) {
+        Some(it) => it,
+        None => return node,
+    };
+    let sibling = match l_curly.next_sibling_or_token() {
+        Some(it) => it,
+        None => return node,
+    };
+    let existing_ws = match sibling.as_token() {
+        None => None,
+        Some(tok) if tok.kind() != WHITESPACE => None,
+        Some(ws) => {
+            if ws.text().contains('\n') {
+                return node;
+            }
+            Some(ws.clone())
+        }
+    };
+
+    let indent = leading_indent(node.syntax()).unwrap_or_default();
+    let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
+    let to_insert = iter::once(ws.ws().into());
+    match existing_ws {
+        None => node.insert_children(InsertPosition::After(l_curly), to_insert),
+        Some(ws) => node.replace_children(single_node(ws), to_insert),
+    }
+}
+
 impl ast::ItemList {
     #[must_use]
     pub fn append_items(&self, items: impl Iterator<Item = ast::ImplItem>) -> ast::ItemList {
         let mut res = self.clone();
         if !self.syntax().text().contains_char('\n') {
-            res = res.make_multiline();
+            res = make_multiline(res);
         }
         items.for_each(|it| res = res.append_item(it));
         res
@@ -81,35 +113,6 @@ impl ast::ItemList {
     fn l_curly(&self) -> Option<SyntaxElement> {
         self.syntax().children_with_tokens().find(|it| it.kind() == T!['{'])
     }
-
-    fn make_multiline(&self) -> ast::ItemList {
-        let l_curly = match self.syntax().children_with_tokens().find(|it| it.kind() == T!['{']) {
-            Some(it) => it,
-            None => return self.clone(),
-        };
-        let sibling = match l_curly.next_sibling_or_token() {
-            Some(it) => it,
-            None => return self.clone(),
-        };
-        let existing_ws = match sibling.as_token() {
-            None => None,
-            Some(tok) if tok.kind() != WHITESPACE => None,
-            Some(ws) => {
-                if ws.text().contains('\n') {
-                    return self.clone();
-                }
-                Some(ws.clone())
-            }
-        };
-
-        let indent = leading_indent(self.syntax()).unwrap_or_default();
-        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
-        let to_insert = iter::once(ws.ws().into());
-        match existing_ws {
-            None => self.insert_children(InsertPosition::After(l_curly), to_insert),
-            Some(ws) => self.replace_children(single_node(ws), to_insert),
-        }
-    }
 }
 
 impl ast::RecordFieldList {
@@ -334,6 +337,69 @@ impl ast::UseTree {
     }
 }
 
+impl ast::MatchArmList {
+    #[must_use]
+    pub fn append_arms(&self, items: impl Iterator<Item = ast::MatchArm>) -> ast::MatchArmList {
+        let mut res = self.clone();
+        res = res.strip_if_only_whitespace();
+        if !res.syntax().text().contains_char('\n') {
+            res = make_multiline(res);
+        }
+        items.for_each(|it| res = res.append_arm(it));
+        res
+    }
+
+    fn strip_if_only_whitespace(&self) -> ast::MatchArmList {
+        let mut iter = self.syntax().children_with_tokens().skip_while(|it| it.kind() != T!['{']);
+        iter.next(); // Eat the curly
+        let mut inner = iter.take_while(|it| it.kind() != T!['}']);
+        if !inner.clone().all(|it| it.kind() == WHITESPACE) {
+            return self.clone();
+        }
+        let start = match inner.next() {
+            Some(s) => s,
+            None => return self.clone(),
+        };
+        let end = match inner.last() {
+            Some(s) => s,
+            None => start.clone(),
+        };
+        self.replace_children(start..=end, &mut iter::empty())
+    }
+
+    #[must_use]
+    pub fn remove_placeholder(&self) -> ast::MatchArmList {
+        let placeholder = self.arms().find(|arm| {
+            if let Some(ast::Pat::PlaceholderPat(_)) = arm.pat() {
+                return true;
+            }
+            false
+        });
+        if let Some(placeholder) = placeholder {
+            let s: SyntaxElement = placeholder.syntax().clone().into();
+            let e = s.clone();
+            self.replace_children(s..=e, &mut iter::empty())
+        } else {
+            self.clone()
+        }
+    }
+
+    #[must_use]
+    pub fn append_arm(&self, item: ast::MatchArm) -> ast::MatchArmList {
+        let r_curly = match self.syntax().children_with_tokens().find(|it| it.kind() == T!['}']) {
+            Some(t) => t,
+            None => return self.clone(),
+        };
+        let position = InsertPosition::Before(r_curly.into());
+        let arm_ws = tokens::WsBuilder::new("    ");
+        let match_indent = &leading_indent(self.syntax()).unwrap_or_default();
+        let match_ws = tokens::WsBuilder::new(&format!("\n{}", match_indent));
+        let to_insert: ArrayVec<[SyntaxElement; 3]> =
+            [arm_ws.ws().into(), item.syntax().clone().into(), match_ws.ws().into()].into();
+        self.insert_children(position, to_insert)
+    }
+}
+
 #[must_use]
 pub fn remove_attrs_and_docs<N: ast::AttrsOwner>(node: &N) -> N {
     N::cast(remove_attrs_and_docs_inner(node.syntax().clone())).unwrap()