<commit_before>extern crate x264_sys as ffi;

use std::mem;
use ffi::x264::*;
use std::ptr::null;
use std::os::raw::c_int;

pub struct Picture {
    pic: x264_picture_t,
    height: i32, // to compute the slice dimension
    native: bool,
}

impl Picture {
    /*
    pub fn new() -> Picture {
        let mut pic = unsafe { mem::uninitialized() };

        unsafe { x264_picture_init(&mut pic as *mut x264_picture_t) };

        Picture { pic: pic }
    }
*/
    pub fn from_param(param: &Param) -> Result<Picture, &'static str> {
        let mut pic = unsafe { mem::uninitialized() };

        let ret = unsafe {
            x264_picture_alloc(&mut pic as *mut x264_picture_t,
                               param.par.i_csp,
                               param.par.i_width,
                               param.par.i_height)
        };
        if ret < 0 {
            Err("Allocation Failure")
        } else {
            Ok(Picture {
                   pic: pic,
                   height: param.par.i_height,
                   native: true,
               })
        }
    }

    pub fn as_slice<'a>(&'a self, plane: usize) -> Result<&'a [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn as_slice_mut<'a>(&'a mut self, plane: usize) -> Result<&'a mut [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts_mut(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn set_timestamp(mut self, pts: i64) -> Picture {
        self.pic.i_pts = pts;
        self
    }
}

impl Drop for Picture {
    fn drop(&mut self) {
        if self.native {
            unsafe { x264_picture_clean(&mut self.pic as *mut x264_picture_t) };
        }
    }
}

// TODO: Provide a builder API instead?
pub struct Param {
    par: x264_param_t,
}

impl Param {
    pub fn new() -> Param {
        let mut par = unsafe { mem::uninitialized() };

        unsafe {
            x264_param_default(&mut par as *mut x264_param_t);
        }

        Param { par: par }
    }
    pub fn default_preset(tune: Option<&str>, preset: Option<&str>) -> Result<Param, &'static str> {
        let mut par = unsafe { mem::uninitialized() };

        match unsafe {
                  x264_param_default_preset(&mut par as *mut x264_param_t,
                                            tune.map_or_else(|| null(),
                                                             |v| v.as_ptr() as *const i8),
                                            preset.map_or_else(|| null(), |v| v.as_ptr()) as
                                            *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(Param { par: par }),
            _ => Err("Unexpected"),
        }
    }
    pub fn apply_profile(mut self, profile: &str) -> Result<Param, &'static str> {
        match unsafe { x264_param_apply_profile(&mut self.par, profile.as_ptr() as *const i8) } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn param_parse<'a>(mut self, name: &str, value: &str) -> Result<Param, &'static str> {
        match unsafe {
                  x264_param_parse(&mut self.par,
                                   name.as_ptr() as *const i8,
                                   value.as_ptr() as *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn set_dimension(mut self, height: usize, width: usize) -> Param {
        self.par.i_height = height as c_int;
        self.par.i_width = width as c_int;

        self
    }
}

// TODO: Expose a NAL abstraction
pub struct NalData {
    vec: Vec<u8>,
}

impl NalData {
    /*
     * x264 functions return x264_nal_t arrays that are valid only until another
     * function of that kind is called.
     *
     * Always copy the data over.
     *
     * TODO: Consider using Bytes as backing store.
     */
    fn from_nals(c_nals: *mut x264_nal_t, nb_nal: usize) -> NalData {
        let mut data = NalData { vec: Vec::new() };

        for i in 0..nb_nal {
            let nal = unsafe { Box::from_raw(c_nals.offset(i as isize)) };

            let payload =
                unsafe { std::slice::from_raw_parts(nal.p_payload, nal.i_payload as usize) };

            data.vec.extend_from_slice(payload);

            mem::forget(payload);
            mem::forget(nal);
        }

        data
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.vec.as_slice()
    }
}

pub struct Encoder {
    enc: *mut x264_t,
}

impl Encoder {
    pub fn open(par: &mut Param) -> Result<Encoder, &'static str> {
        let enc = unsafe { x264_encoder_open(&mut par.par as *mut x264_param_t) };

        if enc.is_null() {
            Err("Out of Memory")
        } else {
            Ok(Encoder { enc: enc })
        }
    }

    pub fn get_headers(&mut self) -> Result<NalData, &'static str> {
        let mut nb_nal: c_int = 0;
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };

        let bytes = unsafe {
            x264_encoder_headers(self.enc,
                                 &mut c_nals as *mut *mut x264_nal_t,
                                 &mut nb_nal as *mut c_int)
        };

        if bytes < 0 {
            Err("Encoding Headers Failed")
        } else {
            Ok(NalData::from_nals(c_nals, nb_nal as usize))
        }
    }

                data.vec.extend_from_slice(payload);

                mem::forget(payload);
                mem::forget(nal);
            }

            Ok(data)
        }
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe { x264_encoder_close(self.enc) };
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_open() {
        let mut par = Param::new().set_dimension(640, 480);

        let mut enc = Encoder::open(&mut par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());
    }

    #[test]
    fn test_picture() {
        let par = Param::new().set_dimension(640, 480);
        {
            let mut pic = Picture::from_param(&par).unwrap();
            {
                let p = pic.as_slice_mut(0).unwrap();
                p[0] = 1;
            }
            let p = pic.as_slice(0).unwrap();

            assert_eq!(p[0], 1);
        }
    }
}
<commit_msg>encode: Provide a safe wrapper for encoding<commit_after>extern crate x264_sys as ffi;

use std::mem;
use ffi::x264::*;
use std::ptr::null;
use std::os::raw::c_int;

pub struct Picture {
    pic: x264_picture_t,
    height: i32, // to compute the slice dimension
    native: bool,
}

impl Picture {
    /*
    pub fn new() -> Picture {
        let mut pic = unsafe { mem::uninitialized() };

        unsafe { x264_picture_init(&mut pic as *mut x264_picture_t) };

        Picture { pic: pic }
    }
*/
    pub fn from_param(param: &Param) -> Result<Picture, &'static str> {
        let mut pic = unsafe { mem::uninitialized() };

        let ret = unsafe {
            x264_picture_alloc(&mut pic as *mut x264_picture_t,
                               param.par.i_csp,
                               param.par.i_width,
                               param.par.i_height)
        };
        if ret < 0 {
            Err("Allocation Failure")
        } else {
            Ok(Picture {
                   pic: pic,
                   height: param.par.i_height,
                   native: true,
               })
        }
    }

    pub fn as_slice<'a>(&'a self, plane: usize) -> Result<&'a [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn as_slice_mut<'a>(&'a mut self, plane: usize) -> Result<&'a mut [u8], &'static str> {
        if plane > self.pic.img.i_plane as usize {
            Err("Invalid Argument")
        } else {
            let size = self.height * self.pic.img.i_stride[plane];
            Ok(unsafe { std::slice::from_raw_parts_mut(self.pic.img.plane[plane], size as usize) })
        }
    }

    pub fn set_timestamp(mut self, pts: i64) -> Picture {
        self.pic.i_pts = pts;
        self
    }
}

impl Drop for Picture {
    fn drop(&mut self) {
        if self.native {
            unsafe { x264_picture_clean(&mut self.pic as *mut x264_picture_t) };
        }
    }
}

// TODO: Provide a builder API instead?
pub struct Param {
    par: x264_param_t,
}

impl Param {
    pub fn new() -> Param {
        let mut par = unsafe { mem::uninitialized() };

        unsafe {
            x264_param_default(&mut par as *mut x264_param_t);
        }

        Param { par: par }
    }
    pub fn default_preset(tune: Option<&str>, preset: Option<&str>) -> Result<Param, &'static str> {
        let mut par = unsafe { mem::uninitialized() };

        match unsafe {
                  x264_param_default_preset(&mut par as *mut x264_param_t,
                                            tune.map_or_else(|| null(),
                                                             |v| v.as_ptr() as *const i8),
                                            preset.map_or_else(|| null(), |v| v.as_ptr()) as
                                            *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(Param { par: par }),
            _ => Err("Unexpected"),
        }
    }
    pub fn apply_profile(mut self, profile: &str) -> Result<Param, &'static str> {
        match unsafe { x264_param_apply_profile(&mut self.par, profile.as_ptr() as *const i8) } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn param_parse<'a>(mut self, name: &str, value: &str) -> Result<Param, &'static str> {
        match unsafe {
                  x264_param_parse(&mut self.par,
                                   name.as_ptr() as *const i8,
                                   value.as_ptr() as *const i8)
              } {
            -1 => Err("Invalid Argument"),
            0 => Ok(self),
            _ => Err("Unexpected"),
        }
    }
    pub fn set_dimension(mut self, height: usize, width: usize) -> Param {
        self.par.i_height = height as c_int;
        self.par.i_width = width as c_int;

        self
    }
}

// TODO: Expose a NAL abstraction
pub struct NalData {
    vec: Vec<u8>,
}

impl NalData {
    /*
     * x264 functions return x264_nal_t arrays that are valid only until another
     * function of that kind is called.
     *
     * Always copy the data over.
     *
     * TODO: Consider using Bytes as backing store.
     */
    fn from_nals(c_nals: *mut x264_nal_t, nb_nal: usize) -> NalData {
        let mut data = NalData { vec: Vec::new() };

        for i in 0..nb_nal {
            let nal = unsafe { Box::from_raw(c_nals.offset(i as isize)) };

            let payload =
                unsafe { std::slice::from_raw_parts(nal.p_payload, nal.i_payload as usize) };

            data.vec.extend_from_slice(payload);

            mem::forget(payload);
            mem::forget(nal);
        }

        data
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.vec.as_slice()
    }
}

pub struct Encoder {
    enc: *mut x264_t,
}

impl Encoder {
    pub fn open(par: &mut Param) -> Result<Encoder, &'static str> {
        let enc = unsafe { x264_encoder_open(&mut par.par as *mut x264_param_t) };

