Bug in `k_smallest_relaxed`?

[Takashiidobe]

If I write this code using k_smallest, and ask for usize::MAX items, this returns all the items in the iterator, which makes sense.

use itertools::Itertools;

#[test]
fn k_smallest_memory_is_fine() {
    let iter = vec![1].into_iter();
    let first = iter.k_smallest(usize::MAX);

    itertools::assert_equal(first, vec![1]);
}

I think this is ok because this line only requires memory for an iterator up to size k

However, when swapping to k_smallest_relaxed given usize::MAX / 2 + 1 items, it panics:

use itertools::Itertools;

#[test]
fn k_smallest_memory_does_not_panic() {
    let iter = vec![1].into_iter();
    let first = iter.k_smallest_relaxed(usize::MAX / 2);

    itertools::assert_equal(first, vec![1]);
}

#[test]
#[should_panic]
fn k_smallest_memory_does_panic() {
    let iter = vec![1].into_iter();
    let first = iter.k_smallest_relaxed(usize::MAX / 2 + 1);

    itertools::assert_equal(first, vec![1]);
}

It seems like the problematic line is here where k is multiplied by 2, unchecked, so usize::MAX / 2 + 1 overflows it and causes the panic.

At the very least, this simple fix causes the test above to pass without panicking, since it does a saturating multiply.

diff --git a/src/k_smallest.rs b/src/k_smallest.rs
index 7e4ace2..c28831b 100644
--- a/src/k_smallest.rs
+++ b/src/k_smallest.rs
@@ -99,7 +99,10 @@ where
     }
 
     let mut iter = iter.fuse();
-    let mut buf = iter.by_ref().take(2 * k).collect::<Vec<_>>();
+    let mut buf = iter
+        .by_ref()
+        .take(2usize.saturating_mul(k))
+        .collect::<Vec<_>>();
 
     if buf.len() < k {
         buf.sort_unstable_by(&mut comparator);

Given that the docs for k_smallest_relaxed explicitly call out that you need 2 * k * sizeof(Self::Item) + O(1) memory, I'm wondering if this change is fine though -- since it explicitly uses more memory than k_smallest which is k * sizeof(Self::Item) + O(1).