<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Weighted index sampling
//!
//! This module provides two implementations for sampling indices:
//!
//! *   [`WeightedIndex`] allows `O(log N)` sampling
//! *   [`alias_method::WeightedIndex`] allows `O(1)` sampling, but with
//!      much greater set-up cost
//!
//! [`alias_method::WeightedIndex`]: alias_method/struct.WeightedIndex.html

pub mod alias_method;

use crate::Rng;
use crate::distributions::Distribution;
use crate::distributions::uniform::{UniformSampler, SampleUniform, SampleBorrow};
use core::cmp::PartialOrd;
use core::fmt;

// Note that this whole module is only imported if feature="alloc" is enabled.
#[cfg(not(feature="std"))] use crate::alloc::vec::Vec;

/// A distribution using weighted sampling to pick a discretely selected
/// item.
///
/// Sampling a `WeightedIndex` distribution returns the index of a randomly
/// selected element from the iterator used when the `WeightedIndex` was
/// created. The chance of a given element being picked is proportional to the
/// value of the element. The weights can use any type `X` for which an
/// implementation of [`Uniform<X>`] exists.
///
/// # Performance
///
/// A `WeightedIndex<X>` contains a `Vec<X>` and a [`Uniform<X>`] and so its
/// size is the sum of the size of those objects, possibly plus some alignment.
///
/// Creating a `WeightedIndex<X>` will allocate enough space to hold `N - 1`
/// weights of type `X`, where `N` is the number of weights. However, since
/// `Vec` doesn't guarantee a particular growth strategy, additional memory
/// might be allocated but not used. Since the `WeightedIndex` object also
/// contains, this might cause additional allocations, though for primitive
/// types, ['Uniform<X>`] doesn't allocate any memory.
///
/// Time complexity of sampling from `WeightedIndex` is `O(log N)` where
/// `N` is the number of weights.
///
/// Sampling from `WeightedIndex` will result in a single call to
/// `Uniform<X>::sample` (method of the [`Distribution`] trait), which typically
/// will request a single value from the underlying [`RngCore`], though the
/// exact number depends on the implementaiton of `Uniform<X>::sample`.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand::distributions::WeightedIndex;
///
/// let choices = ['a', 'b', 'c'];
/// let weights = [2,   1,   1];
/// let dist = WeightedIndex::new(&weights).unwrap();
/// let mut rng = thread_rng();
/// for _ in 0..100 {
///     // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c'
///     println!("{}", choices[dist.sample(&mut rng)]);
/// }
///
/// let items = [('a', 0), ('b', 3), ('c', 7)];
/// let dist2 = WeightedIndex::new(items.iter().map(|item| item.1)).unwrap();
/// for _ in 0..100 {
///     // 0% chance to print 'a', 30% chance to print 'b', 70% chance to print 'c'
///     println!("{}", items[dist2.sample(&mut rng)].0);
/// }
/// ```
///
/// [`Uniform<X>`]: crate::distributions::uniform::Uniform
/// [`RngCore`]: crate::RngCore
#[derive(Debug, Clone)]
pub struct WeightedIndex<X: SampleUniform + PartialOrd> {
    cumulative_weights: Vec<X>,
    total_weight: X,
    weight_distribution: X::Sampler,
}

impl<X: SampleUniform + PartialOrd> WeightedIndex<X> {
    /// Creates a new a `WeightedIndex` [`Distribution`] using the values
    /// in `weights`. The weights can use any type `X` for which an
    /// implementation of [`Uniform<X>`] exists.
    ///
    /// Returns an error if the iterator is empty, if any weight is `< 0`, or
    /// if its total value is 0.
    ///
    /// [`Uniform<X>`]: crate::distributions::uniform::Uniform
    pub fn new<I>(weights: I) -> Result<WeightedIndex<X>, WeightedError>
        where I: IntoIterator,
              I::Item: SampleBorrow<X>,
              X: for<'a> ::core::ops::AddAssign<&'a X> +
                 Clone +
                 Default {
        let mut iter = weights.into_iter();
        let mut total_weight: X = iter.next()
                                      .ok_or(WeightedError::NoItem)?
                                      .borrow()
                                      .clone();

        let zero = <X as Default>::default();
        if total_weight < zero {
            return Err(WeightedError::InvalidWeight);
        }

        let mut weights = Vec::<X>::with_capacity(iter.size_hint().0);
        for w in iter {
            if *w.borrow() < zero {
                return Err(WeightedError::InvalidWeight);
            }
            weights.push(total_weight.clone());
            total_weight += w.borrow();
        }

        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }
        let distr = X::Sampler::new(zero, total_weight.clone());

        Ok(WeightedIndex { cumulative_weights: weights, total_weight, weight_distribution: distr })
    }

    /// Update a subset of weights, without changing the number of weights.
    ///
    /// Using this method instead of `new` might be more efficient if only a small number of
    /// weights is modified. No allocations are performed, unless the weight type `X` uses
    /// allocation internally.
    ///
    /// In case of error, `self` is not modified.
    pub fn update_weights(&mut self, new_weights: &[(usize, &X)]) -> Result<(), WeightedError>
        where X: for<'a> ::core::ops::AddAssign<&'a X> +
                 for<'a> ::core::ops::SubAssign<&'a X> +
                 Clone +
                 Default {
        let zero = <X as Default>::default();

        let mut total_weight = self.total_weight.clone();

        for &(i, w) in new_weights {
            if *w < zero {
                return Err(WeightedError::InvalidWeight);
            }
            if i >= self.cumulative_weights.len() {
                return Err(WeightedError::TooMany);
            }

            // Unfortunately, we will have to calculate the non-cumulative weight a second time, to
            // avoid producing an invalid state of `self`.
            let mut old_w = self.cumulative_weights[i].clone();
            if i > 0 {
                old_w -= &self.cumulative_weights[i - 1];
            }

            total_weight -= &old_w;
            total_weight += w;
        }
        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }

        for &(i, w) in new_weights {
            let mut old_w = self.cumulative_weights[i].clone();
            if i > 0 {
                old_w -= &self.cumulative_weights[i - 1];
            }

            for j in i..self.cumulative_weights.len() {
                self.cumulative_weights[j] -= &old_w;
                self.cumulative_weights[j] += w;
            }
        }
        self.total_weight = total_weight;
        self.weight_distribution = X::Sampler::new(zero, self.total_weight.clone());

        Ok(())
    }
}

impl<X> Distribution<usize> for WeightedIndex<X> where
    X: SampleUniform + PartialOrd {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
        use ::core::cmp::Ordering;
        let chosen_weight = self.weight_distribution.sample(rng);
        // Find the first item which has a weight *higher* than the chosen weight.
        self.cumulative_weights.binary_search_by(
            |w| if *w <= chosen_weight { Ordering::Less } else { Ordering::Greater }).unwrap_err()
    }
}

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

    #[test]
    #[cfg(not(miri))] // Miri is too slow
    fn test_weightedindex() {
        let mut r = crate::test::rng(700);
        const N_REPS: u32 = 5000;
        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
        let total_weight = weights.iter().sum::<u32>() as f32;

        let verify = |result: [i32; 14]| {
            for (i, count) in result.iter().enumerate() {
                let exp = (weights[i] * N_REPS) as f32 / total_weight;
                let mut err = (*count as f32 - exp).abs();
                if err != 0.0 {
                    err /= exp;
                }
                assert!(err <= 0.25);
            }
        };

        // WeightedIndex from vec
        let mut chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.to_vec()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from slice
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(&weights[..]).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from iterator
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.iter()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        for _ in 0..5 {
            assert_eq!(WeightedIndex::new(&[0, 1]).unwrap().sample(&mut r), 1);
            assert_eq!(WeightedIndex::new(&[1, 0]).unwrap().sample(&mut r), 0);
            assert_eq!(WeightedIndex::new(&[0, 0, 0, 0, 10, 0]).unwrap().sample(&mut r), 4);
        }

        assert_eq!(WeightedIndex::new(&[10][0..0]).unwrap_err(), WeightedError::NoItem);
        assert_eq!(WeightedIndex::new(&[0]).unwrap_err(), WeightedError::AllWeightsZero);
        assert_eq!(WeightedIndex::new(&[10, 20, -1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10, 20, 1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10]).unwrap_err(), WeightedError::InvalidWeight);
    }

    #[test]
    fn test_update_weights() {
        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
        let total_weight = weights.iter().sum::<u32>();
        let mut distr = WeightedIndex::new(weights.to_vec()).unwrap();
        assert_eq!(distr.total_weight, total_weight);

        distr.update_weights(&[(2, &4), (5, &1)]).unwrap();
        let expected_weights = [1u32, 2, 4, 0, 5, 1, 7, 1, 2, 3, 4, 5, 6, 7];
        let expected_total_weight = expected_weights.iter().sum::<u32>();
        let expected_distr = WeightedIndex::new(expected_weights.to_vec()).unwrap();
        assert_eq!(distr.total_weight, expected_total_weight);
        assert_eq!(distr.total_weight, expected_distr.total_weight);
        assert_eq!(distr.cumulative_weights, expected_distr.cumulative_weights);
    }
}

/// Error type returned from `WeightedIndex::new`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WeightedError {
    /// The provided weight collection contains no items.
    NoItem,

    /// A weight is either less than zero, greater than the supported maximum or
    /// otherwise invalid.
    InvalidWeight,

    /// All items in the provided weight collection are zero.
    AllWeightsZero,

    /// Too many weights are provided (length greater than `u32::MAX`)
    TooMany,
}

impl WeightedError {
    fn msg(&self) -> &str {
        match *self {
            WeightedError::NoItem => "No weights provided.",
            WeightedError::InvalidWeight => "A weight is invalid.",
            WeightedError::AllWeightsZero => "All weights are zero.",
            WeightedError::TooMany => "Too many weights (hit u32::MAX)",
        }
    }
}

