<commit_before>// Lumol, an extensible molecular simulation engine
// Copyright (C) Lumol's contributors — BSD license

//! Complex type
use std::ops::{Add, Sub, Neg, Mul, Div};
use std::f64;

use types::{Zero, One};

#[derive(Debug, Clone, Copy, PartialEq, Default)]
/// Complex number, with double precision real and imaginary parts.
///
/// `Complex` implements all the usual arithmetic operations:
///
/// ```
/// # use lumol::types::Complex;
///
/// let w = Complex::cartesian(-1.0, 0.5);
/// let z = Complex::cartesian(4.0, 2.0);
///
/// // Addition
/// let c = w + z;
/// assert_eq!(c, Complex::cartesian(3.0, 2.5));
///
/// // Subtraction
/// let c = w - z;
/// assert_eq!(c, Complex::cartesian(-5.0, -1.5));
///
/// // Multiplication
/// let c = w * z;
/// assert_eq!(c, Complex::cartesian(-5.0, 0.0));
///
/// let c = 42.0 * w;
/// assert_eq!(c, Complex::cartesian(-42.0, 21.0));
///
/// // Division
/// let c = z / 2.0;
/// assert_eq!(c, Complex::cartesian(2.0, 1.0));
/// ```
pub struct Complex {
    /// Real part of the complex
    real: f64,
    /// Imaginary part of the complex
    imag: f64,
}

impl Complex {
    /// Create a new `Complex` from a norm `r` and a phase `phi` in radians.
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// # use std::f64;
    /// let z = Complex::polar(3.0, f64::consts::PI);
    /// assert_eq!(z.norm(), 3.0);
    /// ```
    pub fn polar(r: f64, phi: f64) -> Complex {
        Complex{
            real: r * f64::cos(phi),
            imag: r * f64::sin(phi)
        }
    }

    /// Create a complex from Cartesian coordinates
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(3.0, -2.0);
    /// assert_eq!(z.real(), 3.0);
    /// assert_eq!(z.imag(), -2.0);
    /// ```
    pub fn cartesian(x: f64, y: f64) -> Complex {
        Complex{
            real: x,
            imag: y,
        }
    }

    /// Get the real part of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(3.0, -2.0);
    /// assert_eq!(z.real(), 3.0);
    /// ```
    #[inline]
    pub fn real(&self) -> f64 {
        self.real
    }

    /// Get the imaginary part of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(3.0, -2.0);
    /// assert_eq!(z.imag(), -2.0);
    /// ```
    #[inline]
    pub fn imag(&self) -> f64 {
        self.imag
    }

    /// Get the phase of the complex in the [-π, π) interval
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::polar(2.0, 0.3);
    /// assert_eq!(z.phase(), 0.3);
    /// ```
    #[inline]
    pub fn phase(&self) -> f64 {
        f64::atan2(self.imag, self.real)
    }

    /// Get the norm of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// # use std::f64;
    /// let z = Complex::polar(2.0, 0.3);
    /// assert_eq!(z.norm(), 2.0);
    ///
    /// let z = Complex::cartesian(2.0, 1.0);
    /// assert_eq!(z.norm(), f64::sqrt(5.0));
    /// ```
    #[inline]
    pub fn norm(&self) -> f64 {
        f64::sqrt(self.norm2())
    }

    /// Get the square of the norm if this complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(2.0, 1.0);
    /// assert_eq!(z.norm2(), 5.0);
    /// ```
    #[inline]
    pub fn norm2(&self) -> f64 {
        self.real * self.real + self.imag * self.imag
    }

    /// Get the conjugate of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(2.0, 1.0);
    /// assert_eq!(z.conj(), Complex::cartesian(2.0, -1.0));
    /// ```
    #[inline]
    pub fn conj(&self) -> Complex {
        Complex {
            real: self.real,
            imag: -self.imag
        }
    }
}

impl Add<Complex> for Complex {
    type Output = Complex;

    #[inline]
    fn add(self, other: Complex) -> Complex {
        let x = self.real() + other.real();
        let y = self.imag() + other.imag();
        return Complex::cartesian(x, y);
    }
}

impl Sub<Complex> for Complex {
    type Output = Complex;

    #[inline]
    fn sub(self, other: Complex) -> Complex {
        let x = self.real() - other.real();
        let y = self.imag() - other.imag();
        return Complex::cartesian(x, y);
    }
}

impl Neg for Complex {
    type Output = Complex;

    #[inline]
    fn neg(self) -> Complex {
        Complex{
            real: -self.real,
            imag: -self.imag,
        }
    }
}

impl Mul<Complex> for Complex {
    type Output = Complex;

    #[inline]
    fn mul(self, other: Complex) -> Complex {
        let x = self.real() * other.real() - self.imag() * other.imag();
        let y = self.real() * other.imag() + self.imag() * other.real();
        Complex::cartesian(x, y)
    }
}

impl Mul<f64> for Complex {
    type Output = Complex;

    #[inline]
    fn mul(self, other: f64) -> Complex {
        Complex::cartesian(self.real() * other, self.imag() * other)
    }
}

impl Mul<Complex> for f64 {
    type Output = Complex;

    #[inline]
    fn mul(self, other: Complex) -> Complex {
        Complex::cartesian(self * other.real(), self * other.imag())
    }
}

impl Div<Complex> for Complex {
    type Output = Complex;

    #[inline]
    fn div(self, other: Complex) -> Complex {
        let r = other.norm2();
        let x = self.real() * other.real() + self.imag() * other.imag();
        let y = - self.real() * other.imag() + self.imag() * other.real();

        Complex::cartesian(x / r, y / r)
    }
}

impl Div<f64> for Complex {
    type Output = Complex;

    #[inline]
    fn div(self, other: f64) -> Complex {
        let norm = self.norm() / other;
        let phase = self.phase();
        Complex::polar(norm, phase)
    }
}

impl Zero for Complex {
    fn zero() -> Complex {
        Complex::cartesian(0.0, 0.0)
    }

    fn is_zero(&self) -> bool {
        self.norm2() == 0.0
    }
}

impl One for Complex {
    fn one() -> Complex {
        Complex::cartesian(1.0, 0.0)
    }
}

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

    use approx::ApproxEq;
    impl ApproxEq for Complex {
        type Epsilon = <f64 as ApproxEq>::Epsilon;

        fn default_epsilon() -> Self::Epsilon {
            f64::default_epsilon()
        }

        fn default_max_relative() -> Self::Epsilon {
            f64::default_max_relative()
        }

        fn default_max_ulps() -> u32 {
            f64::default_max_ulps()
        }

        fn relative_eq(&self, other: &Self, epsilon: Self::Epsilon, max_relative: Self::Epsilon) -> bool {
            f64::relative_eq(&self.real, &other.real, epsilon, max_relative) &&
            f64::relative_eq(&self.imag, &other.imag, epsilon, max_relative)
        }

        fn ulps_eq(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool {
            f64::ulps_eq(&self.real, &other.real, epsilon, max_ulps) &&
            f64::ulps_eq(&self.imag, &other.imag, epsilon, max_ulps)
        }
    }

    #[test]
    fn norm() {
        let c = Complex::polar(3.0, 5.0);
        assert_ulps_eq!(c.norm(), 3.0);

        let c = Complex::polar(-3.0, 0.0);
        assert_ulps_eq!(c.norm(), 3.0);
    }

    #[test]
    fn phase() {
        // Phase is between 0 and 2π
        for &phase in &[-consts::PI, -3.1, -1.5, 0.0, 0.1, 2.0, 3.1] {
            let c = Complex::polar(1.0, phase);
            assert_ulps_eq!(c.phase(), phase);
        }

        let c = Complex::polar(1.0, -8.0);
        assert_ulps_eq!(c.phase(), -8.0 + 2.0 * consts::PI);

        let c = Complex::polar(1.0, 12.0);
        assert_ulps_eq!(c.phase(), 12.0 - 4.0 * consts::PI, max_ulps=10);

        let c = Complex::polar(1.0, consts::PI);
        assert_eq!(c.phase(), consts::PI);
        let c = Complex::polar(1.0, -consts::PI);
        assert_eq!(c.phase(), -consts::PI);

        let c = Complex::polar(3.0, 0.0);
        assert_eq!(c.phase(), 0.0);

        let c = Complex::polar(-3.0, 0.0);
        assert_eq!(c.phase(), -consts::PI);
    }

    #[test]
    fn conj() {
        let c = Complex::cartesian(3.0, 5.0);
        assert_eq!(c.conj(), Complex::cartesian(3.0, -5.0));

        let c = Complex::cartesian(3.0, -5.0);
        assert_eq!(c.conj(), Complex::cartesian(3.0, 5.0));

        let c = Complex::cartesian(3.0, 0.0);
        assert_eq!(c.conj(), Complex::cartesian(3.0, 0.0));
    }

    #[test]
    fn cartesian() {
        let c = Complex::polar(1.0, 0.0);
        assert_ulps_eq!(c.real(), 1.0);
        assert_ulps_eq!(c.imag(), 0.0);

        let c = Complex::polar(1.0, consts::PI);
        assert_ulps_eq!(c.real(), -1.0);
        assert_ulps_eq!(c.imag(), 0.0);

        let c = Complex::polar(1.0, consts::FRAC_PI_2);
        assert_ulps_eq!(c.real(), 0.0);
        assert_ulps_eq!(c.imag(), 1.0);

        let c = Complex::polar(1.0, consts::FRAC_PI_4);
        assert_ulps_eq!(c.real(), consts::FRAC_1_SQRT_2);
        assert_ulps_eq!(c.imag(), consts::FRAC_1_SQRT_2);

        let c = Complex::cartesian(consts::FRAC_1_SQRT_2, consts::FRAC_1_SQRT_2);
        assert_ulps_eq!(c.norm(), 1.0);
        assert_ulps_eq!(c.phase(), consts::FRAC_PI_4);
    }

    #[test]
    fn add() {
        let a = Complex::polar(2.0, 0.2);
        let b = Complex::polar(1.0, 0.5);
        let c = a + b;

        assert_eq!(c.real(), a.real() + b.real());
        assert_eq!(c.imag(), a.imag() + b.imag());
    }

    #[test]
    fn sub() {
        let a = Complex::polar(2.0, 0.2);
        let b = Complex::polar(1.0, 0.5);
        let c = a - b;

        assert_eq!(c.real(), a.real() - b.real());
        assert_eq!(c.imag(), a.imag() - b.imag());
    }

    #[test]
    fn neg() {
        let a = Complex::polar(2.0, 0.2);
        let c = -a;

        assert_ulps_eq!(c.norm(), a.norm());
        assert_ulps_eq!(c.phase(), a.phase() - consts::PI);

        assert_ulps_eq!(c.real(), - a.real());
        assert_ulps_eq!(c.imag(), - a.imag());
    }

    #[test]
    fn mul() {
        let a = Complex::polar(2.0, 1.4);
        let b = Complex::polar(1.0, 0.5);
        let c = a * b;

        assert_eq!(c.norm(), a.norm() * b.norm());
        assert_eq!(c.phase(), a.phase() + b.phase());

        let c = 3.0 * a;
        assert_eq!(c.norm(), 3.0 * a.norm());
        assert_eq!(c.phase(), a.phase());

        let c = a * 3.0;
        assert_eq!(c.norm(), 3.0 * a.norm());
        assert_eq!(c.phase(), a.phase());

        let c = -2.0 * a;
        assert_eq!(c.norm(), 2.0 * a.norm());
        assert_ulps_eq!(c.phase(), a.phase() - consts::PI);
    }

    #[test]
    fn div() {
        let a = Complex::polar(2.0, 0.2);
        let b = Complex::polar(1.0, 0.5);
        let c = a / b;

        assert_eq!(c.norm(), a.norm() / b.norm());
        assert_eq!(c.phase(), a.phase() - b.phase());

        let c = a / 3.0;
        assert_eq!(c.norm(), a.norm() / 3.0);
        assert_ulps_eq!(c.phase(), a.phase());

        let c = a / (-2.0);
        assert_eq!(c.norm(), a.norm() / 2.0);
        assert_ulps_eq!(c.phase(), a.phase() - consts::PI);
    }
}
<commit_msg>Implement XxxAssign operators for Complex<commit_after>// Lumol, an extensible molecular simulation engine
// Copyright (C) Lumol's contributors — BSD license

