Don't lint vec_init_then_push when further extended

clippy_lints/src/vec_init_then_push.rs

@@ -1,19 +1,30 @@
 use clippy_utils::diagnostics::span_lint_and_sugg;
 use clippy_utils::higher::{get_vec_init_kind, VecInitKind};
 use clippy_utils::source::snippet;
-use clippy_utils::{path_to_local, path_to_local_id};
-use if_chain::if_chain;
+use clippy_utils::visitors::for_each_local_use_after_expr;
+use clippy_utils::{get_parent_expr, path_to_local_id};
+use core::ops::ControlFlow;
 use rustc_errors::Applicability;
-use rustc_hir::{BindingAnnotation, Block, Expr, ExprKind, HirId, Local, PatKind, Stmt, StmtKind};
+use rustc_hir::def::Res;
+use rustc_hir::{
+    BindingAnnotation, Block, Expr, ExprKind, HirId, Local, Mutability, PatKind, QPath, Stmt, StmtKind, UnOp,
+};
 use rustc_lint::{LateContext, LateLintPass, LintContext};
 use rustc_middle::lint::in_external_macro;
 use rustc_session::{declare_tool_lint, impl_lint_pass};
-use rustc_span::Span;
+use rustc_span::{Span, Symbol};
 
 declare_clippy_lint! {
     /// ### What it does
     /// Checks for calls to `push` immediately after creating a new `Vec`.
     ///
+    /// If the `Vec` is created using `with_capacity` this will only lint if the capacity is a
+    /// constant and the number of pushes is greater than or equal to the initial capacity.
+    ///
+    /// If the `Vec` is extended after the initial sequence of pushes and it was default initialized
+    /// then this will only lint after there were at least four pushes. This number may change in
+    /// the future.
+    ///
     /// ### Why is this bad?
     /// The `vec![]` macro is both more performant and easier to read than
     /// multiple `push` calls.
@@ -43,26 +54,88 @@ pub struct VecInitThenPush {
 struct VecPushSearcher {
     local_id: HirId,
     init: VecInitKind,
-    lhs_is_local: bool,
-    lhs_span: Span,
+    lhs_is_let: bool,
+    let_ty_span: Option<Span>,
+    name: Symbol,
     err_span: Span,
-    found: u64,
+    found: u128,
+    last_push_expr: HirId,
 }
 impl VecPushSearcher {
     fn display_err(&self, cx: &LateContext<'_>) {
-        match self.init {
+        let required_pushes_before_extension = match self.init {
             _ if self.found == 0 => return,
-            VecInitKind::WithLiteralCapacity(x) if x > self.found => return,
+            VecInitKind::WithConstCapacity(x) if x > self.found => return,
+            VecInitKind::WithConstCapacity(x) => x,
             VecInitKind::WithExprCapacity(_) => return,
-            _ => (),
+            _ => 3,
         };
 
-        let mut s = if self.lhs_is_local {
+        let mut needs_mut = false;
+        let res = for_each_local_use_after_expr(cx, self.local_id, self.last_push_expr, |e| {
+            let Some(parent) = get_parent_expr(cx, e) else {
+                return ControlFlow::Continue(())
+            };
+            let adjusted_ty = cx.typeck_results().expr_ty_adjusted(e);
+            let adjusted_mut = adjusted_ty.ref_mutability().unwrap_or(Mutability::Not);
+            needs_mut |= adjusted_mut == Mutability::Mut;
+            match parent.kind {
+                ExprKind::AddrOf(_, Mutability::Mut, _) => {
+                    needs_mut = true;
+                    return ControlFlow::Break(true);
+                },
+                ExprKind::Unary(UnOp::Deref, _) | ExprKind::Index(..) if !needs_mut => {
+                    let mut last_place = parent;
+                    while let Some(parent) = get_parent_expr(cx, parent) {
+                        if matches!(parent.kind, ExprKind::Unary(UnOp::Deref, _) | ExprKind::Field(..))
+                            || matches!(parent.kind, ExprKind::Index(e, _) if e.hir_id == last_place.hir_id)
+                        {
+                            last_place = parent;
+                        } else {
+                            break;
+                        }
+                    }
+                    needs_mut |= cx.typeck_results().expr_ty_adjusted(last_place).ref_mutability()
+                        == Some(Mutability::Mut)
+                        || get_parent_expr(cx, last_place)
+                            .map_or(false, |e| matches!(e.kind, ExprKind::AddrOf(_, Mutability::Mut, _)));
+                },
+                ExprKind::MethodCall(_, [recv, ..], _)
+                    if recv.hir_id == e.hir_id
+                        && adjusted_mut == Mutability::Mut
+                        && !adjusted_ty.peel_refs().is_slice() =>
+                {
+                    // No need to set `needs_mut` to true. The receiver will be either explicitly borrowed, or it will
+                    // be implicitly borrowed via an adjustment. Both of these cases are already handled by this point.
+                    return ControlFlow::Break(true);
+                },
+                ExprKind::Assign(lhs, ..) if e.hir_id == lhs.hir_id => {
+                    needs_mut = true;
+                    return ControlFlow::Break(false);
+                },
+                _ => (),
+            }
+            ControlFlow::Continue(())
+        });
+
+        // Avoid allocating small `Vec`s when they'll be extended right after.
+        if res == ControlFlow::Break(true) && self.found <= required_pushes_before_extension {
+            return;
+        }
+
+        let mut s = if self.lhs_is_let {
             String::from("let ")
         } else {
             String::new()
         };
-        s.push_str(&snippet(cx, self.lhs_span, ".."));
+        if needs_mut {
+            s.push_str("mut ");
+        }
+        s.push_str(self.name.as_str());
+        if let Some(span) = self.let_ty_span {
+            s.push_str(": ");
+            s.push_str(&snippet(cx, span, "_"));
+        }
         s.push_str(" = vec![..];");
 
         span_lint_and_sugg(
@@ -83,60 +156,63 @@ impl<'tcx> LateLintPass<'tcx> for VecInitThenPush {
     }
 
     fn check_local(&mut self, cx: &LateContext<'tcx>, local: &'tcx Local<'tcx>) {
-        if_chain! {
-            if !in_external_macro(cx.sess(), local.span);
-            if let Some(init) = local.init;
-            if let PatKind::Binding(BindingAnnotation::Mutable, id, _, None) = local.pat.kind;
-            if let Some(init_kind) = get_vec_init_kind(cx, init);
-            then {
-                self.searcher = Some(VecPushSearcher {
-                        local_id: id,
-                        init: init_kind,
-                        lhs_is_local: true,
-                        lhs_span: local.ty.map_or(local.pat.span, |t| local.pat.span.to(t.span)),
-                        err_span: local.span,
-                        found: 0,
-                    });
-            }
+        if let Some(init_expr) = local.init
+            && let PatKind::Binding(BindingAnnotation::Mutable, id, name, None) = local.pat.kind
+            && !in_external_macro(cx.sess(), local.span)
+            && let Some(init) = get_vec_init_kind(cx, init_expr)
+            && !matches!(init, VecInitKind::WithExprCapacity(_))
+        {
+            self.searcher = Some(VecPushSearcher {
+                local_id: id,
+                init,
+                lhs_is_let: true,
+                name: name.name,
+                let_ty_span: local.ty.map(|ty| ty.span),
+                err_span: local.span,
+                found: 0,
+                last_push_expr: init_expr.hir_id,
+            });
         }
     }
 
     fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
-        if_chain! {
-            if self.searcher.is_none();
-            if !in_external_macro(cx.sess(), expr.span);
-            if let ExprKind::Assign(left, right, _) = expr.kind;
-            if let Some(id) = path_to_local(left);
-            if let Some(init_kind) = get_vec_init_kind(cx, right);
-            then {
-                self.searcher = Some(VecPushSearcher {
-                    local_id: id,
-                    init: init_kind,
-                    lhs_is_local: false,
-                    lhs_span: left.span,
-                    err_span: expr.span,
-                    found: 0,
-                });
-            }
+        if self.searcher.is_none()
+            && let ExprKind::Assign(left, right, _) = expr.kind
+            && let ExprKind::Path(QPath::Resolved(None, path)) = left.kind
+            && let [name] = &path.segments
+            && let Res::Local(id) = path.res
+            && !in_external_macro(cx.sess(), expr.span)
+            && let Some(init) = get_vec_init_kind(cx, right)
+            && !matches!(init, VecInitKind::WithExprCapacity(_))
+        {
+            self.searcher = Some(VecPushSearcher {
+                local_id: id,
+                init,
+                lhs_is_let: false,
+                let_ty_span: None,
+                name: name.ident.name,
+                err_span: expr.span,
+                found: 0,
+                last_push_expr: expr.hir_id,
+            });
         }
     }
 
     fn check_stmt(&mut self, cx: &LateContext<'tcx>, stmt: &'tcx Stmt<'_>) {
         if let Some(searcher) = self.searcher.take() {
-            if_chain! {
-                if let StmtKind::Expr(expr) | StmtKind::Semi(expr) = stmt.kind;
-                if let ExprKind::MethodCall(path, [self_arg, _], _) = expr.kind;
-                if path_to_local_id(self_arg, searcher.local_id);
-                if path.ident.name.as_str() == "push";
-                then {
-                    self.searcher = Some(VecPushSearcher {
-                        found: searcher.found + 1,
-                        err_span: searcher.err_span.to(stmt.span),
-                        .. searcher
-                    });
-                } else {
-                    searcher.display_err(cx);
-                }
+            if let StmtKind::Expr(expr) | StmtKind::Semi(expr) = stmt.kind
+                && let ExprKind::MethodCall(name, [self_arg, _], _) = expr.kind
+                && path_to_local_id(self_arg, searcher.local_id)
+                && name.ident.as_str() == "push"
+            {
+                self.searcher = Some(VecPushSearcher {
+                    found: searcher.found + 1,
+                    err_span: searcher.err_span.to(stmt.span),
+                    last_push_expr: expr.hir_id,
+                    .. searcher
+                });
+            } else {
+                searcher.display_err(cx);
             }
         }
     }

clippy_utils/src/higher.rs

@@ -2,10 +2,11 @@
 
 #![deny(clippy::missing_docs_in_private_items)]
 
+use crate::consts::{constant_simple, Constant};
 use crate::ty::is_type_diagnostic_item;
 use crate::{is_expn_of, match_def_path, paths};
 use if_chain::if_chain;
-use rustc_ast::ast::{self, LitKind};
+use rustc_ast::ast;
 use rustc_hir as hir;
 use rustc_hir::{Arm, Block, Expr, ExprKind, HirId, LoopSource, MatchSource, Node, Pat, QPath};
 use rustc_lint::LateContext;
@@ -431,7 +432,7 @@ pub enum VecInitKind {
     /// `Vec::default()` or `Default::default()`
     Default,
     /// `Vec::with_capacity(123)`
-    WithLiteralCapacity(u64),
+    WithConstCapacity(u128),
     /// `Vec::with_capacity(slice.len())`
     WithExprCapacity(HirId),
 }
@@ -449,15 +450,11 @@ pub fn get_vec_init_kind<'tcx>(cx: &LateContext<'tcx>, expr: &'tcx Expr<'tcx>) -
                     return Some(VecInitKind::Default);
                 } else if name.ident.name.as_str() == "with_capacity" {
                     let arg = args.get(0)?;
-                    if_chain! {
-                        if let ExprKind::Lit(lit) = &arg.kind;
-                        if let LitKind::Int(num, _) = lit.node;
-                        then {
-                            return Some(VecInitKind::WithLiteralCapacity(num.try_into().ok()?));
-                        }
-                    }
-                    return Some(VecInitKind::WithExprCapacity(arg.hir_id));
-                }
+                    return match constant_simple(cx, cx.typeck_results(), arg) {
+                        Some(Constant::Int(num)) => Some(VecInitKind::WithConstCapacity(num)),
+                        _ => Some(VecInitKind::WithExprCapacity(arg.hir_id)),
+                    };
+                };
             },
             ExprKind::Path(QPath::Resolved(_, path))
                 if match_def_path(cx, path.res.opt_def_id()?, &paths::DEFAULT_TRAIT_METHOD)

clippy_utils/src/visitors.rs

@@ -1,4 +1,4 @@
-use crate::path_to_local_id;
+use crate::{get_enclosing_block, path_to_local_id};
 use core::ops::ControlFlow;
 use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
@@ -436,3 +436,61 @@ pub fn for_each_value_source<'tcx, B>(
         _ => f(e),
     }
 }
+
+/// Runs the given function for each path expression referencing the given local which occur after
+/// the given expression.
+pub fn for_each_local_use_after_expr<'tcx, B>(
+    cx: &LateContext<'tcx>,
+    local_id: HirId,
+    expr_id: HirId,
+    f: impl FnMut(&'tcx Expr<'tcx>) -> ControlFlow<B>,
+) -> ControlFlow<B> {
+    struct V<'cx, 'tcx, F, B> {
+        cx: &'cx LateContext<'tcx>,
+        local_id: HirId,
+        expr_id: HirId,
+        found: bool,
+        res: ControlFlow<B>,
+        f: F,
+    }
+    impl<'cx, 'tcx, F: FnMut(&'tcx Expr<'tcx>) -> ControlFlow<B>, B> Visitor<'tcx> for V<'cx, 'tcx, F, B> {
+        type NestedFilter = nested_filter::OnlyBodies;
+        fn nested_visit_map(&mut self) -> Self::Map {
+            self.cx.tcx.hir()
+        }
+
+        fn visit_expr(&mut self, e: &'tcx Expr<'tcx>) {
+            if !self.found {
+                if e.hir_id == self.expr_id {
+                    self.found = true;
+                } else {
+                    walk_expr(self, e);
+                }
+                return;
+            }
+            if self.res.is_break() {
+                return;
+            }
+            if path_to_local_id(e, self.local_id) {
+                self.res = (self.f)(e);
+            } else {
+                walk_expr(self, e);
+            }
+        }
+    }
+
+    if let Some(b) = get_enclosing_block(cx, local_id) {
+        let mut v = V {
+            cx,
+            local_id,
+            expr_id,
+            found: false,
+            res: ControlFlow::Continue(()),
+            f,
+        };
+        v.visit_block(b);
+        v.res
+    } else {
+        ControlFlow::Continue(())
+    }
+}