Add BTreeMap to the HASHMAP_ENTRY rule

Fixes #433

3 files changed, 49 insertions, 39 deletions

diff --git a/src/hashmap.rs b/src/entry.rs
index 095a00e3777..1885bea164b 100644
--- a/src/hashmap.rs
+++ b/src/entry.rs
@@ -2,11 +2,12 @@ use rustc::lint::*;
 use rustc_front::hir::*;
 use syntax::codemap::Span;
 use utils::{get_item_name, is_exp_equal, match_type, snippet, span_help_and_lint, walk_ptrs_ty};
-use utils::HASHMAP_PATH;
+use utils::{BTREEMAP_PATH, HASHMAP_PATH};
 
-/// **What it does:** This lint checks for uses of `contains_key` + `insert` on `HashMap`.
+/// **What it does:** This lint checks for uses of `contains_key` + `insert` on `HashMap` or
+/// `BTreeMap`.
 ///
-/// **Why is this bad?** Using `HashMap::entry` is more efficient.
+/// **Why is this bad?** Using `entry` is more efficient.
 ///
 /// **Known problems:** Some false negatives, eg.:
 /// ```
@@ -23,9 +24,9 @@ use utils::HASHMAP_PATH;
 /// m.entry(k).or_insert(v);
 /// ```
 declare_lint! {
-    pub HASHMAP_ENTRY,
+    pub MAP_ENTRY,
     Warn,
-    "use of `contains_key` followed by `insert` on a `HashMap`"
+    "use of `contains_key` followed by `insert` on a `HashMap` or `BTreeMap`"
 }
 
 #[derive(Copy,Clone)]
@@ -33,7 +34,7 @@ pub struct HashMapLint;
 
 impl LintPass for HashMapLint {
     fn get_lints(&self) -> LintArray {
-        lint_array!(HASHMAP_ENTRY)
+        lint_array!(MAP_ENTRY)
     }
 }
 