//! Complex type
use std::ops::{Add, AddAssign, Sub, SubAssign, Neg, Mul, MulAssign, Div, DivAssign};
use std::f64;

use types::{Zero, One};

#[derive(Debug, Clone, Copy, PartialEq, Default)]
/// Complex number, with double precision real and imaginary parts.
///
/// `Complex` implements all the usual arithmetic operations:
///
/// ```
/// # use lumol::types::Complex;
///
/// let w = Complex::cartesian(-1.0, 0.5);
/// let z = Complex::cartesian(4.0, 2.0);
///
/// // Addition
/// let c = w + z;
/// assert_eq!(c, Complex::cartesian(3.0, 2.5));
///
/// // Subtraction
/// let c = w - z;
/// assert_eq!(c, Complex::cartesian(-5.0, -1.5));
///
/// // Multiplication
/// let c = w * z;
/// assert_eq!(c, Complex::cartesian(-5.0, 0.0));
///
/// let c = 42.0 * w;
/// assert_eq!(c, Complex::cartesian(-42.0, 21.0));
///
/// // Division
/// let c = z / 2.0;
/// assert_eq!(c, Complex::cartesian(2.0, 1.0));
/// ```
pub struct Complex {
    /// Real part of the complex
    real: f64,
    /// Imaginary part of the complex
    imag: f64,
}

impl Complex {
    /// Create a new `Complex` from a norm `r` and a phase `phi` in radians.
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// # use std::f64;
    /// let z = Complex::polar(3.0, f64::consts::PI);
    /// assert_eq!(z.norm(), 3.0);
    /// ```
    pub fn polar(r: f64, phi: f64) -> Complex {
        Complex{
            real: r * f64::cos(phi),
            imag: r * f64::sin(phi)
        }
    }

    /// Create a complex from Cartesian coordinates
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(3.0, -2.0);
    /// assert_eq!(z.real(), 3.0);
    /// assert_eq!(z.imag(), -2.0);
    /// ```
    pub fn cartesian(x: f64, y: f64) -> Complex {
        Complex{
            real: x,
            imag: y,
        }
    }

    /// Get the real part of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(3.0, -2.0);
    /// assert_eq!(z.real(), 3.0);
    /// ```
    #[inline]
    pub fn real(&self) -> f64 {
        self.real
    }

    /// Get the imaginary part of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(3.0, -2.0);
    /// assert_eq!(z.imag(), -2.0);
    /// ```
    #[inline]
    pub fn imag(&self) -> f64 {
        self.imag
    }

    /// Get the phase of the complex in the [-π, π) interval
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::polar(2.0, 0.3);
    /// assert_eq!(z.phase(), 0.3);
    /// ```
    #[inline]
    pub fn phase(&self) -> f64 {
        f64::atan2(self.imag, self.real)
    }

    /// Get the norm of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// # use std::f64;
    /// let z = Complex::polar(2.0, 0.3);
    /// assert_eq!(z.norm(), 2.0);
    ///
    /// let z = Complex::cartesian(2.0, 1.0);
    /// assert_eq!(z.norm(), f64::sqrt(5.0));
    /// ```
    #[inline]
    pub fn norm(&self) -> f64 {
        f64::sqrt(self.norm2())
    }

    /// Get the square of the norm if this complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(2.0, 1.0);
    /// assert_eq!(z.norm2(), 5.0);
    /// ```
    #[inline]
    pub fn norm2(&self) -> f64 {
        self.real * self.real + self.imag * self.imag
    }

    /// Get the conjugate of the complex
    /// # Examples
    /// ```
    /// # use lumol::types::Complex;
    /// let z = Complex::cartesian(2.0, 1.0);
    /// assert_eq!(z.conj(), Complex::cartesian(2.0, -1.0));
    /// ```
    #[inline]
    pub fn conj(&self) -> Complex {
        Complex {
            real: self.real,
            imag: -self.imag
        }
    }
}

impl Add<Complex> for Complex {
    type Output = Complex;

    fn add(self, other: Complex) -> Complex {
        Complex {
            real: self.real + other.real,
            imag: self.imag + other.imag,
        }
    }
}

impl AddAssign<Complex> for Complex {
    fn add_assign(&mut self, other: Complex) {
        self.real += other.real;
        self.imag += other.imag;
    }
}

impl Sub<Complex> for Complex {
    type Output = Complex;

    fn sub(self, other: Complex) -> Complex {
        Complex {
            real: self.real - other.real,
            imag: self.imag - other.imag,
        }
    }
}

impl SubAssign<Complex> for Complex {
    fn sub_assign(&mut self, other: Complex) {
        self.real -= other.real;
        self.imag -= other.imag;
    }
}

impl Neg for Complex {
    type Output = Complex;

    fn neg(self) -> Complex {
        Complex{
            real: -self.real,
            imag: -self.imag,
        }
    }
}

impl Mul<Complex> for Complex {
    type Output = Complex;

    fn mul(self, other: Complex) -> Complex {
        Complex {
            real: self.real * other.real - self.imag * other.imag,
            imag: self.real * other.imag + self.imag * other.real,
        }
    }
}

impl Mul<f64> for Complex {
    type Output = Complex;

    fn mul(self, other: f64) -> Complex {
        Complex {
            real: self.real * other,
            imag: self.imag * other,
        }
    }
}

impl Mul<Complex> for f64 {
    type Output = Complex;

    fn mul(self, other: Complex) -> Complex {
        Complex {
            real: self * other.real,
            imag: self * other.imag,
        }
    }
}

impl MulAssign<Complex> for Complex {
    fn mul_assign(&mut self, other: Complex) {
        let real = self.real * other.real - self.imag * other.imag;
        let imag = self.real * other.imag + self.imag * other.real;

        self.real = real;
        self.imag = imag;
    }
}

impl MulAssign<f64> for Complex {
    fn mul_assign(&mut self, other: f64) {
        self.real *= other;
        self.imag *= other;
    }
}

impl Div<Complex> for Complex {
    type Output = Complex;

    fn div(self, other: Complex) -> Complex {
        let norm = other.norm2();
        let real = self.real * other.real + self.imag * other.imag;
        let imag = - self.real * other.imag + self.imag * other.real;

        Complex {
            real: real / norm,
            imag: imag / norm,
        }
    }
}

impl Div<f64> for Complex {
    type Output = Complex;

    fn div(self, other: f64) -> Complex {
        Complex {
            real: self.real / other,
            imag: self.imag / other,
        }
    }
}

impl DivAssign<Complex> for Complex {
    fn div_assign(&mut self, other: Complex) {
        let norm = other.norm2();
        let real = self.real * other.real + self.imag * other.imag;
        let imag = - self.real * other.imag + self.imag * other.real;

        self.real = real / norm;
        self.imag = imag / norm;
    }
}

impl DivAssign<f64> for Complex {
    fn div_assign(&mut self, other: f64) {
        self.real /= other;
        self.imag /= other;
    }
}

impl Zero for Complex {
    fn zero() -> Complex {
        Complex::cartesian(0.0, 0.0)
    }

    fn is_zero(&self) -> bool {
        self.norm2() == 0.0
    }
}