#[cfg(feature="std")]
impl ::std::error::Error for WeightedError {
    fn description(&self) -> &str {
        self.msg()
    }
    fn cause(&self) -> Option<&dyn (::std::error::Error)> {
        None
    }
}

impl fmt::Display for WeightedError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.msg())
    }
}
<commit_msg>WeightedIndex::update_weights: More efficient implementation<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Weighted index sampling
//!
//! This module provides two implementations for sampling indices:
//!
//! *   [`WeightedIndex`] allows `O(log N)` sampling
//! *   [`alias_method::WeightedIndex`] allows `O(1)` sampling, but with
//!      much greater set-up cost
//!
//! [`alias_method::WeightedIndex`]: alias_method/struct.WeightedIndex.html

pub mod alias_method;

use crate::Rng;
use crate::distributions::Distribution;
use crate::distributions::uniform::{UniformSampler, SampleUniform, SampleBorrow};
use core::cmp::PartialOrd;
use core::fmt;

// Note that this whole module is only imported if feature="alloc" is enabled.
#[cfg(not(feature="std"))] use crate::alloc::vec::Vec;

/// A distribution using weighted sampling to pick a discretely selected
/// item.
///
/// Sampling a `WeightedIndex` distribution returns the index of a randomly
/// selected element from the iterator used when the `WeightedIndex` was
/// created. The chance of a given element being picked is proportional to the
/// value of the element. The weights can use any type `X` for which an
/// implementation of [`Uniform<X>`] exists.
///
/// # Performance
///
/// A `WeightedIndex<X>` contains a `Vec<X>` and a [`Uniform<X>`] and so its
/// size is the sum of the size of those objects, possibly plus some alignment.
///
/// Creating a `WeightedIndex<X>` will allocate enough space to hold `N - 1`
/// weights of type `X`, where `N` is the number of weights. However, since
/// `Vec` doesn't guarantee a particular growth strategy, additional memory
/// might be allocated but not used. Since the `WeightedIndex` object also
/// contains, this might cause additional allocations, though for primitive
/// types, ['Uniform<X>`] doesn't allocate any memory.
///
/// Time complexity of sampling from `WeightedIndex` is `O(log N)` where
/// `N` is the number of weights.
///
/// Sampling from `WeightedIndex` will result in a single call to
/// `Uniform<X>::sample` (method of the [`Distribution`] trait), which typically
/// will request a single value from the underlying [`RngCore`], though the
/// exact number depends on the implementaiton of `Uniform<X>::sample`.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand::distributions::WeightedIndex;
///
/// let choices = ['a', 'b', 'c'];
/// let weights = [2,   1,   1];
/// let dist = WeightedIndex::new(&weights).unwrap();
/// let mut rng = thread_rng();
/// for _ in 0..100 {
///     // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c'
///     println!("{}", choices[dist.sample(&mut rng)]);
/// }
///
/// let items = [('a', 0), ('b', 3), ('c', 7)];
/// let dist2 = WeightedIndex::new(items.iter().map(|item| item.1)).unwrap();
/// for _ in 0..100 {
///     // 0% chance to print 'a', 30% chance to print 'b', 70% chance to print 'c'
///     println!("{}", items[dist2.sample(&mut rng)].0);
/// }
/// ```
///
/// [`Uniform<X>`]: crate::distributions::uniform::Uniform
/// [`RngCore`]: crate::RngCore
#[derive(Debug, Clone)]
pub struct WeightedIndex<X: SampleUniform + PartialOrd> {
    cumulative_weights: Vec<X>,
    total_weight: X,
    weight_distribution: X::Sampler,
}

impl<X: SampleUniform + PartialOrd> WeightedIndex<X> {
    /// Creates a new a `WeightedIndex` [`Distribution`] using the values
    /// in `weights`. The weights can use any type `X` for which an
    /// implementation of [`Uniform<X>`] exists.
    ///
    /// Returns an error if the iterator is empty, if any weight is `< 0`, or
    /// if its total value is 0.
    ///
    /// [`Uniform<X>`]: crate::distributions::uniform::Uniform
    pub fn new<I>(weights: I) -> Result<WeightedIndex<X>, WeightedError>
        where I: IntoIterator,
              I::Item: SampleBorrow<X>,
              X: for<'a> ::core::ops::AddAssign<&'a X> +
                 Clone +
                 Default {
        let mut iter = weights.into_iter();
        let mut total_weight: X = iter.next()
                                      .ok_or(WeightedError::NoItem)?
                                      .borrow()
                                      .clone();

        let zero = <X as Default>::default();
        if total_weight < zero {
            return Err(WeightedError::InvalidWeight);
        }

        let mut weights = Vec::<X>::with_capacity(iter.size_hint().0);
        for w in iter {
            if *w.borrow() < zero {
                return Err(WeightedError::InvalidWeight);
            }
            weights.push(total_weight.clone());
            total_weight += w.borrow();
        }

        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }
        let distr = X::Sampler::new(zero, total_weight.clone());

        Ok(WeightedIndex { cumulative_weights: weights, total_weight, weight_distribution: distr })
    }

    /// Update a subset of weights, without changing the number of weights.
    ///
    /// `new_weights` must be sorted by the index.
    ///
    /// Using this method instead of `new` might be more efficient if only a small number of
    /// weights is modified. No allocations are performed, unless the weight type `X` uses
    /// allocation internally.
    ///
    /// In case of error, `self` is not modified.
    pub fn update_weights(&mut self, new_weights: &[(usize, &X)]) -> Result<(), WeightedError>
        where X: for<'a> ::core::ops::AddAssign<&'a X> +
                 for<'a> ::core::ops::SubAssign<&'a X> +
                 Clone +
                 Default {
        if new_weights.is_empty() {
            return Ok(());
        }

        let zero = <X as Default>::default();

        let mut total_weight = self.total_weight.clone();

        // Check for errors first, so we don't modify `self` in case something
        // goes wrong.
        let mut prev_i = None;
        for &(i, w) in new_weights {
            if let Some(old_i) = prev_i {
                if old_i >= i {
                    return Err(WeightedError::InvalidWeight);
                }
            }
            if *w < zero {
                return Err(WeightedError::InvalidWeight);
            }
            if i >= self.cumulative_weights.len() + 1 {
                return Err(WeightedError::TooMany);
            }

            // Unfortunately, we will have to calculate the non-cumulative
            // weight a second time, to avoid producing an invalid state of
            // `self`.
            let mut old_w = if i < self.cumulative_weights.len() {
                self.cumulative_weights[i].clone()
            } else {
                self.total_weight.clone()
            };
            if i > 0 {
                old_w -= &self.cumulative_weights[i - 1];
            }

            total_weight -= &old_w;
            total_weight += w;
            prev_i = Some(i);
        }
        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }

        // Update the weights. Because we checked all the preconditions in the
        // previous loop, this should never panic.
        let mut iter = new_weights.iter();
        let &(first_new_index, first_new_weight) = iter.next().unwrap();

        // `X` might be an unsigned type, so we have to be careful to avoid
        // negative numbers. This is done by tracking the sign of `change` using
        // `pos_sign`.
        let mut change: X = first_new_weight.clone();
        let mut pos_sign = true;
        if first_new_index > 0 {
            change += &self.cumulative_weights[first_new_index - 1];
        }

        let add = |x: &mut X, pos_sign: &mut bool, y: &X| {
            if !*pos_sign {
                if *x < *y {
                    let tmp = x.clone();
                    *x = y.clone();
                    *x -= &tmp;
                    *pos_sign = !*pos_sign;
                } else {
                    *x -= y;
                }
            } else {
                *x += y;
            }
        };

        let sub = |x: &mut X, pos_sign: &mut bool, y: &X| {
            if *pos_sign {
                if *x < *y {
                    let tmp = x.clone();
                    *x = y.clone();
                    *x -= &tmp;
                    *pos_sign = !*pos_sign;
                } else {
                    *x -= y;
                }
            } else {
                *x += y;
            }
        };

        {
            let first_new_cweight = if first_new_index < self.cumulative_weights.len() {
                &self.cumulative_weights[first_new_index]
            } else {
                &self.total_weight
            };
            sub(&mut change, &mut pos_sign, first_new_cweight);
        }
        let mut next_new_weight = iter.next();
        for i in first_new_index..self.cumulative_weights.len() {
            if let Some(&(j, w)) = next_new_weight {
                // j > first_new_index >= 0
                if i >= j {
                    change = w.clone();
                    pos_sign = true;
                    add(&mut change, &mut pos_sign, &self.cumulative_weights[j - 1]);
                    sub(&mut change, &mut pos_sign, &self.cumulative_weights[j]);
                    next_new_weight = iter.next();
                }
            }
            if pos_sign {
                self.cumulative_weights[i] += &change;
            } else {
                self.cumulative_weights[i] -= &change;
            }
        }

        self.total_weight = total_weight;
        self.weight_distribution = X::Sampler::new(zero, self.total_weight.clone());

        Ok(())
    }
}

