Merge #2226

2226: Use strongly-typed ast building for early-return assist r=matklad a=matklad



Co-authored-by: Aleksey Kladov <[email protected]>

3 files changed, 131 insertions, 70 deletions

diff --git a/crates/ra_assists/src/assists/early_return.rs b/crates/ra_assists/src/assists/early_return.rs
index 570a07a20..264412526 100644
--- a/crates/ra_assists/src/assists/early_return.rs
+++ b/crates/ra_assists/src/assists/early_return.rs
@@ -1,4 +1,4 @@
-use std::ops::RangeInclusive;
+use std::{iter::once, ops::RangeInclusive};
 
 use hir::db::HirDatabase;
 use ra_syntax::{
@@ -38,27 +38,30 @@ use crate::{
 // ```
 pub(crate) fn convert_to_guarded_return(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
     let if_expr: ast::IfExpr = ctx.find_node_at_offset()?;
+    if if_expr.else_branch().is_some() {
+        return None;
+    }
+
     let cond = if_expr.condition()?;
-    let mut if_let_ident: Option<String> = None;
 
     // Check if there is an IfLet that we can handle.
-    match cond.pat() {
-        None => {} // No IfLet, supported.
-        Some(TupleStructPat(ref pat)) if pat.args().count() == 1usize => match &pat.path() {
-            Some(p) => match p.qualifier() {
-                None => if_let_ident = Some(p.syntax().text().to_string()),
-                _ => return None,
-            },
-            _ => return None,
-        },
-        _ => return None, // Unsupported IfLet.
+    let if_let_pat = match cond.pat() {
+        None => None, // No IfLet, supported.
+        Some(TupleStructPat(pat)) if pat.args().count() == 1 => {
+            let path = pat.path()?;
+            match path.qualifier() {
+                None => {
+                    let bound_ident = pat.args().next().unwrap();
+                    Some((path, bound_ident))
+                }
+                Some(_) => return None,
+            }
+        }
+        Some(_) => return None, // Unsupported IfLet.
     };
 
-    let expr = cond.expr()?;
+    let cond_expr = cond.expr()?;
     let then_block = if_expr.then_branch()?.block()?;
-    if if_expr.else_branch().is_some() {
-        return None;
-    }
 
     let parent_block = if_expr.syntax().parent()?.ancestors().find_map(ast::Block::cast)?;
 
@@ -79,11 +82,11 @@ pub(crate) fn convert_to_guarded_return(ctx: AssistCtx<impl HirDatabase>) -> Opt
 
     let parent_container = parent_block.syntax().parent()?.parent()?;
 
-    let early_expression = match parent_container.kind() {
-        WHILE_EXPR | LOOP_EXPR => Some("continue"),
-        FN_DEF => Some("return"),
-        _ => None,
-    }?;
+    let early_expression: ast::Expr = match parent_container.kind() {
+        WHILE_EXPR | LOOP_EXPR => make::expr_continue().into(),
+        FN_DEF => make::expr_return().into(),
+        _ => return None,
+    };
 
     if then_block.syntax().first_child_or_token().map(|t| t.kind() == L_CURLY).is_none() {
         return None;
@@ -94,22 +97,43 @@ pub(crate) fn convert_to_guarded_return(ctx: AssistCtx<impl HirDatabase>) -> Opt
 
     ctx.add_assist(AssistId("convert_to_guarded_return"), "convert to guarded return", |edit| {
         let if_indent_level = IndentLevel::from_node(&if_expr.syntax());
-        let new_block = match if_let_ident {
+        let new_block = match if_let_pat {
             None => {
                 // If.
-                let early_expression = &(early_expression.to_owned() + ";");
-                let new_expr =
-                    if_indent_level.increase_indent(make::if_expression(&expr, early_expression));
-                replace(new_expr, &then_block, &parent_block, &if_expr)
+                let early_expression = &(early_expression.syntax().to_string() + ";");
+                let new_expr = if_indent_level
+                    .increase_indent(make::if_expression(&cond_expr, early_expression));
+                replace(new_expr.syntax(), &then_block, &parent_block, &if_expr)
             }
-            Some(if_let_ident) => {
+            Some((path, bound_ident)) => {
                 // If-let.
-                let new_expr = if_indent_level.increase_indent(make::let_match_early(
-                    expr,
-                    &if_let_ident,
-                    early_expression,
-                ));
-                replace(new_expr, &then_block, &parent_block, &if_expr)
+                let match_expr = {
+                    let happy_arm = make::match_arm(
+                        once(
+                            make::tuple_struct_pat(
+                                path,
+                                once(make::bind_pat(make::name("it")).into()),
+                            )
+                            .into(),
+                        ),
+                        make::expr_path(make::path_from_name_ref(make::name_ref("it"))).into(),
+                    );
+
+                    let sad_arm = make::match_arm(
+                        // FIXME: would be cool to use `None` or `Err(_)` if appropriate
+                        once(make::placeholder_pat().into()),
+                        early_expression.into(),
+                    );
+
+                    make::expr_match(cond_expr, make::match_arm_list(vec![happy_arm, sad_arm]))
+                };
+
+                let let_stmt = make::let_stmt(
+                    make::bind_pat(make::name(&bound_ident.syntax().to_string())).into(),
+                    Some(match_expr.into()),
+                );
+                let let_stmt = if_indent_level.increase_indent(let_stmt);
+                replace(let_stmt.syntax(), &then_block, &parent_block, &if_expr)
             }
         };
         edit.target(if_expr.syntax().text_range());
@@ -117,7 +141,7 @@ pub(crate) fn convert_to_guarded_return(ctx: AssistCtx<impl HirDatabase>) -> Opt
         edit.set_cursor(cursor_position);
 
         fn replace(
-            new_expr: impl AstNode,
+            new_expr: &SyntaxNode,
             then_block: &Block,
             parent_block: &Block,
             if_expr: &ast::IfExpr,
@@ -130,7 +154,7 @@ pub(crate) fn convert_to_guarded_return(ctx: AssistCtx<impl HirDatabase>) -> Opt
                 } else {
                     end_of_then
                 };
-            let mut then_statements = new_expr.syntax().children_with_tokens().chain(
+            let mut then_statements = new_expr.children_with_tokens().chain(
                 then_block_items
                     .syntax()
                     .children_with_tokens()
@@ -151,9 +175,10 @@ pub(crate) fn convert_to_guarded_return(ctx: AssistCtx<impl HirDatabase>) -> Opt
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use crate::helpers::{check_assist, check_assist_not_applicable};
 
+    use super::*;
+
     #[test]
     fn convert_inside_fn() {
         check_assist(
@@ -204,7 +229,7 @@ mod tests {
                 bar();
                 le<|>t n = match n {
                     Some(it) => it,
-                    None => return,
+                    _ => return,
                 };
                 foo(n);
 
@@ -216,6 +241,29 @@ mod tests {
     }
 
     #[test]
+    fn convert_if_let_result() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                if<|> let Ok(x) = Err(92) {
+                    foo(x);
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                le<|>t x = match Err(92) {
+                    Ok(it) => it,
+                    _ => return,
+                };
+                foo(x);
+            }
+            "#,
+        );
+    }
+
+    #[test]
     fn convert_let_ok_inside_fn() {
         check_assist(
             convert_to_guarded_return,
@@ -235,7 +283,7 @@ mod tests {
                 bar();
                 le<|>t n = match n {
                     Ok(it) => it,
-                    None => return,
+                    _ => return,
                 };
                 foo(n);
 
@@ -293,7 +341,7 @@ mod tests {
                 while true {
                     le<|>t n = match n {
                         Some(it) => it,
-                        None => continue,
+                        _ => continue,
                     };
                     foo(n);
                     bar();
@@ -350,7 +398,7 @@ mod tests {
                 loop {
                     le<|>t n = match n {
                         Some(it) => it,
-                        None => continue,
+                        _ => continue,
                     };
                     foo(n);
                     bar();
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
index 47bdbb81a..6f005a2d8 100644
--- a/crates/ra_syntax/src/ast/edit.rs
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -358,7 +358,7 @@ fn replace_children<N: AstNode>(
 fn test_increase_indent() {
     let arm_list = {
         let arm = make::match_arm(iter::once(make::placeholder_pat().into()), make::expr_unit());
-        make::match_arm_list(vec![arm.clone(), arm].into_iter())
+        make::match_arm_list(vec![arm.clone(), arm])
     };
     assert_eq!(
         arm_list.syntax().to_string(),
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index 95062ef6c..9749327fa 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -4,6 +4,10 @@ use itertools::Itertools;
 
 use crate::{ast, AstNode, SourceFile};
 
+pub fn name(text: &str) -> ast::Name {
+    ast_from_text(&format!("mod {};", text))
+}
+
 pub fn name_ref(text: &str) -> ast::NameRef {
     ast_from_text(&format!("fn f() {{ {}; }}", text))
 }
@@ -43,6 +47,21 @@ pub fn expr_unit() -> ast::Expr {
 pub fn expr_unimplemented() -> ast::Expr {
     expr_from_text("unimplemented!()")
 }
+pub fn expr_path(path: ast::Path) -> ast::Expr {
+    expr_from_text(&path.syntax().to_string())
+}
+pub fn expr_continue() -> ast::Expr {
+    expr_from_text("continue")
+}
+pub fn expr_break() -> ast::Expr {
+    expr_from_text("break")
+}
+pub fn expr_return() -> ast::Expr {
+    expr_from_text("return")
+}
+pub fn expr_match(expr: ast::Expr, match_arm_list: ast::MatchArmList) -> ast::Expr {
+    expr_from_text(&format!("match {} {}", expr.syntax(), match_arm_list.syntax()))
+}
 fn expr_from_text(text: &str) -> ast::Expr {
     ast_from_text(&format!("const C: () = {};", text))
 }
@@ -65,9 +84,9 @@ pub fn placeholder_pat() -> ast::PlaceholderPat {
 
 pub fn tuple_struct_pat(
     path: ast::Path,
-    pats: impl Iterator<Item = ast::Pat>,
+    pats: impl IntoIterator<Item = ast::Pat>,
 ) -> ast::TupleStructPat {
-    let pats_str = pats.map(|p| p.syntax().to_string()).join(", ");
+    let pats_str = pats.into_iter().map(|p| p.syntax().to_string()).join(", ");
     return from_text(&format!("{}({})", path.syntax(), pats_str));
 
     fn from_text(text: &str) -> ast::TupleStructPat {
@@ -75,8 +94,8 @@ pub fn tuple_struct_pat(
     }
 }
 
-pub fn record_pat(path: ast::Path, pats: impl Iterator<Item = ast::Pat>) -> ast::RecordPat {
-    let pats_str = pats.map(|p| p.syntax().to_string()).join(", ");
+pub fn record_pat(path: ast::Path, pats: impl IntoIterator<Item = ast::Pat>) -> ast::RecordPat {
+    let pats_str = pats.into_iter().map(|p| p.syntax().to_string()).join(", ");
     return from_text(&format!("{} {{ {} }}", path.syntax(), pats_str));
 
     fn from_text(text: &str) -> ast::RecordPat {
@@ -92,8 +111,8 @@ pub fn path_pat(path: ast::Path) -> ast::PathPat {
     }
 }
 
-pub fn match_arm(pats: impl Iterator<Item = ast::Pat>, expr: ast::Expr) -> ast::MatchArm {
-    let pats_str = pats.map(|p| p.syntax().to_string()).join(" | ");
+pub fn match_arm(pats: impl IntoIterator<Item = ast::Pat>, expr: ast::Expr) -> ast::MatchArm {
+    let pats_str = pats.into_iter().map(|p| p.syntax().to_string()).join(" | ");
     return from_text(&format!("{} => {}", pats_str, expr.syntax()));
 
     fn from_text(text: &str) -> ast::MatchArm {
@@ -101,8 +120,8 @@ pub fn match_arm(pats: impl Iterator<Item = ast::Pat>, expr: ast::Expr) -> ast::
     }
 }
 
-pub fn match_arm_list(arms: impl Iterator<Item = ast::MatchArm>) -> ast::MatchArmList {
-    let arms_str = arms.map(|arm| format!("\n    {}", arm.syntax())).join(",");
+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));
 
     fn from_text(text: &str) -> ast::MatchArmList {
@@ -110,25 +129,11 @@ pub fn match_arm_list(arms: impl Iterator<Item = ast::MatchArm>) -> ast::MatchAr
     }
 }
 
-pub fn let_match_early(expr: ast::Expr, path: &str, early_expression: &str) -> ast::LetStmt {
-    return from_text(&format!(
-        r#"let {} = match {} {{
-    {}(it) => it,
-    None => {},
-}};"#,
-        expr.syntax().text(),
-        expr.syntax().text(),
-        path,
-        early_expression
-    ));
-
-    fn from_text(text: &str) -> ast::LetStmt {
-        ast_from_text(&format!("fn f() {{ {} }}", text))
-    }
-}
-
-pub fn where_pred(path: ast::Path, bounds: impl Iterator<Item = ast::TypeBound>) -> ast::WherePred {
-    let bounds = bounds.map(|b| b.syntax().to_string()).join(" + ");
+pub fn where_pred(
+    path: ast::Path,
+    bounds: impl IntoIterator<Item = ast::TypeBound>,
+) -> ast::WherePred {
+    let bounds = bounds.into_iter().map(|b| b.syntax().to_string()).join(" + ");
     return from_text(&format!("{}: {}", path.syntax(), bounds));
 
     fn from_text(text: &str) -> ast::WherePred {
@@ -136,8 +141,8 @@ pub fn where_pred(path: ast::Path, bounds: impl Iterator<Item = ast::TypeBound>)
     }
 }
 
-pub fn where_clause(preds: impl Iterator<Item = ast::WherePred>) -> ast::WhereClause {
-    let preds = preds.map(|p| p.syntax().to_string()).join(", ");
+pub fn where_clause(preds: impl IntoIterator<Item = ast::WherePred>) -> ast::WhereClause {
+    let preds = preds.into_iter().map(|p| p.syntax().to_string()).join(", ");
     return from_text(preds.as_str());
 
     fn from_text(text: &str) -> ast::WhereClause {
@@ -153,6 +158,14 @@ pub fn if_expression(condition: &ast::Expr, statement: &str) -> ast::IfExpr {
     ))
 }
 
+pub fn let_stmt(pattern: ast::Pat, initializer: Option<ast::Expr>) -> ast::LetStmt {
+    let text = match initializer {
+        Some(it) => format!("let {} = {};", pattern.syntax(), it.syntax()),
+        None => format!("let {};", pattern.syntax()),
+    };
+    ast_from_text(&format!("fn f() {{ {} }}", text))
+}
+
 fn ast_from_text<N: AstNode>(text: &str) -> N {
     let parse = SourceFile::parse(text);
     let res = parse.tree().syntax().descendants().find_map(N::cast).unwrap();