impl One for Complex {
    fn one() -> Complex {
        Complex::cartesian(1.0, 0.0)
    }
}

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

    use approx::ApproxEq;
    impl ApproxEq for Complex {
        type Epsilon = <f64 as ApproxEq>::Epsilon;

        fn default_epsilon() -> Self::Epsilon {
            f64::default_epsilon()
        }

        fn default_max_relative() -> Self::Epsilon {
            f64::default_max_relative()
        }

        fn default_max_ulps() -> u32 {
            f64::default_max_ulps()
        }

        fn relative_eq(&self, other: &Self, epsilon: Self::Epsilon, max_relative: Self::Epsilon) -> bool {
            f64::relative_eq(&self.real, &other.real, epsilon, max_relative) &&
            f64::relative_eq(&self.imag, &other.imag, epsilon, max_relative)
        }

        fn ulps_eq(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool {
            f64::ulps_eq(&self.real, &other.real, epsilon, max_ulps) &&
            f64::ulps_eq(&self.imag, &other.imag, epsilon, max_ulps)
        }
    }

    #[test]
    fn norm() {
        let c = Complex::polar(3.0, 5.0);
        assert_ulps_eq!(c.norm(), 3.0);

        let c = Complex::polar(-3.0, 0.0);
        assert_ulps_eq!(c.norm(), 3.0);
    }

    #[test]
    fn phase() {
        // Phase is between 0 and 2π
        for &phase in &[-consts::PI, -3.1, -1.5, 0.0, 0.1, 2.0, 3.1] {
            let c = Complex::polar(1.0, phase);
            assert_ulps_eq!(c.phase(), phase);
        }

        let c = Complex::polar(1.0, -8.0);
        assert_ulps_eq!(c.phase(), -8.0 + 2.0 * consts::PI);

        let c = Complex::polar(1.0, 12.0);
        assert_ulps_eq!(c.phase(), 12.0 - 4.0 * consts::PI, max_ulps=10);

        let c = Complex::polar(1.0, consts::PI);
        assert_eq!(c.phase(), consts::PI);
        let c = Complex::polar(1.0, -consts::PI);
        assert_eq!(c.phase(), -consts::PI);

        let c = Complex::polar(3.0, 0.0);
        assert_eq!(c.phase(), 0.0);

        let c = Complex::polar(-3.0, 0.0);
        assert_eq!(c.phase(), -consts::PI);
    }

    #[test]
    fn conj() {
        let c = Complex::cartesian(3.0, 5.0);
        assert_eq!(c.conj(), Complex::cartesian(3.0, -5.0));

        let c = Complex::cartesian(3.0, -5.0);
        assert_eq!(c.conj(), Complex::cartesian(3.0, 5.0));

        let c = Complex::cartesian(3.0, 0.0);
        assert_eq!(c.conj(), Complex::cartesian(3.0, 0.0));
    }

    #[test]
    fn cartesian() {
        let c = Complex::polar(1.0, 0.0);
        assert_ulps_eq!(c.real(), 1.0);
        assert_ulps_eq!(c.imag(), 0.0);

        let c = Complex::polar(1.0, consts::PI);
        assert_ulps_eq!(c.real(), -1.0);
        assert_ulps_eq!(c.imag(), 0.0);

        let c = Complex::polar(1.0, consts::FRAC_PI_2);
        assert_ulps_eq!(c.real(), 0.0);
        assert_ulps_eq!(c.imag(), 1.0);

        let c = Complex::polar(1.0, consts::FRAC_PI_4);
        assert_ulps_eq!(c.real(), consts::FRAC_1_SQRT_2);
        assert_ulps_eq!(c.imag(), consts::FRAC_1_SQRT_2);

        let c = Complex::cartesian(consts::FRAC_1_SQRT_2, consts::FRAC_1_SQRT_2);
        assert_ulps_eq!(c.norm(), 1.0);
        assert_ulps_eq!(c.phase(), consts::FRAC_PI_4);
    }

    #[test]
    fn add() {
        let mut a = Complex::polar(2.0, 0.2);
        let b = Complex::polar(1.0, 0.5);
        let c = a + b;

        assert_eq!(c.real(), a.real() + b.real());
        assert_eq!(c.imag(), a.imag() + b.imag());

        a += b;
        assert_eq!(a, c);
    }

    #[test]
    fn sub() {
        let mut a = Complex::polar(2.0, 0.2);
        let b = Complex::polar(1.0, 0.5);
        let c = a - b;

        assert_eq!(c.real(), a.real() - b.real());
        assert_eq!(c.imag(), a.imag() - b.imag());

        a -= b;
        assert_eq!(a, c);
    }

    #[test]
    fn neg() {
        let a = Complex::polar(2.0, 0.2);
        let c = -a;

        assert_ulps_eq!(c.norm(), a.norm());
        assert_ulps_eq!(c.phase(), a.phase() - consts::PI);

        assert_ulps_eq!(c.real(), - a.real());
        assert_ulps_eq!(c.imag(), - a.imag());
    }

    #[test]
    fn mul() {
        let mut a = Complex::polar(2.0, 1.4);
        let b = Complex::polar(1.0, 0.5);
        let c = a * b;

        assert_eq!(c.norm(), a.norm() * b.norm());
        assert_eq!(c.phase(), a.phase() + b.phase());

        let c = 3.0 * a;
        assert_eq!(c.norm(), 3.0 * a.norm());
        assert_eq!(c.phase(), a.phase());

        let c = a * 3.0;
        assert_eq!(c.norm(), 3.0 * a.norm());
        assert_eq!(c.phase(), a.phase());

        let c = -2.0 * a;
        assert_eq!(c.norm(), 2.0 * a.norm());
        assert_ulps_eq!(c.phase(), a.phase() - consts::PI);

        let c = a * b;
        a *= b;
        assert_eq!(a, c);

        let c = 3.0 * a;
        a *= 3.0;
        assert_eq!(a, c);
    }

    #[test]
    fn div() {
        let mut a = Complex::polar(2.0, 0.2);
        let b = Complex::polar(1.0, 0.5);
        let c = a / b;

        assert_eq!(c.norm(), a.norm() / b.norm());
        assert_eq!(c.phase(), a.phase() - b.phase());

        let c = a / 3.0;
        assert_eq!(c.norm(), a.norm() / 3.0);
        assert_ulps_eq!(c.phase(), a.phase());

        let c = a / (-2.0);
        assert_eq!(c.norm(), a.norm() / 2.0);
        assert_ulps_eq!(c.phase(), a.phase() - consts::PI);

        let c = a / b;
        a /= b;
        assert_eq!(a, c);

        let c = a / 3.0;
        a /= 3.0;
        assert_eq!(a, c);
    }
}
<|endoftext|>


<commit_before>use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use function::gamma;
use statistics::*;
use distribution::{Continuous, Distribution};
use result::Result;
use error::StatsError;
use prec;

/// Implements the [Dirichlet](https://en.wikipedia.org/wiki/Dirichlet_distribution)
/// distribution
///
/// # Examples
///
/// ```
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Dirichlet {
    alpha: Vec<f64>,
}

impl Dirichlet {
    /// Constructs a new dirichlet distribution with the given
    /// concenctration parameters (alpha)
    ///
    /// # Errors
    ///
    /// Returns an error if any element `x` in alpha exist
    /// such that `x < = 0.0` or if the length of alpha is
    /// less than 2
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Dirichlet;
    ///
    /// let alpha_ok = [1.0, 2.0, 3.0];
    /// let mut result = Dirichlet::new(&alpha_ok);
    /// assert!(result.is_ok());
    ///
    /// let alpha_err = [0.0];
    /// result = Dirichlet::new(&alpha_err);
    /// assert!(result.is_err());
    /// ```
    pub fn new(alpha: &[f64]) -> Result<Dirichlet> {
        if alpha.len() < 2 {
            return Err(StatsError::BadParams);
        }
        for x in alpha {
            if *x <= 0.0 {
                return Err(StatsError::BadParams);
            }
        }
        Ok(Dirichlet { alpha: alpha.to_vec() })
    }

    /// Returns the concentration parameters of
    /// the dirichlet distribution as a slice
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Dirichlet;
    ///
    /// let n = Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap();
    /// assert_eq!(n.alpha(), [1.0, 2.0, 3.0]);
    /// ```
    pub fn alpha(&self) -> &[f64] {
        &self.alpha
    }

    fn alpha_sum(&self) -> f64 {
        self.alpha.iter().fold(0.0, |acc, x| acc + x)
    }
}

