<commit_before>use std::mem;
use std::num::SignedInt;
use std::raw;

use super::{Ix, Ixs};

/// Calculate offset from `Ix` stride converting sign properly
#[inline]
pub fn stride_offset(n: Ix, stride: Ix) -> isize
{
    (n as isize) * ((stride as Ixs) as isize)
}

/*
#[inline]
pub fn stride_as_int(stride: Ix) -> isize
{
    (stride as Ixs) as isize
}
*/

/// Trait for the shape and index types of arrays.
///
/// unsafe trait due to how the assumptions in the default impls work.
pub unsafe trait Dimension : Clone + Eq {
    fn ndim(&self) -> usize;
    fn slice<'a>(&'a self) -> &'a [Ix] {
        unsafe {
            mem::transmute(raw::Slice {
                data: self as *const _ as *const Ix,
                len: self.ndim(),
            })
        }
    }

    fn slice_mut<'a>(&'a mut self) -> &'a mut [Ix] {
        unsafe {
            mem::transmute(raw::Slice {
                data: self as *mut _ as *const Ix,
                len: self.ndim(),
            })
        }
    }

    fn size(&self) -> usize {
        self.slice().iter().fold(1, |s, &a| s * a as usize)
    }

    fn default_strides(&self) -> Self {
        // Compute default array strides
        // Shape (a, b, c) => Give strides (b * c, c, 1)
        let mut strides = self.clone();
        {
            let mut it = strides.slice_mut().iter_mut().rev();
            // Set first element to 1
            for rs in it.by_ref() {
                *rs = 1;
                break;
            }
            let mut cum_prod = 1;
            for (rs, dim) in it.zip(self.slice().iter().rev()) {
                cum_prod *= *dim;
                *rs = cum_prod;
            }
        }
        strides
    }

    #[inline]
    fn first_index(&self) -> Option<Self>
    {
        for ax in self.slice().iter() {
            if *ax == 0 {
                return None
            }
        }
        let mut index = self.clone();
        for rr in index.slice_mut().iter_mut() {
            *rr = 0;
        }
        Some(index)
    }

    /// Iteration -- Use self as size, and return next index after `index`
    /// or None if there are no more.
    // FIXME: use &Self for index or even &mut?
    #[inline]
    fn next_for(&self, index: Self) -> Option<Self> {
        let mut index = index;
        let mut done = false;
        for (&dim, ix) in self.slice().iter().rev()
                            .zip(index.slice_mut().iter_mut().rev())
        {
            *ix += 1;
            if *ix == dim {
                *ix = 0;
            } else {
                done = true;
                break;
            }
        }
        if done {
            Some(index)
        } else { None }
    }

    /// Return stride offset for index.
    fn stride_offset(index: &Self, strides: &Self) -> isize
    {
        let mut offset = 0;
        for (&i, &s) in index.slice().iter()
                            .zip(strides.slice().iter()) {
            offset += stride_offset(i, s);
        }
        offset
    }

    /// Return stride offset for this dimension and index.
    fn stride_offset_checked(&self, strides: &Self, index: &Self) -> Option<isize>
    {
        let mut offset = 0;
        for ((&d, &i), &s) in self.slice().iter()
                                .zip(index.slice().iter())
                                .zip(strides.slice().iter())
        {
            if i >= d {
                return None;
            }
            offset += stride_offset(i, s);
        }
        Some(offset)
    }

    /// Modify dimension, strides and return data pointer offset
    ///
    /// **Panics** if `slices` does not correspond to the number of axes,
    /// if any stride is 0, or if any index is out of bounds.
    fn do_slices(dim: &mut Self, strides: &mut Self, slices: &[Si]) -> isize
    {
        let mut offset = 0;
        assert!(slices.len() == dim.slice().len());
        for ((dr, sr), &slc) in dim.slice_mut().iter_mut()
                                .zip(strides.slice_mut().iter_mut())
                                .zip(slices.iter())
        {
            let m = *dr;
            let mi = m as Ixs;
            let Si(b1, opt_e1, s1) = slc;
            let e1 = opt_e1.unwrap_or(mi);

            let b1 = abs_index(mi, b1);
            let mut e1 = abs_index(mi, e1);
            if e1 < b1 { e1 = b1; }

            assert!(b1 <= m);
            assert!(e1 <= m);

            let m = e1 - b1;
            // stride
            let s = (*sr) as Ixs;

            // Data pointer offset
            offset += stride_offset(b1, *sr);
            // Adjust for strides
            assert!(s1 != 0);
            // How to implement negative strides:
            //
            // Increase start pointer by
            // old stride * (old dim - 1)
            // to put the pointer completely in the other end
            if s1 < 0 {
                offset += stride_offset(m - 1, *sr);
            }

            let s_prim = s * s1;

            let d = m / s1.abs() as Ix;
            let r = m % s1.abs() as Ix;
            let m_prim = d + if r > 0 { 1 } else { 0 };

            // Update dimension and stride coordinate
            *dr = m_prim;
            *sr = s_prim as Ix;
        }
        offset
    }
}

fn abs_index(len: Ixs, index: Ixs) -> Ix {
    if index < 0 {
        (len + index) as Ix
    } else { index as Ix }
}

/// Collapse axis `axis` and shift so that only subarray `index` is
/// available.
///
/// **Panics** if `index` is larger than the size of the axis
// FIXME: Move to Dimension trait
pub fn do_sub<A, D: Dimension, P: Copy + PtrExt<Target=A>>(dims: &mut D, ptr: &mut P, strides: &D,
                           axis: usize, index: Ix)
{
    let dim = dims.slice()[axis];
    let stride = strides.slice()[axis];
    assert!(index < dim);
    dims.slice_mut()[axis] = 1;
    let off = stride_offset(index, stride);
    unsafe {
        *ptr = ptr.offset(off);
    }
}

unsafe impl Dimension for () {
    // empty product is 1 -> size is 1
    #[inline]
    fn ndim(&self) -> usize { 0 }
    fn slice(&self) -> &[Ix] { &[] }
    fn slice_mut(&mut self) -> &mut [Ix] { &mut [] }
}

unsafe impl Dimension for Ix {
    #[inline]
    fn ndim(&self) -> usize { 1 }
    #[inline]
    fn size(&self) -> usize { *self as usize }
    #[inline]
    fn first_index(&self) -> Option<Ix> {
        if *self != 0 {
            Some(0)
        } else { None }
    }
    #[inline]
    fn next_for(&self, mut index: Ix) -> Option<Ix> {
        index += 1;
        if index < *self {
            Some(index)
        } else { None }
    }

    /// Self is an index, return the stride offset
    #[inline]
    fn stride_offset(index: &Ix, stride: &Ix) -> isize
    {
        stride_offset(*index, *stride)
    }

    /// Return stride offset for this dimension and index.
    fn stride_offset_checked(&self, stride: &Ix, index: &Ix) -> Option<isize>
    {
        if *index < *self {
            Some(stride_offset(*index, *stride))
        } else {
            None
        }
    }
}

unsafe impl Dimension for (Ix, Ix) {
    #[inline]
    fn ndim(&self) -> usize { 2 }
    #[inline]
    fn size(&self) -> usize { let (m, n) = *self; m as usize * n as usize }
    #[inline]
    fn first_index(&self) -> Option<(Ix, Ix)> {
        let (m, n) = *self;
        if m != 0 && n != 0 {
            Some((0, 0))
        } else { None }
    }
    #[inline]
    fn next_for(&self, index: (Ix, Ix)) -> Option<(Ix, Ix)> {
        let (mut i, mut j) = index;
        let (imax, jmax) = *self;
        j += 1;
        if j == jmax {
            j = 0;
            i += 1;
            if i == imax {
                return None;
            }
        }
        Some((i, j))
    }

    /// Self is an index, return the stride offset
    #[inline]
    fn stride_offset(index: &(Ix, Ix), strides: &(Ix, Ix)) -> isize
    {
        let (i, j) = *index;
        let (s, t) = *strides;
        stride_offset(i, s) + stride_offset(j, t)
    }

    /// Return stride offset for this dimension and index.
    fn stride_offset_checked(&self, strides: &(Ix, Ix), index: &(Ix, Ix)) -> Option<isize>
    {
        let (m, n) = *self;
        let (i, j) = *index;
        let (s, t) = *strides;
        if i < m && j < n {
            Some(stride_offset(i, s) + stride_offset(j, t))
        } else {
            None
        }
    }
}

unsafe impl Dimension for (Ix, Ix, Ix) {
    #[inline]
    fn ndim(&self) -> usize { 3 }
    #[inline]
    fn size(&self) -> usize { let (m, n, o) = *self; m as usize * n as usize * o as usize }
    #[inline]
    fn next_for(&self, index: (Ix, Ix, Ix)) -> Option<(Ix, Ix, Ix)> {
        let (mut i, mut j, mut k) = index;
        let (imax, jmax, kmax) = *self;
        k += 1;
        if k == kmax {
            k = 0;
            j += 1;
            if j == jmax {
                j = 0;
                i += 1;
                if i == imax {
                    return None;
                }
            }
        }
        Some((i, j, k))
    }

    /// Self is an index, return the stride offset
    #[inline]
    fn stride_offset(index: &(Ix, Ix, Ix), strides: &(Ix, Ix, Ix)) -> isize
    {
        let (i, j, k) = *index;
        let (s, t, u) = *strides;
        stride_offset(i, s) + stride_offset(j, t) + stride_offset(k, u)
    }
}

unsafe impl Dimension for (Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 4 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 5 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 6 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 7 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 8 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 9 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 10 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 11 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 12 } }

// Vec<Ix> is a "dynamic" index, pretty hard to use when indexing,
// and memory wasteful, but it allows an arbitrary number of dimensions.
//
// NOTE: No Shrink impl for Vec<Ix> yet.
unsafe impl Dimension for Vec<Ix>
{
    fn ndim(&self) -> usize { self.len() }
    fn slice(&self) -> &[Ix] { self }
    fn slice_mut(&mut self) -> &mut [Ix] { self }
}

/// Helper trait to define a larger-than relation for array shapes:
/// removing one axis from *Self* gives smaller dimension *Smaller*.
pub trait RemoveAxis : Dimension {
    type Smaller: Dimension;
    fn remove_axis(&self, axis: usize) -> Self::Smaller;
}

macro_rules! impl_shrink(
    ($from:ident, $($more:ident,)*) => (
impl RemoveAxis for ($from $(,$more)*)
{
    type Smaller = ($($more),*);
    #[allow(unused_parens)]
    fn remove_axis(&self, axis: usize) -> ($($more),*) {
        let mut tup = ($(0 as $more),*);
        {
            let mut it = tup.slice_mut().iter_mut();
            for (i, &d) in self.slice().iter().enumerate() {
                if i == axis {
                    continue;
                }
                for rr in it.by_ref() {
                    *rr = d;
                    break
                }
            }
        }
        tup
    }
}
    )
);

