From a85a3d73edf39f6e4ba6c6ac830743fa31f4cdf3 Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Sat, 9 Oct 2021 22:45:54 +0300
Subject: [PATCH] feat: be more precise in mismatch arg counts highlighting

---
 .../src/handlers/mismatched_arg_count.rs      | 66 +++++++++++++++----
 1 file changed, 54 insertions(+), 12 deletions(-)

diff --git a/crates/ide_diagnostics/src/handlers/mismatched_arg_count.rs b/crates/ide_diagnostics/src/handlers/mismatched_arg_count.rs
index a9b6d3870af..d8f8c316b6a 100644
--- a/crates/ide_diagnostics/src/handlers/mismatched_arg_count.rs
+++ b/crates/ide_diagnostics/src/handlers/mismatched_arg_count.rs
@@ -1,3 +1,10 @@
+use ide_db::base_db::{FileRange, SourceDatabase};
+use syntax::{
+    algo::find_node_at_range,
+    ast::{self, HasArgList},
+    AstNode,
+};
+
 use crate::{Diagnostic, DiagnosticsContext};
 
 // Diagnostic: mismatched-arg-count
@@ -12,10 +19,45 @@ pub(crate) fn mismatched_arg_count(
     Diagnostic::new(
         "mismatched-arg-count",
         message,
-        ctx.sema.diagnostics_display_range(d.call_expr.clone().map(|it| it.into())).range,
+        invalid_args_range(ctx, d).unwrap_or_else(|it| it).range,
     )
 }
 
+fn invalid_args_range(
+    ctx: &DiagnosticsContext<'_>,
+    d: &hir::MismatchedArgCount,
+) -> Result<FileRange, FileRange> {
+    let full_range = ctx.sema.diagnostics_display_range(d.call_expr.clone().map(|it| it.into()));
+
+    let source_file = ctx.sema.db.parse(full_range.file_id);
+    let expr = find_node_at_range::<ast::Expr>(&source_file.syntax_node(), full_range.range)
+        .filter(|it| it.syntax().text_range() == full_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 Err(full_range),
+    };
+    let arg_list_range =
+        FileRange { file_id: full_range.file_id, range: arg_list.syntax().text_range() };
+    if d.found < d.expected {
+        if d.found == 0 {
+            return Ok(arg_list_range);
+        }
+        if let Some(r_paren) = arg_list.r_paren_token() {
+            return Ok(FileRange { file_id: full_range.file_id, range: r_paren.text_range() });
+        }
+    }
+    if d.expected < d.found {
+        return Ok(arg_list_range);
+    }
+
+    Err(full_range)
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::check_diagnostics;
@@ -26,7 +68,7 @@ mod tests {
             r#"
 fn zero() {}
 fn f() { zero(1); }
-       //^^^^^^^ error: expected 0 arguments, found 1
+           //^^^ error: expected 0 arguments, found 1
 "#,
         );
 
@@ -44,7 +86,7 @@ fn f() { zero(); }
             r#"
 fn one(arg: u8) {}
 fn f() { one(); }
-       //^^^^^ error: expected 1 argument, found 0
+          //^^ error: expected 1 argument, found 0
 "#,
         );
 
@@ -65,7 +107,7 @@ impl S { fn method(&self) {} }
 
 fn f() {
     S::method();
-} //^^^^^^^^^^^ error: expected 1 argument, found 0
+}          //^^ error: expected 1 argument, found 0
 "#,
         );
 
@@ -91,7 +133,7 @@ impl S { fn method(&self, arg: u8) {} }
 
             fn f() {
                 S.method();
-            } //^^^^^^^^^^ error: expected 1 argument, found 0
+            }         //^^ error: expected 1 argument, found 0
             "#,
         );
 
@@ -131,7 +173,7 @@ fn f() {
 struct Tup(u8, u16);
 fn f() {
     Tup(0);
-} //^^^^^^ error: expected 2 arguments, found 1
+}      //^ error: expected 2 arguments, found 1
 "#,
         )
     }
@@ -143,7 +185,7 @@ fn f() {
 enum En { Variant(u8, u16), }
 fn f() {
     En::Variant(0);
-} //^^^^^^^^^^^^^^ error: expected 2 arguments, found 1
+}              //^ error: expected 2 arguments, found 1
 "#,
         )
     }
@@ -162,9 +204,9 @@ impl Foo {
     fn new() {
         Foo::Bar(0);
         Foo::Bar(0, 1);
-      //^^^^^^^^^^^^^^ error: expected 1 argument, found 2
+              //^^^^^^ error: expected 1 argument, found 2
         Foo::Bar();
-      //^^^^^^^^^^ error: expected 1 argument, found 0
+              //^^ error: expected 1 argument, found 0
     }
 }
         "#,
@@ -185,7 +227,7 @@ fn f() {
     unsafe {
         fixed(0);
         fixed(0, 1);
-      //^^^^^^^^^^^ error: expected 1 argument, found 2
+           //^^^^^^ error: expected 1 argument, found 2
         varargs(0);
         varargs(0, 1);
         varargs2();
@@ -204,10 +246,10 @@ fn f() {
 fn main() {
     let f = |()| ();
     f();
-  //^^^ error: expected 1 argument, found 0
+   //^^ error: expected 1 argument, found 0
     f(());
     f((), ());
-  //^^^^^^^^^ error: expected 1 argument, found 2
+   //^^^^^^^^ error: expected 1 argument, found 2
 }
 "#,
         )