impl Sample<Vec<f64>> for Dirichlet {
    /// Generate random samples from a dirichlet
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn sample<R: Rng>(&mut self, r: &mut R) -> Vec<f64> {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<Vec<f64>> for Dirichlet {
    /// Generate random independent samples from a dirichlet
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn ind_sample<R: Rng>(&self, r: &mut R) -> Vec<f64> {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<Vec<f64>> for Dirichlet {
    /// Generate random samples from the dirichlet distribution
    /// using `r` as the source of randomness
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Dirichlet, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> Vec<f64> {
        let n = self.alpha.len();
        let mut samples = vec![0.0; n];
        let mut sum = 0.0;
        for i in 0..n {
            samples[i] = super::gamma::sample_unchecked(r, self.alpha[i], 1.0);
            sum += samples[i];
        }
        for i in 0..n {
            samples[i] /= sum
        }
        samples
    }
}

impl Mean<Vec<f64>> for Dirichlet {
    /// Returns the means of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// α_i / α_0
    /// ```
    ///
    /// for the `i`th element where `α_i` is the `i`th concentration parameter
    /// and `α_0` is the sum of all concentration parameters
    fn mean(&self) -> Vec<f64> {
        let sum = self.alpha_sum();
        self.alpha.iter().map(|x| x / sum).collect()
    }
}

impl Variance<Vec<f64>> for Dirichlet {
    /// Returns the variances of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// (α_i * (α_0 - α_i)) / (α_0^2 * (α_0 + 1))
    /// ```
    ///
    /// for the `i`th element where `α_i` is the `i`th concentration parameter
    /// and `α_0` is the sum of all concentration parameters
    fn variance(&self) -> Vec<f64> {
        let sum = self.alpha_sum();
        self.alpha.iter().map(|x| x * (sum - x) / (sum * sum * (sum + 1.0))).collect()
    }

    /// Returns the variances of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sqrt((α_i * (α_0 - α_i)) / (α_0^2 * (α_0 + 1)))
    /// ```
    ///
    /// for the `i`th element where `α_i` is the `i`th concentration parameter
    /// and `α_0` is the sum of all concentration parameters
    fn std_dev(&self) -> Vec<f64> {
        self.variance().iter().map(|x| x.sqrt()).collect()
    }
}

impl Entropy<f64> for Dirichlet {
    /// Returns the entropy of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// ln(B(α)) - (K - α_0)ψ(α_0) - ∑((α_i - 1)ψ(α_i))
    /// ```
    ///
    /// where
    ///
    /// ```ignore
    /// B(α) = ∏(Γ(α_i)) / Γ(∑(α_i))
    /// ```
    ///
    /// `α_0` is the sum of all concentration parameters,
    /// `K` is the number of concentration parameters, `ψ` is the digamma function, `α_i`
    /// is the `i`th concentration parameter, and `∑` is the sum from `1` to `K`
    fn entropy(&self) -> f64 {
        let sum = self.alpha_sum();
        let num = self.alpha.iter().fold(0.0, |acc, &x| acc + (x - 1.0) * gamma::digamma(x));
        gamma::ln_gamma(sum) + (sum - self.alpha.len() as f64) * gamma::digamma(sum) - num
    }
}

impl<'a> Continuous<&'a [f64], f64> for Dirichlet {
    /// Calculates the probabiliy density function for the dirichlet distribution
    /// with given `x`'s corresponding to the concentration parameters for this
    /// distribution
    ///
    /// # Panics
    ///
    /// If any element in `x` is not in `(0, 1)` or if `x` is not the same length
    /// as the vector of concentration parameters for this distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// (1 / B(α)) * ∏(x_i^(α_i - 1))
    /// ```
    ///
    /// where
    ///
    /// ```ignore
    /// B(α) = ∏(Γ(α_i)) / Γ(∑(α_i))
    /// ```
    ///
    /// `α` is the vector of concentration parameters, `α_i` is the `i`th
    /// concentration parameter, `x_i` is the `i`th argument corresponding to
    /// the `i`th concentration parameter, `Γ` is the gamma function,
    /// `∏` is the product from `1` to `K`, `∑` is the sum from `1` to `K`,
    /// and `K` is the number of concentration parameters
    fn pdf(&self, x: &[f64]) -> f64 {
        self.ln_pdf(x).exp()
    }

    /// Calculates the log probabiliy density function for the dirichlet distribution
    /// with given `x`'s corresponding to the concentration parameters for this
    /// distribution
    ///
    /// # Panics
    ///
    /// If any element in `x` is not in `(0, 1)` or if `x` is not the same length
    /// as the vector of concentration parameters for this distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// ln((1 / B(α)) * ∏(x_i^(α_i - 1)))
    /// ```
    ///
    /// where
    ///
    /// ```ignore
    /// B(α) = ∏(Γ(α_i)) / Γ(∑(α_i))
    /// ```
    ///
    /// `α` is the vector of concentration parameters, `α_i` is the `i`th
    /// concentration parameter, `x_i` is the `i`th argument corresponding to
    /// the `i`th concentration parameter, `Γ` is the gamma function,
    /// `∏` is the product from `1` to `K`, `∑` is the sum from `1` to `K`,
    /// and `K` is the number of concentration parameters
    fn ln_pdf(&self, x: &[f64]) -> f64 {
        assert!(self.alpha.len() == x.len(),
                format!("{}", StatsError::ContainersMustBeSameLength));

        let (term, sum_xi, sum_alpha) = x.iter()
            .enumerate()
            .map(|pair| (pair.1, self.alpha[pair.0]))
            .fold((0.0, 0.0, 0.0), |acc, pair| {
                assert!(*pair.0 > 0.0 && *pair.0 < 1.0,
                        format!("{}", StatsError::ArgIntervalExcl("x", 0.0, 1.0)));

                (acc.0 + (pair.1 - 1.0) * pair.0.ln() - gamma::ln_gamma(pair.1),
                 acc.1 + pair.0,
                 acc.2 + pair.1)
            });

        if !prec::almost_eq(sum_xi, 1.0, 1e-8) {
            0.0
        } else {
            term + gamma::ln_gamma(sum_alpha)
        }
    }
}<commit_msg>wip: begin testing dirichlet<commit_after>use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use function::gamma;
use statistics::*;
use distribution::{Continuous, Distribution};
use result::Result;
use error::StatsError;
use prec;

/// Implements the [Dirichlet](https://en.wikipedia.org/wiki/Dirichlet_distribution)
/// distribution
///
/// # Examples
///
/// ```
/// use statrs::distribution::{Dirichlet, Continuous};
/// use statrs::statistics::Mean;
///
/// let n = Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap();
/// assert_eq!(n.mean(), [1.0 / 6.0, 1.0 / 3.0, 0.5]);
/// assert_eq!(n.pdf(&[0.5, 0.5, 0.5]), 1.0);
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Dirichlet {
    alpha: Vec<f64>,
}

impl Dirichlet {
    /// Constructs a new dirichlet distribution with the given
    /// concenctration parameters (alpha)
    ///
    /// # Errors
    ///
    /// Returns an error if any element `x` in alpha exist
    /// such that `x < = 0.0` or `x` is `NaN`, or if the length of alpha is
    /// less than 2
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Dirichlet;
    ///
    /// let alpha_ok = [1.0, 2.0, 3.0];
    /// let mut result = Dirichlet::new(&alpha_ok);
    /// assert!(result.is_ok());
    ///
    /// let alpha_err = [0.0];
    /// result = Dirichlet::new(&alpha_err);
    /// assert!(result.is_err());
    /// ```
    pub fn new(alpha: &[f64]) -> Result<Dirichlet> {
        if alpha.len() < 2 {
            return Err(StatsError::BadParams);
        }
        for x in alpha {
            if *x <= 0.0 || x.is_nan() {
                return Err(StatsError::BadParams);
            }
        }
        Ok(Dirichlet { alpha: alpha.to_vec() })
    }

    /// Constructs a new dirichlet distribution with the given
    /// concenctration parameter (alpha) repeated `n` times
    ///
    /// # Errors
    ///
    /// Returns an error if `alpha < = 0.0` or `alpha` is `NaN`,
    /// or if `n < 2`
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Dirichlet;
    ///
    /// let mut result = Dirichlet::new_with_param(1.0, 3);
    /// assert!(result.is_ok());
    ///
    /// result = Dirichlet::new_with_param(0.0, 1);
    /// assert!(result.is_err());
    /// ```
    pub fn new_with_param(alpha: f64, n: usize) -> Result<Dirichlet> {
        Self::new(&vec![alpha; n])
    }

    /// Returns the concentration parameters of
    /// the dirichlet distribution as a slice
    ///
    /// # Examples
    ///
    /// ```
    /// use statrs::distribution::Dirichlet;
    ///
    /// let n = Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap();
    /// assert_eq!(n.alpha(), [1.0, 2.0, 3.0]);
    /// ```
    pub fn alpha(&self) -> &[f64] {
        &self.alpha
    }

    fn alpha_sum(&self) -> f64 {
        self.alpha.iter().fold(0.0, |acc, x| acc + x)
    }
}

impl Sample<Vec<f64>> for Dirichlet {
    /// Generate random samples from a dirichlet
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn sample<R: Rng>(&mut self, r: &mut R) -> Vec<f64> {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<Vec<f64>> for Dirichlet {
    /// Generate random independent samples from a dirichlet
    /// distribution using `r` as the source of randomness.
    /// Refer [here](#method.sample-1) for implementation details
    fn ind_sample<R: Rng>(&self, r: &mut R) -> Vec<f64> {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<Vec<f64>> for Dirichlet {
    /// Generate random samples from the dirichlet distribution
    /// using `r` as the source of randomness
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Dirichlet, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap();
    /// print!("{:?}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> Vec<f64> {
        let n = self.alpha.len();
        let mut samples = vec![0.0; n];
        let mut sum = 0.0;
        for i in 0..n {
            samples[i] = super::gamma::sample_unchecked(r, self.alpha[i], 1.0);
            sum += samples[i];
        }
        for i in 0..n {
            samples[i] /= sum
        }
        samples
    }
}

impl Mean<Vec<f64>> for Dirichlet {
    /// Returns the means of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// α_i / α_0
    /// ```
    ///
    /// for the `i`th element where `α_i` is the `i`th concentration parameter
    /// and `α_0` is the sum of all concentration parameters
    fn mean(&self) -> Vec<f64> {
        let sum = self.alpha_sum();
        self.alpha.iter().map(|x| x / sum).collect()
    }
}

impl Variance<Vec<f64>> for Dirichlet {
    /// Returns the variances of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// (α_i * (α_0 - α_i)) / (α_0^2 * (α_0 + 1))
    /// ```
    ///
    /// for the `i`th element where `α_i` is the `i`th concentration parameter
    /// and `α_0` is the sum of all concentration parameters
    fn variance(&self) -> Vec<f64> {
        let sum = self.alpha_sum();
        self.alpha.iter().map(|x| x * (sum - x) / (sum * sum * (sum + 1.0))).collect()
    }

    /// Returns the variances of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sqrt((α_i * (α_0 - α_i)) / (α_0^2 * (α_0 + 1)))
    /// ```
    ///
    /// for the `i`th element where `α_i` is the `i`th concentration parameter
    /// and `α_0` is the sum of all concentration parameters
    fn std_dev(&self) -> Vec<f64> {
        self.variance().iter().map(|x| x.sqrt()).collect()
    }
}

impl Entropy<f64> for Dirichlet {
    /// Returns the entropy of the dirichlet distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// ln(B(α)) - (K - α_0)ψ(α_0) - ∑((α_i - 1)ψ(α_i))
    /// ```
    ///
    /// where
    ///
    /// ```ignore
    /// B(α) = ∏(Γ(α_i)) / Γ(∑(α_i))
    /// ```
    ///
    /// `α_0` is the sum of all concentration parameters,
    /// `K` is the number of concentration parameters, `ψ` is the digamma function, `α_i`
    /// is the `i`th concentration parameter, and `∑` is the sum from `1` to `K`
    fn entropy(&self) -> f64 {
        let sum = self.alpha_sum();
        let num = self.alpha.iter().fold(0.0, |acc, &x| acc + (x - 1.0) * gamma::digamma(x));
        gamma::ln_gamma(sum) + (sum - self.alpha.len() as f64) * gamma::digamma(sum) - num
    }
}

impl<'a> Continuous<&'a [f64], f64> for Dirichlet {
    /// Calculates the probabiliy density function for the dirichlet distribution
    /// with given `x`'s corresponding to the concentration parameters for this
    /// distribution
    ///
    /// # Panics
    ///
    /// If any element in `x` is not in `(0, 1)` or if `x` is not the same length
    /// as the vector of concentration parameters for this distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// (1 / B(α)) * ∏(x_i^(α_i - 1))
    /// ```
    ///
    /// where
    ///
    /// ```ignore
    /// B(α) = ∏(Γ(α_i)) / Γ(∑(α_i))
    /// ```
    ///
    /// `α` is the vector of concentration parameters, `α_i` is the `i`th
    /// concentration parameter, `x_i` is the `i`th argument corresponding to
    /// the `i`th concentration parameter, `Γ` is the gamma function,
    /// `∏` is the product from `1` to `K`, `∑` is the sum from `1` to `K`,
    /// and `K` is the number of concentration parameters
    fn pdf(&self, x: &[f64]) -> f64 {
        self.ln_pdf(x).exp()
    }

    /// Calculates the log probabiliy density function for the dirichlet distribution
    /// with given `x`'s corresponding to the concentration parameters for this
    /// distribution
    ///
    /// # Panics
    ///
    /// If any element in `x` is not in `(0, 1)` or if `x` is not the same length
    /// as the vector of concentration parameters for this distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// ln((1 / B(α)) * ∏(x_i^(α_i - 1)))
    /// ```
    ///
    /// where
    ///
    /// ```ignore
    /// B(α) = ∏(Γ(α_i)) / Γ(∑(α_i))
    /// ```
    ///
    /// `α` is the vector of concentration parameters, `α_i` is the `i`th
    /// concentration parameter, `x_i` is the `i`th argument corresponding to
    /// the `i`th concentration parameter, `Γ` is the gamma function,
    /// `∏` is the product from `1` to `K`, `∑` is the sum from `1` to `K`,
    /// and `K` is the number of concentration parameters
    fn ln_pdf(&self, x: &[f64]) -> f64 {
        assert!(self.alpha.len() == x.len(),
                format!("{}", StatsError::ContainersMustBeSameLength));

        let (term, sum_xi, sum_alpha) = x.iter()
            .enumerate()
            .map(|pair| (pair.1, self.alpha[pair.0]))
            .fold((0.0, 0.0, 0.0), |acc, pair| {
                assert!(*pair.0 > 0.0 && *pair.0 < 1.0,
                        format!("{}", StatsError::ArgIntervalExcl("x", 0.0, 1.0)));

                (acc.0 + (pair.1 - 1.0) * pair.0.ln() - gamma::ln_gamma(pair.1),
                 acc.1 + pair.0,
                 acc.2 + pair.1)
            });

        if !prec::almost_eq(sum_xi, 1.0, 1e-8) {
            0.0
        } else {
            term + gamma::ln_gamma(sum_alpha)
        }
    }
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[cfg(test)]
mod test {
    use std::f64;
    use function::gamma;
    use statistics::*;
    use distribution::{Continuous, Dirichlet};

    fn try_create(alpha: &[f64]) -> Dirichlet {
        let n = Dirichlet::new(alpha);
        assert!(n.is_ok());
        n.unwrap()
    }

    fn create_case(alpha: &[f64]) {
        let n = try_create(alpha);
        let a2 = n.alpha();
        for i in 0..alpha.len() {
            assert_eq!(alpha[i], a2[i]);
        }
    }

    fn bad_create_case(alpha: &[f64]) {
        let n = Dirichlet::new(alpha);
        assert!(n.is_err());
    }

    #[test]
    fn test_create() {
        create_case(&[1.0, 2.0, 3.0, 4.0, 5.0]);
        create_case(&[0.001, f64::INFINITY, 3756.0]);
    }

    #[test]
    fn test_bad_create() {
        bad_create_case(&[1.0]);
        bad_create_case(&[1.0, 2.0, 0.0, 4.0, 5.0]);
        bad_create_case(&[1.0, f64::NAN, 3.0, 4.0, 5.0]);
    }

    #[test]
    fn test_mean() {
        let n = Dirichlet::new_with_param(0.3, 5).unwrap();
        let res = n.mean();
        for x in res {
            assert_eq!(x, 0.3 / 1.5);
        }
    }

    #[test]
    fn test_variance() {
        let alpha = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
        let sum = alpha.iter().fold(0.0, |acc, x| acc + x);
        let n = Dirichlet::new(&alpha).unwrap();
        let res = n.variance();
        for i in 1..11 {
            let f = i as f64;
            assert_almost_eq!(res[i-1], f * (sum - f) / (sum * sum * (sum + 1.0)), 1e-15);
        }
    }

    #[test]
    fn test_std_dev() {
        let alpha = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
        let sum = alpha.iter().fold(0.0, |acc, x| acc + x);
        let n = Dirichlet::new(&alpha).unwrap();
        let res = n.std_dev();
        for i in 1..11 {
            let f = i as f64;
            assert_almost_eq!(res[i-1], (f * (sum - f) / (sum * sum * (sum + 1.0))).sqrt(), 1e-15);
        }
    }

    #[test]
    fn test_entropy() {
        let mut alpha = [0.1, 0.3, 0.5, 0.8];
        let mut n = try_create(&alpha);
        let mut sum = alpha.iter().fold(0.0, |acc, x| acc + x);
        let mut num = alpha.iter().fold(0.0, |acc, x| acc + (x - 1.0) * gamma::digamma(*x));
        assert_eq!(n.entropy(), gamma::ln_gamma(sum) + (sum - 4.0) * gamma::digamma(sum) - num);

        alpha = [0.1, 0.2, 0.3, 0.4];
        n = try_create(&alpha);
        sum = alpha.iter().fold(0.0, |acc, x| acc + x);
        num = alpha.iter().fold(0.0, |acc, x| acc + (x - 1.0) * gamma::digamma(*x));
        assert_eq!(n.entropy(), gamma::ln_gamma(sum) + (sum - 4.0) * gamma::digamma(sum) - num);
    }
}<|endoftext|>


<commit_before>extern crate octobuild;
extern crate petgraph;
extern crate tempdir;

use octobuild::common::BuildTask;
use octobuild::cache::Cache;
use octobuild::wincmd;
use octobuild::xg;
use octobuild::utils;
use octobuild::version;
use octobuild::vs::compiler::VsCompiler;
use octobuild::compiler::*;

use petgraph::{Graph, EdgeDirection};
use petgraph::graph::NodeIndex;
use tempdir::TempDir;

use std::fs::File;
use std::env;
use std::io::{BufReader, Error, ErrorKind, Write};
use std::io;
use std::iter::FromIterator;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::process;
use std::thread;

#[derive(Debug)]
struct TaskMessage {
	index: NodeIndex,
	task: BuildTask
}

#[derive(Debug)]
struct ResultMessage {
	index: NodeIndex,
	task: BuildTask,
	worker: usize,
	result: Result<OutputInfo, Error>
}

fn main() {
	println!("XGConsole ({}):", version::full_version());
	let args = Vec::from_iter(env::args());
	for arg in args.iter() {
		println!("  {}", arg);
	}
	process::exit(match execute(&args[1..]) {
		Ok(result) => {
			match result {
				Some(r) => r,
				None => 501
			}
		}
		Err(e) => {
			println!("FATAL ERROR: {:?}", e);
			500
		}
	})
}

fn execute(args: &[String]) -> Result<Option<i32>, Error> {
	let cache = Cache::new();
	let temp_dir = try! (TempDir::new("octobuild"));
	for arg in args.iter() {
		if arg.starts_with("/") {continue}

		let (tx_result, rx_result): (Sender<ResultMessage>, Receiver<ResultMessage>) = channel();
		let (tx_task, rx_task): (Sender<TaskMessage>, Receiver<TaskMessage>) = channel();

		let mutex_rx_task = create_threads(rx_task, tx_result, utils::num_cpus(), |worker_id:usize| {
			let temp_path = temp_dir.path().to_path_buf();
			let temp_cache = cache.clone();
			move |task:TaskMessage| -> ResultMessage {
				execute_task(&temp_cache, &temp_path, worker_id, task)
			}
		});

		let path = Path::new(arg);
		let file = try! (File::open(&path));
		let graph = try! (xg::parser::parse(BufReader::new(file)));
		let validated_graph = try! (validate_graph(graph));
		match try! (execute_graph(&validated_graph, tx_task, mutex_rx_task, rx_result)) {
			Some(v) if v == 0 => {}
			v => {return Ok(v)}
		}
	}
	cache.cleanup(16 * 1024 * 1024 * 1024);
	Ok(Some(0))
}

fn create_threads<R: 'static + Send, T: 'static + Send, Worker:'static + Fn(T) -> R + Send, Factory:Fn(usize) -> Worker>(rx_task: Receiver<T>, tx_result: Sender<R>, num_cpus: usize, factory: Factory) ->  Arc<Mutex<Receiver<T>>> {
	let mutex_rx_task = Arc::new(Mutex::new(rx_task));
	for cpu_id in 0..num_cpus {
		let local_rx_task = mutex_rx_task.clone();
		let local_tx_result = tx_result.clone();
		let worker = factory(cpu_id);
		thread::spawn(move || {
			loop {
				let task: T;
				match local_rx_task.lock().unwrap().recv() {
					Ok(v) => {task = v;}
					Err(_) => {break;}
				}
				match local_tx_result.send(worker(task)) {
					Ok(_) => {}
					Err(_) => {break;}
				}
			}
		});
	}
	mutex_rx_task
}

fn validate_graph(graph: Graph<BuildTask, ()>) -> Result<Graph<BuildTask, ()>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	let mut queue:Vec<NodeIndex> = Vec::new();
	for index in 0 .. graph.node_count() {
		completed.push(false);
		queue.push(NodeIndex::new(index));
	}
	let mut count:usize = 0;
	let mut i:usize = 0;
	while i < queue.len() {
		let index = queue[i];
		if (!completed[index.index()]) && (is_ready(&graph, &completed, &index)) {
			completed[index.index()] = true;
			for neighbor in graph.neighbors_directed(index, EdgeDirection::Incoming) {
				queue.push(neighbor);
			}
			count += 1;
			if count ==completed.len() {
				return Ok(graph);
			}
		}
		i = i + 1;
	}
	Err(Error::new(ErrorKind::InvalidInput, "Found cycles in build dependencies"))
}

fn execute_task(cache: &Cache, temp_dir: &Path, worker: usize, message: TaskMessage) -> ResultMessage {
	let args = wincmd::expand_args(&message.task.args, &|name:&str|->Option<String>{env::var(name).ok()});
	let output = execute_compiler(cache, temp_dir, &message.task.exec, &Path::new(&message.task.working_dir), &args);
	ResultMessage {
		index: message.index,
		task: message.task,
		worker: worker,
		result: output,
	}
}

fn execute_compiler(cache: &Cache, temp_dir: &Path, program: &str, cwd: &Path, args: &[String]) -> Result<OutputInfo, Error> {
	let command = CommandInfo {
		program: Path::new(program).to_path_buf(),
		current_dir: Some(cwd.to_path_buf()),
	};
	if Path::new(program).ends_with("cl.exe") {
		let compiler = VsCompiler::new(cache, temp_dir);
		compiler.compile(command, args)
	} else {
		command.to_command()
			.args(&args)
			.output()
			.map(|o| OutputInfo::new(o))
	}
}

fn execute_graph(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, mutex_rx_task: Arc<Mutex<Receiver<TaskMessage>>>, rx_result: Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	// Run all tasks.
	let result = execute_until_failed(graph, tx_task, &rx_result);
	// Cleanup task queue.
	for _ in mutex_rx_task.lock().unwrap().iter() {
	}
	// Wait for in progress task completion.
	for _ in rx_result.iter() {
	}
	result
}

fn execute_until_failed(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, rx_result: &Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	for _ in 0 .. graph.node_count() {
		completed.push(false);
	}
	for index in graph.without_edges(EdgeDirection::Outgoing) {
		try! (tx_task.send(TaskMessage{
			index: index,
			task: graph.node_weight(index).unwrap().clone(),
		}).map_err(|e| Error::new(ErrorKind::Other, e)));
	}

	let mut count:usize = 0;
	for message in rx_result.iter() {
		assert!(!completed[message.index.index()]);
		count += 1;
		println!("#{} {}/{}: {}", message.worker, count, completed.len(), message.task.title);
		let result = try! (message.result);
		try! (io::stdout().write_all(&result.stdout));
		try! (io::stderr().write_all(&result.stderr));
		if !result.success() {
			return Ok(result.status);
		}
		completed[message.index.index()] = true;

		for source in graph.neighbors_directed(message.index, EdgeDirection::Incoming) {
			if !completed[source.index()] {
				if is_ready(graph, &completed, &source) {
					try! (tx_task.send(TaskMessage{
						index: source,
						task: graph.node_weight(source).unwrap().clone(),
					}).map_err(|e| Error::new(ErrorKind::Other, e)));
				}
			}
		}

		if count == completed.len() {
			return Ok(Some(0));
		}
	}
	panic! ("Unexpected end of result pipe");
}

fn is_ready(graph: &Graph<BuildTask, ()>, completed: &Vec<bool>, source: &NodeIndex) -> bool {
	for neighbor in graph.neighbors_directed(*source, EdgeDirection::Outgoing) {
		if !completed[neighbor.index()]{
			return false
		}
	}
	true
}
<commit_msg>Fix double task execution error<commit_after>extern crate octobuild;
extern crate petgraph;
extern crate tempdir;

use octobuild::common::BuildTask;
use octobuild::cache::Cache;
use octobuild::wincmd;
use octobuild::xg;
use octobuild::utils;
use octobuild::version;
use octobuild::vs::compiler::VsCompiler;
use octobuild::compiler::*;

use petgraph::{Graph, EdgeDirection};
use petgraph::graph::NodeIndex;
use tempdir::TempDir;

use std::fs::File;
use std::env;
use std::io::{BufReader, Error, ErrorKind, Write};
use std::io;
use std::iter::FromIterator;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::process;
use std::thread;

#[derive(Debug)]
struct TaskMessage {
	index: NodeIndex,
	task: BuildTask
}

#[derive(Debug)]
struct ResultMessage {
	index: NodeIndex,
	task: BuildTask,
	worker: usize,
	result: Result<OutputInfo, Error>
}

fn main() {
	println!("XGConsole ({}):", version::full_version());
	let args = Vec::from_iter(env::args());
	for arg in args.iter() {
		println!("  {}", arg);
	}
	process::exit(match execute(&args[1..]) {
		Ok(result) => {
			match result {
				Some(r) => r,
				None => 501
			}
		}
		Err(e) => {
			println!("FATAL ERROR: {:?}", e);
			500
		}
	})
}

fn execute(args: &[String]) -> Result<Option<i32>, Error> {
	let cache = Cache::new();
	let temp_dir = try! (TempDir::new("octobuild"));
	for arg in args.iter() {
		if arg.starts_with("/") {continue}

		let (tx_result, rx_result): (Sender<ResultMessage>, Receiver<ResultMessage>) = channel();
		let (tx_task, rx_task): (Sender<TaskMessage>, Receiver<TaskMessage>) = channel();

		let mutex_rx_task = create_threads(rx_task, tx_result, utils::num_cpus(), |worker_id:usize| {
			let temp_path = temp_dir.path().to_path_buf();
			let temp_cache = cache.clone();
			move |task:TaskMessage| -> ResultMessage {
				execute_task(&temp_cache, &temp_path, worker_id, task)
			}
		});

		let path = Path::new(arg);
		let file = try! (File::open(&path));
		let graph = try! (xg::parser::parse(BufReader::new(file)));
		let validated_graph = try! (validate_graph(graph));
		match try! (execute_graph(&validated_graph, tx_task, mutex_rx_task, rx_result)) {
			Some(v) if v == 0 => {}
			v => {return Ok(v)}
		}
	}
	cache.cleanup(16 * 1024 * 1024 * 1024);
	Ok(Some(0))
}

fn create_threads<R: 'static + Send, T: 'static + Send, Worker:'static + Fn(T) -> R + Send, Factory:Fn(usize) -> Worker>(rx_task: Receiver<T>, tx_result: Sender<R>, num_cpus: usize, factory: Factory) ->  Arc<Mutex<Receiver<T>>> {
	let mutex_rx_task = Arc::new(Mutex::new(rx_task));
	for cpu_id in 0..num_cpus {
		let local_rx_task = mutex_rx_task.clone();
		let local_tx_result = tx_result.clone();
		let worker = factory(cpu_id);
		thread::spawn(move || {
			loop {
				let task: T;
				match local_rx_task.lock().unwrap().recv() {
					Ok(v) => {task = v;}
					Err(_) => {break;}
				}
				match local_tx_result.send(worker(task)) {
					Ok(_) => {}
					Err(_) => {break;}
				}
			}
		});
	}
	mutex_rx_task
}

fn validate_graph(graph: Graph<BuildTask, ()>) -> Result<Graph<BuildTask, ()>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	let mut queue:Vec<NodeIndex> = Vec::new();
	for index in 0 .. graph.node_count() {
		completed.push(false);
		queue.push(NodeIndex::new(index));
	}
	let mut count:usize = 0;
	let mut i:usize = 0;
	while i < queue.len() {
		let index = queue[i];
		if (!completed[index.index()]) && (is_ready(&graph, &completed, &index)) {
			completed[index.index()] = true;
			for neighbor in graph.neighbors_directed(index, EdgeDirection::Incoming) {
				queue.push(neighbor);
			}
			count += 1;
			if count ==completed.len() {
				return Ok(graph);
			}
		}
		i = i + 1;
	}
	Err(Error::new(ErrorKind::InvalidInput, "Found cycles in build dependencies"))
}

fn execute_task(cache: &Cache, temp_dir: &Path, worker: usize, message: TaskMessage) -> ResultMessage {
	let args = wincmd::expand_args(&message.task.args, &|name:&str|->Option<String>{env::var(name).ok()});
	let output = execute_compiler(cache, temp_dir, &message.task.exec, &Path::new(&message.task.working_dir), &args);
	ResultMessage {
		index: message.index,
		task: message.task,
		worker: worker,
		result: output,
	}
}

fn execute_compiler(cache: &Cache, temp_dir: &Path, program: &str, cwd: &Path, args: &[String]) -> Result<OutputInfo, Error> {
	let command = CommandInfo {
		program: Path::new(program).to_path_buf(),
		current_dir: Some(cwd.to_path_buf()),
	};
	if Path::new(program).ends_with("cl.exe") {
		let compiler = VsCompiler::new(cache, temp_dir);
		compiler.compile(command, args)
	} else {
		command.to_command()
			.args(&args)
			.output()
			.map(|o| OutputInfo::new(o))
	}
}

fn execute_graph(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, mutex_rx_task: Arc<Mutex<Receiver<TaskMessage>>>, rx_result: Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	// Run all tasks.
	let result = execute_until_failed(graph, tx_task, &rx_result);
	// Cleanup task queue.
	for _ in mutex_rx_task.lock().unwrap().iter() {
	}
	// Wait for in progress task completion.
	for _ in rx_result.iter() {
	}
	result
}

fn execute_until_failed(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, rx_result: &Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	let mut waiting:Vec<bool> = Vec::new();
	for _ in 0 .. graph.node_count() {
		completed.push(false);
		waiting.push(false);
	}
	for index in graph.without_edges(EdgeDirection::Outgoing) {
		waiting[index.index()] = true;
		try! (tx_task.send(TaskMessage{
			index: index,
			task: graph.node_weight(index).unwrap().clone(),
		}).map_err(|e| Error::new(ErrorKind::Other, e)));
	}

	let mut count:usize = 0;
	for message in rx_result.iter() {
		assert!(!completed[message.index.index()]);
		count += 1;
		println!("#{} {}/{}: {}", message.worker, count, completed.len(), message.task.title);
		let result = try! (message.result);
		try! (io::stdout().write_all(&result.stdout));
		try! (io::stderr().write_all(&result.stderr));
		if !result.success() {
			return Ok(result.status);
		}
		completed[message.index.index()] = true;

		for source in graph.neighbors_directed(message.index, EdgeDirection::Incoming) {
			if is_ready(graph, &completed, &source) {
				if !waiting[source.index()] {
					waiting[source.index()] = true;
					try! (tx_task.send(TaskMessage{
						index: source,
						task: graph.node_weight(source).unwrap().clone(),
					}).map_err(|e| Error::new(ErrorKind::Other, e)));
				}
			}
		}

		if count == completed.len() {
			return Ok(Some(0));
		}
	}
	panic! ("Unexpected end of result pipe");
}

fn is_ready(graph: &Graph<BuildTask, ()>, completed: &Vec<bool>, source: &NodeIndex) -> bool {
	for neighbor in graph.neighbors_directed(*source, EdgeDirection::Outgoing) {
		if !completed[neighbor.index()]{
			return false
		}
	}
	true
}
<|endoftext|>


<commit_before>//! A simple example of hooking up stdin/stdout to a WebSocket stream.
//!
//! This example will connect to a server specified in the argument list and
//! then forward all data read on stdin to the server, printing out all data
//! received on stdout.
//!
//! Note that this is not currently optimized for performance, especially around
//! buffer management. Rather it's intended to show an example of working with a
//! client.
//!
//! You can use this example together with the `server` example.

use std::env;

use futures_util::{future, pin_mut, StreamExt};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio_tungstenite::connect_async;
use tungstenite::protocol::Message;

#[tokio::main]
async fn main() {
    let connect_addr = env::args()
        .nth(1)
        .unwrap_or_else(|| panic!("this program requires at least one argument"));

    let url = url::Url::parse(&connect_addr).unwrap();

    let (stdin_tx, stdin_rx) = futures_channel::mpsc::unbounded();
    tokio::spawn(read_stdin(stdin_tx));

    let (ws_stream, _) = connect_async(url).await.expect("Failed to connect");
    println!("WebSocket handshake has been successfully completed");

    let (write, read) = ws_stream.split();

    let stdin_to_ws = stdin_rx.map(Ok).forward(write);
    let ws_to_stdout = {
        read.for_each(|message| async {
            let data = message.unwrap().into_data();
            tokio::io::stdout().write_all(&data).await.unwrap();
        })
    };

    pin_mut!(stdin_to_ws, ws_to_stdout);
    future::select(stdin_to_ws, ws_to_stdout).await;
}

// Our helper method which will read data from stdin and send it along the
// sender provided.
async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender<Message>) {
    let mut stdin = tokio::io::stdin();
    loop {
        let mut buf = vec![0; 1024];
        let n = match stdin.read(&mut buf).await {
            Err(_) | Ok(0) => break,
            Ok(n) => n,
        };
        buf.truncate(n);
        tx.unbounded_send(Message::binary(buf)).unwrap();
    }
}
<commit_msg>use re-expored tungstenite<commit_after>//! A simple example of hooking up stdin/stdout to a WebSocket stream.
//!
//! This example will connect to a server specified in the argument list and
//! then forward all data read on stdin to the server, printing out all data
//! received on stdout.
//!
//! Note that this is not currently optimized for performance, especially around
//! buffer management. Rather it's intended to show an example of working with a
//! client.
//!
//! You can use this example together with the `server` example.

use std::env;

use futures_util::{future, pin_mut, StreamExt};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio_tungstenite::{connect_async, tungstenite::protocol::Message};

#[tokio::main]
async fn main() {
    let connect_addr = env::args()
        .nth(1)
        .unwrap_or_else(|| panic!("this program requires at least one argument"));

    let url = url::Url::parse(&connect_addr).unwrap();

    let (stdin_tx, stdin_rx) = futures_channel::mpsc::unbounded();
    tokio::spawn(read_stdin(stdin_tx));

    let (ws_stream, _) = connect_async(url).await.expect("Failed to connect");
    println!("WebSocket handshake has been successfully completed");

    let (write, read) = ws_stream.split();

    let stdin_to_ws = stdin_rx.map(Ok).forward(write);
    let ws_to_stdout = {
        read.for_each(|message| async {
            let data = message.unwrap().into_data();
            tokio::io::stdout().write_all(&data).await.unwrap();
        })
    };

    pin_mut!(stdin_to_ws, ws_to_stdout);
    future::select(stdin_to_ws, ws_to_stdout).await;
}

// Our helper method which will read data from stdin and send it along the
// sender provided.
async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender<Message>) {
    let mut stdin = tokio::io::stdin();
    loop {
        let mut buf = vec![0; 1024];
        let n = match stdin.read(&mut buf).await {
            Err(_) | Ok(0) => break,
            Ok(n) => n,
        };
        buf.truncate(n);
        tx.unbounded_send(Message::binary(buf)).unwrap();
    }
}
<|endoftext|>


<commit_before>extern crate dirs;
extern crate itertools;
extern crate symlink;
extern crate winapi;
extern crate winreg;

mod registry;

use std::path::Path;
use std::fs::{remove_file, read_link};
use std::env;
use symlink::symlink_file;

fn main() {
    let args: Vec<String> = env::args().collect();
    let addins_dir = Path::new(&dirs::home_dir().expect("Couldn't find homedir"))
        .join("AppData").join("Roaming").join("Microsoft").join("Addins");

    if args[1] == "--install" {
        let install_dir = Path::new(&args[2]);

        // Excel
        ensure_file_symlink(
            install_dir.join("addins").join("vba-blocks.xlam").as_path(),
            addins_dir.join("vba-blocks.xlam").as_path()
        );
        registry::add_to_open("Excel", "vba-blocks.xlam");
        registry::enable_vbom("Excel");
    } else {
        registry::remove_from_open("Excel", "vba-blocks.xlam");
        remove_symlink(addins_dir.join("vba-blocks.xlam").as_path());
    }
}

fn ensure_file_symlink(target: &Path, path: &Path) {
    remove_symlink(path);
    symlink_file(target, path).expect("Failed to create symlink to add-in");
}

fn remove_symlink(path: &Path) {
    println!("Path exists? {:?} {}", path, path.exists());
    if path.exists() || symlink_exists(path) {
        remove_file(path).expect("Failed to remove symlink to add-in");
    }
}

fn symlink_exists(path: &Path) -> bool {
    match read_link(path) {
        Ok(_) => true,
        Err(_) => false
    }
}
<commit_msg>Log (and swallow) errors in installer actions<commit_after>extern crate dirs;
extern crate itertools;
extern crate symlink;
extern crate winapi;
extern crate winreg;

mod registry;

use std::env;
use std::fs::{remove_file, read_link, OpenOptions};
use std::io::Write;
use std::panic;
use std::path::Path;
use symlink::symlink_file;

fn main() {
    let result = panic::catch_unwind(|| {
        let args: Vec<String> = env::args().collect();
        let addins_dir = Path::new(&dirs::home_dir().expect("Couldn't find homedir"))
            .join("AppData").join("Roaming").join("Microsoft").join("Addins");

        if args[1] == "--install" {
            let install_dir = Path::new(&args[2]);

            // Excel
            ensure_file_symlink(
                install_dir.join("addins").join("vba-blocks.xlam").as_path(),
                addins_dir.join("vba-blocks.xlam").as_path()
            );
            registry::add_to_open("Excel", "vba-blocks.xlam");
            registry::enable_vbom("Excel");
        } else {
            registry::remove_from_open("Excel", "vba-blocks.xlam");
            remove_symlink(addins_dir.join("vba-blocks.xlam").as_path());
        }
    });

    if let Err(cause) = result {
        let message = if let Some(e) = cause.downcast_ref::<&'static str>() {
            format!("Installation failed: {}", e)
        } else {
            "Installation failed with an unknown error".to_string()
        };
        eprintln!("{}", message);

        let mut path = env::temp_dir();
        path.push("vba-blocks.log");

        let mut log = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path);

        if let Err(e) = log {
            eprintln!("Log failed: {:?}", e);
            return;
        }

        if let Err(e) = writeln!(log.unwrap(), "{}", message) {
            eprintln!("Write log failed: {:?}", e);
            return;
        } else {
            eprintln!("Log written to {:?}", path);
        }
    }
}