@@ -55,17 +56,25 @@ impl LateLintPass for HashMapLint {
                 let map = &params[0];
                 let obj_ty = walk_ptrs_ty(cx.tcx.expr_ty(map));
 
-                if match_type(cx, obj_ty, &HASHMAP_PATH) {
-                    let sole_expr = if then.expr.is_some() { 1 } else { 0 } + then.stmts.len() == 1;
+                let kind = if match_type(cx, obj_ty, &BTREEMAP_PATH) {
+                    "BTreeMap"
+                }
+                else if match_type(cx, obj_ty, &HASHMAP_PATH) {
+                    "HashMap"
+                }
+                else {
+                    return
+                };
 
-                    if let Some(ref then) = then.expr {
-                        check_for_insert(cx, expr.span, map, key, then, sole_expr);
-                    }
+                let sole_expr = if then.expr.is_some() { 1 } else { 0 } + then.stmts.len() == 1;
 
-                    for stmt in &then.stmts {
-                        if let StmtSemi(ref stmt, _) = stmt.node {
-                            check_for_insert(cx, expr.span, map, key, stmt, sole_expr);
-                        }
+                if let Some(ref then) = then.expr {
+                    check_for_insert(cx, expr.span, map, key, then, sole_expr, kind);
+                }
+
+                for stmt in &then.stmts {
+                    if let StmtSemi(ref stmt, _) = stmt.node {
+                        check_for_insert(cx, expr.span, map, key, stmt, sole_expr, kind);
                     }
                 }
             }
@@ -73,7 +82,7 @@ impl LateLintPass for HashMapLint {
     }
 }
 
-fn check_for_insert(cx: &LateContext, span: Span, map: &Expr, key: &Expr, expr: &Expr, sole_expr: bool) {
+fn check_for_insert(cx: &LateContext, span: Span, map: &Expr, key: &Expr, expr: &Expr, sole_expr: bool, kind: &str) {
     if_let_chain! {
         [
             let ExprMethodCall(ref name, _, ref params) = expr.node,
@@ -82,21 +91,21 @@ fn check_for_insert(cx: &LateContext, span: Span, map: &Expr, key: &Expr, expr:
             get_item_name(cx, map) == get_item_name(cx, &*params[0]),
             is_exp_equal(cx, key, &params[1])
         ], {
-            if sole_expr {
-                span_help_and_lint(cx, HASHMAP_ENTRY, span,
-                                   "usage of `contains_key` followed by `insert` on `HashMap`",
-                                   &format!("Consider using `{}.entry({}).or_insert({})`",
-                                            snippet(cx, map.span, ".."),
-                                            snippet(cx, params[1].span, ".."),
-                                            snippet(cx, params[2].span, ".."))); 
+            let help = if sole_expr {
+                format!("Consider using `{}.entry({}).or_insert({})`",
+                        snippet(cx, map.span, ".."),
+                        snippet(cx, params[1].span, ".."),
+                        snippet(cx, params[2].span, ".."))
             }
             else {
-                span_help_and_lint(cx, HASHMAP_ENTRY, span,
-                                   "usage of `contains_key` followed by `insert` on `HashMap`",
-                                   &format!("Consider using `{}.entry({})`",
-                                            snippet(cx, map.span, ".."),
-                                            snippet(cx, params[1].span, "..")));
-            }
+                format!("Consider using `{}.entry({})`",
+                        snippet(cx, map.span, ".."),
+                        snippet(cx, params[1].span, ".."))
+            };
+
+            span_help_and_lint(cx, MAP_ENTRY, span,
+                               &format!("usage of `contains_key` followed by `insert` on `{}`", kind),
+                               &help);
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index f8db15605ff..9bb59693795 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -71,7 +71,7 @@ pub mod temporary_assignment;
 pub mod transmute;
 pub mod cyclomatic_complexity;
 pub mod escape;
-pub mod hashmap;
+pub mod entry;
 pub mod misc_early;
 pub mod array_indexing;
 pub mod panic;
@@ -113,7 +113,7 @@ pub fn plugin_registrar(reg: &mut Registry) {
     reg.register_late_lint_pass(box types::UnitCmp);
     reg.register_late_lint_pass(box loops::LoopsPass);
     reg.register_late_lint_pass(box lifetimes::LifetimePass);
-    reg.register_late_lint_pass(box hashmap::HashMapLint);
+    reg.register_late_lint_pass(box entry::HashMapLint);
     reg.register_late_lint_pass(box ranges::StepByZero);
     reg.register_late_lint_pass(box types::CastPass);
     reg.register_late_lint_pass(box types::TypeComplexityPass);
@@ -167,10 +167,10 @@ pub fn plugin_registrar(reg: &mut Registry) {
         block_in_if_condition::BLOCK_IN_IF_CONDITION_STMT,
         collapsible_if::COLLAPSIBLE_IF,
         cyclomatic_complexity::CYCLOMATIC_COMPLEXITY,
+        entry::MAP_ENTRY,
         eq_op::EQ_OP,
         escape::BOXED_LOCAL,
         eta_reduction::REDUNDANT_CLOSURE,
-        hashmap::HASHMAP_ENTRY,
         identity_op::IDENTITY_OP,
         len_zero::LEN_WITHOUT_IS_EMPTY,
         len_zero::LEN_ZERO,
diff --git a/src/utils.rs b/src/utils.rs
index 77e63fe4458..d3ebe809ac9 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -19,17 +19,18 @@ use std::ops::{Deref, DerefMut};
 pub type MethodArgs = HirVec<P<Expr>>;
 
 // module DefPaths for certain structs/enums we check for
+pub const BEGIN_UNWIND: [&'static str; 3] = ["std", "rt", "begin_unwind"];
+pub const BTREEMAP_PATH: [&'static str; 4] = ["collections", "btree", "map", "BTreeMap"];
+pub const CLONE_PATH: [&'static str; 2] = ["Clone", "clone"];
+pub const COW_PATH: [&'static str; 3] = ["collections", "borrow", "Cow"];
+pub const HASHMAP_PATH: [&'static str; 5] = ["std", "collections", "hash", "map", "HashMap"];
+pub const LL_PATH: [&'static str; 3] = ["collections", "linked_list", "LinkedList"];
+pub const MUTEX_PATH: [&'static str; 4] = ["std", "sync", "mutex", "Mutex"];
+pub const OPEN_OPTIONS_PATH: [&'static str; 3] = ["std", "fs", "OpenOptions"];
 pub const OPTION_PATH: [&'static str; 3] = ["core", "option", "Option"];
 pub const RESULT_PATH: [&'static str; 3] = ["core", "result", "Result"];
 pub const STRING_PATH: [&'static str; 3] = ["collections", "string", "String"];
 pub const VEC_PATH: [&'static str; 3] = ["collections", "vec", "Vec"];
-pub const LL_PATH: [&'static str; 3] = ["collections", "linked_list", "LinkedList"];
-pub const HASHMAP_PATH: [&'static str; 5] = ["std", "collections", "hash", "map", "HashMap"];
-pub const OPEN_OPTIONS_PATH: [&'static str; 3] = ["std", "fs", "OpenOptions"];
-pub const MUTEX_PATH: [&'static str; 4] = ["std", "sync", "mutex", "Mutex"];
-pub const CLONE_PATH: [&'static str; 2] = ["Clone", "clone"];
-pub const BEGIN_UNWIND: [&'static str; 3] = ["std", "rt", "begin_unwind"];
-pub const COW_PATH: [&'static str; 3] = ["collections", "borrow", "Cow"];
 
 /// Produce a nested chain of if-lets and ifs from the patterns:
 ///