<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 std::num::Float;
use {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
///
/// ```notrust
/// 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 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 struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    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");

        let repr = 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))
        };
        Gamma { repr: repr }
    }
}

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. / (9. * d).sqrt(),
            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.repr {
            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) * u.powf(&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 ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                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 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 struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // 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.
    DoFExactlyOne,
    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 {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
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.repr {
            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 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 {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    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 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 {
    chi: ChiSquared,
    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 distributions::{Sample, IndependentSample};
    use {Rng, task_rng};
    use super::{ChiSquared, StudentT, FisherF};

    #[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_chi_squared_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 {
    extern crate test;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use {XorShiftRng, RAND_BENCH_N};
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = XorShiftRng::new().unwrap();

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

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = XorShiftRng::new().unwrap();

        b.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }
}
<commit_msg>Have floating point functions take their parameters by value.<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 std::num::Float;
use {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
///
/// ```notrust
/// 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 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 struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    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");

        let repr = 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))
        };
        Gamma { repr: repr }
    }
}

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. / (9. * d).sqrt(),
            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.repr {
            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) * u.powf(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 ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                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 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 struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // 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.
    DoFExactlyOne,
    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 {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
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.repr {
            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 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 {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    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 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 {
    chi: ChiSquared,
    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 distributions::{Sample, IndependentSample};
    use {Rng, task_rng};
    use super::{ChiSquared, StudentT, FisherF};

    #[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_chi_squared_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 {
    extern crate test;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use {XorShiftRng, RAND_BENCH_N};
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = XorShiftRng::new().unwrap();

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

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = XorShiftRng::new().unwrap();

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


<commit_before>//! # Default paths-based resolver
//!
//! Resolves packages based on a set of paths.
//!
//! These paths have the following structure:
//!
//! * `<root>/<package>/<last>.reproto`
//! * `<root>/<package>/<last>/<version>.reproto`
//!
//! The second form is only used when a version requirement is present.

use core::{Object, PathObject, RpPackage, RpRequiredPackage, Version, VersionReq};
use errors::*;
use resolver::{Resolved, ResolvedByPrefix, Resolver};
use std::collections::{BTreeMap, HashMap, LinkedList};
use std::ffi::OsStr;
use std::fs;
use std::path::{Path, PathBuf};

const EXT: &str = "reproto";

pub struct Paths {
    /// Paths to perform lookups in.
    paths: Vec<PathBuf>,
    /// Entries which are locally published.
    published: HashMap<RpPackage, Version>,
}

impl Paths {
    pub fn new(paths: Vec<PathBuf>, published: HashMap<RpPackage, Version>) -> Paths {
        Paths {
            paths: paths,
            published: published,
        }
    }

    fn parse_stem<'a>(&self, stem: &'a str) -> Result<(&'a str, Option<Version>)> {
        let mut it = stem.splitn(2, '-');

        if let (Some(name_base), Some(name_version)) = (it.next(), it.next()) {
            let version = Version::parse(name_version).map_err(
                |_| format!("bad version"),
            )?;

            return Ok((name_base, Some(version)));
        }

        Ok((stem, None))
    }

    /// Finds the published version from most to least specific package.
    pub fn find_published_version(&self, package: &RpPackage) -> Option<&Version> {
        if let Some(version) = self.published.get(package) {
            return Some(version);
        }

        let mut it = package.parts.iter();

        while let Some(_) = it.next_back() {
            let package = RpPackage::new(it.as_slice().to_vec());

            if let Some(version) = self.published.get(&package) {
                return Some(version);
            }
        }

        None
    }

    /// Find any matching versions.
    pub fn find_by_version_req(
        &self,
        path: &Path,
        base: &str,
        package: &RpPackage,
        version_req: &VersionReq,
    ) -> Result<Vec<Resolved>> {
        let mut files: BTreeMap<Option<Version>, Box<Object>> = BTreeMap::new();

        for e in fs::read_dir(path)? {
            let p = e?.path();

            // only look for files
            if !p.is_file() {
                continue;
            }

            if p.extension().map(|ext| ext != EXT).unwrap_or(true) {
                continue;
            }

            if let Some(stem) = p.file_stem().and_then(OsStr::to_str) {
                let (name_base, version) = self.parse_stem(stem).map_err(|m| {
                    format!("{}: {}", p.display(), m)
                })?;

                if name_base != base {
                    continue;
                }

                let version = version.or_else(|| self.find_published_version(package).cloned());

                // versioned matches by requirement.
                if let Some(version) = version {
                    if version_req.matches(&version) {
                        files.insert(Some(version), Box::new(PathObject::new(None, &p)));
                    }

                    continue;
                }

                // unversioned only matches by wildcard.
                if version_req.matches_any() {
                    files.insert(None, Box::new(PathObject::new(None, &p)));
                }
            }
        }

        Ok(files.into_iter().map(Resolved::from_pair).collect())
    }

    /// Load .reproto file from the given package path if present.
    fn load_file(&self, path: &Path, prefix: &RpPackage) -> Option<ResolvedByPrefix> {
        let mut file = path.to_owned();
        file.set_extension(EXT);

        if file.is_file() {
            Some(ResolvedByPrefix {
                package: prefix.clone(),
                object: Box::new(PathObject::new(None, &file)),
            })
        } else {
            None
        }
    }

    /// Load .reproto file from path if valid reproto file and present.
    fn load_from_path(&self, path: &Path, prefix: RpPackage) -> Result<Option<ResolvedByPrefix>> {
        if path.extension().map(|ext| ext != EXT).unwrap_or(true) {
            debug!("skipping wrong file extension: {}", path.display());
            return Ok(None);
        }

        let stem = path.file_stem().and_then(OsStr::to_str).ok_or_else(|| {
            format!("illegal path: {}", path.display())
        })?;

        let (stem, version) = self.parse_stem(stem)?;

        if version.is_some() {
            debug!("skipping versioned: {}", path.display());
            return Ok(None);
        }

        let package = prefix.join_part(stem);

        Ok(Some(ResolvedByPrefix {
            package: package,
            object: Box::new(PathObject::new(None, &path)),
        }))
    }

    /// Find all packages by prefix.
    pub fn find_by_prefix(&self, path: &Path, prefix: &RpPackage) -> Result<Vec<ResolvedByPrefix>> {
        let mut files = Vec::new();

        files.extend(self.load_file(path, prefix));

        if !path.is_dir() {
            return Ok(files);
        }

        let mut queue = LinkedList::new();
        queue.push_back((prefix.clone(), path.to_owned()));

        while let Some((prefix, path)) = queue.pop_front() {
            for entry in fs::read_dir(path)? {
                let entry = entry?;
                let path = entry.path();

                if path.is_file() {
                    files.extend(self.load_from_path(&path, prefix.clone())?);
                    continue;
                }

                if path.is_dir() {
                    let file_name = entry.file_name();

                    let name = file_name.to_str().ok_or_else(|| {
                        format!("illegal path: {}", path.display())
                    })?;

                    queue.push_back((prefix.clone().join_part(name), path));
                    continue;
                }
            }
        }

        Ok(files)
    }
}

impl Resolver for Paths {
    fn resolve(&mut self, package: &RpRequiredPackage) -> Result<Vec<Resolved>> {
        let mut files = Vec::new();

        for path in &self.paths {
            let mut path: PathBuf = path.to_owned();
            let mut it = package.package.parts.iter().peekable();

            while let Some(step) = it.next() {
                if it.peek().is_none() {
                    if path.is_dir() {
                        files.extend(self.find_by_version_req(
                            &path,
                            step,
                            &package.package,
                            &package.version_req,
                        )?);
                    }

                    break;
                }

                path = path.join(step);
            }
        }

        Ok(files)
    }

    fn resolve_by_prefix(&mut self, prefix: &RpPackage) -> Result<Vec<ResolvedByPrefix>> {
        let mut files = Vec::new();

        for path in &self.paths {
            let mut it = prefix.parts.iter();
            let path = it.fold(path.to_owned(), |p, part| p.join(part));
            files.extend(self.find_by_prefix(&path, prefix)?);
        }

        Ok(files)
    }
}
<commit_msg>[repository] removed useless mut<commit_after>//! # Default paths-based resolver
//!
//! Resolves packages based on a set of paths.
//!
//! These paths have the following structure:
//!
//! * `<root>/<package>/<last>.reproto`
//! * `<root>/<package>/<last>/<version>.reproto`
//!
//! The second form is only used when a version requirement is present.

use core::{Object, PathObject, RpPackage, RpRequiredPackage, Version, VersionReq};
use errors::*;
use resolver::{Resolved, ResolvedByPrefix, Resolver};
use std::collections::{BTreeMap, HashMap, LinkedList};
use std::ffi::OsStr;
use std::fs;
use std::path::{Path, PathBuf};

const EXT: &str = "reproto";

pub struct Paths {
    /// Paths to perform lookups in.
    paths: Vec<PathBuf>,
    /// Entries which are locally published.
    published: HashMap<RpPackage, Version>,
}

impl Paths {
    pub fn new(paths: Vec<PathBuf>, published: HashMap<RpPackage, Version>) -> Paths {
        Paths {
            paths: paths,
            published: published,
        }
    }

    fn parse_stem<'a>(&self, stem: &'a str) -> Result<(&'a str, Option<Version>)> {
        let mut it = stem.splitn(2, '-');

        if let (Some(name_base), Some(name_version)) = (it.next(), it.next()) {
            let version = Version::parse(name_version).map_err(
                |_| format!("bad version"),
            )?;

            return Ok((name_base, Some(version)));
        }

        Ok((stem, None))
    }

    /// Finds the published version from most to least specific package.
    pub fn find_published_version(&self, package: &RpPackage) -> Option<&Version> {
        if let Some(version) = self.published.get(package) {
            return Some(version);
        }

        let mut it = package.parts.iter();

        while let Some(_) = it.next_back() {
            let package = RpPackage::new(it.as_slice().to_vec());

            if let Some(version) = self.published.get(&package) {
                return Some(version);
            }
        }

        None
    }

    /// Find any matching versions.
    pub fn find_by_version_req(
        &self,
        path: &Path,
        base: &str,
        package: &RpPackage,
        version_req: &VersionReq,
    ) -> Result<Vec<Resolved>> {
        let mut files: BTreeMap<Option<Version>, Box<Object>> = BTreeMap::new();

        for e in fs::read_dir(path)? {
            let p = e?.path();

            // only look for files
            if !p.is_file() {
                continue;
            }

            if p.extension().map(|ext| ext != EXT).unwrap_or(true) {
                continue;
            }

            if let Some(stem) = p.file_stem().and_then(OsStr::to_str) {
                let (name_base, version) = self.parse_stem(stem).map_err(|m| {
                    format!("{}: {}", p.display(), m)
                })?;

                if name_base != base {
                    continue;
                }

                let version = version.or_else(|| self.find_published_version(package).cloned());

                // versioned matches by requirement.
                if let Some(version) = version {
                    if version_req.matches(&version) {
                        files.insert(Some(version), Box::new(PathObject::new(None, &p)));
                    }

                    continue;
                }

                // unversioned only matches by wildcard.
                if version_req.matches_any() {
                    files.insert(None, Box::new(PathObject::new(None, &p)));
                }
            }
        }

        Ok(files.into_iter().map(Resolved::from_pair).collect())
    }

    /// Load .reproto file from the given package path if present.
    fn load_file(&self, path: &Path, prefix: &RpPackage) -> Option<ResolvedByPrefix> {
        let mut file = path.to_owned();
        file.set_extension(EXT);

        if file.is_file() {
            Some(ResolvedByPrefix {
                package: prefix.clone(),
                object: Box::new(PathObject::new(None, &file)),
            })
        } else {
            None
        }
    }

    /// Load .reproto file from path if valid reproto file and present.
    fn load_from_path(&self, path: &Path, prefix: RpPackage) -> Result<Option<ResolvedByPrefix>> {
        if path.extension().map(|ext| ext != EXT).unwrap_or(true) {
            debug!("skipping wrong file extension: {}", path.display());
            return Ok(None);
        }

        let stem = path.file_stem().and_then(OsStr::to_str).ok_or_else(|| {
            format!("illegal path: {}", path.display())
        })?;

        let (stem, version) = self.parse_stem(stem)?;

        if version.is_some() {
            debug!("skipping versioned: {}", path.display());
            return Ok(None);
        }

        let package = prefix.join_part(stem);

        Ok(Some(ResolvedByPrefix {
            package: package,
            object: Box::new(PathObject::new(None, &path)),
        }))
    }

    /// Find all packages by prefix.
    pub fn find_by_prefix(&self, path: &Path, prefix: &RpPackage) -> Result<Vec<ResolvedByPrefix>> {
        let mut files = Vec::new();

        files.extend(self.load_file(path, prefix));

        if !path.is_dir() {
            return Ok(files);
        }

        let mut queue = LinkedList::new();
        queue.push_back((prefix.clone(), path.to_owned()));

        while let Some((prefix, path)) = queue.pop_front() {
            for entry in fs::read_dir(path)? {
                let entry = entry?;
                let path = entry.path();

                if path.is_file() {
                    files.extend(self.load_from_path(&path, prefix.clone())?);
                    continue;
                }

                if path.is_dir() {
                    let file_name = entry.file_name();

                    let name = file_name.to_str().ok_or_else(|| {
                        format!("illegal path: {}", path.display())
                    })?;

                    queue.push_back((prefix.clone().join_part(name), path));
                    continue;
                }
            }
        }

        Ok(files)
    }
}

impl Resolver for Paths {
    fn resolve(&mut self, package: &RpRequiredPackage) -> Result<Vec<Resolved>> {
        let mut files = Vec::new();

        for path in &self.paths {
            let mut path: PathBuf = path.to_owned();
            let mut it = package.package.parts.iter().peekable();

            while let Some(step) = it.next() {
                if it.peek().is_none() {
                    if path.is_dir() {
                        files.extend(self.find_by_version_req(
                            &path,
                            step,
                            &package.package,
                            &package.version_req,
                        )?);
                    }

                    break;
                }

                path = path.join(step);
            }
        }

        Ok(files)
    }

    fn resolve_by_prefix(&mut self, prefix: &RpPackage) -> Result<Vec<ResolvedByPrefix>> {
        let mut files = Vec::new();

        for path in &self.paths {
            let path = prefix.parts.iter().fold(path.to_owned(), |p, part| p.join(part));
            files.extend(self.find_by_prefix(&path, prefix)?);
        }

        Ok(files)
    }
}
<|endoftext|>


<commit_before>use crate::utils::{
    has_drop, in_macro, is_copy, match_type, paths, snippet_opt, span_lint_hir, span_lint_hir_and_then,
    walk_ptrs_ty_depth,
};
use if_chain::if_chain;
use matches::matches;
use rustc::hir::intravisit::FnKind;
use rustc::hir::{def_id, Body, FnDecl, HirId};
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::mir::{
    self, traversal,
    visit::{MutatingUseContext, PlaceContext, Visitor},
    TerminatorKind,
};
use rustc::ty::{self, Ty};
use rustc::{declare_lint_pass, declare_tool_lint};
use rustc_errors::Applicability;
use std::convert::TryFrom;
use syntax::source_map::{BytePos, Span};

macro_rules! unwrap_or_continue {
    ($x:expr) => {
        match $x {
            Some(x) => x,
            None => continue,
        }
    };
}

declare_clippy_lint! {
    /// **What it does:** Checks for a redudant `clone()` (and its relatives) which clones an owned
    /// value that is going to be dropped without further use.
    ///
    /// **Why is this bad?** It is not always possible for the compiler to eliminate useless
    /// allocations and deallocations generated by redundant `clone()`s.
    ///
    /// **Known problems:**
    ///
    /// * Suggestions made by this lint could require NLL to be enabled.
    /// * False-positive if there is a borrow preventing the value from moving out.
    ///
    /// ```rust
    /// let x = String::new();
    ///
    /// let y = &x;
    ///
    /// foo(x.clone()); // This lint suggests to remove this `clone()`
    /// ```
    ///
    /// **Example:**
    /// ```rust
    /// {
    ///     let x = Foo::new();
    ///     call(x.clone());
    ///     call(x.clone()); // this can just pass `x`
    /// }
    ///
    /// ["lorem", "ipsum"].join(" ").to_string()
    ///
    /// Path::new("/a/b").join("c").to_path_buf()
    /// ```
    pub REDUNDANT_CLONE,
    nursery,
    "`clone()` of an owned value that is going to be dropped immediately"
}

declare_lint_pass!(RedundantClone => [REDUNDANT_CLONE]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for RedundantClone {
    #[allow(clippy::too_many_lines)]
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        _: FnKind<'tcx>,
        _: &'tcx FnDecl,
        body: &'tcx Body,
        _: Span,
        _: HirId,
    ) {
        let def_id = cx.tcx.hir().body_owner_def_id(body.id());
        let mir = cx.tcx.optimized_mir(def_id);

        for (bb, bbdata) in mir.basic_blocks().iter_enumerated() {
            let terminator = bbdata.terminator();

            if in_macro(terminator.source_info.span) {
                continue;
            }

            // Give up on loops
            if terminator.successors().any(|s| *s == bb) {
                continue;
            }

            let (fn_def_id, arg, arg_ty, _) = unwrap_or_continue!(is_call_with_ref_arg(cx, mir, &terminator.kind));

            let from_borrow = cx.match_def_path(fn_def_id, &paths::CLONE_TRAIT_METHOD)
                || cx.match_def_path(fn_def_id, &paths::TO_OWNED_METHOD)
                || (cx.match_def_path(fn_def_id, &paths::TO_STRING_METHOD) && match_type(cx, arg_ty, &paths::STRING));

            let from_deref = !from_borrow
                && (cx.match_def_path(fn_def_id, &paths::PATH_TO_PATH_BUF)
                    || cx.match_def_path(fn_def_id, &paths::OS_STR_TO_OS_STRING));

            if !from_borrow && !from_deref {
                continue;
            }

            // _1 in MIR `{ _2 = &_1; clone(move _2); }` or `{ _2 = _1; to_path_buf(_2); } (from_deref)
            // In case of `from_deref`, `arg` is already a reference since it is `deref`ed in the previous
            // block.
            let (cloned, cannot_move_out) = unwrap_or_continue!(find_stmt_assigns_to(
                cx,
                mir,
                arg,
                from_borrow,
                bbdata.statements.iter()
            ));

            if from_borrow && cannot_move_out {
                continue;
            }

            // _1 in MIR `{ _2 = &_1; _3 = deref(move _2); } -> { _4 = _3; to_path_buf(move _4); }`
            let referent = if from_deref {
                let ps = mir.predecessors_for(bb);
                if ps.len() != 1 {
                    continue;
                }
                let pred_terminator = mir[ps[0]].terminator();

                let pred_arg = if_chain! {
                    if let Some((pred_fn_def_id, pred_arg, pred_arg_ty, Some(res))) =
                        is_call_with_ref_arg(cx, mir, &pred_terminator.kind);
                    if *res == mir::Place::Base(mir::PlaceBase::Local(cloned));
                    if cx.match_def_path(pred_fn_def_id, &paths::DEREF_TRAIT_METHOD);
                    if match_type(cx, pred_arg_ty, &paths::PATH_BUF)
                        || match_type(cx, pred_arg_ty, &paths::OS_STRING);
                    then {
                        pred_arg
                    } else {
                        continue;
                    }
                };

                let (local, cannot_move_out) = unwrap_or_continue!(find_stmt_assigns_to(
                    cx,
                    mir,
                    pred_arg,
                    true,
                    mir[ps[0]].statements.iter()
                ));
                if cannot_move_out {
                    continue;
                }
                local
            } else {
                cloned
            };

            let used_later = traversal::ReversePostorder::new(&mir, bb).skip(1).any(|(tbb, tdata)| {
                // Give up on loops
                if tdata.terminator().successors().any(|s| *s == bb) {
                    return true;
                }

                let mut vis = LocalUseVisitor {
                    local: referent,
                    used_other_than_drop: false,
                };
                vis.visit_basic_block_data(tbb, tdata);
                vis.used_other_than_drop
            });

            if !used_later {
                let span = terminator.source_info.span;
                let node = if let mir::ClearCrossCrate::Set(scope_local_data) = &mir.source_scope_local_data {
                    scope_local_data[terminator.source_info.scope].lint_root
                } else {
                    unreachable!()
                };

                if_chain! {
                    if let Some(snip) = snippet_opt(cx, span);
                    if let Some(dot) = snip.rfind('.');
                    then {
                        let sugg_span = span.with_lo(
                            span.lo() + BytePos(u32::try_from(dot).unwrap())
                        );

                        span_lint_hir_and_then(cx, REDUNDANT_CLONE, node, sugg_span, "redundant clone", |db| {
                            db.span_suggestion(
                                sugg_span,
                                "remove this",
                                String::new(),
                                Applicability::MaybeIncorrect,
                            );
                            db.span_note(
                                span.with_hi(span.lo() + BytePos(u32::try_from(dot).unwrap())),
                                "this value is dropped without further use",
                            );
                        });
                    } else {
                        span_lint_hir(cx, REDUNDANT_CLONE, node, span, "redundant clone");
                    }
                }
            }
        }
    }
}

/// If `kind` is `y = func(x: &T)` where `T: !Copy`, returns `(DefId of func, x, T, y)`.
fn is_call_with_ref_arg<'tcx>(
    cx: &LateContext<'_, 'tcx>,
    mir: &'tcx mir::Mir<'tcx>,
    kind: &'tcx mir::TerminatorKind<'tcx>,
) -> Option<(def_id::DefId, mir::Local, Ty<'tcx>, Option<&'tcx mir::Place<'tcx>>)> {
    if_chain! {
        if let TerminatorKind::Call { func, args, destination, .. } = kind;
        if args.len() == 1;
        if let mir::Operand::Move(mir::Place::Base(mir::PlaceBase::Local(local))) = &args[0];
        if let ty::FnDef(def_id, _) = func.ty(&*mir, cx.tcx).sty;
        if let (inner_ty, 1) = walk_ptrs_ty_depth(args[0].ty(&*mir, cx.tcx));
        if !is_copy(cx, inner_ty);
        then {
            Some((def_id, *local, inner_ty, destination.as_ref().map(|(dest, _)| dest)))
        } else {
            None
        }
    }
}

type CannotMoveOut = bool;

/// Finds the first `to = (&)from`, and returns
/// ``Some((from, [`true` if `from` cannot be moved out]))``.
fn find_stmt_assigns_to<'a, 'tcx: 'a>(
    cx: &LateContext<'_, 'tcx>,
    mir: &mir::Mir<'tcx>,
    to: mir::Local,
    by_ref: bool,
    stmts: impl DoubleEndedIterator<Item = &'a mir::Statement<'tcx>>,
) -> Option<(mir::Local, CannotMoveOut)> {
    stmts
        .rev()
        .find_map(|stmt| {
            if let mir::StatementKind::Assign(mir::Place::Base(mir::PlaceBase::Local(local)), v) = &stmt.kind {
                if *local == to {
                    return Some(v);
                }
            }

            None
        })
        .and_then(|v| {
            if by_ref {
                if let mir::Rvalue::Ref(_, _, ref place) = **v {
                    return base_local_and_movability(cx, mir, place);
                }
            } else if let mir::Rvalue::Use(mir::Operand::Copy(ref place)) = **v {
                return base_local_and_movability(cx, mir, place);
            }
            None
        })
}

/// Extracts and returns the undermost base `Local` of given `place`. Returns `place` itself
/// if it is already a `Local`.
///
/// Also reports whether given `place` cannot be moved out.
fn base_local_and_movability<'tcx>(
    cx: &LateContext<'_, 'tcx>,
    mir: &mir::Mir<'tcx>,
    mut place: &mir::Place<'tcx>,
) -> Option<(mir::Local, CannotMoveOut)> {
    use rustc::mir::Place::*;
    use rustc::mir::PlaceBase;

    // Dereference. You cannot move things out from a borrowed value.
    let mut deref = false;
    // Accessing a field of an ADT that has `Drop`. Moving the field out will cause E0509.
    let mut field = false;

    loop {
        match place {
            Base(PlaceBase::Local(local)) => return Some((*local, deref || field)),
            Projection(proj) => {
                place = &proj.base;
                deref = deref || matches!(proj.elem, mir::ProjectionElem::Deref);
                if !field && matches!(proj.elem, mir::ProjectionElem::Field(..)) {
                    field = has_drop(cx, place.ty(&mir.local_decls, cx.tcx).ty);
                }
            },
            _ => return None,
        }
    }
}