fn ensure_file_symlink(target: &Path, path: &Path) {
    remove_symlink(path);
    symlink_file(target, path).expect("Failed to create symlink to add-in");
}

fn remove_symlink(path: &Path) {
    println!("Path exists? {:?} {}", path, path.exists());
    if path.exists() || symlink_exists(path) {
        remove_file(path).expect("Failed to remove symlink to add-in");
    }
}

fn symlink_exists(path: &Path) -> bool {
    match read_link(path) {
        Ok(_) => true,
        Err(_) => false
    }
}
<|endoftext|>


<commit_before>use flowgger::config::Config;
use openssl::bn::BigNum;
use openssl::dh::DH;
use openssl::ssl::*;
use openssl::x509::X509FileType;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use openssl::x509::X509;

pub mod tls_input;
#[cfg(feature = "coroutines")]
pub mod tlsco_input;

pub use super::Input;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:ECDHE-ECDSA-DES-CBC3-SHA:ECDHE-RSA-DES-CBC3-SHA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SHA";
const DEFAULT_COMPRESSION: bool = false;
const DEFAULT_FRAMING: &'static str = "line";
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_LISTEN: &'static str = "0.0.0.0:6514";
#[cfg(feature = "coroutines")]
const DEFAULT_THREADS: usize = 1;
const DEFAULT_TIMEOUT: u64 = 3600;
const DEFAULT_TLS_METHOD: &'static str = "any";
const DEFAULT_VERIFY_PEER: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;

