From ee44cf18621eb9c6ff2778cebd15b4c4ac3f1a8b Mon Sep 17 00:00:00 2001
From: davidsemakula <hello@davidsemakula.com>
Date: Wed, 31 Jan 2024 06:46:48 +0300
Subject: [PATCH] remove trailing return in trailing if expression

---
 crates/hir-ty/src/diagnostics/expr.rs         |  6 ++
 .../src/handlers/remove_trailing_return.rs    | 60 +++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/crates/hir-ty/src/diagnostics/expr.rs b/crates/hir-ty/src/diagnostics/expr.rs
index 2451cb30ca23..13e2ad782ecd 100644
--- a/crates/hir-ty/src/diagnostics/expr.rs
+++ b/crates/hir-ty/src/diagnostics/expr.rs
@@ -267,6 +267,12 @@ impl ExprValidator {
                     self.check_for_trailing_return(last_stmt, body);
                 }
             }
+            Expr::If { then_branch, else_branch, .. } => {
+                self.check_for_trailing_return(*then_branch, body);
+                if let Some(else_branch) = else_branch {
+                    self.check_for_trailing_return(*else_branch, body);
+                }
+            }
             Expr::Return { .. } => {
                 self.diagnostics.push(BodyValidationDiagnostic::RemoveTrailingReturn {
                     return_expr: body_expr,
diff --git a/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs b/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs
index 6cb5911096f7..ceef9da16c50 100644
--- a/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs
+++ b/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs
@@ -131,6 +131,22 @@ fn foo() -> u8 {
         );
     }
 
+    #[test]
+    fn remove_trailing_return_in_if() {
+        check_diagnostics(
+            r#"
+fn foo(x: usize) -> u8 {
+    if x > 0 {
+        return 1;
+      //^^^^^^^^^ 💡 weak: replace return <expr>; with <expr>
+    } else {
+        return 0;
+    } //^^^^^^^^^ 💡 weak: replace return <expr>; with <expr>
+}
+"#,
+        );
+    }
+
     #[test]
     fn no_diagnostic_if_no_return_keyword() {
         check_diagnostics(
@@ -256,6 +272,50 @@ fn foo() -> u8 {
     };
     bar()
 }
+"#,
+        );
+    }
+
+    #[test]
+    fn replace_in_if() {
+        check_fix(
+            r#"
+fn foo(x: usize) -> u8 {
+    if x > 0 {
+        return$0 1;
+    } else {
+        0
+    }
+}
+"#,
+            r#"
+fn foo(x: usize) -> u8 {
+    if x > 0 {
+        1
+    } else {
+        0
+    }
+}
+"#,
+        );
+        check_fix(
+            r#"
+fn foo(x: usize) -> u8 {
+    if x > 0 {
+        1
+    } else {
+        return$0 0;
+    }
+}
+"#,
+            r#"
+fn foo(x: usize) -> u8 {
+    if x > 0 {
+        1
+    } else {
+        0
+    }
+}
 "#,
         );
     }