struct LocalUseVisitor {
    local: mir::Local,
    used_other_than_drop: bool,
}

impl<'tcx> mir::visit::Visitor<'tcx> for LocalUseVisitor {
    fn visit_basic_block_data(&mut self, block: mir::BasicBlock, data: &mir::BasicBlockData<'tcx>) {
        let statements = &data.statements;
        for (statement_index, statement) in statements.iter().enumerate() {
            self.visit_statement(block, statement, mir::Location { block, statement_index });

            // Once flagged, skip remaining statements
            if self.used_other_than_drop {
                return;
            }
        }

        self.visit_terminator(
            block,
            data.terminator(),
            mir::Location {
                block,
                statement_index: statements.len(),
            },
        );
    }

    fn visit_local(&mut self, local: &mir::Local, ctx: PlaceContext<'tcx>, _: mir::Location) {
        match ctx {
            PlaceContext::MutatingUse(MutatingUseContext::Drop) | PlaceContext::NonUse(_) => return,
            _ => {},
        }

        if *local == self.local {
            self.used_other_than_drop = true;
        }
    }
}
<commit_msg>Auto merge of #4047 - rust-lang:rustup, r=Centril<commit_after>use crate::utils::{
    has_drop, in_macro, is_copy, match_type, paths, snippet_opt, span_lint_hir, span_lint_hir_and_then,
    walk_ptrs_ty_depth,
};
use if_chain::if_chain;
use matches::matches;
use rustc::hir::intravisit::FnKind;
use rustc::hir::{def_id, Body, FnDecl, HirId};
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::mir::{
    self, traversal,
    visit::{MutatingUseContext, PlaceContext, Visitor},
    TerminatorKind,
};
use rustc::ty::{self, Ty};
use rustc::{declare_lint_pass, declare_tool_lint};
use rustc_errors::Applicability;
use std::convert::TryFrom;
use syntax::source_map::{BytePos, Span};