#[derive(Clone)]
pub struct TlsConfig {
    framing: String,
    threads: usize,
    arc_ctx: Arc<SslContext>
}

#[cfg(feature = "ecdh")]
fn set_ecdh(ctx: &mut SslContext) {
    ctx.set_ecdh_auto(true).unwrap();
}

#[cfg(not(feature = "ecdh"))]
fn set_ecdh(ctx: &mut SslContext) {
    let _ = ctx;
}

fn set_fs(ctx: &mut SslContext) {
    let p = BigNum::from_hex_str("87A8E61DB4B6663CFFBBD19C651959998CEEF608660DD0F25D2CEED4435E3B00E00DF8F1D61957D4FAF7DF4561B2AA3016C3D91134096FAA3BF4296D830E9A7C209E0C6497517ABD5A8A9D306BCF67ED91F9E6725B4758C022E0B1EF4275BF7B6C5BFC11D45F9088B941F54EB1E59BB8BC39A0BF12307F5C4FDB70C581B23F76B63ACAE1CAA6B7902D52526735488A0EF13C6D9A51BFA4AB3AD8347796524D8EF6A167B5A41825D967E144E5140564251CCACB83E6B486F6B3CA3F7971506026C0B857F689962856DED4010ABD0BE621C3A3960A54E710C375F26375D7014103A4B54330C198AF126116D2276E11715F693877FAD7EF09CADB094AE91E1A1597").unwrap();
    let g = BigNum::from_hex_str("3FB32C9B73134D0B2E77506660EDBD484CA7B18F21EF205407F4793A1A0BA12510DBC15077BE463FFF4FED4AAC0BB555BE3A6C1B0C6B47B1BC3773BF7E8C6F62901228F8C28CBB18A55AE31341000A650196F931C77A57F2DDF463E5E9EC144B777DE62AAAB8A8628AC376D282D6ED3864E67982428EBC831D14348F6F2F9193B5045AF2767164E1DFC967C1FB3F2E55A4BD1BFFE83B9C80D052B985D182EA0ADB2A3B7313D3FE14C8484B1E052588B9B7D2BBD2DF016199ECD06E1557CD0915B3353BBB64E0EC377FD028370DF92B52C7891428CDC67EB6184B523D1DB246C32F63078490F00EF8D647D148D47954515E2327CFEF98C582664B4C0F6CC41659").unwrap();
    let q = BigNum::from_hex_str("8CF83642A709A097B447997640129DA299B1A47D1EB3750BA308B0FE64F5FBD3").unwrap();
    let dh = DH::from_params(p, g, q).unwrap();
    ctx.set_tmp_dh(dh).unwrap();
    set_ecdh(ctx);
}