        if enc.is_null() {
            Err("Out of Memory")
        } else {
            Ok(Encoder { enc: enc })
        }
    }

    pub fn get_headers(&mut self) -> Result<NalData, &'static str> {
        let mut nb_nal: c_int = 0;
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };

        let bytes = unsafe {
            x264_encoder_headers(self.enc,
                                 &mut c_nals as *mut *mut x264_nal_t,
                                 &mut nb_nal as *mut c_int)
        };

        if bytes < 0 {
            Err("Encoding Headers Failed")
        } else {
            Ok(NalData::from_nals(c_nals, nb_nal as usize))
        }
    }

    pub fn encode<'a, P>(&mut self, pic: P) -> Result<Option<(NalData, i64, i64)>, &'static str>
        where P: Into<Option<&'a Picture>>
    {
        let mut pic_out: x264_picture_t = unsafe { mem::uninitialized() };
        let mut c_nals: *mut x264_nal_t = unsafe { mem::uninitialized() };
        let mut nb_nal: c_int = 0;
        let c_pic = pic.into().map_or_else(|| null(), |v| &v.pic as *const x264_picture_t);

        let ret = unsafe {
            x264_encoder_encode(self.enc,
                                &mut c_nals as *mut *mut x264_nal_t,
                                &mut nb_nal as *mut c_int,
                                c_pic as *mut x264_picture_t,
                                &mut pic_out as *mut x264_picture_t)
        };
        if ret < 0 {
            Err("Error encoding")
        } else {
            if nb_nal > 0 {
                let data = NalData::from_nals(c_nals, nb_nal as usize);
                Ok(Some((data, pic_out.i_pts, pic_out.i_dts)))
            } else {
                Ok(None)
            }
        }
    }

    pub fn delayed_frames(&self) -> bool {
        unsafe { x264_encoder_delayed_frames(self.enc) != 0 }
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe { x264_encoder_close(self.enc) };
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_open() {
        let mut par = Param::new().set_dimension(640, 480);

        let mut enc = Encoder::open(&mut par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());
    }

    #[test]
    fn test_picture() {
        let par = Param::new().set_dimension(640, 480);
        {
            let mut pic = Picture::from_param(&par).unwrap();
            {
                let p = pic.as_slice_mut(0).unwrap();
                p[0] = 1;
            }
            let p = pic.as_slice(0).unwrap();

            assert_eq!(p[0], 1);
        }
    }

    #[test]
    fn test_encode() {
        let mut par = Param::new().set_dimension(640, 480);
        let mut enc = Encoder::open(&mut par).unwrap();
        let mut pic = Picture::from_param(&par).unwrap();

        let headers = enc.get_headers().unwrap();

        println!("Headers len {}", headers.as_bytes().len());

        for pts in 0..5 {
            pic = pic.set_timestamp(pts as i64);
            let ret = enc.encode(&pic).unwrap();
            match ret {
                Some((_, pts, dts)) => println!("Frame pts {}, dts {}", pts, dts),
                _ => (),
            }
        }

        while enc.delayed_frames() {
            let ret = enc.encode(None).unwrap();
            match ret {
                Some((_, pts, dts)) => println!("Frame pts {}, dts {}", pts, dts),
                _ => (),
            }
        }
    }
}
<|endoftext|>


<commit_before>/*!

A basic library for finding primes, providing a basic Iterator over all primes. It is not as fast as
`slow_primes`, but it is meant to be easy to use!

The simplest usage is simply to create an `Iterator`:

```
use primes::PrimeSet;

let mut pset = PrimeSet::new();

for (ix, n) in pset.iter().enumerate().take(10) {
    println!("Prime {}: {}", ix, n);
}
```

This library provides methods for generating primes, testing whether a number is prime, and
factorizing numbers. Most methods generate primes lazily, so only enough primes will be generated
for the given test, and primes are cached for later use.

[*Source*](https://github.com/wackywendell/primes)

# Example: Find the millionth prime

```
use primes::PrimeSet;

let mut pset = PrimeSet::new();
let (ix, n) = pset.find(1_000_000);

println!("Prime {}: {}", ix, n);
```

# Example: Find the first ten primes *after* the thousandth prime
```
use primes::PrimeSet;

let mut pset = PrimeSet::new();
for (ix, n) in pset.iter().enumerate().skip(1_000).take(10) {
    println!("Prime {}: {}", ix, n);
}
```

# Example: Find the first prime greater than 1000
```
use primes::PrimeSet;

let mut pset = PrimeSet::new();
let (ix, n) = pset.find(1_000);
println!("The first prime after 1000 is the {}th prime: {}", ix, n);

assert_eq!(pset.find(n), (ix, n));
```

For more info on use, see `PrimeSet`, a class which handles the Sieve and has multiple methods for
iterating over primes.

This also provides a few functions unconnected to `PrimeSet`, which will be faster for the first
case, but slower in the long term as they do not use any caching of primes.

*/
#![feature(core)]
#![feature(test)]

#[warn(non_camel_case_types)]
#[warn(non_snake_case)]
#[warn(unused_qualifications)]
#[warn(non_upper_case_globals)]
#[warn(missing_docs)]

extern crate test;

use std::ops::Index;
use std::slice;
use std::iter;
use std::num::{Float,cast};
use std::cmp::Ordering::{Equal,Less,Greater};

#[cfg(test)]
use test::Bencher;

fn sqrt_floor<T: std::num::NumCast>(n : T) -> T {
    cast::<f64, T>(
        (cast::<T, f64>(n).unwrap()).sqrt().floor()
    ).unwrap()
}

fn sqrt_ceil<T: std::num::NumCast>(n : T) -> T {
    cast::<f64, T>(
        (cast::<T, f64>(n).unwrap()).sqrt().ceil()
    ).unwrap()
}

/** A prime generator, using the Sieve of Eratosthenes.

Create with `let mut pset = PrimeSet::new()`, and then use `pset.iter()` to iterate over all primes.
**/
pub struct PrimeSet {
    lst : Vec<u64>
}

/// An iterator over generated primes. Created by PrimeSet::iter or
/// PrimeSet::generator
pub struct PrimeSetIter<'a> {
    p : &'a mut PrimeSet,
    n : usize,
    expand : bool
}

impl PrimeSet {
    /// A new prime generator, primed with 2 and 3
    pub fn new() -> PrimeSet {
        PrimeSet{lst:vec!(2,3)}
    }

    /// Finds one more prime, and adds it to the list
    pub fn expand(&mut self) {
        let mut l : u64 = self.lst[self.lst.len()-1] + 2;
        let mut sql = sqrt_floor(l);
        let mut remainder = 0;
        loop {
            for &n in self.lst.iter() {
                remainder = l % n;
                if remainder == 0 || n > sql {
                    break;
                }
            };

            if remainder != 0 {
                self.lst.push(l);
                break;
            };

            l += 2;
            sql = sqrt_floor(l);
        }
    }

    /// Number of primes found so far
    pub fn len(&self) -> usize {
        self.lst.len()
    }

    /// Return all primes found so far as a slice
    pub fn list<'a>(&'a self) -> &'a [u64] {
        &self.lst[..]
    }

    /// Iterator over all primes not yet found
    pub fn generator<'a>(&'a mut self) -> PrimeSetIter<'a> {
        let myn = self.len();
        PrimeSetIter{p:self, n:myn, expand:true}
    }

    /// Iterator over all primes, starting with 2. If you don't care about the "state" of the
    /// PrimeSet, this is what you want!
    pub fn iter<'a>(&'a mut self) -> PrimeSetIter<'a> {
        PrimeSetIter{p:self, n:0, expand:true}
    }

    //~ pub fn iter_once(&'self mut self) -> PrimeSetIter<'self> {
        //~ PrimeSetIter{p:self, n:0, expand:false}
    //~ }

    /// Iterator over just the primes found so far
    pub fn iter_vec<'a>(&'a self) -> slice::Iter<'a, u64> {
        self.lst.iter()
    }

    /// Find the next largest prime from a number
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find(&mut self, n: u64) -> (usize, u64) {
        while n > *(self.lst.last().unwrap_or(&0)){
            self.expand();
        }
        self.find_vec(n).unwrap()
    }

    /// Check if a number is prime
    /// Note that this only requires primes up to n.sqrt() to be generated, and will generate
    /// them as necessary on its own.
    pub fn is_prime(&mut self, n: u64) -> bool {
        if n <= 1 {return false;}
        if n == 2 {return true;} // otherwise we get 2 % 2 == 0!
        for m in self.iter() {
            if n % m == 0 {return false;};
            if m*m > n {return true;};
        }
        panic!("This iterator should not be empty.");
    }

    /// Find the next largest prime from a number, if it is within the already-found list
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find_vec(&self, n: u64) -> Option<(usize, u64)> {
        if n > *(self.lst.last().unwrap_or(&0)){ return None;}

        let mut base : usize = 0;
        let mut lim : usize = self.len();

        // Binary search algorithm
        while lim != 0 {
            let ix = base + (lim >> 1);
            match self.lst[ix].cmp(&n) {
                Equal => return Some((ix, self.lst[ix])),
                Less => {
                    base = ix + 1;
                    lim -= 1;
                }
                Greater => ()
            }
            lim >>= 1;
        }
        return Some((base, self.lst[base]));
    }

    /// Get the nth prime, even if we haven't yet found it
    pub fn get(&mut self, index : &usize) -> &u64 {
		for _ in (0..(*index as isize) + 1 - (self.lst.len() as isize)){
			self.expand();
		}
        self.lst.index(index)
	}

	/// Get the prime factors of a number, starting from 2, including repeats
	pub fn prime_factors(&mut self, n: u64) -> Vec<u64> {
		if n == 1 {return Vec::new();}
		let mut curn = n;
		let mut m = sqrt_ceil(curn);
		let mut lst: Vec<u64> = Vec::new();
		for p in self.iter() {
			while curn % p == 0 {
				lst.push(p);
				curn /= p;
				if curn == 1 {return lst;}
				m = sqrt_ceil(curn);
			}

			if p > m {
				lst.push(p);
				return lst;
			}
		}
		panic!("This should be unreachable.");
	}
}

impl Index<usize> for PrimeSet {
    type Output = u64;
    fn index(&self, index: &usize) -> &u64 {
        self.lst.index(index)
    }
}

impl<'a> Iterator for PrimeSetIter<'a> {
    type Item = u64;
    fn next(&mut self) -> Option<u64> {
        while self.n >= self.p.len(){
            match self.expand {
                true => self.p.expand(),
                false => return None
            }
        }
        self.n += 1;

        let m = self.p.lst[self.n-1];

        Some(m)
    }
}

/// Find the first factor (other than 1) of a number
fn firstfac(x: u64) -> u64 {
    let m = sqrt_ceil(x);
    if x % 2 == 0 { return 2; };
    for n in iter::range_step(3, m + 1, 2) {
        if x % n == 0 { return n; };
    }
    return x;
}

/// Find all prime factors of a number
/// Does not use a PrimeSet, but simply counts upwards
pub fn factors(x: u64) -> Vec<u64> {
    if x <= 1 {return vec!()};
    let mut lst: Vec<u64> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if m == curn { break  } else { curn /= m };
    }
    return lst
}

/// Find all unique prime factors of a number
pub fn factors_uniq(x: u64) -> Vec<u64> {
    if x <= 1 {return vec!()};
    let mut lst: Vec<u64> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if curn == m { break ; }
        while curn % m == 0 { curn /= m; }
        if curn == 1 { break ; }
    }
    return lst
}

/// Test whether a number is prime. Checks every odd number up to sqrt(n).
pub fn is_prime(n : u64) -> bool {
    if n <= 1 {return false;}
    firstfac(n) == n
}

#[test]
fn test_iter(){
    let mut pset = PrimeSet::new();
    let first_few = [2u64,3,5,7,11,13,17,19,23];
    for (m, &n) in pset.iter().zip(first_few.iter()) {
        assert_eq!(m, n);
    }
}

#[test]
fn test_find(){
    let mut pset = PrimeSet::new();

    // pset is empty, so it needs to generate the primes
    assert_eq!(pset.find_vec(1000), None);

    let (ix_exp, n_exp) = (168, 1009);

    assert_eq!(pset.find(1000), (ix_exp, n_exp));
    assert_eq!(pset.find(n_exp), (ix_exp, n_exp));

    // We shouldn't have gone beyond 1009
    {
        let plst = pset.list();
        let plen = plst.len();
        assert_eq!(plen, ix_exp+1);

        assert_eq!(plst[plen-1], n_exp);
    }

    assert_eq!(pset.find_vec(n_exp), Some((ix_exp, n_exp)));
}