macro_rules! unwrap_or_continue {
    ($x:expr) => {
        match $x {
            Some(x) => x,
            None => continue,
        }
    };
}

declare_clippy_lint! {
    /// **What it does:** Checks for a redudant `clone()` (and its relatives) which clones an owned
    /// value that is going to be dropped without further use.
    ///
    /// **Why is this bad?** It is not always possible for the compiler to eliminate useless
    /// allocations and deallocations generated by redundant `clone()`s.
    ///
    /// **Known problems:**
    ///
    /// * Suggestions made by this lint could require NLL to be enabled.
    /// * False-positive if there is a borrow preventing the value from moving out.
    ///
    /// ```rust
    /// let x = String::new();
    ///
    /// let y = &x;
    ///
    /// foo(x.clone()); // This lint suggests to remove this `clone()`
    /// ```
    ///
    /// **Example:**
    /// ```rust
    /// {
    ///     let x = Foo::new();
    ///     call(x.clone());
    ///     call(x.clone()); // this can just pass `x`
    /// }
    ///
    /// ["lorem", "ipsum"].join(" ").to_string()
    ///
    /// Path::new("/a/b").join("c").to_path_buf()
    /// ```
    pub REDUNDANT_CLONE,
    nursery,
    "`clone()` of an owned value that is going to be dropped immediately"
}