#[cfg(feature = "coroutines")]
fn get_default_threads(config: &Config) -> usize {
    config.lookup("input.tls_threads").map_or(DEFAULT_THREADS, |x| x.as_integer().
        expect("input.tls_threads must be an unsigned integer") as usize)
}

#[cfg(not(feature = "coroutines"))]
fn get_default_threads(_config: &Config) -> usize {
    1
}

pub fn config_parse(config: &Config) -> (TlsConfig, String, u64) {
    let listen = config.lookup("input.listen").map_or(DEFAULT_LISTEN, |x| x.as_str().
        expect("input.listen must be an ip:port string")).to_owned();
    let threads = get_default_threads(&config);
    let cert = config.lookup("input.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
        expect("input.tls_cert must be a path to a .pem file")).to_owned();
    let key = config.lookup("input.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
        expect("input.tls_key must be a path to a .pem file")).to_owned();
    let ciphers = config.lookup("input.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
        expect("input.tls_ciphers must be a string with a cipher suite")).to_owned();
    let tls_method = match config.lookup("input.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
        expect("input.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
            "any" | "sslv23" => SslMethod::Sslv23,
            "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
            "tlsv1.1" => SslMethod::Tlsv1_1,
            "tlsv1.2" => SslMethod::Tlsv1_2,
            _ => panic!(r#"TLS method must be "any", "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
    };
    let verify_peer = config.lookup("input.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
        expect("input.tls_verify_peer must be a boolean"));
    let ca_file: Option<PathBuf> = config.lookup("input.tls_ca_file").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("input.tls_ca_file must be a path to a file"))));
    let compression = config.lookup("input.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
        expect("input.tls_compression must be a boolean"));
    let extra_chain_path: Option<PathBuf> = config.lookup("input.tls_cert_chain").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("input.tls_extra_chain_path must be a path to a file"))));
    let timeout = config.lookup("input.timeout").map_or(DEFAULT_TIMEOUT, |x| x.as_integer().
        expect("input.timeout must be an integer") as u64);
    let framing = if config.lookup("input.framed").map_or(false, |x| x.as_bool().
        expect("input.framed must be a boolean")) {
        "syslen"
    } else {
        DEFAULT_FRAMING
    };
    let framing = config.lookup("input.framing").map_or(framing, |x| x.as_str().
        expect(r#"input.framing must be a string set to "line", "nul" or "syslen""#)).to_owned();
    let mut ctx = SslContext::new(tls_method).unwrap();
    if verify_peer == false {
        ctx.set_verify(SSL_VERIFY_NONE, None);
    } else {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        if let Some(ca_file) = ca_file {
            ctx.set_CA_file(&ca_file).expect("Unable to read the trusted CA file");
        }
    }
    if let Some(extra_chain_path) = extra_chain_path {
         let mut extra_chain_file = File::open(extra_chain_path.as_path())
             .ok()
             .expect("Failed to open the extra chain cert file");
         let extra_chain_cert = X509::from_pem(&mut extra_chain_file).ok().expect("unable to load the provided extra chain cert");
         ctx.add_extra_chain_cert(&extra_chain_cert).expect("Unable to read the certificate chain");
    }

    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    set_fs(&mut ctx);
    ctx.set_certificate_file(&Path::new(&cert), X509FileType::PEM).
        expect("Unable to read the TLS certificate");
    ctx.set_private_key_file(&Path::new(&key), X509FileType::PEM).
        expect("Unable to read the TLS key");
    ctx.set_cipher_list(&ciphers).expect("Unsupported cipher suite");
    let arc_ctx = Arc::new(ctx);
    let tls_config = TlsConfig {
        framing: framing,
        threads: threads,
        arc_ctx: arc_ctx
    };
    (tls_config, listen, timeout)
}
<commit_msg>Revert "adding support for cert chain file in tls inputs"<commit_after>use flowgger::config::Config;
use openssl::bn::BigNum;
use openssl::dh::DH;
use openssl::ssl::*;
use openssl::x509::X509FileType;
use std::path::{Path, PathBuf};
use std::sync::Arc;

pub mod tls_input;
#[cfg(feature = "coroutines")]
pub mod tlsco_input;

pub use super::Input;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:ECDHE-ECDSA-DES-CBC3-SHA:ECDHE-RSA-DES-CBC3-SHA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SHA";
const DEFAULT_COMPRESSION: bool = false;
const DEFAULT_FRAMING: &'static str = "line";
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_LISTEN: &'static str = "0.0.0.0:6514";
#[cfg(feature = "coroutines")]
const DEFAULT_THREADS: usize = 1;
const DEFAULT_TIMEOUT: u64 = 3600;
const DEFAULT_TLS_METHOD: &'static str = "any";
const DEFAULT_VERIFY_PEER: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;

#[derive(Clone)]
pub struct TlsConfig {
    framing: String,
    threads: usize,
    arc_ctx: Arc<SslContext>
}

#[cfg(feature = "ecdh")]
fn set_ecdh(ctx: &mut SslContext) {
    ctx.set_ecdh_auto(true).unwrap();
}

#[cfg(not(feature = "ecdh"))]
fn set_ecdh(ctx: &mut SslContext) {
    let _ = ctx;
}

fn set_fs(ctx: &mut SslContext) {
    let p = BigNum::from_hex_str("87A8E61DB4B6663CFFBBD19C651959998CEEF608660DD0F25D2CEED4435E3B00E00DF8F1D61957D4FAF7DF4561B2AA3016C3D91134096FAA3BF4296D830E9A7C209E0C6497517ABD5A8A9D306BCF67ED91F9E6725B4758C022E0B1EF4275BF7B6C5BFC11D45F9088B941F54EB1E59BB8BC39A0BF12307F5C4FDB70C581B23F76B63ACAE1CAA6B7902D52526735488A0EF13C6D9A51BFA4AB3AD8347796524D8EF6A167B5A41825D967E144E5140564251CCACB83E6B486F6B3CA3F7971506026C0B857F689962856DED4010ABD0BE621C3A3960A54E710C375F26375D7014103A4B54330C198AF126116D2276E11715F693877FAD7EF09CADB094AE91E1A1597").unwrap();
    let g = BigNum::from_hex_str("3FB32C9B73134D0B2E77506660EDBD484CA7B18F21EF205407F4793A1A0BA12510DBC15077BE463FFF4FED4AAC0BB555BE3A6C1B0C6B47B1BC3773BF7E8C6F62901228F8C28CBB18A55AE31341000A650196F931C77A57F2DDF463E5E9EC144B777DE62AAAB8A8628AC376D282D6ED3864E67982428EBC831D14348F6F2F9193B5045AF2767164E1DFC967C1FB3F2E55A4BD1BFFE83B9C80D052B985D182EA0ADB2A3B7313D3FE14C8484B1E052588B9B7D2BBD2DF016199ECD06E1557CD0915B3353BBB64E0EC377FD028370DF92B52C7891428CDC67EB6184B523D1DB246C32F63078490F00EF8D647D148D47954515E2327CFEF98C582664B4C0F6CC41659").unwrap();
    let q = BigNum::from_hex_str("8CF83642A709A097B447997640129DA299B1A47D1EB3750BA308B0FE64F5FBD3").unwrap();
    let dh = DH::from_params(p, g, q).unwrap();
    ctx.set_tmp_dh(dh).unwrap();
    set_ecdh(ctx);
}

#[cfg(feature = "coroutines")]
fn get_default_threads(config: &Config) -> usize {
    config.lookup("input.tls_threads").map_or(DEFAULT_THREADS, |x| x.as_integer().
        expect("input.tls_threads must be an unsigned integer") as usize)
}

#[cfg(not(feature = "coroutines"))]
fn get_default_threads(_config: &Config) -> usize {
    1
}

pub fn config_parse(config: &Config) -> (TlsConfig, String, u64) {
    let listen = config.lookup("input.listen").map_or(DEFAULT_LISTEN, |x| x.as_str().
        expect("input.listen must be an ip:port string")).to_owned();
    let threads = get_default_threads(&config);
    let cert = config.lookup("input.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
        expect("input.tls_cert must be a path to a .pem file")).to_owned();
    let key = config.lookup("input.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
        expect("input.tls_key must be a path to a .pem file")).to_owned();
    let ciphers = config.lookup("input.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
        expect("input.tls_ciphers must be a string with a cipher suite")).to_owned();
    let tls_method = match config.lookup("input.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
        expect("input.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
            "any" | "sslv23" => SslMethod::Sslv23,
            "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
            "tlsv1.1" => SslMethod::Tlsv1_1,
            "tlsv1.2" => SslMethod::Tlsv1_2,
            _ => panic!(r#"TLS method must be "any", "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
    };
    let verify_peer = config.lookup("input.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
        expect("input.tls_verify_peer must be a boolean"));
    let ca_file: Option<PathBuf> = config.lookup("input.tls_ca_file").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("input.tls_ca_file must be a path to a file"))));
    let compression = config.lookup("input.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
        expect("input.tls_compression must be a boolean"));
    let timeout = config.lookup("input.timeout").map_or(DEFAULT_TIMEOUT, |x| x.as_integer().
        expect("input.timeout must be an integer") as u64);
    let framing = if config.lookup("input.framed").map_or(false, |x| x.as_bool().
        expect("input.framed must be a boolean")) {
        "syslen"
    } else {
        DEFAULT_FRAMING
    };
    let framing = config.lookup("input.framing").map_or(framing, |x| x.as_str().
        expect(r#"input.framing must be a string set to "line", "nul" or "syslen""#)).to_owned();
    let mut ctx = SslContext::new(tls_method).unwrap();
    if verify_peer == false {
        ctx.set_verify(SSL_VERIFY_NONE, None);
    } else {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        if let Some(ca_file) = ca_file {
            ctx.set_CA_file(&ca_file).expect("Unable to read the trusted CA file");
        }
    }
    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    set_fs(&mut ctx);
    ctx.set_certificate_file(&Path::new(&cert), X509FileType::PEM).
        expect("Unable to read the TLS certificate");
    ctx.set_private_key_file(&Path::new(&key), X509FileType::PEM).
        expect("Unable to read the TLS key");
    ctx.set_cipher_list(&ciphers).expect("Unsupported cipher suite");
    let arc_ctx = Arc::new(ctx);
    let tls_config = TlsConfig {
        framing: framing,
        threads: threads,
        arc_ctx: arc_ctx
    };
    (tls_config, listen, timeout)
}
<|endoftext|>