remove trailing return in trailing match expression

3 files changed, 72 insertions, 8 deletions

diff --git a/crates/hir-ty/src/diagnostics/expr.rs b/crates/hir-ty/src/diagnostics/expr.rs
index 50312af4f21..7f8fb7f4b52 100644
--- a/crates/hir-ty/src/diagnostics/expr.rs
+++ b/crates/hir-ty/src/diagnostics/expr.rs
@@ -272,6 +272,12 @@ impl ExprValidator {
                     self.check_for_trailing_return(*else_branch, body);
                 }
             }
+            Expr::Match { arms, .. } => {
+                for arm in arms.iter() {
+                    let MatchArm { expr, .. } = arm;
+                    self.check_for_trailing_return(*expr, body);
+                }
+            }
             Expr::Return { .. } => {
                 self.diagnostics.push(BodyValidationDiagnostic::RemoveTrailingReturn {
                     return_expr: body_expr,
diff --git a/crates/hir/src/diagnostics.rs b/crates/hir/src/diagnostics.rs
index cec6be8e892..b9f86f34156 100644
--- a/crates/hir/src/diagnostics.rs
+++ b/crates/hir/src/diagnostics.rs
@@ -11,7 +11,7 @@ use cfg::{CfgExpr, CfgOptions};
 use either::Either;
 use hir_def::{body::SyntheticSyntax, hir::ExprOrPatId, path::ModPath, AssocItemId, DefWithBodyId};
 use hir_expand::{name::Name, HirFileId, InFile};
-use syntax::{ast, AstPtr, SyntaxError, SyntaxNodePtr, TextRange};
+use syntax::{ast, AstNode, AstPtr, SyntaxError, SyntaxNodePtr, TextRange};
 
 use crate::{AssocItem, Field, Local, MacroKind, Trait, Type};
 
@@ -459,13 +459,16 @@ impl AnyDiagnostic {
             }
             BodyValidationDiagnostic::RemoveTrailingReturn { return_expr } => {
                 if let Ok(source_ptr) = source_map.expr_syntax(return_expr) {
-                    return Some(
-                        RemoveTrailingReturn {
-                            file_id: source_ptr.file_id,
-                            return_expr: source_ptr.value,
-                        }
-                        .into(),
-                    );
+                    // Filters out desugared return expressions (e.g. desugared try operators).
+                    if ast::ReturnExpr::can_cast(source_ptr.value.kind()) {
+                        return Some(
+                            RemoveTrailingReturn {
+                                file_id: source_ptr.file_id,
+                                return_expr: source_ptr.value,
+                            }
+                            .into(),
+                        );
+                    }
                 }
             }
             BodyValidationDiagnostic::RemoveUnnecessaryElse { if_expr } => {
diff --git a/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs b/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs
index ceef9da16c5..6c4ec7c132f 100644
--- a/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs
+++ b/crates/ide-diagnostics/src/handlers/remove_trailing_return.rs
@@ -148,6 +148,21 @@ fn foo(x: usize) -> u8 {
     }
 
     #[test]
+    fn remove_trailing_return_in_match() {
+        check_diagnostics(
+            r#"
+fn foo<T, E>(x: Result<T, E>) -> u8 {
+    match x {
+        Ok(_) => return 1,
+               //^^^^^^^^ 💡 weak: replace return <expr>; with <expr>
+        Err(_) => return 0,
+    }           //^^^^^^^^ 💡 weak: replace return <expr>; with <expr>
+}
+"#,
+        );
+    }
+
+    #[test]
     fn no_diagnostic_if_no_return_keyword() {
         check_diagnostics(
             r#"
@@ -319,4 +334,44 @@ fn foo(x: usize) -> u8 {
 "#,
         );
     }
+
+    #[test]
+    fn replace_in_match() {
+        check_fix(
+            r#"
+fn foo<T, E>(x: Result<T, E>) -> u8 {
+    match x {
+        Ok(_) => return$0 1,
+        Err(_) => 0,
+    }
+}
+"#,
+            r#"
+fn foo<T, E>(x: Result<T, E>) -> u8 {
+    match x {
+        Ok(_) => 1,
+        Err(_) => 0,
+    }
+}
+"#,
+        );
+        check_fix(
+            r#"
+fn foo<T, E>(x: Result<T, E>) -> u8 {
+    match x {
+        Ok(_) => 1,
+        Err(_) => return$0 0,
+    }
+}
+"#,
+            r#"
+fn foo<T, E>(x: Result<T, E>) -> u8 {
+    match x {
+        Ok(_) => 1,
+        Err(_) => 0,
+    }
+}
+"#,
+        );
+    }
 }