declare_lint_pass!(RedundantClone => [REDUNDANT_CLONE]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for RedundantClone {
    #[allow(clippy::too_many_lines)]
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        _: FnKind<'tcx>,
        _: &'tcx FnDecl,
        body: &'tcx Body,
        _: Span,
        _: HirId,
    ) {
        let def_id = cx.tcx.hir().body_owner_def_id(body.id());
        let mir = cx.tcx.optimized_mir(def_id);

        for (bb, bbdata) in mir.basic_blocks().iter_enumerated() {
            let terminator = bbdata.terminator();

            if in_macro(terminator.source_info.span) {
                continue;
            }

            // Give up on loops
            if terminator.successors().any(|s| *s == bb) {
                continue;
            }

            let (fn_def_id, arg, arg_ty, _) = unwrap_or_continue!(is_call_with_ref_arg(cx, mir, &terminator.kind));

            let from_borrow = cx.match_def_path(fn_def_id, &paths::CLONE_TRAIT_METHOD)
                || cx.match_def_path(fn_def_id, &paths::TO_OWNED_METHOD)
                || (cx.match_def_path(fn_def_id, &paths::TO_STRING_METHOD) && match_type(cx, arg_ty, &paths::STRING));

            let from_deref = !from_borrow
                && (cx.match_def_path(fn_def_id, &paths::PATH_TO_PATH_BUF)
                    || cx.match_def_path(fn_def_id, &paths::OS_STR_TO_OS_STRING));

            if !from_borrow && !from_deref {
                continue;
            }

            // _1 in MIR `{ _2 = &_1; clone(move _2); }` or `{ _2 = _1; to_path_buf(_2); } (from_deref)
            // In case of `from_deref`, `arg` is already a reference since it is `deref`ed in the previous
            // block.
            let (cloned, cannot_move_out) = unwrap_or_continue!(find_stmt_assigns_to(
                cx,
                mir,
                arg,
                from_borrow,
                bbdata.statements.iter()
            ));

            if from_borrow && cannot_move_out {
                continue;
            }

            // _1 in MIR `{ _2 = &_1; _3 = deref(move _2); } -> { _4 = _3; to_path_buf(move _4); }`
            let referent = if from_deref {
                let ps = mir.predecessors_for(bb);
                if ps.len() != 1 {
                    continue;
                }
                let pred_terminator = mir[ps[0]].terminator();

                let pred_arg = if_chain! {
                    if let Some((pred_fn_def_id, pred_arg, pred_arg_ty, Some(res))) =
                        is_call_with_ref_arg(cx, mir, &pred_terminator.kind);
                    if *res == mir::Place::Base(mir::PlaceBase::Local(cloned));
                    if cx.match_def_path(pred_fn_def_id, &paths::DEREF_TRAIT_METHOD);
                    if match_type(cx, pred_arg_ty, &paths::PATH_BUF)
                        || match_type(cx, pred_arg_ty, &paths::OS_STRING);
                    then {
                        pred_arg
                    } else {
                        continue;
                    }
                };

                let (local, cannot_move_out) = unwrap_or_continue!(find_stmt_assigns_to(
                    cx,
                    mir,
                    pred_arg,
                    true,
                    mir[ps[0]].statements.iter()
                ));
                if cannot_move_out {
                    continue;
                }
                local
            } else {
                cloned
            };

            let used_later = traversal::ReversePostorder::new(&mir, bb).skip(1).any(|(tbb, tdata)| {
                // Give up on loops
                if tdata.terminator().successors().any(|s| *s == bb) {
                    return true;
                }

                let mut vis = LocalUseVisitor {
                    local: referent,
                    used_other_than_drop: false,
                };
                vis.visit_basic_block_data(tbb, tdata);
                vis.used_other_than_drop
            });

            if !used_later {
                let span = terminator.source_info.span;
                let node = if let mir::ClearCrossCrate::Set(scope_local_data) = &mir.source_scope_local_data {
                    scope_local_data[terminator.source_info.scope].lint_root
                } else {
                    unreachable!()
                };

                if_chain! {
                    if let Some(snip) = snippet_opt(cx, span);
                    if let Some(dot) = snip.rfind('.');
                    then {
                        let sugg_span = span.with_lo(
                            span.lo() + BytePos(u32::try_from(dot).unwrap())
                        );

                        span_lint_hir_and_then(cx, REDUNDANT_CLONE, node, sugg_span, "redundant clone", |db| {
                            db.span_suggestion(
                                sugg_span,
                                "remove this",
                                String::new(),
                                Applicability::MaybeIncorrect,
                            );
                            db.span_note(
                                span.with_hi(span.lo() + BytePos(u32::try_from(dot).unwrap())),
                                "this value is dropped without further use",
                            );
                        });
                    } else {
                        span_lint_hir(cx, REDUNDANT_CLONE, node, span, "redundant clone");
                    }
                }
            }
        }
    }
}

/// If `kind` is `y = func(x: &T)` where `T: !Copy`, returns `(DefId of func, x, T, y)`.
fn is_call_with_ref_arg<'tcx>(
    cx: &LateContext<'_, 'tcx>,
    mir: &'tcx mir::Mir<'tcx>,
    kind: &'tcx mir::TerminatorKind<'tcx>,
) -> Option<(def_id::DefId, mir::Local, Ty<'tcx>, Option<&'tcx mir::Place<'tcx>>)> {
    if_chain! {
        if let TerminatorKind::Call { func, args, destination, .. } = kind;
        if args.len() == 1;
        if let mir::Operand::Move(mir::Place::Base(mir::PlaceBase::Local(local))) = &args[0];
        if let ty::FnDef(def_id, _) = func.ty(&*mir, cx.tcx).sty;
        if let (inner_ty, 1) = walk_ptrs_ty_depth(args[0].ty(&*mir, cx.tcx));
        if !is_copy(cx, inner_ty);
        then {
            Some((def_id, *local, inner_ty, destination.as_ref().map(|(dest, _)| dest)))
        } else {
            None
        }
    }
}

type CannotMoveOut = bool;

/// Finds the first `to = (&)from`, and returns
/// ``Some((from, [`true` if `from` cannot be moved out]))``.
fn find_stmt_assigns_to<'a, 'tcx: 'a>(
    cx: &LateContext<'_, 'tcx>,
    mir: &mir::Mir<'tcx>,
    to: mir::Local,
    by_ref: bool,
    stmts: impl DoubleEndedIterator<Item = &'a mir::Statement<'tcx>>,
) -> Option<(mir::Local, CannotMoveOut)> {
    stmts
        .rev()
        .find_map(|stmt| {
            if let mir::StatementKind::Assign(mir::Place::Base(mir::PlaceBase::Local(local)), v) = &stmt.kind {
                if *local == to {
                    return Some(v);
                }
            }

            None
        })
        .and_then(|v| {
            if by_ref {
                if let mir::Rvalue::Ref(_, _, ref place) = **v {
                    return base_local_and_movability(cx, mir, place);
                }
            } else if let mir::Rvalue::Use(mir::Operand::Copy(ref place)) = **v {
                return base_local_and_movability(cx, mir, place);
            }
            None
        })
}

