Merge #7097

7097: Add fix to wrap return expression in Some r=matklad a=theotherphil

Fixes https://github.com/rust-analyzer/rust-analyzer/issues/7095

Co-authored-by: Phil Ellison <[email protected]>

7 files changed, 90 insertions, 20 deletions

diff --git a/crates/hir/src/diagnostics.rs b/crates/hir/src/diagnostics.rs
index b1c924167..447faa04f 100644
--- a/crates/hir/src/diagnostics.rs
+++ b/crates/hir/src/diagnostics.rs
@@ -4,6 +4,6 @@ pub use hir_expand::diagnostics::{
     Diagnostic, DiagnosticCode, DiagnosticSink, DiagnosticSinkBuilder,
 };
 pub use hir_ty::diagnostics::{
-    IncorrectCase, MismatchedArgCount, MissingFields, MissingMatchArms, MissingOkInTailExpr,
+    IncorrectCase, MismatchedArgCount, MissingFields, MissingMatchArms, MissingOkOrSomeInTailExpr,
     NoSuchField, RemoveThisSemicolon,
 };
diff --git a/crates/hir_def/src/path.rs b/crates/hir_def/src/path.rs
index e2bf85bbc..3dd7c3cbb 100644
--- a/crates/hir_def/src/path.rs
+++ b/crates/hir_def/src/path.rs
@@ -305,6 +305,7 @@ pub use hir_expand::name as __name;
 macro_rules! __known_path {
     (core::iter::IntoIterator) => {};
     (core::result::Result) => {};
+    (core::option::Option) => {};
     (core::ops::Range) => {};
     (core::ops::RangeFrom) => {};
     (core::ops::RangeFull) => {};
diff --git a/crates/hir_expand/src/name.rs b/crates/hir_expand/src/name.rs
index 2f44876a8..95d853b6d 100644
--- a/crates/hir_expand/src/name.rs
+++ b/crates/hir_expand/src/name.rs
@@ -164,6 +164,7 @@ pub mod known {
         future,
         result,
         boxed,
+        option,
         // Components of known path (type name)
         Iterator,
         IntoIterator,
@@ -172,6 +173,7 @@ pub mod known {
         Ok,
         Future,
         Result,
+        Option,
         Output,
         Target,
         Box,
diff --git a/crates/hir_ty/src/diagnostics.rs b/crates/hir_ty/src/diagnostics.rs
index 14e18f5a1..c67a289f2 100644
--- a/crates/hir_ty/src/diagnostics.rs
+++ b/crates/hir_ty/src/diagnostics.rs
@@ -186,9 +186,10 @@ impl Diagnostic for MissingMatchArms {
     }
 }
 
-// Diagnostic: missing-ok-in-tail-expr
+// Diagnostic: missing-ok-or-some-in-tail-expr
 //
-// This diagnostic is triggered if block that should return `Result` returns a value not wrapped in `Ok`.
+// This diagnostic is triggered if a block that should return `Result` returns a value not wrapped in `Ok`,
+// or if a block that should return `Option` returns a value not wrapped in `Some`.
 //
 // Example:
 //
@@ -198,17 +199,19 @@ impl Diagnostic for MissingMatchArms {
 // }
 // ```
 #[derive(Debug)]
-pub struct MissingOkInTailExpr {
+pub struct MissingOkOrSomeInTailExpr {
     pub file: HirFileId,
     pub expr: AstPtr<ast::Expr>,
+    // `Some` or `Ok` depending on whether the return type is Result or Option
+    pub required: String,
 }
 
-impl Diagnostic for MissingOkInTailExpr {
+impl Diagnostic for MissingOkOrSomeInTailExpr {
     fn code(&self) -> DiagnosticCode {
-        DiagnosticCode("missing-ok-in-tail-expr")
+        DiagnosticCode("missing-ok-or-some-in-tail-expr")
     }
     fn message(&self) -> String {
-        "wrap return expression in Ok".to_string()
+        format!("wrap return expression in {}", self.required)
     }
     fn display_source(&self) -> InFile<SyntaxNodePtr> {
         InFile { file_id: self.file, value: self.expr.clone().into() }
diff --git a/crates/hir_ty/src/diagnostics/expr.rs b/crates/hir_ty/src/diagnostics/expr.rs
index b4e453411..a1c484fdf 100644
--- a/crates/hir_ty/src/diagnostics/expr.rs
+++ b/crates/hir_ty/src/diagnostics/expr.rs
@@ -11,8 +11,8 @@ use crate::{
     db::HirDatabase,
     diagnostics::{
         match_check::{is_useful, MatchCheckCtx, Matrix, PatStack, Usefulness},
-        MismatchedArgCount, MissingFields, MissingMatchArms, MissingOkInTailExpr, MissingPatFields,
-        RemoveThisSemicolon,
+        MismatchedArgCount, MissingFields, MissingMatchArms, MissingOkOrSomeInTailExpr,
+        MissingPatFields, RemoveThisSemicolon,
     },
     utils::variant_data,
     ApplicationTy, InferenceResult, Ty, TypeCtor,
@@ -306,27 +306,40 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
         };
 
         let core_result_path = path![core::result::Result];
+        let core_option_path = path![core::option::Option];
 
         let resolver = self.owner.resolver(db.upcast());
         let core_result_enum = match resolver.resolve_known_enum(db.upcast(), &core_result_path) {
             Some(it) => it,
             _ => return,
         };
+        let core_option_enum = match resolver.resolve_known_enum(db.upcast(), &core_option_path) {
+            Some(it) => it,
+            _ => return,
+        };
 
         let core_result_ctor = TypeCtor::Adt(AdtId::EnumId(core_result_enum));
-        let params = match &mismatch.expected {
+        let core_option_ctor = TypeCtor::Adt(AdtId::EnumId(core_option_enum));
+
+        let (params, required) = match &mismatch.expected {
             Ty::Apply(ApplicationTy { ctor, parameters }) if ctor == &core_result_ctor => {
-                parameters
+                (parameters, "Ok".to_string())
+            }
+            Ty::Apply(ApplicationTy { ctor, parameters }) if ctor == &core_option_ctor => {
+                (parameters, "Some".to_string())
             }
             _ => return,
         };
 
-        if params.len() == 2 && params[0] == mismatch.actual {
+        if params.len() > 0 && params[0] == mismatch.actual {
             let (_, source_map) = db.body_with_source_map(self.owner.into());
 
             if let Ok(source_ptr) = source_map.expr_syntax(id) {
-                self.sink
-                    .push(MissingOkInTailExpr { file: source_ptr.file_id, expr: source_ptr.value });
+                self.sink.push(MissingOkOrSomeInTailExpr {
+                    file: source_ptr.file_id,
+                    expr: source_ptr.value,
+                    required,
+                });
             }
         }
     }
diff --git a/crates/ide/src/diagnostics.rs b/crates/ide/src/diagnostics.rs
index 6931a6190..055c0a79c 100644
--- a/crates/ide/src/diagnostics.rs
+++ b/crates/ide/src/diagnostics.rs
@@ -125,7 +125,7 @@ pub(crate) fn diagnostics(
         .on::<hir::diagnostics::MissingFields, _>(|d| {
             res.borrow_mut().push(diagnostic_with_fix(d, &sema));
         })
-        .on::<hir::diagnostics::MissingOkInTailExpr, _>(|d| {
+        .on::<hir::diagnostics::MissingOkOrSomeInTailExpr, _>(|d| {
             res.borrow_mut().push(diagnostic_with_fix(d, &sema));
         })
         .on::<hir::diagnostics::NoSuchField, _>(|d| {
@@ -305,6 +305,40 @@ mod tests {
     }
 
     #[test]
+    fn test_wrap_return_type_option() {
+        check_fix(
+            r#"
+//- /main.rs crate:main deps:core
+use core::option::Option::{self, Some, None};
+
+fn div(x: i32, y: i32) -> Option<i32> {
+    if y == 0 {
+        return None;
+    }
+    x / y$0
+}
+//- /core/lib.rs crate:core
+pub mod result {
+    pub enum Result<T, E> { Ok(T), Err(E) }
+}
+pub mod option {
+    pub enum Option<T> { Some(T), None }
+}
+"#,
+            r#"
+use core::option::Option::{self, Some, None};
+
+fn div(x: i32, y: i32) -> Option<i32> {
+    if y == 0 {
+        return None;
+    }
+    Some(x / y)
+}
+"#,
+        );
+    }
+
+    #[test]
     fn test_wrap_return_type() {
         check_fix(
             r#"
@@ -321,6 +355,9 @@ fn div(x: i32, y: i32) -> Result<i32, ()> {
 pub mod result {
     pub enum Result<T, E> { Ok(T), Err(E) }
 }
+pub mod option {
+    pub enum Option<T> { Some(T), None }
+}
 "#,
             r#"
 use core::result::Result::{self, Ok, Err};
@@ -352,6 +389,9 @@ fn div<T>(x: T) -> Result<T, i32> {
 pub mod result {
     pub enum Result<T, E> { Ok(T), Err(E) }
 }
+pub mod option {
+    pub enum Option<T> { Some(T), None }
+}
 "#,
             r#"
 use core::result::Result::{self, Ok, Err};
@@ -385,6 +425,9 @@ fn div(x: i32, y: i32) -> MyResult<i32> {
 pub mod result {
     pub enum Result<T, E> { Ok(T), Err(E) }
 }
+pub mod option {
+    pub enum Option<T> { Some(T), None }
+}
 "#,
             r#"
 use core::result::Result::{self, Ok, Err};
@@ -414,12 +457,15 @@ fn foo() -> Result<(), i32> { 0 }
 pub mod result {
     pub enum Result<T, E> { Ok(T), Err(E) }
 }
+pub mod option {
+    pub enum Option<T> { Some(T), None }
+}
 "#,
         );
     }
 
     #[test]
-    fn test_wrap_return_type_not_applicable_when_return_type_is_not_result() {
+    fn test_wrap_return_type_not_applicable_when_return_type_is_not_result_or_option() {
         check_no_diagnostics(
             r#"
 //- /main.rs crate:main deps:core
@@ -433,6 +479,9 @@ fn foo() -> SomeOtherEnum { 0 }
 pub mod result {
     pub enum Result<T, E> { Ok(T), Err(E) }
 }
+pub mod option {
+    pub enum Option<T> { Some(T), None }
+}
 "#,
         );
     }
diff --git a/crates/ide/src/diagnostics/fixes.rs b/crates/ide/src/diagnostics/fixes.rs
index 71ec4df92..d7ad88ed5 100644
--- a/crates/ide/src/diagnostics/fixes.rs
+++ b/crates/ide/src/diagnostics/fixes.rs
@@ -3,7 +3,7 @@
 use hir::{
     db::AstDatabase,
     diagnostics::{
-        Diagnostic, IncorrectCase, MissingFields, MissingOkInTailExpr, NoSuchField,
+        Diagnostic, IncorrectCase, MissingFields, MissingOkOrSomeInTailExpr, NoSuchField,
         RemoveThisSemicolon, UnresolvedModule,
     },
     HasSource, HirDisplay, InFile, Semantics, VariantDef,
@@ -94,15 +94,17 @@ impl DiagnosticWithFix for MissingFields {
     }
 }
 
-impl DiagnosticWithFix for MissingOkInTailExpr {
+impl DiagnosticWithFix for MissingOkOrSomeInTailExpr {
     fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
         let root = sema.db.parse_or_expand(self.file)?;
         let tail_expr = self.expr.to_node(&root);
         let tail_expr_range = tail_expr.syntax().text_range();
-        let edit = TextEdit::replace(tail_expr_range, format!("Ok({})", tail_expr.syntax()));
+        let replacement = format!("{}({})", self.required, tail_expr.syntax());
+        let edit = TextEdit::replace(tail_expr_range, replacement);
         let source_change =
             SourceFileEdit { file_id: self.file.original_file(sema.db), edit }.into();
-        Some(Fix::new("Wrap with ok", source_change, tail_expr_range))
+        let name = if self.required == "Ok" { "Wrap with Ok" } else { "Wrap with Some" };
+        Some(Fix::new(name, source_change, tail_expr_range))
     }
 }