BTreeSet intersection, difference & is_subnet optimizations

1 files changed, 92 insertions, 18 deletions

diff --git a/src/liballoc/tests/btree/set.rs b/src/liballoc/tests/btree/set.rs
index 35db18c39c8..5c611fd21d2 100644
--- a/src/liballoc/tests/btree/set.rs
+++ b/src/liballoc/tests/btree/set.rs
@@ -48,7 +48,9 @@ fn check<F>(a: &[i32], b: &[i32], expected: &[i32], f: F)
     f(&set_a,
       &set_b,
       &mut |&x| {
-          assert_eq!(x, expected[i]);
+          if i < expected.len() {
+              assert_eq!(x, expected[i]);
+          }
           i += 1;
           true
       });
@@ -74,20 +76,20 @@ fn test_intersection() {
         return;
     }
 
-    let large = (0..1000).collect::<Vec<_>>();
+    let large = (0..100).collect::<Vec<_>>();
     check_intersection(&[], &large, &[]);
     check_intersection(&large, &[], &[]);
     check_intersection(&[-1], &large, &[]);
     check_intersection(&large, &[-1], &[]);
     check_intersection(&[0], &large, &[0]);
     check_intersection(&large, &[0], &[0]);
-    check_intersection(&[999], &large, &[999]);
-    check_intersection(&large, &[999], &[999]);
-    check_intersection(&[1000], &large, &[]);
-    check_intersection(&large, &[1000], &[]);
-    check_intersection(&[11, 5000, 1, 3, 77, 8924, 103],
+    check_intersection(&[99], &large, &[99]);
+    check_intersection(&large, &[99], &[99]);
+    check_intersection(&[100], &large, &[]);
+    check_intersection(&large, &[100], &[]);
+    check_intersection(&[11, 5000, 1, 3, 77, 8924],
                        &large,
-                       &[1, 3, 11, 77, 103]);
+                       &[1, 3, 11, 77]);
 }
 
 #[test]
@@ -95,10 +97,15 @@ fn test_intersection_size_hint() {
     let x: BTreeSet<i32> = [3, 4].iter().copied().collect();
     let y: BTreeSet<i32> = [1, 2, 3].iter().copied().collect();
     let mut iter = x.intersection(&y);
-    assert_eq!(iter.size_hint(), (0, Some(2)));
+    assert_eq!(iter.size_hint(), (1, Some(1)));
     assert_eq!(iter.next(), Some(&3));
     assert_eq!(iter.size_hint(), (0, Some(0)));
     assert_eq!(iter.next(), None);
+
+    iter = y.intersection(&y);
+    assert_eq!(iter.size_hint(), (0, Some(3)));
+    assert_eq!(iter.next(), Some(&1));
+    assert_eq!(iter.size_hint(), (0, Some(2)));
 }
 
 #[test]