macro_rules! impl_shrink_recursive(
    ($ix:ident, ) => (impl_shrink!($ix,););
    ($ix1:ident, $($ix:ident,)*) => (
        impl_shrink_recursive!($($ix,)*);
        impl_shrink!($ix1, $($ix,)*);
    )
);

// 12 is the maximum number for having the Eq trait from libstd
impl_shrink_recursive!(Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix,);

impl RemoveAxis for Vec<Ix>
{
    type Smaller = Vec<Ix>;
    fn remove_axis(&self, axis: usize) -> Vec<Ix>
    {
        let mut res = self.clone();
        res.remove(axis);
        res
    }
}

// [a:b:s] syntax for example [:3], [::-1]
// [0,:] -- first row of matrix
// [:,0] -- first column of matrix

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
/// A slice, a description of a range of an array axis.
///
/// Fields are `begin`, `end` and `stride`, where
/// negative `begin` or `end` indexes are counted from the back
/// of the axis.
///
/// If `end` is `None`, the slice extends to the end of the axis.
///
/// ## Examples
///
/// `Si(0, None, 1)` is the full range of an axis.
/// Python equivalent is `[:]`.
///
/// `Si(a, Some(b), 2)` is every second element from `a` until `b`.
/// Python equivalent is `[a:b:2]`.
///
/// `Si(a, None, -1)` is every element, in reverse order, from `a`
/// until the end. Python equivalent is `[a::-1]`
pub struct Si(pub Ixs, pub Option<Ixs>, pub Ixs);

impl Copy for Si {}

/// Slice value for the full range of an axis.
pub const S: Si = Si(0, None, 1);

<commit_msg>Fix use of PtrExt<commit_after>use std::mem;
use std::num::SignedInt;
use std::raw;

use super::{Ix, Ixs};

/// Calculate offset from `Ix` stride converting sign properly
#[inline]
pub fn stride_offset(n: Ix, stride: Ix) -> isize
{
    (n as isize) * ((stride as Ixs) as isize)
}

/*
#[inline]
pub fn stride_as_int(stride: Ix) -> isize
{
    (stride as Ixs) as isize
}
*/

/// Trait for the shape and index types of arrays.
///
/// unsafe trait due to how the assumptions in the default impls work.
pub unsafe trait Dimension : Clone + Eq {
    fn ndim(&self) -> usize;
    fn slice<'a>(&'a self) -> &'a [Ix] {
        unsafe {
            mem::transmute(raw::Slice {
                data: self as *const _ as *const Ix,
                len: self.ndim(),
            })
        }
    }

    fn slice_mut<'a>(&'a mut self) -> &'a mut [Ix] {
        unsafe {
            mem::transmute(raw::Slice {
                data: self as *mut _ as *const Ix,
                len: self.ndim(),
            })
        }
    }

    fn size(&self) -> usize {
        self.slice().iter().fold(1, |s, &a| s * a as usize)
    }

    fn default_strides(&self) -> Self {
        // Compute default array strides
        // Shape (a, b, c) => Give strides (b * c, c, 1)
        let mut strides = self.clone();
        {
            let mut it = strides.slice_mut().iter_mut().rev();
            // Set first element to 1
            for rs in it.by_ref() {
                *rs = 1;
                break;
            }
            let mut cum_prod = 1;
            for (rs, dim) in it.zip(self.slice().iter().rev()) {
                cum_prod *= *dim;
                *rs = cum_prod;
            }
        }
        strides
    }

    #[inline]
    fn first_index(&self) -> Option<Self>
    {
        for ax in self.slice().iter() {
            if *ax == 0 {
                return None
            }
        }
        let mut index = self.clone();
        for rr in index.slice_mut().iter_mut() {
            *rr = 0;
        }
        Some(index)
    }

    /// Iteration -- Use self as size, and return next index after `index`
    /// or None if there are no more.
    // FIXME: use &Self for index or even &mut?
    #[inline]
    fn next_for(&self, index: Self) -> Option<Self> {
        let mut index = index;
        let mut done = false;
        for (&dim, ix) in self.slice().iter().rev()
                            .zip(index.slice_mut().iter_mut().rev())
        {
            *ix += 1;
            if *ix == dim {
                *ix = 0;
            } else {
                done = true;
                break;
            }
        }
        if done {
            Some(index)
        } else { None }
    }

    /// Return stride offset for index.
    fn stride_offset(index: &Self, strides: &Self) -> isize
    {
        let mut offset = 0;
        for (&i, &s) in index.slice().iter()
                            .zip(strides.slice().iter()) {
            offset += stride_offset(i, s);
        }
        offset
    }

    /// Return stride offset for this dimension and index.
    fn stride_offset_checked(&self, strides: &Self, index: &Self) -> Option<isize>
    {
        let mut offset = 0;
        for ((&d, &i), &s) in self.slice().iter()
                                .zip(index.slice().iter())
                                .zip(strides.slice().iter())
        {
            if i >= d {
                return None;
            }
            offset += stride_offset(i, s);
        }
        Some(offset)
    }

    /// Modify dimension, strides and return data pointer offset
    ///
    /// **Panics** if `slices` does not correspond to the number of axes,
    /// if any stride is 0, or if any index is out of bounds.
    fn do_slices(dim: &mut Self, strides: &mut Self, slices: &[Si]) -> isize
    {
        let mut offset = 0;
        assert!(slices.len() == dim.slice().len());
        for ((dr, sr), &slc) in dim.slice_mut().iter_mut()
                                .zip(strides.slice_mut().iter_mut())
                                .zip(slices.iter())
        {
            let m = *dr;
            let mi = m as Ixs;
            let Si(b1, opt_e1, s1) = slc;
            let e1 = opt_e1.unwrap_or(mi);

            let b1 = abs_index(mi, b1);
            let mut e1 = abs_index(mi, e1);
            if e1 < b1 { e1 = b1; }

            assert!(b1 <= m);
            assert!(e1 <= m);

            let m = e1 - b1;
            // stride
            let s = (*sr) as Ixs;

            // Data pointer offset
            offset += stride_offset(b1, *sr);
            // Adjust for strides
            assert!(s1 != 0);
            // How to implement negative strides:
            //
            // Increase start pointer by
            // old stride * (old dim - 1)
            // to put the pointer completely in the other end
            if s1 < 0 {
                offset += stride_offset(m - 1, *sr);
            }

            let s_prim = s * s1;

            let d = m / s1.abs() as Ix;
            let r = m % s1.abs() as Ix;
            let m_prim = d + if r > 0 { 1 } else { 0 };

            // Update dimension and stride coordinate
            *dr = m_prim;
            *sr = s_prim as Ix;
        }
        offset
    }
}

fn abs_index(len: Ixs, index: Ixs) -> Ix {
    if index < 0 {
        (len + index) as Ix
    } else { index as Ix }
}

/// Collapse axis `axis` and shift so that only subarray `index` is
/// available.
///
/// **Panics** if `index` is larger than the size of the axis
// FIXME: Move to Dimension trait
pub fn do_sub<A, D: Dimension>(dims: &mut D, ptr: &mut *mut A, strides: &D,
                               axis: usize, index: Ix)
{
    let dim = dims.slice()[axis];
    let stride = strides.slice()[axis];
    assert!(index < dim);
    dims.slice_mut()[axis] = 1;
    let off = stride_offset(index, stride);
    unsafe {
        *ptr = ptr.offset(off);
    }
}

unsafe impl Dimension for () {
    // empty product is 1 -> size is 1
    #[inline]
    fn ndim(&self) -> usize { 0 }
    fn slice(&self) -> &[Ix] { &[] }
    fn slice_mut(&mut self) -> &mut [Ix] { &mut [] }
}

unsafe impl Dimension for Ix {
    #[inline]
    fn ndim(&self) -> usize { 1 }
    #[inline]
    fn size(&self) -> usize { *self as usize }
    #[inline]
    fn first_index(&self) -> Option<Ix> {
        if *self != 0 {
            Some(0)
        } else { None }
    }
    #[inline]
    fn next_for(&self, mut index: Ix) -> Option<Ix> {
        index += 1;
        if index < *self {
            Some(index)
        } else { None }
    }

    /// Self is an index, return the stride offset
    #[inline]
    fn stride_offset(index: &Ix, stride: &Ix) -> isize
    {
        stride_offset(*index, *stride)
    }

    /// Return stride offset for this dimension and index.
    fn stride_offset_checked(&self, stride: &Ix, index: &Ix) -> Option<isize>
    {
        if *index < *self {
            Some(stride_offset(*index, *stride))
        } else {
            None
        }
    }
}

unsafe impl Dimension for (Ix, Ix) {
    #[inline]
    fn ndim(&self) -> usize { 2 }
    #[inline]
    fn size(&self) -> usize { let (m, n) = *self; m as usize * n as usize }
    #[inline]
    fn first_index(&self) -> Option<(Ix, Ix)> {
        let (m, n) = *self;
        if m != 0 && n != 0 {
            Some((0, 0))
        } else { None }
    }
    #[inline]
    fn next_for(&self, index: (Ix, Ix)) -> Option<(Ix, Ix)> {
        let (mut i, mut j) = index;
        let (imax, jmax) = *self;
        j += 1;
        if j == jmax {
            j = 0;
            i += 1;
            if i == imax {
                return None;
            }
        }
        Some((i, j))
    }

    /// Self is an index, return the stride offset
    #[inline]
    fn stride_offset(index: &(Ix, Ix), strides: &(Ix, Ix)) -> isize
    {
        let (i, j) = *index;
        let (s, t) = *strides;
        stride_offset(i, s) + stride_offset(j, t)
    }

    /// Return stride offset for this dimension and index.
    fn stride_offset_checked(&self, strides: &(Ix, Ix), index: &(Ix, Ix)) -> Option<isize>
    {
        let (m, n) = *self;
        let (i, j) = *index;
        let (s, t) = *strides;
        if i < m && j < n {
            Some(stride_offset(i, s) + stride_offset(j, t))
        } else {
            None
        }
    }
}

