<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Gamma and derived distributions.

use num::Real;
use num;
use rand::{Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub enum Gamma {
    priv Large(GammaLargeShape),
    priv One(Exp),
    priv Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    shape: f64,
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            shape: shape,
            scale: scale,
            c: 1. / num::sqrt(9. * d),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match *self {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * num::pow(u, self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                num::ln(u) < 0.5 * x_sqr + self.d * (1.0 - v + num::ln(v)) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub enum ChiSquared {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    priv DoFExactlyOne,
    priv DoFAnythingElse(Gamma)
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match *self {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    priv numer: ChiSquared,
    priv denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    priv dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    priv chi: ChiSquared,
    priv dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use rand::distributions::*;
    use prelude::*;
    use rand::*;
    use super::*;

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_log_normal_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    use extra::test::BenchHarness;
    use mem::size_of;
    use prelude::*;
    use rand::distributions::IndependentSample;
    use rand::{StdRng, RAND_BENCH_N};
    use super::*;

    #[bench]
    fn bench_gamma_large_shape(bh: &mut BenchHarness) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = StdRng::new();

        bh.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(bh: &mut BenchHarness) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = StdRng::new();

        bh.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }
}
<commit_msg>Add a generic power function<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Gamma and derived distributions.

use num::Real;
use num;
use rand::{Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub enum Gamma {
    priv Large(GammaLargeShape),
    priv One(Exp),
    priv Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    shape: f64,
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            shape: shape,
            scale: scale,
            c: 1. / num::sqrt(9. * d),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match *self {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * num::powf(u, self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                num::ln(u) < 0.5 * x_sqr + self.d * (1.0 - v + num::ln(v)) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub enum ChiSquared {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    priv DoFExactlyOne,
    priv DoFAnythingElse(Gamma)
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match *self {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    priv numer: ChiSquared,
    priv denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    priv dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    priv chi: ChiSquared,
    priv dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use rand::distributions::*;
    use prelude::*;
    use rand::*;
    use super::*;

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_log_normal_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    use extra::test::BenchHarness;
    use mem::size_of;
    use prelude::*;
    use rand::distributions::IndependentSample;
    use rand::{StdRng, RAND_BENCH_N};
    use super::*;

    #[bench]
    fn bench_gamma_large_shape(bh: &mut BenchHarness) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = StdRng::new();

        bh.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(bh: &mut BenchHarness) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = StdRng::new();

        bh.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        bh.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }
}
<|endoftext|>


<commit_before>use hand::BJHand;
use bjaction::BJAction;
use bjaction::BJAction::HIT;
use bjaction::BJAction::STAND;
use bjaction::BJAction::DOUBLE;
use bjaction::BJAction::SURRENDER;
use bjaction::BJAction::SPLIT;
use cards::card::Card;
use cards::value::VALUES;
use cards::value::ACE;
use cards::value::KING;
use cards::value::QUEEN;
use cards::value::JACK;
use cards::value::TEN;
use rules::BJRules;
use shoe::shoe::DirectShoe;
use shoe::randomshoe::SuitPicker;

pub struct ActionCalculator;

impl ActionCalculator {
    pub fn expected_value_best_action(&self, hand: &mut BJHand,
                                  dealer_up_card: &Card, d: &mut DirectShoe,
                                  rules: &BJRules) -> f64 {
        let mut best_result: Option<f64>  = None;
        // TODO: Can I just do something like BJAction.variants ??
        let actions = [STAND, HIT, DOUBLE, SPLIT, SURRENDER];
        for &a in actions.iter() {
            match self.expected_value(hand, dealer_up_card, d, a, rules) {
                Some(r) => {
                    match best_result {
                        Some(b) => {
                            if b < r {
                                best_result = Some(r)
                            }
                        }
                        None => {
                            best_result = Some(r)
                        }
                    }
                }
                _ => {}
            }
        }
        return best_result.unwrap();
    }

    pub fn expected_value(&self, hand: &mut BJHand, dealer_up_card: &Card,
                      d: &mut DirectShoe, action: BJAction,
                      rules: &BJRules) -> Option<f64> {
        if !rules.can_take_action(hand, action) {
            return None;
        }
        return match action {
            HIT => {
                assert!(rules.can_hit(hand));
                let mut final_result = 0.0;
                for &v in VALUES.iter() {
                    let count_of_val = d.count(&v);
                    if count_of_val > 0 {
                        let odds_of_value =
                            count_of_val as f64 / d.len() as f64;
                        let card_from_deck = match d.remove(&v) {
                            Some(c) => c,
                            None => {
                                panic!("Item should exist!");
                            }
                        };
                        hand.add_card(card_from_deck);
                        let ev_with_value =
                            self.expected_value_best_action(
                                hand,dealer_up_card, d, rules);
                        final_result += odds_of_value * ev_with_value;
                        hand.remove_card(card_from_deck);
                        d.insert(&card_from_deck);
                    }
                }
                Some(final_result)
            }
            STAND => {
                assert!(rules.can_stand(hand));
                let this_hands_value = match rules.is_blackjack(hand) {
                    true => {
                        rules.blackjack_payout()
                    }
                    false => {
                        if hand.score() > 21 {
                            -1.0
                        } else {
                            let mut dealer_hand = BJHand::new();
                            dealer_hand.add_card(*dealer_up_card);
                            self.expected_with_dealer(hand,
                                                      &mut dealer_hand,
                                                      d, rules)
                        }
                    }
                };
                //   The flow is backwards, but the math works: We resolve the
                // previous hand, and then try to resolve this hand.  The real
                // flow of the cards is people pick their own hand, THEN the
                // dealer deals her own cards.  This logic is more like:
                // We let the dealer resolve the first split, then pick cards
                // for the second then let the dealer resolve the second split.
                match hand.splits_to_solve() > 0 {
                    false => Some(this_hands_value),
                    true => {
                        // Taking over 'hand' variable
                        let mut hand = hand.create_next_split_hand();
                        Some(this_hands_value +
                             self.expected_value_best_action(
                                 &mut hand, dealer_up_card, d, rules))
                    }
                }
            }
            DOUBLE => {
                assert!(rules.can_double(hand));
                let mut final_result = 0.0;
                for &v in VALUES.iter() {
                    let count_of_val = d.count(&v);
                    if count_of_val > 0 {
                        let odds_of_value =
                            count_of_val as f64 / d.len() as f64;
                        let card_from_deck = match d.remove(&v) {
                            Some(c) => c,
                            None => {
                                panic!("Expect a value!");
                            }
                        };
                        hand.add_card(card_from_deck);
                        hand.add_double_count();
                        let ev_with_value =
                            2.0 * self.expected_value_best_action(
                                hand, dealer_up_card, d, rules);
                        final_result += odds_of_value * ev_with_value;
                        hand.remove_card(card_from_deck);
                        hand.subtract_double_count();
                        d.insert(&card_from_deck);
                    }
                }
                Some(final_result)
            }
            SPLIT => {
                assert!(rules.can_split(hand));
                hand.split();
                // Split the hand, you'll get a hit here (since hit is
                // the only option).
                let final_result = self.expected_value_best_action(
                        hand, dealer_up_card, d, rules);
                hand.unsplit();
                Some(final_result)
            }
            SURRENDER => {
                assert!(rules.can_surrender(hand));
                Some(-0.5)
            }
        }
    }

    pub fn expected_with_dealer(&self, player_hand: &BJHand,
                            dealer_hand: &mut BJHand, d: &mut DirectShoe,
                            rules: &BJRules) -> f64 {
        if !rules.should_hit_dealer_hand(dealer_hand) {
            let dealer_score = dealer_hand.score();
            let player_score = player_hand.score();
            assert!(player_score <= 21);
            if dealer_score > 21 {
                return 1.0;
            } else  if dealer_score > player_score {
                return -1.0;
            } else if dealer_score < player_score {
                return 1.0;
            } else {
                return 0.0;
            }
        } else {
            // The dealer hits ... takes a random card
            let mut final_result = 0.0;
            let number_of_valid_cards = {
                if rules.dealer_blackjack_after_hand() ||
                    dealer_hand.len() != 1 {
                        d.len()
                } else {
                    match dealer_hand.score() {
                        10 => d.len() - d.count(&ACE),
                        11 => d.len() - d.count(&TEN) - d.count(&JACK) - d.count(&QUEEN) -
                           d.count(&KING),
                        _ => d.len()
                    }
                }
            };
            for &v in VALUES.iter() {
                let count_of_val = d.count(&v);
                if count_of_val > 0 {
                    let odds_of_value =
                        count_of_val as f64 /
                        number_of_valid_cards as f64;
                    let card_from_deck = match d.remove(&v) {
                        Some(c) => c,
                        None => {
                            panic!("Count positive, but couldn't remove!");
                        }
                    };
                    assert_eq!(card_from_deck.value().desc(), v.desc());
                    dealer_hand.add_card(card_from_deck);
                    if rules.is_blackjack(dealer_hand) &&
                        !rules.dealer_blackjack_after_hand() {
                            // ignore
                    } else {
                        let ev_with_value = self.expected_with_dealer(
                            player_hand, dealer_hand, d, rules);
                        final_result += odds_of_value * ev_with_value;
                    }

                    dealer_hand.remove_card(card_from_deck);
                    d.insert(&card_from_deck);
                }
            }
            return final_result;
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate test;
    use action_calculator::ActionCalculator;
    use self::test::Bencher;
    use cards::value::Value;
    use shoe::shoe::DirectShoe;
    use hand::BJHand;
    use std::num::Float;
    use cards::value;
    use rules::BJRules;

    fn check_value(dealer_cards: &Vec<Value>, player_cards: &Vec<Value>,
                   expected: f64) {
        use shoe::randomshoe::new_infinite_shoe;
        use rules::BJRules;
        let a = ActionCalculator;
        let rules = BJRules::new();
        let mut shoe = new_infinite_shoe();
        let expansion = 1000000.0f64;
        assert_eq!(
            (a.expected_with_dealer(
                &BJHand::new_from_deck(&mut shoe, player_cards).unwrap(),
                &mut BJHand::new_from_deck(&mut shoe, dealer_cards).unwrap(),
                &mut shoe,
                &rules) * expansion).round() as int,
            (expected * expansion) as int);
    }

    fn check_best_value(dealer_up_card: &Value, player_cards: &Vec<Value>,
                        expected: f64) {
        let rules = BJRules::new();
        check_best_value_rules(dealer_up_card, player_cards, expected, &rules);
    }

    fn check_best_value_rules(dealer_up_card: &Value, player_cards: &Vec<Value>,
                              expected: f64, rules: &BJRules) {
        use shoe::randomshoe::new_infinite_shoe;
        let a = ActionCalculator;
        let mut shoe = new_infinite_shoe();
        let expansion = 1000000.0f64;
        assert_eq!(
            (a.expected_value_best_action(
                &mut BJHand::new_from_deck(&mut shoe, player_cards).unwrap(),
                &shoe.remove(dealer_up_card).unwrap(),
                &mut shoe,
                rules) * expansion).round() as int,
            (expected * expansion) as int);
    }

    #[test]
    fn test_expected_infinite_deck() {
        // Numbers checked against http://wizardofodds.com/games/blackjack/appendix/1/
        check_value(&vec![value::FOUR],  &vec![value::TEN, value::EIGHT]           ,  0.175854);
        check_value(&vec![value::SIX],   &vec![value::TEN, value::SIX]             , -0.153699);
        check_value(&vec![value::SEVEN], &vec![value::TEN, value::TEN]             ,  0.773227);
        check_value(&vec![value::SEVEN], &vec![value::TEN, value::SEVEN]           , -0.106809);
        check_value(&vec![value::NINE],  &vec![value::TEN, value::FIVE]            , -0.543150);
        check_value(&vec![value::TEN],   &vec![value::TEN, value::TEN]             ,  0.554538);
        check_value(&vec![value::TEN],   &vec![value::TEN, value::FIVE, value::SIX],  0.962624);
        check_value(&vec![value::TEN],   &vec![value::TEN, value::SEVEN]           , -0.419721);
        check_value(&vec![value::ACE],   &vec![value::TEN, value::SIX]             , -0.666951);
        check_value(&vec![value::ACE],   &vec![value::TEN, value::EIGHT]           , -0.100199);
    }

    #[test]
    fn test_expected_best_value_surr_16_a() {
        let rules = BJRules::new_complex(true, 1, false, 1);
        check_best_value_rules(&value::ACE,   &vec![value::TEN, value::SIX], -0.5, &rules);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_11_6() {
        check_best_value(&value::SIX,   &vec![value::FIVE, value::SIX], 0.667380);
    }

    #[test]
    #[ignore]
    fn test_expected_best_value_88_10() {
        check_best_value(&value::TEN,   &vec![value::EIGHT, value::EIGHT], -0.480686);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_s18_9() {
        check_best_value(&value::NINE,   &vec![value::ACE, value::SEVEN], -0.100744);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_9_2() {
        // HIT
        check_best_value(&value::TWO,   &vec![value::FIVE, value::FOUR], 0.074446);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_9_3() {
        // DOUBLE
        check_best_value(&value::THREE,   &vec![value::FIVE, value::FOUR], 0.120816);
    }

    #[test]
    fn test_expected_best_value_18_a() {
        check_best_value(&value::ACE,    &vec![value::TEN, value::EIGHT]           , -0.100199);
    }

    #[test]
    fn test_expected_best_value_21_2() {
        check_value(&vec![value::TWO],   &vec![value::TEN, value::FIVE, value::SIX],  0.882007);
        check_value(&vec![value::TWO],   &vec![value::TEN, value::TEN]             ,  0.639987);
    }

    #[test]
    fn test_expected_best_value_15_9() {
        check_best_value(&value::NINE,  &vec![value::TEN, value::FIVE]             , -0.471578);
    }

    #[test]
    fn test_expected_best_value_16_7() {
        check_best_value(&value::SEVEN,  &vec![value::TEN, value::SIX]             , -0.414779);
    }

    #[test]
    fn test_expected_best_value_20_10() {
        check_value(&vec![value::TEN],   &vec![value::TEN, value::TEN]             ,  0.554538);
    }

    #[bench]
    fn bench_with_calc(b: &mut Bencher) {
        use cards::value;
        b.iter(|| {
            check_best_value(&value::SEVEN,  &vec![value::TEN, value::SIX]             , -0.414779);
        });
    }
}
<commit_msg>Include hand hashers<commit_after>use hand::BJHand;
use bjaction::BJAction;
use bjaction::BJAction::HIT;
use bjaction::BJAction::STAND;
use bjaction::BJAction::DOUBLE;
use bjaction::BJAction::SURRENDER;
use bjaction::BJAction::SPLIT;
use cards::card::Card;
use cards::value::VALUES;
use cards::value::ACE;
use cards::value::KING;
use cards::value::QUEEN;
use cards::value::JACK;
use cards::value::TEN;
use rules::BJRules;
use shoe::shoe::DirectShoe;
use shoe::randomshoe::SuitPicker;
use hand_hasher::HandHasher;
use hand_hasher::DeckHasher;

pub struct ActionCalculator<'a> {
    player_hand_hasher: &'a HandHasher + 'a,
    dealer_hand_hasher: &'a HandHasher + 'a,
    deck_hasher: &'a DeckHasher + 'a,
}

impl <'a>ActionCalculator<'a> {
    pub fn expected_value_best_action(&self, hand: &mut BJHand,
                                  dealer_up_card: &Card, d: &mut DirectShoe,
                                  rules: &BJRules) -> f64 {
        let mut best_result: Option<f64>  = None;
        // TODO: Can I just do something like BJAction.variants ??
        let actions = [STAND, HIT, DOUBLE, SPLIT, SURRENDER];
        for &a in actions.iter() {
            match self.expected_value(hand, dealer_up_card, d, a, rules) {
                Some(r) => {
                    match best_result {
                        Some(b) => {
                            if b < r {
                                best_result = Some(r)
                            }
                        }
                        None => {
                            best_result = Some(r)
                        }
                    }
                }
                _ => {}
            }
        }
        return best_result.unwrap();
    }

    pub fn expected_value(&self, hand: &mut BJHand, dealer_up_card: &Card,
                      d: &mut DirectShoe, action: BJAction,
                      rules: &BJRules) -> Option<f64> {
        if !rules.can_take_action(hand, action) {
            return None;
        }
        return match action {
            HIT => {
                assert!(rules.can_hit(hand));
                let mut final_result = 0.0;
                for &v in VALUES.iter() {
                    let count_of_val = d.count(&v);
                    if count_of_val > 0 {
                        let odds_of_value =
                            count_of_val as f64 / d.len() as f64;
                        let card_from_deck = match d.remove(&v) {
                            Some(c) => c,
                            None => {
                                panic!("Item should exist!");
                            }
                        };
                        hand.add_card(card_from_deck);
                        let ev_with_value =
                            self.expected_value_best_action(
                                hand,dealer_up_card, d, rules);
                        final_result += odds_of_value * ev_with_value;
                        hand.remove_card(card_from_deck);
                        d.insert(&card_from_deck);
                    }
                }
                Some(final_result)
            }
            STAND => {
                assert!(rules.can_stand(hand));
                let this_hands_value = match rules.is_blackjack(hand) {
                    true => {
                        rules.blackjack_payout()
                    }
                    false => {
                        if hand.score() > 21 {
                            -1.0
                        } else {
                            let mut dealer_hand = BJHand::new();
                            dealer_hand.add_card(*dealer_up_card);
                            self.expected_with_dealer(hand,
                                                      &mut dealer_hand,
                                                      d, rules)
                        }
                    }
                };
                //   The flow is backwards, but the math works: We resolve the
                // previous hand, and then try to resolve this hand.  The real
                // flow of the cards is people pick their own hand, THEN the
                // dealer deals her own cards.  This logic is more like:
                // We let the dealer resolve the first split, then pick cards
                // for the second then let the dealer resolve the second split.
                match hand.splits_to_solve() > 0 {
                    false => Some(this_hands_value),
                    true => {
                        // Taking over 'hand' variable
                        let mut hand = hand.create_next_split_hand();
                        Some(this_hands_value +
                             self.expected_value_best_action(
                                 &mut hand, dealer_up_card, d, rules))
                    }
                }
            }
            DOUBLE => {
                assert!(rules.can_double(hand));
                let mut final_result = 0.0;
                for &v in VALUES.iter() {
                    let count_of_val = d.count(&v);
                    if count_of_val > 0 {
                        let odds_of_value =
                            count_of_val as f64 / d.len() as f64;
                        let card_from_deck = match d.remove(&v) {
                            Some(c) => c,
                            None => {
                                panic!("Expect a value!");
                            }
                        };
                        hand.add_card(card_from_deck);
                        hand.add_double_count();
                        let ev_with_value =
                            2.0 * self.expected_value_best_action(
                                hand, dealer_up_card, d, rules);
                        final_result += odds_of_value * ev_with_value;
                        hand.remove_card(card_from_deck);
                        hand.subtract_double_count();
                        d.insert(&card_from_deck);
                    }
                }
                Some(final_result)
            }
            SPLIT => {
                assert!(rules.can_split(hand));
                hand.split();
                // Split the hand, you'll get a hit here (since hit is
                // the only option).
                let final_result = self.expected_value_best_action(
                        hand, dealer_up_card, d, rules);
                hand.unsplit();
                Some(final_result)
            }
            SURRENDER => {
                assert!(rules.can_surrender(hand));
                Some(-0.5)
            }
        }
    }

    pub fn expected_with_dealer(&self, player_hand: &BJHand,
                            dealer_hand: &mut BJHand, d: &mut DirectShoe,
                            rules: &BJRules) -> f64 {
        if !rules.should_hit_dealer_hand(dealer_hand) {
            let dealer_score = dealer_hand.score();
            let player_score = player_hand.score();
            assert!(player_score <= 21);
            if dealer_score > 21 {
                return 1.0;
            } else  if dealer_score > player_score {
                return -1.0;
            } else if dealer_score < player_score {
                return 1.0;
            } else {
                return 0.0;
            }
        } else {
            // The dealer hits ... takes a random card
            let mut final_result = 0.0;
            let number_of_valid_cards = {
                if rules.dealer_blackjack_after_hand() ||
                    dealer_hand.len() != 1 {
                        d.len()
                } else {
                    match dealer_hand.score() {
                        10 => d.len() - d.count(&ACE),
                        11 => d.len() - d.count(&TEN) - d.count(&JACK) - d.count(&QUEEN) -
                           d.count(&KING),
                        _ => d.len()
                    }
                }
            };
            for &v in VALUES.iter() {
                let count_of_val = d.count(&v);
                if count_of_val > 0 {
                    let odds_of_value =
                        count_of_val as f64 /
                        number_of_valid_cards as f64;
                    let card_from_deck = match d.remove(&v) {
                        Some(c) => c,
                        None => {
                            panic!("Count positive, but couldn't remove!");
                        }
                    };
                    assert_eq!(card_from_deck.value().desc(), v.desc());
                    dealer_hand.add_card(card_from_deck);
                    if rules.is_blackjack(dealer_hand) &&
                        !rules.dealer_blackjack_after_hand() {
                            // ignore
                    } else {
                        let ev_with_value = self.expected_with_dealer(
                            player_hand, dealer_hand, d, rules);
                        final_result += odds_of_value * ev_with_value;
                    }

                    dealer_hand.remove_card(card_from_deck);
                    d.insert(&card_from_deck);
                }
            }
            return final_result;
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate test;
    use action_calculator::ActionCalculator;
    use self::test::Bencher;
    use cards::value::Value;
    use shoe::shoe::DirectShoe;
    use hand::BJHand;
    use std::num::Float;
    use cards::value;
    use rules::BJRules;
    use hand_hasher::DealerHandHasher;
    use hand_hasher::PlayerHandHasher;
    use hand_hasher::SuitlessDeckHasher;

    fn check_value(dealer_cards: &Vec<Value>, player_cards: &Vec<Value>,
                   expected: f64) {
        use shoe::randomshoe::new_infinite_shoe;
        use rules::BJRules;
        let a = ActionCalculator {
            player_hand_hasher: &PlayerHandHasher,
            dealer_hand_hasher: &DealerHandHasher,
            deck_hasher: &SuitlessDeckHasher,
        };
        let rules = BJRules::new();
        let mut shoe = new_infinite_shoe();
        let expansion = 1000000.0f64;
        assert_eq!(
            (a.expected_with_dealer(
                &BJHand::new_from_deck(&mut shoe, player_cards).unwrap(),
                &mut BJHand::new_from_deck(&mut shoe, dealer_cards).unwrap(),
                &mut shoe,
                &rules) * expansion).round() as int,
            (expected * expansion) as int);
    }

    fn check_best_value(dealer_up_card: &Value, player_cards: &Vec<Value>,
                        expected: f64) {
        let rules = BJRules::new();
        check_best_value_rules(dealer_up_card, player_cards, expected, &rules);
    }

    fn check_best_value_rules(dealer_up_card: &Value, player_cards: &Vec<Value>,
                              expected: f64, rules: &BJRules) {
        use shoe::randomshoe::new_infinite_shoe;
        let a = ActionCalculator {
            player_hand_hasher: &PlayerHandHasher,
            dealer_hand_hasher: &DealerHandHasher,
            deck_hasher: &SuitlessDeckHasher,
        };
        let mut shoe = new_infinite_shoe();
        let expansion = 1000000.0f64;
        assert_eq!(
            (a.expected_value_best_action(
                &mut BJHand::new_from_deck(&mut shoe, player_cards).unwrap(),
                &shoe.remove(dealer_up_card).unwrap(),
                &mut shoe,
                rules) * expansion).round() as int,
            (expected * expansion) as int);
    }

    #[test]
    fn test_expected_infinite_deck() {
        // Numbers checked against http://wizardofodds.com/games/blackjack/appendix/1/
        check_value(&vec![value::FOUR],  &vec![value::TEN, value::EIGHT]           ,  0.175854);
        check_value(&vec![value::SIX],   &vec![value::TEN, value::SIX]             , -0.153699);
        check_value(&vec![value::SEVEN], &vec![value::TEN, value::TEN]             ,  0.773227);
        check_value(&vec![value::SEVEN], &vec![value::TEN, value::SEVEN]           , -0.106809);
        check_value(&vec![value::NINE],  &vec![value::TEN, value::FIVE]            , -0.543150);
        check_value(&vec![value::TEN],   &vec![value::TEN, value::TEN]             ,  0.554538);
        check_value(&vec![value::TEN],   &vec![value::TEN, value::FIVE, value::SIX],  0.962624);
        check_value(&vec![value::TEN],   &vec![value::TEN, value::SEVEN]           , -0.419721);
        check_value(&vec![value::ACE],   &vec![value::TEN, value::SIX]             , -0.666951);
        check_value(&vec![value::ACE],   &vec![value::TEN, value::EIGHT]           , -0.100199);
    }

    #[test]
    fn test_expected_best_value_surr_16_a() {
        let rules = BJRules::new_complex(true, 1, false, 1);
        check_best_value_rules(&value::ACE,   &vec![value::TEN, value::SIX], -0.5, &rules);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_11_6() {
        check_best_value(&value::SIX,   &vec![value::FIVE, value::SIX], 0.667380);
    }

    #[test]
    #[ignore]
    fn test_expected_best_value_88_10() {
        check_best_value(&value::TEN,   &vec![value::EIGHT, value::EIGHT], -0.480686);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_s18_9() {
        check_best_value(&value::NINE,   &vec![value::ACE, value::SEVEN], -0.100744);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_9_2() {
        // HIT
        check_best_value(&value::TWO,   &vec![value::FIVE, value::FOUR], 0.074446);
    }

    #[test]
    #[ignore] // while slow
    fn test_expected_best_value_9_3() {
        // DOUBLE
        check_best_value(&value::THREE,   &vec![value::FIVE, value::FOUR], 0.120816);
    }

    #[test]
    fn test_expected_best_value_18_a() {
        check_best_value(&value::ACE,    &vec![value::TEN, value::EIGHT]           , -0.100199);
    }

    #[test]
    fn test_expected_best_value_21_2() {
        check_value(&vec![value::TWO],   &vec![value::TEN, value::FIVE, value::SIX],  0.882007);
        check_value(&vec![value::TWO],   &vec![value::TEN, value::TEN]             ,  0.639987);
    }

    #[test]
    fn test_expected_best_value_15_9() {
        check_best_value(&value::NINE,  &vec![value::TEN, value::FIVE]             , -0.471578);
    }

    #[test]
    fn test_expected_best_value_16_7() {
        check_best_value(&value::SEVEN,  &vec![value::TEN, value::SIX]             , -0.414779);
    }

    #[test]
    fn test_expected_best_value_20_10() {
        check_value(&vec![value::TEN],   &vec![value::TEN, value::TEN]             ,  0.554538);
    }

    #[bench]
    fn bench_with_calc(b: &mut Bencher) {
        use cards::value;
        b.iter(|| {
            check_best_value(&value::SEVEN,  &vec![value::TEN, value::SIX]             , -0.414779);
        });
    }
}
<|endoftext|>


<commit_before>extern crate gcc;
#[cfg(feature = "bindgen")]
extern crate bindgen;

#[cfg(not(feature = "bindgen"))]
use std::fs;
use std::env;
use std::path::Path;

#[cfg(feature = "bindgen")]
fn generate_binding(out_file: &Path) {
    let _ = bindgen::builder()
        .clang_arg("--std=c++11")
        .clang_arg(if cfg!(feature = "sse42") {
                       "-msse4.2"
                   } else {
                       "-march=native"
                   })
        .header("src/fasthash.hpp")
        .no_unstable_rust()
        .generate_inline_functions(true)
        .disable_name_namespacing()
        .hide_type(".*PCCP.*")
        .whitelisted_function("^CityHash.*")
        .whitelisted_function("^farmhash.*")
        .whitelisted_function("^lookup3.*")
        .whitelisted_function("^metrohash.*")
        .whitelisted_function("^mum_hash.*")
        .whitelisted_function("^MurmurHash.*")
        .whitelisted_function("^SpookyHasher.*")
        .whitelisted_function("^t1ha.*")
        .whitelisted_function("^XXH.*")
        .link_static("fasthash")
        .generate()
        .unwrap()
        .write_to_file(out_file)
        .expect("fail to write bindings");
}

#[cfg(not(feature = "bindgen"))]
fn generate_binding(out_file: &Path) {
    let suffix = if cfg!(target_os = "macos") {
        "macos"
    } else {
        "linux"
    };

    fs::copy(format!("src/fasthash_{}.rs", suffix), out_file).expect("fail to copy bindings");
}

fn main() {
    let mut gcc_config = gcc::Config::new();

    gcc_config
        .file("src/fasthash.cpp")
        .file("src/smhasher/City.cpp")
        .file("src/smhasher/farmhash-c.c")
        .file("src/smhasher/lookup3.cpp")
        .file("src/smhasher/mum.cc")
        .file("src/smhasher/metrohash64.cpp")
        .file("src/smhasher/metrohash128.cpp")
        .file("src/smhasher/MurmurHash1.cpp")
        .file("src/smhasher/MurmurHash2.cpp")
        .file("src/smhasher/MurmurHash3.cpp")
        .file("src/smhasher/Spooky.cpp")
        .file("src/smhasher/t1ha.cc")
        .file("src/smhasher/xxhash.c");

    if cfg!(feature = "sse42") {
        gcc_config
            .flag("-msse4.2")
            .file("src/smhasher/metrohash64crc.cpp")
            .file("src/smhasher/metrohash128crc.cpp");
    }

    gcc_config.compile("libfasthash.a");

    let out_dir = env::var("OUT_DIR").unwrap();
    let out_file = Path::new(&out_dir).join("src/fasthash.rs");

    generate_binding(&out_file);

    if cfg!(target_os = "macos") {
        println!("cargo:rustc-link-lib=dylib=c++");
    } else {
        println!("cargo:rustc-link-lib=dylib=stdc++");
        println!("cargo:rustc-link-lib=dylib=gcc");
    }
}
<commit_msg>Use gcc::Build given deprecation warning<commit_after>extern crate gcc;
#[cfg(feature = "bindgen")]
extern crate bindgen;

#[cfg(not(feature = "bindgen"))]
use std::fs;
use std::env;
use std::path::Path;

#[cfg(feature = "bindgen")]
fn generate_binding(out_file: &Path) {
    let _ = bindgen::builder()
        .clang_arg("--std=c++11")
        .clang_arg(if cfg!(feature = "sse42") {
                       "-msse4.2"
                   } else {
                       "-march=native"
                   })
        .header("src/fasthash.hpp")
        .no_unstable_rust()
        .generate_inline_functions(true)
        .disable_name_namespacing()
        .hide_type(".*PCCP.*")
        .whitelisted_function("^CityHash.*")
        .whitelisted_function("^farmhash.*")
        .whitelisted_function("^lookup3.*")
        .whitelisted_function("^metrohash.*")
        .whitelisted_function("^mum_hash.*")
        .whitelisted_function("^MurmurHash.*")
        .whitelisted_function("^SpookyHasher.*")
        .whitelisted_function("^t1ha.*")
        .whitelisted_function("^XXH.*")
        .link_static("fasthash")
        .generate()
        .unwrap()
        .write_to_file(out_file)
        .expect("fail to write bindings");
}

#[cfg(not(feature = "bindgen"))]
fn generate_binding(out_file: &Path) {
    let suffix = if cfg!(target_os = "macos") {
        "macos"
    } else {
        "linux"
    };

    fs::copy(format!("src/fasthash_{}.rs", suffix), out_file).expect("fail to copy bindings");
}

fn main() {
    let mut gcc_config = gcc::Build::new();

    gcc_config
        .file("src/fasthash.cpp")
        .file("src/smhasher/City.cpp")
        .file("src/smhasher/farmhash-c.c")
        .file("src/smhasher/lookup3.cpp")
        .file("src/smhasher/mum.cc")
        .file("src/smhasher/metrohash64.cpp")
        .file("src/smhasher/metrohash128.cpp")
        .file("src/smhasher/MurmurHash1.cpp")
        .file("src/smhasher/MurmurHash2.cpp")
        .file("src/smhasher/MurmurHash3.cpp")
        .file("src/smhasher/Spooky.cpp")
        .file("src/smhasher/t1ha.cc")
        .file("src/smhasher/xxhash.c");

    if cfg!(feature = "sse42") {
        gcc_config
            .flag("-msse4.2")
            .file("src/smhasher/metrohash64crc.cpp")
            .file("src/smhasher/metrohash128crc.cpp");
    }

    gcc_config.compile("libfasthash.a");

    let out_dir = env::var("OUT_DIR").unwrap();
    let out_file = Path::new(&out_dir).join("src/fasthash.rs");

    generate_binding(&out_file);

    if cfg!(target_os = "macos") {
        println!("cargo:rustc-link-lib=dylib=c++");
    } else {
        println!("cargo:rustc-link-lib=dylib=stdc++");
        println!("cargo:rustc-link-lib=dylib=gcc");
    }
}
<|endoftext|>


<commit_before><commit_msg>(#368): Log when an element is not rendered due to a nonexistent filter<commit_after><|endoftext|>


<commit_before>//! Unix peer credentials.

// NOTE: Code in this file is heavily based on work done in PR 13 from the tokio-uds repository on
//       GitHub.
//
//       For reference, the link is here: https://github.com/tokio-rs/tokio-uds/pull/13
//       Credit to Martin Habovštiak (GitHub username Kixunil) and contributors for this work.

use libc::{gid_t, uid_t};

/// Credentials for a UNIX process for credentials passing.
#[unstable(feature = "peer_credentials_unix_socket", issue = "42839", reason = "unstable")]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct UCred {
    pub uid: uid_t,
    pub gid: gid_t,
}

#[cfg(any(target_os = "android", target_os = "linux"))]
pub use self::impl_linux::peer_cred;

#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "macos",
    target_os = "openbsd"
))]
pub use self::impl_bsd::peer_cred;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod impl_linux {
    use super::UCred;
    use crate::os::unix::io::AsRawFd;
    use crate::os::unix::net::UnixStream;
    use crate::{io, mem};

    pub fn peer_cred(socket: &UnixStream) -> io::Result<UCred> {
        use libc::{c_void, ucred};

        let ucred_size = mem::size_of::<ucred>();

        // Trivial sanity checks.
        assert!(mem::size_of::<u32>() <= mem::size_of::<usize>());
        assert!(ucred_size <= u32::max_value() as usize);

        let mut ucred_size = ucred_size as u32;
        let mut ucred: ucred = ucred { pid: 1, uid: 1, gid: 1 };

        unsafe {
            let ret = libc::getsockopt(
                socket.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_PEERCRED,
                &mut ucred as *mut ucred as *mut c_void,
                &mut ucred_size,
            );

            if ret == 0 && ucred_size as usize == mem::size_of::<ucred>() {
                Ok(UCred { uid: ucred.uid, gid: ucred.gid })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(
    target_os = "dragonfly",
    target_os = "macos",
    target_os = "ios",
    target_os = "freebsd",
    target_os = "openbsd"
))]
pub mod impl_bsd {
    use super::UCred;
    use crate::io;
    use crate::os::unix::io::AsRawFd;
    use crate::os::unix::net::UnixStream;

    pub fn peer_cred(socket: &UnixStream) -> io::Result<UCred> {
        let mut cred = UCred { uid: 1, gid: 1 };
        unsafe {
            let ret = libc::getpeereid(socket.as_raw_fd(), &mut cred.uid, &mut cred.gid);

            if ret == 0 { Ok(cred) } else { Err(io::Error::last_os_error()) }
        }
    }
}
<commit_msg>Add basic test for Unix peer credentials<commit_after>//! Unix peer credentials.

// NOTE: Code in this file is heavily based on work done in PR 13 from the tokio-uds repository on
//       GitHub.
//
//       For reference, the link is here: https://github.com/tokio-rs/tokio-uds/pull/13
//       Credit to Martin Habovštiak (GitHub username Kixunil) and contributors for this work.

use libc::{gid_t, uid_t};

/// Credentials for a UNIX process for credentials passing.
#[unstable(feature = "peer_credentials_unix_socket", issue = "42839", reason = "unstable")]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct UCred {
    pub uid: uid_t,
    pub gid: gid_t,
}

#[cfg(any(target_os = "android", target_os = "linux"))]
pub use self::impl_linux::peer_cred;

#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "macos",
    target_os = "openbsd"
))]
pub use self::impl_bsd::peer_cred;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod impl_linux {
    use super::UCred;
    use crate::os::unix::io::AsRawFd;
    use crate::os::unix::net::UnixStream;
    use crate::{io, mem};

    pub fn peer_cred(socket: &UnixStream) -> io::Result<UCred> {
        use libc::{c_void, ucred};

        let ucred_size = mem::size_of::<ucred>();

        // Trivial sanity checks.
        assert!(mem::size_of::<u32>() <= mem::size_of::<usize>());
        assert!(ucred_size <= u32::max_value() as usize);

        let mut ucred_size = ucred_size as u32;
        let mut ucred: ucred = ucred { pid: 1, uid: 1, gid: 1 };

        unsafe {
            let ret = libc::getsockopt(
                socket.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_PEERCRED,
                &mut ucred as *mut ucred as *mut c_void,
                &mut ucred_size,
            );

            if ret == 0 && ucred_size as usize == mem::size_of::<ucred>() {
                Ok(UCred { uid: ucred.uid, gid: ucred.gid })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(
    target_os = "dragonfly",
    target_os = "macos",
    target_os = "ios",
    target_os = "freebsd",
    target_os = "openbsd"
))]
pub mod impl_bsd {
    use super::UCred;
    use crate::io;
    use crate::os::unix::io::AsRawFd;
    use crate::os::unix::net::UnixStream;

    pub fn peer_cred(socket: &UnixStream) -> io::Result<UCred> {
        let mut cred = UCred { uid: 1, gid: 1 };
        unsafe {
            let ret = libc::getpeereid(socket.as_raw_fd(), &mut cred.uid, &mut cred.gid);

            if ret == 0 { Ok(cred) } else { Err(io::Error::last_os_error()) }
        }
    }
}

#[cfg(test)]
mod test {
    use crate::os::unix::net::UnixStream;
    use libc::{getegid, geteuid};

    #[test]
    fn test_socket_pair() {
        // Create two connected sockets and get their peer credentials. They should be equal.
        let (sock_a, sock_b) = UnixStream::pair().unwrap();
        let (cred_a, cred_b) = (sock_a.peer_cred().unwrap(), sock_b.peer_cred().unwrap());
        assert_eq!(cred_a, cred_b);

        // Check that the UID and GIDs match up.
        let uid = unsafe { geteuid() };
        let gid = unsafe { getegid() };
        assert_eq!(cred_a.uid, uid);
        assert_eq!(cred_a.gid, gid);
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use syntax::{ast, ptr};
use syntax::codemap::Span;
use syntax::parse::classify;

use codemap::SpanUtils;
use expr::{block_contains_comment, is_simple_block, is_unsafe_block, rewrite_cond, ToExpr};
use items::{span_hi_for_arg, span_lo_for_arg};
use lists::{definitive_tactic, itemize_list, write_list, DefinitiveListTactic, ListFormatting,
            ListTactic, Separator, SeparatorPlace, SeparatorTactic};
use rewrite::{Rewrite, RewriteContext};
use shape::Shape;
use utils::{last_line_width, left_most_sub_expr, stmt_expr};

// This module is pretty messy because of the rules around closures and blocks:
// FIXME - the below is probably no longer true in full.
//   * if there is a return type, then there must be braces,
//   * given a closure with braces, whether that is parsed to give an inner block
//     or not depends on if there is a return type and if there are statements
//     in that block,
//   * if the first expression in the body ends with a block (i.e., is a
//     statement without needing a semi-colon), then adding or removing braces
//     can change whether it is treated as an expression or statement.

pub fn rewrite_closure(
    capture: ast::CaptureBy,
    fn_decl: &ast::FnDecl,
    body: &ast::Expr,
    span: Span,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    debug!("rewrite_closure {:?}", body);

    let (prefix, extra_offset) =
        rewrite_closure_fn_decl(capture, fn_decl, body, span, context, shape)?;
    // 1 = space between `|...|` and body.
    let body_shape = shape.offset_left(extra_offset)?;

    if let ast::ExprKind::Block(ref block) = body.node {
        // The body of the closure is an empty block.
        if block.stmts.is_empty() && !block_contains_comment(block, context.codemap) {
            return Some(format!("{} {{}}", prefix));
        }

        let result = match fn_decl.output {
            ast::FunctionRetTy::Default(_) => {
                try_rewrite_without_block(body, &prefix, context, shape, body_shape)
            }
            _ => None,
        };

        result.or_else(|| {
            // Either we require a block, or tried without and failed.
            rewrite_closure_block(block, &prefix, context, body_shape)
        })
    } else {
        rewrite_closure_expr(body, &prefix, context, body_shape).or_else(|| {
            // The closure originally had a non-block expression, but we can't fit on
            // one line, so we'll insert a block.
            rewrite_closure_with_block(body, &prefix, context, body_shape)
        })
    }
}

fn try_rewrite_without_block(
    expr: &ast::Expr,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
    body_shape: Shape,
) -> Option<String> {
    let expr = get_inner_expr(expr, prefix, context);

    if is_block_closure_forced(expr) {
        rewrite_closure_with_block(expr, prefix, context, shape)
    } else {
        rewrite_closure_expr(expr, prefix, context, body_shape)
    }
}

fn get_inner_expr<'a>(
    expr: &'a ast::Expr,
    prefix: &str,
    context: &RewriteContext,
) -> &'a ast::Expr {
    if let ast::ExprKind::Block(ref block) = expr.node {
        if !needs_block(block, prefix, context) {
            // block.stmts.len() == 1
            if let Some(expr) = stmt_expr(&block.stmts[0]) {
                return get_inner_expr(expr, prefix, context);
            }
        }
    }

    expr
}

// Figure out if a block is necessary.
fn needs_block(block: &ast::Block, prefix: &str, context: &RewriteContext) -> bool {
    is_unsafe_block(block) || block.stmts.len() > 1 || context.inside_macro
        || block_contains_comment(block, context.codemap) || prefix.contains('\n')
}

// Rewrite closure with a single expression wrapping its body with block.
fn rewrite_closure_with_block(
    body: &ast::Expr,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    let block = ast::Block {
        stmts: vec![
            ast::Stmt {
                id: ast::NodeId::new(0),
                node: ast::StmtKind::Expr(ptr::P(body.clone())),
                span: body.span,
            },
        ],
        id: ast::NodeId::new(0),
        rules: ast::BlockCheckMode::Default,
        span: body.span,
    };
    let block = ::expr::rewrite_block_with_visitor(context, "", &block, shape, false)?;
    Some(format!("{} {}", prefix, block))
}

// Rewrite closure with a single expression without wrapping its body with block.
fn rewrite_closure_expr(
    expr: &ast::Expr,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    let mut rewrite = expr.rewrite(context, shape);
    if classify::expr_requires_semi_to_be_stmt(left_most_sub_expr(expr)) {
        rewrite = and_one_line(rewrite);
    }
    rewrite = rewrite.and_then(|rw| {
        if context.config.force_multiline_blocks() && rw.contains('\n') {
            None
        } else {
            Some(rw)
        }
    });
    rewrite.map(|rw| format!("{} {}", prefix, rw))
}

// Rewrite closure whose body is block.
fn rewrite_closure_block(
    block: &ast::Block,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    let block_shape = shape.block();
    let block_str = block.rewrite(context, block_shape)?;
    Some(format!("{} {}", prefix, block_str))
}

// Return type is (prefix, extra_offset)
fn rewrite_closure_fn_decl(
    capture: ast::CaptureBy,
    fn_decl: &ast::FnDecl,
    body: &ast::Expr,
    span: Span,
    context: &RewriteContext,
    shape: Shape,
) -> Option<(String, usize)> {
    let mover = if capture == ast::CaptureBy::Value {
        "move "
    } else {
        ""
    };
    // 4 = "|| {".len(), which is overconservative when the closure consists of
    // a single expression.
    let nested_shape = shape.shrink_left(mover.len())?.sub_width(4)?;

    // 1 = |
    let argument_offset = nested_shape.indent + 1;
    let arg_shape = nested_shape.offset_left(1)?.visual_indent(0);
    let ret_str = fn_decl.output.rewrite(context, arg_shape)?;

    let arg_items = itemize_list(
        context.codemap,
        fn_decl.inputs.iter(),
        "|",
        ",",
        |arg| span_lo_for_arg(arg),
        |arg| span_hi_for_arg(context, arg),
        |arg| arg.rewrite(context, arg_shape),
        context.codemap.span_after(span, "|"),
        body.span.lo(),
        false,
    );
    let item_vec = arg_items.collect::<Vec<_>>();
    // 1 = space between arguments and return type.
    let horizontal_budget = nested_shape
        .width
        .checked_sub(ret_str.len() + 1)
        .unwrap_or(0);
    let tactic = definitive_tactic(
        &item_vec,
        ListTactic::HorizontalVertical,
        Separator::Comma,
        horizontal_budget,
    );
    let arg_shape = match tactic {
        DefinitiveListTactic::Horizontal => arg_shape.sub_width(ret_str.len() + 1)?,
        _ => arg_shape,
    };

    let fmt = ListFormatting {
        tactic: tactic,
        separator: ",",
        trailing_separator: SeparatorTactic::Never,
        separator_place: SeparatorPlace::Back,
        shape: arg_shape,
        ends_with_newline: false,
        preserve_newline: true,
        config: context.config,
    };
    let list_str = write_list(&item_vec, &fmt)?;
    let mut prefix = format!("{}|{}|", mover, list_str);

    if !ret_str.is_empty() {
        if prefix.contains('\n') {
            prefix.push('\n');
            prefix.push_str(&argument_offset.to_string(context.config));
        } else {
            prefix.push(' ');
        }
        prefix.push_str(&ret_str);
    }
    // 1 = space between `|...|` and body.
    let extra_offset = last_line_width(&prefix) + 1;

    Some((prefix, extra_offset))
}

// Rewriting closure which is placed at the end of the function call's arg.
// Returns `None` if the reformatted closure 'looks bad'.
pub fn rewrite_last_closure(
    context: &RewriteContext,
    expr: &ast::Expr,
    shape: Shape,
) -> Option<String> {
    if let ast::ExprKind::Closure(capture, ref fn_decl, ref body, _) = expr.node {
        let body = match body.node {
            ast::ExprKind::Block(ref block)
                if !is_unsafe_block(block) && is_simple_block(block, context.codemap) =>
            {
                stmt_expr(&block.stmts[0]).unwrap_or(body)
            }
            _ => body,
        };
        let (prefix, extra_offset) =
            rewrite_closure_fn_decl(capture, fn_decl, body, expr.span, context, shape)?;
        // If the closure goes multi line before its body, do not overflow the closure.
        if prefix.contains('\n') {
            return None;
        }
        // If we are inside macro, we do not want to add or remove block from closure body.
        if context.inside_macro {
            return expr.rewrite(context, shape);
        }

        let body_shape = shape.offset_left(extra_offset)?;

        // We force to use block for the body of the closure for certain kinds of expressions.
        if is_block_closure_forced(body) {
            return rewrite_closure_with_block(body, &prefix, context, body_shape).and_then(
                |body_str| {
                    // If the expression can fit in a single line, we need not force block closure.
                    if body_str.lines().count() <= 7 {
                        match rewrite_closure_expr(body, &prefix, context, shape) {
                            Some(ref single_line_body_str)
                                if !single_line_body_str.contains('\n') =>
                            {
                                Some(single_line_body_str.clone())
                            }
                            _ => Some(body_str),
                        }
                    } else {
                        Some(body_str)
                    }
                },
            );
        }

        // When overflowing the closure which consists of a single control flow expression,
        // force to use block if its condition uses multi line.
        let is_multi_lined_cond = rewrite_cond(context, body, body_shape)
            .map(|cond| cond.contains('\n') || cond.len() > body_shape.width)
            .unwrap_or(false);
        if is_multi_lined_cond {
            return rewrite_closure_with_block(body, &prefix, context, body_shape);
        }

        // Seems fine, just format the closure in usual manner.
        return expr.rewrite(context, shape);
    }
    None
}

/// Returns true if the given vector of arguments has more than one `ast::ExprKind::Closure`.
pub fn args_have_many_closure<T>(args: &[&T]) -> bool
where
    T: ToExpr,
{
    args.iter()
        .filter(|arg| {
            arg.to_expr()
                .map(|e| match e.node {
                    ast::ExprKind::Closure(..) => true,
                    _ => false,
                })
                .unwrap_or(false)
        })
        .count() > 1
}

fn is_block_closure_forced(expr: &ast::Expr) -> bool {
    match expr.node {
        ast::ExprKind::If(..)
        | ast::ExprKind::IfLet(..)
        | ast::ExprKind::Loop(..)
        | ast::ExprKind::While(..)
        | ast::ExprKind::WhileLet(..)
        | ast::ExprKind::ForLoop(..) => true,
        ast::ExprKind::AddrOf(_, ref expr)
        | ast::ExprKind::Box(ref expr)
        | ast::ExprKind::Try(ref expr)
        | ast::ExprKind::Unary(_, ref expr)
        | ast::ExprKind::Cast(ref expr, _) => is_block_closure_forced(expr),
        _ => false,
    }
}

fn and_one_line(x: Option<String>) -> Option<String> {
    x.and_then(|x| if x.contains('\n') { None } else { Some(x) })
}
<commit_msg>Verify whether adding block is safe in rewrite_closure_with_block()<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use syntax::{ast, ptr};
use syntax::codemap::Span;
use syntax::parse::classify;

use codemap::SpanUtils;
use expr::{block_contains_comment, is_simple_block, is_unsafe_block, rewrite_cond, ToExpr};
use items::{span_hi_for_arg, span_lo_for_arg};
use lists::{definitive_tactic, itemize_list, write_list, DefinitiveListTactic, ListFormatting,
            ListTactic, Separator, SeparatorPlace, SeparatorTactic};
use rewrite::{Rewrite, RewriteContext};
use shape::Shape;
use utils::{last_line_width, left_most_sub_expr, stmt_expr};

// This module is pretty messy because of the rules around closures and blocks:
// FIXME - the below is probably no longer true in full.
//   * if there is a return type, then there must be braces,
//   * given a closure with braces, whether that is parsed to give an inner block
//     or not depends on if there is a return type and if there are statements
//     in that block,
//   * if the first expression in the body ends with a block (i.e., is a
//     statement without needing a semi-colon), then adding or removing braces
//     can change whether it is treated as an expression or statement.

pub fn rewrite_closure(
    capture: ast::CaptureBy,
    fn_decl: &ast::FnDecl,
    body: &ast::Expr,
    span: Span,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    debug!("rewrite_closure {:?}", body);

    let (prefix, extra_offset) =
        rewrite_closure_fn_decl(capture, fn_decl, body, span, context, shape)?;
    // 1 = space between `|...|` and body.
    let body_shape = shape.offset_left(extra_offset)?;

    if let ast::ExprKind::Block(ref block) = body.node {
        // The body of the closure is an empty block.
        if block.stmts.is_empty() && !block_contains_comment(block, context.codemap) {
            return Some(format!("{} {{}}", prefix));
        }

        let result = match fn_decl.output {
            ast::FunctionRetTy::Default(_) => {
                try_rewrite_without_block(body, &prefix, context, shape, body_shape)
            }
            _ => None,
        };

        result.or_else(|| {
            // Either we require a block, or tried without and failed.
            rewrite_closure_block(block, &prefix, context, body_shape)
        })
    } else {
        rewrite_closure_expr(body, &prefix, context, body_shape).or_else(|| {
            // The closure originally had a non-block expression, but we can't fit on
            // one line, so we'll insert a block.
            rewrite_closure_with_block(body, &prefix, context, body_shape)
        })
    }
}

fn try_rewrite_without_block(
    expr: &ast::Expr,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
    body_shape: Shape,
) -> Option<String> {
    let expr = get_inner_expr(expr, prefix, context);

    if is_block_closure_forced(expr) {
        rewrite_closure_with_block(expr, prefix, context, shape)
    } else {
        rewrite_closure_expr(expr, prefix, context, body_shape)
    }
}

fn get_inner_expr<'a>(
    expr: &'a ast::Expr,
    prefix: &str,
    context: &RewriteContext,
) -> &'a ast::Expr {
    if let ast::ExprKind::Block(ref block) = expr.node {
        if !needs_block(block, prefix, context) {
            // block.stmts.len() == 1
            if let Some(expr) = stmt_expr(&block.stmts[0]) {
                return get_inner_expr(expr, prefix, context);
            }
        }
    }

    expr
}

// Figure out if a block is necessary.
fn needs_block(block: &ast::Block, prefix: &str, context: &RewriteContext) -> bool {
    is_unsafe_block(block) || block.stmts.len() > 1 || context.inside_macro
        || block_contains_comment(block, context.codemap) || prefix.contains('\n')
}

// Rewrite closure with a single expression wrapping its body with block.
fn rewrite_closure_with_block(
    body: &ast::Expr,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    let left_most = left_most_sub_expr(body);
    let veto_block = left_most != body && !classify::expr_requires_semi_to_be_stmt(left_most);
    if veto_block {
        return None;
    }

    let block = ast::Block {
        stmts: vec![
            ast::Stmt {
                id: ast::NodeId::new(0),
                node: ast::StmtKind::Expr(ptr::P(body.clone())),
                span: body.span,
            },
        ],
        id: ast::NodeId::new(0),
        rules: ast::BlockCheckMode::Default,
        span: body.span,
    };
    let block = ::expr::rewrite_block_with_visitor(context, "", &block, shape, false)?;
    Some(format!("{} {}", prefix, block))
}

// Rewrite closure with a single expression without wrapping its body with block.
fn rewrite_closure_expr(
    expr: &ast::Expr,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    let mut rewrite = expr.rewrite(context, shape);
    rewrite = rewrite.and_then(|rw| {
        if context.config.force_multiline_blocks() && rw.contains('\n') {
            None
        } else {
            Some(rw)
        }
    });
    rewrite.map(|rw| format!("{} {}", prefix, rw))
}

// Rewrite closure whose body is block.
fn rewrite_closure_block(
    block: &ast::Block,
    prefix: &str,
    context: &RewriteContext,
    shape: Shape,
) -> Option<String> {
    let block_shape = shape.block();
    let block_str = block.rewrite(context, block_shape)?;
    Some(format!("{} {}", prefix, block_str))
}

// Return type is (prefix, extra_offset)
fn rewrite_closure_fn_decl(
    capture: ast::CaptureBy,
    fn_decl: &ast::FnDecl,
    body: &ast::Expr,
    span: Span,
    context: &RewriteContext,
    shape: Shape,
) -> Option<(String, usize)> {
    let mover = if capture == ast::CaptureBy::Value {
        "move "
    } else {
        ""
    };
    // 4 = "|| {".len(), which is overconservative when the closure consists of
    // a single expression.
    let nested_shape = shape.shrink_left(mover.len())?.sub_width(4)?;

    // 1 = |
    let argument_offset = nested_shape.indent + 1;
    let arg_shape = nested_shape.offset_left(1)?.visual_indent(0);
    let ret_str = fn_decl.output.rewrite(context, arg_shape)?;

    let arg_items = itemize_list(
        context.codemap,
        fn_decl.inputs.iter(),
        "|",
        ",",
        |arg| span_lo_for_arg(arg),
        |arg| span_hi_for_arg(context, arg),
        |arg| arg.rewrite(context, arg_shape),
        context.codemap.span_after(span, "|"),
        body.span.lo(),
        false,
    );
    let item_vec = arg_items.collect::<Vec<_>>();
    // 1 = space between arguments and return type.
    let horizontal_budget = nested_shape
        .width
        .checked_sub(ret_str.len() + 1)
        .unwrap_or(0);
    let tactic = definitive_tactic(
        &item_vec,
        ListTactic::HorizontalVertical,
        Separator::Comma,
        horizontal_budget,
    );
    let arg_shape = match tactic {
        DefinitiveListTactic::Horizontal => arg_shape.sub_width(ret_str.len() + 1)?,
        _ => arg_shape,
    };

    let fmt = ListFormatting {
        tactic: tactic,
        separator: ",",
        trailing_separator: SeparatorTactic::Never,
        separator_place: SeparatorPlace::Back,
        shape: arg_shape,
        ends_with_newline: false,
        preserve_newline: true,
        config: context.config,
    };
    let list_str = write_list(&item_vec, &fmt)?;
    let mut prefix = format!("{}|{}|", mover, list_str);

    if !ret_str.is_empty() {
        if prefix.contains('\n') {
            prefix.push('\n');
            prefix.push_str(&argument_offset.to_string(context.config));
        } else {
            prefix.push(' ');
        }
        prefix.push_str(&ret_str);
    }
    // 1 = space between `|...|` and body.
    let extra_offset = last_line_width(&prefix) + 1;

    Some((prefix, extra_offset))
}

// Rewriting closure which is placed at the end of the function call's arg.
// Returns `None` if the reformatted closure 'looks bad'.
pub fn rewrite_last_closure(
    context: &RewriteContext,
    expr: &ast::Expr,
    shape: Shape,
) -> Option<String> {
    if let ast::ExprKind::Closure(capture, ref fn_decl, ref body, _) = expr.node {
        let body = match body.node {
            ast::ExprKind::Block(ref block)
                if !is_unsafe_block(block) && is_simple_block(block, context.codemap) =>
            {
                stmt_expr(&block.stmts[0]).unwrap_or(body)
            }
            _ => body,
        };
        let (prefix, extra_offset) =
            rewrite_closure_fn_decl(capture, fn_decl, body, expr.span, context, shape)?;
        // If the closure goes multi line before its body, do not overflow the closure.
        if prefix.contains('\n') {
            return None;
        }
        // If we are inside macro, we do not want to add or remove block from closure body.
        if context.inside_macro {
            return expr.rewrite(context, shape);
        }

        let body_shape = shape.offset_left(extra_offset)?;

        // We force to use block for the body of the closure for certain kinds of expressions.
        if is_block_closure_forced(body) {
            return rewrite_closure_with_block(body, &prefix, context, body_shape).and_then(
                |body_str| {
                    // If the expression can fit in a single line, we need not force block closure.
                    if body_str.lines().count() <= 7 {
                        match rewrite_closure_expr(body, &prefix, context, shape) {
                            Some(ref single_line_body_str)
                                if !single_line_body_str.contains('\n') =>
                            {
                                Some(single_line_body_str.clone())
                            }
                            _ => Some(body_str),
                        }
                    } else {
                        Some(body_str)
                    }
                },
            );
        }

        // When overflowing the closure which consists of a single control flow expression,
        // force to use block if its condition uses multi line.
        let is_multi_lined_cond = rewrite_cond(context, body, body_shape)
            .map(|cond| cond.contains('\n') || cond.len() > body_shape.width)
            .unwrap_or(false);
        if is_multi_lined_cond {
            return rewrite_closure_with_block(body, &prefix, context, body_shape);
        }

        // Seems fine, just format the closure in usual manner.
        return expr.rewrite(context, shape);
    }
    None
}

/// Returns true if the given vector of arguments has more than one `ast::ExprKind::Closure`.
pub fn args_have_many_closure<T>(args: &[&T]) -> bool
where
    T: ToExpr,
{
    args.iter()
        .filter(|arg| {
            arg.to_expr()
                .map(|e| match e.node {
                    ast::ExprKind::Closure(..) => true,
                    _ => false,
                })
                .unwrap_or(false)
        })
        .count() > 1
}

fn is_block_closure_forced(expr: &ast::Expr) -> bool {
    match expr.node {
        ast::ExprKind::If(..)
        | ast::ExprKind::IfLet(..)
        | ast::ExprKind::Loop(..)
        | ast::ExprKind::While(..)
        | ast::ExprKind::WhileLet(..)
        | ast::ExprKind::ForLoop(..) => true,
        ast::ExprKind::AddrOf(_, ref expr)
        | ast::ExprKind::Box(ref expr)
        | ast::ExprKind::Try(ref expr)
        | ast::ExprKind::Unary(_, ref expr)
        | ast::ExprKind::Cast(ref expr, _) => is_block_closure_forced(expr),
        _ => false,
    }
}

fn and_one_line(x: Option<String>) -> Option<String> {
    x.and_then(|x| if x.contains('\n') { None } else { Some(x) })
}
<|endoftext|>