impl<X> Distribution<usize> for WeightedIndex<X> where
    X: SampleUniform + PartialOrd {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
        use ::core::cmp::Ordering;
        let chosen_weight = self.weight_distribution.sample(rng);
        // Find the first item which has a weight *higher* than the chosen weight.
        self.cumulative_weights.binary_search_by(
            |w| if *w <= chosen_weight { Ordering::Less } else { Ordering::Greater }).unwrap_err()
    }
}

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

    #[test]
    #[cfg(not(miri))] // Miri is too slow
    fn test_weightedindex() {
        let mut r = crate::test::rng(700);
        const N_REPS: u32 = 5000;
        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
        let total_weight = weights.iter().sum::<u32>() as f32;

        let verify = |result: [i32; 14]| {
            for (i, count) in result.iter().enumerate() {
                let exp = (weights[i] * N_REPS) as f32 / total_weight;
                let mut err = (*count as f32 - exp).abs();
                if err != 0.0 {
                    err /= exp;
                }
                assert!(err <= 0.25);
            }
        };

        // WeightedIndex from vec
        let mut chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.to_vec()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from slice
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(&weights[..]).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from iterator
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.iter()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        for _ in 0..5 {
            assert_eq!(WeightedIndex::new(&[0, 1]).unwrap().sample(&mut r), 1);
            assert_eq!(WeightedIndex::new(&[1, 0]).unwrap().sample(&mut r), 0);
            assert_eq!(WeightedIndex::new(&[0, 0, 0, 0, 10, 0]).unwrap().sample(&mut r), 4);
        }

        assert_eq!(WeightedIndex::new(&[10][0..0]).unwrap_err(), WeightedError::NoItem);
        assert_eq!(WeightedIndex::new(&[0]).unwrap_err(), WeightedError::AllWeightsZero);
        assert_eq!(WeightedIndex::new(&[10, 20, -1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10, 20, 1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10]).unwrap_err(), WeightedError::InvalidWeight);
    }

    #[test]
    fn test_update_weights() {
        let data = [
            (&[1u32, 2, 3, 4][..],
             &[(1, &100), (2, &4)][..],  // positive change
             &[1, 100, 4, 4][..]),
            (&[1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7][..],
             &[(2, &1), (5, &1), (13, &100)][..],  // negative change and last element
             &[1u32, 2, 1, 0, 5, 1, 7, 1, 2, 3, 4, 5, 6, 100][..]),
        ];

        for (weights, update, expected_weights) in data.into_iter() {
            let total_weight = weights.iter().sum::<u32>();
            let mut distr = WeightedIndex::new(weights.to_vec()).unwrap();
            assert_eq!(distr.total_weight, total_weight);

            distr.update_weights(update).unwrap();
            let expected_total_weight = expected_weights.iter().sum::<u32>();
            let expected_distr = WeightedIndex::new(expected_weights.to_vec()).unwrap();
            assert_eq!(distr.total_weight, expected_total_weight);
            assert_eq!(distr.total_weight, expected_distr.total_weight);
            assert_eq!(distr.cumulative_weights, expected_distr.cumulative_weights);
        }
    }
}

/// Error type returned from `WeightedIndex::new`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WeightedError {
    /// The provided weight collection contains no items.
    NoItem,

    /// A weight is either less than zero, greater than the supported maximum or
    /// otherwise invalid.
    InvalidWeight,

    /// All items in the provided weight collection are zero.
    AllWeightsZero,

    /// Too many weights are provided (length greater than `u32::MAX`)
    TooMany,
}

impl WeightedError {
    fn msg(&self) -> &str {
        match *self {
            WeightedError::NoItem => "No weights provided.",
            WeightedError::InvalidWeight => "A weight is invalid.",
            WeightedError::AllWeightsZero => "All weights are zero.",
            WeightedError::TooMany => "Too many weights (hit u32::MAX)",
        }
    }
}

#[cfg(feature="std")]
impl ::std::error::Error for WeightedError {
    fn description(&self) -> &str {
        self.msg()
    }
    fn cause(&self) -> Option<&dyn (::std::error::Error)> {
        None
    }
}

impl fmt::Display for WeightedError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.msg())
    }
}
<|endoftext|>


<commit_before>use std::io::{Result, Write};

use env::Env;
use library;
use nameutil;
use super::ffi_type::*;
use traits::*;

pub fn generate_classes_funcs<W: Write>(w: &mut W, env: &Env, classes: &[&library::Class]) -> Result<()> {
    for klass in classes {
        try!(generate_object_funcs(w, env, &klass.c_type,
            &klass.glib_get_type, &klass.functions));
    }

    Ok(())
}

pub fn generate_interfaces_funcs<W: Write>(w: &mut W, env: &Env, interfaces: &[&library::Interface]) -> Result<()> {
    for interface in interfaces {
        try!(generate_object_funcs(w, env,  &interface.c_type,
            &interface.glib_get_type, &interface.functions));
    }

    Ok(())
}

fn generate_object_funcs<W: Write>(w: &mut W, env: &Env, c_type: &str,
    glib_get_type: &str, functions: &[library::Function]) -> Result<()> {
    try!(writeln!(w, ""));
    try!(writeln!(w, "    //========================================================================="));
    try!(writeln!(w, "    // {}", c_type));
    try!(writeln!(w, "    //========================================================================="));
    try!(writeln!(w, "    pub fn {:<36}() -> GType;", glib_get_type));

    for func in functions {
        let (commented, sig) = function_signature(env, func, false);
        let comment = if commented { "//" } else { "" };
        try!(writeln!(w, "    {}pub fn {:<36}{};",
                      comment, func.c_identifier.as_ref().unwrap(), sig));
    }

    Ok(())
}

pub fn generate_callbacks<W: Write>(w: &mut W, env: &Env, callbacks: &[&library::Function]) -> Result<()> {
    for func in callbacks {
        let (commented, sig) = function_signature(env, func, true);
        let comment = if commented { "//" } else { "" };
        try!(writeln!(w, "{}pub type {} = unsafe extern \"C\" fn{};",
                      comment, func.c_identifier.as_ref().unwrap(), sig));
    }

    Ok(())
}

pub fn function_signature(env: &Env, func: &library::Function, bare: bool) -> (bool, String) {
    let (mut commented, ret_str) = function_return_value(env, func);

    let mut parameter_strs: Vec<String> = Vec::new();
    for par in &func.parameters {
        let (c, par_str) = function_parameter(env, par, bare);
        parameter_strs.push(par_str);
        if c { commented = true; }
    }

    (commented, format!("({}){}", parameter_strs.connect(", "), ret_str))
}

fn function_return_value(env: &Env, func: &library::Function) -> (bool, String) {
    if func.ret.typ == Default::default() { return (false, String::new()) }
    let ffi_type = ffi_type(env, func.ret.typ, &func.ret.c_type);
    let commented = ffi_type.is_err();
    (commented, format!(" -> {}", ffi_type.as_str()))
}

fn function_parameter(env: &Env, par: &library::Parameter, bare: bool) -> (bool, String) {
    if let &library::Type::Fundamental(library::Fundamental::VarArgs) = env.library.type_(par.typ) {
        return (false, "...".into());
    }
    let ffi_type = ffi_type(env, par.typ, &par.c_type);
    let commented = ffi_type.is_err();
    let res = if bare {
        ffi_type.as_str().into()
    }
    else {
        format!("{}: {}", nameutil::mangle_keywords(&*par.name), ffi_type.as_str())
    };
    (commented, res)
}
<commit_msg>Don't add internal get_types<commit_after>use std::io::{Result, Write};

use env::Env;
use library;
use nameutil;
use super::ffi_type::*;
use traits::*;

pub fn generate_classes_funcs<W: Write>(w: &mut W, env: &Env, classes: &[&library::Class]) -> Result<()> {
    for klass in classes {
        try!(generate_object_funcs(w, env, &klass.c_type,
            &klass.glib_get_type, &klass.functions));
    }

    Ok(())
}

pub fn generate_interfaces_funcs<W: Write>(w: &mut W, env: &Env, interfaces: &[&library::Interface]) -> Result<()> {
    for interface in interfaces {
        try!(generate_object_funcs(w, env,  &interface.c_type,
            &interface.glib_get_type, &interface.functions));
    }

    Ok(())
}

fn generate_object_funcs<W: Write>(w: &mut W, env: &Env, c_type: &str,
    glib_get_type: &str, functions: &[library::Function]) -> Result<()> {
    try!(writeln!(w, ""));
    try!(writeln!(w, "    //========================================================================="));
    try!(writeln!(w, "    // {}", c_type));
    try!(writeln!(w, "    //========================================================================="));
    if glib_get_type != "intern" {
        try!(writeln!(w, "    pub fn {:<36}() -> GType;", glib_get_type));
    }

    for func in functions {
        let (commented, sig) = function_signature(env, func, false);
        let comment = if commented { "//" } else { "" };
        try!(writeln!(w, "    {}pub fn {:<36}{};",
                      comment, func.c_identifier.as_ref().unwrap(), sig));
    }

    Ok(())
}

pub fn generate_callbacks<W: Write>(w: &mut W, env: &Env, callbacks: &[&library::Function]) -> Result<()> {
    for func in callbacks {
        let (commented, sig) = function_signature(env, func, true);
        let comment = if commented { "//" } else { "" };
        try!(writeln!(w, "{}pub type {} = unsafe extern \"C\" fn{};",
                      comment, func.c_identifier.as_ref().unwrap(), sig));
    }

    Ok(())
}

pub fn function_signature(env: &Env, func: &library::Function, bare: bool) -> (bool, String) {
    let (mut commented, ret_str) = function_return_value(env, func);

    let mut parameter_strs: Vec<String> = Vec::new();
    for par in &func.parameters {
        let (c, par_str) = function_parameter(env, par, bare);
        parameter_strs.push(par_str);
        if c { commented = true; }
    }

    (commented, format!("({}){}", parameter_strs.connect(", "), ret_str))
}

fn function_return_value(env: &Env, func: &library::Function) -> (bool, String) {
    if func.ret.typ == Default::default() { return (false, String::new()) }
    let ffi_type = ffi_type(env, func.ret.typ, &func.ret.c_type);
    let commented = ffi_type.is_err();
    (commented, format!(" -> {}", ffi_type.as_str()))
}

