<commit_before>use std::f64;
use rand::Rng;
use error::StatsError;
use function;
use result::Result;
use super::{Distribution, Univariate, Continuous};
use super::normal;

#[derive(Debug, Copy, Clone, PartialEq)]
pub struct StudentT {
    location: f64,
    scale: f64,
    freedom: f64,
    n: normal::Normal,
}

impl StudentT {
    pub fn new(location: f64, scale: f64, freedom: f64) -> Result<StudentT> {
        let is_nan = location.is_nan() || scale.is_nan() || freedom.is_nan();
        if is_nan || scale <= 0.0 || freedom <= 0.0 {
            Err(StatsError::BadParams)
        } else {
            Ok(StudentT {
                location: location,
                scale: scale,
                freedom: freedom,
                n: normal::Normal::new(location, scale).unwrap(),
            })
        }
    }

    pub fn location(&self) -> f64 {
        self.location
    }

    pub fn scale(&self) -> f64 {
        self.scale
    }

    pub fn freedom(&self) -> f64 {
        self.freedom
    }
}

impl Distribution for StudentT {
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let gamma = super::gamma::sample_unchecked(r, 0.5 * self.freedom, 0.5);
        normal::sample_unchecked(r, self.location, self.scale * (self.freedom / gamma).sqrt())
    }
}

impl Univariate for StudentT {
    fn mean(&self) -> f64 {
        if self.freedom <= 1.0 {
            panic!("Cannot calculate mean for StudentT distribution with freedom <= 1.0");
        }
        self.location
    }

    fn variance(&self) -> f64 {
        if self.freedom <= 1.0 {
            panic!("Cannot calculate variance for StudentT distribution with freedom <= 1.0");
        }

        if self.freedom == f64::INFINITY {
            self.scale * self.scale
        } else if self.freedom > 2.0 {
            self.freedom * self.scale * self.scale / (self.freedom - 2.0)
        } else {
            f64::INFINITY
        }
    }

    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }

    fn entropy(&self) -> f64 {
        if self.location != 0.0 || self.scale != 1.0 {
            panic!("Cannot calculate entropy for StudentT distribution where location is not 0 \
                    and scale is not 1");
        }

        (self.freedom + 1.0) / 2.0 *
        (function::gamma::digamma((self.freedom + 1.0) / 2.0).unwrap() -
         function::gamma::digamma(self.freedom / 2.0).unwrap()) +
        (self.freedom.sqrt() * function::beta::beta(self.freedom / 2.0, 0.5).unwrap()).ln()
    }

    fn skewness(&self) -> f64 {
        if self.freedom <= 3.0 {
            panic!("Cannot calculate skewness for StudentT distribution where freedom <= 3");
        }
        0.0
    }

    fn median(&self) -> f64 {
        self.location
    }

    fn cdf(&self, x: f64) -> f64 {
        if self.freedom == f64::INFINITY {
            self.n.cdf(x)
        } else {
            let k = (x - self.location) / self.scale;
            let h = self.freedom / (self.freedom + k * k);
            let ib = 0.5 * function::beta::beta_reg(self.freedom / 2.0, 0.5, h).unwrap();
            if x <= self.location {
                ib
            } else {
                1.0 - ib
            }
        }
    }
}

impl Continuous for StudentT {
    fn mode(&self) -> f64 {
        self.location
    }

    fn min(&self) -> f64 {
        f64::NEG_INFINITY
    }

    fn max(&self) -> f64 {
        f64::INFINITY
    }

    fn pdf(&self, x: f64) -> f64 {
        if self.freedom >= 1e8 {
            self.n.pdf(x)
        } else {
            let d = (x - self.location) / self.scale;
            (function::gamma::ln_gamma((self.freedom + 1.0) / 2.0) -
             function::gamma::ln_gamma(self.freedom / 2.0))
                .exp() *
            (1.0 + d * d / self.freedom).powf(-0.5 * (self.freedom + 1.0)) /
            (self.freedom * f64::consts::PI).sqrt() / self.scale
        }
    }

    fn ln_pdf(&self, x: f64) -> f64 {
        if self.freedom >= 1e8 {
            self.n.ln_pdf(x)
        } else {
            let d = (x - self.location) / self.scale;
            function::gamma::ln_gamma((self.freedom + 1.0) / 2.0) -
            0.5 * ((self.freedom + 1.0) * (1.0 + d * d / self.freedom).ln()) -
            function::gamma::ln_gamma(self.freedom / 2.0) -
            0.5 * (self.freedom * f64::consts::PI).ln() - self.scale.ln()
        }
    }
}
<commit_msg>remove full path import for student t functions since gamma uses full path import<commit_after>use std::f64;
use rand::Rng;
use error::StatsError;
use function::{beta, gamma};
use result::Result;
use super::{Distribution, Univariate, Continuous};
use super::normal;

#[derive(Debug, Copy, Clone, PartialEq)]
pub struct StudentT {
    location: f64,
    scale: f64,
    freedom: f64,
    n: normal::Normal,
}

impl StudentT {
    pub fn new(location: f64, scale: f64, freedom: f64) -> Result<StudentT> {
        let is_nan = location.is_nan() || scale.is_nan() || freedom.is_nan();
        if is_nan || scale <= 0.0 || freedom <= 0.0 {
            Err(StatsError::BadParams)
        } else {
            Ok(StudentT {
                location: location,
                scale: scale,
                freedom: freedom,
                n: normal::Normal::new(location, scale).unwrap(),
            })
        }
    }

    pub fn location(&self) -> f64 {
        self.location
    }

    pub fn scale(&self) -> f64 {
        self.scale
    }

    pub fn freedom(&self) -> f64 {
        self.freedom
    }
}

impl Distribution for StudentT {
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let gamma = super::gamma::sample_unchecked(r, 0.5 * self.freedom, 0.5);
        normal::sample_unchecked(r, self.location, self.scale * (self.freedom / gamma).sqrt())
    }
}

impl Univariate for StudentT {
    fn mean(&self) -> f64 {
        if self.freedom <= 1.0 {
            panic!("Cannot calculate mean for StudentT distribution with freedom <= 1.0");
        }
        self.location
    }

    fn variance(&self) -> f64 {
        if self.freedom <= 1.0 {
            panic!("Cannot calculate variance for StudentT distribution with freedom <= 1.0");
        }

        if self.freedom == f64::INFINITY {
            self.scale * self.scale
        } else if self.freedom > 2.0 {
            self.freedom * self.scale * self.scale / (self.freedom - 2.0)
        } else {
            f64::INFINITY
        }
    }

    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }

    fn entropy(&self) -> f64 {
        if self.location != 0.0 || self.scale != 1.0 {
            panic!("Cannot calculate entropy for StudentT distribution where location is not 0 \
                    and scale is not 1");
        }

        (self.freedom + 1.0) / 2.0 *
        (gamma::digamma((self.freedom + 1.0) / 2.0).unwrap() -
         gamma::digamma(self.freedom / 2.0).unwrap()) +
        (self.freedom.sqrt() * beta::beta(self.freedom / 2.0, 0.5).unwrap()).ln()
    }

    fn skewness(&self) -> f64 {
        if self.freedom <= 3.0 {
            panic!("Cannot calculate skewness for StudentT distribution where freedom <= 3");
        }
        0.0
    }

    fn median(&self) -> f64 {
        self.location
    }

    fn cdf(&self, x: f64) -> f64 {
        if self.freedom == f64::INFINITY {
            self.n.cdf(x)
        } else {
            let k = (x - self.location) / self.scale;
            let h = self.freedom / (self.freedom + k * k);
            let ib = 0.5 * beta::beta_reg(self.freedom / 2.0, 0.5, h).unwrap();
            if x <= self.location {
                ib
            } else {
                1.0 - ib
            }
        }
    }
}

impl Continuous for StudentT {
    fn mode(&self) -> f64 {
        self.location
    }

    fn min(&self) -> f64 {
        f64::NEG_INFINITY
    }

    fn max(&self) -> f64 {
        f64::INFINITY
    }

    fn pdf(&self, x: f64) -> f64 {
        if self.freedom >= 1e8 {
            self.n.pdf(x)
        } else {
            let d = (x - self.location) / self.scale;
            (gamma::ln_gamma((self.freedom + 1.0) / 2.0) - gamma::ln_gamma(self.freedom / 2.0))
                .exp() *
            (1.0 + d * d / self.freedom).powf(-0.5 * (self.freedom + 1.0)) /
            (self.freedom * f64::consts::PI).sqrt() / self.scale
        }
    }

    fn ln_pdf(&self, x: f64) -> f64 {
        if self.freedom >= 1e8 {
            self.n.ln_pdf(x)
        } else {
            let d = (x - self.location) / self.scale;
            gamma::ln_gamma((self.freedom + 1.0) / 2.0) -
            0.5 * ((self.freedom + 1.0) * (1.0 + d * d / self.freedom).ln()) -
            gamma::ln_gamma(self.freedom / 2.0) -
            0.5 * (self.freedom * f64::consts::PI).ln() - self.scale.ln()
        }
    }
}
<|endoftext|>


<commit_before>//! Provides BLAS for a Native backend.

