4 files changed, 55 insertions, 32 deletions

diff --git a/crates/ra_assists/src/assists/replace_if_let_with_match.rs b/crates/ra_assists/src/assists/replace_if_let_with_match.rs
index c9b62e5ff..c8b13b7b3 100644
--- a/crates/ra_assists/src/assists/replace_if_let_with_match.rs
+++ b/crates/ra_assists/src/assists/replace_if_let_with_match.rs
@@ -1,9 +1,12 @@
-use format_buf::format;
 use hir::db::HirDatabase;
-use ra_fmt::extract_trivial_expression;
-use ra_syntax::{ast, AstNode};
+use ra_fmt::unwrap_trivial_block;
+use ra_syntax::{
+    ast::{self, make},
+    AstNode,
+};
 
 use crate::{Assist, AssistCtx, AssistId};
+use ast::edit::IndentLevel;
 
 // Assist: replace_if_let_with_match
 //
@@ -43,32 +46,24 @@ pub(crate) fn replace_if_let_with_match(ctx: AssistCtx<impl HirDatabase>) -> Opt
     };
 
     ctx.add_assist(AssistId("replace_if_let_with_match"), "Replace with match", |edit| {
-        let match_expr = build_match_expr(expr, pat, then_block, else_block);
-        edit.target(if_expr.syntax().text_range());
-        edit.replace_node_and_indent(if_expr.syntax(), match_expr);
-        edit.set_cursor(if_expr.syntax().text_range().start())
-    })
-}
+        let match_expr = {
+            let then_arm = {
+                let then_expr = unwrap_trivial_block(then_block);
+                make::match_arm(vec![pat], then_expr)
+            };
+            let else_arm = {
+                let else_expr = unwrap_trivial_block(else_block);
+                make::match_arm(vec![make::placeholder_pat().into()], else_expr)
+            };
+            make::expr_match(expr, make::match_arm_list(vec![then_arm, else_arm]))
+        };
 
-fn build_match_expr(
-    expr: ast::Expr,
-    pat1: ast::Pat,
-    arm1: ast::BlockExpr,
-    arm2: ast::BlockExpr,
-) -> String {
-    let mut buf = String::new();
-    format!(buf, "match {} {{\n", expr.syntax().text());
-    format!(buf, "    {} => {}\n", pat1.syntax().text(), format_arm(&arm1));
-    format!(buf, "    _ => {}\n", format_arm(&arm2));
-    buf.push_str("}");
-    buf
-}
+        let match_expr = IndentLevel::from_node(if_expr.syntax()).increase_indent(match_expr);
 
-fn format_arm(block: &ast::BlockExpr) -> String {
-    match extract_trivial_expression(block) {
-        Some(e) if !e.syntax().text().contains_char('\n') => format!("{},", e.syntax().text()),
-        _ => block.syntax().text().to_string(),
-    }
+        edit.target(if_expr.syntax().text_range());
+        edit.set_cursor(if_expr.syntax().text_range().start());
+        edit.replace_ast::<ast::Expr>(if_expr.into(), match_expr.into());
+    })
 }
 
 #[cfg(test)]
diff --git a/crates/ra_fmt/src/lib.rs b/crates/ra_fmt/src/lib.rs
index 10f592257..4bca27b5c 100644
--- a/crates/ra_fmt/src/lib.rs
+++ b/crates/ra_fmt/src/lib.rs
@@ -35,8 +35,14 @@ fn prev_tokens(token: SyntaxToken) -> impl Iterator<Item = SyntaxToken> {
     successors(token.prev_token(), |token| token.prev_token())
 }
 
-pub fn extract_trivial_expression(expr: &ast::BlockExpr) -> Option<ast::Expr> {
-    let block = expr.block()?;
+pub fn unwrap_trivial_block(block: ast::BlockExpr) -> ast::Expr {
+    extract_trivial_expression(&block)
+        .filter(|expr| !expr.syntax().text().contains_char('\n'))
+        .unwrap_or_else(|| block.into())
+}
+
+pub fn extract_trivial_expression(block: &ast::BlockExpr) -> Option<ast::Expr> {
+    let block = block.block()?;
     let expr = block.expr()?;
     let non_trivial_children = block.syntax().children().filter(|it| match it.kind() {
         WHITESPACE | T!['{'] | T!['}'] => false,
diff --git a/crates/ra_syntax/src/ast/expr_extensions.rs b/crates/ra_syntax/src/ast/expr_extensions.rs
index 539759450..2e50a095c 100644
--- a/crates/ra_syntax/src/ast/expr_extensions.rs
+++ b/crates/ra_syntax/src/ast/expr_extensions.rs
@@ -7,6 +7,21 @@ use crate::{
     SyntaxToken, T,
 };
 
+impl ast::Expr {
+    pub fn is_block_like(&self) -> bool {
+        match self {
+            ast::Expr::IfExpr(_)
+            | ast::Expr::LoopExpr(_)
+            | ast::Expr::ForExpr(_)
+            | ast::Expr::WhileExpr(_)
+            | ast::Expr::BlockExpr(_)
+            | ast::Expr::MatchExpr(_)
+            | ast::Expr::TryBlockExpr(_) => true,
+            _ => false,
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ElseBranch {
     Block(ast::BlockExpr),
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index 38c0e9a66..629503dc5 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -122,11 +122,18 @@ pub fn match_arm(pats: impl IntoIterator<Item = ast::Pat>, expr: ast::Expr) -> a
 }
 
 pub fn match_arm_list(arms: impl IntoIterator<Item = ast::MatchArm>) -> ast::MatchArmList {
-    let arms_str = arms.into_iter().map(|arm| format!("\n    {}", arm.syntax())).join(",");
-    return from_text(&format!("{},\n", arms_str));
+    let arms_str = arms
+        .into_iter()
+        .map(|arm| {
+            let needs_comma = arm.expr().map_or(true, |it| !it.is_block_like());
+            let comma = if needs_comma { "," } else { "" };
+            format!("    {}{}\n", arm.syntax(), comma)
+        })
+        .collect::<String>();
+    return from_text(&format!("{}", arms_str));
 
     fn from_text(text: &str) -> ast::MatchArmList {
-        ast_from_text(&format!("fn f() {{ match () {{{}}} }}", text))
+        ast_from_text(&format!("fn f() {{ match () {{\n{}}} }}", text))
     }
 }