2 files changed, 48 insertions, 16 deletions

diff --git a/crates/hir_ty/src/diagnostics.rs b/crates/hir_ty/src/diagnostics.rs
index 6eaa1beb8..323c5f963 100644
--- a/crates/hir_ty/src/diagnostics.rs
+++ b/crates/hir_ty/src/diagnostics.rs
@@ -670,21 +670,49 @@ fn foo() { break; }
         );
     }
 
+    // Register the required standard library types to make the tests work
+    fn add_filter_map_with_find_next_boilerplate(body: &str) -> String {
+        let prefix = r#"
+        //- /main.rs crate:main deps:core
+        use core::iter::Iterator;
+        use core::option::Option::{self, Some, None};
+        "#;
+        let suffix = r#"
+        //- /core/lib.rs crate:core
+        pub mod option {
+            pub enum Option<T> { Some(T), None }
+        }
+        pub mod iter {
+            pub trait Iterator {
+                type Item;
+                fn filter_map<B, F>(self, f: F) -> FilterMap where F: FnMut(Self::Item) -> Option<B> { FilterMap }
+                fn next(&mut self) -> Option<Self::Item>;
+            }
+            pub struct FilterMap {}
+            impl Iterator for FilterMap {
+                type Item = i32;
+                fn next(&mut self) -> i32 { 7 }
+            }
+        }
+        "#;
+        format!("{}{}{}", prefix, body, suffix)
+    }
+
     #[test]
-    fn replace_filter_map_next_with_find_map() {
-        check_diagnostics(
+    fn replace_filter_map_next_with_find_map2() {
+        check_diagnostics(&add_filter_map_with_find_next_boilerplate(
             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(..)
             }
-            "#,
-        );
+        "#,
+        ));
     }
 
     #[test]
     fn replace_filter_map_next_with_find_map_no_diagnostic_without_next() {
-        check_diagnostics(
+        check_diagnostics(&add_filter_map_with_find_next_boilerplate(
             r#"
             fn foo() {
                 let m = [1, 2, 3]
@@ -693,12 +721,12 @@ fn foo() { break; }
                     .len();
             }
             "#,
-        );
+        ));
     }
 
     #[test]
     fn replace_filter_map_next_with_find_map_no_diagnostic_with_intervening_methods() {
-        check_diagnostics(
+        check_diagnostics(&add_filter_map_with_find_next_boilerplate(
             r#"
             fn foo() {
                 let m = [1, 2, 3]
@@ -708,12 +736,12 @@ fn foo() { break; }
                     .len();
             }
             "#,
-        );
+        ));
     }
 
     #[test]
     fn replace_filter_map_next_with_find_map_no_diagnostic_if_not_in_chain() {
-        check_diagnostics(
+        check_diagnostics(&add_filter_map_with_find_next_boilerplate(
             r#"
             fn foo() {
                 let m = [1, 2, 3]
@@ -722,6 +750,6 @@ fn foo() { break; }
                 let n = m.next();
             }
             "#,
-        );
+        ));
     }
 }
diff --git a/crates/hir_ty/src/diagnostics/expr.rs b/crates/hir_ty/src/diagnostics/expr.rs
index 16bbd48fb..d740b7265 100644
--- a/crates/hir_ty/src/diagnostics/expr.rs
+++ b/crates/hir_ty/src/diagnostics/expr.rs
@@ -2,7 +2,9 @@
 
 use std::sync::Arc;
 
-use hir_def::{AdtId, AssocItemId, DefWithBodyId, expr::Statement, path::path, resolver::HasResolver};
+use hir_def::{
+    expr::Statement, path::path, resolver::HasResolver, AdtId, AssocItemId, DefWithBodyId,
+};
 use hir_expand::{diagnostics::DiagnosticSink, name};
 use rustc_hash::FxHashSet;
 use syntax::{ast, AstPtr};
@@ -163,11 +165,13 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
             None => return,
         };
         let iterator_trait_items = &db.trait_data(iterator_trait_id).items;
-        let filter_map_function_id = match iterator_trait_items.iter().find(|item| item.0 == name![filter_map]) {
-            Some((_, AssocItemId::FunctionId(id))) => id,
-            _ => return,
-        };
-        let next_function_id = match iterator_trait_items.iter().find(|item| item.0 == name![next]) {
+        let filter_map_function_id =
+            match iterator_trait_items.iter().find(|item| item.0 == name![filter_map]) {
+                Some((_, AssocItemId::FunctionId(id))) => id,
+                _ => return,
+            };
+        let next_function_id = match iterator_trait_items.iter().find(|item| item.0 == name![next])
+        {
             Some((_, AssocItemId::FunctionId(id))) => id,
             _ => return,
         };