Improve manual_map

Jarcho · Jarcho · commit 8dbc4b70c371 · 2021-08-04T14:34:40.000-04:00
In some cases check if a borrow made in the scrutinee expression would prevent creating the closure used by `map`
diff --git a/clippy_lints/src/manual_map.rs b/clippy_lints/src/manual_map.rs
@@ -4,12 +4,14 @@ use clippy_utils::source::{snippet_with_applicability, snippet_with_context};
 use clippy_utils::ty::{is_type_diagnostic_item, peel_mid_ty_refs_is_mutable};
 use clippy_utils::{
     can_move_expr_to_closure, in_constant, is_else_clause, is_lang_ctor, is_lint_allowed, path_to_local_id,
-    peel_hir_expr_refs,
+    peel_hir_expr_refs, peel_hir_expr_while, CaptureKind,
 };
 use rustc_ast::util::parser::PREC_POSTFIX;
 use rustc_errors::Applicability;
 use rustc_hir::LangItem::{OptionNone, OptionSome};
-use rustc_hir::{Arm, BindingAnnotation, Block, Expr, ExprKind, HirId, MatchSource, Mutability, Pat, PatKind};
+use rustc_hir::{
+    def::Res, Arm, BindingAnnotation, Block, Expr, ExprKind, HirId, MatchSource, Mutability, Pat, PatKind, Path, QPath,
+};
 use rustc_lint::{LateContext, LateLintPass, LintContext};
 use rustc_middle::lint::in_external_macro;
 use rustc_session::{declare_lint_pass, declare_tool_lint};
@@ -111,10 +113,6 @@ impl LateLintPass<'_> for ManualMap {
                 return;
             }
 
-            if !can_move_expr_to_closure(cx, some_expr) {
-                return;
-            }
-
             // Determine which binding mode to use.
             let explicit_ref = some_pat.contains_explicit_ref_binding();
             let binding_ref = explicit_ref.or_else(|| (ty_ref_count != pat_ref_count).then(|| ty_mutability));
@@ -125,6 +123,32 @@ impl LateLintPass<'_> for ManualMap {
                 None => "",
             };
 
+            match can_move_expr_to_closure(cx, some_expr) {
+                Some(captures) => {
+                    // Check if captures the closure will need conflict with borrows made in the scrutinee.
+                    // TODO: check all the references made in the scrutinee expression. This will require interacting
+                    // with the borrow checker. Currently only `<local>[.<field>]*` is checked for.
+                    if let Some(binding_ref_mutability) = binding_ref {
+                        let e = peel_hir_expr_while(scrutinee, |e| match e.kind {
+                            ExprKind::Field(e, _) | ExprKind::AddrOf(_, _, e) => Some(e),
+                            _ => None,
+                        });
+                        if let ExprKind::Path(QPath::Resolved(None, Path { res: Res::Local(l), .. })) = e.kind {
+                            match captures.get(l) {
+                                Some(CaptureKind::Value | CaptureKind::Ref(Mutability::Mut)) => return,
+                                Some(CaptureKind::Ref(Mutability::Not))
+                                    if binding_ref_mutability == Mutability::Mut =>
+                                {
+                                    return;
+                                }
+                                Some(CaptureKind::Ref(Mutability::Not)) | None => (),
+                            }
+                        }
+                    }
+                },
+                None => return,
+            };
+
             let mut app = Applicability::MachineApplicable;
 
             // Remove address-of expressions from the scrutinee. Either `as_ref` will be called, or
diff --git a/clippy_utils/src/lib.rs b/clippy_utils/src/lib.rs
@@ -67,19 +67,20 @@ use rustc_data_structures::unhash::UnhashMap;
 use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::DefId;
