Reuse `is_expr_identity_function` for `flat_map_identity`

4 files changed, 28 insertions, 38 deletions

diff --git a/clippy_lints/src/methods/flat_map_identity.rs b/clippy_lints/src/methods/flat_map_identity.rs
index 25f8434cb94..6f911d79d0b 100644
--- a/clippy_lints/src/methods/flat_map_identity.rs
+++ b/clippy_lints/src/methods/flat_map_identity.rs
@@ -1,6 +1,5 @@
 use clippy_utils::diagnostics::span_lint_and_sugg;
-use clippy_utils::{is_expr_path_def_path, is_trait_method, paths};
-use if_chain::if_chain;
+use clippy_utils::{is_expr_identity_function, is_trait_method};
 use rustc_errors::Applicability;
 use rustc_hir as hir;
 use rustc_lint::LateContext;
@@ -15,36 +14,15 @@ pub(super) fn check<'tcx>(
     flat_map_arg: &'tcx hir::Expr<'_>,
     flat_map_span: Span,
 ) {
-    if is_trait_method(cx, expr, sym::Iterator) {
-        let apply_lint = |message: &str| {
-            span_lint_and_sugg(
-                cx,
-                FLAT_MAP_IDENTITY,
-                flat_map_span.with_hi(expr.span.hi()),
-                message,
-                "try",
-                "flatten()".to_string(),
-                Applicability::MachineApplicable,
-            );
-        };
-
-        if_chain! {
-            if let hir::ExprKind::Closure(_, _, body_id, _, _) = flat_map_arg.kind;
-            let body = cx.tcx.hir().body(body_id);
-
-            if let hir::PatKind::Binding(_, _, binding_ident, _) = body.params[0].pat.kind;
-            if let hir::ExprKind::Path(hir::QPath::Resolved(_, path)) = body.value.kind;
-
-            if path.segments.len() == 1;
-            if path.segments[0].ident.name == binding_ident.name;
-
-            then {
-                apply_lint("called `flat_map(|x| x)` on an `Iterator`");
-            }
-        }
-
-        if is_expr_path_def_path(cx, flat_map_arg, &paths::CONVERT_IDENTITY) {
-            apply_lint("called `flat_map(std::convert::identity)` on an `Iterator`");
-        }
+    if is_trait_method(cx, expr, sym::Iterator) && is_expr_identity_function(cx, flat_map_arg) {
+        span_lint_and_sugg(
+            cx,
+            FLAT_MAP_IDENTITY,
+            flat_map_span.with_hi(expr.span.hi()),
+            "use of `flat_map` with an identity function",
+            "try",
+            "flatten()".to_string(),
+            Applicability::MachineApplicable,
+        );
     }
 }
diff --git a/tests/ui/flat_map_identity.fixed b/tests/ui/flat_map_identity.fixed
index dfe3bd47e13..1f4b880ef5b 100644
--- a/tests/ui/flat_map_identity.fixed
+++ b/tests/ui/flat_map_identity.fixed
@@ -1,6 +1,6 @@
 // run-rustfix
 
-#![allow(unused_imports)]
+#![allow(unused_imports, clippy::needless_return)]
 #![warn(clippy::flat_map_identity)]
 
 use std::convert;
@@ -11,4 +11,7 @@ fn main() {
 
     let iterator = [[0, 1], [2, 3], [4, 5]].iter();
     let _ = iterator.flatten();
+
+    let iterator = [[0, 1], [2, 3], [4, 5]].iter();
+    let _ = iterator.flatten();
 }
diff --git a/tests/ui/flat_map_identity.rs b/tests/ui/flat_map_identity.rs
index 393b9569255..de14a06d4e6 100644
--- a/tests/ui/flat_map_identity.rs
+++ b/tests/ui/flat_map_identity.rs
@@ -1,6 +1,6 @@
 // run-rustfix
 
-#![allow(unused_imports)]
+#![allow(unused_imports, clippy::needless_return)]
 #![warn(clippy::flat_map_identity)]
 
 use std::convert;
@@ -11,4 +11,7 @@ fn main() {
 
     let iterator = [[0, 1], [2, 3], [4, 5]].iter();
     let _ = iterator.flat_map(convert::identity);
+
+    let iterator = [[0, 1], [2, 3], [4, 5]].iter();
+    let _ = iterator.flat_map(|x| return x);
 }
diff --git a/tests/ui/flat_map_identity.stderr b/tests/ui/flat_map_identity.stderr
index e4686ae5a54..e776c9fdf51 100644
--- a/tests/ui/flat_map_identity.stderr
+++ b/tests/ui/flat_map_identity.stderr
@@ -1,4 +1,4 @@
-error: called `flat_map(|x| x)` on an `Iterator`
+error: use of `flat_map` with an identity function
   --> $DIR/flat_map_identity.rs:10:22
    |
 LL |     let _ = iterator.flat_map(|x| x);
@@ -6,11 +6,17 @@ LL |     let _ = iterator.flat_map(|x| x);
    |
    = note: `-D clippy::flat-map-identity` implied by `-D warnings`
 
-error: called `flat_map(std::convert::identity)` on an `Iterator`
+error: use of `flat_map` with an identity function
   --> $DIR/flat_map_identity.rs:13:22
    |
 LL |     let _ = iterator.flat_map(convert::identity);
    |                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: try: `flatten()`
 
-error: aborting due to 2 previous errors
+error: use of `flat_map` with an identity function
+  --> $DIR/flat_map_identity.rs:16:22
+   |
+LL |     let _ = iterator.flat_map(|x| return x);
+   |                      ^^^^^^^^^^^^^^^^^^^^^^ help: try: `flatten()`
+
+error: aborting due to 3 previous errors