/// Extracts and returns the undermost base `Local` of given `place`. Returns `place` itself
/// if it is already a `Local`.
///
/// Also reports whether given `place` cannot be moved out.
fn base_local_and_movability<'tcx>(
    cx: &LateContext<'_, 'tcx>,
    mir: &mir::Mir<'tcx>,
    mut place: &mir::Place<'tcx>,
) -> Option<(mir::Local, CannotMoveOut)> {
    use rustc::mir::Place::*;
    use rustc::mir::PlaceBase;

    // Dereference. You cannot move things out from a borrowed value.
    let mut deref = false;
    // Accessing a field of an ADT that has `Drop`. Moving the field out will cause E0509.
    let mut field = false;

    loop {
        match place {
            Base(PlaceBase::Local(local)) => return Some((*local, deref || field)),
            Projection(proj) => {
                place = &proj.base;
                deref = deref || matches!(proj.elem, mir::ProjectionElem::Deref);
                if !field && matches!(proj.elem, mir::ProjectionElem::Field(..)) {
                    field = has_drop(cx, place.ty(&mir.local_decls, cx.tcx).ty);
                }
            },
            _ => return None,
        }
    }
}

struct LocalUseVisitor {
    local: mir::Local,
    used_other_than_drop: bool,
}

impl<'tcx> mir::visit::Visitor<'tcx> for LocalUseVisitor {
    fn visit_basic_block_data(&mut self, block: mir::BasicBlock, data: &mir::BasicBlockData<'tcx>) {
        let statements = &data.statements;
        for (statement_index, statement) in statements.iter().enumerate() {
            self.visit_statement(statement, mir::Location { block, statement_index });

            // Once flagged, skip remaining statements
            if self.used_other_than_drop {
                return;
            }
        }

        self.visit_terminator(
            data.terminator(),
            mir::Location {
                block,
                statement_index: statements.len(),
            },
        );
    }

