internal: extract common code for adjusting diagnostic range

crates/ide-diagnostics/src/handlers/mismatched_arg_count.rs

@@ -1,11 +1,9 @@
-use ide_db::base_db::{FileRange, SourceDatabase};
 use syntax::{
-    algo::find_node_at_range,
     ast::{self, HasArgList},
     AstNode, TextRange,
 };
 
-use crate::{Diagnostic, DiagnosticsContext};
+use crate::{adjusted_display_range, Diagnostic, DiagnosticsContext};
 
 // Diagnostic: mismatched-arg-count
 //
@@ -20,40 +18,32 @@ pub(crate) fn mismatched_arg_count(
 }
 
 fn invalid_args_range(ctx: &DiagnosticsContext<'_>, d: &hir::MismatchedArgCount) -> TextRange {
-    let FileRange { file_id, range } =
-        ctx.sema.diagnostics_display_range(d.call_expr.clone().map(|it| it.into()));
+    adjusted_display_range::<ast::Expr>(ctx, d.call_expr.clone().map(|it| it.into()), &|expr| {
+        let arg_list = match expr {
+            ast::Expr::CallExpr(call) => call.arg_list()?,
+            ast::Expr::MethodCallExpr(call) => call.arg_list()?,
+            _ => return None,
+        };
+        if d.found < d.expected {
+            if d.found == 0 {
+                return Some(arg_list.syntax().text_range());
+            }
+            if let Some(r_paren) = arg_list.r_paren_token() {
+                return Some(r_paren.text_range());
+            }
+        }
+        if d.expected < d.found {
+            if d.expected == 0 {
+                return Some(arg_list.syntax().text_range());
+            }
+            let zip = arg_list.args().nth(d.expected).zip(arg_list.r_paren_token());
+            if let Some((arg, r_paren)) = zip {
+                return Some(arg.syntax().text_range().cover(r_paren.text_range()));
+            }
+        }
 
-    let source_file = ctx.sema.db.parse(file_id);
-    let expr = find_node_at_range::<ast::Expr>(&source_file.syntax_node(), range)
-        .filter(|it| it.syntax().text_range() == range);
-    let arg_list = match expr {
-        Some(ast::Expr::CallExpr(call)) => call.arg_list(),
-        Some(ast::Expr::MethodCallExpr(call)) => call.arg_list(),
-        _ => None,
-    };
-    let arg_list = match arg_list {
-        Some(it) => it,
-        None => return range,
-    };
-    if d.found < d.expected {
-        if d.found == 0 {
-            return arg_list.syntax().text_range();
-        }
-        if let Some(r_paren) = arg_list.r_paren_token() {
-            return r_paren.text_range();
-        }
-    }
-    if d.expected < d.found {
-        if d.expected == 0 {
-            return arg_list.syntax().text_range();
-        }
-        let zip = arg_list.args().nth(d.expected).zip(arg_list.r_paren_token());
-        if let Some((arg, r_paren)) = zip {
-            return arg.syntax().text_range().cover(r_paren.text_range());
-        }
-    }
-
-    range
+        None
+    })
 }
 
 #[cfg(test)]

crates/ide-diagnostics/src/handlers/type_mismatch.rs

@@ -1,34 +1,27 @@
 use hir::{db::AstDatabase, HirDisplay, Type};
-use ide_db::{
-    base_db::{FileRange, SourceDatabase},
-    famous_defs::FamousDefs,
-    source_change::SourceChange,
-};
+use ide_db::{famous_defs::FamousDefs, source_change::SourceChange};
 use syntax::{
-    algo::find_node_at_range,
     ast::{self, BlockExpr, ExprStmt},
     AstNode,
 };
 use text_edit::TextEdit;
 
-use crate::{fix, Assist, Diagnostic, DiagnosticsContext};
+use crate::{adjusted_display_range, fix, Assist, Diagnostic, DiagnosticsContext};
 
 // Diagnostic: type-mismatch
 //
 // This diagnostic is triggered when the type of an expression does not match
 // the expected type.
 pub(crate) fn type_mismatch(ctx: &DiagnosticsContext<'_>, d: &hir::TypeMismatch) -> Diagnostic {
-    let FileRange { file_id, range } =
-        ctx.sema.diagnostics_display_range(d.expr.clone().map(|it| it.into()));
-
-    let source_file = ctx.sema.db.parse(file_id);
-    let block = find_node_at_range::<ast::BlockExpr>(&source_file.syntax_node(), range)
-        .filter(|it| it.syntax().text_range() == range);
-    let display_range = block
-        .and_then(|it| it.stmt_list())
-        .and_then(|it| it.r_curly_token())
-        .map(|it| it.text_range())
-        .unwrap_or(range);
+    let display_range = adjusted_display_range::<ast::BlockExpr>(
+        ctx,
+        d.expr.clone().map(|it| it.into()),
+        &|block| {
+            let r_curly_range = block.stmt_list()?.r_curly_token()?.text_range();
+            cov_mark::hit!(type_mismatch_on_block);
+            Some(r_curly_range)
+        },
+    );
 
     let mut diag = Diagnostic::new(
         "type-mismatch",
@@ -565,6 +558,7 @@ fn test() -> String {
 
     #[test]
     fn type_mismatch_on_block() {
+        cov_mark::check!(type_mismatch_on_block);
         check_diagnostics(
             r#"
 fn f() -> i32 {

crates/ide-diagnostics/src/lib.rs

@@ -55,15 +55,15 @@ mod handlers {
 #[cfg(test)]
 mod tests;
 
-use hir::{diagnostics::AnyDiagnostic, Semantics};
+use hir::{diagnostics::AnyDiagnostic, InFile, Semantics};
 use ide_db::{
     assists::{Assist, AssistId, AssistKind, AssistResolveStrategy},
-    base_db::{FileId, SourceDatabase},
+    base_db::{FileId, FileRange, SourceDatabase},
     label::Label,
     source_change::SourceChange,
     FxHashSet, RootDatabase,
 };
-use syntax::{ast::AstNode, TextRange};
+use syntax::{algo::find_node_at_range, ast::AstNode, SyntaxNodePtr, TextRange};
 
 #[derive(Copy, Clone, Debug, PartialEq)]
 pub struct DiagnosticCode(pub &'static str);
@@ -244,3 +244,17 @@ fn unresolved_fix(id: &'static str, label: &str, target: TextRange) -> Assist {
         trigger_signature_help: false,
     }
 }
+
+fn adjusted_display_range<N: AstNode>(
+    ctx: &DiagnosticsContext<'_>,
+    diag_ptr: InFile<SyntaxNodePtr>,
+    adj: &dyn Fn(N) -> Option<TextRange>,
+) -> TextRange {
+    let FileRange { file_id, range } = ctx.sema.diagnostics_display_range(diag_ptr);
+
+    let source_file = ctx.sema.db.parse(file_id);
+    find_node_at_range::<N>(&source_file.syntax_node(), range)
+        .filter(|it| it.syntax().text_range() == range)
+        .and_then(adj)
+        .unwrap_or(range)
+}