5 files changed, 44 insertions, 58 deletions

diff --git a/crates/hir/src/diagnostics.rs b/crates/hir/src/diagnostics.rs
index b144bb335..9afee0b90 100644
--- a/crates/hir/src/diagnostics.rs
+++ b/crates/hir/src/diagnostics.rs
@@ -37,6 +37,7 @@ diagnostics![
     MacroError,
     MismatchedArgCount,
     MissingFields,
+    MissingOkOrSomeInTailExpr,
     MissingUnsafe,
     NoSuchField,
     RemoveThisSemicolon,
@@ -157,41 +158,13 @@ pub struct RemoveThisSemicolon {
     pub expr: InFile<AstPtr<ast::Expr>>,
 }
 
-// Diagnostic: missing-ok-or-some-in-tail-expr
-//
-// 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:
-//
-// ```rust
-// fn foo() -> Result<u8, ()> {
-//     10
-// }
-// ```
 #[derive(Debug)]
 pub struct MissingOkOrSomeInTailExpr {
-    pub file: HirFileId,
-    pub expr: AstPtr<ast::Expr>,
+    pub expr: InFile<AstPtr<ast::Expr>>,
     // `Some` or `Ok` depending on whether the return type is Result or Option
     pub required: String,
 }
 
-impl Diagnostic for MissingOkOrSomeInTailExpr {
-    fn code(&self) -> DiagnosticCode {
-        DiagnosticCode("missing-ok-or-some-in-tail-expr")
-    }
-    fn message(&self) -> String {
-        format!("wrap return expression in {}", self.required)
-    }
-    fn display_source(&self) -> InFile<SyntaxNodePtr> {
-        InFile { file_id: self.file, value: self.expr.clone().into() }
-    }
-    fn as_any(&self) -> &(dyn Any + Send + 'static) {
-        self
-    }
-}
-
 // Diagnostic: missing-match-arm
 //
 // This diagnostic is triggered if `match` block is missing one or more match arms.
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index cb9bf60b8..aaab5336a 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -1190,11 +1190,7 @@ impl Function {
                 }
                 BodyValidationDiagnostic::MissingOkOrSomeInTailExpr { expr, required } => {
                     match source_map.expr_syntax(expr) {
-                        Ok(source_ptr) => sink.push(MissingOkOrSomeInTailExpr {
-                            file: source_ptr.file_id,
-                            expr: source_ptr.value,
-                            required,
-                        }),
+                        Ok(expr) => acc.push(MissingOkOrSomeInTailExpr { expr, required }.into()),
                         Err(SyntheticSyntax) => (),
                     }
                 }
diff --git a/crates/ide/src/diagnostics.rs b/crates/ide/src/diagnostics.rs
index 3ced08f30..af282db0c 100644
--- a/crates/ide/src/diagnostics.rs
+++ b/crates/ide/src/diagnostics.rs
@@ -9,6 +9,7 @@ mod inactive_code;
 mod macro_error;
 mod mismatched_arg_count;
 mod missing_fields;
+mod missing_ok_or_some_in_tail_expr;
 mod missing_unsafe;
 mod no_such_field;
 mod remove_this_semicolon;
@@ -163,9 +164,6 @@ pub(crate) fn diagnostics(
     }
     let res = RefCell::new(res);
     let sink_builder = DiagnosticSinkBuilder::new()
-        .on::<hir::diagnostics::MissingOkOrSomeInTailExpr, _>(|d| {
-            res.borrow_mut().push(diagnostic_with_fix(d, &sema, resolve));
-        })
         .on::<hir::diagnostics::IncorrectCase, _>(|d| {
             res.borrow_mut().push(warning_with_fix(d, &sema, resolve));
         })
@@ -223,6 +221,7 @@ pub(crate) fn diagnostics(
             AnyDiagnostic::MacroError(d) => macro_error::macro_error(&ctx, &d),
             AnyDiagnostic::MismatchedArgCount(d) => mismatched_arg_count::mismatched_arg_count(&ctx, &d),
             AnyDiagnostic::MissingFields(d) => missing_fields::missing_fields(&ctx, &d),
+            AnyDiagnostic::MissingOkOrSomeInTailExpr(d) => missing_ok_or_some_in_tail_expr::missing_ok_or_some_in_tail_expr(&ctx, &d),
             AnyDiagnostic::MissingUnsafe(d) => missing_unsafe::missing_unsafe(&ctx, &d),
             AnyDiagnostic::NoSuchField(d) => no_such_field::no_such_field(&ctx, &d),
             AnyDiagnostic::RemoveThisSemicolon(d) => remove_this_semicolon::remove_this_semicolon(&ctx, &d),
diff --git a/crates/ide/src/diagnostics/fixes.rs b/crates/ide/src/diagnostics/fixes.rs
index 70f17881e..350575b3a 100644
--- a/crates/ide/src/diagnostics/fixes.rs
+++ b/crates/ide/src/diagnostics/fixes.rs
@@ -2,7 +2,6 @@
 //! The same module also has all curret custom fixes for the diagnostics implemented.
 mod change_case;
 mod replace_with_find_map;
-mod wrap_tail_expr;
 
 use hir::{diagnostics::Diagnostic, Semantics};
 use ide_assists::AssistResolveStrategy;
diff --git a/crates/ide/src/diagnostics/fixes/wrap_tail_expr.rs b/crates/ide/src/diagnostics/missing_ok_or_some_in_tail_expr.rs
index dcb21e037..e27b54e66 100644
--- a/crates/ide/src/diagnostics/fixes/wrap_tail_expr.rs
+++ b/crates/ide/src/diagnostics/missing_ok_or_some_in_tail_expr.rs
@@ -1,26 +1,45 @@
-use hir::{db::AstDatabase, diagnostics::MissingOkOrSomeInTailExpr, Semantics};
-use ide_assists::{Assist, AssistResolveStrategy};
-use ide_db::{source_change::SourceChange, RootDatabase};
+use hir::{db::AstDatabase, Semantics};
+use ide_assists::Assist;
+use ide_db::source_change::SourceChange;
 use syntax::AstNode;
 use text_edit::TextEdit;
 
-use crate::diagnostics::{fix, DiagnosticWithFixes};
-
-impl DiagnosticWithFixes for MissingOkOrSomeInTailExpr {
-    fn fixes(
-        &self,
-        sema: &Semantics<RootDatabase>,
-        _resolve: &AssistResolveStrategy,
-    ) -> Option<Vec<Assist>> {
-        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 replacement = format!("{}({})", self.required, tail_expr.syntax());
-        let edit = TextEdit::replace(tail_expr_range, replacement);
-        let source_change = SourceChange::from_text_edit(self.file.original_file(sema.db), edit);
-        let name = if self.required == "Ok" { "Wrap with Ok" } else { "Wrap with Some" };
-        Some(vec![fix("wrap_tail_expr", name, source_change, tail_expr_range)])
-    }
+use crate::diagnostics::{fix, Diagnostic, DiagnosticsContext};
+
+// Diagnostic: missing-ok-or-some-in-tail-expr
+//
+// 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:
+//
+// ```rust
+// fn foo() -> Result<u8, ()> {
+//     10
+// }
+// ```
+pub(super) fn missing_ok_or_some_in_tail_expr(
+    ctx: &DiagnosticsContext<'_>,
+    d: &hir::MissingOkOrSomeInTailExpr,
+) -> Diagnostic {
+    Diagnostic::new(
+        "missing-ok-or-some-in-tail-expr",
+        format!("wrap return expression in {}", d.required),
+        ctx.sema.diagnostics_display_range(d.expr.clone().map(|it| it.into())).range,
+    )
+    .with_fixes(fixes(ctx, d))
+}
+
+fn fixes(ctx: &DiagnosticsContext<'_>, d: &hir::MissingOkOrSomeInTailExpr) -> Option<Vec<Assist>> {
+    let root = ctx.sema.db.parse_or_expand(d.expr.file_id)?;
+    let tail_expr = d.expr.value.to_node(&root);
+    let tail_expr_range = tail_expr.syntax().text_range();
+    let replacement = format!("{}({})", d.required, tail_expr.syntax());
+    let edit = TextEdit::replace(tail_expr_range, replacement);
+    let source_change =
+        SourceChange::from_text_edit(d.expr.file_id.original_file(ctx.sema.db), edit);
+    let name = if d.required == "Ok" { "Wrap with Ok" } else { "Wrap with Some" };
+    Some(vec![fix("wrap_tail_expr", name, source_change, tail_expr_range)])
 }
 
 #[cfg(test)]