@@ -111,6 +118,9 @@ fn test_difference() {
     check_difference(&[1, 12], &[], &[1, 12]);
     check_difference(&[], &[1, 2, 3, 9], &[]);
     check_difference(&[1, 3, 5, 9, 11], &[3, 9], &[1, 5, 11]);
+    check_difference(&[1, 3, 5, 9, 11], &[3, 6, 9], &[1, 5, 11]);
+    check_difference(&[1, 3, 5, 9, 11], &[0, 1], &[3, 5, 9, 11]);
+    check_difference(&[1, 3, 5, 9, 11], &[11, 12], &[1, 3, 5, 9]);
     check_difference(&[-5, 11, 22, 33, 40, 42],
                      &[-12, -5, 14, 23, 34, 38, 39, 50],
                      &[11, 22, 33, 40, 42]);
@@ -119,18 +129,82 @@ fn test_difference() {
         return;
     }
 
-    let large = (0..1000).collect::<Vec<_>>();
+    let large = (0..100).collect::<Vec<_>>();
     check_difference(&[], &large, &[]);
     check_difference(&[-1], &large, &[-1]);
     check_difference(&[0], &large, &[]);
-    check_difference(&[999], &large, &[]);
-    check_difference(&[1000], &large, &[1000]);
-    check_difference(&[11, 5000, 1, 3, 77, 8924, 103],
+    check_difference(&[99], &large, &[]);
+    check_difference(&[100], &large, &[100]);
+    check_difference(&[11, 5000, 1, 3, 77, 8924],
                      &large,
                      &[5000, 8924]);
     check_difference(&large, &[], &large);
     check_difference(&large, &[-1], &large);
-    check_difference(&large, &[1000], &large);
+    check_difference(&large, &[100], &large);
+}
+
+#[test]
+fn test_difference_size_hint() {
+    let s246: BTreeSet<i32> = [2, 4, 6].iter().copied().collect();
+    let s23456: BTreeSet<i32> = (2..=6).collect();
+    let mut iter = s246.difference(&s23456);
+    assert_eq!(iter.size_hint(), (0, Some(3)));
+    assert_eq!(iter.next(), None);
+
+    let s12345: BTreeSet<i32> = (1..=5).collect();
+    iter = s246.difference(&s12345);
+    assert_eq!(iter.size_hint(), (0, Some(3)));
+    assert_eq!(iter.next(), Some(&6));
+    assert_eq!(iter.size_hint(), (0, Some(0)));
+    assert_eq!(iter.next(), None);
+
+    let s34567: BTreeSet<i32> = (3..=7).collect();
+    iter = s246.difference(&s34567);
+    assert_eq!(iter.size_hint(), (0, Some(3)));
+    assert_eq!(iter.next(), Some(&2));
+    assert_eq!(iter.size_hint(), (0, Some(2)));
+    assert_eq!(iter.next(), None);
+
+    let s1: BTreeSet<i32> = (-9..=1).collect();
+    iter = s246.difference(&s1);
+    assert_eq!(iter.size_hint(), (3, Some(3)));
+
+    let s2: BTreeSet<i32> = (-9..=2).collect();
+    iter = s246.difference(&s2);
+    assert_eq!(iter.size_hint(), (2, Some(2)));
+    assert_eq!(iter.next(), Some(&4));
+    assert_eq!(iter.size_hint(), (1, Some(1)));
+
+    let s23: BTreeSet<i32> = (2..=3).collect();
+    iter = s246.difference(&s23);
+    assert_eq!(iter.size_hint(), (1, Some(3)));
+    assert_eq!(iter.next(), Some(&4));
+    assert_eq!(iter.size_hint(), (1, Some(1)));
+
+    let s4: BTreeSet<i32> = (4..=4).collect();
+    iter = s246.difference(&s4);
+    assert_eq!(iter.size_hint(), (2, Some(3)));
+    assert_eq!(iter.next(), Some(&2));
+    assert_eq!(iter.size_hint(), (1, Some(2)));
+    assert_eq!(iter.next(), Some(&6));
+    assert_eq!(iter.size_hint(), (0, Some(0)));
+    assert_eq!(iter.next(), None);
+
+    let s56: BTreeSet<i32> = (5..=6).collect();
+    iter = s246.difference(&s56);
+    assert_eq!(iter.size_hint(), (1, Some(3)));
+    assert_eq!(iter.next(), Some(&2));
+    assert_eq!(iter.size_hint(), (0, Some(2)));
+
+    let s6: BTreeSet<i32> = (6..=19).collect();
+    iter = s246.difference(&s6);
+    assert_eq!(iter.size_hint(), (2, Some(2)));
+    assert_eq!(iter.next(), Some(&2));
+    assert_eq!(iter.size_hint(), (1, Some(1)));
+
+    let s7: BTreeSet<i32> = (7..=19).collect();
+    iter = s246.difference(&s7);
+    assert_eq!(iter.size_hint(), (3, Some(3)));
 }
 
 #[test]
@@ -188,23 +262,23 @@ fn test_is_subset() {
     assert_eq!(is_subset(&[1, 2], &[1, 2]), true);
     assert_eq!(is_subset(&[1, 2], &[2, 3]), false);
     assert_eq!(is_subset(&[-5, 11, 22, 33, 40, 42],
-                         &[-12, -5, 14, 23, 11, 34, 22, 38, 33, 42, 39, 40]),
+                         &[-12, -5, 11, 14, 22, 23, 33, 34, 38, 39, 40, 42]),
                true);
     assert_eq!(is_subset(&[-5, 11, 22, 33, 40, 42],
-                         &[-12, -5, 14, 23, 34, 38, 22, 11]),
+                         &[-12, -5, 11, 14, 22, 23, 34, 38]),
                false);
 
     if cfg!(miri) { // Miri is too slow
         return;
     }
 
-    let large = (0..1000).collect::<Vec<_>>();
+    let large = (0..100).collect::<Vec<_>>();
     assert_eq!(is_subset(&[], &large), true);
     assert_eq!(is_subset(&large, &[]), false);
     assert_eq!(is_subset(&[-1], &large), false);
     assert_eq!(is_subset(&[0], &large), true);
     assert_eq!(is_subset(&[1, 2], &large), true);
-    assert_eq!(is_subset(&[999, 1000], &large), false);
+    assert_eq!(is_subset(&[99, 100], &large), false);
 }
 
 #[test]