unsafe impl Dimension for (Ix, Ix, Ix) {
    #[inline]
    fn ndim(&self) -> usize { 3 }
    #[inline]
    fn size(&self) -> usize { let (m, n, o) = *self; m as usize * n as usize * o as usize }
    #[inline]
    fn next_for(&self, index: (Ix, Ix, Ix)) -> Option<(Ix, Ix, Ix)> {
        let (mut i, mut j, mut k) = index;
        let (imax, jmax, kmax) = *self;
        k += 1;
        if k == kmax {
            k = 0;
            j += 1;
            if j == jmax {
                j = 0;
                i += 1;
                if i == imax {
                    return None;
                }
            }
        }
        Some((i, j, k))
    }

    /// Self is an index, return the stride offset
    #[inline]
    fn stride_offset(index: &(Ix, Ix, Ix), strides: &(Ix, Ix, Ix)) -> isize
    {
        let (i, j, k) = *index;
        let (s, t, u) = *strides;
        stride_offset(i, s) + stride_offset(j, t) + stride_offset(k, u)
    }
}

unsafe impl Dimension for (Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 4 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 5 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 6 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 7 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 8 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 9 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 10 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 11 } }
unsafe impl Dimension for (Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix) { fn ndim(&self) -> usize { 12 } }

// Vec<Ix> is a "dynamic" index, pretty hard to use when indexing,
// and memory wasteful, but it allows an arbitrary number of dimensions.
//
// NOTE: No Shrink impl for Vec<Ix> yet.
unsafe impl Dimension for Vec<Ix>
{
    fn ndim(&self) -> usize { self.len() }
    fn slice(&self) -> &[Ix] { self }
    fn slice_mut(&mut self) -> &mut [Ix] { self }
}

/// Helper trait to define a larger-than relation for array shapes:
/// removing one axis from *Self* gives smaller dimension *Smaller*.
pub trait RemoveAxis : Dimension {
    type Smaller: Dimension;
    fn remove_axis(&self, axis: usize) -> Self::Smaller;
}

macro_rules! impl_shrink(
    ($from:ident, $($more:ident,)*) => (
impl RemoveAxis for ($from $(,$more)*)
{
    type Smaller = ($($more),*);
    #[allow(unused_parens)]
    fn remove_axis(&self, axis: usize) -> ($($more),*) {
        let mut tup = ($(0 as $more),*);
        {
            let mut it = tup.slice_mut().iter_mut();
            for (i, &d) in self.slice().iter().enumerate() {
                if i == axis {
                    continue;
                }
                for rr in it.by_ref() {
                    *rr = d;
                    break
                }
            }
        }
        tup
    }
}
    )
);

macro_rules! impl_shrink_recursive(
    ($ix:ident, ) => (impl_shrink!($ix,););
    ($ix1:ident, $($ix:ident,)*) => (
        impl_shrink_recursive!($($ix,)*);
        impl_shrink!($ix1, $($ix,)*);
    )
);

// 12 is the maximum number for having the Eq trait from libstd
impl_shrink_recursive!(Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix, Ix,);

impl RemoveAxis for Vec<Ix>
{
    type Smaller = Vec<Ix>;
    fn remove_axis(&self, axis: usize) -> Vec<Ix>
    {
        let mut res = self.clone();
        res.remove(axis);
        res
    }
}

// [a:b:s] syntax for example [:3], [::-1]
// [0,:] -- first row of matrix
// [:,0] -- first column of matrix

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
/// A slice, a description of a range of an array axis.
///
/// Fields are `begin`, `end` and `stride`, where
/// negative `begin` or `end` indexes are counted from the back
/// of the axis.
///
/// If `end` is `None`, the slice extends to the end of the axis.
///
/// ## Examples
///
/// `Si(0, None, 1)` is the full range of an axis.
/// Python equivalent is `[:]`.
///
/// `Si(a, Some(b), 2)` is every second element from `a` until `b`.
/// Python equivalent is `[a:b:2]`.
///
/// `Si(a, None, -1)` is every element, in reverse order, from `a`
/// until the end. Python equivalent is `[a::-1]`
pub struct Si(pub Ixs, pub Option<Ixs>, pub Ixs);

impl Copy for Si {}

/// Slice value for the full range of an axis.
pub const S: Si = Si(0, None, 1);

<|endoftext|>


<commit_before>extern crate sfunc;

use std::rand::Rng;
use std::rand::distributions::{Gamma, IndependentSample};

use super::Distribution;

/// A beta distribution.
pub struct Beta {
    /// The first shape parameter.
    pub alpha: f64,
    /// The second shape parameter.
    pub beta: f64,
    /// The left endpoint of the support.
    pub a: f64,
    /// The right endpoint of the support.
    pub b: f64,

    gamma_alpha: Gamma,
    gamma_beta: Gamma,
}

impl Beta {
    /// Creates a beta distribution with the shape parameters `alpha` and
    /// `beta` on the interval `[a, b]`.
    #[inline]
    pub fn new(alpha: f64, beta: f64, a: f64, b: f64) -> Beta {
        Beta {
            alpha: alpha,
            beta: beta,
            a: a,
            b: b,
            gamma_alpha: Gamma::new(alpha, 1.0),
            gamma_beta: Gamma::new(beta, 1.0),
        }
    }
}

impl Distribution<f64> for Beta {
    fn cdf(&self, x: f64) -> f64 {
        use self::sfunc::{inc_beta, ln_beta};
        inc_beta((x - self.a) / (self.b - self.a), self.alpha, self.beta,
                 ln_beta(self.alpha, self.beta))
    }

    #[inline]
    fn inv_cdf(&self, p: f64) -> f64 {
        use self::sfunc::{inv_inc_beta, ln_beta};
        self.a + (self.b - self.a) * inv_inc_beta(p, self.alpha, self.beta,
                                                  ln_beta(self.alpha, self.beta))
    }

    #[inline]
    fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let x = self.gamma_alpha.ind_sample(rng);
        let y = self.gamma_beta.ind_sample(rng);
        self.a + (self.b - self.a) * x / (x + y)
    }
}

#[cfg(test)]
mod test {
    use super::super::{Distribution, Sampler};
    use super::Beta;

    macro_rules! assert_almost_eq(
        ($x:expr, $y:expr) => ({
            let e: f64 = ::std::f64::EPSILON.sqrt();
            let x: Vec<f64> = $x;
            let y: Vec<f64> = $y;
            for i in range(0u, x.len()) {
                assert!(::std::num::abs(x[i] - y[i]) < e,
                        "expected {:e} ~ {:e}", x[i], y[i]);
            }
        });
    )

    #[test]
    fn cdf() {
        let dist = Beta::new(2.0, 3.0, -1.0, 2.0);

        let x = vec![-1.00, -0.85, -0.70, -0.55, -0.40, -0.25, -0.10, 0.05,
                     0.20, 0.35, 0.50, 0.65, 0.80, 0.95, 1.10, 1.25, 1.40,
                     1.55, 1.70, 1.85, 2.00];
        let p = vec![0.000000000000000e+00, 1.401875000000000e-02,
                     5.230000000000002e-02, 1.095187500000000e-01,
                     1.807999999999999e-01, 2.617187500000001e-01,
                     3.483000000000000e-01, 4.370187500000001e-01,
                     5.248000000000003e-01, 6.090187500000001e-01,
                     6.875000000000000e-01, 7.585187500000001e-01,
                     8.208000000000000e-01, 8.735187499999999e-01,
                     9.163000000000000e-01, 9.492187500000000e-01,
                     9.728000000000000e-01, 9.880187500000001e-01,
                     9.963000000000000e-01, 9.995187500000000e-01,
                     1.000000000000000e+00];

        assert_almost_eq!(x.iter().map(|&x| dist.cdf(x)).collect(), p);
    }

    #[test]
    fn inv_cdf() {
        let dist = Beta::new(1.0, 2.0, 3.0, 4.0);

        let p = vec![0.00, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40,
                     0.45, 0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85,
                     0.90, 0.95, 1.00];
        let x = vec![3.000000000000000e+00, 3.025320565519104e+00,
                     3.051316701949486e+00, 3.078045554270711e+00,
                     3.105572809000084e+00, 3.133974596215561e+00,
                     3.163339973465924e+00, 3.193774225170145e+00,
                     3.225403330758517e+00, 3.258380151290432e+00,
                     3.292893218813452e+00, 3.329179606750063e+00,
                     3.367544467966324e+00, 3.408392021690038e+00,
                     3.452277442494834e+00, 3.500000000000000e+00,
                     3.552786404500042e+00, 3.612701665379257e+00,
                     3.683772233983162e+00, 3.776393202250021e+00,
                     4.000000000000000e+00];

        assert_almost_eq!(p.iter().map(|&p| dist.inv_cdf(p)).collect(), x);
    }

    #[test]
    fn sample() {
        let mut rng = ::std::rand::task_rng();
        for x in Sampler(&Beta::new(1.0, 2.0, 7.0, 42.0), &mut rng).take(100) {
            assert!(7.0 <= x && x <= 42.0);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use super::super::{Distribution, Sampler, Uniform};
    use super::Beta;

    #[bench]
    fn cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Beta::new(0.5, 1.5, 0.0, 1.0);
        let x = Sampler(&Uniform::new(0.0, 1.0), &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(x.iter().map(|&x| dist.cdf(x)).collect::<Vec<_>>())
        });
    }

    #[bench]
    fn inv_cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Beta::new(0.5, 1.5, 0.0, 1.0);
        let p = Sampler(&Uniform::new(0.0, 1.0), &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(p.iter().map(|&p| dist.inv_cdf(p)).collect::<Vec<_>>())
        });
    }
}
<commit_msg>Improved cdf and inv_cdf of Beta<commit_after>extern crate sfunc;

use std::rand::Rng;
use std::rand::distributions::{Gamma, IndependentSample};

use super::Distribution;

/// A beta distribution.
pub struct Beta {
    /// The first shape parameter.
    pub alpha: f64,
    /// The second shape parameter.
    pub beta: f64,
    /// The left endpoint of the support.
    pub a: f64,
    /// The right endpoint of the support.
    pub b: f64,

    ln_beta: f64,
    gamma_alpha: Gamma,
    gamma_beta: Gamma,
}

impl Beta {
    /// Creates a beta distribution with the shape parameters `alpha` and
    /// `beta` on the interval `[a, b]`.
    #[inline]
    pub fn new(alpha: f64, beta: f64, a: f64, b: f64) -> Beta {
        Beta {
            alpha: alpha,
            beta: beta,
            a: a,
            b: b,
            ln_beta: sfunc::ln_beta(alpha, beta),
            gamma_alpha: Gamma::new(alpha, 1.0),
            gamma_beta: Gamma::new(beta, 1.0),
        }
    }
}

impl Distribution<f64> for Beta {
    fn cdf(&self, x: f64) -> f64 {
        use self::sfunc::inc_beta;
        inc_beta((x - self.a) / (self.b - self.a), self.alpha, self.beta, self.ln_beta)
    }