fn function_parameter(env: &Env, par: &library::Parameter, bare: bool) -> (bool, String) {
    if let &library::Type::Fundamental(library::Fundamental::VarArgs) = env.library.type_(par.typ) {
        return (false, "...".into());
    }
    let ffi_type = ffi_type(env, par.typ, &par.c_type);
    let commented = ffi_type.is_err();
    let res = if bare {
        ffi_type.as_str().into()
    }
    else {
        format!("{}: {}", nameutil::mangle_keywords(&*par.name), ffi_type.as_str())
    };
    (commented, res)
}
<|endoftext|>


<commit_before><commit_msg>Added optimizations for mpc<commit_after><|endoftext|>


<commit_before>use core::convert::TryFrom;
use core::num::{IntErrorKind, NonZeroI32, NonZeroI8, NonZeroU32, NonZeroU8, TryFromIntError};
use core::option::Option::{self, None, Some};
use std::mem::size_of;

#[test]
fn test_create_nonzero_instance() {
    let _a = unsafe { NonZeroU32::new_unchecked(21) };
}

#[test]
fn test_size_nonzero_in_option() {
    assert_eq!(size_of::<NonZeroU32>(), size_of::<Option<NonZeroU32>>());
    assert_eq!(size_of::<NonZeroI32>(), size_of::<Option<NonZeroI32>>());
}

#[test]
fn test_match_on_nonzero_option() {
    let a = Some(unsafe { NonZeroU32::new_unchecked(42) });
    match a {
        Some(val) => assert_eq!(val.get(), 42),
        None => panic!("unexpected None while matching on Some(NonZeroU32(_))"),
    }

    match unsafe { Some(NonZeroU32::new_unchecked(43)) } {
        Some(val) => assert_eq!(val.get(), 43),
        None => panic!("unexpected None while matching on Some(NonZeroU32(_))"),
    }
}

#[test]
fn test_match_option_empty_vec() {
    let a: Option<Vec<isize>> = Some(vec![]);
    match a {
        None => panic!("unexpected None while matching on Some(vec![])"),
        _ => {}
    }
}

#[test]
fn test_match_option_vec() {
    let a = Some(vec![1, 2, 3, 4]);
    match a {
        Some(v) => assert_eq!(v, [1, 2, 3, 4]),
        None => panic!("unexpected None while matching on Some(vec![1, 2, 3, 4])"),
    }
}

#[test]
fn test_match_option_rc() {
    use std::rc::Rc;

    let five = Rc::new(5);
    match Some(five) {
        Some(r) => assert_eq!(*r, 5),
        None => panic!("unexpected None while matching on Some(Rc::new(5))"),
    }
}

#[test]
fn test_match_option_arc() {
    use std::sync::Arc;

    let five = Arc::new(5);
    match Some(five) {
        Some(a) => assert_eq!(*a, 5),
        None => panic!("unexpected None while matching on Some(Arc::new(5))"),
    }
}

#[test]
fn test_match_option_empty_string() {
    let a = Some(String::new());
    match a {
        None => panic!("unexpected None while matching on Some(String::new())"),
        _ => {}
    }
}

#[test]
fn test_match_option_string() {
    let five = "Five".to_string();
    match Some(five) {
        Some(s) => assert_eq!(s, "Five"),
        None => panic!("unexpected None while matching on Some(String { ... })"),
    }
}

mod atom {
    use core::num::NonZeroU32;

    #[derive(PartialEq, Eq)]
    pub struct Atom {
        index: NonZeroU32, // private
    }
    pub const FOO_ATOM: Atom = Atom { index: unsafe { NonZeroU32::new_unchecked(7) } };
}

macro_rules! atom {
    ("foo") => {
        atom::FOO_ATOM
    };
}

#[test]
fn test_match_nonzero_const_pattern() {
    match atom!("foo") {
        // Using as a pattern is supported by the compiler:
        atom!("foo") => {}
        _ => panic!("Expected the const item as a pattern to match."),
    }
}

#[test]
fn test_from_nonzero() {
    let nz = NonZeroU32::new(1).unwrap();
    let num: u32 = nz.into();
    assert_eq!(num, 1u32);
}

#[test]
fn test_from_signed_nonzero() {
    let nz = NonZeroI32::new(1).unwrap();
    let num: i32 = nz.into();
    assert_eq!(num, 1i32);
}

#[test]
fn test_from_str() {
    assert_eq!("123".parse::<NonZeroU8>(), Ok(NonZeroU8::new(123).unwrap()));
    assert_eq!("0".parse::<NonZeroU8>().err().map(|e| e.kind().clone()), Some(IntErrorKind::Zero));
    assert_eq!(
        "-1".parse::<NonZeroU8>().err().map(|e| e.kind().clone()),
        Some(IntErrorKind::InvalidDigit)
    );
    assert_eq!(
        "-129".parse::<NonZeroI8>().err().map(|e| e.kind().clone()),
        Some(IntErrorKind::Underflow)
    );
    assert_eq!(
        "257".parse::<NonZeroU8>().err().map(|e| e.kind().clone()),
        Some(IntErrorKind::Overflow)
    );
}

#[test]
fn test_nonzero_bitor() {
    let nz_alt = NonZeroU8::new(0b1010_1010).unwrap();
    let nz_low = NonZeroU8::new(0b0000_1111).unwrap();

    let both_nz: NonZeroU8 = nz_alt | nz_low;
    assert_eq!(both_nz.get(), 0b1010_1111);

    let rhs_int: NonZeroU8 = nz_low | 0b1100_0000u8;
    assert_eq!(rhs_int.get(), 0b1100_1111);

    let rhs_zero: NonZeroU8 = nz_alt | 0u8;
    assert_eq!(rhs_zero.get(), 0b1010_1010);

    let lhs_int: NonZeroU8 = 0b0110_0110u8 | nz_alt;
    assert_eq!(lhs_int.get(), 0b1110_1110);

    let lhs_zero: NonZeroU8 = 0u8 | nz_low;
    assert_eq!(lhs_zero.get(), 0b0000_1111);
}

#[test]
fn test_nonzero_bitor_assign() {
    let mut target = NonZeroU8::new(0b1010_1010).unwrap();

    target |= NonZeroU8::new(0b0000_1111).unwrap();
    assert_eq!(target.get(), 0b1010_1111);

    target |= 0b0001_0000;
    assert_eq!(target.get(), 0b1011_1111);

    target |= 0;
    assert_eq!(target.get(), 0b1011_1111);
}

#[test]
fn test_nonzero_from_int_on_success() {
    assert_eq!(NonZeroU8::try_from(5), Ok(NonZeroU8::new(5)));
    assert_eq!(NonZeroU32::try_from(5), Ok(NonZeroU32::new(5)));

    assert_eq!(NonZeroI8::try_from(-5), Ok(NonZeroI8::new(-5)));
    assert_eq!(NonZeroI32::try_from(-5), Ok(NonZeroI32::new(-5)));
}

#[test]
fn test_nonzero_from_int_on_err() {
    assert_eq!(NonZeroU8::try_from(0), Err(TryFromIntError(())));
    assert_eq!(NonZeroU32::try_from(0), Err(TryFromIntError(())));

    assert_eq!(NonZeroI8::try_from(0), Err(TryFromIntError(())));
    assert_eq!(NonZeroI32::try_from(0), Err(TryFromIntError(())));
}
<commit_msg>Fix test call for error case<commit_after>use core::convert::TryFrom;
use core::num::{IntErrorKind, NonZeroI32, NonZeroI8, NonZeroU32, NonZeroU8};
use core::option::Option::{self, None, Some};
use std::mem::size_of;

#[test]
fn test_create_nonzero_instance() {
    let _a = unsafe { NonZeroU32::new_unchecked(21) };
}

#[test]
fn test_size_nonzero_in_option() {
    assert_eq!(size_of::<NonZeroU32>(), size_of::<Option<NonZeroU32>>());
    assert_eq!(size_of::<NonZeroI32>(), size_of::<Option<NonZeroI32>>());
}

#[test]
fn test_match_on_nonzero_option() {
    let a = Some(unsafe { NonZeroU32::new_unchecked(42) });
    match a {
        Some(val) => assert_eq!(val.get(), 42),
        None => panic!("unexpected None while matching on Some(NonZeroU32(_))"),
    }

    match unsafe { Some(NonZeroU32::new_unchecked(43)) } {
        Some(val) => assert_eq!(val.get(), 43),
        None => panic!("unexpected None while matching on Some(NonZeroU32(_))"),
    }
}

#[test]
fn test_match_option_empty_vec() {
    let a: Option<Vec<isize>> = Some(vec![]);
    match a {
        None => panic!("unexpected None while matching on Some(vec![])"),
        _ => {}
    }
}

#[test]
fn test_match_option_vec() {
    let a = Some(vec![1, 2, 3, 4]);
    match a {
        Some(v) => assert_eq!(v, [1, 2, 3, 4]),
        None => panic!("unexpected None while matching on Some(vec![1, 2, 3, 4])"),
    }
}

#[test]
fn test_match_option_rc() {
    use std::rc::Rc;

    let five = Rc::new(5);
    match Some(five) {
        Some(r) => assert_eq!(*r, 5),
        None => panic!("unexpected None while matching on Some(Rc::new(5))"),
    }
}

#[test]
fn test_match_option_arc() {
    use std::sync::Arc;

    let five = Arc::new(5);
    match Some(five) {
        Some(a) => assert_eq!(*a, 5),
        None => panic!("unexpected None while matching on Some(Arc::new(5))"),
    }
}

#[test]
fn test_match_option_empty_string() {
    let a = Some(String::new());
    match a {
        None => panic!("unexpected None while matching on Some(String::new())"),
        _ => {}
    }
}