use ::operation::*;
use ::plugin::*;
use ::transpose::*;
use collenchyma::backend::Backend;
use collenchyma::memory::MemoryType;
use collenchyma::frameworks::native::Native;
use collenchyma::plugin::Error;
use rblas::math::mat::Mat;
use rblas::matrix::Matrix;
use rblas::{Asum, Axpy, Copy, Dot, Nrm2, Scal, Swap}; // Level 1
use rblas::Gemm; // Level 3

macro_rules! impl_asum_for {
    ($t:ident, $b:ty) => (
        impl IOperationAsum<$t> for $b {
            fn compute(&self, x: &MemoryType, result: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let mut r_slice = try!(result.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `result`."))).as_mut_slice::<$t>();
                r_slice[0] = Asum::asum(x_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_axpy_for {
    ($t:ident, $b:ty) => (
        impl IOperationAxpy<$t> for $b {
            fn compute(&self, a: &MemoryType, x: &MemoryType, y: &mut MemoryType) -> Result<(), Error> {
                let a_slice = try!(a.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `a`."))).as_slice::<$t>();
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let y_slice = try!(y.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_mut_slice::<$t>();
                Axpy::axpy(&a_slice[0], x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_copy_for {
    ($t:ident, $b:ty) => (
        impl IOperationCopy<$t> for $b {
            fn compute(&self, x: &MemoryType, y: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let y_slice = try!(y.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_mut_slice::<$t>();
                Copy::copy(x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_dot_for {
    ($t:ident, $b:ty) => (
        impl IOperationDot<$t> for $b {
            fn compute(&self, x: &MemoryType, y: &MemoryType, result: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let y_slice = try!(y.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_slice::<$t>();
                let mut r_slice = try!(result.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `result`."))).as_mut_slice::<$t>();
                r_slice[0] = Dot::dot(x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_nrm2_for {
    ($t:ident, $b:ty) => (
        impl IOperationNrm2<$t> for $b {
            fn compute(&self, x: &MemoryType, result: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let mut r_slice = try!(result.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `result`."))).as_mut_slice::<$t>();
                r_slice[0] = Nrm2::nrm2(x_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_scale_for {
    ($t:ident, $b:ty) => (
        impl IOperationScale<$t> for $b {
            fn compute(&self, a: &MemoryType, x: &mut MemoryType) -> Result<(), Error> {
                let a_slice = try!(a.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `a`."))).as_slice::<$t>();
                let mut x_slice = try!(x.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_mut_slice::<$t>();
                Scal::scal(&a_slice[0], x_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_swap_for {
    ($t:ident, $b:ty) => (
        impl IOperationSwap<$t> for $b {
            fn compute(&self, x: &mut MemoryType, y: &mut MemoryType) -> Result<(), Error> {
                let mut x_slice = try!(x.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_mut_slice::<$t>();
                let mut y_slice = try!(y.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_mut_slice::<$t>();
                Swap::swap(x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_gemm_for {
    ($t:ident, $b:ty) => (
        impl IOperationGemm<$t> for $b {
            fn compute(&self, alpha: &MemoryType, at: Transpose, a_dims: &[usize], a: &MemoryType, bt: Transpose, b_dims: &[usize], b: &MemoryType, beta: &MemoryType, c_dims: &[usize], c: &mut MemoryType) -> Result<(), ::collenchyma::error::Error> {
                let alpha_slice = try!(alpha.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `alpha`."))).as_slice::<$t>();
                let a_slice = try!(a.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `a`."))).as_slice::<$t>();
                let beta_slice = try!(beta.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `beta`."))).as_slice::<$t>();
                let b_slice = try!(b.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `b`."))).as_slice::<$t>();
                let mut c_slice = try!(c.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `c`."))).as_mut_slice::<$t>();

                let a_matrix = as_matrix(a_slice, a_dims);
                let b_matrix = as_matrix(b_slice, b_dims);
                let mut c_matrix = as_matrix(c_slice, c_dims);
                Gemm::gemm(&alpha_slice[0], at.to_rblas(), &a_matrix, bt.to_rblas(), &b_matrix, &beta_slice[0], &mut c_matrix);
                read_from_matrix(&c_matrix, c_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_iblas_for {
    ($t:ident, $b:ty) => (
        impl_asum_for!($t, $b);
        impl_axpy_for!($t, $b);
        impl_copy_for!($t, $b);
        impl_dot_for!($t, $b);
        impl_nrm2_for!($t, $b);
        impl_scale_for!($t, $b);
        impl_swap_for!($t, $b);

        impl_gemm_for!($t, $b);

        impl IBlas<$t> for $b {
            iblas_asum_for!($t, $b);
            iblas_axpy_for!($t, $b);
            iblas_copy_for!($t, $b);
            iblas_dot_for!($t, $b);
            iblas_nrm2_for!($t, $b);
            iblas_scale_for!($t, $b);
            iblas_swap_for!($t, $b);

            iblas_gemm_for!($t, $b);
        }
    );
}

impl_iblas_for!(f32, Backend<Native>);
impl_iblas_for!(f64, Backend<Native>);

/// Create a rblas-Matrix from a slice and dimensions.
fn as_matrix<T: Clone + ::std::fmt::Debug>(slice: &[T], dims: &[usize]) -> Mat<T> {
    let n = dims[0];
    let m = dims[1];
    let mut mat = Mat::new(n, m);
    for i in 0..n {
        for j in 0..m {
            unsafe {
                *mat.as_mut_ptr().offset((n*i + j) as isize) = slice[n*i + j].clone();
            }
        }
    }

    mat
}

fn read_from_matrix<T: Clone>(mat: &Mat<T>, slice: &mut [T]) {
    let n = mat.rows();
    let m = mat.cols();
    for i in 0..n {
        for j in 0..m {
            slice[n*i + j] = mat[i][j].clone();
        }
    }
}
<commit_msg>fix/as_matrix: fix filling index in as_matrix<commit_after>//! Provides BLAS for a Native backend.

use ::operation::*;
use ::plugin::*;
use ::transpose::*;
use collenchyma::backend::Backend;
use collenchyma::memory::MemoryType;
use collenchyma::frameworks::native::Native;
use collenchyma::plugin::Error;
use rblas::math::mat::Mat;
use rblas::matrix::Matrix;
use rblas::{Asum, Axpy, Copy, Dot, Nrm2, Scal, Swap}; // Level 1
use rblas::Gemm; // Level 3

macro_rules! impl_asum_for {
    ($t:ident, $b:ty) => (
        impl IOperationAsum<$t> for $b {
            fn compute(&self, x: &MemoryType, result: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let mut r_slice = try!(result.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `result`."))).as_mut_slice::<$t>();
                r_slice[0] = Asum::asum(x_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_axpy_for {
    ($t:ident, $b:ty) => (
        impl IOperationAxpy<$t> for $b {
            fn compute(&self, a: &MemoryType, x: &MemoryType, y: &mut MemoryType) -> Result<(), Error> {
                let a_slice = try!(a.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `a`."))).as_slice::<$t>();
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let y_slice = try!(y.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_mut_slice::<$t>();
                Axpy::axpy(&a_slice[0], x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_copy_for {
    ($t:ident, $b:ty) => (
        impl IOperationCopy<$t> for $b {
            fn compute(&self, x: &MemoryType, y: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let y_slice = try!(y.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_mut_slice::<$t>();
                Copy::copy(x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_dot_for {
    ($t:ident, $b:ty) => (
        impl IOperationDot<$t> for $b {
            fn compute(&self, x: &MemoryType, y: &MemoryType, result: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let y_slice = try!(y.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_slice::<$t>();
                let mut r_slice = try!(result.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `result`."))).as_mut_slice::<$t>();
                r_slice[0] = Dot::dot(x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_nrm2_for {
    ($t:ident, $b:ty) => (
        impl IOperationNrm2<$t> for $b {
            fn compute(&self, x: &MemoryType, result: &mut MemoryType) -> Result<(), Error> {
                let x_slice = try!(x.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_slice::<$t>();
                let mut r_slice = try!(result.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `result`."))).as_mut_slice::<$t>();
                r_slice[0] = Nrm2::nrm2(x_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_scale_for {
    ($t:ident, $b:ty) => (
        impl IOperationScale<$t> for $b {
            fn compute(&self, a: &MemoryType, x: &mut MemoryType) -> Result<(), Error> {
                let a_slice = try!(a.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `a`."))).as_slice::<$t>();
                let mut x_slice = try!(x.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_mut_slice::<$t>();
                Scal::scal(&a_slice[0], x_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_swap_for {
    ($t:ident, $b:ty) => (
        impl IOperationSwap<$t> for $b {
            fn compute(&self, x: &mut MemoryType, y: &mut MemoryType) -> Result<(), Error> {
                let mut x_slice = try!(x.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `x`."))).as_mut_slice::<$t>();
                let mut y_slice = try!(y.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `y`."))).as_mut_slice::<$t>();
                Swap::swap(x_slice, y_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_gemm_for {
    ($t:ident, $b:ty) => (
        impl IOperationGemm<$t> for $b {
            fn compute(&self, alpha: &MemoryType, at: Transpose, a_dims: &[usize], a: &MemoryType, bt: Transpose, b_dims: &[usize], b: &MemoryType, beta: &MemoryType, c_dims: &[usize], c: &mut MemoryType) -> Result<(), ::collenchyma::error::Error> {
                let alpha_slice = try!(alpha.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `alpha`."))).as_slice::<$t>();
                let a_slice = try!(a.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `a`."))).as_slice::<$t>();
                let beta_slice = try!(beta.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `beta`."))).as_slice::<$t>();
                let b_slice = try!(b.as_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `b`."))).as_slice::<$t>();
                let mut c_slice = try!(c.as_mut_native().ok_or(Error::MissingMemoryForDevice("Unable to receive native memory for `c`."))).as_mut_slice::<$t>();

                let a_matrix = as_matrix(a_slice, a_dims);
                let b_matrix = as_matrix(b_slice, b_dims);
                let mut c_matrix = as_matrix(c_slice, c_dims);
                Gemm::gemm(&alpha_slice[0], at.to_rblas(), &a_matrix, bt.to_rblas(), &b_matrix, &beta_slice[0], &mut c_matrix);
                read_from_matrix(&c_matrix, c_slice);
                Ok(())
            }
        }
    );
}

macro_rules! impl_iblas_for {
    ($t:ident, $b:ty) => (
        impl_asum_for!($t, $b);
        impl_axpy_for!($t, $b);
        impl_copy_for!($t, $b);
        impl_dot_for!($t, $b);
        impl_nrm2_for!($t, $b);
        impl_scale_for!($t, $b);
        impl_swap_for!($t, $b);

        impl_gemm_for!($t, $b);

        impl IBlas<$t> for $b {
            iblas_asum_for!($t, $b);
            iblas_axpy_for!($t, $b);
            iblas_copy_for!($t, $b);
            iblas_dot_for!($t, $b);
            iblas_nrm2_for!($t, $b);
            iblas_scale_for!($t, $b);
            iblas_swap_for!($t, $b);

            iblas_gemm_for!($t, $b);
        }
    );
}

impl_iblas_for!(f32, Backend<Native>);
impl_iblas_for!(f64, Backend<Native>);

/// Create a rblas-Matrix from a slice and dimensions.
fn as_matrix<T: Clone + ::std::fmt::Debug>(slice: &[T], dims: &[usize]) -> Mat<T> {
    let n = dims[0];
    let m = dims.iter().skip(1).fold(1, |prod, i| prod * i);
    let mut mat = Mat::new(n, m);
    for i in 0..n {
        for j in 0..m {
            let index = m * i + j;
            unsafe {
                *mat.as_mut_ptr().offset(index as isize) = slice[index].clone();
            }
        }
    }

    mat
}

fn read_from_matrix<T: Clone>(mat: &Mat<T>, slice: &mut [T]) {
    let n = mat.rows();
    let m = mat.cols();
    for i in 0..n {
        for j in 0..m {
            slice[n*i + j] = mat[i][j].clone();
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The xi-editor Authors.
//
// 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.

#![allow(
    clippy::if_same_then_else,
    clippy::needless_bool,
    clippy::needless_pass_by_value,
    clippy::ptr_arg
)]

#[cfg(all(test, feature = "benchmarks"))]
extern crate test;

use super::Sample;
use serde_json;
use std::io::{Error as IOError, ErrorKind as IOErrorKind, Read, Write};

#[derive(Debug)]
pub enum Error {
    Io(IOError),
    Json(serde_json::Error),
    DecodingFormat(String),
}

impl From<IOError> for Error {
    fn from(e: IOError) -> Error {
        Error::Io(e)
    }
}

impl From<serde_json::Error> for Error {
    fn from(e: serde_json::Error) -> Error {
        Error::Json(e)
    }
}

impl From<String> for Error {
    fn from(e: String) -> Error {
        Error::DecodingFormat(e)
    }
}

impl Error {
    pub fn already_exists() -> Error {
        Error::Io(IOError::from(IOErrorKind::AlreadyExists))
    }
}

#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
enum ChromeTraceArrayEntries {
    Array(Vec<Sample>),
}

/// This serializes the samples into the [Chrome trace event format](https://www.google.com/url?sa=t&rct=j&q=&esrc=s&source=web&cd=1&ved=0ahUKEwiJlZmDguXYAhUD4GMKHVmEDqIQFggpMAA&url=https%3A%2F%2Fdocs.google.com%2Fdocument%2Fd%2F1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU%2Fpreview&usg=AOvVaw0tBFlVbDVBikdzLqgrWK3g).
///
/// # Arguments
/// `samples` - Something that can be converted into an iterator of sample
/// references.
/// `format` - Which trace format to save the data in.  There are four total
/// formats described in the document.
/// `output` - Where to write the serialized result.
///
/// # Returns
/// A `Result<(), Error>` that indicates if serialization was successful or the
/// details of any error that occured.
///
/// # Examples
/// ```norun
/// let samples = xi_trace::samples_cloned_sorted();
/// let mut serialized = Vec::<u8>::new();
/// serialize(samples.iter(), serialized);
/// ```
pub fn serialize<W>(samples: &Vec<Sample>, output: W) -> Result<(), Error>
where
    W: Write,
{
    serde_json::to_writer(output, samples).map_err(Error::Json)
}

pub fn to_value(samples: &Vec<Sample>) -> Result<serde_json::Value, Error> {
    serde_json::to_value(samples).map_err(Error::Json)
}

pub fn decode(samples: serde_json::Value) -> Result<Vec<Sample>, Error> {
    serde_json::from_value(samples).map_err(Error::Json)
}

pub fn deserialize<R>(input: R) -> Result<Vec<Sample>, Error>
where
    R: Read,
{
    serde_json::from_reader(input).map_err(Error::Json)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[cfg(feature = "benchmarks")]
    use test::Bencher;

    #[cfg(not(feature = "json_payload"))]
    fn to_payload(value: &'static str) -> &'static str {
        value
    }

    #[cfg(feature = "json_payload")]
    fn to_payload(value: &'static str) -> TracePayloadT {
        json!({ "test": value })
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_serialization() {
        use super::super::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload(
            "sample4",
            &["test", "chrome"],
            || {
                let _guard = trace.block("sample5", &["test,chrome"]);
            },
            to_payload("payload 4"),
        );

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();

        let result = serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let decoded_result: Vec<serde_json::Value> = serde_json::from_slice(&serialized).unwrap();
        assert_eq!(decoded_result.len(), 8);
        assert_eq!(decoded_result[0]["name"].as_str().unwrap(), "process_name");
        assert_eq!(decoded_result[1]["name"].as_str().unwrap(), "thread_name");
        for i in 2..5 {
            assert_eq!(decoded_result[i]["name"].as_str().unwrap(), samples[i].name);
            assert_eq!(decoded_result[i]["cat"].as_str().unwrap(), "test,chrome");
            assert_eq!(decoded_result[i]["ph"].as_str().unwrap(), "i");
            assert_eq!(decoded_result[i]["ts"], samples[i].timestamp_us);
            let nth_sample = &samples[i];
            let nth_args = nth_sample.args.as_ref().unwrap();
            assert_eq!(decoded_result[i]["args"]["xi_payload"], json!(nth_args.payload.as_ref()));
        }
        assert_eq!(decoded_result[5]["ph"], "B");
        assert_eq!(decoded_result[6]["ph"], "E");
        assert_eq!(decoded_result[7]["ph"], "X");
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_deserialization() {
        use super::super::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload("sample4", &["test", "chrome"], || (), to_payload("payload 4"));

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();
        let result = serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let deserialized_samples = deserialize(serialized.as_slice()).unwrap();
        assert_eq!(deserialized_samples, samples);
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_one_element(b: &mut Bencher) {
        use super::*;

        let mut serialized = Vec::<u8>::new();
        let samples = vec![super::Sample::new_instant("trace1", &["benchmark", "test"], None)];
        b.iter(|| {
            serialized.clear();
            serialize(&samples, &mut serialized).unwrap();
        });
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_multiple_elements(b: &mut Bencher) {
        use super::super::*;
        use super::*;

        let mut serialized = Vec::<u8>::new();
        let samples = vec![
            Sample::new_instant("trace1", &["benchmark", "test"], None),
            Sample::new_instant("trace2", &["benchmark"], None),
            Sample::new_duration("trace3", &["benchmark"], Some(to_payload("some payload")), 0, 0),
            Sample::new_instant("trace4", &["benchmark"], None),
        ];

        b.iter(|| {
            serialized.clear();
            serialize(&samples, &mut serialized).unwrap();
        });
    }
}
<commit_msg>trace: fix tests w/ chrome_trace_dump<commit_after>// Copyright 2018 The xi-editor Authors.
//
// 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.

#![allow(
    clippy::if_same_then_else,
    clippy::needless_bool,
    clippy::needless_pass_by_value,
    clippy::ptr_arg
)]

#[cfg(all(test, feature = "benchmarks"))]
extern crate test;

use super::Sample;
use serde_json;
use std::io::{Error as IOError, ErrorKind as IOErrorKind, Read, Write};

#[derive(Debug)]
pub enum Error {
    Io(IOError),
    Json(serde_json::Error),
    DecodingFormat(String),
}

impl From<IOError> for Error {
    fn from(e: IOError) -> Error {
        Error::Io(e)
    }
}

impl From<serde_json::Error> for Error {
    fn from(e: serde_json::Error) -> Error {
        Error::Json(e)
    }
}

impl From<String> for Error {
    fn from(e: String) -> Error {
        Error::DecodingFormat(e)
    }
}

impl Error {
    pub fn already_exists() -> Error {
        Error::Io(IOError::from(IOErrorKind::AlreadyExists))
    }
}

#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
enum ChromeTraceArrayEntries {
    Array(Vec<Sample>),
}

/// This serializes the samples into the [Chrome trace event format](https://www.google.com/url?sa=t&rct=j&q=&esrc=s&source=web&cd=1&ved=0ahUKEwiJlZmDguXYAhUD4GMKHVmEDqIQFggpMAA&url=https%3A%2F%2Fdocs.google.com%2Fdocument%2Fd%2F1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU%2Fpreview&usg=AOvVaw0tBFlVbDVBikdzLqgrWK3g).
///
/// # Arguments
/// `samples` - Something that can be converted into an iterator of sample
/// references.
/// `format` - Which trace format to save the data in.  There are four total
/// formats described in the document.
/// `output` - Where to write the serialized result.
///
/// # Returns
/// A `Result<(), Error>` that indicates if serialization was successful or the
/// details of any error that occured.
///
/// # Examples
/// ```norun
/// let samples = xi_trace::samples_cloned_sorted();
/// let mut serialized = Vec::<u8>::new();
/// serialize(samples.iter(), serialized);
/// ```
pub fn serialize<W>(samples: &Vec<Sample>, output: W) -> Result<(), Error>
where
    W: Write,
{
    serde_json::to_writer(output, samples).map_err(Error::Json)
}

pub fn to_value(samples: &Vec<Sample>) -> Result<serde_json::Value, Error> {
    serde_json::to_value(samples).map_err(Error::Json)
}

pub fn decode(samples: serde_json::Value) -> Result<Vec<Sample>, Error> {
    serde_json::from_value(samples).map_err(Error::Json)
}

pub fn deserialize<R>(input: R) -> Result<Vec<Sample>, Error>
where
    R: Read,
{
    serde_json::from_reader(input).map_err(Error::Json)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::TracePayloadT;
    #[cfg(feature = "benchmarks")]
    use test::Bencher;

    #[cfg(not(feature = "json_payload"))]
    fn to_payload(value: &'static str) -> &'static str {
        value
    }

    #[cfg(feature = "json_payload")]
    fn to_payload(value: &'static str) -> TracePayloadT {
        json!({ "test": value })
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_serialization() {
        use super::super::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload(
            "sample4",
            &["test", "chrome"],
            || {
                let _guard = trace.block("sample5", &["test,chrome"]);
            },
            to_payload("payload 4"),
        );

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();

        let result = serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let decoded_result: Vec<serde_json::Value> = serde_json::from_slice(&serialized).unwrap();
        assert_eq!(decoded_result.len(), 8);
        assert_eq!(decoded_result[0]["name"].as_str().unwrap(), "process_name");
        assert_eq!(decoded_result[1]["name"].as_str().unwrap(), "thread_name");
        for i in 2..5 {
            assert_eq!(decoded_result[i]["name"].as_str().unwrap(), samples[i].name);
            assert_eq!(decoded_result[i]["cat"].as_str().unwrap(), "test,chrome");
            assert_eq!(decoded_result[i]["ph"].as_str().unwrap(), "i");
            assert_eq!(decoded_result[i]["ts"], samples[i].timestamp_us);
            let nth_sample = &samples[i];
            let nth_args = nth_sample.args.as_ref().unwrap();
            assert_eq!(decoded_result[i]["args"]["xi_payload"], json!(nth_args.payload.as_ref()));
        }
        assert_eq!(decoded_result[5]["ph"], "B");
        assert_eq!(decoded_result[6]["ph"], "E");
        assert_eq!(decoded_result[7]["ph"], "X");
    }

    #[cfg(feature = "chrome_trace_event")]
    #[test]
    fn test_chrome_trace_deserialization() {
        use super::super::*;

        let trace = Trace::enabled(Config::with_limit_count(10));
        trace.instant("sample1", &["test", "chrome"]);
        trace.instant_payload("sample2", &["test", "chrome"], to_payload("payload 2"));
        trace.instant_payload("sample3", &["test", "chrome"], to_payload("payload 3"));
        trace.closure_payload("sample4", &["test", "chrome"], || (), to_payload("payload 4"));

        let samples = trace.samples_cloned_unsorted();

        let mut serialized = Vec::<u8>::new();
        let result = serialize(&samples, &mut serialized);
        assert!(result.is_ok(), "{:?}", result);

        let deserialized_samples = deserialize(serialized.as_slice()).unwrap();
        assert_eq!(deserialized_samples, samples);
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_one_element(b: &mut Bencher) {
        use super::*;

        let mut serialized = Vec::<u8>::new();
        let samples = vec![super::Sample::new_instant("trace1", &["benchmark", "test"], None)];
        b.iter(|| {
            serialized.clear();
            serialize(&samples, &mut serialized).unwrap();
        });
    }

    #[cfg(all(feature = "chrome_trace_event", feature = "benchmarks"))]
    #[bench]
    fn bench_chrome_trace_serialization_multiple_elements(b: &mut Bencher) {
        use super::super::*;
        use super::*;

        let mut serialized = Vec::<u8>::new();
        let samples = vec![
            Sample::new_instant("trace1", &["benchmark", "test"], None),
            Sample::new_instant("trace2", &["benchmark"], None),
            Sample::new_duration("trace3", &["benchmark"], Some(to_payload("some payload")), 0, 0),
            Sample::new_instant("trace4", &["benchmark"], None),
        ];

        b.iter(|| {
            serialized.clear();
            serialize(&samples, &mut serialized).unwrap();
        });
    }
}
<|endoftext|>


<commit_before>extern crate tempfile;
extern crate zmq;

use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use std::io;
use std::io::ErrorKind;
use std::io::Write;
use std::ops::Deref;
use std::path::Path;
use std::sync::Arc;
use std::sync::Mutex;

use time;

use dispatcher::server;
use helpers;
use helpers::{TaskProgress, TaskReport, TaskStatus};
use models::Service;

/// Specifies the binding and operation parameters for a ZMQ sink component
pub struct Sink {
  /// port to listen on
  pub port: usize,
  /// the size of the dispatch queue
  /// (also the batch size for Task store queue requests)
  pub queue_size: usize,
  /// size of an individual message chunk sent via zeromq
  /// (keep this small to avoid large RAM use, increase to reduce network bandwidth)
  pub message_size: usize,
  /// address for the Task store postgres endpoint
  pub backend_address: String,
}

impl Sink {
  /// Starts a receiver/sink `Server` (ZMQ Pull), to accept processing responses.
  /// The sink shares state with other manager threads via queues for tasks in progress,
  /// as well as a queue for completed tasks pending persisting to disk.
  /// A job limit can be provided as a termination condition for the sink server.
  pub fn start(
    &self,
    services_arc: Arc<Mutex<HashMap<String, Option<Service>>>>,
    progress_queue_arc: Arc<Mutex<HashMap<i64, TaskProgress>>>,
    done_queue_arc: Arc<Mutex<Vec<TaskReport>>>,
    job_limit: Option<usize>,
  ) -> Result<(), Box<Error>>
  {
    // Ok, let's bind to a port and start broadcasting
    let context = zmq::Context::new();
    let sink = context.socket(zmq::PULL)?;
    let port_str = self.port.to_string();
    let address = "tcp://*:".to_string() + &port_str;
    assert!(sink.bind(&address).is_ok());

    let mut sink_job_count: usize = 0;

    loop {
      let mut recv_msg = zmq::Message::new()?;
      let mut taskid_msg = zmq::Message::new()?;
      let mut service_msg = zmq::Message::new()?;

      sink.recv(&mut service_msg, 0)?;
      let service_name = match service_msg.as_str() {
        Some(some_name) => some_name,
        None => "_unknown_",
      };

      sink.recv(&mut taskid_msg, 0)?;
      let taskid_str = match taskid_msg.as_str() {
        Some(some_id) => some_id,
        None => "-1",
      };
      let taskid = match taskid_str.parse::<i64>() {
        Ok(some_id) => some_id,
        Err(_) => -1,
      };
      // We have a job, count it
      sink_job_count += 1;
      let mut total_incoming = 0;
      let request_time = time::get_time();
      println!(
        "Incoming sink job {:?} for service {:?}, taskid: {}",
        sink_job_count, service_name, taskid_str
      );

      if let Some(task_progress) = server::pop_progress_task(&progress_queue_arc, taskid) {
        let task = task_progress.task;
        match server::get_service(service_name, &services_arc) {
          None => {
            return Err(Box::new(io::Error::new(
              ErrorKind::Other,
              "TODO: Server::get_service found nothing.",
            )));
          }, // TODO: Handle errors
          Some(service) => {
            if service.id == task.service_id {
              // println!("Task and Service match up.");
              if service.id == 1 {
                // No payload needed for init
                sink.recv(&mut recv_msg, 0)?;
                let done_report = TaskReport {
                  task: task.clone(),
                  status: TaskStatus::NoProblem,
                  messages: Vec::new(),
                };
                server::push_done_queue(&done_queue_arc, done_report);
              } else {
                // Receive the rest of the input in the correct file
                match Path::new(&task.entry.clone()).parent() {
                  None => {
                    println!("-- Error TODO: Path::new(&task.entry).parent() failed.");
                  },
                  Some(recv_dir) => {
                    match recv_dir.to_str() {
                      None => {
                        println!("-- Error TODO: recv_dir.to_str() failed");
                      },
                      Some(recv_dir_str) => {
                        let recv_dir_string = recv_dir_str.to_string();
                        let recv_pathname = recv_dir_string + "/" + &service.name + ".zip";
                        let recv_path = Path::new(&recv_pathname);
                        // println!("Will write to {:?}", recv_path);
                        {
                          // Explicitly scope file, so that we drop it the moment we are done
                          // writing.
                          let mut file = match File::create(recv_path) {
                            Ok(f) => f,
                            Err(e) => {
                              println!("-- Error TODO: File::create(recv_path): {:?}", e);
                              continue;
                            },
                          };
                          while let Ok(_) = sink.recv(&mut recv_msg, 0) {
                            match file.write(recv_msg.deref()) {
                              Ok(written_bytes) => total_incoming += written_bytes,
                              Err(e) => {
                                println!(
                                  "-- Error TODO: file.write(recv_msg.deref()) failed: {:?}",
                                  e
                                );
                                break;
                              },
                            };
                            match sink.get_rcvmore() {
                              Ok(true) => {}, // keep receiving
                              _ => break,     /* println!("Error TODO: sink.get_rcvmore failed:
                                                * {:?}", e); */
                            };
                          }
                          drop(file);
                        }
                        // Then mark the task done. This can be in a new thread later on
                        let done_report = helpers::generate_report(task, recv_path);
                        server::push_done_queue(&done_queue_arc, done_report);
                      },
                    }
                  },
                }
              }
            } else {
              // Otherwise just discard the rest of the message
              println!(
                "-- Mismatch between requested service id {:?} and task's service id {:?} for task {:?}, discarding response",
                service.id, task.service_id, taskid
              );
              while let Ok(_) = sink.recv(&mut recv_msg, 0) {
                if !sink.get_rcvmore()? {
                  break;
                }
              }
            }
          },
        };
      } else {
        // No such task, just discard the next message from the sink
        println!("-- No such task id found in dispatcher queue: {:?}", taskid);
        while let Ok(_) = sink.recv(&mut recv_msg, 0) {
          if !sink.get_rcvmore()? {
            break;
          }
        }
      }
      let responded_time = time::get_time();
      let request_duration = (responded_time - request_time).num_milliseconds();
      println!(
        "Sink job {}, message size: {}, took {}ms.",
        sink_job_count, total_incoming, request_duration
      );
      if let Some(limit_number) = job_limit {
        if sink_job_count >= limit_number {
          println!(
            "Manager job limit of {:?} reached, terminating Sink thread...",
            limit_number
          );
          break;
        }
      }
    }
    Ok(())
  }
}
<commit_msg>minor, print the final taskid int<commit_after>extern crate tempfile;
extern crate zmq;

use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use std::io;
use std::io::ErrorKind;
use std::io::Write;
use std::ops::Deref;
use std::path::Path;
use std::sync::Arc;
use std::sync::Mutex;

use time;

use dispatcher::server;
use helpers;
use helpers::{TaskProgress, TaskReport, TaskStatus};
use models::Service;

/// Specifies the binding and operation parameters for a ZMQ sink component
pub struct Sink {
  /// port to listen on
  pub port: usize,
  /// the size of the dispatch queue
  /// (also the batch size for Task store queue requests)
  pub queue_size: usize,
  /// size of an individual message chunk sent via zeromq
  /// (keep this small to avoid large RAM use, increase to reduce network bandwidth)
  pub message_size: usize,
  /// address for the Task store postgres endpoint
  pub backend_address: String,
}

impl Sink {
  /// Starts a receiver/sink `Server` (ZMQ Pull), to accept processing responses.
  /// The sink shares state with other manager threads via queues for tasks in progress,
  /// as well as a queue for completed tasks pending persisting to disk.
  /// A job limit can be provided as a termination condition for the sink server.
  pub fn start(
    &self,
    services_arc: Arc<Mutex<HashMap<String, Option<Service>>>>,
    progress_queue_arc: Arc<Mutex<HashMap<i64, TaskProgress>>>,
    done_queue_arc: Arc<Mutex<Vec<TaskReport>>>,
    job_limit: Option<usize>,
  ) -> Result<(), Box<Error>>
  {
    // Ok, let's bind to a port and start broadcasting
    let context = zmq::Context::new();
    let sink = context.socket(zmq::PULL)?;
    let port_str = self.port.to_string();
    let address = "tcp://*:".to_string() + &port_str;
    assert!(sink.bind(&address).is_ok());

    let mut sink_job_count: usize = 0;

    loop {
      let mut recv_msg = zmq::Message::new()?;
      let mut taskid_msg = zmq::Message::new()?;
      let mut service_msg = zmq::Message::new()?;

      sink.recv(&mut service_msg, 0)?;
      let service_name = match service_msg.as_str() {
        Some(some_name) => some_name,
        None => "_unknown_",
      };

      sink.recv(&mut taskid_msg, 0)?;
      let taskid_str = match taskid_msg.as_str() {
        Some(some_id) => some_id,
        None => "-1",
      };
      let taskid = match taskid_str.parse::<i64>() {
        Ok(some_id) => some_id,
        Err(_) => -1,
      };
      // We have a job, count it
      sink_job_count += 1;
      let mut total_incoming = 0;
      let request_time = time::get_time();
      println!(
        "Incoming sink job {:?} for service {:?}, taskid: {}",
        sink_job_count, service_name, taskid
      );

      if let Some(task_progress) = server::pop_progress_task(&progress_queue_arc, taskid) {
        let task = task_progress.task;
        match server::get_service(service_name, &services_arc) {
          None => {
            return Err(Box::new(io::Error::new(
              ErrorKind::Other,
              "TODO: Server::get_service found nothing.",
            )));
          }, // TODO: Handle errors
          Some(service) => {
            if service.id == task.service_id {
              // println!("Task and Service match up.");
              if service.id == 1 {
                // No payload needed for init
                sink.recv(&mut recv_msg, 0)?;
                let done_report = TaskReport {
                  task: task.clone(),
                  status: TaskStatus::NoProblem,
                  messages: Vec::new(),
                };
                server::push_done_queue(&done_queue_arc, done_report);
              } else {
                // Receive the rest of the input in the correct file
                match Path::new(&task.entry.clone()).parent() {
                  None => {
                    println!("-- Error TODO: Path::new(&task.entry).parent() failed.");
                  },
                  Some(recv_dir) => {
                    match recv_dir.to_str() {
                      None => {
                        println!("-- Error TODO: recv_dir.to_str() failed");
                      },
                      Some(recv_dir_str) => {
                        let recv_dir_string = recv_dir_str.to_string();
                        let recv_pathname = recv_dir_string + "/" + &service.name + ".zip";
                        let recv_path = Path::new(&recv_pathname);
                        // println!("Will write to {:?}", recv_path);
                        {
                          // Explicitly scope file, so that we drop it the moment we are done
                          // writing.
                          let mut file = match File::create(recv_path) {
                            Ok(f) => f,
                            Err(e) => {
                              println!("-- Error TODO: File::create(recv_path): {:?}", e);
                              continue;
                            },
                          };
                          while let Ok(_) = sink.recv(&mut recv_msg, 0) {
                            match file.write(recv_msg.deref()) {
                              Ok(written_bytes) => total_incoming += written_bytes,
                              Err(e) => {
                                println!(
                                  "-- Error TODO: file.write(recv_msg.deref()) failed: {:?}",
                                  e
                                );
                                break;
                              },
                            };
                            match sink.get_rcvmore() {
                              Ok(true) => {}, // keep receiving
                              _ => break,     /* println!("Error TODO: sink.get_rcvmore failed:
                                                * {:?}", e); */
                            };
                          }
                          drop(file);
                        }
                        // Then mark the task done. This can be in a new thread later on
                        let done_report = helpers::generate_report(task, recv_path);
                        server::push_done_queue(&done_queue_arc, done_report);
                      },
                    }
                  },
                }
              }
            } else {
              // Otherwise just discard the rest of the message
              println!(
                "-- Mismatch between requested service id {:?} and task's service id {:?} for task {:?}, discarding response",
                service.id, task.service_id, taskid
              );
              while let Ok(_) = sink.recv(&mut recv_msg, 0) {
                if !sink.get_rcvmore()? {
                  break;
                }
              }
            }
          },
        };
      } else {
        // No such task, just discard the next message from the sink
        println!("-- No such task id found in dispatcher queue: {:?}", taskid);
        while let Ok(_) = sink.recv(&mut recv_msg, 0) {
          if !sink.get_rcvmore()? {
            break;
          }
        }
      }
      let responded_time = time::get_time();
      let request_duration = (responded_time - request_time).num_milliseconds();
      println!(
        "Sink job {}, message size: {}, took {}ms.",
        sink_job_count, total_incoming, request_duration
      );
      if let Some(limit_number) = job_limit {
        if sink_job_count >= limit_number {
          println!(
            "Manager job limit of {:?} reached, terminating Sink thread...",
            limit_number
          );
          break;
        }
      }
    }
    Ok(())
  }
}
<|endoftext|>


<commit_before>#![feature(specialization)]
#![feature(conservative_impl_trait)]
#![feature(cfg_target_feature)]
#![feature(asm)]

#![allow(unused_imports)]

extern crate core;
extern crate typenum;
extern crate momms;

use std::time::{Instant};
use typenum::{U1, B0, UInt, Unsigned};

use momms::kern::{KernelNM,KernelMN};
use momms::matrix::{Scalar, Mat, ColumnPanelMatrix, RowPanelMatrix, Matrix, Hierarch};
use momms::composables::{GemmNode, AlgorithmStep, PartM, PartN, PartK, PackA, PackB, SpawnThreads, ParallelM, ParallelN, TheRest};
use momms::thread_comm::ThreadInfo;
use momms::util;

fn test_algorithm<T: Scalar, Mr: Unsigned, Nr: Unsigned, Kc:Unsigned, CLRS: Unsigned, CLCS: Unsigned,
    S: GemmNode<T, Hierarch<T, Mr, Kc, U1, Mr>, Hierarch<T, Kc, Nr, Nr, U1>, Hierarch<T, Mr, Nr, CLRS, CLCS>>>
    ( m:usize, n: usize, k: usize, algo: &mut S, flusher: &mut Vec<f64>, n_reps: usize ) -> (f64, T)
{
    let algo_desc = S::hierarchy_description();
    let mut best_time: f64 = 9999999999.0;
    let mut worst_err: T = T::zero();

    for _ in 0..n_reps {
        //Create matrices.
        let mut a : Hierarch<T, Mr, Kc, U1, Mr> = Hierarch::new(m, k, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::K{bsz: 0});
        let mut b : Hierarch<T, Kc, Nr, Nr, U1> = Hierarch::new(k, n, &algo_desc, AlgorithmStep::K{bsz: 0}, AlgorithmStep::N{bsz: 0});
        let mut c : Hierarch<T, Mr, Nr, CLRS, CLCS> = Hierarch::new(m, n, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::N{bsz: 0});

        //Fill the matrices
        a.fill_rand(); c.fill_zero(); b.fill_rand();

        //Read a buffer so that A, B, and C are cold in cache.
        for i in flusher.iter_mut() { *i += 1.0; }

        //Time and run algorithm
        let start = Instant::now();
        unsafe{ algo.run( &mut a, &mut b, &mut c, &ThreadInfo::single_thread() ); }
        best_time = best_time.min(util::dur_seconds(start));
        let err = util::test_c_eq_a_b( &mut a, &mut b, &mut c);
        worst_err = worst_err.max(err);
    }
    (best_time, worst_err)
}

fn test(m_selector: isize, n_selector: isize, k_selector: isize) {
    //Goto's algorithm
    type U3000 = UInt<UInt<typenum::U750, B0>, B0>;
    type Nc = U3000;
    type Kc = typenum::U192;
    type Mc = typenum::U120;
    type Mr = typenum::U4;
    type Nr = typenum::U12;

    type Goto<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, Nc,
          PartK<T, MTA, MTB, MTC, Kc,
          PartM<T, MTA, MTB, MTC, Mc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>;

    type RootS3 = typenum::U768;
    //Resident A algorithm
    type NcL2 = typenum::U120;
    type L3A<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartM<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartN<T, MTA, MTB, MTC, NcL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelM<T, MTA, MTB, MTC, Nr, TheRest, //This algorithm uses a 12x4 ukernel instead of 4x12 like the other two
          KernelMN<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //Resident B algorithm
    type McL2 = typenum::U120;
    type L3B<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartM<T, MTA, MTB, MTC, McL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    type L3CNc = typenum::U780;
    type L3CKc = typenum::U156;
    //Resident C algorithm
    type L3C<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, L3CNc,
          PartM<T, MTA, MTB, MTC, L3CNc,
          PartK<T, MTA, MTB, MTC, L3CKc,
          PartM<T, MTA, MTB, MTC, L3CKc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //We can use the same matrix type for all three algorithms
    type HierA<T> = Hierarch<T, Mr, Kc, U1, Mr>;
    type HierB<T> = Hierarch<T, Kc, Nr, Nr, U1>;
    type HierC<T> = Hierarch<T, Mr, Nr, Nr, U1>;

    type HierAL3a<T> = Hierarch<T, Nr, Kc, U1, Nr>;
    type HierBL3a<T> = Hierarch<T, Kc, Mr, Mr, U1>;
    type HierCL3a<T> = Hierarch<T, Nr, Mr, U1, Nr>;

    type HierAL3c<T> = Hierarch<T, Mr, L3CKc, U1, Mr>;
    type HierBL3c<T> = Hierarch<T, L3CKc, Nr, Nr, U1>;
    type HierCL3c<T> = Hierarch<T, Mr, Nr, Nr, U1>;

    let mut goto = <Goto<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3a = <L3A<f64, HierAL3a<f64>, HierBL3a<f64>, HierCL3a<f64>>>::new();
    let mut l3b = <L3B<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3c = <L3C<f64, HierAL3c<f64>, HierBL3c<f64>, HierCL3c<f64>>>::new();

    goto.set_n_threads(4);
    l3a.set_n_threads(4);
    l3b.set_n_threads(4);
    l3c.set_n_threads(4);

    let flusher_len = 2*1024*1024; //16MB
    let mut flusher : Vec<f64> = Vec::with_capacity(flusher_len);
    for _ in 0..flusher_len {
        flusher.push(0.0);
    }

    println!("m\tn\tk\t{: <13}{: <13}{: <13}{: <13}{: <15}{: <15}{: <15}{: <15}", "goto", "l3a", "l3b", "l3c", "goto", "l3a", "l3b", "l3c");
    for index in 01..81 {
        let size = index*50;
//        let (m, n, k) = (size, size, size);
        let m = if m_selector < 0 { size * isize::abs(m_selector) } else { m_selector } as usize;
        let n = if n_selector < 0 { size * isize::abs(n_selector) } else { n_selector } as usize;
        let k = if k_selector < 0 { size * isize::abs(k_selector) } else { k_selector } as usize;

        let n_reps = 5;
        let (goto_time, goto_err) = test_algorithm(m, n, k, &mut goto, &mut flusher, n_reps);
        let (l3a_time, l3a_err) = test_algorithm(m, n, k, &mut l3a, &mut flusher, n_reps);
        let (l3b_time, l3b_err) = test_algorithm(m, n, k, &mut l3b, &mut flusher, n_reps);
        let (l3c_time, l3c_err) = test_algorithm(m, n, k, &mut l3c, &mut flusher, n_reps);

        println!("{}\t{}\t{}\t{}{}{}{}{}{}{}{}",
                 m, n, k,
                 format!("{: <13.5}", util::gflops(m,n,k,goto_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3a_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3b_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3c_time)),
                 format!("{: <15.5e}", goto_err.sqrt()),
                 format!("{: <15.5e}", l3a_err.sqrt()),
                 format!("{: <15.5e}", l3b_err.sqrt()),
                 format!("{: <15.5e}", l3c_err.sqrt()));
    }

    let mut sum = 0.0;
    for a in flusher.iter() {
        sum += *a;
    }
    println!("Flush value {}", sum);
}

fn main() {
    //Best for Resident A
    test(768,-1,768);
    //Best for Resident B
    test(-1,768,768);
    //Best for Resident C
    test(768,768,-1);
}
<commit_msg>Adjust blocksize for L3C2B algorithm<commit_after>#![feature(specialization)]
#![feature(conservative_impl_trait)]
#![feature(cfg_target_feature)]
#![feature(asm)]

#![allow(unused_imports)]

extern crate core;
extern crate typenum;
extern crate momms;

use std::time::{Instant};
use typenum::{U1, B0, UInt, Unsigned};

use momms::kern::{KernelNM,KernelMN};
use momms::matrix::{Scalar, Mat, ColumnPanelMatrix, RowPanelMatrix, Matrix, Hierarch};
use momms::composables::{GemmNode, AlgorithmStep, PartM, PartN, PartK, PackA, PackB, SpawnThreads, ParallelM, ParallelN, TheRest};
use momms::thread_comm::ThreadInfo;
use momms::util;

fn test_algorithm<T: Scalar, Mr: Unsigned, Nr: Unsigned, Kc:Unsigned, CLRS: Unsigned, CLCS: Unsigned,
    S: GemmNode<T, Hierarch<T, Mr, Kc, U1, Mr>, Hierarch<T, Kc, Nr, Nr, U1>, Hierarch<T, Mr, Nr, CLRS, CLCS>>>
    ( m:usize, n: usize, k: usize, algo: &mut S, flusher: &mut Vec<f64>, n_reps: usize ) -> (f64, T)
{
    let algo_desc = S::hierarchy_description();
    let mut best_time: f64 = 9999999999.0;
    let mut worst_err: T = T::zero();

    for _ in 0..n_reps {
        //Create matrices.
        let mut a : Hierarch<T, Mr, Kc, U1, Mr> = Hierarch::new(m, k, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::K{bsz: 0});
        let mut b : Hierarch<T, Kc, Nr, Nr, U1> = Hierarch::new(k, n, &algo_desc, AlgorithmStep::K{bsz: 0}, AlgorithmStep::N{bsz: 0});
        let mut c : Hierarch<T, Mr, Nr, CLRS, CLCS> = Hierarch::new(m, n, &algo_desc, AlgorithmStep::M{bsz: 0}, AlgorithmStep::N{bsz: 0});

        //Fill the matrices
        a.fill_rand(); c.fill_zero(); b.fill_rand();

        //Read a buffer so that A, B, and C are cold in cache.
        for i in flusher.iter_mut() { *i += 1.0; }

        //Time and run algorithm
        let start = Instant::now();
        unsafe{ algo.run( &mut a, &mut b, &mut c, &ThreadInfo::single_thread() ); }
        best_time = best_time.min(util::dur_seconds(start));
        let err = util::test_c_eq_a_b( &mut a, &mut b, &mut c);
        worst_err = worst_err.max(err);
    }
    (best_time, worst_err)
}

fn test(m_selector: isize, n_selector: isize, k_selector: isize) {
    //Goto's algorithm
    type U3000 = UInt<UInt<typenum::U750, B0>, B0>;
    type Nc = U3000;
    type Kc = typenum::U192;
    type Mc = typenum::U120;
    type Mr = typenum::U4;
    type Nr = typenum::U12;

    type Goto<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, Nc,
          PartK<T, MTA, MTB, MTC, Kc,
          PartM<T, MTA, MTB, MTC, Mc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>;

    type RootS3 = typenum::U768;
    //Resident A algorithm
    type NcL2 = typenum::U120;
    type L3A<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartM<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartN<T, MTA, MTB, MTC, NcL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelM<T, MTA, MTB, MTC, Nr, TheRest, //This algorithm uses a 12x4 ukernel instead of 4x12 like the other two
          KernelMN<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //Resident B algorithm
    type McL2 = typenum::U120;
    type L3B<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, RootS3,
          PartK<T, MTA, MTB, MTC, RootS3,
          PartM<T, MTA, MTB, MTC, McL2,
          PartK<T, MTA, MTB, MTC, Kc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

//    type L3CNc = typenum::U780; //Too Big!
    type L3CNc = typenum::U624;
    type L3CKc = typenum::U156;
    //Resident C algorithm
    type L3C<T,MTA,MTB,MTC>
        = SpawnThreads<T, MTA, MTB, MTC,
          PartN<T, MTA, MTB, MTC, L3CNc,
          PartM<T, MTA, MTB, MTC, L3CNc,
          PartK<T, MTA, MTB, MTC, L3CKc,
          PartM<T, MTA, MTB, MTC, L3CKc,
          ParallelN<T, MTA, MTB, MTC, Nr, TheRest,
          KernelNM<T, MTA, MTB, MTC, Nr, Mr>>>>>>>;

    //We can use the same matrix type for all three algorithms
    type HierA<T> = Hierarch<T, Mr, Kc, U1, Mr>;
    type HierB<T> = Hierarch<T, Kc, Nr, Nr, U1>;
    type HierC<T> = Hierarch<T, Mr, Nr, Nr, U1>;

    type HierAL3a<T> = Hierarch<T, Nr, Kc, U1, Nr>;
    type HierBL3a<T> = Hierarch<T, Kc, Mr, Mr, U1>;
    type HierCL3a<T> = Hierarch<T, Nr, Mr, U1, Nr>;

    type HierAL3c<T> = Hierarch<T, Mr, L3CKc, U1, Mr>;
    type HierBL3c<T> = Hierarch<T, L3CKc, Nr, Nr, U1>;
    type HierCL3c<T> = Hierarch<T, Mr, Nr, Nr, U1>;

    let mut goto = <Goto<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3a = <L3A<f64, HierAL3a<f64>, HierBL3a<f64>, HierCL3a<f64>>>::new();
    let mut l3b = <L3B<f64, HierA<f64>, HierB<f64>, HierC<f64>>>::new();
    let mut l3c = <L3C<f64, HierAL3c<f64>, HierBL3c<f64>, HierCL3c<f64>>>::new();

    goto.set_n_threads(4);
    l3a.set_n_threads(4);
    l3b.set_n_threads(4);
    l3c.set_n_threads(4);

    let flusher_len = 2*1024*1024; //16MB
    let mut flusher : Vec<f64> = Vec::with_capacity(flusher_len);
    for _ in 0..flusher_len {
        flusher.push(0.0);
    }

    println!("m\tn\tk\t{: <13}{: <13}{: <13}{: <13}{: <15}{: <15}{: <15}{: <15}", "goto", "l3a", "l3b", "l3c", "goto", "l3a", "l3b", "l3c");
    for index in 01..81 {
        let size = index*50;
//        let (m, n, k) = (size, size, size);
        let m = if m_selector < 0 { size * isize::abs(m_selector) } else { m_selector } as usize;
        let n = if n_selector < 0 { size * isize::abs(n_selector) } else { n_selector } as usize;
        let k = if k_selector < 0 { size * isize::abs(k_selector) } else { k_selector } as usize;

        let n_reps = 5;
        let (goto_time, goto_err) = test_algorithm(m, n, k, &mut goto, &mut flusher, n_reps);
        let (l3a_time, l3a_err) = test_algorithm(m, n, k, &mut l3a, &mut flusher, n_reps);
        let (l3b_time, l3b_err) = test_algorithm(m, n, k, &mut l3b, &mut flusher, n_reps);
        let (l3c_time, l3c_err) = test_algorithm(m, n, k, &mut l3c, &mut flusher, n_reps);

        println!("{}\t{}\t{}\t{}{}{}{}{}{}{}{}",
                 m, n, k,
                 format!("{: <13.5}", util::gflops(m,n,k,goto_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3a_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3b_time)),
                 format!("{: <13.5}", util::gflops(m,n,k,l3c_time)),
                 format!("{: <15.5e}", goto_err.sqrt()),
                 format!("{: <15.5e}", l3a_err.sqrt()),
                 format!("{: <15.5e}", l3b_err.sqrt()),
                 format!("{: <15.5e}", l3c_err.sqrt()));
    }

    let mut sum = 0.0;
    for a in flusher.iter() {
        sum += *a;
    }
    println!("Flush value {}", sum);
}

fn main() {
    let smalldim = 600;
    //2 small 1 large
    println!("Best for L3B");
    test(-1,smalldim,smalldim);
    println!("Best for L3A");
    test(smalldim,-1,smalldim);
    println!("Best for L3C");
    test(smalldim,smalldim,-1);

    //square
    println!("Square");
    test(-1,-1,-1);

    //2 large 1 small
    println!("Best for L3A or L3C");
    test(smalldim,-1,-1);
    println!("Best for L3B or L3C");
    test(-1,smalldim,-1);
    println!("Best for L3A or L3B");
    test(-1,-1,smalldim);
}
<|endoftext|>


<commit_before>// pest. Elegant, efficient grammars
// Copyright (C) 2016  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::iter::Peekable;
use std::str::{self, Chars};

use super::super::Input;

/// A `struct` useful for matching in-memory `String`s.
///
/// # Examples
///
/// ```
/// # use pest::Input;
/// # use pest::StringInput;
/// let mut input = StringInput::new("asdasdf");
///
/// assert!(input.match_string("asd"));
/// assert!(input.match_string("asdf"));
/// assert!(!input.match_string("nope"));
/// ```
pub struct StringInput {
    string: String,
    pos: usize
}

impl StringInput {
    /// Creates a new `StringInput` from a `&str`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use pest::Input;
    /// # use pest::StringInput;
    /// let mut input = StringInput::new("asd");
    ///
    /// assert_eq!(input.len(), 3);
    /// ```
    pub fn new(string: &str) -> StringInput {
        StringInput {
            string: string.to_owned(),
            pos : 0
        }
    }
}

impl Input for StringInput {
    #[inline]
    fn len(&self) -> usize {
        self.string.len()
    }

    #[inline]
    fn is_empty(&self) -> bool {
        self.string.is_empty()
    }

    #[inline]
    fn pos(&self) -> usize {
        self.pos
    }

    #[inline]
    fn set_pos(&mut self, pos: usize) {
        self.pos = pos
    }

    #[inline]
    fn slice(&self, start: usize, end: usize) -> &str {
        &self.string[start..end]
    }

    #[inline]
    fn line_col(&self, pos: usize) -> (usize, usize) {
        fn find(chars: &mut Peekable<Chars>, pos: usize,
                current: (usize, usize)) -> (usize, usize) {
            if pos == 0 {
                current
            } else {
                match chars.next() {
                    Some('\r') => {
                        if let Some(&'\n') = chars.peek() {
                            chars.next();

                            if pos == 1 {
                                find(chars, pos - 1, (current.0 + 1, 1))
                            } else {
                                find(chars, pos - 2, (current.0 + 1, 1))
                            }
                        } else {
                            find(chars, pos - 1, (current.0 + 1, 1))
                        }
                    },
                    Some('\n') => find(chars, pos - 1, (current.0 + 1, 1)),
                    Some(_)    => find(chars, pos - 1, (current.0, current.1 + 1)),
                    None       => unreachable!()
                }
            }
        }

        if pos > self.string.len() {
            panic!("position out of bounds");
        }

        find(&mut self.string.chars().peekable(), pos, (1, 1))
    }

    #[inline]
    fn match_string(&mut self, string: &str) -> bool {
        let to = self.pos + string.len();

        if to <= self.string.len() {
            let slice = unsafe { self.string.slice_unchecked(self.pos, to) };
            let result = slice == string;

            if result {
                self.pos = to;
            }

            result
        } else {
            false
        }
    }

    #[inline]
    fn match_range(&mut self, left: char, right: char) -> bool {
        let len = left.len_utf8();

        if len != right.len_utf8() {
            panic!("ranges should have same-sized UTF-8 limits");
        }

        let to = self.pos + len;

        if to <= self.string.len() {
            if let Ok(string) = str::from_utf8(&self.string.as_bytes()[self.pos..to]) {
                let c = string.chars().next().unwrap();

                let result = left <= c && c <= right;

                if result {
                    self.pos += len;
                }

                result
            } else {
                false
            }
        } else {
            false
        }
    }
}

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

    #[test]
    fn empty() {
        let mut input = StringInput::new("");

        assert!(input.match_string(""));
        assert!(!input.match_string("a"));
    }

    #[test]
    fn parts() {
        let mut input = StringInput::new("asdasdf");

        assert!(input.match_string("asd"));
        assert!(input.match_string("asdf"));
    }

    #[test]
    fn len() {
        assert_eq!(StringInput::new("asdasdf").len(), 7);
    }

    #[test]
    fn pos() {
        let mut input = StringInput::new("asdasdf");

        assert_eq!(input.pos(), 0);
        assert!(input.match_string("asd"));
        assert_eq!(input.pos(), 3);
        assert!(input.match_string("asdf"));
        assert_eq!(input.pos(), 7);

        input.set_pos(3);

        assert_eq!(input.pos(), 3);
        assert!(input.match_string("asdf"));
        assert_eq!(input.pos(), 7);
    }

    #[test]
    fn slice() {
        let input = StringInput::new("asdasdf");

        assert_eq!(input.slice(1, 3), "sd");
    }

    #[test]
    fn line_col() {
        let input = StringInput::new("a\rb\nc\r\nd");

        assert_eq!(input.line_col(0), (1, 1));
        assert_eq!(input.line_col(1), (1, 2));
        assert_eq!(input.line_col(2), (2, 1));
        assert_eq!(input.line_col(3), (2, 2));
        assert_eq!(input.line_col(4), (3, 1));
        assert_eq!(input.line_col(5), (3, 2));
        assert_eq!(input.line_col(6), (4, 1));
        assert_eq!(input.line_col(7), (4, 1));
        assert_eq!(input.line_col(8), (4, 2));
    }

    #[test]
    fn match_range() {
        let mut input = StringInput::new("bbbb");

        assert!(input.match_range('a', 'c'));
        assert!(input.match_range('b', 'b'));
        assert!(!input.match_range('a', 'a'));
        assert!(!input.match_range('c', 'c'));

        assert_eq!(input.pos(), 2);
    }
}
<commit_msg>Added is_empty tests.<commit_after>// pest. Elegant, efficient grammars
// Copyright (C) 2016  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::iter::Peekable;
use std::str::{self, Chars};

use super::super::Input;

/// A `struct` useful for matching in-memory `String`s.
///
/// # Examples
///
/// ```
/// # use pest::Input;
/// # use pest::StringInput;
/// let mut input = StringInput::new("asdasdf");
///
/// assert!(input.match_string("asd"));
/// assert!(input.match_string("asdf"));
/// assert!(!input.match_string("nope"));
/// ```
pub struct StringInput {
    string: String,
    pos: usize
}

impl StringInput {
    /// Creates a new `StringInput` from a `&str`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use pest::Input;
    /// # use pest::StringInput;
    /// let mut input = StringInput::new("asd");
    ///
    /// assert_eq!(input.len(), 3);
    /// ```
    pub fn new(string: &str) -> StringInput {
        StringInput {
            string: string.to_owned(),
            pos : 0
        }
    }
}

impl Input for StringInput {
    #[inline]
    fn len(&self) -> usize {
        self.string.len()
    }

    #[inline]
    fn is_empty(&self) -> bool {
        self.string.is_empty()
    }

    #[inline]
    fn pos(&self) -> usize {
        self.pos
    }

    #[inline]
    fn set_pos(&mut self, pos: usize) {
        self.pos = pos
    }

    #[inline]
    fn slice(&self, start: usize, end: usize) -> &str {
        &self.string[start..end]
    }

    #[inline]
    fn line_col(&self, pos: usize) -> (usize, usize) {
        fn find(chars: &mut Peekable<Chars>, pos: usize,
                current: (usize, usize)) -> (usize, usize) {
            if pos == 0 {
                current
            } else {
                match chars.next() {
                    Some('\r') => {
                        if let Some(&'\n') = chars.peek() {
                            chars.next();

                            if pos == 1 {
                                find(chars, pos - 1, (current.0 + 1, 1))
                            } else {
                                find(chars, pos - 2, (current.0 + 1, 1))
                            }
                        } else {
                            find(chars, pos - 1, (current.0 + 1, 1))
                        }
                    },
                    Some('\n') => find(chars, pos - 1, (current.0 + 1, 1)),
                    Some(_)    => find(chars, pos - 1, (current.0, current.1 + 1)),
                    None       => unreachable!()
                }
            }
        }

        if pos > self.string.len() {
            panic!("position out of bounds");
        }

        find(&mut self.string.chars().peekable(), pos, (1, 1))
    }

    #[inline]
    fn match_string(&mut self, string: &str) -> bool {
        let to = self.pos + string.len();

        if to <= self.string.len() {
            let slice = unsafe { self.string.slice_unchecked(self.pos, to) };
            let result = slice == string;

            if result {
                self.pos = to;
            }

            result
        } else {
            false
        }
    }

    #[inline]
    fn match_range(&mut self, left: char, right: char) -> bool {
        let len = left.len_utf8();

        if len != right.len_utf8() {
            panic!("ranges should have same-sized UTF-8 limits");
        }

        let to = self.pos + len;

        if to <= self.string.len() {
            if let Ok(string) = str::from_utf8(&self.string.as_bytes()[self.pos..to]) {
                let c = string.chars().next().unwrap();

                let result = left <= c && c <= right;

                if result {
                    self.pos += len;
                }

                result
            } else {
                false
            }
        } else {
            false
        }
    }
}

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

    #[test]
    fn empty() {
        let mut input = StringInput::new("");

        assert!(input.is_empty());
        assert!(input.match_string(""));
        assert!(!input.match_string("a"));
    }

    #[test]
    fn parts() {
        let mut input = StringInput::new("asdasdf");

        assert!(!input.is_empty());
        assert!(input.match_string("asd"));
        assert!(input.match_string("asdf"));
    }

    #[test]
    fn len() {
        assert_eq!(StringInput::new("asdasdf").len(), 7);
    }

    #[test]
    fn pos() {
        let mut input = StringInput::new("asdasdf");

        assert_eq!(input.pos(), 0);
        assert!(input.match_string("asd"));
        assert_eq!(input.pos(), 3);
        assert!(input.match_string("asdf"));
        assert_eq!(input.pos(), 7);

        input.set_pos(3);

        assert_eq!(input.pos(), 3);
        assert!(input.match_string("asdf"));
        assert_eq!(input.pos(), 7);
    }

    #[test]
    fn slice() {
        let input = StringInput::new("asdasdf");

        assert_eq!(input.slice(1, 3), "sd");
    }

    #[test]
    fn line_col() {
        let input = StringInput::new("a\rb\nc\r\nd");

        assert_eq!(input.line_col(0), (1, 1));
        assert_eq!(input.line_col(1), (1, 2));
        assert_eq!(input.line_col(2), (2, 1));
        assert_eq!(input.line_col(3), (2, 2));
        assert_eq!(input.line_col(4), (3, 1));
        assert_eq!(input.line_col(5), (3, 2));
        assert_eq!(input.line_col(6), (4, 1));
        assert_eq!(input.line_col(7), (4, 1));
        assert_eq!(input.line_col(8), (4, 2));
    }

    #[test]
    fn match_range() {
        let mut input = StringInput::new("bbbb");

        assert!(input.match_range('a', 'c'));
        assert!(input.match_range('b', 'b'));
        assert!(!input.match_range('a', 'a'));
        assert!(!input.match_range('c', 'c'));

        assert_eq!(input.pos(), 2);
    }
}
<|endoftext|>