Implement fix, add tests

4 files changed, 66 insertions, 55 deletions

diff --git a/crates/hir/src/diagnostics.rs b/crates/hir/src/diagnostics.rs
index f7fd3e237..5343a036c 100644
--- a/crates/hir/src/diagnostics.rs
+++ b/crates/hir/src/diagnostics.rs
@@ -7,12 +7,3 @@ pub use hir_ty::diagnostics::{
     IncorrectCase, MismatchedArgCount, MissingFields, MissingMatchArms, MissingOkOrSomeInTailExpr,
     NoSuchField, RemoveThisSemicolon, ReplaceFilterMapNextWithFindMap,
 };
-
-// PHIL:
-// hir/src/diagnostics.rs - just pub uses the type from hir_ty::diagnostics (DONE)
-// hir_ty/src/diagnostics.rs - defines the type (DONE)
-// hir_ty/src/diagnostics.rs - plus a test (DONE) <--- one example found, need to copy the not-applicable tests from the assist version
-// ide/src/diagnostics.rs - define handler for when this diagnostic is raised (DONE)
-
-// ide/src/diagnostics/fixes.rs - pulls in type from hir, and impls DiagnosticWithFix (TODO)
-// hir_ty/src/diagnostics/expr.rs - do the real work (TODO)
diff --git a/crates/hir_ty/src/diagnostics.rs b/crates/hir_ty/src/diagnostics.rs
index b4cf81c10..3d7663f6a 100644
--- a/crates/hir_ty/src/diagnostics.rs
+++ b/crates/hir_ty/src/diagnostics.rs
@@ -392,7 +392,7 @@ impl Diagnostic for IncorrectCase {
 #[derive(Debug)]
 pub struct ReplaceFilterMapNextWithFindMap {
     pub file: HirFileId,
-    pub filter_map_expr: AstPtr<ast::Expr>,
+    /// This expression is the whole method chain up to and including `.filter_map(..).next()`.
     pub next_expr: AstPtr<ast::Expr>,
 }
 
@@ -404,7 +404,7 @@ impl Diagnostic for ReplaceFilterMapNextWithFindMap {
         "replace filter_map(..).next() with find_map(..)".to_string()
     }
     fn display_source(&self) -> InFile<SyntaxNodePtr> {
-        InFile { file_id: self.file, value: self.filter_map_expr.clone().into() }
+        InFile { file_id: self.file, value: self.next_expr.clone().into() }
     }
     fn as_any(&self) -> &(dyn Any + Send + 'static) {
         self
@@ -671,15 +671,55 @@ fn foo() { break; }
     }
 
     #[test]
-    fn replace_missing_filter_next_with_find_map() {
+    fn replace_filter_next_with_find_map() {
         check_diagnostics(
             r#"
             fn foo() {
-            let m = [1, 2, 3]
-                .iter()
-                .filter_map(|x| if *x == 2 { Some (4) } else { None })
-                .next();
-                //^^^ Replace .filter_map(..).next() with .find_map(..)
+                let m = [1, 2, 3].iter().filter_map(|x| if *x == 2 { Some (4) } else { None }).next();
+                      //^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ replace filter_map(..).next() with find_map(..)
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn replace_filter_next_with_find_map_no_diagnostic_without_next() {
+        check_diagnostics(
+            r#"
+            fn foo() {
+                let m = [1, 2, 3]
+                    .iter()
+                    .filter_map(|x| if *x == 2 { Some (4) } else { None })
+                    .len();
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn replace_filter_next_with_find_map_no_diagnostic_with_intervening_methods() {
+        check_diagnostics(
+            r#"
+            fn foo() {
+                let m = [1, 2, 3]
+                    .iter()
+                    .filter_map(|x| if *x == 2 { Some (4) } else { None })
+                    .map(|x| x + 2)
+                    .len();
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn replace_filter_next_with_find_map_no_diagnostic_if_not_in_chain() {
+        check_diagnostics(
+            r#"
+            fn foo() {
+                let m = [1, 2, 3]
+                    .iter()
+                    .filter_map(|x| if *x == 2 { Some (4) } else { None });
+                let n = m.next();
             }
             "#,
         );
diff --git a/crates/hir_ty/src/diagnostics/expr.rs b/crates/hir_ty/src/diagnostics/expr.rs
index 170d23178..b87557ff5 100644
--- a/crates/hir_ty/src/diagnostics/expr.rs
+++ b/crates/hir_ty/src/diagnostics/expr.rs
@@ -41,16 +41,7 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
         ExprValidator { owner, infer, sink }
     }
 
-    fn bar() {
-        // LOOK FOR THIS
-        let m = [1, 2, 3]
-            .iter()
-            .filter_map(|x| if *x == 2 { Some(4) } else { None })
-            .next();
-    }
-
     pub(super) fn validate_body(&mut self, db: &dyn HirDatabase) {
-        // DO NOT MERGE: just getting something working for now
         self.check_for_filter_map_next(db);
 
         let body = db.body(self.owner.into());
@@ -169,24 +160,20 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
 
         for (id, expr) in body.exprs.iter() {
             if let Expr::MethodCall { receiver, method_name, args, .. } = expr {
-                let method_name_hack_do_not_merge = format!("{}", method_name);
+                let method_name = format!("{}", method_name);
 
-                if method_name_hack_do_not_merge == "filter_map" && args.len() == 1 {
-                    prev = Some((id, args[0]));
+                if method_name == "filter_map" && args.len() == 1 {
+                    prev = Some(id);
                     continue;
                 }
 
-                if method_name_hack_do_not_merge == "next" {
-                    if let Some((filter_map_id, filter_map_args)) = prev {
+                if method_name == "next" {
+                    if let Some(filter_map_id) = prev {
                         if *receiver == filter_map_id {
                             let (_, source_map) = db.body_with_source_map(self.owner.into());
-                            if let (Ok(filter_map_source_ptr), Ok(next_source_ptr)) = (
-                                source_map.expr_syntax(filter_map_id),
-                                source_map.expr_syntax(id),
-                            ) {
+                            if let Ok(next_source_ptr) = source_map.expr_syntax(id) {
                                 self.sink.push(ReplaceFilterMapNextWithFindMap {
-                                    file: filter_map_source_ptr.file_id,
-                                    filter_map_expr: filter_map_source_ptr.value,
+                                    file: next_source_ptr.file_id,
                                     next_expr: next_source_ptr.value,
                                 });
                             }
diff --git a/crates/ide/src/diagnostics/fixes.rs b/crates/ide/src/diagnostics/fixes.rs
index eafbac43a..7bbf1d8c7 100644
--- a/crates/ide/src/diagnostics/fixes.rs
+++ b/crates/ide/src/diagnostics/fixes.rs
@@ -1,5 +1,6 @@
 //! Provides a way to attach fixes to the diagnostics.
 //! The same module also has all curret custom fixes for the diagnostics implemented.
+use ast::MethodCallExpr;
 use hir::{
     db::AstDatabase,
     diagnostics::{
@@ -13,11 +14,7 @@ use ide_db::{
     source_change::{FileSystemEdit, SourceChange},
     RootDatabase,
 };
-use syntax::{
-    algo,
-    ast::{self, edit::IndentLevel, make},
-    AstNode,
-};
+use syntax::{AstNode, TextRange, algo, ast::{self, ArgList, edit::IndentLevel, make}};
 use text_edit::TextEdit;
 
 use crate::{diagnostics::Fix, references::rename::rename_with_semantics, FilePosition};
@@ -144,25 +141,21 @@ impl DiagnosticWithFix for IncorrectCase {
     }
 }
 
-// Bugs:
-//  * Action is applicable for both iter() and filter_map() rows
-//  * Action deletes the entire method chain
 impl DiagnosticWithFix for ReplaceFilterMapNextWithFindMap {
     fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
         let root = sema.db.parse_or_expand(self.file)?;
-
         let next_expr = self.next_expr.to_node(&root);
-        let next_expr_range = next_expr.syntax().text_range();
+        let next_call = MethodCallExpr::cast(next_expr.syntax().clone())?;
 
-        let filter_map_expr = self.filter_map_expr.to_node(&root);
-        let filter_map_expr_range = filter_map_expr.syntax().text_range();
+        let filter_map_call = MethodCallExpr::cast(next_call.receiver()?.syntax().clone())?;
+        let filter_map_name_range = filter_map_call.name_ref()?.ident_token()?.text_range();
+        let filter_map_args = filter_map_call.syntax().children().find_map(ArgList::cast)?;
 
-        let edit = TextEdit::delete(next_expr_range);
+        let range_to_replace = TextRange::new(filter_map_name_range.start(), next_expr.syntax().text_range().end());
+        let replacement = format!("find_map{}", filter_map_args.syntax().text());
+        let trigger_range = next_expr.syntax().text_range();
 
-        // This is the entire method chain, including the array literal
-        eprintln!("NEXT EXPR: {:#?}", next_expr);
-        // This is the entire method chain except for the final next()
-        eprintln!("FILTER MAP EXPR: {:#?}", filter_map_expr);
+        let edit = TextEdit::replace(range_to_replace, replacement);
 
         let source_change =
             SourceFileEdit { file_id: self.file.original_file(sema.db), edit }.into();
@@ -170,7 +163,7 @@ impl DiagnosticWithFix for ReplaceFilterMapNextWithFindMap {
         Some(Fix::new(
             "Replace filter_map(..).next() with find_map()",
             source_change,
-            filter_map_expr_range,
+            trigger_range,
         ))
     }
 }