#[test]
fn test_match_option_string() {
    let five = "Five".to_string();
    match Some(five) {
        Some(s) => assert_eq!(s, "Five"),
        None => panic!("unexpected None while matching on Some(String { ... })"),
    }
}

mod atom {
    use core::num::NonZeroU32;

    #[derive(PartialEq, Eq)]
    pub struct Atom {
        index: NonZeroU32, // private
    }
    pub const FOO_ATOM: Atom = Atom { index: unsafe { NonZeroU32::new_unchecked(7) } };
}

macro_rules! atom {
    ("foo") => {
        atom::FOO_ATOM
    };
}

#[test]
fn test_match_nonzero_const_pattern() {
    match atom!("foo") {
        // Using as a pattern is supported by the compiler:
        atom!("foo") => {}
        _ => panic!("Expected the const item as a pattern to match."),
    }
}

#[test]
fn test_from_nonzero() {
    let nz = NonZeroU32::new(1).unwrap();
    let num: u32 = nz.into();
    assert_eq!(num, 1u32);
}

#[test]
fn test_from_signed_nonzero() {
    let nz = NonZeroI32::new(1).unwrap();
    let num: i32 = nz.into();
    assert_eq!(num, 1i32);
}

#[test]
fn test_from_str() {
    assert_eq!("123".parse::<NonZeroU8>(), Ok(NonZeroU8::new(123).unwrap()));
    assert_eq!("0".parse::<NonZeroU8>().err().map(|e| e.kind().clone()), Some(IntErrorKind::Zero));
    assert_eq!(
        "-1".parse::<NonZeroU8>().err().map(|e| e.kind().clone()),
        Some(IntErrorKind::InvalidDigit)
    );
    assert_eq!(
        "-129".parse::<NonZeroI8>().err().map(|e| e.kind().clone()),
        Some(IntErrorKind::Underflow)
    );
    assert_eq!(
        "257".parse::<NonZeroU8>().err().map(|e| e.kind().clone()),
        Some(IntErrorKind::Overflow)
    );
}

#[test]
fn test_nonzero_bitor() {
    let nz_alt = NonZeroU8::new(0b1010_1010).unwrap();
    let nz_low = NonZeroU8::new(0b0000_1111).unwrap();

    let both_nz: NonZeroU8 = nz_alt | nz_low;
    assert_eq!(both_nz.get(), 0b1010_1111);

    let rhs_int: NonZeroU8 = nz_low | 0b1100_0000u8;
    assert_eq!(rhs_int.get(), 0b1100_1111);

    let rhs_zero: NonZeroU8 = nz_alt | 0u8;
    assert_eq!(rhs_zero.get(), 0b1010_1010);

    let lhs_int: NonZeroU8 = 0b0110_0110u8 | nz_alt;
    assert_eq!(lhs_int.get(), 0b1110_1110);

    let lhs_zero: NonZeroU8 = 0u8 | nz_low;
    assert_eq!(lhs_zero.get(), 0b0000_1111);
}

#[test]
fn test_nonzero_bitor_assign() {
    let mut target = NonZeroU8::new(0b1010_1010).unwrap();

    target |= NonZeroU8::new(0b0000_1111).unwrap();
    assert_eq!(target.get(), 0b1010_1111);

    target |= 0b0001_0000;
    assert_eq!(target.get(), 0b1011_1111);

    target |= 0;
    assert_eq!(target.get(), 0b1011_1111);
}

#[test]
fn test_nonzero_from_int_on_success() {
    assert_eq!(NonZeroU8::try_from(5), Ok(NonZeroU8::new(5).unwrap()));
    assert_eq!(NonZeroU32::try_from(5), Ok(NonZeroU32::new(5).unwrap()));

    assert_eq!(NonZeroI8::try_from(-5), Ok(NonZeroI8::new(-5).unwrap()));
    assert_eq!(NonZeroI32::try_from(-5), Ok(NonZeroI32::new(-5).unwrap()));
}

#[test]
fn test_nonzero_from_int_on_err() {
    assert!(NonZeroU8::try_from(0).is_err());
    assert!(NonZeroU32::try_from(0).is_err());

    assert!(NonZeroI8::try_from(0).is_err());
    assert!(NonZeroI32::try_from(0).is_err());
}
<|endoftext|>


<commit_before>// Copyright 2016 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.

//! Script to check the validity of `href` links in our HTML documentation.
//!
//! In the past we've been quite error prone to writing in broken links as most
//! of them are manually rather than automatically added. As files move over
//! time or apis change old links become stale or broken. The purpose of this
//! script is to check all relative links in our documentation to make sure they
//! actually point to a valid place.
//!
//! Currently this doesn't actually do any HTML parsing or anything fancy like
//! that, it just has a simple "regex" to search for `href` and `id` tags.
//! These values are then translated to file URLs if possible and then the
//! destination is asserted to exist.
//!
//! A few whitelisted exceptions are allowed as there's known bugs in rustdoc,
//! but this should catch the majority of "broken link" cases.

extern crate url;

use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Entry;

use url::{Url, UrlParser};

use Redirect::*;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {:?}", stringify!($e), e),
    })
}

fn main() {
    let docs = env::args().nth(1).unwrap();
    let docs = env::current_dir().unwrap().join(docs);
    let mut url = Url::from_file_path(&docs).unwrap();
    let mut errors = false;
    walk(&mut HashMap::new(), &docs, &docs, &mut url, &mut errors);
    if errors {
        panic!("found some broken links");
    }
}

#[derive(Debug)]
pub enum LoadError {
    IOError(std::io::Error),
    BrokenRedirect(PathBuf, std::io::Error),
    IsRedirect,
}

enum Redirect {
    SkipRedirect,
    FromRedirect(bool),
}

struct FileEntry {
    source: String,
    ids: HashSet<String>,
}

type Cache = HashMap<PathBuf, FileEntry>;

impl FileEntry {
    fn parse_ids(&mut self, file: &Path, contents: &str, errors: &mut bool) {
        if self.ids.is_empty() {
            with_attrs_in_source(contents, " id", |fragment, i| {
                let frag = fragment.trim_left_matches("#").to_owned();
                if !self.ids.insert(frag) {
                    *errors = true;
                    println!("{}:{}: id is not unique: `{}`", file.display(), i, fragment);
                }
            });
        }
    }
}

fn walk(cache: &mut Cache, root: &Path, dir: &Path, url: &mut Url, errors: &mut bool) {
    for entry in t!(dir.read_dir()).map(|e| t!(e)) {
        let path = entry.path();
        let kind = t!(entry.file_type());
        url.path_mut().unwrap().push(entry.file_name().into_string().unwrap());
        if kind.is_dir() {
            walk(cache, root, &path, url, errors);
        } else {
            let pretty_path = check(cache, root, &path, url, errors);
            if let Some(pretty_path) = pretty_path {
                let entry = cache.get_mut(&pretty_path).unwrap();
                // we don't need the source anymore,
                // so drop to reduce memory-usage
                entry.source = String::new();
            }
        }
        url.path_mut().unwrap().pop();
    }
}

fn check(cache: &mut Cache,
         root: &Path,
         file: &Path,
         base: &Url,
         errors: &mut bool)
         -> Option<PathBuf> {
    // ignore js files as they are not prone to errors as the rest of the
    // documentation is and they otherwise bring up false positives.
    if file.extension().and_then(|s| s.to_str()) == Some("js") {
        return None;
    }

    // Unfortunately we're not 100% full of valid links today to we need a few
    // whitelists to get this past `make check` today.
    // FIXME(#32129)
    if file.ends_with("std/string/struct.String.html") {
        return None;
    }
    // FIXME(#32553)
    if file.ends_with("collections/string/struct.String.html") {
        return None;
    }
    // FIXME(#32130)
    if file.ends_with("btree_set/struct.BTreeSet.html") ||
       file.ends_with("collections/struct.BTreeSet.html") ||
       file.ends_with("collections/btree_map/struct.BTreeMap.html") ||
       file.ends_with("collections/hash_map/struct.HashMap.html") {
        return None;
    }

    if file.ends_with("std/sys/ext/index.html") {
        return None;
    }

    if let Some(file) = file.to_str() {
        // FIXME(#31948)
        if file.contains("ParseFloatError") {
            return None;
        }
        // weird reexports, but this module is on its way out, so chalk it up to
        // "rustdoc weirdness" and move on from there
        if file.contains("scoped_tls") {
            return None;
        }
    }

    let mut parser = UrlParser::new();
    parser.base_url(base);

    let res = load_file(cache, root, PathBuf::from(file), SkipRedirect);
    let (pretty_file, contents) = match res {
        Ok(res) => res,
        Err(_) => return None,
    };
    {
        cache.get_mut(&pretty_file)
             .unwrap()
             .parse_ids(&pretty_file, &contents, errors);
    }

    // Search for anything that's the regex 'href[ ]*=[ ]*".*?"'
    with_attrs_in_source(&contents, " href", |url, i| {
        // Once we've plucked out the URL, parse it using our base url and
        // then try to extract a file path. If either of these fail then we
        // just keep going.
        let (parsed_url, path) = match url_to_file_path(&parser, url) {
            Some((url, path)) => (url, PathBuf::from(path)),
            None => return,
        };

        // Alright, if we've found a file name then this file had better
        // exist! If it doesn't then we register and print an error.
        if path.exists() {
            if path.is_dir() {
                // Links to directories show as directory listings when viewing
                // the docs offline so it's best to avoid them.
                *errors = true;
                let pretty_path = path.strip_prefix(root).unwrap_or(&path);
                println!("{}:{}: directory link - {}",
                         pretty_file.display(),
                         i + 1,
                         pretty_path.display());
                return;
            }
            let res = load_file(cache, root, path.clone(), FromRedirect(false));
            let (pretty_path, contents) = match res {
                Ok(res) => res,
                Err(LoadError::IOError(err)) => panic!(format!("{}", err)),
                Err(LoadError::BrokenRedirect(target, _)) => {
                    print!("{}:{}: broken redirect to {}",
                           pretty_file.display(),
                           i + 1,
                           target.display());
                    return;
                }
                Err(LoadError::IsRedirect) => unreachable!(),
            };

            if let Some(ref fragment) = parsed_url.fragment {
                // Fragments like `#1-6` are most likely line numbers to be
                // interpreted by javascript, so we're ignoring these
                if fragment.splitn(2, '-')
                           .all(|f| f.chars().all(|c| c.is_numeric())) {
                    return;
                }

                let entry = &mut cache.get_mut(&pretty_path).unwrap();
                entry.parse_ids(&pretty_path, &contents, errors);

                if !entry.ids.contains(fragment) {
                    *errors = true;
                    print!("{}:{}: broken link fragment  ",
                           pretty_file.display(),
                           i + 1);
                    println!("`#{}` pointing to `{}`", fragment, pretty_path.display());
                };
            }
        } else {
            *errors = true;
            print!("{}:{}: broken link - ", pretty_file.display(), i + 1);
            let pretty_path = path.strip_prefix(root).unwrap_or(&path);
            println!("{}", pretty_path.display());
        }
    });
    Some(pretty_file)
}