    #[inline]
    fn inv_cdf(&self, p: f64) -> f64 {
        use self::sfunc::inv_inc_beta;
        self.a + (self.b - self.a) * inv_inc_beta(p, self.alpha, self.beta, self.ln_beta)
    }

    #[inline]
    fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let x = self.gamma_alpha.ind_sample(rng);
        let y = self.gamma_beta.ind_sample(rng);
        self.a + (self.b - self.a) * x / (x + y)
    }
}

#[cfg(test)]
mod test {
    use super::super::{Distribution, Sampler};
    use super::Beta;

    macro_rules! assert_almost_eq(
        ($x:expr, $y:expr) => ({
            let e: f64 = ::std::f64::EPSILON.sqrt();
            let x: Vec<f64> = $x;
            let y: Vec<f64> = $y;
            for i in range(0u, x.len()) {
                assert!(::std::num::abs(x[i] - y[i]) < e,
                        "expected {:e} ~ {:e}", x[i], y[i]);
            }
        });
    )

    #[test]
    fn cdf() {
        let dist = Beta::new(2.0, 3.0, -1.0, 2.0);

        let x = vec![-1.00, -0.85, -0.70, -0.55, -0.40, -0.25, -0.10, 0.05,
                     0.20, 0.35, 0.50, 0.65, 0.80, 0.95, 1.10, 1.25, 1.40,
                     1.55, 1.70, 1.85, 2.00];
        let p = vec![0.000000000000000e+00, 1.401875000000000e-02,
                     5.230000000000002e-02, 1.095187500000000e-01,
                     1.807999999999999e-01, 2.617187500000001e-01,
                     3.483000000000000e-01, 4.370187500000001e-01,
                     5.248000000000003e-01, 6.090187500000001e-01,
                     6.875000000000000e-01, 7.585187500000001e-01,
                     8.208000000000000e-01, 8.735187499999999e-01,
                     9.163000000000000e-01, 9.492187500000000e-01,
                     9.728000000000000e-01, 9.880187500000001e-01,
                     9.963000000000000e-01, 9.995187500000000e-01,
                     1.000000000000000e+00];

        assert_almost_eq!(x.iter().map(|&x| dist.cdf(x)).collect(), p);
    }

    #[test]
    fn inv_cdf() {
        let dist = Beta::new(1.0, 2.0, 3.0, 4.0);

        let p = vec![0.00, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40,
                     0.45, 0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85,
                     0.90, 0.95, 1.00];
        let x = vec![3.000000000000000e+00, 3.025320565519104e+00,
                     3.051316701949486e+00, 3.078045554270711e+00,
                     3.105572809000084e+00, 3.133974596215561e+00,
                     3.163339973465924e+00, 3.193774225170145e+00,
                     3.225403330758517e+00, 3.258380151290432e+00,
                     3.292893218813452e+00, 3.329179606750063e+00,
                     3.367544467966324e+00, 3.408392021690038e+00,
                     3.452277442494834e+00, 3.500000000000000e+00,
                     3.552786404500042e+00, 3.612701665379257e+00,
                     3.683772233983162e+00, 3.776393202250021e+00,
                     4.000000000000000e+00];

        assert_almost_eq!(p.iter().map(|&p| dist.inv_cdf(p)).collect(), x);
    }

    #[test]
    fn sample() {
        let mut rng = ::std::rand::task_rng();
        for x in Sampler(&Beta::new(1.0, 2.0, 7.0, 42.0), &mut rng).take(100) {
            assert!(7.0 <= x && x <= 42.0);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use super::super::{Distribution, Sampler, Uniform};
    use super::Beta;

    #[bench]
    fn cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Beta::new(0.5, 1.5, 0.0, 1.0);
        let x = Sampler(&Uniform::new(0.0, 1.0), &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(x.iter().map(|&x| dist.cdf(x)).collect::<Vec<_>>())
        });
    }

    #[bench]
    fn inv_cdf(bench: &mut test::Bencher) {
        let mut rng = ::std::rand::task_rng();
        let dist = Beta::new(0.5, 1.5, 0.0, 1.0);
        let p = Sampler(&Uniform::new(0.0, 1.0), &mut rng).take(1000).collect::<Vec<_>>();

        bench.iter(|| {
            test::black_box(p.iter().map(|&p| dist.inv_cdf(p)).collect::<Vec<_>>())
        });
    }
}
<|endoftext|>


<commit_before>use std::slice;
use num::Integer;
use std::num::FromPrimitive;
use std::f32::consts::PI;
use calx::V2;

/// Hex grid directions.
#[deriving(Eq, PartialEq, Clone, Show, FromPrimitive, Encodable, Decodable)]
pub enum Dir6 {
    North = 0,
    NorthEast,
    SouthEast,
    South,
    SouthWest,
    NorthWest,
}

impl Dir6 {
    /// Convert a vector into the closest hex direction.
    pub fn from_v2(v: V2<int>) -> Dir6 {
        let hexadecant = {
            let width = PI / 8.0;
            let mut radian = (v.0 as f32).atan2(-v.1 as f32);
            if radian < 0.0 { radian += 2.0 * PI }
            (radian / width).floor() as int
        };

        Dir6::from_int(match hexadecant {
            13 | 14 => 0,
            15 | 0 | 1 => 1,
            2 | 3 | 4 => 2,
            5 | 6 => 3,
            7 | 8 | 9 => 4,
            10 | 11 | 12 => 5,
            _ => fail!("Bad hexadecant")
        })
    }

    /// Convert an integer to a hex dir using modular arithmetic.
    pub fn from_int(i: int) -> Dir6 {
        FromPrimitive::from_int(i.mod_floor(&6)).unwrap()
    }

    /// Convert a hex dir into the corresponding unit vector.
    pub fn to_v2(&self) -> V2<int> {
        [V2(-1, -1),
         V2( 0, -1),
         V2( 1,  0),
         V2( 1,  1),
         V2( 0,  1),
         V2(-1,  0)][*self as uint]
    }

    /// Iterate through the six hex dirs in the standard order.
    pub fn iter() -> slice::Items<'static, Dir6> {
        static DIRS: [Dir6, ..6] = [
            North,
            NorthEast,
            SouthEast,
            South,
            SouthWest,
            NorthWest];

        DIRS.iter()
    }
}

#[cfg(test)]
mod test {
    use calx::V2;
    // XXX: Why doesn't super::* work here?
    use super::{Dir6, North, NorthEast, SouthEast, South, SouthWest, NorthWest};

    #[test]
    fn test_dir6() {
        assert_eq!(North, Dir6::from_int(0));
        assert_eq!(NorthWest, Dir6::from_int(-1));
        assert_eq!(NorthWest, Dir6::from_int(5));
        assert_eq!(NorthEast, Dir6::from_int(1));

        for i in range(0, 6) {
            let d = Dir6::from_int(i);
            let v = d.to_v2();
            let v1 = Dir6::from_int(i - 1).to_v2();
            let v2 = Dir6::from_int(i + 1).to_v2();
            assert_eq!(d, Dir6::from_v2(v));
            assert_eq!(Dir6::from_int(i + 3), Dir6::from_v2(-v));
            assert_eq!(d, Dir6::from_v2(v * 6));
            // Test deviated longer vecs.
            assert_eq!(d, Dir6::from_v2(v * 3 + v1));
            assert_eq!(d, Dir6::from_v2(v * 3 + v2));
        }
        /*
        assert_eq!(NorthEast, Dir6::from_v2(V2(20, -21)));
        assert_eq!(SouthEast, Dir6::from_v2(V2(20, -10)));
        assert_eq!(North, Dir6::from_v2(V2(-10, -10)));
        assert_eq!(South, Dir6::from_v2(V2(1, 1)));
        */
    }
}
<commit_msg>Dir6 commenting and unit tests<commit_after>use std::slice;
use num::Integer;
use std::num::FromPrimitive;
use std::f32::consts::PI;
use calx::V2;

/// Hex grid directions.
#[deriving(Eq, PartialEq, Clone, Show, FromPrimitive, Encodable, Decodable)]
pub enum Dir6 {
    North = 0,
    NorthEast,
    SouthEast,
    South,
    SouthWest,
    NorthWest,
}

impl Dir6 {
    /// Convert a vector into the closest hex direction.
    pub fn from_v2(v: V2<int>) -> Dir6 {
        let hexadecant = {
            let width = PI / 8.0;
            let mut radian = (v.0 as f32).atan2(-v.1 as f32);
            if radian < 0.0 { radian += 2.0 * PI }
            (radian / width).floor() as int
        };

        // The hexadecants (00 to 15) and the hex directions (*0* to *5*) around
        // the origin.
        //
        //    *0*       *1*
        //       \ 14 15 | 00 01
        //       13\     |      02
        //           \   |
        //     12      \ |        03
        // *5* ----------O-X------- *2*
        //     11        Y \      04
        //               |   \
        //       10      |     \05
        //         09 08 | 07 06 \
        //              *4*       *3*
        //
        // Vectors that are in a space between two hex direction vectors are
        // snapped to the hex direction vector nearer to them. The hexadecants
        // are a convenience structure that lets us assign every vector using
        // the single hexadecant value.
        Dir6::from_int(match hexadecant {
            13 | 14 => 0,
            15 | 0 | 1 => 1,
            2 | 3 | 4 => 2,
            5 | 6 => 3,
            7 | 8 | 9 => 4,
            10 | 11 | 12 => 5,
            _ => fail!("Bad hexadecant")
        })
    }

    /// Convert an integer to a hex dir using modular arithmetic.
    pub fn from_int(i: int) -> Dir6 {
        FromPrimitive::from_int(i.mod_floor(&6)).unwrap()
    }

    /// Convert a hex dir into the corresponding unit vector.
    pub fn to_v2(&self) -> V2<int> {
        [V2(-1, -1),
         V2( 0, -1),
         V2( 1,  0),
         V2( 1,  1),
         V2( 0,  1),
         V2(-1,  0)][*self as uint]
    }

    /// Iterate through the six hex dirs in the standard order.
    pub fn iter() -> slice::Items<'static, Dir6> {
        static DIRS: [Dir6, ..6] = [
            North,
            NorthEast,
            SouthEast,
            South,
            SouthWest,
            NorthWest];

        DIRS.iter()
    }
}

#[cfg(test)]
mod test {
    use calx::V2;
    // XXX: Why doesn't super::* work here?
    use super::{Dir6, North, NorthEast, SouthEast, South, SouthWest, NorthWest};

