From 7306504b82cfc1d39f44c3bd59aceca483e37eac Mon Sep 17 00:00:00 2001
From: Ryan Mehri <ryan.mehri1@gmail.com>
Date: Sun, 24 Sep 2023 15:39:12 -0700
Subject: [PATCH] fix panic with wrapping/unwrapping result return type assists

---
 .../src/handlers/unwrap_result_return_type.rs | 24 +++++++++++++++----
 .../handlers/wrap_return_type_in_result.rs    | 24 +++++++++++++++----
 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 f235b554e61f4..03e6dfebebfb2 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()),
     }
@@ -800,6 +798,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(
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 61e9bcdcc5162..b68ed00f77210 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()),
     }
@@ -732,6 +730,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(