k_smallest test now includes variation functions

Author: ejmount

tests/test_std.rs

@@ -418,22 +418,44 @@ fn sorted_by() {
 }
 
 qc::quickcheck! {
-    fn k_smallest_range(n: u64, m: u16, k: u16) -> () {
+    fn k_smallest_range(n: i64, m: u16, k: u16) -> () {
         // u16 is used to constrain k and m to 0..2¹⁶,
         //  otherwise the test could use too much memory.
-        let (k, m) = (k as u64, m as u64);
+        let (k, m) = (k as usize, m as u64);
 
+        let mut v: Vec<_> = (n..n.saturating_add(m as _)).collect();
         // Generate a random permutation of n..n+m
-        let i = {
-            let mut v: Vec<u64> = (n..n.saturating_add(m)).collect();
-            v.shuffle(&mut thread_rng());
-            v.into_iter()
-        };
+        v.shuffle(&mut thread_rng());
 
-        // Check that taking the k smallest elements yields n..n+min(k, m)
+        // Construct the right answers for the top and bottom elements
+        let mut sorted = v.clone();
+        sorted.sort();
+        // how many elements are we checking
+        let num_elements = min(k, m as _);
+
+        // Compute the top and bottom k in various combinations
+        let smallest_by_smallest : Vec<_> = v.iter().cloned().k_smallest(k).collect();
+        let smallest_by_largest : Vec<_> = v.iter().cloned().k_largest_by_key(k, |x| -x).collect();
+
+        let largest_by_smallest : Vec<_> = v.iter().cloned().k_smallest_by(k, |a,b| b.cmp(a)).collect();
+        let largest_by_largest : Vec<_> = v.iter().cloned().k_largest_by(k, Ord::cmp).collect();
+
+        // Check the variations produce the same answers and that they're right
+        assert_eq!(
+            smallest_by_smallest,
+            smallest_by_largest
+        );
+        assert_eq!(
+            largest_by_smallest,
+            largest_by_largest
+        );
+        assert_eq!(
+            sorted[..num_elements],
+            smallest_by_smallest[..], 
+        );
         it::assert_equal(
-            i.k_smallest(k as usize),
-            n..n.saturating_add(min(k, m))
+            sorted[sorted.len()-num_elements..].iter().rev(),
+            largest_by_smallest[..].iter()
         );
     }
 }