    #[test]
    fn test_dir6() {
        assert_eq!(North, Dir6::from_int(0));
        assert_eq!(NorthWest, Dir6::from_int(-1));
        assert_eq!(NorthWest, Dir6::from_int(5));
        assert_eq!(NorthEast, Dir6::from_int(1));

        assert_eq!(NorthEast, Dir6::from_v2(V2(20, -21)));
        assert_eq!(SouthEast, Dir6::from_v2(V2(20, -10)));
        assert_eq!(North, Dir6::from_v2(V2(-10, -10)));
        assert_eq!(South, Dir6::from_v2(V2(1, 1)));

        for i in range(0, 6) {
            let d = Dir6::from_int(i);
            let v = d.to_v2();
            let v1 = Dir6::from_int(i - 1).to_v2();
            let v2 = Dir6::from_int(i + 1).to_v2();

            // Test static iter
            assert_eq!(Some(d), Dir6::iter().nth(i as uint).map(|&x| x));

            // Test vector mapping.
            assert_eq!(d, Dir6::from_v2(v));

            // Test opposite dir vector mapping.
            assert_eq!(Dir6::from_int(i + 3), Dir6::from_v2(-v));

            // Test approximation of longer vectors.
            assert_eq!(d, Dir6::from_v2(v * 3));
            assert_eq!(d, Dir6::from_v2(v * 3 + v1));
            assert_eq!(d, Dir6::from_v2(v * 3 + v2));
        }
    }
}
<|endoftext|>


<commit_before>use nom::{IResult, anychar, is_space};
use std::fmt::Debug;
use std::str::{FromStr, from_utf8};

/// Parses a token (i.e. a terminal symbol or a non-terminal symbol).
/// A *token* can be of one of the following two forms:
///
/// * It is a string containing neither of the symbols `'"'`, `' '`, `'-'`, `'→'`, `','`, `';'`, `')'`, `']'`, `'%'`.
/// * It is delimited by the symbol `'"'` on both sides and each occurrence of `'\\'` or `'"'` inside the delimiters is escaped.
pub fn parse_token<A>(input: &[u8])
        -> IResult<&[u8], A>
    where A: FromStr,
          A::Err: Debug,
{
    named!(
        parse_token_s<&str>,
        map_res!(
            alt!(
                delimited!(
                    tag!("\""),
                    escaped!(is_not!("\"\\"), '\\', one_of!("\\\"")),
                    tag!("\"")
                ) |
                is_not!(" \\\"-→,;)]%#")
            ),
            from_utf8
        )
    );

    do_parse!(
        input,
        output: parse_token_s >>
        token: expr_res!(output.parse()) >>
        (token)
    )
}