fn load_file(cache: &mut Cache,
             root: &Path,
             file: PathBuf,
             redirect: Redirect)
             -> Result<(PathBuf, String), LoadError> {
    let mut contents = String::new();
    let pretty_file = PathBuf::from(file.strip_prefix(root).unwrap_or(&file));

    let maybe_redirect = match cache.entry(pretty_file.clone()) {
        Entry::Occupied(entry) => {
            contents = entry.get().source.clone();
            None
        }
        Entry::Vacant(entry) => {
            let mut fp = try!(File::open(file.clone()).map_err(|err| {
                if let FromRedirect(true) = redirect {
                    LoadError::BrokenRedirect(file.clone(), err)
                } else {
                    LoadError::IOError(err)
                }
            }));
            try!(fp.read_to_string(&mut contents)
                   .map_err(|err| LoadError::IOError(err)));

            let maybe = maybe_redirect(&contents);
            if maybe.is_some() {
                if let SkipRedirect = redirect {
                    return Err(LoadError::IsRedirect);
                }
            } else {
                entry.insert(FileEntry {
                    source: contents.clone(),
                    ids: HashSet::new(),
                });
            }
            maybe
        }
    };
    let base = Url::from_file_path(&file).unwrap();
    let mut parser = UrlParser::new();
    parser.base_url(&base);

    match maybe_redirect.and_then(|url| url_to_file_path(&parser, &url)) {
        Some((_, redirect_file)) => {
            let path = PathBuf::from(redirect_file);
            load_file(cache, root, path, FromRedirect(true))
        }
        None => Ok((pretty_file, contents)),
    }
}

fn maybe_redirect(source: &str) -> Option<String> {
    const REDIRECT: &'static str = "
Redirecting to  l,
        None => return None,
    };

    redirect_line.find(REDIRECT).map(|i| {
        let rest = &redirect_line[(i + REDIRECT.len() + 1)..];
        let pos_quote = rest.find('"').unwrap();
        rest[..pos_quote].to_owned()
    })
}

fn url_to_file_path(parser: &UrlParser, url: &str) -> Option<(Url, PathBuf)> {
    parser.parse(url)
          .ok()
          .and_then(|parsed_url| parsed_url.to_file_path().ok().map(|f| (parsed_url, f)))
}

fn with_attrs_in_source<F: FnMut(&str, usize)>(contents: &str, attr: &str, mut f: F) {
    for (i, mut line) in contents.lines().enumerate() {
        while let Some(j) = line.find(attr) {
            let rest = &line[j + attr.len()..];
            line = rest;
            let pos_equals = match rest.find("=") {
                Some(i) => i,
                None => continue,
            };
            if rest[..pos_equals].trim_left_matches(" ") != "" {
                continue;
            }

            let rest = &rest[pos_equals + 1..];

            let pos_quote = match rest.find(&['"', '\''][..]) {
                Some(i) => i,
                None => continue,
            };
            let quote_delim = rest.as_bytes()[pos_quote] as char;

            if rest[..pos_quote].trim_left_matches(" ") != "" {
                continue;
            }
            let rest = &rest[pos_quote + 1..];
            let url = match rest.find(quote_delim) {
                Some(i) => &rest[..i],
                None => continue,
            };
            f(url, i)
        }
    }
}
<commit_msg>Rollup merge of #34566 - ollie27:linkchecker_invalid_urls, r=alexcrichton<commit_after>// Copyright 2016 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.

//! Script to check the validity of `href` links in our HTML documentation.
//!
//! In the past we've been quite error prone to writing in broken links as most
//! of them are manually rather than automatically added. As files move over
//! time or apis change old links become stale or broken. The purpose of this
//! script is to check all relative links in our documentation to make sure they
//! actually point to a valid place.
//!
//! Currently this doesn't actually do any HTML parsing or anything fancy like
//! that, it just has a simple "regex" to search for `href` and `id` tags.
//! These values are then translated to file URLs if possible and then the
//! destination is asserted to exist.
//!
//! A few whitelisted exceptions are allowed as there's known bugs in rustdoc,
//! but this should catch the majority of "broken link" cases.

extern crate url;

use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Entry;

use url::{Url, UrlParser};

use Redirect::*;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {:?}", stringify!($e), e),
    })
}

fn main() {
    let docs = env::args().nth(1).unwrap();
    let docs = env::current_dir().unwrap().join(docs);
    let mut url = Url::from_file_path(&docs).unwrap();
    let mut errors = false;
    walk(&mut HashMap::new(), &docs, &docs, &mut url, &mut errors);
    if errors {
        panic!("found some broken links");
    }
}

#[derive(Debug)]
pub enum LoadError {
    IOError(std::io::Error),
    BrokenRedirect(PathBuf, std::io::Error),
    IsRedirect,
}

enum Redirect {
    SkipRedirect,
    FromRedirect(bool),
}

struct FileEntry {
    source: String,
    ids: HashSet<String>,
}

type Cache = HashMap<PathBuf, FileEntry>;

impl FileEntry {
    fn parse_ids(&mut self, file: &Path, contents: &str, errors: &mut bool) {
        if self.ids.is_empty() {
            with_attrs_in_source(contents, " id", |fragment, i| {
                let frag = fragment.trim_left_matches("#").to_owned();
                if !self.ids.insert(frag) {
                    *errors = true;
                    println!("{}:{}: id is not unique: `{}`", file.display(), i, fragment);
                }
            });
        }
    }
}

fn walk(cache: &mut Cache, root: &Path, dir: &Path, url: &mut Url, errors: &mut bool) {
    for entry in t!(dir.read_dir()).map(|e| t!(e)) {
        let path = entry.path();
        let kind = t!(entry.file_type());
        url.path_mut().unwrap().push(entry.file_name().into_string().unwrap());
        if kind.is_dir() {
            walk(cache, root, &path, url, errors);
        } else {
            let pretty_path = check(cache, root, &path, url, errors);
            if let Some(pretty_path) = pretty_path {
                let entry = cache.get_mut(&pretty_path).unwrap();
                // we don't need the source anymore,
                // so drop to reduce memory-usage
                entry.source = String::new();
            }
        }
        url.path_mut().unwrap().pop();
    }
}

