Merge pull request #689 from mcarton/entries

Improve the `MAP_ENTRY` lint

src/entry.rs

@@ -1,5 +1,6 @@
 use rustc::lint::*;
 use rustc_front::hir::*;
+use rustc_front::intravisit::{Visitor, walk_expr, walk_block};
 use syntax::codemap::Span;
 use utils::SpanlessEq;
 use utils::{BTREEMAP_PATH, HASHMAP_PATH};
@@ -41,73 +42,108 @@ impl LintPass for HashMapLint {
 
 impl LateLintPass for HashMapLint {
     fn check_expr(&mut self, cx: &LateContext, expr: &Expr) {
-        if_let_chain! {
+        if let ExprIf(ref check, ref then_block, ref else_block) = expr.node {
-            [
+            if let ExprUnary(UnOp::UnNot, ref check) = check.node {
-                let ExprIf(ref check, ref then, _) = expr.node,
+                if let Some((ty, map, key)) = check_cond(cx, check) {
-                let ExprUnary(UnOp::UnNot, ref check) = check.node,
+                    // in case of `if !m.contains_key(&k) { m.insert(k, v); }`
-                let ExprMethodCall(ref name, _, ref params) = check.node,
+                    // we can give a better error message
-                params.len() >= 2,
+                    let sole_expr = else_block.is_none() &&
-                name.node.as_str() == "contains_key"
+                        if then_block.expr.is_some() { 1 } else { 0 } + then_block.stmts.len() == 1;
-            ], {
+
-                let key = match params[1].node {
+                    let mut visitor = InsertVisitor {
-                    ExprAddrOf(_, ref key) => key,
+                        cx: cx,
-                    _ => return
+                        span: expr.span,
+                        ty: ty,
+                        map: map,
+                        key: key,
+                        sole_expr: sole_expr,
                     };
 
+                    walk_block(&mut visitor, then_block);
+                }
+            } else if let Some(ref else_block) = *else_block {
+                if let Some((ty, map, key)) = check_cond(cx, check) {
+                    let mut visitor = InsertVisitor {
+                        cx: cx,
+                        span: expr.span,
+                        ty: ty,
+                        map: map,
+                        key: key,
+                        sole_expr: false,
+                    };
+
+                    walk_expr(&mut visitor, else_block);
+                }
+            }
+        }
+    }
+}
+
+fn check_cond<'a, 'tcx, 'b>(cx: &'a LateContext<'a, 'tcx>, check: &'b Expr) -> Option<(&'static str, &'b Expr, &'b Expr)> {
+    if_let_chain! {[
+        let ExprMethodCall(ref name, _, ref params) = check.node,
+        params.len() >= 2,
+        name.node.as_str() == "contains_key",
+        let ExprAddrOf(_, ref key) = params[1].node
+    ], {
         let map = &params[0];
         let obj_ty = walk_ptrs_ty(cx.tcx.expr_ty(map));
 
-                let kind = if match_type(cx, obj_ty, &BTREEMAP_PATH) {
+        return if match_type(cx, obj_ty, &BTREEMAP_PATH) {
-                    "BTreeMap"
+            Some(("BTreeMap", map, key))
         }
         else if match_type(cx, obj_ty, &HASHMAP_PATH) {
-                    "HashMap"
+            Some(("HashMap", map, key))
         }
         else {
-                    return
+            None
         };
+    }}
 
-                let sole_expr = if then.expr.is_some() { 1 } else { 0 } + then.stmts.len() == 1;
+    None
-
-                if let Some(ref then) = then.expr {
-                    check_for_insert(cx, expr.span, map, key, then, sole_expr, kind);
 }
 
-                for stmt in &then.stmts {
+struct InsertVisitor<'a, 'tcx: 'a, 'b> {
-                    if let StmtSemi(ref stmt, _) = stmt.node {
+    cx: &'a LateContext<'a, 'tcx>,
-                        check_for_insert(cx, expr.span, map, key, stmt, sole_expr, kind);
+    span: Span,
-                    }
+    ty: &'static str,
-                }
+    map: &'b Expr,
-            }
+    key: &'b Expr,
-        }
+    sole_expr: bool,
-    }
 }
 
-fn check_for_insert(cx: &LateContext, span: Span, map: &Expr, key: &Expr, expr: &Expr, sole_expr: bool, kind: &str) {
+impl<'a, 'tcx, 'v, 'b> Visitor<'v> for InsertVisitor<'a, 'tcx, 'b> {
-    if_let_chain! {
+    fn visit_expr(&mut self, expr: &'v Expr) {
-        [
+        if_let_chain! {[
             let ExprMethodCall(ref name, _, ref params) = expr.node,
             params.len() == 3,
             name.node.as_str() == "insert",
-            get_item_name(cx, map) == get_item_name(cx, &*params[0]),
+            get_item_name(self.cx, self.map) == get_item_name(self.cx, &*params[0]),
-            SpanlessEq::new(cx).eq_expr(key, &params[1])
+            SpanlessEq::new(self.cx).eq_expr(self.key, &params[1])
         ], {
-            let help = if sole_expr {
+
-                format!("{}.entry({}).or_insert({})",
+            span_lint_and_then(self.cx, MAP_ENTRY, self.span,
-                        snippet(cx, map.span, "map"),
+                               &format!("usage of `contains_key` followed by `insert` on `{}`", self.ty), |db| {
-                        snippet(cx, params[1].span, ".."),
+                if self.sole_expr {
-                        snippet(cx, params[2].span, ".."))
+                    let help = format!("{}.entry({}).or_insert({})",
+                                       snippet(self.cx, self.map.span, "map"),
+                                       snippet(self.cx, params[1].span, ".."),
+                                       snippet(self.cx, params[2].span, ".."));
+
+                    db.span_suggestion(self.span, "Consider using", help);
                 }
                 else {
-                format!("{}.entry({})",
+                    let help = format!("Consider using `{}.entry({})`",
-                        snippet(cx, map.span, "map"),
+                                       snippet(self.cx, self.map.span, "map"),
-                        snippet(cx, params[1].span, ".."))
+                                       snippet(self.cx, params[1].span, ".."));
-            };
 
-            span_lint_and_then(cx, MAP_ENTRY, span,
+                    db.span_note(self.span, &help);
-                               &format!("usage of `contains_key` followed by `insert` on `{}`", kind), |db| {
+                }
-                db.span_suggestion(span, "Consider using", help);
             });
+        }}
+
+        if !self.sole_expr {
+            walk_expr(self, expr);
         }
     }
 }

tests/compile-fail/entry.rs

@@ -19,29 +19,37 @@ fn insert_if_absent0<K: Eq + Hash, V>(m: &mut HashMap<K, V>, k: K, v: V) {
 fn insert_if_absent1<K: Eq + Hash, V>(m: &mut HashMap<K, V>, k: K, v: V) {
     if !m.contains_key(&k) { foo(); m.insert(k, v); }
     //~^ ERROR usage of `contains_key` followed by `insert` on `HashMap`
-    //~| HELP Consider
+    //~| NOTE Consider using `m.entry(k)`
-    //~| SUGGESTION m.entry(k)
 }
 
 fn insert_if_absent2<K: Eq + Hash, V>(m: &mut HashMap<K, V>, k: K, v: V) {
     if !m.contains_key(&k) { m.insert(k, v) } else { None };
     //~^ ERROR usage of `contains_key` followed by `insert` on `HashMap`
-    //~| HELP Consider
+    //~| NOTE Consider using `m.entry(k)`
-    //~| SUGGESTION m.entry(k).or_insert(v)
+}
+
+fn insert_if_present2<K: Eq + Hash, V>(m: &mut HashMap<K, V>, k: K, v: V) {
+    if m.contains_key(&k) { None } else { m.insert(k, v) };
+    //~^ ERROR usage of `contains_key` followed by `insert` on `HashMap`
+    //~| NOTE Consider using `m.entry(k)`
 }
 
 fn insert_if_absent3<K: Eq + Hash, V>(m: &mut HashMap<K, V>, k: K, v: V) {
     if !m.contains_key(&k) { foo(); m.insert(k, v) } else { None };
     //~^ ERROR usage of `contains_key` followed by `insert` on `HashMap`
-    //~| HELP Consider
+    //~| NOTE Consider using `m.entry(k)`
-    //~| SUGGESTION m.entry(k)
+}
+
+fn insert_if_present3<K: Eq + Hash, V>(m: &mut HashMap<K, V>, k: K, v: V) {
+    if m.contains_key(&k) { None } else { foo(); m.insert(k, v) };
+    //~^ ERROR usage of `contains_key` followed by `insert` on `HashMap`
+    //~| NOTE Consider using `m.entry(k)`
 }
 
 fn insert_in_btreemap<K: Ord, V>(m: &mut BTreeMap<K, V>, k: K, v: V) {
     if !m.contains_key(&k) { foo(); m.insert(k, v) } else { None };
     //~^ ERROR usage of `contains_key` followed by `insert` on `BTreeMap`
-    //~| HELP Consider
+    //~| NOTE Consider using `m.entry(k)`
-    //~| SUGGESTION m.entry(k)
 }
 
 fn insert_other_if_absent<K: Eq + Hash, V>(m: &mut HashMap<K, V>, k: K, o: K, v: V) {