Improve while_let_on_iterator suggestion inside an FnOnce closure

Jarcho · Jarcho · commit 9fa12def3c23 · 2022-07-07T22:10:24.000-04:00
diff --git a/clippy_lints/src/loops/while_let_on_iterator.rs b/clippy_lints/src/loops/while_let_on_iterator.rs
@@ -3,13 +3,15 @@ use clippy_utils::diagnostics::span_lint_and_sugg;
 use clippy_utils::higher;
 use clippy_utils::source::snippet_with_applicability;
 use clippy_utils::{
-    get_enclosing_loop_or_closure, is_refutable, is_trait_method, match_def_path, paths, visitors::is_res_used,
+    get_enclosing_loop_or_multi_call_closure, is_refutable, is_trait_method, match_def_path, paths,
+    visitors::is_res_used,
 };
 use if_chain::if_chain;
 use rustc_errors::Applicability;
 use rustc_hir::intravisit::{walk_expr, Visitor};
 use rustc_hir::{def::Res, Expr, ExprKind, HirId, Local, Mutability, PatKind, QPath, UnOp};
 use rustc_lint::LateContext;
+use rustc_middle::hir::nested_filter::OnlyBodies;
 use rustc_middle::ty::adjustment::Adjust;
 use rustc_span::{symbol::sym, Symbol};
 
@@ -249,6 +251,11 @@ fn needs_mutable_borrow(cx: &LateContext<'_>, iter_expr: &IterExpr, loop_expr: &
         used_iter: bool,
     }
     impl<'tcx> Visitor<'tcx> for AfterLoopVisitor<'_, '_, 'tcx> {