fn check(cache: &mut Cache,
         root: &Path,
         file: &Path,
         base: &Url,
         errors: &mut bool)
         -> Option<PathBuf> {
    // ignore js files as they are not prone to errors as the rest of the
    // documentation is and they otherwise bring up false positives.
    if file.extension().and_then(|s| s.to_str()) == Some("js") {
        return None;
    }

    // Unfortunately we're not 100% full of valid links today to we need a few
    // whitelists to get this past `make check` today.
    // FIXME(#32129)
    if file.ends_with("std/string/struct.String.html") {
        return None;
    }
    // FIXME(#32553)
    if file.ends_with("collections/string/struct.String.html") {
        return None;
    }
    // FIXME(#32130)
    if file.ends_with("btree_set/struct.BTreeSet.html") ||
       file.ends_with("collections/struct.BTreeSet.html") ||
       file.ends_with("collections/btree_map/struct.BTreeMap.html") ||
       file.ends_with("collections/hash_map/struct.HashMap.html") {
        return None;
    }

    let mut parser = UrlParser::new();
    parser.base_url(base);

    let res = load_file(cache, root, PathBuf::from(file), SkipRedirect);
    let (pretty_file, contents) = match res {
        Ok(res) => res,
        Err(_) => return None,
    };
    {
        cache.get_mut(&pretty_file)
             .unwrap()
             .parse_ids(&pretty_file, &contents, errors);
    }

    // Search for anything that's the regex 'href[ ]*=[ ]*".*?"'
    with_attrs_in_source(&contents, " href", |url, i| {
        // Ignore external URLs
        if url.starts_with("http:") || url.starts_with("https:") ||
           url.starts_with("javascript:") || url.starts_with("ftp:") ||
           url.starts_with("irc:") || url.starts_with("data:") {
            return;
        }
        // Once we've plucked out the URL, parse it using our base url and
        // then try to extract a file path.
        let (parsed_url, path) = match url_to_file_path(&parser, url) {
            Some((url, path)) => (url, PathBuf::from(path)),
            None => {
                *errors = true;
                println!("{}:{}: invalid link - {}",
                         pretty_file.display(),
                         i + 1,
                         url);
                return;
            }
        };

        // Alright, if we've found a file name then this file had better
        // exist! If it doesn't then we register and print an error.
        if path.exists() {
            if path.is_dir() {
                // Links to directories show as directory listings when viewing
                // the docs offline so it's best to avoid them.
                *errors = true;
                let pretty_path = path.strip_prefix(root).unwrap_or(&path);
                println!("{}:{}: directory link - {}",
                         pretty_file.display(),
                         i + 1,
                         pretty_path.display());
                return;
            }
            let res = load_file(cache, root, path.clone(), FromRedirect(false));
            let (pretty_path, contents) = match res {
                Ok(res) => res,
                Err(LoadError::IOError(err)) => panic!(format!("{}", err)),
                Err(LoadError::BrokenRedirect(target, _)) => {
                    *errors = true;
                    println!("{}:{}: broken redirect to {}",
                             pretty_file.display(),
                             i + 1,
                             target.display());
                    return;
                }
                Err(LoadError::IsRedirect) => unreachable!(),
            };

            if let Some(ref fragment) = parsed_url.fragment {
                // Fragments like `#1-6` are most likely line numbers to be
                // interpreted by javascript, so we're ignoring these
                if fragment.splitn(2, '-')
                           .all(|f| f.chars().all(|c| c.is_numeric())) {
                    return;
                }

                let entry = &mut cache.get_mut(&pretty_path).unwrap();
                entry.parse_ids(&pretty_path, &contents, errors);

                if !entry.ids.contains(fragment) {
                    *errors = true;
                    print!("{}:{}: broken link fragment  ",
                           pretty_file.display(),
                           i + 1);
                    println!("`#{}` pointing to `{}`", fragment, pretty_path.display());
                };
            }
        } else {
            *errors = true;
            print!("{}:{}: broken link - ", pretty_file.display(), i + 1);
            let pretty_path = path.strip_prefix(root).unwrap_or(&path);
            println!("{}", pretty_path.display());
        }
    });
    Some(pretty_file)
}

fn load_file(cache: &mut Cache,
             root: &Path,
             file: PathBuf,
             redirect: Redirect)
             -> Result<(PathBuf, String), LoadError> {
    let mut contents = String::new();
    let pretty_file = PathBuf::from(file.strip_prefix(root).unwrap_or(&file));

    let maybe_redirect = match cache.entry(pretty_file.clone()) {
        Entry::Occupied(entry) => {
            contents = entry.get().source.clone();
            None
        }
        Entry::Vacant(entry) => {
            let mut fp = try!(File::open(file.clone()).map_err(|err| {
                if let FromRedirect(true) = redirect {
                    LoadError::BrokenRedirect(file.clone(), err)
                } else {
                    LoadError::IOError(err)
                }
            }));
            try!(fp.read_to_string(&mut contents)
                   .map_err(|err| LoadError::IOError(err)));

            let maybe = maybe_redirect(&contents);
            if maybe.is_some() {
                if let SkipRedirect = redirect {
                    return Err(LoadError::IsRedirect);
                }
            } else {
                entry.insert(FileEntry {
                    source: contents.clone(),
                    ids: HashSet::new(),
                });
            }
            maybe
        }
    };
    let base = Url::from_file_path(&file).unwrap();
    let mut parser = UrlParser::new();
    parser.base_url(&base);

    match maybe_redirect.and_then(|url| url_to_file_path(&parser, &url)) {
        Some((_, redirect_file)) => {
            let path = PathBuf::from(redirect_file);
            load_file(cache, root, path, FromRedirect(true))
        }
        None => Ok((pretty_file, contents)),
    }
}

fn maybe_redirect(source: &str) -> Option<String> {
    const REDIRECT: &'static str = "
Redirecting to  l,
        None => return None,
    };

    redirect_line.find(REDIRECT).map(|i| {
        let rest = &redirect_line[(i + REDIRECT.len() + 1)..];
        let pos_quote = rest.find('"').unwrap();
        rest[..pos_quote].to_owned()
    })
}

fn url_to_file_path(parser: &UrlParser, url: &str) -> Option<(Url, PathBuf)> {
    parser.parse(url)
          .ok()
          .and_then(|parsed_url| parsed_url.to_file_path().ok().map(|f| (parsed_url, f)))
}

fn with_attrs_in_source<F: FnMut(&str, usize)>(contents: &str, attr: &str, mut f: F) {
    for (i, mut line) in contents.lines().enumerate() {
        while let Some(j) = line.find(attr) {
            let rest = &line[j + attr.len()..];
            line = rest;
            let pos_equals = match rest.find("=") {
                Some(i) => i,
                None => continue,
            };
            if rest[..pos_equals].trim_left_matches(" ") != "" {
                continue;
            }

            let rest = &rest[pos_equals + 1..];

            let pos_quote = match rest.find(&['"', '\''][..]) {
                Some(i) => i,
                None => continue,
            };
            let quote_delim = rest.as_bytes()[pos_quote] as char;

            if rest[..pos_quote].trim_left_matches(" ") != "" {
                continue;
            }
            let rest = &rest[pos_quote + 1..];
            let url = match rest.find(quote_delim) {
                Some(i) => &rest[..i],
                None => continue,
            };
            f(url, i)
        }
    }
}
<|endoftext|>


<commit_before>//! The sendmail transport sends the email using the local sendmail command.
//!
//! ## Sync example
//!
//! ```rust
//! # use std::error::Error;
//!
//! # #[cfg(all(feature = "sendmail-transport", feature = "builder"))]
//! # fn main() -> Result<(), Box<dyn Error>> {
//! # use lettre::{Message, Transport, SendmailTransport};
//!
//! let email = Message::builder()
//!     .from("NoBody <nobody@domain.tld>".parse()?)
//!     .reply_to("Yuin <yuin@domain.tld>".parse()?)
//!     .to("Hei <hei@domain.tld>".parse()?)
//!     .subject("Happy new year")
//!     .body("Be happy!")?;
//!
//! let sender = SendmailTransport::new();
//! let result = sender.send(&email);
//! assert!(result.is_ok());
//! # Ok(())
//! # }
//!
//! # #[cfg(not(all(feature = "sendmail-transport", feature = "builder")))]
//! # fn main() {}
//! ```
//!
//! ## Async tokio 0.2 example
//!
//! ```rust,no_run
//! # use std::error::Error;
//!
//! # #[cfg(all(feature = "tokio02", feature = "sendmail-transport", feature = "builder"))]
//! # async fn run() -> Result<(), Box<dyn Error>> {
//! use lettre::{Message, Tokio02Transport, SendmailTransport};
//!
//! let email = Message::builder()
//!     .from("NoBody <nobody@domain.tld>".parse()?)
//!     .reply_to("Yuin <yuin@domain.tld>".parse()?)
//!     .to("Hei <hei@domain.tld>".parse()?)
//!     .subject("Happy new year")
//!     .body("Be happy!")?;
//!
//! let sender = SendmailTransport::new();
//! let result = sender.send(email).await;
//! assert!(result.is_ok());
//! # Ok(())
//! # }
//! ```
//!
//! ## Async async-std 1.x example
//!
//!```rust,no_run
//! # use std::error::Error;
//!
//! # #[cfg(all(feature = "async-std1", feature = "sendmail-transport", feature = "builder"))]
//! # async fn run() -> Result<(), Box<dyn Error>> {
//! use lettre::{Message, AsyncStd1Transport, SendmailTransport};
//!
//! let email = Message::builder()
//!     .from("NoBody <nobody@domain.tld>".parse()?)
//!     .reply_to("Yuin <yuin@domain.tld>".parse()?)
//!     .to("Hei <hei@domain.tld>".parse()?)
//!     .subject("Happy new year")
//!     .body("Be happy!")?;
//!
//! let sender = SendmailTransport::new();
//! let result = sender.send(email).await;
//! assert!(result.is_ok());
//! # Ok(())
//! # }
//! ```

pub use self::error::Error;
use crate::address::Envelope;
#[cfg(feature = "async-std1")]
use crate::AsyncStd1Transport;
#[cfg(feature = "tokio02")]
use crate::Tokio02Transport;
#[cfg(feature = "tokio03")]
use crate::Tokio03Transport;
use crate::Transport;
#[cfg(any(feature = "async-std1", feature = "tokio02", feature = "tokio03"))]
use async_trait::async_trait;
use std::{
    ffi::OsString,
    io::prelude::*,
    process::{Command, Stdio},
};

mod error;

const DEFAUT_SENDMAIL: &str = "/usr/sbin/sendmail";

/// Sends an email using the `sendmail` command
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct SendmailTransport {
    command: OsString,
}

impl SendmailTransport {
    /// Creates a new transport with the default `/usr/sbin/sendmail` command
    pub fn new() -> SendmailTransport {
        SendmailTransport {
            command: DEFAUT_SENDMAIL.into(),
        }
    }

    /// Creates a new transport to the given sendmail command
    pub fn new_with_command<S: Into<OsString>>(command: S) -> SendmailTransport {
        SendmailTransport {
            command: command.into(),
        }
    }