    fn visit_local(&mut self, local: &mir::Local, ctx: PlaceContext, _: mir::Location) {
        match ctx {
            PlaceContext::MutatingUse(MutatingUseContext::Drop) | PlaceContext::NonUse(_) => return,
            _ => {},
        }

        if *local == self.local {
            self.used_other_than_drop = true;
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use ipc_channel::ipc::IpcSharedMemory;
use png;
use stb_image::image as stb_image2;
use std::mem;
use util::vec::byte_swap;

// FIXME: Images must not be copied every frame. Instead we should atomically
// reference count them.

#[derive(Deserialize, Serialize)]
pub enum PixelFormat {
    K8,         // Luminance channel only
    KA8,        // Luminance + alpha
    RGB8,       // RGB, 8 bits per channel
    RGBA8,      // RGB + alpha, 8 bits per channel
}

#[derive(Deserialize, Serialize)]
pub struct Image {
    pub width: u32,
    pub height: u32,
    pub format: PixelFormat,
    pub bytes: IpcSharedMemory,
}

// TODO(pcwalton): Speed up with SIMD, or better yet, find some way to not do this.
fn byte_swap_and_premultiply(data: &mut [u8]) {
    let length = data.len();
    for i in (0..length).step_by(4) {
        let r = data[i + 2];
        let g = data[i + 1];
        let b = data[i + 0];
        let a = data[i + 3];
        data[i + 0] = ((r as u32) * (a as u32) / 255) as u8;
        data[i + 1] = ((g as u32) * (a as u32) / 255) as u8;
        data[i + 2] = ((b as u32) * (a as u32) / 255) as u8;
    }
}

pub fn load_from_memory(buffer: &[u8]) -> Option<Image> {
    if buffer.len() == 0 {
        return None;
    }

    if png::is_png(buffer) {
        match png::load_png_from_memory(buffer) {
            Ok(mut png_image) => {
                let (bytes, format) = match png_image.pixels {
                    png::PixelsByColorType::K8(ref mut data) => {
                        (data, PixelFormat::K8)
                    }
                    png::PixelsByColorType::KA8(ref mut data) => {
                        (data, PixelFormat::KA8)
                    }
                    png::PixelsByColorType::RGB8(ref mut data) => {
                        byte_swap(data);
                        (data, PixelFormat::RGB8)
                    }
                    png::PixelsByColorType::RGBA8(ref mut data) => {
                        byte_swap_and_premultiply(data);
                        (data, PixelFormat::RGBA8)
                    }
                };

                let bytes = mem::replace(bytes, Vec::new());
                let bytes = IpcSharedMemory::from_bytes(&bytes[..]);
                let image = Image {
                    width: png_image.width,
                    height: png_image.height,
                    format: format,
                    bytes: bytes,
                };

                Some(image)
            }
            Err(_err) => None,
        }
    } else {
        // For non-png images, we use stb_image
        // Can't remember why we do this. Maybe it's what cairo wants
        static FORCE_DEPTH: usize = 4;

        match stb_image2::load_from_memory_with_depth(buffer, FORCE_DEPTH, true) {
            stb_image2::LoadResult::ImageU8(mut image) => {
                assert!(image.depth == 4);
                // handle gif separately because the alpha-channel has to be premultiplied
                if is_gif(buffer) {
                    byte_swap_and_premultiply(&mut image.data);
                } else {
                    byte_swap(&mut image.data);
                }
                Some(Image {
                    width: image.width as u32,
                    height: image.height as u32,
                    format: PixelFormat::RGBA8,
                    bytes: IpcSharedMemory::from_bytes(&image.data[..]),
                })
            }
            stb_image2::LoadResult::ImageF32(_image) => {
                error!("HDR images not implemented");
                None
            }
            stb_image2::LoadResult::Error(e) => {
                error!("stb_image failed: {}", e);
                None
            }
        }
    }
}

fn is_gif(buffer: &[u8]) -> bool {
    match buffer {
        [b'G',b'I',b'F',b'8', n, b'a', ..] if n == b'7' || n == b'9' => true,
        _ => false
    }
}

<commit_msg>Be less vocal about image parsing failures.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use ipc_channel::ipc::IpcSharedMemory;
use png;
use stb_image::image as stb_image2;
use std::mem;
use util::vec::byte_swap;

// FIXME: Images must not be copied every frame. Instead we should atomically
// reference count them.

#[derive(Deserialize, Serialize)]
pub enum PixelFormat {
    K8,         // Luminance channel only
    KA8,        // Luminance + alpha
    RGB8,       // RGB, 8 bits per channel
    RGBA8,      // RGB + alpha, 8 bits per channel
}

#[derive(Deserialize, Serialize)]
pub struct Image {
    pub width: u32,
    pub height: u32,
    pub format: PixelFormat,
    pub bytes: IpcSharedMemory,
}

// TODO(pcwalton): Speed up with SIMD, or better yet, find some way to not do this.
fn byte_swap_and_premultiply(data: &mut [u8]) {
    let length = data.len();
    for i in (0..length).step_by(4) {
        let r = data[i + 2];
        let g = data[i + 1];
        let b = data[i + 0];
        let a = data[i + 3];
        data[i + 0] = ((r as u32) * (a as u32) / 255) as u8;
        data[i + 1] = ((g as u32) * (a as u32) / 255) as u8;
        data[i + 2] = ((b as u32) * (a as u32) / 255) as u8;
    }
}

pub fn load_from_memory(buffer: &[u8]) -> Option<Image> {
    if buffer.len() == 0 {
        return None;
    }

    if png::is_png(buffer) {
        match png::load_png_from_memory(buffer) {
            Ok(mut png_image) => {
                let (bytes, format) = match png_image.pixels {
                    png::PixelsByColorType::K8(ref mut data) => {
                        (data, PixelFormat::K8)
                    }
                    png::PixelsByColorType::KA8(ref mut data) => {
                        (data, PixelFormat::KA8)
                    }
                    png::PixelsByColorType::RGB8(ref mut data) => {
                        byte_swap(data);
                        (data, PixelFormat::RGB8)
                    }
                    png::PixelsByColorType::RGBA8(ref mut data) => {
                        byte_swap_and_premultiply(data);
                        (data, PixelFormat::RGBA8)
                    }
                };

                let bytes = mem::replace(bytes, Vec::new());
                let bytes = IpcSharedMemory::from_bytes(&bytes[..]);
                let image = Image {
                    width: png_image.width,
                    height: png_image.height,
                    format: format,
                    bytes: bytes,
                };

                Some(image)
            }
            Err(_err) => None,
        }
    } else {
        // For non-png images, we use stb_image
        // Can't remember why we do this. Maybe it's what cairo wants
        static FORCE_DEPTH: usize = 4;

        match stb_image2::load_from_memory_with_depth(buffer, FORCE_DEPTH, true) {
            stb_image2::LoadResult::ImageU8(mut image) => {
                assert!(image.depth == 4);
                // handle gif separately because the alpha-channel has to be premultiplied
                if is_gif(buffer) {
                    byte_swap_and_premultiply(&mut image.data);
                } else {
                    byte_swap(&mut image.data);
                }
                Some(Image {
                    width: image.width as u32,
                    height: image.height as u32,
                    format: PixelFormat::RGBA8,
                    bytes: IpcSharedMemory::from_bytes(&image.data[..]),
                })
            }
            stb_image2::LoadResult::ImageF32(_image) => {
                debug!("HDR images not implemented");
                None
            }
            stb_image2::LoadResult::Error(e) => {
                debug!("stb_image failed: {}", e);
                None
            }
        }
    }
}

fn is_gif(buffer: &[u8]) -> bool {
    match buffer {
        [b'G',b'I',b'F',b'8', n, b'a', ..] if n == b'7' || n == b'9' => true,
        _ => false
    }
}

<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    config::{Error, RoleType, SecureBackend},
    keys::ConfigKey,
    network_id::NetworkId,
    utils,
};
use diem_crypto::{x25519, Uniform};
use diem_network_address::NetworkAddress;
use diem_network_address_encryption::Encryptor;
use diem_secure_storage::{CryptoStorage, KVStorage, Storage};
use diem_types::{transaction::authenticator::AuthenticationKey, PeerId};
use rand::{
    rngs::{OsRng, StdRng},
    Rng, SeedableRng,
};
use serde::{Deserialize, Serialize};
use std::{
    collections::{HashMap, HashSet},
    convert::TryFrom,
    string::ToString,
};

// TODO: We could possibly move these constants somewhere else, but since they are defaults for the
//   configurations of the system, we'll leave it here for now.
/// Current supported protocol negotiation handshake version.
///
/// See [`perform_handshake`] in `network/src/transport.rs`
pub const HANDSHAKE_VERSION: u8 = 0;
pub const NETWORK_CHANNEL_SIZE: usize = 1024;
pub const PING_INTERVAL_MS: u64 = 1000;
pub const PING_TIMEOUT_MS: u64 = 10_000;
pub const PING_FAILURES_TOLERATED: u64 = 5;
pub const CONNECTIVITY_CHECK_INTERVAL_MS: u64 = 5000;
pub const MAX_CONCURRENT_NETWORK_REQS: usize = 100;
pub const MAX_CONCURRENT_NETWORK_NOTIFS: usize = 100;
pub const MAX_CONNECTION_DELAY_MS: u64 = 60_000; /* 1 minute */
pub const MAX_FULLNODE_OUTBOUND_CONNECTIONS: usize = 3;
pub const MAX_INBOUND_CONNECTIONS: usize = 100;
pub const MAX_FRAME_SIZE: usize = 8 * 1024 * 1024; /* 8 MiB */
pub const CONNECTION_BACKOFF_BASE: u64 = 2;
pub const INBOUND_IP_BYTE_BUCKET_RATE: usize = MAX_FRAME_SIZE;
pub const INBOUND_IP_BYTE_BUCKET_SIZE: usize = INBOUND_IP_BYTE_BUCKET_RATE;

pub type SeedPublicKeys = HashMap<PeerId, HashSet<x25519::PublicKey>>;
pub type SeedAddresses = HashMap<PeerId, Vec<NetworkAddress>>;

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct NetworkConfig {
    // Maximum backoff delay for connecting outbound to peers
    pub max_connection_delay_ms: u64,
    // Base for outbound connection backoff
    pub connection_backoff_base: u64,
    // Rate to check connectivity to connected peers
    pub connectivity_check_interval_ms: u64,
    // Size of all network channels
    pub network_channel_size: usize,
    // Maximum number of concurrent network requests
    pub max_concurrent_network_reqs: usize,
    // Maximum number of concurrent network notifications
    pub max_concurrent_network_notifs: usize,
    // Choose a protocol to discover and dial out to other peers on this network.
    // `DiscoveryMethod::None` disables discovery and dialing out (unless you have
    // seed peers configured).
    pub discovery_method: DiscoveryMethod,
    pub identity: Identity,
    // TODO: Add support for multiple listen/advertised addresses in config.
    // The address that this node is listening on for new connections.
    pub listen_address: NetworkAddress,
    // Select this to enforce that both peers should authenticate each other, otherwise
    // authentication only occurs for outgoing connections.
    pub mutual_authentication: bool,
    // Used to store network address encryption keys for validator nodes
    pub network_address_key_backend: Option<SecureBackend>,
    pub network_id: NetworkId,
    // Addresses of initial peers to connect to. In a mutual_authentication network,
    // we will extract the public keys from these addresses to set our initial
    // trusted peers set.
    pub seed_addrs: SeedAddresses,
    // Backup for public keys of peers that we'll accept connections from in a
    // mutual_authentication network. This config field is intended as a fallback
    // in case some peers don't have well defined addresses.
    pub seed_pubkeys: SeedPublicKeys,
    // The maximum size of an inbound or outbound request frame
    pub max_frame_size: usize,
    // Enables proxy protocol on incoming connections to get original source addresses
    pub enable_proxy_protocol: bool,
    // Interval to send healthcheck pings to peers
    pub ping_interval_ms: u64,
    // Timeout until a healthcheck ping is rejected
    pub ping_timeout_ms: u64,
    // Number of failed healthcheck pings until a peer is marked unhealthy
    pub ping_failures_tolerated: u64,
    // Maximum number of outbound connections, limited by ConnectivityManager
    pub max_outbound_connections: usize,
    // Maximum number of outbound connections, limited by PeerManager
    pub max_inbound_connections: usize,
    // Inbound rate limiting configuration, if not specified, no rate limiting
    pub inbound_rate_limit_config: Option<RateLimitConfig>,
    // Outbound rate limiting configuration, if not specified, no rate limiting
    pub outbound_rate_limit_config: Option<RateLimitConfig>,
}

impl Default for NetworkConfig {
    fn default() -> Self {
        NetworkConfig::network_with_id(NetworkId::default())
    }
}

impl NetworkConfig {
    pub fn network_with_id(network_id: NetworkId) -> NetworkConfig {
        let mut config = Self {
            discovery_method: DiscoveryMethod::None,
            identity: Identity::None,
            listen_address: "/ip4/0.0.0.0/tcp/6180".parse().unwrap(),
            mutual_authentication: false,
            network_address_key_backend: None,
            network_id,
            seed_pubkeys: HashMap::default(),
            seed_addrs: HashMap::default(),
            max_frame_size: MAX_FRAME_SIZE,
            enable_proxy_protocol: false,
            max_connection_delay_ms: MAX_CONNECTION_DELAY_MS,
            connectivity_check_interval_ms: CONNECTIVITY_CHECK_INTERVAL_MS,
            network_channel_size: NETWORK_CHANNEL_SIZE,
            max_concurrent_network_reqs: MAX_CONCURRENT_NETWORK_REQS,
            max_concurrent_network_notifs: MAX_CONCURRENT_NETWORK_NOTIFS,
            connection_backoff_base: CONNECTION_BACKOFF_BASE,
            ping_interval_ms: PING_INTERVAL_MS,
            ping_timeout_ms: PING_TIMEOUT_MS,
            ping_failures_tolerated: PING_FAILURES_TOLERATED,
            max_outbound_connections: MAX_FULLNODE_OUTBOUND_CONNECTIONS,
            max_inbound_connections: MAX_INBOUND_CONNECTIONS,
            inbound_rate_limit_config: Some(RateLimitConfig::default()),
            outbound_rate_limit_config: Some(RateLimitConfig::default()),
        };
        config.prepare_identity();
        config
    }
}

impl NetworkConfig {
    pub fn identity_key(&self) -> x25519::PrivateKey {
        let key = match &self.identity {
            Identity::FromConfig(config) => Some(config.key.clone().key),
            Identity::FromStorage(config) => {
                let storage: Storage = (&config.backend).into();
                let key = storage
                    .export_private_key(&config.key_name)
                    .expect("Unable to read key");
                let key = x25519::PrivateKey::from_ed25519_private_bytes(&key.to_bytes())
                    .expect("Unable to convert key");
                Some(key)
            }
            Identity::None => None,
        };
        key.expect("identity key should be present")
    }

    pub fn identity_from_storage(&self) -> IdentityFromStorage {
        if let Identity::FromStorage(identity) = self.identity.clone() {
            identity
        } else {
            panic!("Invalid identity found, expected a storage identity.");
        }
    }

    pub fn encryptor(&self) -> Encryptor {
        if let Some(backend) = self.network_address_key_backend.as_ref() {
            let storage = backend.into();
            Encryptor::new(storage)
        } else {
            Encryptor::for_testing()
        }
    }

    pub fn load(&mut self, role: RoleType) -> Result<(), Error> {
        if self.listen_address.to_string().is_empty() {
            self.listen_address = utils::get_local_ip()
                .ok_or_else(|| Error::InvariantViolation("No local IP".to_string()))?;
        }

        if role == RoleType::Validator {
            self.network_id = NetworkId::Validator;
        } else if self.network_id == NetworkId::Validator {
            return Err(Error::InvariantViolation(
                "Set NetworkId::Validator network for a non-validator network".to_string(),
            ));
        }

        self.prepare_identity();
        Ok(())
    }

    pub fn peer_id(&self) -> PeerId {
        match &self.identity {
            Identity::FromConfig(config) => Some(config.peer_id),
            Identity::FromStorage(config) => {
                let storage: Storage = (&config.backend).into();
                let peer_id = storage
                    .get::<PeerId>(&config.peer_id_name)
                    .expect("Unable to read peer id")
                    .value;
                Some(peer_id)
            }
            Identity::None => None,
        }
        .expect("peer id should be present")
    }

    fn prepare_identity(&mut self) {
        match &mut self.identity {
            Identity::FromStorage(_) => (),
            Identity::None => {
                let mut rng = StdRng::from_seed(OsRng.gen());
                let key = x25519::PrivateKey::generate(&mut rng);
                let peer_id = PeerId::from_identity_public_key(key.public_key());
                self.identity = Identity::from_config(key, peer_id);
            }
            Identity::FromConfig(config) => {
                let peer_id = PeerId::from_identity_public_key(config.key.public_key());
                if config.peer_id == PeerId::ZERO {
                    config.peer_id = peer_id;
                }
            }
        };
    }

    pub fn random(&mut self, rng: &mut StdRng) {
        self.random_with_peer_id(rng, None);
    }

    pub fn random_with_peer_id(&mut self, rng: &mut StdRng, peer_id: Option<PeerId>) {
        let identity_key = x25519::PrivateKey::generate(rng);
        let peer_id = if let Some(peer_id) = peer_id {
            peer_id
        } else {
            AuthenticationKey::try_from(identity_key.public_key().as_slice())
                .unwrap()
                .derived_address()
        };
        self.identity = Identity::from_config(identity_key, peer_id);
    }

    /// Check that all seed peer addresses look like canonical DiemNet addresses
    pub fn verify_seed_addrs(&self) -> Result<(), Error> {
        for (peer_id, addrs) in self.seed_addrs.iter() {
            for addr in addrs {
                crate::config::invariant(
                    addr.is_diemnet_addr(),
                    format!(
                        "Unexpected seed peer address format: peer_id: {}, addr: '{}'",
                        peer_id.short_str(),
                        addr,
                    ),
                )?;
            }
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum DiscoveryMethod {
    Onchain,
    None,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum Identity {
    FromConfig(IdentityFromConfig),
    FromStorage(IdentityFromStorage),
    None,
}

impl Identity {
    pub fn from_config(key: x25519::PrivateKey, peer_id: PeerId) -> Self {
        let key = ConfigKey::new(key);
        Identity::FromConfig(IdentityFromConfig { key, peer_id })
    }

    pub fn from_storage(key_name: String, peer_id_name: String, backend: SecureBackend) -> Self {
        Identity::FromStorage(IdentityFromStorage {
            backend,
            key_name,
            peer_id_name,
        })
    }
}

/// The identity is stored within the config.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct IdentityFromConfig {
    #[serde(flatten)]
    pub key: ConfigKey<x25519::PrivateKey>,
    pub peer_id: PeerId,
}

/// This represents an identity in a secure-storage as defined in NodeConfig::secure.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct IdentityFromStorage {
    pub backend: SecureBackend,
    pub key_name: String,
    pub peer_id_name: String,
}

#[derive(Copy, Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct RateLimitConfig {
    /// Maximum number of bytes/s for an IP
    pub ip_byte_bucket_rate: usize,
    /// Maximum burst of bytes for an IP
    pub ip_byte_bucket_size: usize,
    /// Initial amount of tokens initially in the bucket
    pub initial_bucket_fill_percentage: u8,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            ip_byte_bucket_rate: INBOUND_IP_BYTE_BUCKET_RATE,
            ip_byte_bucket_size: INBOUND_IP_BYTE_BUCKET_SIZE,
            initial_bucket_fill_percentage: 25,
        }
    }
}
<commit_msg>[config] Make rate limiter optional for Validators<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    config::{Error, RoleType, SecureBackend},
    keys::ConfigKey,
    network_id::NetworkId,
    utils,
};
use diem_crypto::{x25519, Uniform};
use diem_network_address::NetworkAddress;
use diem_network_address_encryption::Encryptor;
use diem_secure_storage::{CryptoStorage, KVStorage, Storage};
use diem_types::{transaction::authenticator::AuthenticationKey, PeerId};
use rand::{
    rngs::{OsRng, StdRng},
    Rng, SeedableRng,
};
use serde::{Deserialize, Serialize};
use std::{
    collections::{HashMap, HashSet},
    convert::TryFrom,
    string::ToString,
};

