clippy_lints/src/loops/needless_collect.rs

   1 use super::NEEDLESS_COLLECT;
   2 use clippy_utils::diagnostics::{span_lint_and_sugg, span_lint_hir_and_then};
   3 use clippy_utils::higher;
   4 use clippy_utils::source::{snippet, snippet_with_applicability};
   5 use clippy_utils::sugg::Sugg;
   6 use clippy_utils::ty::is_type_diagnostic_item;
   7 use clippy_utils::{can_move_expr_to_closure, is_trait_method, path_to_local, path_to_local_id, CaptureKind};
   8 use if_chain::if_chain;
   9 use rustc_data_structures::fx::FxHashMap;
  10 use rustc_errors::{Applicability, MultiSpan};
  11 use rustc_hir::intravisit::{walk_block, walk_expr, Visitor};
  12 use rustc_hir::{Block, Expr, ExprKind, HirId, HirIdSet, Local, Mutability, Node, PatKind, Stmt, StmtKind};
  13 use rustc_lint::LateContext;
  14 use rustc_middle::hir::nested_filter;
  15 use rustc_middle::ty::subst::GenericArgKind;
  16 use rustc_middle::ty::{self, Ty};
  17 use rustc_span::sym;
  18 use rustc_span::Span;
  19
  20 const NEEDLESS_COLLECT_MSG: &str = "avoid using `collect()` when not needed";
  21
  22 pub(super) fn check<'tcx>(expr: &'tcx Expr<'_>, cx: &LateContext<'tcx>) {
  23     check_needless_collect_direct_usage(expr, cx);
  24     check_needless_collect_indirect_usage(expr, cx);
  25 }
  26 fn check_needless_collect_direct_usage<'tcx>(expr: &'tcx Expr<'_>, cx: &LateContext<'tcx>) {
  27     if_chain! {
  28         if let ExprKind::MethodCall(method, receiver, args, _) = expr.kind;
  29         if let ExprKind::MethodCall(chain_method, ..) = receiver.kind;
  30         if chain_method.ident.name == sym!(collect) && is_trait_method(cx, receiver, sym::Iterator);
  31         then {
  32             let ty = cx.typeck_results().expr_ty(receiver);
  33             let mut applicability = Applicability::MaybeIncorrect;
  34             let is_empty_sugg = "next().is_none()".to_string();
  35             let method_name = method.ident.name.as_str();
  36             let sugg = if is_type_diagnostic_item(cx, ty, sym::Vec) ||
  37                         is_type_diagnostic_item(cx, ty, sym::VecDeque) ||
  38                         is_type_diagnostic_item(cx, ty, sym::LinkedList) ||
  39                         is_type_diagnostic_item(cx, ty, sym::BinaryHeap) {
  40                 match method_name {
  41                     "len" => "count()".to_string(),
  42                     "is_empty" => is_empty_sugg,
  43                     "contains" => {
  44                         let contains_arg = snippet_with_applicability(cx, args[0].span, "??", &mut applicability);
  45                         let (arg, pred) = contains_arg
  46                             .strip_prefix('&')
  47                             .map_or(("&x", &*contains_arg), |s| ("x", s));
  48                         format!("any(|{arg}| x == {pred})")
  49                     }
  50                     _ => return,
  51                 }
  52             }
  53             else if is_type_diagnostic_item(cx, ty, sym::BTreeMap) ||
  54                 is_type_diagnostic_item(cx, ty, sym::HashMap) {
  55                 match method_name {
  56                     "is_empty" => is_empty_sugg,
  57                     _ => return,
  58                 }
  59             }
  60             else {
  61                 return;
  62             };
  63             span_lint_and_sugg(
  64                 cx,
  65                 NEEDLESS_COLLECT,
  66                 chain_method.ident.span.with_hi(expr.span.hi()),
  67                 NEEDLESS_COLLECT_MSG,
  68                 "replace with",
  69                 sugg,
  70                 applicability,
  71             );
  72         }
  73     }
  74 }
  75
  76 fn check_needless_collect_indirect_usage<'tcx>(expr: &'tcx Expr<'_>, cx: &LateContext<'tcx>) {
  77     if let ExprKind::Block(block, _) = expr.kind {
  78         for stmt in block.stmts {
  79             if_chain! {
  80                 if let StmtKind::Local(local) = stmt.kind;
  81                 if let PatKind::Binding(_, id, ..) = local.pat.kind;
  82                 if let Some(init_expr) = local.init;
  83                 if let ExprKind::MethodCall(method_name, iter_source, [], ..) = init_expr.kind;
  84                 if method_name.ident.name == sym!(collect) && is_trait_method(cx, init_expr, sym::Iterator);
  85                 let ty = cx.typeck_results().expr_ty(init_expr);
  86                 if is_type_diagnostic_item(cx, ty, sym::Vec) ||
  87                     is_type_diagnostic_item(cx, ty, sym::VecDeque) ||
  88                     is_type_diagnostic_item(cx, ty, sym::BinaryHeap) ||
  89                     is_type_diagnostic_item(cx, ty, sym::LinkedList);
  90                 let iter_ty = cx.typeck_results().expr_ty(iter_source);
  91                 if let Some(iter_calls) = detect_iter_and_into_iters(block, id, cx, get_captured_ids(cx, iter_ty));
  92                 if let [iter_call] = &*iter_calls;
  93                 then {
  94                     let mut used_count_visitor = UsedCountVisitor {
  95                         cx,
  96                         id,
  97                         count: 0,
  98                     };
  99                     walk_block(&mut used_count_visitor, block);
 100                     if used_count_visitor.count > 1 {
 101                         return;
 102                     }
 103
 104                     // Suggest replacing iter_call with iter_replacement, and removing stmt
 105                     let mut span = MultiSpan::from_span(method_name.ident.span);
 106                     span.push_span_label(iter_call.span, "the iterator could be used here instead");
 107                     span_lint_hir_and_then(
 108                         cx,
 109                         super::NEEDLESS_COLLECT,
 110                         init_expr.hir_id,
 111                         span,
 112                         NEEDLESS_COLLECT_MSG,
 113                         |diag| {
 114                             let iter_replacement = format!("{}{}", Sugg::hir(cx, iter_source, ".."), iter_call.get_iter_method(cx));
 115                             diag.multipart_suggestion(
 116                                 iter_call.get_suggestion_text(),
 117                                 vec![
 118                                     (stmt.span, String::new()),
 119                                     (iter_call.span, iter_replacement)
 120                                 ],
 121                                 Applicability::MaybeIncorrect,
 122                             );
 123                         },
 124                     );
 125                 }
 126             }
 127         }
 128     }
 129 }
 130
 131 struct IterFunction {
 132     func: IterFunctionKind,
 133     span: Span,
 134 }
 135 impl IterFunction {
 136     fn get_iter_method(&self, cx: &LateContext<'_>) -> String {
 137         match &self.func {
 138             IterFunctionKind::IntoIter => String::new(),
 139             IterFunctionKind::Len => String::from(".count()"),
 140             IterFunctionKind::IsEmpty => String::from(".next().is_none()"),
 141             IterFunctionKind::Contains(span) => {
 142                 let s = snippet(cx, *span, "..");
 143                 if let Some(stripped) = s.strip_prefix('&') {
 144                     format!(".any(|x| x == {stripped})")
 145                 } else {
 146                     format!(".any(|x| x == *{s})")
 147                 }
 148             },
 149         }
 150     }
 151     fn get_suggestion_text(&self) -> &'static str {
 152         match &self.func {
 153             IterFunctionKind::IntoIter => {
 154                 "use the original Iterator instead of collecting it and then producing a new one"
 155             },
 156             IterFunctionKind::Len => {
 157                 "take the original Iterator's count instead of collecting it and finding the length"
 158             },
 159             IterFunctionKind::IsEmpty => {
 160                 "check if the original Iterator has anything instead of collecting it and seeing if it's empty"
 161             },
 162             IterFunctionKind::Contains(_) => {
 163                 "check if the original Iterator contains an element instead of collecting then checking"
 164             },
 165         }
 166     }
 167 }
 168 enum IterFunctionKind {
 169     IntoIter,
 170     Len,
 171     IsEmpty,
 172     Contains(Span),
 173 }
 174
 175 struct IterFunctionVisitor<'a, 'tcx> {
 176     illegal_mutable_capture_ids: HirIdSet,
 177     current_mutably_captured_ids: HirIdSet,
 178     cx: &'a LateContext<'tcx>,
 179     uses: Vec<Option<IterFunction>>,
 180     hir_id_uses_map: FxHashMap<HirId, usize>,
 181     current_statement_hir_id: Option<HirId>,
 182     seen_other: bool,
 183     target: HirId,
 184 }
 185 impl<'tcx> Visitor<'tcx> for IterFunctionVisitor<'_, 'tcx> {
 186     fn visit_block(&mut self, block: &'tcx Block<'tcx>) {
 187         for (expr, hir_id) in block.stmts.iter().filter_map(get_expr_and_hir_id_from_stmt) {
 188             if check_loop_kind(expr).is_some() {
 189                 continue;
 190             }
 191             self.visit_block_expr(expr, hir_id);
 192         }
 193         if let Some(expr) = block.expr {
 194             if let Some(loop_kind) = check_loop_kind(expr) {
 195                 if let LoopKind::Conditional(block_expr) = loop_kind {
 196                     self.visit_block_expr(block_expr, None);
 197                 }
 198             } else {
 199                 self.visit_block_expr(expr, None);
 200             }
 201         }
 202     }
 203
 204     fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
 205         // Check function calls on our collection
 206         if let ExprKind::MethodCall(method_name, recv, [args @ ..], _) = &expr.kind {
 207             if method_name.ident.name == sym!(collect) && is_trait_method(self.cx, expr, sym::Iterator) {
 208                 self.current_mutably_captured_ids = get_captured_ids(self.cx, self.cx.typeck_results().expr_ty(recv));
 209                 self.visit_expr(recv);
 210                 return;
 211             }
 212
 213             if path_to_local_id(recv, self.target) {
 214                 if self
 215                     .illegal_mutable_capture_ids
 216                     .intersection(&self.current_mutably_captured_ids)
 217                     .next()
 218                     .is_none()
 219                 {
 220                     if let Some(hir_id) = self.current_statement_hir_id {
 221                         self.hir_id_uses_map.insert(hir_id, self.uses.len());
 222                     }
 223                     match method_name.ident.name.as_str() {
 224                         "into_iter" => self.uses.push(Some(IterFunction {
 225                             func: IterFunctionKind::IntoIter,
 226                             span: expr.span,
 227                         })),
 228                         "len" => self.uses.push(Some(IterFunction {
 229                             func: IterFunctionKind::Len,
 230                             span: expr.span,
 231                         })),
 232                         "is_empty" => self.uses.push(Some(IterFunction {
 233                             func: IterFunctionKind::IsEmpty,
 234                             span: expr.span,
 235                         })),
 236                         "contains" => self.uses.push(Some(IterFunction {
 237                             func: IterFunctionKind::Contains(args[0].span),
 238                             span: expr.span,
 239                         })),
 240                         _ => {
 241                             self.seen_other = true;
 242                             if let Some(hir_id) = self.current_statement_hir_id {
 243                                 self.hir_id_uses_map.remove(&hir_id);
 244                             }
 245                         },
 246                     }
 247                 }
 248                 return;
 249             }
 250
 251             if let Some(hir_id) = path_to_local(recv) {
 252                 if let Some(index) = self.hir_id_uses_map.remove(&hir_id) {
 253                     if self
 254                         .illegal_mutable_capture_ids
 255                         .intersection(&self.current_mutably_captured_ids)
 256                         .next()
 257                         .is_none()
 258                     {
 259                         if let Some(hir_id) = self.current_statement_hir_id {
 260                             self.hir_id_uses_map.insert(hir_id, index);
 261                         }
 262                     } else {
 263                         self.uses[index] = None;
 264                     }
 265                 }
 266             }
 267         }
 268         // Check if the collection is used for anything else
 269         if path_to_local_id(expr, self.target) {
 270             self.seen_other = true;
 271         } else {
 272             walk_expr(self, expr);
 273         }
 274     }
 275 }
 276
 277 enum LoopKind<'tcx> {
 278     Conditional(&'tcx Expr<'tcx>),
 279     Loop,
 280 }
 281
 282 fn check_loop_kind<'tcx>(expr: &Expr<'tcx>) -> Option<LoopKind<'tcx>> {
 283     if let Some(higher::WhileLet { let_expr, .. }) = higher::WhileLet::hir(expr) {
 284         return Some(LoopKind::Conditional(let_expr));
 285     }
 286     if let Some(higher::While { condition, .. }) = higher::While::hir(expr) {
 287         return Some(LoopKind::Conditional(condition));
 288     }
 289     if let Some(higher::ForLoop { arg, .. }) = higher::ForLoop::hir(expr) {
 290         return Some(LoopKind::Conditional(arg));
 291     }
 292     if let ExprKind::Loop { .. } = expr.kind {
 293         return Some(LoopKind::Loop);
 294     }
 295
 296     None
 297 }
 298
 299 impl<'tcx> IterFunctionVisitor<'_, 'tcx> {
 300     fn visit_block_expr(&mut self, expr: &'tcx Expr<'tcx>, hir_id: Option<HirId>) {
 301         self.current_statement_hir_id = hir_id;
 302         self.current_mutably_captured_ids = get_captured_ids(self.cx, self.cx.typeck_results().expr_ty(expr));
 303         self.visit_expr(expr);
 304     }
 305 }
 306
 307 fn get_expr_and_hir_id_from_stmt<'v>(stmt: &'v Stmt<'v>) -> Option<(&'v Expr<'v>, Option<HirId>)> {
 308     match stmt.kind {
 309         StmtKind::Expr(expr) | StmtKind::Semi(expr) => Some((expr, None)),
 310         StmtKind::Item(..) => None,
 311         StmtKind::Local(Local { init, pat, .. }) => {
 312             if let PatKind::Binding(_, hir_id, ..) = pat.kind {
 313                 init.map(|init_expr| (init_expr, Some(hir_id)))
 314             } else {
 315                 init.map(|init_expr| (init_expr, None))
 316             }
 317         },
 318     }
 319 }
 320
 321 struct UsedCountVisitor<'a, 'tcx> {
 322     cx: &'a LateContext<'tcx>,
 323     id: HirId,
 324     count: usize,
 325 }
 326
 327 impl<'a, 'tcx> Visitor<'tcx> for UsedCountVisitor<'a, 'tcx> {
 328     type NestedFilter = nested_filter::OnlyBodies;
 329
 330     fn visit_expr(&mut self, expr: &'tcx Expr<'_>) {
 331         if path_to_local_id(expr, self.id) {
 332             self.count += 1;
 333         } else {
 334             walk_expr(self, expr);
 335         }
 336     }
 337
 338     fn nested_visit_map(&mut self) -> Self::Map {
 339         self.cx.tcx.hir()
 340     }
 341 }
 342
 343 /// Detect the occurrences of calls to `iter` or `into_iter` for the
 344 /// given identifier
 345 fn detect_iter_and_into_iters<'tcx: 'a, 'a>(
 346     block: &'tcx Block<'tcx>,
 347     id: HirId,
 348     cx: &'a LateContext<'tcx>,
 349     captured_ids: HirIdSet,
 350 ) -> Option<Vec<IterFunction>> {
 351     let mut visitor = IterFunctionVisitor {
 352         uses: Vec::new(),
 353         target: id,
 354         seen_other: false,
 355         cx,
 356         current_mutably_captured_ids: HirIdSet::default(),
 357         illegal_mutable_capture_ids: captured_ids,
 358         hir_id_uses_map: FxHashMap::default(),
 359         current_statement_hir_id: None,
 360     };
 361     visitor.visit_block(block);
 362     if visitor.seen_other {
 363         None
 364     } else {
 365         Some(visitor.uses.into_iter().flatten().collect())
 366     }
 367 }
 368
 369 fn get_captured_ids(cx: &LateContext<'_>, ty: Ty<'_>) -> HirIdSet {
 370     fn get_captured_ids_recursive(cx: &LateContext<'_>, ty: Ty<'_>, set: &mut HirIdSet) {
 371         match ty.kind() {
 372             ty::Adt(_, generics) => {
 373                 for generic in *generics {
 374                     if let GenericArgKind::Type(ty) = generic.unpack() {
 375                         get_captured_ids_recursive(cx, ty, set);
 376                     }
 377                 }
 378             },
 379             ty::Closure(def_id, _) => {
 380                 let closure_hir_node = cx.tcx.hir().get_if_local(*def_id).unwrap();
 381                 if let Node::Expr(closure_expr) = closure_hir_node {
 382                     can_move_expr_to_closure(cx, closure_expr)
 383                         .unwrap()
 384                         .into_iter()
 385                         .for_each(|(hir_id, capture_kind)| {
 386                             if matches!(capture_kind, CaptureKind::Ref(Mutability::Mut)) {
 387                                 set.insert(hir_id);
 388                             }
 389                         });
 390                 }
 391             },
 392             _ => (),
 393         }
 394     }
 395
 396     let mut set = HirIdSet::default();
 397
 398     get_captured_ids_recursive(cx, ty, &mut set);
 399
 400     set
 401 }