#[test]
fn test_primes(){
    let mut pset = PrimeSet::new();

    // note: some are repeated, because the pset list grows as it goes

    assert!(!pset.is_prime(1));
    assert!(!is_prime(1));
    assert!(pset.is_prime(2));
    assert!(is_prime(2));
    assert!(pset.is_prime(13));
    assert!(is_prime(13));
    assert!(!pset.is_prime(45));
    assert!(!is_prime(45));
    assert!(!pset.is_prime(13*13));
    assert!(!is_prime(13*13));
    assert!(pset.is_prime(13));
    assert!(pset.is_prime(7));
    assert!(is_prime(7));
    assert!(!pset.is_prime(9));
    assert!(!is_prime(9));
    assert!(pset.is_prime(5));
    assert!(is_prime(5));
}

#[test]
fn test_factors(){
	let mut pset = PrimeSet::new();

	let ns = [  (1, vec!()),
                (2, vec!(2)),
                (3, vec!(3)),
                (4, vec!(2,2)),
                (5, vec!(5)),
                (6, vec!(2,3)),
	            (9, vec!(3,3)),
	            (12, vec!(2,2,3)),
	            (121, vec!(11,11)),
	            (144, vec!(2,2,2,2,3,3)),
	            (10_000_000, vec!(2,2,2,2,2,2,2,5,5,5,5,5,5,5)),
                (100, vec!(2,2,5,5)),
                (121, vec!(11, 11)),
                ];

    // Test unique factors
    for &(n, ref v) in ns.iter(){
        assert_eq!(pset.prime_factors(n), *v);
        assert_eq!(factors(n), *v);

        let unique_factors = factors_uniq(n);

        // Get unique factors from the lists we made above
        let mut unique_factors_exp : Vec<u64> = v.iter().map(|&x| {x}).collect();
        unique_factors_exp.dedup();

        assert_eq!(unique_factors, unique_factors_exp);
    }

	pset = PrimeSet::new();
	assert_eq!(pset.prime_factors(12), vec!(2,2,3));

}

#[bench]
fn bench_primes(b : &mut Bencher){
	b.iter(|| {
		let mut pset = PrimeSet::new();
		 let (_, _) = pset.find(1_000_000);
		 //~ let (idx, n) = pset.find(1_000_000);
		 //~ println!("Prime {}: {}", idx, n);
		 })
}
<commit_msg>Updated to newest Rust<commit_after>/*!

A basic library for finding primes, providing a basic Iterator over all primes. It is not as fast as
`slow_primes`, but it is meant to be easy to use!

The simplest usage is simply to create an `Iterator`:

```
use primes::PrimeSet;

let mut pset = PrimeSet::new();

for (ix, n) in pset.iter().enumerate().take(10) {
    println!("Prime {}: {}", ix, n);
}
```

This library provides methods for generating primes, testing whether a number is prime, and
factorizing numbers. Most methods generate primes lazily, so only enough primes will be generated
for the given test, and primes are cached for later use.

[*Source*](https://github.com/wackywendell/primes)

# Example: Find the millionth prime

```
use primes::PrimeSet;

let mut pset = PrimeSet::new();
let (ix, n) = pset.find(1_000_000);

println!("Prime {}: {}", ix, n);
```

# Example: Find the first ten primes *after* the thousandth prime
```
use primes::PrimeSet;

let mut pset = PrimeSet::new();
for (ix, n) in pset.iter().enumerate().skip(1_000).take(10) {
    println!("Prime {}: {}", ix, n);
}
```

# Example: Find the first prime greater than 1000
```
use primes::PrimeSet;

let mut pset = PrimeSet::new();
let (ix, n) = pset.find(1_000);
println!("The first prime after 1000 is the {}th prime: {}", ix, n);

assert_eq!(pset.find(n), (ix, n));
```

For more info on use, see `PrimeSet`, a class which handles the Sieve and has multiple methods for
iterating over primes.

This also provides a few functions unconnected to `PrimeSet`, which will be faster for the first
case, but slower in the long term as they do not use any caching of primes.

*/
#![feature(core,step_by,test)]

#[warn(non_camel_case_types)]
#[warn(non_snake_case)]
#[warn(unused_qualifications)]
#[warn(non_upper_case_globals)]
#[warn(missing_docs)]

extern crate test;

use std::ops::Index;
use std::slice;
use std::num::{Float,cast};
use std::cmp::Ordering::{Equal,Less,Greater};

#[cfg(test)]
use test::Bencher;

fn sqrt_floor<T: std::num::NumCast>(n : T) -> T {
    cast::<f64, T>(
        (cast::<T, f64>(n).unwrap()).sqrt().floor()
    ).unwrap()
}

fn sqrt_ceil<T: std::num::NumCast>(n : T) -> T {
    cast::<f64, T>(
        (cast::<T, f64>(n).unwrap()).sqrt().ceil()
    ).unwrap()
}

/** A prime generator, using the Sieve of Eratosthenes.

Create with `let mut pset = PrimeSet::new()`, and then use `pset.iter()` to iterate over all primes.
**/
pub struct PrimeSet {
    lst : Vec<u64>
}

/// An iterator over generated primes. Created by PrimeSet::iter or
/// PrimeSet::generator
pub struct PrimeSetIter<'a> {
    p : &'a mut PrimeSet,
    n : usize,
    expand : bool
}

impl PrimeSet {
    /// A new prime generator, primed with 2 and 3
    pub fn new() -> PrimeSet {
        PrimeSet{lst:vec!(2,3)}
    }

    /// Finds one more prime, and adds it to the list
    pub fn expand(&mut self) {
        let mut l : u64 = self.lst[self.lst.len()-1] + 2;
        let mut sql = sqrt_floor(l);
        let mut remainder = 0;
        loop {
            for &n in self.lst.iter() {
                remainder = l % n;
                if remainder == 0 || n > sql {
                    break;
                }
            };

            if remainder != 0 {
                self.lst.push(l);
                break;
            };

            l += 2;
            sql = sqrt_floor(l);
        }
    }

    /// Number of primes found so far
    pub fn len(&self) -> usize {
        self.lst.len()
    }

    /// Return all primes found so far as a slice
    pub fn list<'a>(&'a self) -> &'a [u64] {
        &self.lst[..]
    }

    /// Iterator over all primes not yet found
    pub fn generator<'a>(&'a mut self) -> PrimeSetIter<'a> {
        let myn = self.len();
        PrimeSetIter{p:self, n:myn, expand:true}
    }

    /// Iterator over all primes, starting with 2. If you don't care about the "state" of the
    /// PrimeSet, this is what you want!
    pub fn iter<'a>(&'a mut self) -> PrimeSetIter<'a> {
        PrimeSetIter{p:self, n:0, expand:true}
    }

    //~ pub fn iter_once(&'self mut self) -> PrimeSetIter<'self> {
        //~ PrimeSetIter{p:self, n:0, expand:false}
    //~ }

    /// Iterator over just the primes found so far
    pub fn iter_vec<'a>(&'a self) -> slice::Iter<'a, u64> {
        self.lst.iter()
    }

    /// Find the next largest prime from a number
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find(&mut self, n: u64) -> (usize, u64) {
        while n > *(self.lst.last().unwrap_or(&0)){
            self.expand();
        }
        self.find_vec(n).unwrap()
    }

    /// Check if a number is prime
    /// Note that this only requires primes up to n.sqrt() to be generated, and will generate
    /// them as necessary on its own.
    pub fn is_prime(&mut self, n: u64) -> bool {
        if n <= 1 {return false;}
        if n == 2 {return true;} // otherwise we get 2 % 2 == 0!
        for m in self.iter() {
            if n % m == 0 {return false;};
            if m*m > n {return true;};
        }
        panic!("This iterator should not be empty.");
    }

    /// Find the next largest prime from a number, if it is within the already-found list
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find_vec(&self, n: u64) -> Option<(usize, u64)> {
        if n > *(self.lst.last().unwrap_or(&0)){ return None;}

        let mut base : usize = 0;
        let mut lim : usize = self.len();

        // Binary search algorithm
        while lim != 0 {
            let ix = base + (lim >> 1);
            match self.lst[ix].cmp(&n) {
                Equal => return Some((ix, self.lst[ix])),
                Less => {
                    base = ix + 1;
                    lim -= 1;
                }
                Greater => ()
            }
            lim >>= 1;
        }
        return Some((base, self.lst[base]));
    }

    /// Get the nth prime, even if we haven't yet found it
    pub fn get(&mut self, index : usize) -> u64 {
		for _ in (0..(index as isize) + 1 - (self.lst.len() as isize)){
			self.expand();
		}
        self.lst[index]
	}

	/// Get the prime factors of a number, starting from 2, including repeats
	pub fn prime_factors(&mut self, n: u64) -> Vec<u64> {
		if n == 1 {return Vec::new();}
		let mut curn = n;
		let mut m = sqrt_ceil(curn);
		let mut lst: Vec<u64> = Vec::new();
		for p in self.iter() {
			while curn % p == 0 {
				lst.push(p);
				curn /= p;
				if curn == 1 {return lst;}
				m = sqrt_ceil(curn);
			}

			if p > m {
				lst.push(p);
				return lst;
			}
		}
		panic!("This should be unreachable.");
	}
}

impl Index<usize> for PrimeSet {
    type Output = u64;
    fn index(&self, index: usize) -> &u64 {
        &self.lst[index]
    }
}

impl<'a> Iterator for PrimeSetIter<'a> {
    type Item = u64;
    fn next(&mut self) -> Option<u64> {
        while self.n >= self.p.len(){
            match self.expand {
                true => self.p.expand(),
                false => return None
            }
        }
        self.n += 1;

        let m = self.p.lst[self.n-1];

        Some(m)
    }
}

/// Find the first factor (other than 1) of a number
fn firstfac(x: u64) -> u64 {
    let m = sqrt_ceil(x);
    if x % 2 == 0 { return 2; };
    for n in (3..m + 1).step_by(2) {
        if x % n == 0 { return n; };
    }
    return x;
}

/// Find all prime factors of a number
/// Does not use a PrimeSet, but simply counts upwards
pub fn factors(x: u64) -> Vec<u64> {
    if x <= 1 {return vec!()};
    let mut lst: Vec<u64> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if m == curn { break  } else { curn /= m };
    }
    return lst
}

/// Find all unique prime factors of a number
pub fn factors_uniq(x: u64) -> Vec<u64> {
    if x <= 1 {return vec!()};
    let mut lst: Vec<u64> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if curn == m { break ; }
        while curn % m == 0 { curn /= m; }
        if curn == 1 { break ; }
    }
    return lst
}

/// Test whether a number is prime. Checks every odd number up to sqrt(n).
pub fn is_prime(n : u64) -> bool {
    if n <= 1 {return false;}
    firstfac(n) == n
}

#[test]
fn test_iter(){
    let mut pset = PrimeSet::new();
    let first_few = [2u64,3,5,7,11,13,17,19,23];
    for (m, &n) in pset.iter().zip(first_few.iter()) {
        assert_eq!(m, n);
    }
}