// TODO: We could possibly move these constants somewhere else, but since they are defaults for the
//   configurations of the system, we'll leave it here for now.
/// Current supported protocol negotiation handshake version.
///
/// See [`perform_handshake`] in `network/src/transport.rs`
pub const HANDSHAKE_VERSION: u8 = 0;
pub const NETWORK_CHANNEL_SIZE: usize = 1024;
pub const PING_INTERVAL_MS: u64 = 1000;
pub const PING_TIMEOUT_MS: u64 = 10_000;
pub const PING_FAILURES_TOLERATED: u64 = 5;
pub const CONNECTIVITY_CHECK_INTERVAL_MS: u64 = 5000;
pub const MAX_CONCURRENT_NETWORK_REQS: usize = 100;
pub const MAX_CONCURRENT_NETWORK_NOTIFS: usize = 100;
pub const MAX_CONNECTION_DELAY_MS: u64 = 60_000; /* 1 minute */
pub const MAX_FULLNODE_OUTBOUND_CONNECTIONS: usize = 3;
pub const MAX_INBOUND_CONNECTIONS: usize = 100;
pub const MAX_FRAME_SIZE: usize = 8 * 1024 * 1024; /* 8 MiB */
pub const CONNECTION_BACKOFF_BASE: u64 = 2;
pub const INBOUND_IP_BYTE_BUCKET_RATE: usize = MAX_FRAME_SIZE;
pub const INBOUND_IP_BYTE_BUCKET_SIZE: usize = INBOUND_IP_BYTE_BUCKET_RATE;

pub type SeedPublicKeys = HashMap<PeerId, HashSet<x25519::PublicKey>>;
pub type SeedAddresses = HashMap<PeerId, Vec<NetworkAddress>>;

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct NetworkConfig {
    // Maximum backoff delay for connecting outbound to peers
    pub max_connection_delay_ms: u64,
    // Base for outbound connection backoff
    pub connection_backoff_base: u64,
    // Rate to check connectivity to connected peers
    pub connectivity_check_interval_ms: u64,
    // Size of all network channels
    pub network_channel_size: usize,
    // Maximum number of concurrent network requests
    pub max_concurrent_network_reqs: usize,
    // Maximum number of concurrent network notifications
    pub max_concurrent_network_notifs: usize,
    // Choose a protocol to discover and dial out to other peers on this network.
    // `DiscoveryMethod::None` disables discovery and dialing out (unless you have
    // seed peers configured).
    pub discovery_method: DiscoveryMethod,
    pub identity: Identity,
    // TODO: Add support for multiple listen/advertised addresses in config.
    // The address that this node is listening on for new connections.
    pub listen_address: NetworkAddress,
    // Select this to enforce that both peers should authenticate each other, otherwise
    // authentication only occurs for outgoing connections.
    pub mutual_authentication: bool,
    // Used to store network address encryption keys for validator nodes
    pub network_address_key_backend: Option<SecureBackend>,
    pub network_id: NetworkId,
    // Addresses of initial peers to connect to. In a mutual_authentication network,
    // we will extract the public keys from these addresses to set our initial
    // trusted peers set.
    pub seed_addrs: SeedAddresses,
    // Backup for public keys of peers that we'll accept connections from in a
    // mutual_authentication network. This config field is intended as a fallback
    // in case some peers don't have well defined addresses.
    pub seed_pubkeys: SeedPublicKeys,
    // The maximum size of an inbound or outbound request frame
    pub max_frame_size: usize,
    // Enables proxy protocol on incoming connections to get original source addresses
    pub enable_proxy_protocol: bool,
    // Interval to send healthcheck pings to peers
    pub ping_interval_ms: u64,
    // Timeout until a healthcheck ping is rejected
    pub ping_timeout_ms: u64,
    // Number of failed healthcheck pings until a peer is marked unhealthy
    pub ping_failures_tolerated: u64,
    // Maximum number of outbound connections, limited by ConnectivityManager
    pub max_outbound_connections: usize,
    // Maximum number of outbound connections, limited by PeerManager
    pub max_inbound_connections: usize,
    // Inbound rate limiting configuration, if not specified, no rate limiting
    pub inbound_rate_limit_config: Option<RateLimitConfig>,
    // Outbound rate limiting configuration, if not specified, no rate limiting
    pub outbound_rate_limit_config: Option<RateLimitConfig>,
}

