2 files changed, 207 insertions, 29 deletions

diff --git a/crates/ra_assists/src/assists/early_return.rs b/crates/ra_assists/src/assists/early_return.rs
index ad6c5695a..570a07a20 100644
--- a/crates/ra_assists/src/assists/early_return.rs
+++ b/crates/ra_assists/src/assists/early_return.rs
@@ -3,9 +3,10 @@ use std::ops::RangeInclusive;
 use hir::db::HirDatabase;
 use ra_syntax::{
     algo::replace_children,
-    ast::{self, edit::IndentLevel, make},
+    ast::{self, edit::IndentLevel, make, Block, Pat::TupleStructPat},
     AstNode,
     SyntaxKind::{FN_DEF, LOOP_EXPR, L_CURLY, R_CURLY, WHILE_EXPR, WHITESPACE},
+    SyntaxNode,
 };
 
 use crate::{
@@ -37,7 +38,23 @@ 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()?;
-    let expr = if_expr.condition()?.expr()?;
+    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 expr = cond.expr()?;
     let then_block = if_expr.then_branch()?.block()?;
     if if_expr.else_branch().is_some() {
         return None;
@@ -63,8 +80,8 @@ 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;"),
+        WHILE_EXPR | LOOP_EXPR => Some("continue"),
+        FN_DEF => Some("return"),
         _ => None,
     }?;
 
@@ -77,34 +94,58 @@ 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_if_expr =
-            if_indent_level.increase_indent(make::if_expression(&expr, early_expression));
-        let then_block_items = IndentLevel::from(1).decrease_indent(then_block.clone());
-        let end_of_then = then_block_items.syntax().last_child_or_token().unwrap();
-        let end_of_then =
-            if end_of_then.prev_sibling_or_token().map(|n| n.kind()) == Some(WHITESPACE) {
-                end_of_then.prev_sibling_or_token().unwrap()
-            } else {
-                end_of_then
-            };
-        let mut new_if_and_then_statements = new_if_expr.syntax().children_with_tokens().chain(
-            then_block_items
-                .syntax()
-                .children_with_tokens()
-                .skip(1)
-                .take_while(|i| *i != end_of_then),
-        );
-        let new_block = replace_children(
-            &parent_block.syntax(),
-            RangeInclusive::new(
-                if_expr.clone().syntax().clone().into(),
-                if_expr.syntax().clone().into(),
-            ),
-            &mut new_if_and_then_statements,
-        );
+        let new_block = match if_let_ident {
+            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)
+            }
+            Some(if_let_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)
+            }
+        };
         edit.target(if_expr.syntax().text_range());
         edit.replace_ast(parent_block, ast::Block::cast(new_block).unwrap());
         edit.set_cursor(cursor_position);
+
+        fn replace(
+            new_expr: impl AstNode,
+            then_block: &Block,
+            parent_block: &Block,
+            if_expr: &ast::IfExpr,
+        ) -> SyntaxNode {
+            let then_block_items = IndentLevel::from(1).decrease_indent(then_block.clone());
+            let end_of_then = then_block_items.syntax().last_child_or_token().unwrap();
+            let end_of_then =
+                if end_of_then.prev_sibling_or_token().map(|n| n.kind()) == Some(WHITESPACE) {
+                    end_of_then.prev_sibling_or_token().unwrap()
+                } else {
+                    end_of_then
+                };
+            let mut then_statements = new_expr.syntax().children_with_tokens().chain(
+                then_block_items
+                    .syntax()
+                    .children_with_tokens()
+                    .skip(1)
+                    .take_while(|i| *i != end_of_then),
+            );
+            replace_children(
+                &parent_block.syntax(),
+                RangeInclusive::new(
+                    if_expr.clone().syntax().clone().into(),
+                    if_expr.syntax().clone().into(),
+                ),
+                &mut then_statements,
+            )
+        }
     })
 }
 
@@ -144,6 +185,68 @@ mod tests {
     }
 
     #[test]
+    fn convert_let_inside_fn() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main(n: Option<String>) {
+                bar();
+                if<|> let Some(n) = n {
+                    foo(n);
+
+                    //comment
+                    bar();
+                }
+            }
+            "#,
+            r#"
+            fn main(n: Option<String>) {
+                bar();
+                le<|>t n = match n {
+                    Some(it) => it,
+                    None => return,
+                };
+                foo(n);
+
+                //comment
+                bar();
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn convert_let_ok_inside_fn() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main(n: Option<String>) {
+                bar();
+                if<|> let Ok(n) = n {
+                    foo(n);
+
+                    //comment
+                    bar();
+                }
+            }
+            "#,
+            r#"
+            fn main(n: Option<String>) {
+                bar();
+                le<|>t n = match n {
+                    Ok(it) => it,
+                    None => return,
+                };
+                foo(n);
+
+                //comment
+                bar();
+            }
+            "#,
+        );
+    }
+
+    #[test]
     fn convert_inside_while() {
         check_assist(
             convert_to_guarded_return,
@@ -172,6 +275,35 @@ mod tests {
     }
 
     #[test]
+    fn convert_let_inside_while() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                while true {
+                    if<|> let Some(n) = n {
+                        foo(n);
+                        bar();
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                while true {
+                    le<|>t n = match n {
+                        Some(it) => it,
+                        None => continue,
+                    };
+                    foo(n);
+                    bar();
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
     fn convert_inside_loop() {
         check_assist(
             convert_to_guarded_return,
@@ -200,6 +332,35 @@ mod tests {
     }
 
     #[test]
+    fn convert_let_inside_loop() {
+        check_assist(
+            convert_to_guarded_return,
+            r#"
+            fn main() {
+                loop {
+                    if<|> let Some(n) = n {
+                        foo(n);
+                        bar();
+                    }
+                }
+            }
+            "#,
+            r#"
+            fn main() {
+                loop {
+                    le<|>t n = match n {
+                        Some(it) => it,
+                        None => continue,
+                    };
+                    foo(n);
+                    bar();
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
     fn ignore_already_converted_if() {
         check_assist_not_applicable(
             convert_to_guarded_return,
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index 3d5f18bfa..95062ef6c 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -110,6 +110,23 @@ 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(" + ");
     return from_text(&format!("{}: {}", path.syntax(), bounds));