/// Parses the `input` into a `Vec<A>` given an `inner_parser` for type `A`, an `opening` delimiter, a `closing` delimiter, and a `separator`.
/// The `inner_parser` must not consume the `separator`s or the `closing` delimiter of the given `input`.
pub fn parse_vec<'a, A, P>(input: &'a [u8], inner_parser: P, opening: &str, closing: &str, separator: &str)
        -> IResult<&'a [u8], Vec<A>>
    where P: Fn(&'a [u8]) -> IResult<&'a [u8], A>
{
    do_parse!(
        input,
        tag!(opening) >>
        take_while!(is_space) >>
        result: many0!(
            do_parse!(
                opt!(tag!(separator)) >>
                take_while!(is_space) >>
                the_token: inner_parser >>
                take_while!(is_space) >>
                (the_token)
            )
        ) >>
        tag!(closing) >>
        (result)
    )
}

/// Parses a string of the form `initials: ⟨token⟩` as a the initial symbol of type `I`.
pub fn parse_initial<I>(input: &[u8])
                         -> IResult<&[u8], I>
where I: FromStr,
      I::Err: Debug,
{
    do_parse!(
        input,
        tag!("initial:") >>
            take_while!(is_space) >>
            result: parse_token >>
            take_while!(|_| true) >>
            (result)
    )
}

/// Parses a string of the form `initials: [...]` as a vector of initial symbols of type `I`.
pub fn parse_initials<I>(input: &[u8])
        -> IResult<&[u8], Vec<I>>
    where I: FromStr,
          I::Err: Debug,
{
    do_parse!(
        input,
        tag!("initial:") >>
            take_while!(is_space) >>
            result: call!(|x| parse_vec(x, parse_token, "[", "]", ",")) >>
            (result)
    )
}

/// Consumes any string that begins with the character `%`.
pub fn parse_comment(input: &[u8]) -> IResult<&[u8], ()> {
    do_parse!(
        input,
        tag!("%") >>
            take_while!(|_| true) >>
            (())
    )
}

/// Parses a string as a grammar of initial symbols of type `I` and rules of type `R`.
/// The syntax of the initials must comply with `parse_initials`; the syntax of the rules is solely
/// determined by their type.
/// If the string contains multiple definitions of initials, then the initials of the grammar are
/// going to be the union of all defined initials.
pub fn initial_rule_grammar_from_str<I, R>(s: &str)
        -> Result<((Vec<I>, Vec<R>)), String>
    where I: FromStr,
          I::Err: Debug,
          R: FromStr,
          String: From<R::Err>,
{
    let mut it = s.lines();
    let mut initial = Vec::new();
    let mut rules: Vec<R> = Vec::new();

    while let Some(l) = it.next() {
        if l.trim_left().starts_with("initial:") {
            match parse_initials(l.as_bytes()) {
                IResult::Done(_, mut result) =>
                    initial.append(&mut result),
                _ =>
                    return Err(format!("Malformed declaration of initial nonterminals: \'{}\'", l))
            }
        } else if !l.is_empty() && !l.trim_left().starts_with("%") {
            rules.push(l.trim().parse()?);
        }
    }

    Ok((initial, rules))
}

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

    #[test]
    fn test_parse_comment_legal_input() {
        let legal_inputs = vec![
            "% some comment",
            "%% another comment",
        ];

        for legal_input in legal_inputs {
            assert_eq!(
                ("".as_bytes(), ()),
                parse_comment(legal_input.as_bytes()).unwrap(),
            )
        }
    }

    #[test]
    fn test_parse_comment_illegal_input() {
        let illegal_inputs = vec![
            " % not a comment",
            "something else % not a comment",
        ];

        for illegal_input in illegal_inputs {
            match parse_comment(illegal_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_token_legal_input() {
        let legal_inputs = vec![
            ("abcxyz", "", String::from("abcxyz")),
            ("abc,xyz", ",xyz", String::from("abc")),
            ("\"abc\"xyz", "xyz", String::from("abc")),
            ("\"a\\\\b\\\"c\"xyz", "xyz", String::from("a\\\\b\\\"c")),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_token::<String>(legal_input.as_bytes()).unwrap()
            );
        }
    }

    #[test]
    fn test_parse_token_illegal_input() {
        let illegal_inputs = vec![
            " xyz",
            "-xyz",
            "→xyz",
            ",xyz",
            ";xyz",
            ")xyz",
            "]xyz",
            "\"\\\"",
            "\"\"\"",
            " \"a\"",
        ];

        for illegal_input in illegal_inputs {
            match parse_token::<String>(illegal_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }

        let illegal_inputs = vec![
            "a",
            "\"a\"",
        ];

        for illegal_input in illegal_inputs {
            match parse_token::<u8>(illegal_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_token_incomplete_input() {
        let incomplete_inputs = vec![
            "\"a",
        ];

        for incomplete_input in incomplete_inputs {
            match parse_token::<String>(incomplete_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Error(_) =>
                    panic!("The input was not handled as incomplete: \'{}\'", incomplete_input),
                IResult::Incomplete(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_vec_legal_input() {
        let legal_inputs = vec![
            ("[]xyz", "xyz", vec![]),
            ("[a, bc, d]xyz", "xyz",
                vec![String::from("a"), String::from("bc"), String::from("d")]),
            ("[  a, b ,c]xyz", "xyz",
                vec![String::from("a"), String::from("b"), String::from("c")]),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_vec(legal_input.as_bytes(), parse_token, "[", "]", ",").unwrap()
            );
        }
    }

    #[test]
    fn test_parse_vec_illegal_input() {
        let illegal_inputs = vec![
            "(0)",
            "[0]",
            "[0, 1; 2)",
            " []",
            "[a)",
        ];

        for illegal_input in illegal_inputs {
            match parse_vec::<u8, _>(illegal_input.as_bytes(), parse_token, "[", ")", ",") {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_vec_incomplete_input() {
        let incomplete_inputs = vec![
            "[a",
            "["
        ];

        for incomplete_input in incomplete_inputs {
            match parse_vec::<String, _>(incomplete_input.as_bytes(), parse_token, "[", "]", ",") {
                IResult::Done(_, _) | IResult::Error(_) =>
                    panic!("The input was not handled as incomplete: \'{}\'", incomplete_input),
                IResult::Incomplete(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_initial_legal_input() {
        let legal_inputs = vec![
            ("initial: xyz % some comment", "", String::from("xyz")),
            ("initial: \"% xy\"", "", String::from("% xy")),
            ("initial:  xyz something else", "", String::from("xyz")),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_initial(legal_input.as_bytes()).unwrap()
            );
        }
    }

    #[test]
    fn test_parse_initials_legal_input() {
        let legal_inputs = vec![
            ("initial: [a]xyz", "xyz", vec![String::from("a")]),
            ("initial:  []xyz", "xyz", vec![]),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_initials(legal_input.as_bytes()).unwrap()
            );
        }
    }

    #[test]
    fn test_parse_initials_illegal_input() {
        let illegal_inputs = vec![
            "initials: []xyz",
            " initial: []xyz",
        ];

        for illegal_input in illegal_inputs {
            match parse_initials::<String>(illegal_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_initials_incomplete_input() {
        let incomplete_inputs = vec![
            "init",
            "initial: [",
        ];

        for incomplete_input in incomplete_inputs {
            match parse_initials::<String>(incomplete_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Error(_) =>
                    panic!("The input was not handled as incomplete: \'{}\'", incomplete_input),
                IResult::Incomplete(_) => (),
            }
        }
    }
}
<commit_msg>src/util/parsing.rs: made test_parse_token_illegal_input succeed<commit_after>use nom::{IResult, is_space};
use std::fmt::Debug;
use std::str::{FromStr, from_utf8};

/// Parses a token (i.e. a terminal symbol or a non-terminal symbol).
/// A *token* can be of one of the following two forms:
///
/// * It is a string containing neither of the symbols `'"'`, `' '`, `'-'`, `'→'`, `','`, `';'`, `')'`, `']'`, `'%'`.
/// * It is delimited by the symbol `'"'` on both sides and each occurrence of `'\\'` or `'"'` inside the delimiters is escaped.
pub fn parse_token<A>(input: &[u8])
        -> IResult<&[u8], A>
    where A: FromStr,
          A::Err: Debug,
{
    named!(
        parse_token_s<&str>,
        map_res!(
            alt!(
                delimited!(
                    tag!("\""),
                    escaped!(is_not!("\"\\"), '\\', one_of!("\\\"")),
                    tag!("\"")
                ) |
                is_not!(" \\\"-→,;)]%#")
            ),
            from_utf8
        )
    );

    do_parse!(
        input,
        output: parse_token_s >>
        token: expr_res!(output.parse()) >>
        (token)
    )
}

/// Parses the `input` into a `Vec<A>` given an `inner_parser` for type `A`, an `opening` delimiter, a `closing` delimiter, and a `separator`.
/// The `inner_parser` must not consume the `separator`s or the `closing` delimiter of the given `input`.
pub fn parse_vec<'a, A, P>(input: &'a [u8], inner_parser: P, opening: &str, closing: &str, separator: &str)
        -> IResult<&'a [u8], Vec<A>>
    where P: Fn(&'a [u8]) -> IResult<&'a [u8], A>
{
    do_parse!(
        input,
        tag!(opening) >>
        take_while!(is_space) >>
        result: many0!(
            do_parse!(
                opt!(tag!(separator)) >>
                take_while!(is_space) >>
                the_token: inner_parser >>
                take_while!(is_space) >>
                (the_token)
            )
        ) >>
        tag!(closing) >>
        (result)
    )
}

/// Parses a string of the form `initials: ⟨token⟩` as a the initial symbol of type `I`.
pub fn parse_initial<I>(input: &[u8])
                         -> IResult<&[u8], I>
where I: FromStr,
      I::Err: Debug,
{
    do_parse!(
        input,
        tag!("initial:") >>
            take_while!(is_space) >>
            result: parse_token >>
            take_while!(|_| true) >>
            (result)
    )
}

/// Parses a string of the form `initials: [...]` as a vector of initial symbols of type `I`.
pub fn parse_initials<I>(input: &[u8])
        -> IResult<&[u8], Vec<I>>
    where I: FromStr,
          I::Err: Debug,
{
    do_parse!(
        input,
        tag!("initial:") >>
            take_while!(is_space) >>
            result: call!(|x| parse_vec(x, parse_token, "[", "]", ",")) >>
            (result)
    )
}

/// Consumes any string that begins with the character `%`.
pub fn parse_comment(input: &[u8]) -> IResult<&[u8], ()> {
    do_parse!(
        input,
        tag!("%") >>
            take_while!(|_| true) >>
            (())
    )
}

/// Parses a string as a grammar of initial symbols of type `I` and rules of type `R`.
/// The syntax of the initials must comply with `parse_initials`; the syntax of the rules is solely
/// determined by their type.
/// If the string contains multiple definitions of initials, then the initials of the grammar are
/// going to be the union of all defined initials.
pub fn initial_rule_grammar_from_str<I, R>(s: &str)
        -> Result<((Vec<I>, Vec<R>)), String>
    where I: FromStr,
          I::Err: Debug,
          R: FromStr,
          String: From<R::Err>,
{
    let mut it = s.lines();
    let mut initial = Vec::new();
    let mut rules: Vec<R> = Vec::new();

    while let Some(l) = it.next() {
        if l.trim_left().starts_with("initial:") {
            match parse_initials(l.as_bytes()) {
                IResult::Done(_, mut result) =>
                    initial.append(&mut result),
                _ =>
                    return Err(format!("Malformed declaration of initial nonterminals: \'{}\'", l))
            }
        } else if !l.is_empty() && !l.trim_left().starts_with("%") {
            rules.push(l.trim().parse()?);
        }
    }

    Ok((initial, rules))
}

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

    #[test]
    fn test_parse_comment_legal_input() {
        let legal_inputs = vec![
            "% some comment",
            "%% another comment",
        ];

        for legal_input in legal_inputs {
            assert_eq!(
                ("".as_bytes(), ()),
                parse_comment(legal_input.as_bytes()).unwrap(),
            )
        }
    }

    #[test]
    fn test_parse_comment_illegal_input() {
        let illegal_inputs = vec![
            " % not a comment",
            "something else % not a comment",
        ];

        for illegal_input in illegal_inputs {
            match parse_comment(illegal_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_token_legal_input() {
        let legal_inputs = vec![
            ("abcxyz", "", String::from("abcxyz")),
            ("abc,xyz", ",xyz", String::from("abc")),
            ("\"abc\"xyz", "xyz", String::from("abc")),
            ("\"a\\\\b\\\"c\"xyz", "xyz", String::from("a\\\\b\\\"c")),
            ("\"\"\"", "\"", String::from("")),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_token::<String>(legal_input.as_bytes()).unwrap()
            );
        }
    }

    #[test]
    fn test_parse_token_illegal_input() {
        let illegal_inputs = vec![
            " xyz",
            "-xyz",
            "→xyz",
            ",xyz",
            ";xyz",
            ")xyz",
            "]xyz",
            "\"\\\"",
            " \"a\"",
        ];

        for illegal_input in illegal_inputs {
            match parse_token::<String>(illegal_input.as_bytes()) {
                IResult::Done(_, _) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) | IResult::Incomplete(_) => (),
            }
        }

        let illegal_inputs = vec![
            "a",
            "\"a\"",
        ];

        for illegal_input in illegal_inputs {
            match parse_token::<u8>(illegal_input.as_bytes()) {
                IResult::Done(_, _) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) | IResult::Incomplete(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_token_incomplete_input() {
        let incomplete_inputs = vec![
            "\"a",
        ];

        for incomplete_input in incomplete_inputs {
            match parse_token::<String>(incomplete_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Error(_) =>
                    panic!("The input was not handled as incomplete: \'{}\'", incomplete_input),
                IResult::Incomplete(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_vec_legal_input() {
        let legal_inputs = vec![
            ("[]xyz", "xyz", vec![]),
            ("[a, bc, d]xyz", "xyz",
                vec![String::from("a"), String::from("bc"), String::from("d")]),
            ("[  a, b ,c]xyz", "xyz",
                vec![String::from("a"), String::from("b"), String::from("c")]),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_vec(legal_input.as_bytes(), parse_token, "[", "]", ",").unwrap()
            );
        }
    }

    #[test]
    fn test_parse_vec_illegal_input() {
        let illegal_inputs = vec![
            "(0)",
            "[0]",
            "[0, 1; 2)",
            " []",
            "[a)",
        ];

        for illegal_input in illegal_inputs {
            match parse_vec::<u8, _>(illegal_input.as_bytes(), parse_token, "[", ")", ",") {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_vec_incomplete_input() {
        let incomplete_inputs = vec![
            "[a",
            "["
        ];

        for incomplete_input in incomplete_inputs {
            match parse_vec::<String, _>(incomplete_input.as_bytes(), parse_token, "[", "]", ",") {
                IResult::Done(_, _) | IResult::Error(_) =>
                    panic!("The input was not handled as incomplete: \'{}\'", incomplete_input),
                IResult::Incomplete(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_initial_legal_input() {
        let legal_inputs = vec![
            ("initial: xyz % some comment", "", String::from("xyz")),
            ("initial: \"% xy\"", "", String::from("% xy")),
            ("initial:  xyz something else", "", String::from("xyz")),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_initial(legal_input.as_bytes()).unwrap()
            );
        }
    }

    #[test]
    fn test_parse_initials_legal_input() {
        let legal_inputs = vec![
            ("initial: [a]xyz", "xyz", vec![String::from("a")]),
            ("initial:  []xyz", "xyz", vec![]),
        ];

        for (legal_input, control_rest, control_parsed) in legal_inputs {
            assert_eq!(
                (control_rest.as_bytes(), control_parsed),
                parse_initials(legal_input.as_bytes()).unwrap()
            );
        }
    }

    #[test]
    fn test_parse_initials_illegal_input() {
        let illegal_inputs = vec![
            "initials: []xyz",
            " initial: []xyz",
        ];

        for illegal_input in illegal_inputs {
            match parse_initials::<String>(illegal_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Incomplete(_) =>
                    panic!("Was able to parse the illegal input \'{}\'", illegal_input),
                IResult::Error(_) => (),
            }
        }
    }

    #[test]
    fn test_parse_initials_incomplete_input() {
        let incomplete_inputs = vec![
            "init",
            "initial: [",
        ];

        for incomplete_input in incomplete_inputs {
            match parse_initials::<String>(incomplete_input.as_bytes()) {
                IResult::Done(_, _) | IResult::Error(_) =>
                    panic!("The input was not handled as incomplete: \'{}\'", incomplete_input),
                IResult::Incomplete(_) => (),
            }
        }
    }
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Boden Garman <bpgarman@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

#[macro_use]
extern crate uucore;

mod count_bytes;
mod countable;
mod word_count;
use count_bytes::count_bytes_fast;
use countable::WordCountable;
use word_count::{TitledWordCount, WordCount};

use clap::{crate_version, App, Arg, ArgMatches};
use thiserror::Error;

use std::fs::{self, File};
use std::io::{self, ErrorKind, Write};
use std::path::Path;

/// The minimum character width for formatting counts when reading from stdin.
const MINIMUM_WIDTH: usize = 7;

#[derive(Error, Debug)]
pub enum WcError {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("Expected a file, found directory {0}")]
    IsDirectory(String),
}

type WcResult<T> = Result<T, WcError>;

struct Settings {
    show_bytes: bool,
    show_chars: bool,
    show_lines: bool,
    show_words: bool,
    show_max_line_length: bool,
}

impl Settings {
    fn new(matches: &ArgMatches) -> Settings {
        let settings = Settings {
            show_bytes: matches.is_present(options::BYTES),
            show_chars: matches.is_present(options::CHAR),
            show_lines: matches.is_present(options::LINES),
            show_words: matches.is_present(options::WORDS),
            show_max_line_length: matches.is_present(options::MAX_LINE_LENGTH),
        };

        if settings.show_bytes
            || settings.show_chars
            || settings.show_lines
            || settings.show_words
            || settings.show_max_line_length
        {
            return settings;
        }

        Settings {
            show_bytes: true,
            show_chars: false,
            show_lines: true,
            show_words: true,
            show_max_line_length: false,
        }
    }

    fn number_enabled(&self) -> u32 {
        let mut result = 0;
        result += self.show_bytes as u32;
        result += self.show_chars as u32;
        result += self.show_lines as u32;
        result += self.show_max_line_length as u32;
        result += self.show_words as u32;
        result
    }
}

static ABOUT: &str = "Display newline, word, and byte counts for each FILE, and a total line if
more than one FILE is specified.";

pub mod options {
    pub static BYTES: &str = "bytes";
    pub static CHAR: &str = "chars";
    pub static LINES: &str = "lines";
    pub static MAX_LINE_LENGTH: &str = "max-line-length";
    pub static WORDS: &str = "words";
}

static ARG_FILES: &str = "files";

fn get_usage() -> String {
    format!(
        "{0} [OPTION]... [FILE]...
 With no FILE, or when FILE is -, read standard input.",
        executable!()
    )
}

enum StdinKind {
    /// Stdin specified on command-line with "-".
    Explicit,

    /// Stdin implicitly specified on command-line by not passing any positional argument.
    Implicit,
}

/// Supported inputs.
enum Input {
    /// A regular file.
    Path(String),

    /// Standard input.
    Stdin(StdinKind),
}

impl Input {
    /// Converts input to title that appears in stats.
    fn to_title(&self) -> Option<&str> {
        match self {
            Input::Path(path) => Some(path),
            Input::Stdin(StdinKind::Explicit) => Some("-"),
            Input::Stdin(StdinKind::Implicit) => None,
        }
    }
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();

    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::BYTES)
                .short("c")
                .long(options::BYTES)
                .help("print the byte counts"),
        )
        .arg(
            Arg::with_name(options::CHAR)
                .short("m")
                .long(options::CHAR)
                .help("print the character counts"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("l")
                .long(options::LINES)
                .help("print the newline counts"),
        )
        .arg(
            Arg::with_name(options::MAX_LINE_LENGTH)
                .short("L")
                .long(options::MAX_LINE_LENGTH)
                .help("print the length of the longest line"),
        )
        .arg(
            Arg::with_name(options::WORDS)
                .short("w")
                .long(options::WORDS)
                .help("print the word counts"),
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true))
        .get_matches_from(args);

    let mut inputs: Vec<Input> = matches
        .values_of(ARG_FILES)
        .map(|v| {
            v.map(|i| {
                if i == "-" {
                    Input::Stdin(StdinKind::Explicit)
                } else {
                    Input::Path(ToString::to_string(i))
                }
            })
            .collect()
        })
        .unwrap_or_default();

    if inputs.is_empty() {
        inputs.push(Input::Stdin(StdinKind::Implicit));
    }

    let settings = Settings::new(&matches);

    if wc(inputs, &settings).is_ok() {
        0
    } else {
        1
    }
}

fn word_count_from_reader<T: WordCountable>(
    mut reader: T,
    settings: &Settings,
    path: &str,
) -> WcResult<WordCount> {
    let only_count_bytes = settings.show_bytes
        && (!(settings.show_chars
            || settings.show_lines
            || settings.show_max_line_length
            || settings.show_words));
    if only_count_bytes {
        return Ok(WordCount {
            bytes: count_bytes_fast(&mut reader)?,
            ..WordCount::default()
        });
    }

    // we do not need to decode the byte stream if we're only counting bytes/newlines
    let decode_chars = settings.show_chars || settings.show_words || settings.show_max_line_length;

    // Sum the WordCount for each line. Show a warning for each line
    // that results in an IO error when trying to read it.
    let total = reader
        .lines()
        .filter_map(|res| match res {
            Ok(line) => Some(line),
            Err(e) => {
                show_warning!("Error while reading {}: {}", path, e);
                None
            }
        })
        .map(|line| WordCount::from_line(&line, decode_chars))
        .sum();
    Ok(total)
}

fn word_count_from_input(input: &Input, settings: &Settings) -> WcResult<WordCount> {
    match input {
        Input::Stdin(_) => {
            let stdin = io::stdin();
            let stdin_lock = stdin.lock();
            word_count_from_reader(stdin_lock, settings, "-")
        }
        Input::Path(path) => {
            let path_obj = Path::new(path);
            if path_obj.is_dir() {
                Err(WcError::IsDirectory(path.to_owned()))
            } else {
                let file = File::open(path)?;
                word_count_from_reader(file, settings, path)
            }
        }
    }
}

/// Print a message appropriate for the particular error to `stderr`.
///
/// # Examples
///
/// This will print `wc: /tmp: Is a directory` to `stderr`.
///
/// ```rust,ignore
/// show_error(Input::Path("/tmp"), WcError::IsDirectory("/tmp"))
/// ```
fn show_error(input: &Input, err: WcError) {
    match (input, err) {
        (_, WcError::IsDirectory(path)) => {
            show_error_custom_description!(path, "Is a directory");
        }
        (Input::Path(path), WcError::Io(e)) if e.kind() == ErrorKind::NotFound => {
            show_error_custom_description!(path, "No such file or directory");
        }
        (_, e) => {
            show_error!("{}", e);
        }
    };
}

/// Compute the number of digits needed to represent any count for this input.
///
/// If `input` is [`Input::Stdin`], then this function returns
/// [`MINIMUM_WIDTH`]. Otherwise, if metadata could not be read from
/// `input` then this function returns 1.
///
/// # Errors
///
/// This function will return an error if `input` is a [`Input::Path`]
/// and there is a problem accessing the metadata of the given `input`.
///
/// # Examples
///
/// A [`Input::Stdin`] gets a default minimum width:
///
/// ```rust,ignore
/// let input = Input::Stdin(StdinKind::Explicit);
/// assert_eq!(7, digit_width(input));
/// ```
fn digit_width(input: &Input) -> WcResult<Option<usize>> {
    match input {
        Input::Stdin(_) => Ok(Some(MINIMUM_WIDTH)),
        Input::Path(filename) => {
            let path = Path::new(filename);
            let metadata = fs::metadata(path)?;
            if metadata.is_file() {
                // TODO We are now computing the number of bytes in a file
                // twice: once here and once in `WordCount::from_line()` (or
                // in `count_bytes_fast()` if that function is called
                // instead). See GitHub issue #2201.
                let num_bytes = metadata.len();
                let num_digits = num_bytes.to_string().len();
                Ok(Some(num_digits))
            } else {
                Ok(None)
            }
        }
    }
}

/// Compute the number of digits needed to represent all counts in all inputs.
///
/// `inputs` may include zero or more [`Input::Stdin`] entries, each of
/// which represents reading from `stdin`. The presence of any such
/// entry causes this function to return a width that is at least
/// [`MINIMUM_WIDTH`].
///
/// If `input` is empty, then this function returns 1. If file metadata
/// could not be read from any of the [`Input::Path`] inputs and there
/// are no [`Input::Stdin`] inputs, then this function returns 1.
///
/// If there is a problem accessing the metadata, this function will
/// silently ignore the error and assume that the number of digits
/// needed to display the counts for that file is 1.
///
/// # Examples
///
/// An empty slice implies a width of 1:
///
/// ```rust,ignore
/// assert_eq!(1, max_width(&vec![]));
/// ```
///
/// The presence of [`Input::Stdin`] implies a minimum width:
///
/// ```rust,ignore
/// let inputs = vec![Input::Stdin(StdinKind::Explicit)];
/// assert_eq!(7, max_width(&inputs));
/// ```
fn max_width(inputs: &[Input]) -> usize {
    let mut result = 1;
    for input in inputs {
        match digit_width(input) {
            Ok(maybe_n) => {
                if let Some(n) = maybe_n {
                    result = result.max(n);
                }
            }
            Err(_) => continue,
        }
    }
    result
}

fn wc(inputs: Vec<Input>, settings: &Settings) -> Result<(), u32> {
    // Compute the width, in digits, to use when formatting counts.
    //
    // The width is the number of digits needed to print the number of
    // bytes in the largest file. This is true regardless of whether
    // the `settings` indicate that the bytes will be displayed.
    let mut error_count = 0;
    let max_width = max_width(&inputs);

    let mut total_word_count = WordCount::default();

    let num_inputs = inputs.len();

    for input in &inputs {
        let word_count = word_count_from_input(&input, settings).unwrap_or_else(|err| {
            show_error(&input, err);
            error_count += 1;
            WordCount::default()
        });
        total_word_count += word_count;
        let result = word_count.with_title(input.to_title());
        if let Err(err) = print_stats(settings, &result, max_width) {
            show_warning!(
                "failed to print result for {}: {}",
                result.title.unwrap_or("<stdin>"),
                err
            );
            error_count += 1;
        }
    }

    if num_inputs > 1 {
        let total_result = total_word_count.with_title(Some("total"));
        if let Err(err) = print_stats(settings, &total_result, max_width) {
            show_warning!("failed to print total: {}", err);
            error_count += 1;
        }
    }

    if error_count == 0 {
        Ok(())
    } else {
        Err(error_count)
    }
}

fn print_stats(
    settings: &Settings,
    result: &TitledWordCount,
    mut min_width: usize,
) -> WcResult<()> {
    let stdout = io::stdout();
    let mut stdout_lock = stdout.lock();

    if settings.number_enabled() <= 1 {
        // Prevent a leading space in case we only need to display a single
        // number.
        min_width = 0;
    }

    let mut is_first: bool = true;

    if settings.show_lines {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.lines, min_width)?;
        is_first = false;
    }
    if settings.show_words {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.words, min_width)?;
        is_first = false;
    }
    if settings.show_bytes {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.bytes, min_width)?;
        is_first = false;
    }
    if settings.show_chars {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.chars, min_width)?;
        is_first = false;
    }
    if settings.show_max_line_length {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(
            stdout_lock,
            "{:1$}",
            result.count.max_line_length, min_width
        )?;
    }

    if let Some(title) = result.title {
        writeln!(stdout_lock, " {}", title)?;
    } else {
        writeln!(stdout_lock)?;
    }

    Ok(())
}
<commit_msg>refactor/wc ~ fix `cargo clippy` complaint (clippy::needless_borrow)<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Boden Garman <bpgarman@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

#[macro_use]
extern crate uucore;

mod count_bytes;
mod countable;
mod word_count;
use count_bytes::count_bytes_fast;
use countable::WordCountable;
use word_count::{TitledWordCount, WordCount};

use clap::{crate_version, App, Arg, ArgMatches};
use thiserror::Error;

use std::fs::{self, File};
use std::io::{self, ErrorKind, Write};
use std::path::Path;

/// The minimum character width for formatting counts when reading from stdin.
const MINIMUM_WIDTH: usize = 7;

#[derive(Error, Debug)]
pub enum WcError {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("Expected a file, found directory {0}")]
    IsDirectory(String),
}

type WcResult<T> = Result<T, WcError>;

struct Settings {
    show_bytes: bool,
    show_chars: bool,
    show_lines: bool,
    show_words: bool,
    show_max_line_length: bool,
}

impl Settings {
    fn new(matches: &ArgMatches) -> Settings {
        let settings = Settings {
            show_bytes: matches.is_present(options::BYTES),
            show_chars: matches.is_present(options::CHAR),
            show_lines: matches.is_present(options::LINES),
            show_words: matches.is_present(options::WORDS),
            show_max_line_length: matches.is_present(options::MAX_LINE_LENGTH),
        };

        if settings.show_bytes
            || settings.show_chars
            || settings.show_lines
            || settings.show_words
            || settings.show_max_line_length
        {
            return settings;
        }

        Settings {
            show_bytes: true,
            show_chars: false,
            show_lines: true,
            show_words: true,
            show_max_line_length: false,
        }
    }

    fn number_enabled(&self) -> u32 {
        let mut result = 0;
        result += self.show_bytes as u32;
        result += self.show_chars as u32;
        result += self.show_lines as u32;
        result += self.show_max_line_length as u32;
        result += self.show_words as u32;
        result
    }
}

static ABOUT: &str = "Display newline, word, and byte counts for each FILE, and a total line if
more than one FILE is specified.";

pub mod options {
    pub static BYTES: &str = "bytes";
    pub static CHAR: &str = "chars";
    pub static LINES: &str = "lines";
    pub static MAX_LINE_LENGTH: &str = "max-line-length";
    pub static WORDS: &str = "words";
}

static ARG_FILES: &str = "files";

fn get_usage() -> String {
    format!(
        "{0} [OPTION]... [FILE]...
 With no FILE, or when FILE is -, read standard input.",
        executable!()
    )
}

enum StdinKind {
    /// Stdin specified on command-line with "-".
    Explicit,

    /// Stdin implicitly specified on command-line by not passing any positional argument.
    Implicit,
}

/// Supported inputs.
enum Input {
    /// A regular file.
    Path(String),

    /// Standard input.
    Stdin(StdinKind),
}

impl Input {
    /// Converts input to title that appears in stats.
    fn to_title(&self) -> Option<&str> {
        match self {
            Input::Path(path) => Some(path),
            Input::Stdin(StdinKind::Explicit) => Some("-"),
            Input::Stdin(StdinKind::Implicit) => None,
        }
    }
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();

    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::BYTES)
                .short("c")
                .long(options::BYTES)
                .help("print the byte counts"),
        )
        .arg(
            Arg::with_name(options::CHAR)
                .short("m")
                .long(options::CHAR)
                .help("print the character counts"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("l")
                .long(options::LINES)
                .help("print the newline counts"),
        )
        .arg(
            Arg::with_name(options::MAX_LINE_LENGTH)
                .short("L")
                .long(options::MAX_LINE_LENGTH)
                .help("print the length of the longest line"),
        )
        .arg(
            Arg::with_name(options::WORDS)
                .short("w")
                .long(options::WORDS)
                .help("print the word counts"),
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true))
        .get_matches_from(args);

    let mut inputs: Vec<Input> = matches
        .values_of(ARG_FILES)
        .map(|v| {
            v.map(|i| {
                if i == "-" {
                    Input::Stdin(StdinKind::Explicit)
                } else {
                    Input::Path(ToString::to_string(i))
                }
            })
            .collect()
        })
        .unwrap_or_default();

    if inputs.is_empty() {
        inputs.push(Input::Stdin(StdinKind::Implicit));
    }

    let settings = Settings::new(&matches);

    if wc(inputs, &settings).is_ok() {
        0
    } else {
        1
    }
}

fn word_count_from_reader<T: WordCountable>(
    mut reader: T,
    settings: &Settings,
    path: &str,
) -> WcResult<WordCount> {
    let only_count_bytes = settings.show_bytes
        && (!(settings.show_chars
            || settings.show_lines
            || settings.show_max_line_length
            || settings.show_words));
    if only_count_bytes {
        return Ok(WordCount {
            bytes: count_bytes_fast(&mut reader)?,
            ..WordCount::default()
        });
    }

    // we do not need to decode the byte stream if we're only counting bytes/newlines
    let decode_chars = settings.show_chars || settings.show_words || settings.show_max_line_length;

    // Sum the WordCount for each line. Show a warning for each line
    // that results in an IO error when trying to read it.
    let total = reader
        .lines()
        .filter_map(|res| match res {
            Ok(line) => Some(line),
            Err(e) => {
                show_warning!("Error while reading {}: {}", path, e);
                None
            }
        })
        .map(|line| WordCount::from_line(&line, decode_chars))
        .sum();
    Ok(total)
}

fn word_count_from_input(input: &Input, settings: &Settings) -> WcResult<WordCount> {
    match input {
        Input::Stdin(_) => {
            let stdin = io::stdin();
            let stdin_lock = stdin.lock();
            word_count_from_reader(stdin_lock, settings, "-")
        }
        Input::Path(path) => {
            let path_obj = Path::new(path);
            if path_obj.is_dir() {
                Err(WcError::IsDirectory(path.to_owned()))
            } else {
                let file = File::open(path)?;
                word_count_from_reader(file, settings, path)
            }
        }
    }
}

/// Print a message appropriate for the particular error to `stderr`.
///
/// # Examples
///
/// This will print `wc: /tmp: Is a directory` to `stderr`.
///
/// ```rust,ignore
/// show_error(Input::Path("/tmp"), WcError::IsDirectory("/tmp"))
/// ```
fn show_error(input: &Input, err: WcError) {
    match (input, err) {
        (_, WcError::IsDirectory(path)) => {
            show_error_custom_description!(path, "Is a directory");
        }
        (Input::Path(path), WcError::Io(e)) if e.kind() == ErrorKind::NotFound => {
            show_error_custom_description!(path, "No such file or directory");
        }
        (_, e) => {
            show_error!("{}", e);
        }
    };
}

/// Compute the number of digits needed to represent any count for this input.
///
/// If `input` is [`Input::Stdin`], then this function returns
/// [`MINIMUM_WIDTH`]. Otherwise, if metadata could not be read from
/// `input` then this function returns 1.
///
/// # Errors
///
/// This function will return an error if `input` is a [`Input::Path`]
/// and there is a problem accessing the metadata of the given `input`.
///
/// # Examples
///
/// A [`Input::Stdin`] gets a default minimum width:
///
/// ```rust,ignore
/// let input = Input::Stdin(StdinKind::Explicit);
/// assert_eq!(7, digit_width(input));
/// ```
fn digit_width(input: &Input) -> WcResult<Option<usize>> {
    match input {
        Input::Stdin(_) => Ok(Some(MINIMUM_WIDTH)),
        Input::Path(filename) => {
            let path = Path::new(filename);
            let metadata = fs::metadata(path)?;
            if metadata.is_file() {
                // TODO We are now computing the number of bytes in a file
                // twice: once here and once in `WordCount::from_line()` (or
                // in `count_bytes_fast()` if that function is called
                // instead). See GitHub issue #2201.
                let num_bytes = metadata.len();
                let num_digits = num_bytes.to_string().len();
                Ok(Some(num_digits))
            } else {
                Ok(None)
            }
        }
    }
}

/// Compute the number of digits needed to represent all counts in all inputs.
///
/// `inputs` may include zero or more [`Input::Stdin`] entries, each of
/// which represents reading from `stdin`. The presence of any such
/// entry causes this function to return a width that is at least
/// [`MINIMUM_WIDTH`].
///
/// If `input` is empty, then this function returns 1. If file metadata
/// could not be read from any of the [`Input::Path`] inputs and there
/// are no [`Input::Stdin`] inputs, then this function returns 1.
///
/// If there is a problem accessing the metadata, this function will
/// silently ignore the error and assume that the number of digits
/// needed to display the counts for that file is 1.
///
/// # Examples
///
/// An empty slice implies a width of 1:
///
/// ```rust,ignore
/// assert_eq!(1, max_width(&vec![]));
/// ```
///
/// The presence of [`Input::Stdin`] implies a minimum width:
///
/// ```rust,ignore
/// let inputs = vec![Input::Stdin(StdinKind::Explicit)];
/// assert_eq!(7, max_width(&inputs));
/// ```
fn max_width(inputs: &[Input]) -> usize {
    let mut result = 1;
    for input in inputs {
        match digit_width(input) {
            Ok(maybe_n) => {
                if let Some(n) = maybe_n {
                    result = result.max(n);
                }
            }
            Err(_) => continue,
        }
    }
    result
}

fn wc(inputs: Vec<Input>, settings: &Settings) -> Result<(), u32> {
    // Compute the width, in digits, to use when formatting counts.
    //
    // The width is the number of digits needed to print the number of
    // bytes in the largest file. This is true regardless of whether
    // the `settings` indicate that the bytes will be displayed.
    let mut error_count = 0;
    let max_width = max_width(&inputs);

    let mut total_word_count = WordCount::default();

    let num_inputs = inputs.len();

    for input in &inputs {
        let word_count = word_count_from_input(input, settings).unwrap_or_else(|err| {
            show_error(input, err);
            error_count += 1;
            WordCount::default()
        });
        total_word_count += word_count;
        let result = word_count.with_title(input.to_title());
        if let Err(err) = print_stats(settings, &result, max_width) {
            show_warning!(
                "failed to print result for {}: {}",
                result.title.unwrap_or("<stdin>"),
                err
            );
            error_count += 1;
        }
    }

    if num_inputs > 1 {
        let total_result = total_word_count.with_title(Some("total"));
        if let Err(err) = print_stats(settings, &total_result, max_width) {
            show_warning!("failed to print total: {}", err);
            error_count += 1;
        }
    }

    if error_count == 0 {
        Ok(())
    } else {
        Err(error_count)
    }
}

fn print_stats(
    settings: &Settings,
    result: &TitledWordCount,
    mut min_width: usize,
) -> WcResult<()> {
    let stdout = io::stdout();
    let mut stdout_lock = stdout.lock();

    if settings.number_enabled() <= 1 {
        // Prevent a leading space in case we only need to display a single
        // number.
        min_width = 0;
    }

    let mut is_first: bool = true;

    if settings.show_lines {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.lines, min_width)?;
        is_first = false;
    }
    if settings.show_words {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.words, min_width)?;
        is_first = false;
    }
    if settings.show_bytes {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.bytes, min_width)?;
        is_first = false;
    }
    if settings.show_chars {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.chars, min_width)?;
        is_first = false;
    }
    if settings.show_max_line_length {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(
            stdout_lock,
            "{:1$}",
            result.count.max_line_length, min_width
        )?;
    }

    if let Some(title) = result.title {
        writeln!(stdout_lock, " {}", title)?;
    } else {
        writeln!(stdout_lock)?;
    }

    Ok(())
}
<|endoftext|>