Add `BTreeSet` to set_contains_or_insert

* Detect `BTreeSet::contains` + `BTreeSet::insert` usage in the same way as with the `HashSet`.

4 files changed, 160 insertions, 30 deletions

diff --git a/clippy_lints/src/lib.rs b/clippy_lints/src/lib.rs
index ecfb6cdbdf0..5549c8df101 100644
--- a/clippy_lints/src/lib.rs
+++ b/clippy_lints/src/lib.rs
@@ -908,7 +908,7 @@ pub fn register_lints(store: &mut rustc_lint::LintStore, conf: &'static Conf) {
     store.register_late_pass(move |_| Box::new(macro_metavars_in_unsafe::ExprMetavarsInUnsafe::new(conf)));
     store.register_late_pass(move |_| Box::new(string_patterns::StringPatterns::new(conf)));
     store.register_early_pass(|| Box::new(field_scoped_visibility_modifiers::FieldScopedVisibilityModifiers));
-    store.register_late_pass(|_| Box::new(set_contains_or_insert::HashsetInsertAfterContains));
+    store.register_late_pass(|_| Box::new(set_contains_or_insert::SetContainsOrInsert));
     store.register_early_pass(|| Box::new(byte_char_slices::ByteCharSlice));
     store.register_early_pass(|| Box::new(cfg_not_test::CfgNotTest));
     // add lints here, do not remove this comment, it's used in `new_lint`
diff --git a/clippy_lints/src/set_contains_or_insert.rs b/clippy_lints/src/set_contains_or_insert.rs
index 5e65b9fa517..e6fe7649397 100644
--- a/clippy_lints/src/set_contains_or_insert.rs
+++ b/clippy_lints/src/set_contains_or_insert.rs
@@ -12,8 +12,8 @@ use rustc_span::{sym, Span};
 
 declare_clippy_lint! {
     /// ### What it does
-    /// Checks for usage of `contains` to see if a value is not
-    /// present on `HashSet` followed by a `insert`.
+    /// Checks for usage of `contains` to see if a value is not present
+    /// in a set like `HashSet` or `BTreeSet`, followed by an `insert`.
     ///
     /// ### Why is this bad?
     /// Using just `insert` and checking the returned `bool` is more efficient.
@@ -45,12 +45,12 @@ declare_clippy_lint! {
     #[clippy::version = "1.80.0"]
     pub SET_CONTAINS_OR_INSERT,
     nursery,
-    "call to `HashSet::contains` followed by `HashSet::insert`"
+    "call to `<set>::contains` followed by `<set>::insert`"
 }
 
-declare_lint_pass!(HashsetInsertAfterContains => [SET_CONTAINS_OR_INSERT]);
+declare_lint_pass!(SetContainsOrInsert => [SET_CONTAINS_OR_INSERT]);
 
-impl<'tcx> LateLintPass<'tcx> for HashsetInsertAfterContains {
+impl<'tcx> LateLintPass<'tcx> for SetContainsOrInsert {
     fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
         if !expr.span.from_expansion()
             && let Some(higher::If {
@@ -58,14 +58,14 @@ impl<'tcx> LateLintPass<'tcx> for HashsetInsertAfterContains {
                 then: then_expr,
                 ..
             }) = higher::If::hir(expr)
-            && let Some(contains_expr) = try_parse_op_call(cx, cond_expr, sym!(contains))//try_parse_contains(cx, cond_expr)
+            && let Some((contains_expr, sym)) = try_parse_op_call(cx, cond_expr, sym!(contains))//try_parse_contains(cx, cond_expr)
             && let Some(insert_expr) = find_insert_calls(cx, &contains_expr, then_expr)
         {
             span_lint(
                 cx,
                 SET_CONTAINS_OR_INSERT,
                 vec![contains_expr.span, insert_expr.span],
-                "usage of `HashSet::insert` after `HashSet::contains`",
+                format!("usage of `{sym}::insert` after `{sym}::contains`"),
             );
         }
     }
@@ -77,7 +77,11 @@ struct OpExpr<'tcx> {
     span: Span,
 }
 
-fn try_parse_op_call<'tcx>(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>, symbol: Symbol) -> Option<OpExpr<'tcx>> {
+fn try_parse_op_call<'tcx>(
+    cx: &LateContext<'tcx>,
+    expr: &'tcx Expr<'_>,
+    symbol: Symbol,
+) -> Option<(OpExpr<'tcx>, Symbol)> {
     let expr = peel_hir_expr_while(expr, |e| {
         if let ExprKind::Unary(UnOp::Not, e) = e.kind {
             Some(e)
@@ -97,11 +101,12 @@ fn try_parse_op_call<'tcx>(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>, symbol:
         });
         let receiver = receiver.peel_borrows();
         let receiver_ty = cx.typeck_results().expr_ty(receiver).peel_refs();
-        if value.span.eq_ctxt(expr.span)
-            && is_type_diagnostic_item(cx, receiver_ty, sym::HashSet)
-            && path.ident.name == symbol
-        {
-            return Some(OpExpr { receiver, value, span });
+        if value.span.eq_ctxt(expr.span) && path.ident.name == symbol {
+            for sym in &[sym::HashSet, sym::BTreeSet] {
+                if is_type_diagnostic_item(cx, receiver_ty, *sym) {
+                    return Some((OpExpr { receiver, value, span }, *sym));
+                }
+            }
         }
     }
     None
@@ -113,7 +118,7 @@ fn find_insert_calls<'tcx>(
     expr: &'tcx Expr<'_>,
 ) -> Option<OpExpr<'tcx>> {
     for_each_expr(cx, expr, |e| {
-        if let Some(insert_expr) = try_parse_op_call(cx, e, sym!(insert))
+        if let Some((insert_expr, _)) = try_parse_op_call(cx, e, sym!(insert))
             && SpanlessEq::new(cx).eq_expr(contains_expr.receiver, insert_expr.receiver)
             && SpanlessEq::new(cx).eq_expr(contains_expr.value, insert_expr.value)
         {
diff --git a/tests/ui/set_contains_or_insert.rs b/tests/ui/set_contains_or_insert.rs
index 8465007402a..d3a3e1c878b 100644
--- a/tests/ui/set_contains_or_insert.rs
+++ b/tests/ui/set_contains_or_insert.rs
@@ -3,17 +3,78 @@
 #![allow(clippy::needless_borrow)]
 #![warn(clippy::set_contains_or_insert)]
 
-use std::collections::HashSet;
+use std::collections::{BTreeSet, HashSet};
 
-fn main() {
-    should_warn_cases();
+fn should_warn_hashset() {
+    let mut set = HashSet::new();
+    let value = 5;
+
+    if !set.contains(&value) {
+        set.insert(value);
+        println!("Just a comment");
+    }
+
+    if set.contains(&value) {
+        set.insert(value);
+        println!("Just a comment");
+    }
+
+    if !set.contains(&value) {
+        set.insert(value);
+    }
 
-    should_not_warn_cases();
+    if !!set.contains(&value) {
+        set.insert(value);
+        println!("Just a comment");
+    }
+
+    if (&set).contains(&value) {
+        set.insert(value);
+    }
+
+    let borrow_value = &6;
+    if !set.contains(borrow_value) {
+        set.insert(*borrow_value);
+    }
+
+    let borrow_set = &mut set;
+    if !borrow_set.contains(&value) {
+        borrow_set.insert(value);
+    }
 }
 
-fn should_warn_cases() {
+fn should_not_warn_hashset() {
     let mut set = HashSet::new();
     let value = 5;
+    let another_value = 6;
+
+    if !set.contains(&value) {
+        set.insert(another_value);
+    }
+
+    if !set.contains(&value) {
+        println!("Just a comment");
+    }
+
+    if simply_true() {
+        set.insert(value);
+    }
+
+    if !set.contains(&value) {
+        set.replace(value); //it is not insert
+        println!("Just a comment");
+    }
+
+    if set.contains(&value) {
+        println!("value is already in set");
+    } else {
+        set.insert(value);
+    }
+}
+
+fn should_warn_btreeset() {
+    let mut set = BTreeSet::new();
+    let value = 5;
 
     if !set.contains(&value) {
         set.insert(value);
@@ -49,8 +110,8 @@ fn should_warn_cases() {
     }
 }
 
-fn should_not_warn_cases() {
-    let mut set = HashSet::new();
+fn should_not_warn_btreeset() {
+    let mut set = BTreeSet::new();
     let value = 5;
     let another_value = 6;
 
@@ -81,3 +142,11 @@ fn should_not_warn_cases() {
 fn simply_true() -> bool {
     true
 }
+
+// This is placed last in order to be able to add new tests without changing line numbers
+fn main() {
+    should_warn_hashset();
+    should_warn_btreeset();
+    should_not_warn_hashset();
+    should_not_warn_btreeset();
+}
diff --git a/tests/ui/set_contains_or_insert.stderr b/tests/ui/set_contains_or_insert.stderr
index 507e20964fc..14ad6300544 100644
--- a/tests/ui/set_contains_or_insert.stderr
+++ b/tests/ui/set_contains_or_insert.stderr
@@ -1,5 +1,5 @@
 error: usage of `HashSet::insert` after `HashSet::contains`
-  --> tests/ui/set_contains_or_insert.rs:18:13
+  --> tests/ui/set_contains_or_insert.rs:12:13
    |
 LL |     if !set.contains(&value) {
    |             ^^^^^^^^^^^^^^^^
@@ -10,7 +10,7 @@ LL |         set.insert(value);
    = help: to override `-D warnings` add `#[allow(clippy::set_contains_or_insert)]`
 
 error: usage of `HashSet::insert` after `HashSet::contains`
-  --> tests/ui/set_contains_or_insert.rs:23:12
+  --> tests/ui/set_contains_or_insert.rs:17:12
    |
 LL |     if set.contains(&value) {
    |            ^^^^^^^^^^^^^^^^
@@ -18,7 +18,7 @@ LL |         set.insert(value);
    |             ^^^^^^^^^^^^^
 
 error: usage of `HashSet::insert` after `HashSet::contains`
-  --> tests/ui/set_contains_or_insert.rs:28:13
+  --> tests/ui/set_contains_or_insert.rs:22:13
    |
 LL |     if !set.contains(&value) {
    |             ^^^^^^^^^^^^^^^^
@@ -26,7 +26,7 @@ LL |         set.insert(value);
    |             ^^^^^^^^^^^^^
 
 error: usage of `HashSet::insert` after `HashSet::contains`
-  --> tests/ui/set_contains_or_insert.rs:32:14
+  --> tests/ui/set_contains_or_insert.rs:26:14
    |
 LL |     if !!set.contains(&value) {
    |              ^^^^^^^^^^^^^^^^
@@ -34,7 +34,7 @@ LL |         set.insert(value);
    |             ^^^^^^^^^^^^^
 
 error: usage of `HashSet::insert` after `HashSet::contains`
-  --> tests/ui/set_contains_or_insert.rs:37:15
+  --> tests/ui/set_contains_or_insert.rs:31:15
    |
 LL |     if (&set).contains(&value) {
    |               ^^^^^^^^^^^^^^^^
@@ -42,7 +42,7 @@ LL |         set.insert(value);
    |             ^^^^^^^^^^^^^
 
 error: usage of `HashSet::insert` after `HashSet::contains`
-  --> tests/ui/set_contains_or_insert.rs:42:13
+  --> tests/ui/set_contains_or_insert.rs:36:13
    |
 LL |     if !set.contains(borrow_value) {
    |             ^^^^^^^^^^^^^^^^^^^^^^
@@ -50,12 +50,68 @@ LL |         set.insert(*borrow_value);
    |             ^^^^^^^^^^^^^^^^^^^^^
 
 error: usage of `HashSet::insert` after `HashSet::contains`
-  --> tests/ui/set_contains_or_insert.rs:47:20
+  --> tests/ui/set_contains_or_insert.rs:41:20
    |
 LL |     if !borrow_set.contains(&value) {
    |                    ^^^^^^^^^^^^^^^^
 LL |         borrow_set.insert(value);
    |                    ^^^^^^^^^^^^^
 
-error: aborting due to 7 previous errors
+error: usage of `BTreeSet::insert` after `BTreeSet::contains`
+  --> tests/ui/set_contains_or_insert.rs:79:13
+   |
+LL |     if !set.contains(&value) {
+   |             ^^^^^^^^^^^^^^^^
+LL |         set.insert(value);
+   |             ^^^^^^^^^^^^^
+
+error: usage of `BTreeSet::insert` after `BTreeSet::contains`
+  --> tests/ui/set_contains_or_insert.rs:84:12
+   |
+LL |     if set.contains(&value) {
+   |            ^^^^^^^^^^^^^^^^
+LL |         set.insert(value);
+   |             ^^^^^^^^^^^^^
+
+error: usage of `BTreeSet::insert` after `BTreeSet::contains`
+  --> tests/ui/set_contains_or_insert.rs:89:13
+   |
+LL |     if !set.contains(&value) {
+   |             ^^^^^^^^^^^^^^^^
+LL |         set.insert(value);
+   |             ^^^^^^^^^^^^^
+
+error: usage of `BTreeSet::insert` after `BTreeSet::contains`
+  --> tests/ui/set_contains_or_insert.rs:93:14
+   |
+LL |     if !!set.contains(&value) {
+   |              ^^^^^^^^^^^^^^^^
+LL |         set.insert(value);
+   |             ^^^^^^^^^^^^^
+
+error: usage of `BTreeSet::insert` after `BTreeSet::contains`
+  --> tests/ui/set_contains_or_insert.rs:98:15
+   |
+LL |     if (&set).contains(&value) {
+   |               ^^^^^^^^^^^^^^^^
+LL |         set.insert(value);
+   |             ^^^^^^^^^^^^^
+
+error: usage of `BTreeSet::insert` after `BTreeSet::contains`
+  --> tests/ui/set_contains_or_insert.rs:103:13
+   |
+LL |     if !set.contains(borrow_value) {
+   |             ^^^^^^^^^^^^^^^^^^^^^^
+LL |         set.insert(*borrow_value);
+   |             ^^^^^^^^^^^^^^^^^^^^^
+
+error: usage of `BTreeSet::insert` after `BTreeSet::contains`
+  --> tests/ui/set_contains_or_insert.rs:108:20
+   |
+LL |     if !borrow_set.contains(&value) {
+   |                    ^^^^^^^^^^^^^^^^
+LL |         borrow_set.insert(value);
+   |                    ^^^^^^^^^^^^^
+
+error: aborting due to 14 previous errors