From 69ffbe25a1ed473678ac142ad5ce852f82d041c6 Mon Sep 17 00:00:00 2001
From: Alex Vasilev <vaspiring@gmail.com>
Date: Sat, 21 Jan 2023 00:17:23 +0530
Subject: [PATCH] feat: add braces assist

---
 crates/ide-assists/src/handlers/add_braces.rs | 157 ++++++++++++++++++
 crates/ide-assists/src/lib.rs                 |   2 +
 2 files changed, 159 insertions(+)
 create mode 100644 crates/ide-assists/src/handlers/add_braces.rs

diff --git a/crates/ide-assists/src/handlers/add_braces.rs b/crates/ide-assists/src/handlers/add_braces.rs
new file mode 100644
index 00000000000..90bde1bab9b
--- /dev/null
+++ b/crates/ide-assists/src/handlers/add_braces.rs
@@ -0,0 +1,157 @@
+use syntax::{
+    ast::{self, edit::AstNodeEdit, make},
+    AstNode,
+};
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+enum ParentType {
+    MatchArmExpr,
+    ClosureExpr,
+}
+
+fn get_replacement_node(ctx: &AssistContext<'_>) -> Option<(ParentType, ast::Expr)> {
+    if let Some(match_arm) = ctx.find_node_at_offset::<ast::MatchArm>() {
+        let match_arm_expr = match_arm.syntax().children().find_map(ast::Expr::cast)?;
+
+        if matches!(match_arm_expr, ast::Expr::BlockExpr(_)) {
+            return None;
+        }
+
+        return Some((ParentType::MatchArmExpr, match_arm_expr));
+    } else if let Some(closure_expr) = ctx.find_node_at_offset::<ast::ClosureExpr>() {
+        let body = closure_expr.body()?;
+
+        if matches!(body, ast::Expr::BlockExpr(_)) {
+            return None;
+        }
+
+        return Some((ParentType::ClosureExpr, body));
+    }
+
+    None
+}
+
+// Assist: add_braces
+//
+// Adds braces to lamda and match arm expressions
+//
+// ```
+// fn foo(n: i32) -> i32 {
+//     match n {
+//         1 =>$0 n + 1,
+//         _ => 0
+//     }
+// }
+// ```
+// ->
+// ```
+// fn foo(n: i32) -> i32 {
+//     match n {
+//         1 => {
+//             n + 1
+//         },
+//         _ => 0
+//     }
+// }
+// ```
+pub(crate) fn add_braces(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let (expr_type, expr) = get_replacement_node(ctx)?;
+
+    acc.add(
+        AssistId("wrap_with_braces", AssistKind::RefactorRewrite),
+        match expr_type {
+            ParentType::ClosureExpr => "Add braces to lamda expression",
+            ParentType::MatchArmExpr => "Add braces to arm expression",
+        },
+        expr.syntax().text_range(),
+        |builder| {
+            let block_expr = AstNodeEdit::indent(
+                &make::block_expr(None, Some(expr.clone())),
+                AstNodeEdit::indent_level(&expr),
+            );
+
+            builder.replace(expr.syntax().text_range(), block_expr.syntax().text());
+        },
+    );
+
+    Some(())
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn suggest_add_braces_for_closure() {
+        check_assist(
+            add_braces,
+            r#"
+fn foo() {
+    t(|n|$0 n + 100);
+}
+"#,
+            r#"
+fn foo() {
+    t(|n| {
+        n + 100
+    });
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn no_assist_for_closures_with_braces() {
+        check_assist_not_applicable(
+            add_braces,
+            r#"
+fn foo() {
+    t(|n|$0 { n + 100 });
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn suggest_add_braces_for_match() {
+        check_assist(
+            add_braces,
+            r#"
+fn foo() {
+    match n {
+        Some(n) $0=> 29,
+        _ => ()
+    };
+}
+"#,
+            r#"
+fn foo() {
+    match n {
+        Some(n) => {
+            29
+        },
+        _ => ()
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn no_assist_for_match_with_braces() {
+        check_assist_not_applicable(
+            add_braces,
+            r#"
+fn foo() {
+    match n {
+        Some(n) $0=> { return 29; },
+        _ => ()
+    };
+}
+"#,
+        );
+    }
+}
diff --git a/crates/ide-assists/src/lib.rs b/crates/ide-assists/src/lib.rs
index 546ef96260f..276cf5f5dd0 100644
--- a/crates/ide-assists/src/lib.rs
+++ b/crates/ide-assists/src/lib.rs
@@ -106,6 +106,7 @@ mod handlers {
 
     pub(crate) type Handler = fn(&mut Assists, &AssistContext<'_>) -> Option<()>;
 
+    mod add_braces;
     mod add_explicit_type;
     mod add_label_to_loop;
     mod add_lifetime_to_type;
@@ -209,6 +210,7 @@ mod handlers {
     pub(crate) fn all() -> &'static [Handler] {
         &[
             // These are alphabetic for the foolish consistency
+            add_braces::add_braces,
             add_explicit_type::add_explicit_type,
             add_label_to_loop::add_label_to_loop,
             add_missing_match_arms::add_missing_match_arms,