-use rustc_hir::hir_id::HirIdSet;
+use rustc_hir::hir_id::{HirIdMap, HirIdSet};
 use rustc_hir::intravisit::{self, walk_expr, ErasedMap, FnKind, NestedVisitorMap, Visitor};
 use rustc_hir::LangItem::{ResultErr, ResultOk};
 use rustc_hir::{
     def, Arm, BindingAnnotation, Block, Body, Constness, Destination, Expr, ExprKind, FnDecl, GenericArgs, HirId, Impl,
-    ImplItem, ImplItemKind, IsAsync, Item, ItemKind, LangItem, Local, MatchSource, Node, Param, Pat, PatKind, Path,
-    PathSegment, PrimTy, QPath, Stmt, StmtKind, TraitItem, TraitItemKind, TraitRef, TyKind, UnOp,
+    ImplItem, ImplItemKind, IsAsync, Item, ItemKind, LangItem, Local, MatchSource, Mutability, Node, Param, Pat,
+    PatKind, Path, PathSegment, PrimTy, QPath, Stmt, StmtKind, TraitItem, TraitItemKind, TraitRef, TyKind, UnOp,
 };
 use rustc_lint::{LateContext, Level, Lint, LintContext};
 use rustc_middle::hir::exports::Export;
 use rustc_middle::hir::map::Map;
 use rustc_middle::ty as rustc_ty;
-use rustc_middle::ty::{layout::IntegerExt, DefIdTree, Ty, TyCtxt, TypeFoldable};
+use rustc_middle::ty::adjustment::{Adjust, Adjustment, AutoBorrow};
+use rustc_middle::ty::{layout::IntegerExt, DefIdTree, Ty, TyCtxt, TypeAndMut, TypeFoldable};
 use rustc_semver::RustcVersion;
 use rustc_session::Session;
 use rustc_span::hygiene::{ExpnKind, MacroKind};
@@ -603,8 +604,82 @@ pub fn can_move_expr_to_closure_no_visit(
     }
 }
 
