From a521702d9c6ff38500abcbad481c669a68ec2640 Mon Sep 17 00:00:00 2001
From: roife <roifewu@gmail.com>
Date: Tue, 15 Oct 2024 14:22:20 +0800
Subject: [PATCH] fix: autofix for missing wrapped unit in return expr

---
 .../src/handlers/type_mismatch.rs             | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/crates/ide-diagnostics/src/handlers/type_mismatch.rs b/crates/ide-diagnostics/src/handlers/type_mismatch.rs
index 5cce7c4aed..90f88d6705 100644
--- a/crates/ide-diagnostics/src/handlers/type_mismatch.rs
+++ b/crates/ide-diagnostics/src/handlers/type_mismatch.rs
@@ -130,6 +130,7 @@ fn add_missing_ok_or_some(
     if d.actual.is_unit() {
         if let Expr::BlockExpr(block) = &expr {
             if block.tail_expr().is_none() {
+                // Fix for forms like `fn foo() -> Result<(), String> {}`
                 let mut builder = TextEdit::builder();
                 let block_indent = block.indent_level();
 
@@ -156,6 +157,20 @@ fn add_missing_ok_or_some(
                 acc.push(fix("insert_wrapped_unit", &name, source_change, expr_range));
             }
             return Some(());
+        } else if let Expr::ReturnExpr(ret_expr) = &expr {
+            // Fix for forms like `fn foo() -> Result<(), String> { return; }`
+            if ret_expr.expr().is_none() {
+                let mut builder = TextEdit::builder();
+                builder
+                    .insert(ret_expr.syntax().text_range().end(), format!(" {variant_name}(())"));
+                let source_change = SourceChange::from_text_edit(
+                    expr_ptr.file_id.original_file(ctx.sema.db),
+                    builder.finish(),
+                );
+                let name = format!("Insert {variant_name}(()) as the return value");
+                acc.push(fix("insert_wrapped_unit", &name, source_change, expr_range));
+            }
+            return Some(());
         }
     }
 
@@ -603,6 +618,29 @@ fn foo() -> Result<(), ()> {
         );
     }
 
+    #[test]
+    fn test_wrapped_unit_as_return_expr() {
+        check_fix(
+            r#"
+//- minicore: result
+fn foo(b: bool) -> Result<(), String> {
+    if b {
+        return$0;
+    }
+
+    Err("oh dear".to_owned())
+}"#,
+            r#"
+fn foo(b: bool) -> Result<(), String> {
+    if b {
+        return Ok(());
+    }
+
+    Err("oh dear".to_owned())
+}"#,
+        );
+    }
+
     #[test]
     fn test_in_const_and_static() {
         check_fix(