+        type NestedFilter = OnlyBodies;
+        fn nested_visit_map(&mut self) -> Self::Map {
+            self.cx.tcx.hir()
+        }
+
         fn visit_expr(&mut self, e: &'tcx Expr<'_>) {
             if self.used_iter {
                 return;
@@ -283,6 +290,11 @@ fn needs_mutable_borrow(cx: &LateContext<'_>, iter_expr: &IterExpr, loop_expr: &
         used_after: bool,
     }
     impl<'a, 'b, 'tcx> Visitor<'tcx> for NestedLoopVisitor<'a, 'b, 'tcx> {
+        type NestedFilter = OnlyBodies;
+        fn nested_visit_map(&mut self) -> Self::Map {
+            self.cx.tcx.hir()
+        }
+
         fn visit_local(&mut self, l: &'tcx Local<'_>) {
             if !self.after_loop {
                 l.pat.each_binding_or_first(&mut |_, id, _, _| {
@@ -320,10 +332,7 @@ fn needs_mutable_borrow(cx: &LateContext<'_>, iter_expr: &IterExpr, loop_expr: &
         }
     }
 
-    if let Some(e) = get_enclosing_loop_or_closure(cx.tcx, loop_expr) {
-        // The iterator expression will be used on the next iteration (for loops), or on the next call (for
-        // closures) unless it is declared within the enclosing expression. TODO: Check for closures
-        // used where an `FnOnce` type is expected.
+    if let Some(e) = get_enclosing_loop_or_multi_call_closure(cx, loop_expr) {
         let local_id = match iter_expr.path {
             Res::Local(id) => id,
             _ => return true,
diff --git a/clippy_utils/src/lib.rs b/clippy_utils/src/lib.rs
@@ -93,7 +93,9 @@ use rustc_middle::ty::fast_reject::SimplifiedTypeGen::{
     ArraySimplifiedType, BoolSimplifiedType, CharSimplifiedType, FloatSimplifiedType, IntSimplifiedType,
     PtrSimplifiedType, SliceSimplifiedType, StrSimplifiedType, UintSimplifiedType,
 };
-use rustc_middle::ty::{layout::IntegerExt, BorrowKind, DefIdTree, Ty, TyCtxt, TypeAndMut, TypeFoldable, UpvarCapture};
+use rustc_middle::ty::{
+    layout::IntegerExt, BorrowKind, ClosureKind, DefIdTree, Ty, TyCtxt, TypeAndMut, TypeFoldable, UpvarCapture,
+};
 use rustc_middle::ty::{FloatTy, IntTy, UintTy};
 use rustc_semver::RustcVersion;
 use rustc_session::Session;
@@ -105,7 +107,7 @@ use rustc_span::{Span, DUMMY_SP};
 use rustc_target::abi::Integer;
 
 use crate::consts::{constant, Constant};
-use crate::ty::{can_partially_move_ty, is_copy, is_recursively_primitive_type};
+use crate::ty::{can_partially_move_ty, expr_sig, is_copy, is_recursively_primitive_type, ty_is_fn_once_param};
 use crate::visitors::expr_visitor_no_bodies;
 
 pub fn parse_msrv(msrv: &str, sess: Option<&Session>, span: Option<Span>) -> Option<RustcVersion> {
@@ -1197,16 +1199,54 @@ pub fn get_enclosing_block<'tcx>(cx: &LateContext<'tcx>, hir_id: HirId) -> Optio
 }
 
 /// Gets the loop or closure enclosing the given expression, if any.
-pub fn get_enclosing_loop_or_closure<'tcx>(tcx: TyCtxt<'tcx>, expr: &Expr<'_>) -> Option<&'tcx Expr<'tcx>> {
-    for (_, node) in tcx.hir().parent_iter(expr.hir_id) {
+pub fn get_enclosing_loop_or_multi_call_closure<'tcx>(
+    cx: &LateContext<'tcx>,
+    expr: &Expr<'_>,
+) -> Option<&'tcx Expr<'tcx>> {
+    for (_, node) in cx.tcx.hir().parent_iter(expr.hir_id) {
         match node {
-            Node::Expr(
-                e @ Expr {
-                    kind: ExprKind::Loop(..) | ExprKind::Closure { .. },
-                    ..
+            Node::Expr(e) => match e.kind {
+                ExprKind::Closure { .. } => {
+                    if let rustc_ty::Closure(_, subs) = cx.typeck_results().expr_ty(e).kind()
+                        && subs.as_closure().kind() == ClosureKind::FnOnce
+                    {
+                        continue;
+                    }
+                    let is_once = walk_to_expr_usage(cx, e, |node, id| {
+                        let Node::Expr(e) = node else {
+                            return None;
+                        };
+                        match e.kind {
+                            ExprKind::Call(f, _) if f.hir_id == id => Some(()),
+                            ExprKind::Call(f, args) => {
+                                let i = args.iter().position(|arg| arg.hir_id == id)?;
+                                let sig = expr_sig(cx, f)?;
+                                let predicates = sig
+                                    .predicates_id()
+                                    .map_or(cx.param_env, |id| cx.tcx.param_env(id))
+                                    .caller_bounds();
+                                sig.input(i).and_then(|ty| {
+                                    ty_is_fn_once_param(cx.tcx, ty.skip_binder(), predicates).then_some(())
+                                })
+                            },
+                            ExprKind::MethodCall(_, args, _) => {
+                                let i = args.iter().position(|arg| arg.hir_id == id)?;
+                                let id = cx.typeck_results().type_dependent_def_id(e.hir_id)?;
+                                let ty = cx.tcx.fn_sig(id).skip_binder().inputs()[i];
+                                ty_is_fn_once_param(cx.tcx, ty, cx.tcx.param_env(id).caller_bounds()).then_some(())
+                            },
+                            _ => None,
+                        }
+                    })
+                    .is_some();
+                    if !is_once {
+                        return Some(e);
+                    }
                 },
-            ) => return Some(e),
-            Node::Expr(_) | Node::Stmt(_) | Node::Block(_) | Node::Local(_) | Node::Arm(_) => (),
+                ExprKind::Loop(..) => return Some(e),
+                _ => (),
+            },
+            Node::Stmt(_) | Node::Block(_) | Node::Local(_) | Node::Arm(_) => (),
             _ => break,
         }
     }
diff --git a/clippy_utils/src/ty.rs b/clippy_utils/src/ty.rs
@@ -501,7 +501,7 @@ pub fn all_predicates_of(tcx: TyCtxt<'_>, id: DefId) -> impl Iterator<Item = &(P
 /// A signature for a function like type.
 #[derive(Clone, Copy)]
 pub enum ExprFnSig<'tcx> {
-    Sig(Binder<'tcx, FnSig<'tcx>>),
+    Sig(Binder<'tcx, FnSig<'tcx>>, Option<DefId>),
     Closure(Option<&'tcx FnDecl<'tcx>>, Binder<'tcx, FnSig<'tcx>>),
     Trait(Binder<'tcx, Ty<'tcx>>, Option<Binder<'tcx, Ty<'tcx>>>),
 }
@@ -510,7 +510,7 @@ impl<'tcx> ExprFnSig<'tcx> {
     /// bounds only for variadic functions, otherwise this will panic.
     pub fn input(self, i: usize) -> Option<Binder<'tcx, Ty<'tcx>>> {
         match self {
-            Self::Sig(sig) => {
+            Self::Sig(sig, _) => {
                 if sig.c_variadic() {
                     sig.inputs().map_bound(|inputs| inputs.get(i).copied()).transpose()
                 } else {
@@ -527,7 +527,7 @@ impl<'tcx> ExprFnSig<'tcx> {
     /// functions, otherwise this will panic.
     pub fn input_with_hir(self, i: usize) -> Option<(Option<&'tcx hir::Ty<'tcx>>, Binder<'tcx, Ty<'tcx>>)> {
         match self {
-            Self::Sig(sig) => {
+            Self::Sig(sig, _) => {
                 if sig.c_variadic() {
                     sig.inputs()
                         .map_bound(|inputs| inputs.get(i).copied())
@@ -549,16 +549,20 @@ impl<'tcx> ExprFnSig<'tcx> {
     /// specified.
     pub fn output(self) -> Option<Binder<'tcx, Ty<'tcx>>> {
         match self {
-            Self::Sig(sig) | Self::Closure(_, sig) => Some(sig.output()),
+            Self::Sig(sig, _) | Self::Closure(_, sig) => Some(sig.output()),
             Self::Trait(_, output) => output,
         }
     }
+
+    pub fn predicates_id(&self) -> Option<DefId> {
+        if let ExprFnSig::Sig(_, id) = *self { id } else { None }
+    }
 }
 
 /// If the expression is function like, get the signature for it.
 pub fn expr_sig<'tcx>(cx: &LateContext<'tcx>, expr: &Expr<'_>) -> Option<ExprFnSig<'tcx>> {
     if let Res::Def(DefKind::Fn | DefKind::Ctor(_, CtorKind::Fn) | DefKind::AssocFn, id) = path_res(cx, expr) {
-        Some(ExprFnSig::Sig(cx.tcx.fn_sig(id)))
+        Some(ExprFnSig::Sig(cx.tcx.fn_sig(id), Some(id)))
     } else {
         ty_sig(cx, cx.typeck_results().expr_ty_adjusted(expr).peel_refs())
     }
@@ -575,9 +579,9 @@ fn ty_sig<'tcx>(cx: &LateContext<'tcx>, ty: Ty<'tcx>) -> Option<ExprFnSig<'tcx>>
                 .and_then(|id| cx.tcx.hir().fn_decl_by_hir_id(cx.tcx.hir().local_def_id_to_hir_id(id)));
             Some(ExprFnSig::Closure(decl, subs.as_closure().sig()))
         },
-        ty::FnDef(id, subs) => Some(ExprFnSig::Sig(cx.tcx.bound_fn_sig(id).subst(cx.tcx, subs))),
+        ty::FnDef(id, subs) => Some(ExprFnSig::Sig(cx.tcx.bound_fn_sig(id).subst(cx.tcx, subs), Some(id))),
         ty::Opaque(id, _) => ty_sig(cx, cx.tcx.type_of(id)),
-        ty::FnPtr(sig) => Some(ExprFnSig::Sig(sig)),
+        ty::FnPtr(sig) => Some(ExprFnSig::Sig(sig, None)),
         ty::Dynamic(bounds, _) => {
             let lang_items = cx.tcx.lang_items();
             match bounds.principal() {
@@ -793,3 +797,33 @@ pub fn variant_of_res<'tcx>(cx: &LateContext<'tcx>, res: Res) -> Option<&'tcx Va
         _ => None,
     }
 }
+
+/// Checks if the type is a type parameter implementing `FnOnce`, but not `FnMut`.
+pub fn ty_is_fn_once_param<'tcx>(tcx: TyCtxt<'_>, ty: Ty<'tcx>, predicates: &'tcx [Predicate<'_>]) -> bool {
+    let ty::Param(ty) = *ty.kind() else {
+        return false;
+    };
+    let lang = tcx.lang_items();
+    let (Some(fn_once_id), Some(fn_mut_id), Some(fn_id))
+        = (lang.fn_once_trait(), lang.fn_mut_trait(), lang.fn_trait())
+    else {
+        return false;
+    };
+    predicates
+        .iter()
+        .try_fold(false, |found, p| {
+            if let PredicateKind::Trait(p) = p.kind().skip_binder()
+            && let ty::Param(self_ty) = p.trait_ref.self_ty().kind()
+            && ty.index == self_ty.index
+        {
+            // This should use `super_traits_of`, but that's a private function.
+            if p.trait_ref.def_id == fn_once_id {
+                return Some(true);
+            } else if p.trait_ref.def_id == fn_mut_id || p.trait_ref.def_id == fn_id {
+                return None;
+            }
+        }
+            Some(found)
+        })
+        .unwrap_or(false)
+}
diff --git a/tests/ui/while_let_on_iterator.fixed b/tests/ui/while_let_on_iterator.fixed
@@ -7,7 +7,8 @@
     unused_mut,
     dead_code,
     clippy::equatable_if_let,
-    clippy::manual_find
+    clippy::manual_find,
+    clippy::redundant_closure_call
 )]
 
 fn base() {
@@ -259,7 +260,7 @@ fn issue1924() {
         fn f(&mut self) -> Option<u32> {
             // Used as a field.
             for i in self.0.by_ref() {
-                if !(3..=7).contains(&i) {
+                if !(3..8).contains(&i) {
                     return Some(i);
                 }
             }
@@ -403,6 +404,47 @@ fn issue_8113() {
     }
 }
 
+fn fn_once_closure() {
+    let mut it = 0..10;
+    (|| {
+        for x in it {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    })();
+
+    fn f(_: impl FnOnce()) {}
+    let mut it = 0..10;
+    f(|| {
+        for x in it {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    });
+
+    fn f2(_: impl FnMut()) {}
+    let mut it = 0..10;
+    f2(|| {
+        for x in it.by_ref() {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    });
+
+    fn f3(_: fn()) {}
+    f3(|| {
+        let mut it = 0..10;
+        for x in it {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    })
+}
+
 fn main() {
     let mut it = 0..20;
     for _ in it {
diff --git a/tests/ui/while_let_on_iterator.rs b/tests/ui/while_let_on_iterator.rs
@@ -7,7 +7,8 @@
     unused_mut,
     dead_code,
     clippy::equatable_if_let,
-    clippy::manual_find
+    clippy::manual_find,
+    clippy::redundant_closure_call
 )]
 
 fn base() {
@@ -259,7 +260,7 @@ fn issue1924() {
         fn f(&mut self) -> Option<u32> {
             // Used as a field.
             while let Some(i) = self.0.next() {
-                if i < 3 || i > 7 {
+                if !(3..8).contains(&i) {
                     return Some(i);
                 }
             }
@@ -403,6 +404,47 @@ fn issue_8113() {
     }
 }
 
+fn fn_once_closure() {
+    let mut it = 0..10;
+    (|| {
+        while let Some(x) = it.next() {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    })();
+
+    fn f(_: impl FnOnce()) {}
+    let mut it = 0..10;
+    f(|| {
+        while let Some(x) = it.next() {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    });
+
+    fn f2(_: impl FnMut()) {}
+    let mut it = 0..10;
+    f2(|| {
+        while let Some(x) = it.next() {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    });
+
+    fn f3(_: fn()) {}
+    f3(|| {
+        let mut it = 0..10;
+        while let Some(x) = it.next() {
+            if x % 2 == 0 {
+                break;
+            }
+        }
+    })
+}
+
 fn main() {
     let mut it = 0..20;
     while let Some(..) = it.next() {
diff --git a/tests/ui/while_let_on_iterator.stderr b/tests/ui/while_let_on_iterator.stderr
