diff --git a/crates/ra_assists/src/handlers/fill_match_arms.rs b/crates/ra_assists/src/handlers/fill_match_arms.rs index add82e5b1b4d..5f279d25ad3a 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 { } 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 2304e00cf96e..baf9a1b4b20d 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(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) -> 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 { 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) -> 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(node: &N) -> N { N::cast(remove_attrs_and_docs_inner(node.syntax().clone())).unwrap()