impl Default for NetworkConfig {
    fn default() -> Self {
        NetworkConfig::network_with_id(NetworkId::default())
    }
}

impl NetworkConfig {
    pub fn network_with_id(network_id: NetworkId) -> NetworkConfig {
        let mut config = Self {
            discovery_method: DiscoveryMethod::None,
            identity: Identity::None,
            listen_address: "/ip4/0.0.0.0/tcp/6180".parse().unwrap(),
            mutual_authentication: false,
            network_address_key_backend: None,
            network_id,
            seed_pubkeys: HashMap::default(),
            seed_addrs: HashMap::default(),
            max_frame_size: MAX_FRAME_SIZE,
            enable_proxy_protocol: false,
            max_connection_delay_ms: MAX_CONNECTION_DELAY_MS,
            connectivity_check_interval_ms: CONNECTIVITY_CHECK_INTERVAL_MS,
            network_channel_size: NETWORK_CHANNEL_SIZE,
            max_concurrent_network_reqs: MAX_CONCURRENT_NETWORK_REQS,
            max_concurrent_network_notifs: MAX_CONCURRENT_NETWORK_NOTIFS,
            connection_backoff_base: CONNECTION_BACKOFF_BASE,
            ping_interval_ms: PING_INTERVAL_MS,
            ping_timeout_ms: PING_TIMEOUT_MS,
            ping_failures_tolerated: PING_FAILURES_TOLERATED,
            max_outbound_connections: MAX_FULLNODE_OUTBOUND_CONNECTIONS,
            max_inbound_connections: MAX_INBOUND_CONNECTIONS,
            inbound_rate_limit_config: None,
            outbound_rate_limit_config: None,
        };
        config.prepare_identity();
        config
    }
}

impl NetworkConfig {
    pub fn identity_key(&self) -> x25519::PrivateKey {
        let key = match &self.identity {
            Identity::FromConfig(config) => Some(config.key.clone().key),
            Identity::FromStorage(config) => {
                let storage: Storage = (&config.backend).into();
                let key = storage
                    .export_private_key(&config.key_name)
                    .expect("Unable to read key");
                let key = x25519::PrivateKey::from_ed25519_private_bytes(&key.to_bytes())
                    .expect("Unable to convert key");
                Some(key)
            }
            Identity::None => None,
        };
        key.expect("identity key should be present")
    }

    pub fn identity_from_storage(&self) -> IdentityFromStorage {
        if let Identity::FromStorage(identity) = self.identity.clone() {
            identity
        } else {
            panic!("Invalid identity found, expected a storage identity.");
        }
    }

    pub fn encryptor(&self) -> Encryptor {
        if let Some(backend) = self.network_address_key_backend.as_ref() {
            let storage = backend.into();
            Encryptor::new(storage)
        } else {
            Encryptor::for_testing()
        }
    }

    pub fn load(&mut self, role: RoleType) -> Result<(), Error> {
        if self.listen_address.to_string().is_empty() {
            self.listen_address = utils::get_local_ip()
                .ok_or_else(|| Error::InvariantViolation("No local IP".to_string()))?;
        }

        if role == RoleType::Validator {
            self.network_id = NetworkId::Validator;
        } else if self.network_id == NetworkId::Validator {
            return Err(Error::InvariantViolation(
                "Set NetworkId::Validator network for a non-validator network".to_string(),
            ));
        }

        // Fullnodes require a rate limiting config
        if role == RoleType::FullNode {
            self.inbound_rate_limit_config = Some(
                self.inbound_rate_limit_config
                    .map_or(RateLimitConfig::default(), |config| config),
            );
            self.outbound_rate_limit_config = Some(
                self.outbound_rate_limit_config
                    .map_or(RateLimitConfig::default(), |config| config),
            );
        }

        self.prepare_identity();
        Ok(())
    }

    pub fn peer_id(&self) -> PeerId {
        match &self.identity {
            Identity::FromConfig(config) => Some(config.peer_id),
            Identity::FromStorage(config) => {
                let storage: Storage = (&config.backend).into();
                let peer_id = storage
                    .get::<PeerId>(&config.peer_id_name)
                    .expect("Unable to read peer id")
                    .value;
                Some(peer_id)
            }
            Identity::None => None,
        }
        .expect("peer id should be present")
    }

    fn prepare_identity(&mut self) {
        match &mut self.identity {
            Identity::FromStorage(_) => (),
            Identity::None => {
                let mut rng = StdRng::from_seed(OsRng.gen());
                let key = x25519::PrivateKey::generate(&mut rng);
                let peer_id = PeerId::from_identity_public_key(key.public_key());
                self.identity = Identity::from_config(key, peer_id);
            }
            Identity::FromConfig(config) => {
                let peer_id = PeerId::from_identity_public_key(config.key.public_key());
                if config.peer_id == PeerId::ZERO {
                    config.peer_id = peer_id;
                }
            }
        };
    }

    pub fn random(&mut self, rng: &mut StdRng) {
        self.random_with_peer_id(rng, None);
    }

    pub fn random_with_peer_id(&mut self, rng: &mut StdRng, peer_id: Option<PeerId>) {
        let identity_key = x25519::PrivateKey::generate(rng);
        let peer_id = if let Some(peer_id) = peer_id {
            peer_id
        } else {
            AuthenticationKey::try_from(identity_key.public_key().as_slice())
                .unwrap()
                .derived_address()
        };
        self.identity = Identity::from_config(identity_key, peer_id);
    }

    /// Check that all seed peer addresses look like canonical DiemNet addresses
    pub fn verify_seed_addrs(&self) -> Result<(), Error> {
        for (peer_id, addrs) in self.seed_addrs.iter() {
            for addr in addrs {
                crate::config::invariant(
                    addr.is_diemnet_addr(),
                    format!(
                        "Unexpected seed peer address format: peer_id: {}, addr: '{}'",
                        peer_id.short_str(),
                        addr,
                    ),
                )?;
            }
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum DiscoveryMethod {
    Onchain,
    None,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum Identity {
    FromConfig(IdentityFromConfig),
    FromStorage(IdentityFromStorage),
    None,
}

impl Identity {
    pub fn from_config(key: x25519::PrivateKey, peer_id: PeerId) -> Self {
        let key = ConfigKey::new(key);
        Identity::FromConfig(IdentityFromConfig { key, peer_id })
    }

    pub fn from_storage(key_name: String, peer_id_name: String, backend: SecureBackend) -> Self {
        Identity::FromStorage(IdentityFromStorage {
            backend,
            key_name,
            peer_id_name,
        })
    }
}

/// The identity is stored within the config.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct IdentityFromConfig {
    #[serde(flatten)]
    pub key: ConfigKey<x25519::PrivateKey>,
    pub peer_id: PeerId,
}

/// This represents an identity in a secure-storage as defined in NodeConfig::secure.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct IdentityFromStorage {
    pub backend: SecureBackend,
    pub key_name: String,
    pub peer_id_name: String,
}

#[derive(Copy, Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct RateLimitConfig {
    /// Maximum number of bytes/s for an IP
    pub ip_byte_bucket_rate: usize,
    /// Maximum burst of bytes for an IP
    pub ip_byte_bucket_size: usize,
    /// Initial amount of tokens initially in the bucket
    pub initial_bucket_fill_percentage: u8,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            ip_byte_bucket_rate: INBOUND_IP_BYTE_BUCKET_RATE,
            ip_byte_bucket_size: INBOUND_IP_BYTE_BUCKET_SIZE,
            initial_bucket_fill_percentage: 25,
        }
    }
}
<|endoftext|>