#[test]
fn test_find(){
    let mut pset = PrimeSet::new();

    // pset is empty, so it needs to generate the primes
    assert_eq!(pset.find_vec(1000), None);

    let (ix_exp, n_exp) = (168, 1009);

    assert_eq!(pset.find(1000), (ix_exp, n_exp));
    assert_eq!(pset.find(n_exp), (ix_exp, n_exp));

    // We shouldn't have gone beyond 1009
    {
        let plst = pset.list();
        let plen = plst.len();
        assert_eq!(plen, ix_exp+1);

        assert_eq!(plst[plen-1], n_exp);
    }

    assert_eq!(pset.find_vec(n_exp), Some((ix_exp, n_exp)));
}

#[test]
fn test_primes(){
    let mut pset = PrimeSet::new();

    // note: some are repeated, because the pset list grows as it goes

    assert!(!pset.is_prime(1));
    assert!(!is_prime(1));
    assert!(pset.is_prime(2));
    assert!(is_prime(2));
    assert!(pset.is_prime(13));
    assert!(is_prime(13));
    assert!(!pset.is_prime(45));
    assert!(!is_prime(45));
    assert!(!pset.is_prime(13*13));
    assert!(!is_prime(13*13));
    assert!(pset.is_prime(13));
    assert!(pset.is_prime(7));
    assert!(is_prime(7));
    assert!(!pset.is_prime(9));
    assert!(!is_prime(9));
    assert!(pset.is_prime(5));
    assert!(is_prime(5));
}

#[test]
fn test_factors(){
	let mut pset = PrimeSet::new();

	let ns = [  (1, vec!()),
                (2, vec!(2)),
                (3, vec!(3)),
                (4, vec!(2,2)),
                (5, vec!(5)),
                (6, vec!(2,3)),
	            (9, vec!(3,3)),
	            (12, vec!(2,2,3)),
	            (121, vec!(11,11)),
	            (144, vec!(2,2,2,2,3,3)),
	            (10_000_000, vec!(2,2,2,2,2,2,2,5,5,5,5,5,5,5)),
                (100, vec!(2,2,5,5)),
                (121, vec!(11, 11)),
                ];

    // Test unique factors
    for &(n, ref v) in ns.iter(){
        assert_eq!(pset.prime_factors(n), *v);
        assert_eq!(factors(n), *v);

        let unique_factors = factors_uniq(n);

        // Get unique factors from the lists we made above
        let mut unique_factors_exp : Vec<u64> = v.iter().map(|&x| {x}).collect();
        unique_factors_exp.dedup();

        assert_eq!(unique_factors, unique_factors_exp);
    }

	pset = PrimeSet::new();
	assert_eq!(pset.prime_factors(12), vec!(2,2,3));

}

#[bench]
fn bench_primes(b : &mut Bencher){
	b.iter(|| {
		let mut pset = PrimeSet::new();
		 let (_, _) = pset.find(1_000_000);
		 //~ let (idx, n) = pset.find(1_000_000);
		 //~ println!("Prime {}: {}", idx, n);
		 })
}
<|endoftext|>


<commit_before>// a Shiny test framework
// Copyright 2014 Vladimir "farcaller" Pouzanov <farcaller@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![crate_name = "shiny"]
#![crate_type = "dylib"]
#![feature(plugin_registrar, managed_boxes)]

extern crate rustc;
extern crate syntax;

use rustc::plugin::Registry;
use syntax::abi;
use syntax::ptr::P;
use syntax::ast;
use syntax::ast_util::empty_generics;
use syntax::codemap::DUMMY_SP;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacResult};
use syntax::ext::build::AstBuilder;
use syntax::parse::token;
use syntax::parse::tts_to_parser;
use syntax::util::small_vector::SmallVector;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("describe", macro_describe);
}

pub fn macro_describe(cx: &mut ExtCtxt, _: Span, tts: &[ast::TokenTree]) -> Box<MacResult+'static> {
    let sess = cx.parse_sess();
    let ttsvec = tts.iter().map(|x| (*x).clone()).collect();

    let mut parser = tts_to_parser(sess, ttsvec, cx.cfg());

    let mut before_block = None;
    let mut test_blocks = vec!();

    loop {
        if parser.token == token::EOF {
            break;
        }

        let ident = parser.parse_ident();
        match ident.as_str() {
            "before_each" => {
                if before_block.is_some() {
                    fail!("only one before_each block is allowed");
                }
                before_block = Some(parser.parse_block());
            },
            "it" => {
                let (name, _) = parser.parse_str();
                let block = parser.parse_block();
                test_blocks.push((name.get().to_string(), block));
            }
            other => {
                let span = parser.span;
                parser.span_fatal(span, format!("expected one of: {} but found `{}`",
                    "`before_each`, `it`",
                    other).as_slice());
            },
        }
    }

    let mut funcs = vec!();

    for &(ref name, ref block) in test_blocks.iter() {
        let body = match before_block {
            None => block.clone(),
            Some(ref before) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items,
                    stmts: before.stmts + block.stmts,

                    ..block.deref().clone()
                })
            }
        };

        let attr_test = cx.attribute(DUMMY_SP,
            cx.meta_word(DUMMY_SP, token::InternedString::new("test")));

        let func = P(ast::Item {
            ident: cx.ident_of(name.replace(" ", "_").as_slice()),
            attrs: vec!(attr_test),
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                cx.fn_decl(Vec::new(), cx.ty_nil()),
                ast::NormalFn,
                abi::Rust,
                empty_generics(),
                body),
            vis: ast::Inherited,
            span: DUMMY_SP,
        });
        funcs.push(func);
    }

    MacItems::new(funcs)
}

pub struct MacItems {
    items: Vec<P<ast::Item>>
}

impl MacItems {
    pub fn new(items: Vec<P<ast::Item>>) -> Box<MacResult+'static> {
        box MacItems { items: items } as Box<MacResult>
    }
}

impl MacResult for MacItems {
    fn make_items(self: Box<MacItems>) -> Option<SmallVector<P<ast::Item>>> {
        Some(SmallVector::many(self.items.clone()))
    }
}
<commit_msg>managed_boxes are gone<commit_after>// a Shiny test framework
// Copyright 2014 Vladimir "farcaller" Pouzanov <farcaller@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![crate_name = "shiny"]
#![crate_type = "dylib"]
#![feature(plugin_registrar)]

extern crate rustc;
extern crate syntax;

use rustc::plugin::Registry;
use syntax::abi;
use syntax::ptr::P;
use syntax::ast;
use syntax::ast_util::empty_generics;
use syntax::codemap::DUMMY_SP;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacResult};
use syntax::ext::build::AstBuilder;
use syntax::parse::token;
use syntax::parse::tts_to_parser;
use syntax::util::small_vector::SmallVector;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("describe", macro_describe);
}

pub fn macro_describe(cx: &mut ExtCtxt, _: Span, tts: &[ast::TokenTree]) -> Box<MacResult+'static> {
    let sess = cx.parse_sess();
    let ttsvec = tts.iter().map(|x| (*x).clone()).collect();

    let mut parser = tts_to_parser(sess, ttsvec, cx.cfg());

    let mut before_block = None;
    let mut test_blocks = vec!();

    loop {
        if parser.token == token::EOF {
            break;
        }

        let ident = parser.parse_ident();
        match ident.as_str() {
            "before_each" => {
                if before_block.is_some() {
                    fail!("only one before_each block is allowed");
                }
                before_block = Some(parser.parse_block());
            },
            "it" => {
                let (name, _) = parser.parse_str();
                let block = parser.parse_block();
                test_blocks.push((name.get().to_string(), block));
            }
            other => {
                let span = parser.span;
                parser.span_fatal(span, format!("expected one of: {} but found `{}`",
                    "`before_each`, `it`",
                    other).as_slice());
            },
        }
    }

    let mut funcs = vec!();

    for &(ref name, ref block) in test_blocks.iter() {
        let body = match before_block {
            None => block.clone(),
            Some(ref before) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items,
                    stmts: before.stmts + block.stmts,

                    ..block.deref().clone()
                })
            }
        };

        let attr_test = cx.attribute(DUMMY_SP,
            cx.meta_word(DUMMY_SP, token::InternedString::new("test")));

        let func = P(ast::Item {
            ident: cx.ident_of(name.replace(" ", "_").as_slice()),
            attrs: vec!(attr_test),
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                cx.fn_decl(Vec::new(), cx.ty_nil()),
                ast::NormalFn,
                abi::Rust,
                empty_generics(),
                body),
            vis: ast::Inherited,
            span: DUMMY_SP,
        });
        funcs.push(func);
    }

    MacItems::new(funcs)
}

pub struct MacItems {
    items: Vec<P<ast::Item>>
}

impl MacItems {
    pub fn new(items: Vec<P<ast::Item>>) -> Box<MacResult+'static> {
        box MacItems { items: items } as Box<MacResult>
    }
}

