fix panic with wrapping/unwrapping result return type assists

2 files changed, 40 insertions, 8 deletions

diff --git a/crates/ide-assists/src/handlers/unwrap_result_return_type.rs b/crates/ide-assists/src/handlers/unwrap_result_return_type.rs
index f235b554e61..03e6dfebebf 100644
--- a/crates/ide-assists/src/handlers/unwrap_result_return_type.rs
+++ b/crates/ide-assists/src/handlers/unwrap_result_return_type.rs
@@ -123,10 +123,8 @@ fn tail_cb_impl(acc: &mut Vec<ast::Expr>, e: &ast::Expr) {
                 for_each_tail_expr(&break_expr_arg, &mut |e| tail_cb_impl(acc, e))
             }
         }
-        Expr::ReturnExpr(ret_expr) => {
-            if let Some(ret_expr_arg) = &ret_expr.expr() {
-                for_each_tail_expr(ret_expr_arg, &mut |e| tail_cb_impl(acc, e));
-            }
+        Expr::ReturnExpr(_) => {
+            // all return expressions have already been handled by the walk loop
         }
         e => acc.push(e.clone()),
     }
@@ -801,6 +799,24 @@ fn foo() -> i32 {
     }
 
     #[test]
+    fn wrap_return_in_tail_position() {
+        check_assist(
+            unwrap_result_return_type,
+            r#"
+//- minicore: result
+fn foo(num: i32) -> $0Result<i32, String> {
+    return Ok(num)
+}
+"#,
+            r#"
+fn foo(num: i32) -> i32 {
+    return num
+}
+"#,
+        );
+    }
+
+    #[test]
     fn unwrap_result_return_type_simple_with_closure() {
         check_assist(
             unwrap_result_return_type,
diff --git a/crates/ide-assists/src/handlers/wrap_return_type_in_result.rs b/crates/ide-assists/src/handlers/wrap_return_type_in_result.rs
index 61e9bcdcc51..b68ed00f772 100644
--- a/crates/ide-assists/src/handlers/wrap_return_type_in_result.rs
+++ b/crates/ide-assists/src/handlers/wrap_return_type_in_result.rs
@@ -98,10 +98,8 @@ fn tail_cb_impl(acc: &mut Vec<ast::Expr>, e: &ast::Expr) {
                 for_each_tail_expr(&break_expr_arg, &mut |e| tail_cb_impl(acc, e))
             }
         }
-        Expr::ReturnExpr(ret_expr) => {
-            if let Some(ret_expr_arg) = &ret_expr.expr() {
-                for_each_tail_expr(ret_expr_arg, &mut |e| tail_cb_impl(acc, e));
-            }
+        Expr::ReturnExpr(_) => {
+            // all return expressions have already been handled by the walk loop
         }
         e => acc.push(e.clone()),
     }
@@ -733,6 +731,24 @@ fn foo() -> Result<i32, ${0:_}> {
     }
 
     #[test]
+    fn wrap_return_in_tail_position() {
+        check_assist(
+            wrap_return_type_in_result,
+            r#"
+//- minicore: result
+fn foo(num: i32) -> $0i32 {
+    return num
+}
+"#,
+            r#"
+fn foo(num: i32) -> Result<i32, ${0:_}> {
+    return Ok(num)
+}
+"#,
+        );
+    }
+
+    #[test]
     fn wrap_return_type_in_result_simple_with_closure() {
         check_assist(
             wrap_return_type_in_result,