Quickselect in Rust

Cargo.toml

[package]
name = "quickselect"
version = "1.0.0"
authors = ["Graham Wren <[email protected]>"]
publish = false

[lib]
name="quickselect"
path="./quickselect.rs"

[dependencies]
rand = "0.6.5"

quickselect.rs

#![feature(test)]

extern crate rand;
extern crate test;

fn quickselect<'a, T: Ord>(v: &'a mut [T], n: usize) -> &'a T {
    assert!(n < v.len(), "Select value must be less than list length.");

    let right = v.len() - 1;
    return quickselect_helper(v, 0, right, n);
}

fn quickselect_helper<T: Ord>(v: &mut [T], left: usize, right: usize, n: usize) -> &T {
    assert!(left <= right, "Left index must be less than or equal to right index.");

    if left == right { return &v[left]; }

    let mut pivot = (rand::random::<usize>() % (right - left)) + left;
    assert!(left <= pivot && pivot <= right, "Random pivot was invalid");
    
    pivot = partition(v, left, right, pivot);
    assert!(left <= pivot && pivot <= right, "Selected pivot was invalid");

    if pivot == n { return &v[pivot]; }
    if pivot > n { return quickselect_helper(v, left, pivot - 1, n); }
    return quickselect_helper(v, pivot + 1, right, n);
}

fn partition<T: Ord>(v: &mut [T], left: usize, right: usize, pivot: usize) -> usize {
  v.swap(pivot, right); // pivot is now the right-most value

  let mut store_i = left;
  for n in left..right {
    // move each value less than the pivot value left of store_i
    if v[n] < v[right] {
      v.swap(n, store_i);
      store_i += 1;
    }
  }
  v.swap(right, store_i); // swap pivot with store_i
  return store_i;
}

#[cfg(test)]
mod tests {
    use super::quickselect;
    use super::test::Bencher;
    use super::rand::{thread_rng, Rng};

    #[test]
    fn select_0() {
        let mut v = [9, 8, 7, 6, 5, 4, 3, 2, 1, 0];
        assert_eq!(*quickselect(&mut v, 0), 0);
    }

    #[test]
    fn select_each() {
        let mut v = [2,3,5,1,3,6,2,8,9,3,4,5,2];
        let mut output = v;
        output.sort();
        for n in 0..output.len() {
            assert_eq!(*quickselect(&mut v, n), output[n]);
        }
    }

    #[test]
    #[should_panic(expected = "less than list length")]
    fn select_longer_than_list() {
        let mut v = [1, 2];
        quickselect(&mut v, 3);
    }

    #[bench]
    fn bench_quickselect(b: &mut Bencher) {
        let mut a: [u16; 1_000_000] = [0; 1_000_000];
        let mut rng = thread_rng();
        for n in 0..a.len() {
            let x: u16 = rng.gen();
            a[n] = x;
        }
        b.iter(|| {
            let i: usize = rng.gen();
            let mut test_a = a;
            quickselect(&mut test_a, i % a.len());
        });
    }

    #[bench]
    fn bench_sorted_select(b: &mut Bencher) {
        let mut a: [u16; 1_000_000] = [0; 1_000_000];
        let mut rng = thread_rng();
        for n in 0..a.len() {
            let x: u16 = rng.gen();
            a[n] = x;
        }
        b.iter(|| {
            let i: usize = rng.gen();
            let mut test_a = a;
            test_a.sort();
            let _res = test_a[i % a.len()];
        });
    }
}