impl MacResult for MacItems {
    fn make_items(self: Box<MacItems>) -> Option<SmallVector<P<ast::Item>>> {
        Some(SmallVector::many(self.items.clone()))
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

#![feature(plugin)]
#![plugin(clippy)]

extern crate test;
extern crate rand;

use rand::distributions::{IndependentSample, Range};

pub struct Settings<R>
    where R: rand::Rng
{
    pub min_pos: Vec<f32>,
    pub max_pos: Vec<f32>,

    pub cr_min: f32,
    pub cr_max: f32,
    pub cr_change_probability: f32,

    pub f_min: f32,
    pub f_max: f32,
    pub f_change_probability: f32,

    pub pop_size: usize,
    pub rng: R,
}

impl<R> Settings<R>
    where R: rand::Rng
{
    pub fn min_max_rng(min_pos: Vec<f32>, max_pos: Vec<f32>, rng: R) -> Settings<R> {
        // create settings for the algorithm
        Settings {
            min_pos: min_pos,
            max_pos: max_pos,

            cr_min: 0.0,
            cr_max: 1.0,
            cr_change_probability: 0.1,

            f_min: 0.1,
            f_max: 1.0,
            f_change_probability: 0.1,

            pop_size: 50,
            rng: rng,
        }
    }
}

impl Settings<rand::XorShiftRng> {
    pub fn new(min_pos: Vec<f32>, max_pos: Vec<f32>) -> Settings<rand::XorShiftRng> {
        Settings::min_max_rng(min_pos, max_pos, rand::weak_rng())
    }
}

#[derive(Clone,Debug)]
pub struct Individual {
    pub pos: Vec<f32>,
    // the lower, the better.
    pub cost: Option<f32>,

    // control parameters
    cr: f32,
    f: f32,
}

pub struct Population<R>
    where R: rand::Rng
{
    // TODO use a single vector for curr and best and controlparameters?
    pub curr: Vec<Individual>,
    best: Vec<Individual>,

    settings: Settings<R>,
    best_idx: Option<usize>,
    dim: usize,
    between_popsize: Range<usize>,
    between_01: Range<f32>,
    between_dim: Range<usize>,
    between_cr: Range<f32>,
    between_f: Range<f32>,
}

impl<R> Population<R>
    where R: rand::Rng
{
    // Creates a new population based on the given settings.
    pub fn new(s: Settings<R>) -> Population<R> {
        assert_eq!(s.min_pos.len(),
                   s.max_pos.len(),
                   "min_pos and max_pos need to have the same number of elements");
        assert!(s.min_pos.len() >= 1,
                "need at least one element to optimize");

        // create a vector of randomly initialized individuals for current.
        let dim = s.min_pos.len();

        // Empty individual, with no cost value (yet)
        let dummy_individual = Individual {
            pos: vec![0.0; dim],
            cost: None,
            cr: 0.0,
            f: 0.0,
        };

        // creates all the empty individuals
        let mut pop = Population {
            curr: vec![dummy_individual.clone(); s.pop_size],
            best: vec![dummy_individual; s.pop_size],
            best_idx: None,
            dim: dim,
            between_popsize: Range::new(0, s.pop_size),
            between_01: Range::new(0.0, 1.0),
            between_dim: Range::new(0, dim),
            between_cr: Range::new(s.cr_min, s.cr_max),
            between_f: Range::new(s.f_min, s.f_max),
            settings: s,
        };

        for ind in &mut pop.curr {
            // init control parameters
            ind.cr = pop.between_cr.ind_sample(&mut pop.settings.rng);
            ind.f = pop.between_f.ind_sample(&mut pop.settings.rng);

            // random range for each dimension
            for d in 0..dim {
                let between_min_max = Range::new(pop.settings.min_pos[d], pop.settings.max_pos[d]);
                ind.pos[d] = between_min_max.ind_sample(&mut pop.settings.rng);
            }
        }

        pop
    }

    fn update_best(&mut self) -> bool {
        let last_best_cost: Option<f32>;
        if let Some(last_best_idx) = self.best_idx {
            last_best_cost = self.best[last_best_idx].cost;
        } else {
            last_best_cost = None;
        }

        let mut new_best_cost = last_best_cost;
        for i in 0..self.curr.len() {
            let curr = &mut self.curr[i];
            let best = &mut self.best[i];

            // we use <= here, so that the individual moves even if the cost
            // stays the same.
            if best.cost.is_none() || curr.cost.unwrap() <= best.cost.unwrap() {
                // find global best. We use < here so that global only updates when fitness changes.
                if new_best_cost.is_none() || curr.cost.unwrap() < new_best_cost.unwrap() {
                    self.best_idx = Some(i);
                    new_best_cost = curr.cost;
                }

                // replace individual's best. swap is *much* faster than clone.
                std::mem::swap(curr, best);
            }
        }

        // got a new best?
        last_best_cost.is_none() || new_best_cost.unwrap() < last_best_cost.unwrap()
    }

    fn update_positions(&mut self) {
        let best_idx = self.best_idx.unwrap();
        let rng = &mut self.settings.rng;
        for i in 0..self.curr.len() {
            let mut id1 = i;
            while id1 == i {
                id1 = self.between_popsize.ind_sample(rng);
            }

            let mut id2 = i;
            while id2 == i || id2 == id1 {
                id2 = self.between_popsize.ind_sample(rng);
            }

            let curr = &mut self.curr[i];
            let best = &self.best[i];

            // see "Self-Adapting Control Parameters in Differential Evolution:
            // A Comparative Study on Numerical Benchmark Problems"
            curr.cr = best.cr;
            curr.f = best.f;
            if self.between_01.ind_sample(rng) < self.settings.cr_change_probability {
                curr.cr = self.between_cr.ind_sample(rng);
            }
            if self.between_01.ind_sample(rng) < self.settings.f_change_probability {
                curr.f = self.between_f.ind_sample(rng);
            }

            let forced_mutation_dim = self.between_dim.ind_sample(rng);
            for d in 0..self.dim {
                if d == forced_mutation_dim || self.between_01.ind_sample(rng) < curr.cr {

                    curr.pos[d] = self.best[best_idx].pos[d] +
                                  curr.f * (self.best[id1].pos[d] - self.best[id2].pos[d]);
                } else {
                    curr.pos[d] = self.best[i].pos[d];
                }
            }

            // reset cost, has to be updated by the user.
            curr.cost = None;
        }
    }

    // Uses updated cost values to update positions of individuals.
    pub fn evolve(&mut self) -> Option<&Individual> {
        let found_new_best = self.update_best();
        self.update_positions();

        if found_new_best {
            Some(&self.best[self.best_idx.unwrap()])
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use test::{black_box, Bencher};
    use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng};
    use rand::{OsRng, weak_rng};

    fn setup<R: rand::Rng>(dim: usize, rng: R) -> Population<R> {
        let settings = Settings::min_max_rng(vec![-20.0; dim], vec![20.0; dim], rng);
        Population::new(settings)
    }

    fn dummy_fitness(individuals: &mut Vec<Individual>) {
        for ind in individuals {
            ind.cost = Some(1.234);
        }
    }

    #[bench]
    fn bench_square_fitness_opt(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            for ind in &mut pop.curr {
                let mut f = 0.0;
                for x in &ind.pos {
                    f += x * x;
                }
                ind.cost = Some(f);
            }
            black_box(pop.evolve());
        });
    }

    #[bench]
    fn rand_thread_rng(b: &mut Bencher) {
        let mut pop = setup(5, rand::thread_rng());
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            black_box(pop.evolve());
        });
    }

    #[bench]
    fn rand_xor_shift(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac(b: &mut Bencher) {
        let rng: IsaacRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac64(b: &mut Bencher) {
        let rng: Isaac64Rng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_std(b: &mut Bencher) {
        let rng: StdRng = StdRng::new().unwrap();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_weak_rng(b: &mut Bencher) {
        let rng = weak_rng();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

}
<commit_msg>speedup: use rng.next_f32() instead of range<commit_after>#![feature(test)]

#![feature(plugin)]
#![plugin(clippy)]

extern crate test;
extern crate rand;

use rand::distributions::{IndependentSample, Range};

pub struct Settings<R>
    where R: rand::Rng
{
    pub min_pos: Vec<f32>,
    pub max_pos: Vec<f32>,

    pub cr_min: f32,
    pub cr_max: f32,
    pub cr_change_probability: f32,

    pub f_min: f32,
    pub f_max: f32,
    pub f_change_probability: f32,

    pub pop_size: usize,
    pub rng: R,
}

impl<R> Settings<R>
    where R: rand::Rng
{
    pub fn min_max_rng(min_pos: Vec<f32>, max_pos: Vec<f32>, rng: R) -> Settings<R> {
        // create settings for the algorithm
        Settings {
            min_pos: min_pos,
            max_pos: max_pos,

            cr_min: 0.0,
            cr_max: 1.0,
            cr_change_probability: 0.1,

            f_min: 0.1,
            f_max: 1.0,
            f_change_probability: 0.1,

            pop_size: 50,
            rng: rng,
        }
    }
}

impl Settings<rand::XorShiftRng> {
    pub fn new(min_pos: Vec<f32>, max_pos: Vec<f32>) -> Settings<rand::XorShiftRng> {
        Settings::min_max_rng(min_pos, max_pos, rand::weak_rng())
    }
}

#[derive(Clone,Debug)]
pub struct Individual {
    pub pos: Vec<f32>,
    // the lower, the better.
    pub cost: Option<f32>,

    // control parameters
    cr: f32,
    f: f32,
}

pub struct Population<R>
    where R: rand::Rng
{
    // TODO use a single vector for curr and best and controlparameters?
    pub curr: Vec<Individual>,
    best: Vec<Individual>,

    settings: Settings<R>,
    best_idx: Option<usize>,
    dim: usize,
    between_popsize: Range<usize>,
    between_dim: Range<usize>,
    between_cr: Range<f32>,
    between_f: Range<f32>,
}

impl<R> Population<R>
    where R: rand::Rng
{
    // Creates a new population based on the given settings.
    pub fn new(s: Settings<R>) -> Population<R> {
        assert_eq!(s.min_pos.len(),
                   s.max_pos.len(),
                   "min_pos and max_pos need to have the same number of elements");
        assert!(s.min_pos.len() >= 1,
                "need at least one element to optimize");

        // create a vector of randomly initialized individuals for current.
        let dim = s.min_pos.len();

        // Empty individual, with no cost value (yet)
        let dummy_individual = Individual {
            pos: vec![0.0; dim],
            cost: None,
            cr: 0.0,
            f: 0.0,
        };

        // creates all the empty individuals
        let mut pop = Population {
            curr: vec![dummy_individual.clone(); s.pop_size],
            best: vec![dummy_individual; s.pop_size],
            best_idx: None,
            dim: dim,
            between_popsize: Range::new(0, s.pop_size),
            between_dim: Range::new(0, dim),
            between_cr: Range::new(s.cr_min, s.cr_max),
            between_f: Range::new(s.f_min, s.f_max),
            settings: s,
        };

        for ind in &mut pop.curr {
            // init control parameters
            ind.cr = pop.between_cr.ind_sample(&mut pop.settings.rng);
            ind.f = pop.between_f.ind_sample(&mut pop.settings.rng);

            // random range for each dimension
            for d in 0..dim {
                let between_min_max = Range::new(pop.settings.min_pos[d], pop.settings.max_pos[d]);
                ind.pos[d] = between_min_max.ind_sample(&mut pop.settings.rng);
            }
        }

        pop
    }

    fn update_best(&mut self) -> bool {
        let last_best_cost: Option<f32>;
        if let Some(last_best_idx) = self.best_idx {
            last_best_cost = self.best[last_best_idx].cost;
        } else {
            last_best_cost = None;
        }

        let mut new_best_cost = last_best_cost;
        for i in 0..self.curr.len() {
            let curr = &mut self.curr[i];
            let best = &mut self.best[i];

            // we use <= here, so that the individual moves even if the cost
            // stays the same.
            if best.cost.is_none() || curr.cost.unwrap() <= best.cost.unwrap() {
                // find global best. We use < here so that global only updates when fitness changes.
                if new_best_cost.is_none() || curr.cost.unwrap() < new_best_cost.unwrap() {
                    self.best_idx = Some(i);
                    new_best_cost = curr.cost;
                }

                // replace individual's best. swap is *much* faster than clone.
                std::mem::swap(curr, best);
            }
        }

        // got a new best?
        last_best_cost.is_none() || new_best_cost.unwrap() < last_best_cost.unwrap()
    }

    fn update_positions(&mut self) {
        let global_best_pos = &self.best[self.best_idx.unwrap()].pos;
        let rng = &mut self.settings.rng;
        for i in 0..self.curr.len() {
            let mut id1 = i;
            while id1 == i {
                id1 = self.between_popsize.ind_sample(rng);
            }

            let mut id2 = i;
            while id2 == i || id2 == id1 {
                id2 = self.between_popsize.ind_sample(rng);
            }

            let curr = &mut self.curr[i];
            let best = &self.best[i];

            // see "Self-Adapting Control Parameters in Differential Evolution:
            // A Comparative Study on Numerical Benchmark Problems"
            if rng.next_f32() < self.settings.cr_change_probability {
                curr.cr = self.between_cr.ind_sample(rng);
            } else {
                curr.cr = best.cr;
            }
            if rng.next_f32() < self.settings.f_change_probability {
                curr.f = self.between_f.ind_sample(rng);
            } else {
                curr.f = best.f;
            }

            let curr_pos = &mut curr.pos;
            let best_pos = &best.pos;
            let best1_pos = &self.best[id1].pos;
            let best2_pos = &self.best[id2].pos;

            let forced_mutation_dim = self.between_dim.ind_sample(rng);
            for d in 0..self.dim {
                if d == forced_mutation_dim || rng.next_f32() < curr.cr {
                    curr_pos[d] = global_best_pos[d] + curr.f * (best1_pos[d] - best2_pos[d]);
                } else {
                    curr_pos[d] = best_pos[d];
                }
            }

            // reset cost, has to be updated by the user.
            curr.cost = None;
        }
    }

    // Uses updated cost values to update positions of individuals.
    pub fn evolve(&mut self) -> Option<&Individual> {
        let found_new_best = self.update_best();
        self.update_positions();

        if found_new_best {
            Some(&self.best[self.best_idx.unwrap()])
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use test::Bencher;
    use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng};
    use rand::{OsRng, weak_rng};

    fn setup<R: rand::Rng>(dim: usize, rng: R) -> Population<R> {
        let settings = Settings::min_max_rng(vec![-20.0; dim], vec![20.0; dim], rng);
        Population::new(settings)
    }

    fn dummy_fitness(individuals: &mut Vec<Individual>) {
        for ind in individuals {
            ind.cost = Some(1.234);
        }
    }

    #[bench]
    fn bench_square_fitness_opt(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            for ind in &mut pop.curr {
                let mut f = 0.0;
                for x in &ind.pos {
                    f += x * x;
                }
                ind.cost = Some(f);
            }
            pop.evolve();
        });
    }

    #[bench]
    fn rand_thread_rng(b: &mut Bencher) {
        let mut pop = setup(5, rand::thread_rng());
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_xor_shift(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac(b: &mut Bencher) {
        let rng: IsaacRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_isaac64(b: &mut Bencher) {
        let rng: Isaac64Rng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_std(b: &mut Bencher) {
        let rng: StdRng = StdRng::new().unwrap();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

    #[bench]
    fn rand_weak_rng(b: &mut Bencher) {
        let rng = weak_rng();
        let mut pop = setup(5, rng);
        b.iter(|| {
            dummy_fitness(&mut pop.curr);
            pop.evolve();
        });
    }

}
<|endoftext|>


<commit_before>//! Simple resampling library in pure Rust.
//!
//! # Examples
//!
//! ```
//! use resize::Pixel::RGB24;
//! use resize::Type::Lanczos3;
//!
//! // Downscale by 2x.
//! let (w1, h1) = (640, 480);
//! let (w2, h2) = (320, 240);
//! // Don't forget to fill `src` with image data (RGB24).
//! let src = vec![0;w1*h1*3];
//! // Destination buffer. Must be mutable.
//! let mut dst = vec![0;w2*h2*3];
//! // Create reusable instance.
//! let mut resizer = resize::new(w1, h1, w2, h2, RGB24, Lanczos3);
//! // Do resize without heap allocations.
//! // Might be executed multiple times for different `src` or `dst`.
//! resizer.resize(&src, &mut dst);
//! ```
// Current implementation is based on:
// * https://github.com/sekrit-twc/zimg/tree/master/src/zimg/resize
// * https://github.com/PistonDevelopers/image/blob/master/src/imageops/sample.rs
#![deny(missing_docs)]

use std::sync::Arc;
use std::collections::HashMap;
use std::f32;

mod px;
pub use px::*;

/// Resizing type to use.
pub enum Type {
    /// Point resizing.
    Point,
    /// Triangle (bilinear) resizing.
    Triangle,
    /// Catmull-Rom (bicubic) resizing.
    Catrom,
    /// Resize using Mitchell-Netravali filter.
    Mitchell,
    /// Resize using Sinc-windowed Sinc with radius of 3.
    Lanczos3,
    /// Resize with custom filter.
    Custom(Filter),
}

/// Resampling filter.
pub struct Filter {
    kernel: Box<dyn Fn(f32) -> f32>,
    support: f32,
}

impl Filter {
    /// Create a new filter.
    ///
    /// # Examples
    ///
    /// ```
    /// use resize::Filter;
    /// fn kernel(x: f32) -> f32 { f32::max(1.0 - x.abs(), 0.0) }
    /// let filter = Filter::new(Box::new(kernel), 1.0);
    /// ```
    #[must_use]
    pub fn new(kernel: Box<dyn Fn(f32) -> f32>, support: f32) -> Self {
        Self { kernel, support }
    }

    /// Helper to create Cubic filter with custom B and C parameters.
    #[must_use]
    pub fn new_cubic(b: f32, c: f32) -> Self {
        Self::new(Box::new(move |x| cubic_bc(b, c, x)), 2.0)
    }

    /// Helper to create Lanczos filter with custom radius.
    #[must_use]
    pub fn new_lanczos(radius: f32) -> Self {
        Self::new(Box::new(move |x| lanczos(radius, x)), radius)
    }
}

#[inline]
fn point_kernel(_: f32) -> f32 {
    1.0
}

#[inline]
fn triangle_kernel(x: f32) -> f32 {
    f32::max(1.0 - x.abs(), 0.0)
}

// Taken from
// https://github.com/PistonDevelopers/image/blob/2921cd7/src/imageops/sample.rs#L68
// TODO(Kagami): Could be optimized for known B and C, see e.g.
// https://github.com/sekrit-twc/zimg/blob/1a606c0/src/zimg/resize/filter.cpp#L149
#[inline]
fn cubic_bc(b: f32, c: f32, x: f32) -> f32 {
    let a = x.abs();
    let k = if a < 1.0 {
        (12.0 - 9.0 * b - 6.0 * c) * a.powi(3) +
        (-18.0 + 12.0 * b + 6.0 * c) * a.powi(2) +
        (6.0 - 2.0 * b)
    } else if a < 2.0 {
        (-b - 6.0 * c) * a.powi(3) +
        (6.0 * b + 30.0 * c) * a.powi(2) +
        (-12.0 * b - 48.0 * c) * a +
        (8.0 * b + 24.0 * c)
    } else {
        0.0
    };
    k / 6.0
}

#[inline]
fn sinc(x: f32) -> f32 {
    if x == 0.0 {
        1.0
    } else {
        let a = x * f32::consts::PI;
        a.sin() / a
    }
}

#[inline]
fn lanczos(taps: f32, x: f32) -> f32 {
    if x.abs() < taps {
        sinc(x) * sinc(x / taps)
    } else {
        0.0
    }
}

/// Supported pixel formats.
// TODO(Kagami): YUV planes?
#[allow(non_snake_case)]
pub mod Pixel {
    /// Grayscale, 8-bit.
    #[derive(Debug, Clone, Copy)]
    pub struct Gray8;
    /// Grayscale, 16-bit, native endian.
    #[derive(Debug, Clone, Copy)]
    pub struct Gray16;
    /// RGB, 8-bit per component.
    #[derive(Debug, Clone, Copy)]
    pub struct RGB24;
    /// RGB, 16-bit per component, native endian.
    #[derive(Debug, Clone, Copy)]
    pub struct RGB48;
    /// RGBA, 8-bit per component.
    #[derive(Debug, Clone, Copy)]
    pub struct RGBA;
    /// RGBA, 16-bit per component, native endian.
    #[derive(Debug, Clone, Copy)]
    pub struct RGBA64;
}

/// Resampler with preallocated buffers and coeffecients for the given
/// dimensions and filter type.
#[derive(Debug)]
pub struct Resizer<Format: PixelFormat> {
    // Source/target dimensions.
    w1: usize,
    h1: usize,
    w2: usize,
    h2: usize,
    pix_fmt: Format,
    // Temporary/preallocated stuff.
    tmp: Vec<f32>,
    coeffs_w: Vec<CoeffsLine>,
    coeffs_h: Vec<CoeffsLine>,
}

#[derive(Debug, Clone)]
struct CoeffsLine {
    start: usize,
    coeffs: Arc<[f32]>,
}

impl<Format: PixelFormat> Resizer<Format> {
    /// Create a new resizer instance.
    pub fn new(source_width: usize, source_heigth: usize, dest_width: usize, dest_height: usize, pixel_format: Format, filter_type: Type) -> Self {
        let filter = match filter_type {
            Type::Point => Filter::new(Box::new(point_kernel), 0.0),
            Type::Triangle => Filter::new(Box::new(triangle_kernel), 1.0),
            Type::Catrom => Filter::new_cubic(0.0, 0.5),
            Type::Mitchell => Filter::new_cubic(1.0/3.0, 1.0/3.0),
            Type::Lanczos3 => Filter::new_lanczos(3.0),
            Type::Custom(f) => f,
        };
        // filters very often create repeating patterns,
        // so overall memory used by them can be reduced
        // which should save some cache space
        let mut recycled_coeffs = HashMap::new();

        let coeffs_w = Self::calc_coeffs(source_width, dest_width, &filter, &mut recycled_coeffs);
        let coeffs_h = if source_heigth == source_width && dest_height == dest_width {
            coeffs_w.clone()
        } else {
            Self::calc_coeffs(source_heigth, dest_height, &filter, &mut recycled_coeffs)
        };
        Self {
            w1: source_width,
            h1: source_heigth,
            w2: dest_width,
            h2: dest_height,
            tmp: Vec::with_capacity(source_width * dest_height * pixel_format.get_ncomponents()),
            pix_fmt: pixel_format,
            coeffs_w,
            coeffs_h,
        }
    }

    fn calc_coeffs(s1: usize, s2: usize, f: &Filter, recycled_coeffs: &mut HashMap<(usize, u64), Arc<[f32]>>) -> Vec<CoeffsLine> {
        let ratio = s1 as f32 / s2 as f32;
        // Scale the filter when downsampling.
        let filter_scale = ratio.max(1.);
        let filter_radius = (f.support * filter_scale).ceil();
        (0..s2).map(|x2| {
            let x1 = (x2 as f32 + 0.5) * ratio - 0.5;
            let start = (x1 - filter_radius).ceil() as isize;
            let start = Self::clamp(start, 0, s1 as isize - 1) as usize;
            let end = (x1 + filter_radius).floor() as isize;
            let end = Self::clamp(end, 0, s1 as isize - 1) as usize;
            let sum: f32 = (start..=end).map(|i| (f.kernel)((i as f32 - x1) / filter_scale)).sum();
            let key = (end - start, ((x1 - start as f32) * 100_000.) as u64);
            let coeffs = recycled_coeffs.entry(key).or_insert_with(|| {
                (start..=end).map(|i| {
                    let v = (f.kernel)((i as f32 - x1) / filter_scale);
                    v / sum
                }).collect::<Arc<[_]>>()
            }).clone();
            CoeffsLine { start, coeffs }
        }).collect()
    }

    #[inline]
    fn clamp<N: PartialOrd>(input: N, min: N, max: N) -> N {
        if input > max {
            max
        } else if input < min {
            min
        } else {
            input
        }
    }

    // Resample W1xH1 to W1xH2.
    // Stride is a length of the source row (>= W1)
    fn sample_rows(&mut self, src: &[Format::Subpixel], stride: usize) {
        let ncomp = self.pix_fmt.get_ncomponents();
        self.tmp.clear();
        assert!(self.tmp.capacity() <= self.w1 * self.h2 * ncomp); // no reallocations
        for x1 in 0..self.w1 {
            let h2 = self.h2;
            let coeffs_h = &self.coeffs_h[0..h2];
            for y2 in 0..h2 {
                let mut accum = Format::new_accum();
                let line = &coeffs_h[y2];
                let src = &src[(line.start * stride + x1) * ncomp..];
                for (i, coeff) in line.coeffs.iter().copied().enumerate() {
                    let base = (i * stride) * ncomp;
                    let src = &src[base..base + ncomp];
                    for (acc, s) in accum.as_mut().iter_mut().zip(src) {
                        *acc += Format::from_subpixel(s) * coeff;
                    }
                }
                for &v in accum.as_ref().iter() {
                    self.tmp.push(v);
                }
            }
        }
    }

    // Resample W1xH2 to W2xH2.
    fn sample_cols(&mut self, dst: &mut [Format::Subpixel]) {
        let ncomp = self.pix_fmt.get_ncomponents();
        let mut offset = 0;
        // Assert that dst is large enough
        let dst = &mut dst[0..self.h2 * self.w2 * ncomp];
        for y2 in 0..self.h2 {
            let w2 = self.w2;
            let coeffs_w = &self.coeffs_w[0..w2];
            for x2 in 0..w2 {
                let mut accum = Format::new_accum();
                let line = &coeffs_w[x2];
                for (i, coeff) in line.coeffs.iter().copied().enumerate() {
                    let x0 = line.start + i;
                    let base = (x0 * self.h2 + y2) * ncomp;
                    let tmp = &self.tmp[base..base + ncomp];
                    for (acc, &p) in accum.as_mut().iter_mut().zip(tmp) {
                        *acc += p * coeff;
                    }
                }
                for &v in accum.as_ref().iter() {
                    dst[offset] = Format::into_subpixel(v);
                    offset += 1;
                }
            }
        }
    }

    /// Resize `src` image data into `dst`.
    pub fn resize(&mut self, src: &[Format::Subpixel], dst: &mut [Format::Subpixel]) {
        let stride = self.w1;
        self.resize_stride(src, stride, dst)
    }

    /// Resize `src` image data into `dst`, skipping `stride` pixels each row.
    pub fn resize_stride(&mut self, src: &[Format::Subpixel], src_stride: usize, dst: &mut [Format::Subpixel]) {
        // TODO(Kagami):
        // * Multi-thread
        // * Bound checkings
        // * SIMD
        assert!(self.w1 <= src_stride);
        assert!(src.len() >= src_stride * self.h1 * self.pix_fmt.get_ncomponents());
        assert_eq!(dst.len(), self.w2 * self.h2 * self.pix_fmt.get_ncomponents());
        self.sample_rows(src, src_stride);
        self.sample_cols(dst)
    }
}

/// Create a new resizer instance. Alias for `Resizer::new`.
pub fn new<Format: PixelFormat>(src_width: usize, src_height: usize, dest_width: usize, dest_height: usize, pixel_format: Format, filter_type: Type) -> Resizer<Format> {
    Resizer::new(src_width, src_height, dest_width, dest_height, pixel_format, filter_type)
}

/// Resize image data to the new dimension in a single step.
///
/// **NOTE:** If you need to resize to the same dimension multiple times,
/// consider creating an resizer instance since it's faster.
#[deprecated(note="Use resize::new().resize()")]
#[allow(deprecated)]
pub fn resize<Format: PixelFormat>(
    src_width: usize, src_height: usize, dest_width: usize, dest_height: usize,
    pixel_format: Format, filter_type: Type,
    src: &[Format::Subpixel], dst: &mut [Format::Subpixel],
) {
    Resizer::new(src_width, src_height, dest_width, dest_height, pixel_format, filter_type).resize(src, dst)
}

#[test]
fn pixel_sizes() {
    assert_eq!(Pixel::RGB24.get_ncomponents(), 3);
    assert_eq!(Pixel::RGB24.get_size(), 3 * 1);
    assert_eq!(Pixel::RGBA.get_size(), 4 * 1);

    assert_eq!(Pixel::RGB48.get_ncomponents(), 3);
    assert_eq!(Pixel::RGB48.get_size(), 3 * 2);
    assert_eq!(Pixel::RGBA64.get_ncomponents(), 4);
    assert_eq!(Pixel::RGBA64.get_size(), 4 * 2);
}

#[test]
fn resize_stride() {
    let mut r = new(2, 2, 3, 4, Pixel::Gray16, Type::Triangle);
    let mut dst = vec![0; 12];
    r.resize_stride(&[
        65535,65535,1,2,
        65535,65535,3,4,
    ], 4, &mut dst);
    assert_eq!(&dst, &[65535; 12]);
}
<commit_msg>Fix coeffs reuse<commit_after>//! Simple resampling library in pure Rust.
//!
//! # Examples
//!
//! ```
//! use resize::Pixel::RGB24;
//! use resize::Type::Lanczos3;
//!
//! // Downscale by 2x.
//! let (w1, h1) = (640, 480);
//! let (w2, h2) = (320, 240);
//! // Don't forget to fill `src` with image data (RGB24).
//! let src = vec![0;w1*h1*3];
//! // Destination buffer. Must be mutable.
//! let mut dst = vec![0;w2*h2*3];
//! // Create reusable instance.
//! let mut resizer = resize::new(w1, h1, w2, h2, RGB24, Lanczos3);
//! // Do resize without heap allocations.
//! // Might be executed multiple times for different `src` or `dst`.
//! resizer.resize(&src, &mut dst);
//! ```
// Current implementation is based on:
// * https://github.com/sekrit-twc/zimg/tree/master/src/zimg/resize
// * https://github.com/PistonDevelopers/image/blob/master/src/imageops/sample.rs
#![deny(missing_docs)]

use std::sync::Arc;
use std::collections::HashMap;
use std::f32;

mod px;
pub use px::*;

/// Resizing type to use.
pub enum Type {
    /// Point resizing.
    Point,
    /// Triangle (bilinear) resizing.
    Triangle,
    /// Catmull-Rom (bicubic) resizing.
    Catrom,
    /// Resize using Mitchell-Netravali filter.
    Mitchell,
    /// Resize using Sinc-windowed Sinc with radius of 3.
    Lanczos3,
    /// Resize with custom filter.
    Custom(Filter),
}

/// Resampling filter.
pub struct Filter {
    kernel: Box<dyn Fn(f32) -> f32>,
    support: f32,
}

impl Filter {
    /// Create a new filter.
    ///
    /// # Examples
    ///
    /// ```
    /// use resize::Filter;
    /// fn kernel(x: f32) -> f32 { f32::max(1.0 - x.abs(), 0.0) }
    /// let filter = Filter::new(Box::new(kernel), 1.0);
    /// ```
    #[must_use]
    pub fn new(kernel: Box<dyn Fn(f32) -> f32>, support: f32) -> Self {
        Self { kernel, support }
    }

    /// Helper to create Cubic filter with custom B and C parameters.
    #[must_use]
    pub fn new_cubic(b: f32, c: f32) -> Self {
        Self::new(Box::new(move |x| cubic_bc(b, c, x)), 2.0)
    }

    /// Helper to create Lanczos filter with custom radius.
    #[must_use]
    pub fn new_lanczos(radius: f32) -> Self {
        Self::new(Box::new(move |x| lanczos(radius, x)), radius)
    }
}

#[inline]
fn point_kernel(_: f32) -> f32 {
    1.0
}

#[inline]
fn triangle_kernel(x: f32) -> f32 {
    f32::max(1.0 - x.abs(), 0.0)
}

// Taken from
// https://github.com/PistonDevelopers/image/blob/2921cd7/src/imageops/sample.rs#L68
// TODO(Kagami): Could be optimized for known B and C, see e.g.
// https://github.com/sekrit-twc/zimg/blob/1a606c0/src/zimg/resize/filter.cpp#L149
#[inline]
fn cubic_bc(b: f32, c: f32, x: f32) -> f32 {
    let a = x.abs();
    let k = if a < 1.0 {
        (12.0 - 9.0 * b - 6.0 * c) * a.powi(3) +
        (-18.0 + 12.0 * b + 6.0 * c) * a.powi(2) +
        (6.0 - 2.0 * b)
    } else if a < 2.0 {
        (-b - 6.0 * c) * a.powi(3) +
        (6.0 * b + 30.0 * c) * a.powi(2) +
        (-12.0 * b - 48.0 * c) * a +
        (8.0 * b + 24.0 * c)
    } else {
        0.0
    };
    k / 6.0
}

#[inline]
fn sinc(x: f32) -> f32 {
    if x == 0.0 {
        1.0
    } else {
        let a = x * f32::consts::PI;
        a.sin() / a
    }
}

#[inline]
fn lanczos(taps: f32, x: f32) -> f32 {
    if x.abs() < taps {
        sinc(x) * sinc(x / taps)
    } else {
        0.0
    }
}

/// Supported pixel formats.
// TODO(Kagami): YUV planes?
#[allow(non_snake_case)]
pub mod Pixel {
    /// Grayscale, 8-bit.
    #[derive(Debug, Clone, Copy)]
    pub struct Gray8;
    /// Grayscale, 16-bit, native endian.
    #[derive(Debug, Clone, Copy)]
    pub struct Gray16;
    /// RGB, 8-bit per component.
    #[derive(Debug, Clone, Copy)]
    pub struct RGB24;
    /// RGB, 16-bit per component, native endian.
    #[derive(Debug, Clone, Copy)]
    pub struct RGB48;
    /// RGBA, 8-bit per component.
    #[derive(Debug, Clone, Copy)]
    pub struct RGBA;
    /// RGBA, 16-bit per component, native endian.
    #[derive(Debug, Clone, Copy)]
    pub struct RGBA64;
}

/// Resampler with preallocated buffers and coeffecients for the given
/// dimensions and filter type.
#[derive(Debug)]
pub struct Resizer<Format: PixelFormat> {
    // Source/target dimensions.
    w1: usize,
    h1: usize,
    w2: usize,
    h2: usize,
    pix_fmt: Format,
    // Temporary/preallocated stuff.
    tmp: Vec<f32>,
    coeffs_w: Vec<CoeffsLine>,
    coeffs_h: Vec<CoeffsLine>,
}

#[derive(Debug, Clone)]
struct CoeffsLine {
    start: usize,
    coeffs: Arc<[f32]>,
}

impl<Format: PixelFormat> Resizer<Format> {
    /// Create a new resizer instance.
    pub fn new(source_width: usize, source_heigth: usize, dest_width: usize, dest_height: usize, pixel_format: Format, filter_type: Type) -> Self {
        let filter = match filter_type {
            Type::Point => Filter::new(Box::new(point_kernel), 0.0),
            Type::Triangle => Filter::new(Box::new(triangle_kernel), 1.0),
            Type::Catrom => Filter::new_cubic(0.0, 0.5),
            Type::Mitchell => Filter::new_cubic(1.0/3.0, 1.0/3.0),
            Type::Lanczos3 => Filter::new_lanczos(3.0),
            Type::Custom(f) => f,
        };
        // filters very often create repeating patterns,
        // so overall memory used by them can be reduced
        // which should save some cache space
        let mut recycled_coeffs = HashMap::new();

        let coeffs_w = Self::calc_coeffs(source_width, dest_width, &filter, &mut recycled_coeffs);
        let coeffs_h = if source_heigth == source_width && dest_height == dest_width {
            coeffs_w.clone()
        } else {
            Self::calc_coeffs(source_heigth, dest_height, &filter, &mut recycled_coeffs)
        };
        Self {
            w1: source_width,
            h1: source_heigth,
            w2: dest_width,
            h2: dest_height,
            tmp: Vec::with_capacity(source_width * dest_height * pixel_format.get_ncomponents()),
            pix_fmt: pixel_format,
            coeffs_w,
            coeffs_h,
        }
    }

    fn calc_coeffs(s1: usize, s2: usize, f: &Filter, recycled_coeffs: &mut HashMap<(usize, [u8; 4], [u8; 4]), Arc<[f32]>>) -> Vec<CoeffsLine> {
        let ratio = s1 as f32 / s2 as f32;
        // Scale the filter when downsampling.
        let filter_scale = ratio.max(1.);
        let filter_radius = (f.support * filter_scale).ceil();
        (0..s2).map(|x2| {
            let x1 = (x2 as f32 + 0.5) * ratio - 0.5;
            let start = (x1 - filter_radius).ceil() as isize;
            let start = Self::clamp(start, 0, s1 as isize - 1) as usize;
            let end = (x1 + filter_radius).floor() as isize;
            let end = Self::clamp(end, 0, s1 as isize - 1) as usize;
            let sum: f32 = (start..=end).map(|i| (f.kernel)((i as f32 - x1) / filter_scale)).sum();
            let key = (end - start, filter_scale.to_ne_bytes(), (x1 - start as f32).to_ne_bytes());
            let coeffs = recycled_coeffs.entry(key).or_insert_with(|| {
                (start..=end).map(|i| {
                    let v = (f.kernel)((i as f32 - x1) / filter_scale);
                    v / sum
                }).collect::<Arc<[_]>>()
            }).clone();
            CoeffsLine { start, coeffs }
        }).collect()
    }

    #[inline]
    fn clamp<N: PartialOrd>(input: N, min: N, max: N) -> N {
        if input > max {
            max
        } else if input < min {
            min
        } else {
            input
        }
    }

    // Resample W1xH1 to W1xH2.
    // Stride is a length of the source row (>= W1)
    fn sample_rows(&mut self, src: &[Format::Subpixel], stride: usize) {
        let ncomp = self.pix_fmt.get_ncomponents();
        self.tmp.clear();
        assert!(self.tmp.capacity() <= self.w1 * self.h2 * ncomp); // no reallocations
        for x1 in 0..self.w1 {
            let h2 = self.h2;
            let coeffs_h = &self.coeffs_h[0..h2];
            for y2 in 0..h2 {
                let mut accum = Format::new_accum();
                let line = &coeffs_h[y2];
                let src = &src[(line.start * stride + x1) * ncomp..];
                for (i, coeff) in line.coeffs.iter().copied().enumerate() {
                    let base = (i * stride) * ncomp;
                    let src = &src[base..base + ncomp];
                    for (acc, s) in accum.as_mut().iter_mut().zip(src) {
                        *acc += Format::from_subpixel(s) * coeff;
                    }
                }
                for &v in accum.as_ref().iter() {
                    self.tmp.push(v);
                }
            }
        }
    }

    // Resample W1xH2 to W2xH2.
    fn sample_cols(&mut self, dst: &mut [Format::Subpixel]) {
        let ncomp = self.pix_fmt.get_ncomponents();
        let mut offset = 0;
        // Assert that dst is large enough
        let dst = &mut dst[0..self.h2 * self.w2 * ncomp];
        for y2 in 0..self.h2 {
            let w2 = self.w2;
            let coeffs_w = &self.coeffs_w[0..w2];
            for x2 in 0..w2 {
                let mut accum = Format::new_accum();
                let line = &coeffs_w[x2];
                for (i, coeff) in line.coeffs.iter().copied().enumerate() {
                    let x0 = line.start + i;
                    let base = (x0 * self.h2 + y2) * ncomp;
                    let tmp = &self.tmp[base..base + ncomp];
                    for (acc, &p) in accum.as_mut().iter_mut().zip(tmp) {
                        *acc += p * coeff;
                    }
                }
                for &v in accum.as_ref().iter() {
                    dst[offset] = Format::into_subpixel(v);
                    offset += 1;
                }
            }
        }
    }

    /// Resize `src` image data into `dst`.
    pub fn resize(&mut self, src: &[Format::Subpixel], dst: &mut [Format::Subpixel]) {
        let stride = self.w1;
        self.resize_stride(src, stride, dst)
    }

    /// Resize `src` image data into `dst`, skipping `stride` pixels each row.
    pub fn resize_stride(&mut self, src: &[Format::Subpixel], src_stride: usize, dst: &mut [Format::Subpixel]) {
        // TODO(Kagami):
        // * Multi-thread
        // * Bound checkings
        // * SIMD
        assert!(self.w1 <= src_stride);
        assert!(src.len() >= src_stride * self.h1 * self.pix_fmt.get_ncomponents());
        assert_eq!(dst.len(), self.w2 * self.h2 * self.pix_fmt.get_ncomponents());
        self.sample_rows(src, src_stride);
        self.sample_cols(dst)
    }
}

/// Create a new resizer instance. Alias for `Resizer::new`.
pub fn new<Format: PixelFormat>(src_width: usize, src_height: usize, dest_width: usize, dest_height: usize, pixel_format: Format, filter_type: Type) -> Resizer<Format> {
    Resizer::new(src_width, src_height, dest_width, dest_height, pixel_format, filter_type)
}

/// Resize image data to the new dimension in a single step.
///
/// **NOTE:** If you need to resize to the same dimension multiple times,
/// consider creating an resizer instance since it's faster.
#[deprecated(note="Use resize::new().resize()")]
#[allow(deprecated)]
pub fn resize<Format: PixelFormat>(
    src_width: usize, src_height: usize, dest_width: usize, dest_height: usize,
    pixel_format: Format, filter_type: Type,
    src: &[Format::Subpixel], dst: &mut [Format::Subpixel],
) {
    Resizer::new(src_width, src_height, dest_width, dest_height, pixel_format, filter_type).resize(src, dst)
}

#[test]
fn pixel_sizes() {
    assert_eq!(Pixel::RGB24.get_ncomponents(), 3);
    assert_eq!(Pixel::RGB24.get_size(), 3 * 1);
    assert_eq!(Pixel::RGBA.get_size(), 4 * 1);

    assert_eq!(Pixel::RGB48.get_ncomponents(), 3);
    assert_eq!(Pixel::RGB48.get_size(), 3 * 2);
    assert_eq!(Pixel::RGBA64.get_ncomponents(), 4);
    assert_eq!(Pixel::RGBA64.get_size(), 4 * 2);
}

#[test]
fn resize_stride() {
    let mut r = new(2, 2, 3, 4, Pixel::Gray16, Type::Triangle);
    let mut dst = vec![0; 12];
    r.resize_stride(&[
        65535,65535,1,2,
        65535,65535,3,4,
    ], 4, &mut dst);
    assert_eq!(&dst, &[65535; 12]);
}
<|endoftext|>


<commit_before>extern crate futures;

use futures::{Async, Poll, Future, Stream};

pub enum StreamEvent<I, S> {
    Next(I),
    Done(S),
}

pub trait StateStream {
    type Item;
    type State;
    type Error;

    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error>;

    #[inline]
    fn into_future(self) -> IntoFuture<Self>
        where Self: Sized
    {
        IntoFuture(Some(self))
    }
}

pub fn stream<S>(stream: S) -> FromStream<S>
    where S: Stream
{
    FromStream(stream)
}

pub struct FromStream<S>(S);

impl<S> StateStream for FromStream<S>
    where S: Stream
{
    type Item = S::Item;
    type State = ();
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, ()>, S::Error> {
        self.0.poll().map(|a| {
            match a {
                Async::Ready(Some(i)) => Async::Ready(StreamEvent::Next(i)),
                Async::Ready(None) => Async::Ready(StreamEvent::Done(())),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct IntoFuture<S>(Option<S>);

impl<S> Future for IntoFuture<S>
    where S: StateStream
{
    type Item = (StreamEvent<S::Item, S::State>, S);
    type Error = (S::Error, S);

    #[inline]
    fn poll(&mut self) -> Poll<(StreamEvent<S::Item, S::State>, S), (S::Error, S)> {
        let item = match self.0.as_mut().expect("polling IntoFuture twice").poll() {
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Ok(Async::Ready(i)) => Ok(i),
            Err(e) => Err(e),
        };
        let stream = self.0.take().unwrap();
        match item {
            Ok(i) => Ok(Async::Ready((i, stream))),
            Err(e) => Err((e, stream)),
        }
    }
}<commit_msg>Add into_stream<commit_after>extern crate futures;

use futures::{Async, Poll, Future, Stream};

pub enum StreamEvent<I, S> {
    Next(I),
    Done(S),
}

pub trait StateStream {
    type Item;
    type State;
    type Error;

    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error>;

    #[inline]
    fn into_future(self) -> IntoFuture<Self>
        where Self: Sized
    {
        IntoFuture(Some(self))
    }

    #[inline]
    fn into_stream(self) -> IntoStream<Self>
        where Self: Sized
    {
        IntoStream(self)
    }
}

pub fn stream<S>(stream: S) -> FromStream<S>
    where S: Stream
{
    FromStream(stream)
}

pub struct FromStream<S>(S);

impl<S> StateStream for FromStream<S>
    where S: Stream
{
    type Item = S::Item;
    type State = ();
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, ()>, S::Error> {
        self.0.poll().map(|a| {
            match a {
                Async::Ready(Some(i)) => Async::Ready(StreamEvent::Next(i)),
                Async::Ready(None) => Async::Ready(StreamEvent::Done(())),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct IntoFuture<S>(Option<S>);

impl<S> Future for IntoFuture<S>
    where S: StateStream
{
    type Item = (StreamEvent<S::Item, S::State>, S);
    type Error = (S::Error, S);

    #[inline]
    fn poll(&mut self) -> Poll<(StreamEvent<S::Item, S::State>, S), (S::Error, S)> {
        let item = match self.0.as_mut().expect("polling IntoFuture twice").poll() {
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Ok(Async::Ready(i)) => Ok(i),
            Err(e) => Err(e),
        };
        let stream = self.0.take().unwrap();
        match item {
            Ok(i) => Ok(Async::Ready((i, stream))),
            Err(e) => Err((e, stream)),
        }
    }
}

pub struct IntoStream<S>(S);

impl<S> Stream for IntoStream<S>
    where S: StateStream
{
    type Item = S::Item;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
        self.0.poll().map(|a| {
            match a {
                Async::Ready(StreamEvent::Next(i)) => Async::Ready(Some(i)),
                Async::Ready(StreamEvent::Done(_)) => Async::Ready(None),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}<|endoftext|>