map_entry improvements

Suggest using `or_insert_with` when possible

Author: Jarcho

clippy_lints/src/entry.rs

@@ -1,4 +1,5 @@
 use clippy_utils::{
+    can_move_expr_to_closure_no_visit,
     diagnostics::span_lint_and_sugg,
     is_expr_final_block_expr, is_expr_used_or_unified, match_def_path, paths, peel_hir_expr_while,
     source::{reindent_multiline, snippet_indent, snippet_with_applicability, snippet_with_context},
@@ -7,7 +8,7 @@ use clippy_utils::{
 use rustc_errors::Applicability;
 use rustc_hir::{
     intravisit::{walk_expr, ErasedMap, NestedVisitorMap, Visitor},
-    Expr, ExprKind, Guard, Local, Stmt, StmtKind, UnOp,
+    Block, Expr, ExprKind, Guard, HirId, Local, Stmt, StmtKind, UnOp,
 };
 use rustc_lint::{LateContext, LateLintPass};
 use rustc_session::{declare_lint_pass, declare_tool_lint};
@@ -78,13 +79,13 @@ impl<'tcx> LateLintPass<'tcx> for HashMapPass {
         let sugg = if let Some(else_expr) = else_expr {
             // if .. { .. } else { .. }
             let else_search = match find_insert_calls(cx, &contains_expr, else_expr) {
-                Some(search) if !(then_search.insertions.is_empty() && search.insertions.is_empty()) => search,
+                Some(search) if !(then_search.edits.is_empty() && search.edits.is_empty()) => search,
                 _ => return,
             };
 
-            if then_search.insertions.is_empty() || else_search.insertions.is_empty() {
+            if then_search.edits.is_empty() || else_search.edits.is_empty() {
                 // if .. { insert } else { .. } or if .. { .. } else { then } of
-                let (then_str, else_str, entry_kind) = if else_search.insertions.is_empty() {
+                let (then_str, else_str, entry_kind) = if else_search.edits.is_empty() {
                     if contains_expr.negated {
                         (
                             then_search.snippet_vacant(cx, then_expr.span, &mut app),
@@ -152,37 +153,48 @@ impl<'tcx> LateLintPass<'tcx> for HashMapPass {
                     indent = indent_str,
                 )
             }
-        } else if then_search.insertions.is_empty() || !contains_expr.negated {
+        } else if then_search.edits.is_empty() {
+            // no insertions
             return;
         } else {
             // if .. { insert }
-            match then_search.as_single_insertion() {
-                Some(insertion) if !insertion.value.can_have_side_effects() => {
-                    format!(
-                        "{}.entry({}).or_insert({});",
-                        map_str,
-                        key_str,
-                        snippet_with_context(cx, insertion.value.span, insertion.call.span.ctxt(), "..", &mut app).0,
+            if !then_search.allow_insert_closure {
+                let (body_str, entry_kind) = if contains_expr.negated {
+                    (then_search.snippet_vacant(cx, then_expr.span, &mut app), "Vacant(e)")
+                } else {
+                    (
+                        then_search.snippet_occupied(cx, then_expr.span, &mut app),
+                        "Occupied(mut e)",
                     )
-                },
-                _ => {
-                    let (body_str, entry_kind) = if contains_expr.negated {
-                        (then_search.snippet_vacant(cx, then_expr.span, &mut app), "Vacant(e)")
+                };
+                format!(
+                    "if let {}::{} = {}.entry({}) {}",
+                    map_ty.entry_path(),
+                    entry_kind,
+                    map_str,
+                    key_str,
+                    body_str,
+                )
+            } else if let Some(insertion) = then_search.as_single_insertion() {
+                let value_str = snippet_with_context(cx, insertion.value.span, then_expr.span.ctxt(), "..", &mut app).0;
+                if contains_expr.negated {
+                    if insertion.value.can_have_side_effects() {
+                        format!("{}.entry({}).or_insert_with(|| {});", map_str, key_str, value_str)
                     } else {
-                        (
-                            then_search.snippet_occupied(cx, then_expr.span, &mut app),
-                            "Occupied(mut e)",
-                        )
-                    };
-                    format!(
-                        "if let {}::{} = {}.entry({}) {}",
-                        map_ty.entry_path(),
-                        entry_kind,
-                        map_str,
-                        key_str,
-                        body_str,
-                    )
-                },
+                        format!("{}.entry({}).or_insert({});", map_str, key_str, value_str)
+                    }
+                } else {
+                    // Todo: if let Some(v) = map.get_mut(k)
+                    return;
+                }
+            } else {
+                let block_str = then_search.snippet_closure(cx, then_expr.span, &mut app);
+                if contains_expr.negated {
+                    format!("{}.entry({}).or_insert_with(|| {});", map_str, key_str, block_str)
+                } else {
+                    // Todo: if let Some(v) = map.get_mut(k)
+                    return;
+                }
             }
         };
 
@@ -281,6 +293,19 @@ fn try_parse_insert(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) -> Option<Inse
     }
 }
 
+/// An edit that will need to be made to move the expression to use the entry api
+#[derive(Clone, Copy)]
+enum Edit<'tcx> {
+    /// A semicolon that needs to be removed. Used to create a closure for `insert_with`.
+    RemoveSemi(Span),
+    /// An insertion into the map.
+    Insertion(Insertion<'tcx>),
+}
+impl Edit<'tcx> {
+    fn as_insertion(self) -> Option<Insertion<'tcx>> {
+        if let Self::Insertion(i) = self { Some(i) } else { None }
+    }
+}
 #[derive(Clone, Copy)]
 struct Insertion<'tcx> {
     call: &'tcx Expr<'tcx>,
@@ -303,24 +328,40 @@ struct InsertSearcher<'cx, 'i, 'tcx> {
     key: &'tcx Expr<'tcx>,
     /// The context of the top level block. All insert calls must be in the same context.
     ctxt: SyntaxContext,
+    /// Whether this expression can be safely moved into a closure.
+    allow_insert_closure: bool,
     /// Whether this expression can use the entry api.
     can_use_entry: bool,
+    /// Whether this expression is the final expression in this code path. This may be a statement.
+    in_tail_pos: bool,
     // A single insert expression has a slightly different suggestion.
     is_single_insert: bool,
     is_map_used: bool,
-    insertions: &'i mut Vec<Insertion<'tcx>>,
+    edits: &'i mut Vec<Edit<'tcx>>,
+    loops: Vec<HirId>,
 }
 impl<'tcx> InsertSearcher<'_, '_, 'tcx> {
     /// Visit the expression as a branch in control flow. Multiple insert calls can be used, but
     /// only if they are on separate code paths. This will return whether the map was used in the
     /// given expression.
     fn visit_cond_arm(&mut self, e: &'tcx Expr<'_>) -> bool {
         let is_map_used = self.is_map_used;
+        let in_tail_pos = self.in_tail_pos;
         self.visit_expr(e);
         let res = self.is_map_used;
         self.is_map_used = is_map_used;
+        self.in_tail_pos = in_tail_pos;
         res
     }
+
+    /// Visits an expression which is not itself in a tail position, but other sibling expressions
+    /// may be. e.g. if conditions
+    fn visit_non_tail_expr(&mut self, e: &'tcx Expr<'_>) {
+        let in_tail_pos = self.in_tail_pos;
+        self.in_tail_pos = false;
+        self.visit_expr(e);
+        self.in_tail_pos = in_tail_pos;
+    }
 }
 impl<'tcx> Visitor<'tcx> for InsertSearcher<'_, '_, 'tcx> {
     type Map = ErasedMap<'tcx>;
@@ -330,17 +371,63 @@ impl<'tcx> Visitor<'tcx> for InsertSearcher<'_, '_, 'tcx> {
 
     fn visit_stmt(&mut self, stmt: &'tcx Stmt<'_>) {
         match stmt.kind {
-            StmtKind::Semi(e) | StmtKind::Expr(e) => self.visit_expr(e),
+            StmtKind::Semi(e) => {
+                self.visit_expr(e);
+
+                if self.in_tail_pos && self.allow_insert_closure {
+                    // The spans are used to slice the top level expression into multiple parts. This requires that
+                    // they all come from the same part of the source code.
+                    if stmt.span.ctxt() == self.ctxt && e.span.ctxt() == self.ctxt {
+                        self.edits
+                            .push(Edit::RemoveSemi(stmt.span.trim_start(e.span).unwrap_or(DUMMY_SP)));
+                    } else {
+                        self.allow_insert_closure = false;
+                    }
+                }
+            },
+            StmtKind::Expr(e) => self.visit_expr(e),
             StmtKind::Local(Local { init: Some(e), .. }) => {
+                self.allow_insert_closure &= !self.in_tail_pos;
+                self.in_tail_pos = false;
                 self.is_single_insert = false;
                 self.visit_expr(e);
             },
             _ => {
+                self.allow_insert_closure &= !self.in_tail_pos;
                 self.is_single_insert = false;
             },
         }
     }
 
+    fn visit_block(&mut self, block: &'tcx Block<'_>) {
+        // If the block is in a tail position, then the last expression (possibly a statement) is in the
+        // tail position. The rest, however, are not.
+        match (block.stmts, block.expr) {
+            ([], None) => {
+                self.allow_insert_closure &= !self.in_tail_pos;
+            },
+            ([], Some(expr)) => self.visit_expr(expr),
+            (stmts, Some(expr)) => {
+                let in_tail_pos = self.in_tail_pos;
+                self.in_tail_pos = false;
+                for stmt in stmts {
+                    self.visit_stmt(stmt);
+                }
+                self.in_tail_pos = in_tail_pos;
+                self.visit_expr(expr);
+            },
+            ([stmts @ .., stmt], None) => {
+                let in_tail_pos = self.in_tail_pos;
+                self.in_tail_pos = false;
+                for stmt in stmts {
+                    self.visit_stmt(stmt);
+                }
+                self.in_tail_pos = in_tail_pos;
+                self.visit_stmt(stmt);
+            },
+        }
+    }
+
     fn visit_expr(&mut self, expr: &'tcx Expr<'_>) {
         if !self.can_use_entry {
             return;
@@ -357,15 +444,16 @@ impl<'tcx> Visitor<'tcx> for InsertSearcher<'_, '_, 'tcx> {
                     return;
                 }
 
-                self.insertions.push(Insertion {
+                self.edits.push(Edit::Insertion(Insertion {
                     call: expr,
                     value: insert_expr.value,
-                });
+                }));
                 self.is_map_used = true;
+                self.allow_insert_closure &= self.in_tail_pos;
 
                 // The value doesn't affect whether there is only a single insert expression.
                 let is_single_insert = self.is_single_insert;
-                self.visit_expr(insert_expr.value);
+                self.visit_non_tail_expr(insert_expr.value);
                 self.is_single_insert = is_single_insert;
             },
             _ if SpanlessEq::new(self.cx).eq_expr(self.map, expr) => {
@@ -374,39 +462,46 @@ impl<'tcx> Visitor<'tcx> for InsertSearcher<'_, '_, 'tcx> {
             _ => match expr.kind {
                 ExprKind::If(cond_expr, then_expr, Some(else_expr)) => {
                     self.is_single_insert = false;
-                    self.visit_expr(cond_expr);
+                    self.visit_non_tail_expr(cond_expr);
                     // Each branch may contain it's own insert expression.
                     let mut is_map_used = self.visit_cond_arm(then_expr);
                     is_map_used |= self.visit_cond_arm(else_expr);
                     self.is_map_used = is_map_used;
                 },
                 ExprKind::Match(scrutinee_expr, arms, _) => {
                     self.is_single_insert = false;
-                    self.visit_expr(scrutinee_expr);
+                    self.visit_non_tail_expr(scrutinee_expr);
                     // Each branch may contain it's own insert expression.
                     let mut is_map_used = self.is_map_used;
                     for arm in arms {
                         if let Some(Guard::If(guard) | Guard::IfLet(_, guard)) = arm.guard {
-                            self.visit_expr(guard)
+                            self.visit_non_tail_expr(guard)
                         }
                         is_map_used |= self.visit_cond_arm(arm.body);
                     }
                     self.is_map_used = is_map_used;
                 },
                 ExprKind::Loop(block, ..) => {
+                    self.loops.push(expr.hir_id);
+                    self.allow_insert_closure &= !self.in_tail_pos;
                     // Don't allow insertions inside of a loop.
-                    let insertions_len = self.insertions.len();
+                    let edit_len = self.edits.len();
                     self.visit_block(block);
-                    if self.insertions.len() != insertions_len {
+                    if self.edits.len() != edit_len {
                         self.can_use_entry = false;
                     }
+                    self.loops.pop();
                 },
                 ExprKind::Block(block, _) => self.visit_block(block),
                 ExprKind::InlineAsm(_) | ExprKind::LlvmInlineAsm(_) => {
                     self.can_use_entry = false;
                 },
                 _ => {
+                    self.allow_insert_closure &= !self.in_tail_pos;
+                    self.allow_insert_closure &= can_move_expr_to_closure_no_visit(self.cx, expr, &self.loops);
+                    // Sub expressions are no longer in the tail position.
                     self.is_single_insert = false;
+                    self.in_tail_pos = false;
                     walk_expr(self, expr);
                 },
             },
@@ -415,18 +510,19 @@ impl<'tcx> Visitor<'tcx> for InsertSearcher<'_, '_, 'tcx> {
 }
 
 struct InsertSearchResults<'tcx> {
-    insertions: Vec<Insertion<'tcx>>,
+    edits: Vec<Edit<'tcx>>,
+    allow_insert_closure: bool,
     is_single_insert: bool,
 }
 impl InsertSearchResults<'tcx> {
     fn as_single_insertion(&self) -> Option<Insertion<'tcx>> {
-        self.is_single_insert.then(|| self.insertions[0])
+        self.is_single_insert.then(|| self.edits[0].as_insertion().unwrap())
     }
 
     fn snippet_occupied(&self, cx: &LateContext<'_>, mut span: Span, app: &mut Applicability) -> String {
         let ctxt = span.ctxt();
         let mut res = String::new();
-        for insertion in self.insertions.iter() {
+        for insertion in self.edits.iter().filter_map(|e| e.as_insertion()) {
             res.push_str(&snippet_with_applicability(
                 cx,
                 span.until(insertion.call.span),
@@ -451,7 +547,7 @@ impl InsertSearchResults<'tcx> {
     fn snippet_vacant(&self, cx: &LateContext<'_>, mut span: Span, app: &mut Applicability) -> String {
         let ctxt = span.ctxt();
         let mut res = String::new();
-        for insertion in self.insertions.iter() {
+        for insertion in self.edits.iter().filter_map(|e| e.as_insertion()) {
             res.push_str(&snippet_with_applicability(
                 cx,
                 span.until(insertion.call.span),
@@ -485,27 +581,57 @@ impl InsertSearchResults<'tcx> {
         res.push_str(&snippet_with_applicability(cx, span, "..", app));
         res
     }
+
+    fn snippet_closure(&self, cx: &LateContext<'_>, mut span: Span, app: &mut Applicability) -> String {
+        let ctxt = span.ctxt();
+        let mut res = String::new();
+        for edit in &self.edits {
+            match *edit {
+                Edit::Insertion(insertion) => {
+                    res.push_str(&snippet_with_applicability(
+                        cx,
+                        span.until(insertion.call.span),
+                        "..",
+                        app,
+                    ));
+                    res.push_str(&snippet_with_context(cx, insertion.value.span, ctxt, "..", app).0);
+                    span = span.trim_start(insertion.call.span).unwrap_or(DUMMY_SP);
+                },
+                Edit::RemoveSemi(semi_span) => {
+                    res.push_str(&snippet_with_applicability(cx, span.until(semi_span), "..", app));
+                    span = span.trim_start(semi_span).unwrap_or(DUMMY_SP);
+                },
+            }
+        }
+        res.push_str(&snippet_with_applicability(cx, span, "..", app));
+        res
+    }
 }
 fn find_insert_calls(
     cx: &LateContext<'tcx>,
     contains_expr: &ContainsExpr<'tcx>,
     expr: &'tcx Expr<'_>,
 ) -> Option<InsertSearchResults<'tcx>> {
-    let mut insertions = Vec::new();
+    let mut edits = Vec::new();
     let mut s = InsertSearcher {
         cx,
         map: contains_expr.map,
         key: contains_expr.key,
         ctxt: expr.span.ctxt(),
-        insertions: &mut insertions,
+        edits: &mut edits,
         is_map_used: false,
+        allow_insert_closure: true,
         can_use_entry: true,
+        in_tail_pos: true,
         is_single_insert: true,
+        loops: Vec::new(),
     };
     s.visit_expr(expr);
+    let allow_insert_closure = s.allow_insert_closure;
     let is_single_insert = s.is_single_insert;
     s.can_use_entry.then(|| InsertSearchResults {
-        insertions,
+        edits,
+        allow_insert_closure,
         is_single_insert,
     })
 }