    fn command(&self, envelope: &Envelope) -> Command {
        let mut c = Command::new(&self.command);
        c.arg("-i")
            .arg("-f")
            .arg(envelope.from().map(|f| f.as_ref()).unwrap_or("\"\""))
            .args(envelope.to())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped());
        c
    }

    #[cfg(feature = "tokio02")]
    fn tokio02_command(&self, envelope: &Envelope) -> tokio02_crate::process::Command {
        use tokio02_crate::process::Command;

        let mut c = Command::new(&self.command);
        c.kill_on_drop(true);
        c.arg("-i")
            .arg("-f")
            .arg(envelope.from().map(|f| f.as_ref()).unwrap_or("\"\""))
            .args(envelope.to())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped());
        c
    }

    #[cfg(feature = "tokio03")]
    fn tokio03_command(&self, envelope: &Envelope) -> tokio03_crate::process::Command {
        use tokio03_crate::process::Command;

        let mut c = Command::new(&self.command);
        c.kill_on_drop(true);
        c.arg("-i")
            .arg("-f")
            .arg(envelope.from().map(|f| f.as_ref()).unwrap_or("\"\""))
            .args(envelope.to())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped());
        c
    }
}

impl Default for SendmailTransport {
    fn default() -> Self {
        Self::new()
    }
}

impl Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        // Spawn the sendmail command
        let mut process = self.command(envelope).spawn()?;

        process.stdin.as_mut().unwrap().write_all(email)?;
        let output = process.wait_with_output()?;

        if output.status.success() {
            Ok(())
        } else {
            Err(error::Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}

#[cfg(feature = "async-std1")]
#[async_trait]
impl AsyncStd1Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    async fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        let mut command = self.command(envelope);
        let email = email.to_vec();

        // TODO: Convert to real async, once async-std has a process implementation.
        let output = async_std::task::spawn_blocking(move || {
            // Spawn the sendmail command
            let mut process = command.spawn()?;

            process.stdin.as_mut().unwrap().write_all(&email)?;
            process.wait_with_output()
        })
        .await?;

        if output.status.success() {
            Ok(())
        } else {
            Err(Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}

#[cfg(feature = "tokio02")]
#[async_trait]
impl Tokio02Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    async fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        use tokio02_crate::io::AsyncWriteExt;

        let mut command = self.tokio02_command(envelope);

        // Spawn the sendmail command
        let mut process = command.spawn()?;

        process.stdin.as_mut().unwrap().write_all(&email).await?;
        let output = process.wait_with_output().await?;

        if output.status.success() {
            Ok(())
        } else {
            Err(Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}

#[cfg(feature = "tokio03")]
#[async_trait]
impl Tokio03Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    async fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        use tokio03_crate::io::AsyncWriteExt;

        let mut command = self.tokio03_command(envelope);

        // Spawn the sendmail command
        let mut process = command.spawn()?;

        process.stdin.as_mut().unwrap().write_all(&email).await?;
        let output = process.wait_with_output().await?;

        if output.status.success() {
            Ok(())
        } else {
            Err(Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}
<commit_msg>fix(transport-sendmail): Stop argument parsing before destination addresses<commit_after>//! The sendmail transport sends the email using the local sendmail command.
//!
//! ## Sync example
//!
//! ```rust
//! # use std::error::Error;
//!
//! # #[cfg(all(feature = "sendmail-transport", feature = "builder"))]
//! # fn main() -> Result<(), Box<dyn Error>> {
//! # use lettre::{Message, Transport, SendmailTransport};
//!
//! let email = Message::builder()
//!     .from("NoBody <nobody@domain.tld>".parse()?)
//!     .reply_to("Yuin <yuin@domain.tld>".parse()?)
//!     .to("Hei <hei@domain.tld>".parse()?)
//!     .subject("Happy new year")
//!     .body("Be happy!")?;
//!
//! let sender = SendmailTransport::new();
//! let result = sender.send(&email);
//! assert!(result.is_ok());
//! # Ok(())
//! # }
//!
//! # #[cfg(not(all(feature = "sendmail-transport", feature = "builder")))]
//! # fn main() {}
//! ```
//!
//! ## Async tokio 0.2 example
//!
//! ```rust,no_run
//! # use std::error::Error;
//!
//! # #[cfg(all(feature = "tokio02", feature = "sendmail-transport", feature = "builder"))]
//! # async fn run() -> Result<(), Box<dyn Error>> {
//! use lettre::{Message, Tokio02Transport, SendmailTransport};
//!
//! let email = Message::builder()
//!     .from("NoBody <nobody@domain.tld>".parse()?)
//!     .reply_to("Yuin <yuin@domain.tld>".parse()?)
//!     .to("Hei <hei@domain.tld>".parse()?)
//!     .subject("Happy new year")
//!     .body("Be happy!")?;
//!
//! let sender = SendmailTransport::new();
//! let result = sender.send(email).await;
//! assert!(result.is_ok());
//! # Ok(())
//! # }
//! ```
//!
//! ## Async async-std 1.x example
//!
//!```rust,no_run
//! # use std::error::Error;
//!
//! # #[cfg(all(feature = "async-std1", feature = "sendmail-transport", feature = "builder"))]
//! # async fn run() -> Result<(), Box<dyn Error>> {
//! use lettre::{Message, AsyncStd1Transport, SendmailTransport};
//!
//! let email = Message::builder()
//!     .from("NoBody <nobody@domain.tld>".parse()?)
//!     .reply_to("Yuin <yuin@domain.tld>".parse()?)
//!     .to("Hei <hei@domain.tld>".parse()?)
//!     .subject("Happy new year")
//!     .body("Be happy!")?;
//!
//! let sender = SendmailTransport::new();
//! let result = sender.send(email).await;
//! assert!(result.is_ok());
//! # Ok(())
//! # }
//! ```

pub use self::error::Error;
use crate::address::Envelope;
#[cfg(feature = "async-std1")]
use crate::AsyncStd1Transport;
#[cfg(feature = "tokio02")]
use crate::Tokio02Transport;
#[cfg(feature = "tokio03")]
use crate::Tokio03Transport;
use crate::Transport;
#[cfg(any(feature = "async-std1", feature = "tokio02", feature = "tokio03"))]
use async_trait::async_trait;
use std::{
    ffi::OsString,
    io::prelude::*,
    process::{Command, Stdio},
};

mod error;

const DEFAUT_SENDMAIL: &str = "/usr/sbin/sendmail";

/// Sends an email using the `sendmail` command
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct SendmailTransport {
    command: OsString,
}

impl SendmailTransport {
    /// Creates a new transport with the default `/usr/sbin/sendmail` command
    pub fn new() -> SendmailTransport {
        SendmailTransport {
            command: DEFAUT_SENDMAIL.into(),
        }
    }

    /// Creates a new transport to the given sendmail command
    pub fn new_with_command<S: Into<OsString>>(command: S) -> SendmailTransport {
        SendmailTransport {
            command: command.into(),
        }
    }

    fn command(&self, envelope: &Envelope) -> Command {
        let mut c = Command::new(&self.command);
        c.arg("-i")
            .arg("-f")
            .arg(envelope.from().map(|f| f.as_ref()).unwrap_or("\"\""))
            .arg("--")
            .args(envelope.to())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped());
        c
    }

    #[cfg(feature = "tokio02")]
    fn tokio02_command(&self, envelope: &Envelope) -> tokio02_crate::process::Command {
        use tokio02_crate::process::Command;

        let mut c = Command::new(&self.command);
        c.kill_on_drop(true);
        c.arg("-i")
            .arg("-f")
            .arg(envelope.from().map(|f| f.as_ref()).unwrap_or("\"\""))
            .arg("--")
            .args(envelope.to())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped());
        c
    }

    #[cfg(feature = "tokio03")]
    fn tokio03_command(&self, envelope: &Envelope) -> tokio03_crate::process::Command {
        use tokio03_crate::process::Command;

        let mut c = Command::new(&self.command);
        c.kill_on_drop(true);
        c.arg("-i")
            .arg("-f")
            .arg(envelope.from().map(|f| f.as_ref()).unwrap_or("\"\""))
            .arg("--")
            .args(envelope.to())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped());
        c
    }
}

impl Default for SendmailTransport {
    fn default() -> Self {
        Self::new()
    }
}

impl Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        // Spawn the sendmail command
        let mut process = self.command(envelope).spawn()?;

        process.stdin.as_mut().unwrap().write_all(email)?;
        let output = process.wait_with_output()?;

        if output.status.success() {
            Ok(())
        } else {
            Err(error::Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}

#[cfg(feature = "async-std1")]
#[async_trait]
impl AsyncStd1Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    async fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        let mut command = self.command(envelope);
        let email = email.to_vec();

        // TODO: Convert to real async, once async-std has a process implementation.
        let output = async_std::task::spawn_blocking(move || {
            // Spawn the sendmail command
            let mut process = command.spawn()?;

            process.stdin.as_mut().unwrap().write_all(&email)?;
            process.wait_with_output()
        })
        .await?;

        if output.status.success() {
            Ok(())
        } else {
            Err(Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}

#[cfg(feature = "tokio02")]
#[async_trait]
impl Tokio02Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    async fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        use tokio02_crate::io::AsyncWriteExt;

        let mut command = self.tokio02_command(envelope);

        // Spawn the sendmail command
        let mut process = command.spawn()?;

        process.stdin.as_mut().unwrap().write_all(&email).await?;
        let output = process.wait_with_output().await?;

        if output.status.success() {
            Ok(())
        } else {
            Err(Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}

#[cfg(feature = "tokio03")]
#[async_trait]
impl Tokio03Transport for SendmailTransport {
    type Ok = ();
    type Error = Error;

    async fn send_raw(&self, envelope: &Envelope, email: &[u8]) -> Result<Self::Ok, Self::Error> {
        use tokio03_crate::io::AsyncWriteExt;

        let mut command = self.tokio03_command(envelope);

        // Spawn the sendmail command
        let mut process = command.spawn()?;

        process.stdin.as_mut().unwrap().write_all(&email).await?;
        let output = process.wait_with_output().await?;

        if output.status.success() {
            Ok(())
        } else {
            Err(Error::Client(String::from_utf8(output.stderr)?))
        }
    }
}
<|endoftext|>