-/// Checks if the expression can be moved into a closure as is.
-pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) -> bool {
+/// How a local is captured by a closure
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CaptureKind {
+    Value,
+    Ref(Mutability),
+}
+impl std::ops::BitOr for CaptureKind {
+    type Output = Self;
+    fn bitor(self, rhs: Self) -> Self::Output {
+        match (self, rhs) {
+            (CaptureKind::Value, _) | (_, CaptureKind::Value) => CaptureKind::Value,
+            (CaptureKind::Ref(Mutability::Mut), CaptureKind::Ref(_))
+            | (CaptureKind::Ref(_), CaptureKind::Ref(Mutability::Mut)) => CaptureKind::Ref(Mutability::Mut),
+            (CaptureKind::Ref(Mutability::Not), CaptureKind::Ref(Mutability::Not)) => CaptureKind::Ref(Mutability::Not),
+        }
+    }
+}
+impl std::ops::BitOrAssign for CaptureKind {
+    fn bitor_assign(&mut self, rhs: Self) {
+        *self = *self | rhs;
+    }
+}
+
+/// Given an expression referencing a local, determines how it would be captured in a closure.
+/// Note as this will walk up to parent expressions until the capture can be determined it should
+/// only be used while making a closure somewhere a value is consumed. e.g. a block, match arm, or
+/// function argument (other than a receiver).
+pub fn capture_local_usage(cx: &LateContext<'tcx>, e: &Expr<'_>) -> CaptureKind {
+    debug_assert!(matches!(
+        e.kind,
+        ExprKind::Path(QPath::Resolved(None, Path { res: Res::Local(_), .. }))
+    ));
+
+    let map = cx.tcx.hir();
+    let mut child_id = e.hir_id;
+    let mut capture = CaptureKind::Value;
+
+    for (parent_id, parent) in map.parent_iter(e.hir_id) {
+        if let [Adjustment {
+            kind: Adjust::Deref(_) | Adjust::Borrow(AutoBorrow::Ref(..)),
+            target,
+        }, ref adjust @ ..] = *cx
+            .typeck_results()
+            .adjustments()
+            .get(child_id)
+            .map_or(&[][..], |x| &**x)
+        {
+            if let rustc_ty::RawPtr(TypeAndMut { mutbl: mutability, .. }) | rustc_ty::Ref(_, _, mutability) =
+                *adjust.last().map_or(target, |a| a.target).kind()
+            {
+                return CaptureKind::Ref(mutability);
+            }
+        }
+
+        if let Node::Expr(e) = parent {
+            match e.kind {
+                ExprKind::AddrOf(_, mutability, _) => return CaptureKind::Ref(mutability),
+                ExprKind::Index(..) | ExprKind::Unary(UnOp::Deref, _) => capture = CaptureKind::Ref(Mutability::Not),
+                ExprKind::Assign(lhs, ..) | ExprKind::Assign(_, lhs, _) if lhs.hir_id == child_id => {
+                    return CaptureKind::Ref(Mutability::Mut);
+                },
+                _ => break,
+            }
+        } else {
+            break;
+        }
+
+        child_id = parent_id;
+    }
+
+    capture
+}
+
+/// Checks if the expression can be moved into a closure as is. This will return a list of captures
+/// if so, otherwise, `None`.
+pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) -> Option<HirIdMap<CaptureKind>> {
     struct V<'cx, 'tcx> {
         cx: &'cx LateContext<'tcx>,
         // Stack of potential break targets contained in the expression.
@@ -613,6 +688,9 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
         locals: HirIdSet,
         /// Whether this expression can be turned into a closure.
         allow_closure: bool,
+        /// Locals which need to be captured, and whether they need to be by value, reference, or
+        /// mutable reference.
+        captures: HirIdMap<CaptureKind>,
     }
     impl Visitor<'tcx> for V<'_, 'tcx> {
         type Map = ErasedMap<'tcx>;
@@ -624,13 +702,23 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
             if !self.allow_closure {
                 return;
             }
-            if let ExprKind::Loop(b, ..) = e.kind {
-                self.loops.push(e.hir_id);
-                self.visit_block(b);
-                self.loops.pop();
-            } else {
-                self.allow_closure &= can_move_expr_to_closure_no_visit(self.cx, e, &self.loops, &self.locals);
-                walk_expr(self, e);
+
+            match e.kind {
+                ExprKind::Path(QPath::Resolved(None, &Path { res: Res::Local(l), .. })) => {
+                    if !self.locals.contains(&l) {
+                        let cap = capture_local_usage(self.cx, e);
+                        self.captures.entry(l).and_modify(|e| *e |= cap).or_insert(cap);
+                    }
+                },
+                ExprKind::Loop(b, ..) => {
+                    self.loops.push(e.hir_id);
+                    self.visit_block(b);
+                    self.loops.pop();
+                },
+                _ => {
+                    self.allow_closure &= can_move_expr_to_closure_no_visit(self.cx, e, &self.loops, &self.locals);
+                    walk_expr(self, e);
+                },
             }
         }
 
@@ -646,9 +734,10 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
         allow_closure: true,
         loops: Vec::new(),
         locals: HirIdSet::default(),
+        captures: HirIdMap::default(),
     };
     v.visit_expr(expr);
-    v.allow_closure
+    v.allow_closure.then(|| v.captures)
 }
 
 /// Returns the method names and argument list of nested method call expressions that make up
diff --git a/tests/ui/manual_map_option_2.fixed b/tests/ui/manual_map_option_2.fixed
@@ -7,4 +7,10 @@ fn main() {
             let y = (String::new(), String::new());
             (x, y.0)
         });
+
+    let s = Some(String::new());
+    let _ = match &s {
+        Some(x) => Some((x.clone(), s)),
+        None => None,
+    };
 }
diff --git a/tests/ui/manual_map_option_2.rs b/tests/ui/manual_map_option_2.rs
@@ -10,4 +10,10 @@ fn main() {
         }),
         None => None,
     };
+
+    let s = Some(String::new());
+    let _ = match &s {
+        Some(x) => Some((x.clone(), s)),
+        None => None,
+    };
 }
