<commit_before>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use statistics::*;
use distribution::{Univariate, Discrete, Distribution};
use {Result, StatsError};

/// Implements the [Categorical](https://en.wikipedia.org/wiki/Categorical_distribution)
/// distribution, also known as the generalized Bernoulli or discrete distribution
///
/// # Examples
///
/// ```
/// use statrs::distribution::{Categorical, Discrete};
/// use statrs::statistics::Mode;
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Categorical {
    norm_pmf: Vec<f64>,
    cdf: Vec<f64>,
}

impl Categorical {
    pub fn new(prob_mass: &[f64]) -> Result<Categorical> {
        if !is_valid_prob_mass(prob_mass) {
            Err(StatsError::BadParams)
        } else {
            // extract un-normalized cdf
            let mut cdf = vec![0.0; prob_mass.len()];
            cdf[0] = prob_mass[0];
            for i in 1..prob_mass.len() {
                unsafe {
                    let val = cdf.get_unchecked(i - 1) + prob_mass.get_unchecked(i);
                    let elem = cdf.get_unchecked_mut(i);
                    *elem = val;
                }
            }

            // extract normalized probability mass
            let sum = cdf[cdf.len() - 1];
            let mut norm_pmf = vec![0.0; prob_mass.len()];
            for i in 0..prob_mass.len() {
                unsafe {
                    let elem = norm_pmf.get_unchecked_mut(i);
                    *elem = prob_mass.get_unchecked(i) / sum;
                }
            }
            Ok(Categorical {
                norm_pmf: norm_pmf,
                cdf: cdf,
            })
        }
    }

    fn cdf_max(&self) -> f64 {
        *unsafe { self.cdf.get_unchecked(self.cdf.len() - 1) }
    }
}

impl Sample<f64> for Categorical {
    /// Generate a random sample from a categorical
    /// 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) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for Categorical {
    /// Generate a random independent sample from a categorical
    /// 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) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for Categorical {
    /// Generate a random sample from the categorical distribution
    /// using `r` as the source of randomness
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Categorical, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Categorical::new(&[1.0, 2.0, 3.0]).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let draw = r.next_f64() * self.cdf_max();
        let mut idx = 0;

        if draw == 0.0 {
            // skip zero-probability categories
            let mut el = unsafe { self.cdf.get_unchecked(idx) };
            while *el == 0.0 {
                // don't need bounds checking because we do not allow
                // creating Categorical distributions with all 0.0 probs
                idx += 1;
                el = unsafe { self.cdf.get_unchecked(idx) }
            }
        }
        let mut el = unsafe { self.cdf.get_unchecked(idx) };
        while draw > *el {
            idx += 1;
            el = unsafe { self.cdf.get_unchecked(idx) };
        }
        return idx as f64;
    }
}

impl Univariate<u64, f64> for Categorical {
    /// Calculates the cumulative distribution function for the categorical
    /// distribution at `x`
    ///
    /// # Panics
    ///
    /// If `x < 0.0` or `x > k` where `k` is the number of categories
    /// (i.e. the length of the `prob_mass` slice passed to the constructor)
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(p_j) from 0..x
    /// ```
    ///
    /// where `p_j` is the probability mass for the `j`th category
    fn cdf(&self, x: f64) -> f64 {
        assert!(x >= 0.0 && x <= self.cdf.len() as f64,
                format!("{}",
                        StatsError::ArgIntervalIncl("x", 0.0, self.cdf.len() as f64)));
        if x == self.cdf.len() as f64 {
            1.0
        } else {
            unsafe { self.cdf.get_unchecked(x as usize) / self.cdf_max() }
        }
    }
}

impl Min<u64> for Categorical {
    /// Returns the minimum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> u64 {
        0
    }
}

impl Max<u64> for Categorical {
    /// Returns the maximum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// n
    /// ```
    fn max(&self) -> u64 {
        self.cdf.len() as u64 - 1
    }
}

impl Mean<f64> for Categorical {
    /// Returns the mean of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// E[X] = sum(j * p_j) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass and `k` is the number
    /// of categoires
    fn mean(&self) -> f64 {
        self.norm_pmf.iter().enumerate().fold(0.0, |acc, (idx, &val)| acc + idx as f64 * val)
    }
}

// determines if `p` is a valid probability mass array
// for the Categorical distribution
fn is_valid_prob_mass(p: &[f64]) -> bool {
    !p.iter().any(|&x| x < 0.0 || x.is_nan()) && !p.iter().all(|&x| x == 0.0)
}

#[test]
fn test_is_valid_prob_mass() {
    let invalid = [1.0, f64::NAN, 3.0];
    assert!(!is_valid_prob_mass(&invalid));
    let invalid2 = [-2.0, 5.0, 1.0, 6.2];
    assert!(!is_valid_prob_mass(&invalid2));
    let invalid3 = [0.0, 0.0, 0.0];
    assert!(!is_valid_prob_mass(&invalid3));
    let invalid4: [f64; 0] = [];
    assert!(!is_valid_prob_mass(&invalid4));
    let valid = [5.2, 0.00001, 1e-15, 1000000.12];
    assert!(is_valid_prob_mass(&valid));
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[cfg(test)]
mod test {
    use std::fmt::Debug;
    use statistics::*;
    use distribution::{Univariate, Categorical};

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

    fn create_case(prob_mass: &[f64]) {
        try_create(prob_mass);
    }

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

    fn test_case<T, F>(prob_mass: &[f64], expected: T, eval: F)
        where T: PartialEq + Debug,
              F: Fn(Categorical) -> T
    {
        let n = try_create(prob_mass);
        let x = eval(n);
        assert_eq!(expected, x);
    }

    #[test]
    fn test_create() {
        create_case(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
    }

    #[test]
    fn test_bad_create() {
        bad_create_case(&[-1.0, 1.0]);
        bad_create_case(&[0.0, 0.0]);
    }

    #[test]
    fn test_mean() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 2.0, |x| x.mean());
        test_case(&[0.0, 1.0, 2.0, 1.0], 2.0, |x| x.mean());
        test_case(&[0.0, 0.5, 0.5], 1.5, |x| x.mean());
        test_case(&[0.75, 0.25], 0.25, |x| x.mean());
        test_case(&[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], 5.0, |x| x.mean());

    }

    #[test]
    fn test_min_max() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 0, |x| x.min());
        test_case(&[4.0, 2.5, 2.5, 1.0], 3, |x| x.max());
    }

    #[test]
    fn test_cdf() {
        test_case(&[0.0, 3.0, 1.0, 1.0], 3.0 / 5.0, |x| x.cdf(1.5));
        test_case(&[1.0, 1.0, 1.0, 1.0], 0.25, |x| x.cdf(0.0));
        test_case(&[4.0, 2.5, 2.5, 1.0], 0.4, |x| x.cdf(0.8));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(3.2));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(4.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_low() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(-1.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_high() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(4.5));
    }
}<commit_msg>feat: impl variance for categorical<commit_after>use std::f64;
use rand::Rng;
use rand::distributions::{Sample, IndependentSample};
use statistics::*;
use distribution::{Univariate, Discrete, Distribution};
use {Result, StatsError};

/// Implements the [Categorical](https://en.wikipedia.org/wiki/Categorical_distribution)
/// distribution, also known as the generalized Bernoulli or discrete distribution
///
/// # Examples
///
/// ```
/// use statrs::distribution::{Categorical, Discrete};
/// use statrs::statistics::Mode;
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Categorical {
    norm_pmf: Vec<f64>,
    cdf: Vec<f64>,
}

impl Categorical {
    pub fn new(prob_mass: &[f64]) -> Result<Categorical> {
        if !is_valid_prob_mass(prob_mass) {
            Err(StatsError::BadParams)
        } else {
            // extract un-normalized cdf
            let mut cdf = vec![0.0; prob_mass.len()];
            cdf[0] = prob_mass[0];
            for i in 1..prob_mass.len() {
                unsafe {
                    let val = cdf.get_unchecked(i - 1) + prob_mass.get_unchecked(i);
                    let elem = cdf.get_unchecked_mut(i);
                    *elem = val;
                }
            }

            // extract normalized probability mass
            let sum = cdf[cdf.len() - 1];
            let mut norm_pmf = vec![0.0; prob_mass.len()];
            for i in 0..prob_mass.len() {
                unsafe {
                    let elem = norm_pmf.get_unchecked_mut(i);
                    *elem = prob_mass.get_unchecked(i) / sum;
                }
            }
            Ok(Categorical {
                norm_pmf: norm_pmf,
                cdf: cdf,
            })
        }
    }

    fn cdf_max(&self) -> f64 {
        *unsafe { self.cdf.get_unchecked(self.cdf.len() - 1) }
    }
}

impl Sample<f64> for Categorical {
    /// Generate a random sample from a categorical
    /// 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) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl IndependentSample<f64> for Categorical {
    /// Generate a random independent sample from a categorical
    /// 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) -> f64 {
        super::Distribution::sample(self, r)
    }
}

impl Distribution<f64> for Categorical {
    /// Generate a random sample from the categorical distribution
    /// using `r` as the source of randomness
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate rand;
    /// # extern crate statrs;
    /// use rand::StdRng;
    /// use statrs::distribution::{Categorical, Distribution};
    ///
    /// # fn main() {
    /// let mut r = rand::StdRng::new().unwrap();
    /// let n = Categorical::new(&[1.0, 2.0, 3.0]).unwrap();
    /// print!("{}", n.sample::<StdRng>(&mut r));
    /// # }
    /// ```
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        let draw = r.next_f64() * self.cdf_max();
        let mut idx = 0;

        if draw == 0.0 {
            // skip zero-probability categories
            let mut el = unsafe { self.cdf.get_unchecked(idx) };
            while *el == 0.0 {
                // don't need bounds checking because we do not allow
                // creating Categorical distributions with all 0.0 probs
                idx += 1;
                el = unsafe { self.cdf.get_unchecked(idx) }
            }
        }
        let mut el = unsafe { self.cdf.get_unchecked(idx) };
        while draw > *el {
            idx += 1;
            el = unsafe { self.cdf.get_unchecked(idx) };
        }
        return idx as f64;
    }
}

impl Univariate<u64, f64> for Categorical {
    /// Calculates the cumulative distribution function for the categorical
    /// distribution at `x`
    ///
    /// # Panics
    ///
    /// If `x < 0.0` or `x > k` where `k` is the number of categories
    /// (i.e. the length of the `prob_mass` slice passed to the constructor)
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(p_j) from 0..x
    /// ```
    ///
    /// where `p_j` is the probability mass for the `j`th category
    fn cdf(&self, x: f64) -> f64 {
        assert!(x >= 0.0 && x <= self.cdf.len() as f64,
                format!("{}",
                        StatsError::ArgIntervalIncl("x", 0.0, self.cdf.len() as f64)));
        if x == self.cdf.len() as f64 {
            1.0
        } else {
            unsafe { self.cdf.get_unchecked(x as usize) / self.cdf_max() }
        }
    }
}

impl Min<u64> for Categorical {
    /// Returns the minimum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// 0
    /// ```
    fn min(&self) -> u64 {
        0
    }
}

impl Max<u64> for Categorical {
    /// Returns the maximum value in the domain of the
    /// categorical distribution representable by a 64-bit
    /// integer
    ///
    /// # Formula
    ///
    /// ```ignore
    /// n
    /// ```
    fn max(&self) -> u64 {
        self.cdf.len() as u64 - 1
    }
}

impl Mean<f64> for Categorical {
    /// Returns the mean of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(j * p_j) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass and `k` is the number
    /// of categories
    fn mean(&self) -> f64 {
        self.norm_pmf.iter().enumerate().fold(0.0, |acc, (idx, &val)| acc + idx as f64 * val)
    }
}

impl Variance<f64> for Categorical {
    /// Returns the variance of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sum(p_j * (j - μ)^2) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass, `k` is the number
    /// of categories, and `μ` is the mean
    fn variance(&self) -> f64 {
        let mu = self.mean();
        self.norm_pmf.iter().enumerate().fold(0.0, |acc, (idx, &val)| {
            let r = idx as f64 - mu;
            acc + r * r * val
        })
    }

    /// Returns the standard deviation of the categorical distribution
    ///
    /// # Formula
    ///
    /// ```ignore
    /// sqrt(sum(p_j * (j - μ)^2)) for j in 0..k-1
    /// ```
    ///
    /// where `p_j` is the `j`th probability mass, `k` is the number
    /// of categories, and `μ` is the mean
    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }
}

// determines if `p` is a valid probability mass array
// for the Categorical distribution
fn is_valid_prob_mass(p: &[f64]) -> bool {
    !p.iter().any(|&x| x < 0.0 || x.is_nan()) && !p.iter().all(|&x| x == 0.0)
}

#[test]
fn test_is_valid_prob_mass() {
    let invalid = [1.0, f64::NAN, 3.0];
    assert!(!is_valid_prob_mass(&invalid));
    let invalid2 = [-2.0, 5.0, 1.0, 6.2];
    assert!(!is_valid_prob_mass(&invalid2));
    let invalid3 = [0.0, 0.0, 0.0];
    assert!(!is_valid_prob_mass(&invalid3));
    let invalid4: [f64; 0] = [];
    assert!(!is_valid_prob_mass(&invalid4));
    let valid = [5.2, 0.00001, 1e-15, 1000000.12];
    assert!(is_valid_prob_mass(&valid));
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[cfg(test)]
mod test {
    use std::fmt::Debug;
    use statistics::*;
    use distribution::{Univariate, Categorical};

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

    fn create_case(prob_mass: &[f64]) {
        try_create(prob_mass);
    }

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

    fn test_case<T, F>(prob_mass: &[f64], expected: T, eval: F)
        where T: PartialEq + Debug,
              F: Fn(Categorical) -> T
    {
        let n = try_create(prob_mass);
        let x = eval(n);
        assert_eq!(expected, x);
    }

    #[test]
    fn test_create() {
        create_case(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
    }

    #[test]
    fn test_bad_create() {
        bad_create_case(&[-1.0, 1.0]);
        bad_create_case(&[0.0, 0.0]);
    }

    #[test]
    fn test_mean() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 2.0, |x| x.mean());
        test_case(&[0.0, 1.0, 2.0, 1.0], 2.0, |x| x.mean());
        test_case(&[0.0, 0.5, 0.5], 1.5, |x| x.mean());
        test_case(&[0.75, 0.25], 0.25, |x| x.mean());
        test_case(&[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], 5.0, |x| x.mean());
    }

    #[test]
    fn test_variance() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.5, |x| x.variance());
        test_case(&[0.0, 1.0, 2.0, 1.0], 0.5, |x| x.variance());
        test_case(&[0.0, 0.5, 0.5], 0.25, |x| x.variance());
        test_case(&[0.75, 0.25], 0.1875, |x| x.variance());
        test_case(&[1.0, 0.0, 1.0], 1.0, |x| x.variance());
    }

    #[test]
    fn test_std_dev() {
        test_case(&[0.0, 0.25, 0.5, 0.25], 0.70710678118654752440084436210485, |x| x.std_dev());
        test_case(&[0.0, 1.0, 2.0, 1.0], 0.70710678118654752440084436210485, |x| x.std_dev());
        test_case(&[0.0, 0.5, 0.5], 0.5, |x| x.std_dev());
        test_case(&[0.75, 0.25], 0.43301270189221932338186158537647, |x| x.std_dev());
        test_case(&[1.0, 0.0, 1.0], 1.0, |x| x.std_dev());
    }

    #[test]
    fn test_min_max() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 0, |x| x.min());
        test_case(&[4.0, 2.5, 2.5, 1.0], 3, |x| x.max());
    }

    #[test]
    fn test_cdf() {
        test_case(&[0.0, 3.0, 1.0, 1.0], 3.0 / 5.0, |x| x.cdf(1.5));
        test_case(&[1.0, 1.0, 1.0, 1.0], 0.25, |x| x.cdf(0.0));
        test_case(&[4.0, 2.5, 2.5, 1.0], 0.4, |x| x.cdf(0.8));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(3.2));
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(4.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_low() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(-1.0));
    }

    #[test]
    #[should_panic]
    fn test_cdf_input_high() {
        test_case(&[4.0, 2.5, 2.5, 1.0], 1.0, |x| x.cdf(4.5));
    }
}<|endoftext|>


<commit_before>use antidote::RwLock;
use irc::client::data::User;
use irc::client::prelude::*;

use chrono::NaiveDateTime;
use humantime::format_duration;
use std::time::Duration;
use time;

use plugin::*;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::Fail;
use failure::ResultExt;

pub mod database;
use self::database::Database;

#[derive(PluginName)]
pub struct Tell<T: Database> {
    tells: RwLock<T>,
}

impl<T: Database> Tell<T> {
    pub fn new(db: T) -> Tell<T> {
        Tell {
            tells: RwLock::new(db),
        }
    }

    fn tell_command(
        &self,
        client: &IrcClient,
        command: PluginCommand,
    ) -> Result<String, TellError> {
        if command.tokens.len() < 2 {
            return Ok(self.invalid_command().to_owned());
        }

        let mut online = Vec::new();

        let receivers = command.tokens[0].split(',').filter(|&s| !s.is_empty());
        let sender = command.source;

        let mut no_receiver = true;
        for receiver in receivers {
            if receiver.eq_ignore_ascii_case(client.current_nickname())
                || receiver.eq_ignore_ascii_case(&sender)
            {
                online.push(receiver);
                continue;
            }

            let channels = client
                .list_channels()
                .expect("The irc crate should not be compiled with the \"nochanlists\" feature");

            let find_receiver = |option: Option<Vec<User>>| {
                option.and_then(|users| {
                    users
                        .into_iter()
                        .find(|user| user.get_nickname().eq_ignore_ascii_case(&receiver))
                })
            };

            if channels
                .iter()
                .map(|channel| client.list_users(&channel))
                .map(find_receiver)
                .any(|option| option.is_some())
            {
                online.push(receiver);
                continue;
            }

            let tm = time::now().to_timespec();
            let message = command.tokens[1..].join(" ");
            let tell = database::NewTellMessage {
                sender: &sender,
                receiver: &receiver.to_lowercase(),
                time: NaiveDateTime::from_timestamp(tm.sec, 0u32),
                message: &message,
            };

            debug!("Saving tell for {:?}", receiver);
            self.tells.write().insert_tell(&tell)?;
            no_receiver = false;
        }

        Ok(if no_receiver && online.is_empty() {
            String::from("Invalid receiver.")
        } else {
            match online.len() {
                0 => format!("Got it!"),
                1 => format!("{} is currently online.", online[0]),
                _ => format!("{} are currently online.", online.join(", ")),
            }
        })
    }

    fn on_namelist(&self, client: &IrcClient, channel: &str) -> Result<(), FrippyError> {
        let receivers = self.tells
            .read()
            .get_receivers()
            .context(FrippyErrorKind::Tell)?;

        if let Some(users) = client.list_users(channel) {
            debug!("Outstanding tells for {:?}", receivers);

            for receiver in users
                .iter()
                .map(|u| u.get_nickname())
                .filter(|u| receivers.iter().any(|r| r == &u.to_lowercase()))
            {
                self.send_tells(client, receiver)?;
            }

            Ok(())
        } else {
            Ok(())
        }
    }

    fn send_tells(&self, client: &IrcClient, receiver: &str) -> Result<(), FrippyError> {
        if client.current_nickname() == receiver {
            return Ok(());
        }

        let mut tells = self.tells.write();

        let tell_messages = match tells.get_tells(&receiver.to_lowercase()) {
            Ok(t) => t,
            Err(e) => {
                // This warning only occurs if frippy is built without a database
                #[allow(unreachable_patterns)]
                return match e.kind() {
                    ErrorKind::NotFound => Ok(()),
                    _ => Err(e.context(FrippyErrorKind::Tell))?,
                };
            }
        };

        for tell in tell_messages {
            let now = Duration::new(time::now().to_timespec().sec as u64, 0);
            let dur = now - Duration::new(tell.time.timestamp() as u64, 0);
            let human_dur = format_duration(dur);

            client
                .send_notice(
                    receiver,
                    &format!(
                        "Tell from {} {} ago: {}",
                        tell.sender, human_dur, tell.message
                    ),
                )
                .context(FrippyErrorKind::Connection)?;

            debug!(
                "Sent {:?} from {:?} to {:?}",
                tell.message, tell.sender, receiver
            );
        }

        tells
            .delete_tells(&receiver.to_lowercase())
            .context(FrippyErrorKind::Tell)?;

        Ok(())
    }

    fn invalid_command(&self) -> &str {
        "Incorrect Command. \
         Send \"tell help\" for help."
    }

    fn help(&self) -> &str {
        "usage: tell user message\r\n\
         example: tell Foobar Hello!"
    }
}

impl<T: Database> Plugin for Tell<T> {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        let res = match message.command {
            Command::JOIN(_, _, _) => self.send_tells(client, message.source_nickname().unwrap()),
            Command::NICK(ref nick) => self.send_tells(client, nick),
            Command::Response(resp, ref chan_info, _) => {
                if resp == Response::RPL_NAMREPLY {
                    debug!("NAMREPLY info: {:?}", chan_info);

                    self.on_namelist(client, &chan_info[chan_info.len() - 1])
                } else {
                    Ok(())
                }
            }
            _ => Ok(()),
        };

        match res {
            Ok(_) => ExecutionStatus::Done,
            Err(e) => ExecutionStatus::Err(e),
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Tell should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            client
                .send_notice(&command.source, &self.invalid_command())
                .context(FrippyErrorKind::Connection)?;
            return Ok(());
        }

        let sender = command.source.to_owned();

        match command.tokens[0].as_ref() {
            "help" => client
                .send_notice(&command.source, &self.help())
                .context(FrippyErrorKind::Connection),
            _ => match self.tell_command(client, command) {
                Ok(msg) => client
                    .send_notice(&sender, &msg)
                    .context(FrippyErrorKind::Connection),
                Err(e) => client
                    .send_notice(&sender, &e.to_string())
                    .context(FrippyErrorKind::Connection),
            },
        }?;

        Ok(())
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from("This Plugin does not implement any commands."))
    }
}

use std::fmt;
impl<T: Database> fmt::Debug for Tell<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Tell {{ ... }}")
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "TellError"]
    pub enum ErrorKind {
        /// Not found command error
        #[fail(display = "Tell was not found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<commit_msg>Tell: Allow messages to online users<commit_after>use antidote::RwLock;
use irc::client::data::User;
use irc::client::prelude::*;

use chrono::NaiveDateTime;
use humantime::format_duration;
use std::time::Duration;
use time;

use plugin::*;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::Fail;
use failure::ResultExt;

pub mod database;
use self::database::Database;

#[derive(PluginName)]
pub struct Tell<T: Database> {
    tells: RwLock<T>,
}

impl<T: Database> Tell<T> {
    pub fn new(db: T) -> Tell<T> {
        Tell {
            tells: RwLock::new(db),
        }
    }

    fn tell_command(
        &self,
        client: &IrcClient,
        command: PluginCommand,
    ) -> Result<String, TellError> {
        if command.tokens.len() < 2 {
            return Ok(self.invalid_command().to_owned());
        }

        let mut online = Vec::new();

        let receivers = command.tokens[0].split(',').filter(|&s| !s.is_empty());
        let sender = command.source;

        let mut no_receiver = true;
        for receiver in receivers {
            if receiver.eq_ignore_ascii_case(client.current_nickname())
                || receiver.eq_ignore_ascii_case(&sender)
            {
                online.push(receiver);
                // TODO Change this when https://github.com/aatxe/irc/issues/136 gets resolved
                // continue;
            }

            let channels = client
                .list_channels()
                .expect("The irc crate should not be compiled with the \"nochanlists\" feature");

            let find_receiver = |option: Option<Vec<User>>| {
                option.and_then(|users| {
                    users
                        .into_iter()
                        .find(|user| user.get_nickname().eq_ignore_ascii_case(&receiver))
                })
            };

            if channels
                .iter()
                .map(|channel| client.list_users(&channel))
                .map(find_receiver)
                .any(|option| option.is_some())
            {
                online.push(receiver);
                continue;
            }

            let tm = time::now().to_timespec();
            let message = command.tokens[1..].join(" ");
            let tell = database::NewTellMessage {
                sender: &sender,
                receiver: &receiver.to_lowercase(),
                time: NaiveDateTime::from_timestamp(tm.sec, 0u32),
                message: &message,
            };

            debug!("Saving tell for {:?}", receiver);
            self.tells.write().insert_tell(&tell)?;
            no_receiver = false;
        }

        Ok(if no_receiver && online.is_empty() {
            String::from("Invalid receiver.")
        } else {
            match online.len() {
                0 => format!("Got it!"),
                1 => format!("{} is currently online.", online[0]),
                _ => format!("{} are currently online.", online.join(", ")),
            }
        })
    }

    fn on_namelist(&self, client: &IrcClient, channel: &str) -> Result<(), FrippyError> {
        let receivers = self.tells
            .read()
            .get_receivers()
            .context(FrippyErrorKind::Tell)?;

        if let Some(users) = client.list_users(channel) {
            debug!("Outstanding tells for {:?}", receivers);

            for receiver in users
                .iter()
                .map(|u| u.get_nickname())
                .filter(|u| receivers.iter().any(|r| r == &u.to_lowercase()))
            {
                self.send_tells(client, receiver)?;
            }

            Ok(())
        } else {
            Ok(())
        }
    }

    fn send_tells(&self, client: &IrcClient, receiver: &str) -> Result<(), FrippyError> {
        if client.current_nickname() == receiver {
            return Ok(());
        }

        let mut tells = self.tells.write();

        let tell_messages = match tells.get_tells(&receiver.to_lowercase()) {
            Ok(t) => t,
            Err(e) => {
                // This warning only occurs if frippy is built without a database
                #[allow(unreachable_patterns)]
                return match e.kind() {
                    ErrorKind::NotFound => Ok(()),
                    _ => Err(e.context(FrippyErrorKind::Tell))?,
                };
            }
        };

        for tell in tell_messages {
            let now = Duration::new(time::now().to_timespec().sec as u64, 0);
            let dur = now - Duration::new(tell.time.timestamp() as u64, 0);
            let human_dur = format_duration(dur);

            client
                .send_notice(
                    receiver,
                    &format!(
                        "Tell from {} {} ago: {}",
                        tell.sender, human_dur, tell.message
                    ),
                )
                .context(FrippyErrorKind::Connection)?;

            debug!(
                "Sent {:?} from {:?} to {:?}",
                tell.message, tell.sender, receiver
            );
        }

        tells
            .delete_tells(&receiver.to_lowercase())
            .context(FrippyErrorKind::Tell)?;

        Ok(())
    }

    fn invalid_command(&self) -> &str {
        "Incorrect Command. \
         Send \"tell help\" for help."
    }

    fn help(&self) -> &str {
        "Used to send messages to offline users which they will receive when they come online.\r\n
         usage: tell user message\r\n\
         example: tell Foobar Hello!"
    }
}

impl<T: Database> Plugin for Tell<T> {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        let res = match message.command {
            Command::JOIN(_, _, _) => self.send_tells(client, message.source_nickname().unwrap()),
            Command::NICK(ref nick) => self.send_tells(client, nick),
            Command::PRIVMSG(ref nick, _) => self.send_tells(client, nick),
            Command::Response(resp, ref chan_info, _) => {
                if resp == Response::RPL_NAMREPLY {
                    debug!("NAMREPLY info: {:?}", chan_info);

                    self.on_namelist(client, &chan_info[chan_info.len() - 1])
                } else {
                    Ok(())
                }
            }
            _ => Ok(()),
        };

        match res {
            Ok(_) => ExecutionStatus::Done,
            Err(e) => ExecutionStatus::Err(e),
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Tell should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        if command.tokens.is_empty() {
            client
                .send_notice(&command.source, &self.invalid_command())
                .context(FrippyErrorKind::Connection)?;
            return Ok(());
        }

        let sender = command.source.to_owned();

        match command.tokens[0].as_ref() {
            "help" => client
                .send_notice(&command.source, &self.help())
                .context(FrippyErrorKind::Connection),
            _ => match self.tell_command(client, command) {
                Ok(msg) => client
                    .send_notice(&sender, &msg)
                    .context(FrippyErrorKind::Connection),
                Err(e) => client
                    .send_notice(&sender, &e.to_string())
                    .context(FrippyErrorKind::Connection),
            },
        }?;

        Ok(())
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from("This Plugin does not implement any commands."))
    }
}

use std::fmt;
impl<T: Database> fmt::Debug for Tell<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Tell {{ ... }}")
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "TellError"]
    pub enum ErrorKind {
        /// Not found command error
        #[fail(display = "Tell was not found")]
        NotFound,

        /// MySQL error
        #[cfg(feature = "mysql")]
        #[fail(display = "Failed to execute MySQL Query")]
        MysqlError,

        /// No connection error
        #[cfg(feature = "mysql")]
        #[fail(display = "No connection to the database")]
        NoConnection,
    }
}
<|endoftext|>


<commit_before>use std::process;
use std::fs::File;
use std::io::{Read, BufReader};

use serde_json;
use std::collections::{HashMap, HashSet};

use crate::data_types::{APIMaps, IdGeneSubsetMap, APIGeneSummary, APIAlleleDetails,
                GeneDetails, TermDetails, GenotypeDetails, ReferenceDetails,
                InteractionType, OntAnnotationMap, IdOntAnnotationDetailMap,
                TermShort, TermShortOptionMap, ChromosomeDetails,
                ReferenceShort, ReferenceShortOptionMap,
                GeneShort, GeneShortOptionMap, GeneQueryData};
use crate::web::config::{Config, TermAndName};
use crate::api::query::{SingleOrMultiAllele, QueryExpressionFilter};

use flate2::read::GzDecoder;

use crate::types::{TermId, GeneUniquename};

use pombase_rc_string::RcString;

pub struct APIData {
    maps: APIMaps,
    secondary_identifiers_map: HashMap<TermId, TermId>,
}

fn make_secondary_identifiers_map(terms: &HashMap<TermId, TermDetails>)
                                  -> HashMap<TermId, TermId>
{
    let mut ret_map = HashMap::new();

    for term_details in terms.values() {
        for secondary_identifier in &term_details.secondary_identifiers {
            ret_map.insert(secondary_identifier.clone(),
                           term_details.termid.clone());
        }
    }

    ret_map
}

impl APIData {
    pub fn new_from_file(config: &Config, search_maps_file_name: &str)
               -> APIData
    {
        let file = match File::open(search_maps_file_name) {
            Ok(file) => file,
            Err(err) => {
                eprint!("Failed to read {}: {}\n", search_maps_file_name, err);
                process::exit(1);
            }
        };
        let reader = BufReader::new(file);

        let mut decoder = GzDecoder::new(reader);
        let mut decoded_json = String::new();
        decoder.read_to_string(&mut decoded_json).unwrap();

        let maps: APIMaps =
            match serde_json::from_str(&decoded_json) {
                Ok(results) => results,
                Err(err) => {
                    eprint!("failed to parse {}: {}", search_maps_file_name, err);
                    process::exit(1);
                },
            };

        APIData::new(config, &maps)
    }

    pub fn new(config: &Config, maps: &APIMaps)
               ->APIData
    {
        let mut maps = maps.clone();
        let mut new_entries: IdGeneSubsetMap = HashMap::new();

        let prefixes_to_remove: Vec<String> =
            config.server.subsets.prefixes_to_remove
            .iter().map(|prefix| prefix.clone() + ":").collect();

        // strip prefixes and add to map
        for (subset_name, subset_details) in &maps.gene_subsets {
            for prefix in &prefixes_to_remove {
                if subset_name.starts_with(prefix) {
                    let new_subset_name = &subset_name[prefix.len()..];
                    new_entries.insert(RcString::from(new_subset_name), subset_details.clone());
                }
            }
        }

        maps.gene_subsets.extend(new_entries);

        let secondary_identifiers_map =
            make_secondary_identifiers_map(&maps.terms);

        APIData {
            secondary_identifiers_map,
            maps
        }
    }

    pub fn gene_uniquename_of_id(&self, id: &RcString) -> Option<GeneUniquename> {
        if self.maps.gene_summaries.contains_key(id) {
            Some(id.clone())
        } else {
            if let Some(gene_uniquename) =
                self.maps.gene_name_gene_map.get(id) {
                    Some(gene_uniquename.clone())
                } else {
                    let lower_case_id = id.to_lowercase();
                    for gene_uniquename in self.maps.gene_summaries.keys() {
                        if gene_uniquename.to_lowercase() == lower_case_id {
                            return Some(gene_uniquename.clone());
                        }
                    }
                    None
                }
        }
    }

    pub fn get_gene_summary(&self, gene_uniquename: &str) -> Option<&APIGeneSummary> {
        self.maps.gene_summaries.get(gene_uniquename)
    }

    pub fn get_gene_details(&self, gene_uniquename: &str) -> Option<&GeneDetails> {
        self.maps.genes.get(gene_uniquename)
    }

    pub fn get_gene_query_data(&self, gene_uniquename: &str) -> Option<&GeneQueryData> {
        self.maps.gene_query_data_map.get(gene_uniquename)
    }

    pub fn genes_of_termid(&self, term_id: &str) -> Vec<GeneUniquename> {
        match self.maps.termid_genes.get(term_id) {
            Some(gene_uniquenames) => {
                gene_uniquenames.iter().cloned().collect::<Vec<_>>()
            },
            None => vec![],
        }
    }

    pub fn genes_of_subset(&self, search_name: &str) -> Vec<GeneUniquename> {
        if search_name.starts_with('!') || search_name.ends_with('*') {
            let mut trimmed_search_name = search_name.to_owned();
            let invert_search = search_name.starts_with('!');
            let wildcard = search_name.ends_with('*');
            if invert_search {
                trimmed_search_name.remove(0);
            }
            if wildcard {
                trimmed_search_name.pop();
            }
            let mut genes = HashSet::new();
            for (subset_name, subset_details) in &self.maps.gene_subsets {
                let name_matches =
                    wildcard && subset_name.starts_with(&trimmed_search_name) ||
                    trimmed_search_name.eq(subset_name);

                if !invert_search && name_matches || invert_search && !name_matches {
                    genes.extend(subset_details.elements.iter().cloned());
                }
            }
            genes.iter().cloned().collect::<Vec<_>>()
        } else {
            match self.maps.gene_subsets.get(search_name) {
                Some(subset_details) => {
                    subset_details.elements.iter().cloned().collect::<Vec<_>>()
                },
                None => vec![],
            }
        }
    }

    pub fn genes_of_genotypes(&self, term_id: &str,
                              single_or_multi_allele: &SingleOrMultiAllele,
                              expression_filter: &Option<QueryExpressionFilter>,
                              conditions_filter: &HashSet<TermAndName>,
                              excluded_conditions_filter: &HashSet<TermAndName>)
                              -> Vec<GeneUniquename>
    {
        if let Some(annotations) = self.maps.termid_genotype_annotation.get(term_id) {
            let mut genes: HashSet<GeneUniquename> = HashSet::new();
            for annotation in annotations {

                let mut add_single = false;
                let mut add_multi = false;

                match *single_or_multi_allele {
                    SingleOrMultiAllele::Single => add_single = true,
                    SingleOrMultiAllele::Multi => add_multi = true,
                    SingleOrMultiAllele::Both => {
                        add_single = true;
                        add_multi = true;
                    },
                }

                let expression_matches =
                    |allele_details: &APIAlleleDetails| {
                        let expression = &allele_details.expression;
                        let allele_type = &allele_details.allele_type;
                        if let Some(ref expression_filter) = *expression_filter {
                            if *expression_filter == QueryExpressionFilter::Any {
                                return true;
                            }
                            if *expression_filter == QueryExpressionFilter::Null {
                                if allele_type == "deletion" {
                                    return true;
                                }
                                if let Some(ref expression) = *expression {
                                    expression == "Null"
                                } else {
                                    false
                                }
                            } else {
                                if *expression_filter == QueryExpressionFilter::WtOverexpressed {
                                    if let Some(ref expression) = *expression {
                                        expression == "Overexpression" &&
                                            allele_type == "wild_type"
                                    } else {
                                        false
                                    }
                                } else {
                                    false
                                }
                            }
                        } else {
                            true
                        }
                    };

                let condition_matches =
                    |conditions: &HashSet<TermAndName>, test_conditions_filter: &HashSet<TermAndName>| {
                        if test_conditions_filter.len() == 0 {
                            // don't filter in this case
                            return true;
                        }

                        for filter_condition in test_conditions_filter {
                            if conditions.iter().find(
                                |cond| {
                                    filter_condition.termid == cond.termid
                                })
                                .is_some() {
                                return true
                            }
                        }

                        false
                    };

                if annotation.is_multi && add_multi ||
                    !annotation.is_multi && add_single &&
                    expression_matches(&annotation.alleles[0]) &&
                    condition_matches(&annotation.conditions, conditions_filter) &&
                    (excluded_conditions_filter.len() == 0 ||
                     !condition_matches(&annotation.conditions, excluded_conditions_filter)) {
                        for allele_details in &annotation.alleles {
                            genes.insert(allele_details.gene.clone());
                        }
                }
            }
            genes.into_iter().collect::<Vec<_>>()
        } else {
            vec![]
        }
    }

    pub fn filter_genes(&self, p: & dyn Fn(&APIGeneSummary) -> bool)
                        -> Vec<GeneUniquename>
    {
        self.maps.gene_summaries.values()
            .filter(|summ| p(summ))
            .map(|summ| summ.uniquename.clone())
            .collect()
    }

    fn detail_map_of_cv_annotations(&self, ont_annotation_map: &OntAnnotationMap)
                                    -> IdOntAnnotationDetailMap
    {
        let mut details_map = HashMap::new();

        for term_annotations in ont_annotation_map.values() {
            for term_annotation in term_annotations {
                for annotation_detail_id in &term_annotation.annotations {
                    let details =
                        &self.maps.annotation_details[annotation_detail_id];
                    details_map.insert(*annotation_detail_id, details.clone());
                }
            }
        }

        details_map
    }

    fn fill_term_map(&self, term_map: &TermShortOptionMap) -> TermShortOptionMap {
        let mut ret = term_map.clone();
        for termid in term_map.keys() {
            if let Some(term_details) = self.maps.terms.get(termid) {
                let term_short = TermShort::from_term_details(term_details);
                ret.insert(termid.clone(), Some(term_short));
            } else {
                eprint!("WARNING missing short info for: {}\n", termid);
                ret.insert(termid.clone(), None);
            }
        }
        ret
    }

    fn fill_gene_map(&self, gene_map: &GeneShortOptionMap) -> GeneShortOptionMap {
        let mut ret = gene_map.clone();
        for gene_uniquename in gene_map.keys() {
            let gene_details = &self.maps.genes[gene_uniquename];
            let gene_short = GeneShort::from_gene_details(gene_details);
            ret.insert(gene_uniquename.clone(), Some(gene_short));
        }
        ret
    }

    fn fill_reference_map(&self, reference_map: &ReferenceShortOptionMap) -> ReferenceShortOptionMap {
        let mut ret = reference_map.clone();
        for reference_uniquename in reference_map.keys() {
            let reference_details = &self.maps.references[reference_uniquename];
            let reference_short = ReferenceShort::from_reference_details(reference_details);
            ret.insert(reference_uniquename.clone(), Some(reference_short));
        }
        ret
    }

    // return a GeneDetails object with the term, genes and references maps filled in
    pub fn get_full_gene_details(&self, gene_uniquename: &str) -> Option<GeneDetails> {
        if let Some(gene_ref) = self.maps.genes.get(gene_uniquename) {
            let mut gene = gene_ref.clone();
            let details_map = self.detail_map_of_cv_annotations(&gene.cv_annotations);
            gene.annotation_details = details_map;
            gene.terms_by_termid = self.fill_term_map(&gene.terms_by_termid);
            gene.genes_by_uniquename = self.fill_gene_map(&gene.genes_by_uniquename);
            gene.references_by_uniquename =
                self.fill_reference_map(&gene.references_by_uniquename);
            Some(gene)
        } else {
            None
        }
    }

    pub fn get_genotype_details(&self, genotype_uniquename: &str) -> Option<GenotypeDetails> {
        if let Some(genotype_ref) = self.maps.genotypes.get(genotype_uniquename) {
            let mut genotype = genotype_ref.clone();
            let details_map = self.detail_map_of_cv_annotations(&genotype.cv_annotations);
            genotype.annotation_details = details_map;
            genotype.terms_by_termid = self.fill_term_map(&genotype.terms_by_termid);
            genotype.genes_by_uniquename = self.fill_gene_map(&genotype.genes_by_uniquename);
            genotype.references_by_uniquename =
                self.fill_reference_map(&genotype.references_by_uniquename);
            Some(genotype)
        } else {
            None
        }
    }

    fn term_details_helper(&self, term_ref: &TermDetails) -> Option<TermDetails> {
        let mut term = term_ref.clone();
        let details_map = self.detail_map_of_cv_annotations(&term.cv_annotations);
        term.annotation_details = details_map;
        term.terms_by_termid = self.fill_term_map(&term.terms_by_termid);
        term.genes_by_uniquename = self.fill_gene_map(&term.genes_by_uniquename);
        term.references_by_uniquename =
            self.fill_reference_map(&term.references_by_uniquename);
        Some(term)
    }

    pub fn get_term_details(&self, termid: &str) -> Option<TermDetails> {
        let termid =
            if let Some(real_termid) = self.secondary_identifiers_map.get(termid) {
                real_termid
            } else {
                termid
            };
        if let Some(term_ref) = self.maps.terms.get(termid) {
            self.term_details_helper(term_ref)
        } else {
            None
        }
    }

    pub fn get_reference_details(&self, reference_uniquename: &str) -> Option<ReferenceDetails> {
        if let Some(reference_ref) = self.maps.references.get(reference_uniquename) {
            let mut reference = reference_ref.clone();
            let details_map = self.detail_map_of_cv_annotations(&reference.cv_annotations);
            reference.annotation_details = details_map;
            reference.terms_by_termid = self.fill_term_map(&reference.terms_by_termid);
            reference.genes_by_uniquename = self.fill_gene_map(&reference.genes_by_uniquename);
            Some(reference)
        } else {
            None
        }
    }

    pub fn interactors_of_genes(&self, gene_uniquename: &GeneUniquename,
                                interaction_type: InteractionType)
                                -> Vec<GeneUniquename> {
        if let Some(interactors) = self.maps.interactors_of_genes.get(gene_uniquename) {
            interactors.iter()
                .filter(|interactor| {
                    interactor.interaction_type == interaction_type
                })
                .map(|interactor| interactor.interactor_uniquename.clone())
                .collect::<Vec<_>>()
        } else {
            vec![]
        }
    }

    pub fn get_chr_details(&self, chr_name: &str) -> Option<&ChromosomeDetails> {
        self.maps.chromosomes.get(chr_name)
    }
}
<commit_msg>Add example impl for reading JSON directly<commit_after>use std::process;
use std::fs::File;
use std::io::{Read, BufReader};

use serde_json;
use std::collections::{HashMap, HashSet};

use crate::data_types::{APIMaps, IdGeneSubsetMap, APIGeneSummary, APIAlleleDetails,
                GeneDetails, TermDetails, GenotypeDetails, ReferenceDetails,
                InteractionType, OntAnnotationMap, IdOntAnnotationDetailMap,
                TermShort, TermShortOptionMap, ChromosomeDetails,
                ReferenceShort, ReferenceShortOptionMap,
                GeneShort, GeneShortOptionMap, GeneQueryData};
use crate::web::config::{Config, TermAndName};
use crate::api::query::{SingleOrMultiAllele, QueryExpressionFilter};

use flate2::read::GzDecoder;

use crate::types::{TermId, GeneUniquename};

use pombase_rc_string::RcString;

pub struct APIData {
    maps: APIMaps,
    secondary_identifiers_map: HashMap<TermId, TermId>,
}

fn make_secondary_identifiers_map(terms: &HashMap<TermId, TermDetails>)
                                  -> HashMap<TermId, TermId>
{
    let mut ret_map = HashMap::new();

    for term_details in terms.values() {
        for secondary_identifier in &term_details.secondary_identifiers {
            ret_map.insert(secondary_identifier.clone(),
                           term_details.termid.clone());
        }
    }

    ret_map
}

impl APIData {
    pub fn new_from_file(config: &Config, search_maps_file_name: &str)
               -> APIData
    {
        let file = match File::open(search_maps_file_name) {
            Ok(file) => file,
            Err(err) => {
                eprint!("Failed to read {}: {}\n", search_maps_file_name, err);
                process::exit(1);
            }
        };

        let reader = BufReader::new(file);
        let mut decoder = GzDecoder::new(reader);

//  this uses less peak memory but is 4X slower
//  See: https://github.com/serde-rs/json/issues/160
//        let serde_result = serde_json::de::from_reader(&mut decoder);

        let mut decoded_json = String::new();
        decoder.read_to_string(&mut decoded_json).unwrap();
        let serde_result = serde_json::from_str(&decoded_json);

        let maps: APIMaps =
            match serde_result {
                Ok(results) => results,
                Err(err) => {
                    eprint!("failed to parse {}: {}", search_maps_file_name, err);
                    process::exit(1);
                },
            };

        APIData::new(config, &maps)
    }

    pub fn new(config: &Config, maps: &APIMaps)
               ->APIData
    {
        let mut maps = maps.clone();
        let mut new_entries: IdGeneSubsetMap = HashMap::new();

        let prefixes_to_remove: Vec<String> =
            config.server.subsets.prefixes_to_remove
            .iter().map(|prefix| prefix.clone() + ":").collect();

        // strip prefixes and add to map
        for (subset_name, subset_details) in &maps.gene_subsets {
            for prefix in &prefixes_to_remove {
                if subset_name.starts_with(prefix) {
                    let new_subset_name = &subset_name[prefix.len()..];
                    new_entries.insert(RcString::from(new_subset_name), subset_details.clone());
                }
            }
        }

        maps.gene_subsets.extend(new_entries);

        let secondary_identifiers_map =
            make_secondary_identifiers_map(&maps.terms);

        APIData {
            secondary_identifiers_map,
            maps
        }
    }

    pub fn gene_uniquename_of_id(&self, id: &RcString) -> Option<GeneUniquename> {
        if self.maps.gene_summaries.contains_key(id) {
            Some(id.clone())
        } else {
            if let Some(gene_uniquename) =
                self.maps.gene_name_gene_map.get(id) {
                    Some(gene_uniquename.clone())
                } else {
                    let lower_case_id = id.to_lowercase();
                    for gene_uniquename in self.maps.gene_summaries.keys() {
                        if gene_uniquename.to_lowercase() == lower_case_id {
                            return Some(gene_uniquename.clone());
                        }
                    }
                    None
                }
        }
    }

    pub fn get_gene_summary(&self, gene_uniquename: &str) -> Option<&APIGeneSummary> {
        self.maps.gene_summaries.get(gene_uniquename)
    }

    pub fn get_gene_details(&self, gene_uniquename: &str) -> Option<&GeneDetails> {
        self.maps.genes.get(gene_uniquename)
    }

    pub fn get_gene_query_data(&self, gene_uniquename: &str) -> Option<&GeneQueryData> {
        self.maps.gene_query_data_map.get(gene_uniquename)
    }

    pub fn genes_of_termid(&self, term_id: &str) -> Vec<GeneUniquename> {
        match self.maps.termid_genes.get(term_id) {
            Some(gene_uniquenames) => {
                gene_uniquenames.iter().cloned().collect::<Vec<_>>()
            },
            None => vec![],
        }
    }

    pub fn genes_of_subset(&self, search_name: &str) -> Vec<GeneUniquename> {
        if search_name.starts_with('!') || search_name.ends_with('*') {
            let mut trimmed_search_name = search_name.to_owned();
            let invert_search = search_name.starts_with('!');
            let wildcard = search_name.ends_with('*');
            if invert_search {
                trimmed_search_name.remove(0);
            }
            if wildcard {
                trimmed_search_name.pop();
            }
            let mut genes = HashSet::new();
            for (subset_name, subset_details) in &self.maps.gene_subsets {
                let name_matches =
                    wildcard && subset_name.starts_with(&trimmed_search_name) ||
                    trimmed_search_name.eq(subset_name);

                if !invert_search && name_matches || invert_search && !name_matches {
                    genes.extend(subset_details.elements.iter().cloned());
                }
            }
            genes.iter().cloned().collect::<Vec<_>>()
        } else {
            match self.maps.gene_subsets.get(search_name) {
                Some(subset_details) => {
                    subset_details.elements.iter().cloned().collect::<Vec<_>>()
                },
                None => vec![],
            }
        }
    }

    pub fn genes_of_genotypes(&self, term_id: &str,
                              single_or_multi_allele: &SingleOrMultiAllele,
                              expression_filter: &Option<QueryExpressionFilter>,
                              conditions_filter: &HashSet<TermAndName>,
                              excluded_conditions_filter: &HashSet<TermAndName>)
                              -> Vec<GeneUniquename>
    {
        if let Some(annotations) = self.maps.termid_genotype_annotation.get(term_id) {
            let mut genes: HashSet<GeneUniquename> = HashSet::new();
            for annotation in annotations {

                let mut add_single = false;
                let mut add_multi = false;

                match *single_or_multi_allele {
                    SingleOrMultiAllele::Single => add_single = true,
                    SingleOrMultiAllele::Multi => add_multi = true,
                    SingleOrMultiAllele::Both => {
                        add_single = true;
                        add_multi = true;
                    },
                }

                let expression_matches =
                    |allele_details: &APIAlleleDetails| {
                        let expression = &allele_details.expression;
                        let allele_type = &allele_details.allele_type;
                        if let Some(ref expression_filter) = *expression_filter {
                            if *expression_filter == QueryExpressionFilter::Any {
                                return true;
                            }
                            if *expression_filter == QueryExpressionFilter::Null {
                                if allele_type == "deletion" {
                                    return true;
                                }
                                if let Some(ref expression) = *expression {
                                    expression == "Null"
                                } else {
                                    false
                                }
                            } else {
                                if *expression_filter == QueryExpressionFilter::WtOverexpressed {
                                    if let Some(ref expression) = *expression {
                                        expression == "Overexpression" &&
                                            allele_type == "wild_type"
                                    } else {
                                        false
                                    }
                                } else {
                                    false
                                }
                            }
                        } else {
                            true
                        }
                    };

                let condition_matches =
                    |conditions: &HashSet<TermAndName>, test_conditions_filter: &HashSet<TermAndName>| {
                        if test_conditions_filter.len() == 0 {
                            // don't filter in this case
                            return true;
                        }

                        for filter_condition in test_conditions_filter {
                            if conditions.iter().find(
                                |cond| {
                                    filter_condition.termid == cond.termid
                                })
                                .is_some() {
                                return true
                            }
                        }

                        false
                    };

                if annotation.is_multi && add_multi ||
                    !annotation.is_multi && add_single &&
                    expression_matches(&annotation.alleles[0]) &&
                    condition_matches(&annotation.conditions, conditions_filter) &&
                    (excluded_conditions_filter.len() == 0 ||
                     !condition_matches(&annotation.conditions, excluded_conditions_filter)) {
                        for allele_details in &annotation.alleles {
                            genes.insert(allele_details.gene.clone());
                        }
                }
            }
            genes.into_iter().collect::<Vec<_>>()
        } else {
            vec![]
        }
    }

    pub fn filter_genes(&self, p: & dyn Fn(&APIGeneSummary) -> bool)
                        -> Vec<GeneUniquename>
    {
        self.maps.gene_summaries.values()
            .filter(|summ| p(summ))
            .map(|summ| summ.uniquename.clone())
            .collect()
    }

    fn detail_map_of_cv_annotations(&self, ont_annotation_map: &OntAnnotationMap)
                                    -> IdOntAnnotationDetailMap
    {
        let mut details_map = HashMap::new();

        for term_annotations in ont_annotation_map.values() {
            for term_annotation in term_annotations {
                for annotation_detail_id in &term_annotation.annotations {
                    let details =
                        &self.maps.annotation_details[annotation_detail_id];
                    details_map.insert(*annotation_detail_id, details.clone());
                }
            }
        }

        details_map
    }

    fn fill_term_map(&self, term_map: &TermShortOptionMap) -> TermShortOptionMap {
        let mut ret = term_map.clone();
        for termid in term_map.keys() {
            if let Some(term_details) = self.maps.terms.get(termid) {
                let term_short = TermShort::from_term_details(term_details);
                ret.insert(termid.clone(), Some(term_short));
            } else {
                eprint!("WARNING missing short info for: {}\n", termid);
                ret.insert(termid.clone(), None);
            }
        }
        ret
    }

    fn fill_gene_map(&self, gene_map: &GeneShortOptionMap) -> GeneShortOptionMap {
        let mut ret = gene_map.clone();
        for gene_uniquename in gene_map.keys() {
            let gene_details = &self.maps.genes[gene_uniquename];
            let gene_short = GeneShort::from_gene_details(gene_details);
            ret.insert(gene_uniquename.clone(), Some(gene_short));
        }
        ret
    }

    fn fill_reference_map(&self, reference_map: &ReferenceShortOptionMap) -> ReferenceShortOptionMap {
        let mut ret = reference_map.clone();
        for reference_uniquename in reference_map.keys() {
            let reference_details = &self.maps.references[reference_uniquename];
            let reference_short = ReferenceShort::from_reference_details(reference_details);
            ret.insert(reference_uniquename.clone(), Some(reference_short));
        }
        ret
    }

    // return a GeneDetails object with the term, genes and references maps filled in
    pub fn get_full_gene_details(&self, gene_uniquename: &str) -> Option<GeneDetails> {
        if let Some(gene_ref) = self.maps.genes.get(gene_uniquename) {
            let mut gene = gene_ref.clone();
            let details_map = self.detail_map_of_cv_annotations(&gene.cv_annotations);
            gene.annotation_details = details_map;
            gene.terms_by_termid = self.fill_term_map(&gene.terms_by_termid);
            gene.genes_by_uniquename = self.fill_gene_map(&gene.genes_by_uniquename);
            gene.references_by_uniquename =
                self.fill_reference_map(&gene.references_by_uniquename);
            Some(gene)
        } else {
            None
        }
    }

    pub fn get_genotype_details(&self, genotype_uniquename: &str) -> Option<GenotypeDetails> {
        if let Some(genotype_ref) = self.maps.genotypes.get(genotype_uniquename) {
            let mut genotype = genotype_ref.clone();
            let details_map = self.detail_map_of_cv_annotations(&genotype.cv_annotations);
            genotype.annotation_details = details_map;
            genotype.terms_by_termid = self.fill_term_map(&genotype.terms_by_termid);
            genotype.genes_by_uniquename = self.fill_gene_map(&genotype.genes_by_uniquename);
            genotype.references_by_uniquename =
                self.fill_reference_map(&genotype.references_by_uniquename);
            Some(genotype)
        } else {
            None
        }
    }

    fn term_details_helper(&self, term_ref: &TermDetails) -> Option<TermDetails> {
        let mut term = term_ref.clone();
        let details_map = self.detail_map_of_cv_annotations(&term.cv_annotations);
        term.annotation_details = details_map;
        term.terms_by_termid = self.fill_term_map(&term.terms_by_termid);
        term.genes_by_uniquename = self.fill_gene_map(&term.genes_by_uniquename);
        term.references_by_uniquename =
            self.fill_reference_map(&term.references_by_uniquename);
        Some(term)
    }

    pub fn get_term_details(&self, termid: &str) -> Option<TermDetails> {
        let termid =
            if let Some(real_termid) = self.secondary_identifiers_map.get(termid) {
                real_termid
            } else {
                termid
            };
        if let Some(term_ref) = self.maps.terms.get(termid) {
            self.term_details_helper(term_ref)
        } else {
            None
        }
    }

    pub fn get_reference_details(&self, reference_uniquename: &str) -> Option<ReferenceDetails> {
        if let Some(reference_ref) = self.maps.references.get(reference_uniquename) {
            let mut reference = reference_ref.clone();
            let details_map = self.detail_map_of_cv_annotations(&reference.cv_annotations);
            reference.annotation_details = details_map;
            reference.terms_by_termid = self.fill_term_map(&reference.terms_by_termid);
            reference.genes_by_uniquename = self.fill_gene_map(&reference.genes_by_uniquename);
            Some(reference)
        } else {
            None
        }
    }

    pub fn interactors_of_genes(&self, gene_uniquename: &GeneUniquename,
                                interaction_type: InteractionType)
                                -> Vec<GeneUniquename> {
        if let Some(interactors) = self.maps.interactors_of_genes.get(gene_uniquename) {
            interactors.iter()
                .filter(|interactor| {
                    interactor.interaction_type == interaction_type
                })
                .map(|interactor| interactor.interactor_uniquename.clone())
                .collect::<Vec<_>>()
        } else {
            vec![]
        }
    }

    pub fn get_chr_details(&self, chr_name: &str) -> Option<&ChromosomeDetails> {
        self.maps.chromosomes.get(chr_name)
    }
}
<|endoftext|>


<commit_before>//! This crates contains the types and methods needed to interact with the DICOM upper-layer protocol.
//!
//! It is very low level and not very usable as is.
//!
//! Eventually, a finite-state-machine and higher-level SCU/SCP helpers will be added that will make
//! interacting with these types more idiomatic and friendly.
//!
//! The [`pdu`](crate::pdu) module
//! provides data structures representing _protocol data units_,
//! which are passed around as part of the DICOM network communication support.

pub mod pdu;
pub mod association;
<commit_msg>[ul] Improve root level docs<commit_after>//! This crates contains the types and methods needed to interact
//! with DICOM nodes through the upper layer protocol.
//!
//! This crate can be used as a base
//! for finite-state machines and higher-level helpers,
//! enabling the creation of concrete service class users (SCUs)
//! and service class providers (SCPs).
//!
//! - The [`pdu`](crate::pdu) module
//! provides data structures representing _protocol data units_,
//! which are passed around as part of the DICOM network communication support.
//! - The [`association`](crate::association) module
//! comprises abstractions for establishing and negotiating associations
//! between application entities,
//! via the upper layer protocol by TCP.

pub mod association;
pub mod pdu;
<|endoftext|>


<commit_before>use std;
use slog::Logger;
use std::time::{SystemTime, UNIX_EPOCH};
use chrono::{Utc, TimeZone, DateTime};
use hostname;

// The complete execution context of Qork.
pub struct Context {
    pub logger: Logger,
    pub args: Vec<String>,
    pub exe_path: Option<std::path::PathBuf>,
    pub exe_meta_data: Option<std::fs::Metadata>,
    pub hostname: Option<String>
}

fn system_time_to_date_time(t: SystemTime) -> DateTime<Utc> {
    let (sec, nsec) = match t.duration_since(UNIX_EPOCH) {
        Ok(dur) => (dur.as_secs() as i64, dur.subsec_nanos()),
        Err(e) => { // unlikely but should be handled
            let dur = e.duration();
            let (sec, nsec) = (dur.as_secs() as i64, dur.subsec_nanos());
            if nsec == 0 {
                (-sec, 0)
            } else {
                (-sec - 1, 1_000_000_000 - nsec)
            }
        },
    };
    Utc.timestamp(sec, nsec)
}

impl Context {
    pub fn new(logger: Logger) -> Context {
        let exe = std::env::current_exe().ok();
        let md = exe.as_ref().and_then(|e| e.metadata().ok());

        Context {
            logger: logger,
            args: std::env::args().collect(),
            exe_path: exe,
            exe_meta_data: md,
            hostname: hostname::get_hostname(),
        }
    }

    pub fn version(&self) -> &'static str {
        "0.1.0"
    }

    pub fn log_created_message(&self) -> () {
        let p = self.exe_path.as_ref()
            .and_then(|p| p.to_str())
            .unwrap_or("unknown");

        let bytes = self.exe_meta_data.as_ref()
            .map(|m| m.len().to_string())
            .unwrap_or("unknown".to_string());

        let mdate = self.exe_meta_data.as_ref()
            .and_then(|md| md.modified().ok())
            .map(|md| system_time_to_date_time(md))
            .map(|md| md.format("%Y-%m-%d %H:%M:%S%.3f UTC").to_string())
            .unwrap_or("unknown".to_string());

        info!(self.logger, "Created Context";
               "version" => self.version(),
               "hostname" => &self.hostname,
               "exe_modified" => mdate,
               "exe_bytes" => bytes,
               "exe_path" => p
         );
    }
}<commit_msg>Write a function to determine where .qork.d is.<commit_after>use std;
use slog::Logger;
use std::time::{SystemTime, UNIX_EPOCH};
use chrono::{Utc, TimeZone, DateTime};
use hostname;
use std::path::PathBuf;
use std::env;

// The complete execution context of Qork.
pub struct Context {
    pub logger: Logger,
    pub args: Vec<String>,
    pub exe_path: Option<std::path::PathBuf>,
    pub exe_meta_data: Option<std::fs::Metadata>,
    pub hostname: Option<String>,
    pub config_directory: PathBuf
}

fn system_time_to_date_time(t: SystemTime) -> DateTime<Utc> {
    let (sec, nsec) = match t.duration_since(UNIX_EPOCH) {
        Ok(dur) => (dur.as_secs() as i64, dur.subsec_nanos()),
        Err(e) => { // unlikely but should be handled
            let dur = e.duration();
            let (sec, nsec) = (dur.as_secs() as i64, dur.subsec_nanos());
            if nsec == 0 {
                (-sec, 0)
            } else {
                (-sec - 1, 1_000_000_000 - nsec)
            }
        },
    };
    Utc.timestamp(sec, nsec)
}

fn get_config_directory(logger: &Logger) -> PathBuf {
    const CONFIG_DIR_KEY : &'static str = "config_dir";
    const CONFIG_DIR : &'static str = ".qork.d";

    // Command line has highest priority.
    // let msg = "Configuration Directory determined from command line arguments.";

    // Next is an environment variable.
    if let Ok(env_var) = std::env::var("QORK_CONFIG_DIR") {
        if !env_var.is_empty() {
            info!(logger, "Configuration Directory determined from QORK_CONFIG_DIR environment variable"; CONFIG_DIR_KEY => &env_var);
            return PathBuf::from(env_var);
        }
    }

    // If still no luck, try for '.qork.d' in the user's home directory - if we can
    // determine the user's home directory, that is.
    if let Some(mut home_dir) = env::home_dir() {
        home_dir.push(CONFIG_DIR);
        info!(logger, "Configuration Directory defaulted to '{}' in user's home directory", CONFIG_DIR; CONFIG_DIR_KEY => %&home_dir.display());
        return home_dir;
    }

    // Still stuck? Try .qork.d in the exe's directroy.
    if let Ok(mut exe_path) = std::env::current_exe() {
        exe_path.pop();
        exe_path.push(CONFIG_DIR);
        info!(logger, "Configuration Directory defaulted to '{}' in directory of the exe", CONFIG_DIR; CONFIG_DIR_KEY => %&exe_path.display());
        return exe_path;
    }

    // Last, just default to .qork.d in the current working directory.
    warn!(logger, "Unable to determine Configuration Directory from command line parameter, home dir, or exe dir: defaulting to '{}' in the current working directory", CONFIG_DIR;
        CONFIG_DIR_KEY => CONFIG_DIR);
    PathBuf::from(CONFIG_DIR)
}

impl Context {
    pub fn new(logger: Logger) -> Context {
        let exe = std::env::current_exe().ok();
        let md = exe.as_ref().and_then(|e| e.metadata().ok());
        let cd = get_config_directory(&logger);

        Context {
            logger: logger,
            args: std::env::args().collect(),
            exe_path: exe,
            exe_meta_data: md,
            hostname: hostname::get_hostname(),
            config_directory: cd
        }
    }

    pub fn version(&self) -> &'static str {
        "0.1.0"
    }

    pub fn log_created_message(&self) -> () {
        let p = self.exe_path.as_ref()
            .and_then(|p| p.to_str())
            .unwrap_or("unknown");

        let bytes = self.exe_meta_data.as_ref()
            .map(|m| m.len().to_string())
            .unwrap_or("unknown".to_string());

        let mdate = self.exe_meta_data.as_ref()
            .and_then(|md| md.modified().ok())
            .map(|md| system_time_to_date_time(md))
            .map(|md| md.format("%Y-%m-%d %H:%M:%S%.3f UTC").to_string())
            .unwrap_or("unknown".to_string());

        info!(self.logger, "Created Context";
               "config_directory" => &self.config_directory.to_str(),
               "version" => self.version(),
               "hostname" => &self.hostname,
               "exe_modified" => mdate,
               "exe_bytes" => bytes,
               "exe_path" => p
         );
    }
}<|endoftext|>


<commit_before>#[doc = "Finds docs for reexported items and duplicates them"];

import std::map;
import rustc::syntax::ast;
import rustc::syntax::ast_util;
import rustc::util::common;

export mk_pass;

fn mk_pass() -> pass {
    run
}

type def_set = map::set<ast::def_id>;
type def_map = map::hashmap<ast::def_id, doc::itemtag>;
type path_map = map::hashmap<str, [(str, doc::itemtag)]>;

fn run(srv: astsrv::srv, doc: doc::cratedoc) -> doc::cratedoc {

    // First gather the set of defs that are used as reexports
    let def_set = build_reexport_def_set(srv);

    // Now find the docs that go with those defs
    let def_map = build_reexport_def_map(srv, doc, def_set);

    // Now create a map that tells us where to insert the duplicated
    // docs into the existing doc tree
    let path_map = build_reexport_path_map(srv, def_map);

    // Finally update the doc tree
    merge_reexports(doc, path_map)
}

// Hash maps are not sendable so converting them back and forth
// to association lists. Yuck.
fn to_assoc_list<K:copy, V:copy>(
    map: map::hashmap<K, V>
) -> [(K, V)] {

    let vec = [];
    map.items {|k, v|
        vec += [(k, v)];
    }
    ret vec;
}

fn from_assoc_list<K:copy, V:copy>(
    list: [(K, V)],
    new_hash: fn() -> map::hashmap<K, V>
) -> map::hashmap<K, V> {

    let map = new_hash();
    vec::iter(list) {|elt|
        let (k, v) = elt;
        map.insert(k, v);
    }
    ret map;
}

fn from_def_assoc_list<V:copy>(
    list: [(ast::def_id, V)]
) -> map::hashmap<ast::def_id, V> {
    from_assoc_list(list, bind common::new_def_hash())
}

fn from_str_assoc_list<V:copy>(
    list: [(str, V)]
) -> map::hashmap<str, V> {
    from_assoc_list(list, bind map::new_str_hash())
}

fn build_reexport_def_set(srv: astsrv::srv) -> def_set {
    let assoc_list = astsrv::exec(srv) {|ctxt|
        let def_set = common::new_def_hash();
        ctxt.exp_map.items {|_path, defs|
            for def in *defs {
                let def_id = ast_util::def_id_of_def(def);
                def_set.insert(def_id, ());
            }
        }
        to_assoc_list(def_set)
    };

    from_def_assoc_list(assoc_list)
}

fn build_reexport_def_map(
    srv: astsrv::srv,
    doc: doc::cratedoc,
    def_set: def_set
) -> def_map {

    type ctxt = {
        srv: astsrv::srv,
        def_set: def_set,
        def_map: def_map
    };

    let ctxt = {
        srv: srv,
        def_set: def_set,
        def_map: common::new_def_hash()
    };

    // FIXME: Do a parallel fold
    let fold = fold::fold({
        fold_mod: fold_mod
        with *fold::default_seq_fold(ctxt)
    });

    fold.fold_crate(fold, doc);

    ret ctxt.def_map;

    fn fold_mod(fold: fold::fold<ctxt>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        for item in doc.items {
            let def_id = ast_util::local_def(item.id());
            if fold.ctxt.def_set.contains_key(def_id) {
                fold.ctxt.def_map.insert(def_id, item);
            }
        }

        ret doc;
    }
}

fn build_reexport_path_map(srv: astsrv::srv, -def_map: def_map) -> path_map {

    // This is real unfortunate. Lots of copying going on here
    let def_assoc_list = to_assoc_list(def_map);
    #debug("def_map: %?", def_assoc_list);

    let assoc_list = astsrv::exec(srv) {|ctxt|

        let def_map = from_def_assoc_list(def_assoc_list);
        let path_map = map::new_str_hash();

        ctxt.exp_map.items {|path, defs|

            let path = str::split_str(path, "::");
            let modpath = str::connect(vec::init(path), "::");
            let name = option::get(vec::last(path));

            let reexportdocs = [];
            for def in *defs {
                let def_id = ast_util::def_id_of_def(def);
                alt def_map.find(def_id) {
                  some(itemtag) {
                    reexportdocs += [(name, itemtag)];
                  }
                  none { }
                }
            }

            if vec::is_not_empty(reexportdocs) {
                let prevdocs = alt path_map.find(modpath) {
                  some(docs) { docs }
                  none { [] }
                };
                let reexportdocs = prevdocs + reexportdocs;
                path_map.insert(modpath, reexportdocs);
                #debug("path_map entry: %? - %?",
                       modpath, (name, reexportdocs));
            }
        }

        to_assoc_list(path_map)
    };

    from_str_assoc_list(assoc_list)
}

fn merge_reexports(
    doc: doc::cratedoc,
    path_map: path_map
) -> doc::cratedoc {

    let fold = fold::fold({
        fold_mod: fold_mod
        with *fold::default_seq_fold(path_map)
    });

    ret fold.fold_crate(fold, doc);

    fn fold_mod(fold: fold::fold<path_map>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        let is_topmod = doc.id() == rustc::syntax::ast::crate_node_id;

        // In the case of the top mod, it really doesn't have a name;
        // the name we have here is actually the crate name
        let path = if is_topmod {
            doc.path()
        } else {
            doc.path() + [doc.name()]
        };

        let new_items = get_new_items(path, fold.ctxt);
        #debug("merging into %?: %?", path, new_items);

        {
            items: (doc.items + new_items)
            with doc
        }
    }

    fn get_new_items(path: [str], path_map: path_map) -> [doc::itemtag] {
        #debug("looking for reexports in path %?", path);
        alt path_map.find(str::connect(path, "::")) {
          some(name_docs) {
            vec::foldl([], name_docs) {|v, name_doc|
                let (name, doc) = name_doc;
                v + [reexport_doc(doc, name)]
            }
          }
          none { [] }
        }
    }

    fn reexport_doc(doc: doc::itemtag, name: str) -> doc::itemtag {
        alt doc {
          doc::modtag(doc @ {item, _}) {
            doc::modtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::nmodtag(_) { fail }
          doc::consttag(doc @ {item, _}) {
            doc::consttag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::fntag(doc @ {item, _}) {
            doc::fntag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::enumtag(doc @ {item, _}) {
            doc::enumtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::restag(doc @ {item, _}) {
            doc::restag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::ifacetag(doc @ {item, _}) {
            doc::ifacetag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::impltag(doc @ {item, _}) {
            doc::impltag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::tytag(doc @ {item, _}) {
            doc::tytag({
                item: reexport(item, name)
                with doc
            })
          }
        }
    }

    fn reexport(doc: doc::itemdoc, name: str) -> doc::itemdoc {
        {
            name: name,
            reexport: true
            with doc
        }
    }
}

#[test]
fn should_duplicate_reexported_items() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.topmod.mods()[1].fns()[0].name() == "b";
}

#[test]
fn should_mark_reepxorts_as_such() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.topmod.mods()[1].fns()[0].item.reexport == true;
}

#[test]
fn should_duplicate_multiple_reexported_items() {
    let source = "mod a { \
                  export b; export c; \
                  fn b() { } fn c() { } \
                  } \
                  mod d { \
                  import a::b; import a::c; \
                  export b; export c; \
                  }";
    astsrv::from_str(source) {|srv|
        let doc = extract::from_srv(srv, "");
        let doc = path_pass::mk_pass()(srv, doc);
        let doc = run(srv, doc);
        // Reexports may not be in any specific order
        let doc = sort_item_name_pass::mk_pass()(srv, doc);
        assert doc.topmod.mods()[1].fns()[0].name() == "b";
        assert doc.topmod.mods()[1].fns()[1].name() == "c";
    }
}

#[test]
fn should_rename_items_reexported_with_different_names() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import x = a::b; export x; }";
    let doc = test::mk_doc(source);
    assert doc.topmod.mods()[1].fns()[0].name() == "x";
}

#[test]
fn should_reexport_in_topmod() {
    fn mk_doc(source: str) -> doc::cratedoc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "core");
            let doc = path_pass::mk_pass()(srv, doc);
            run(srv, doc)
        }
    }
    let source = "import option::{some, none}; \
                  import option = option::t; \
                  export option, some, none; \
                  mod option { \
                  enum t { some, none } \
                  }";
    let doc = mk_doc(source);
    assert doc.topmod.enums()[0].name() == "option";
}

#[test]
fn should_not_reexport_multiple_times() {
    let source = "import option = option::t; \
                  export option; \
                  export option; \
                  mod option { \
                  enum t { none, some } \
                  }";
    let doc = test::mk_doc(source);
    assert vec::len(doc.topmod.enums()) == 1u;
}

#[cfg(test)]
mod test {
    fn mk_doc(source: str) -> doc::cratedoc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "");
            let doc = path_pass::mk_pass()(srv, doc);
            run(srv, doc)
        }
    }
}
<commit_msg>rustdoc: Reexport native functions<commit_after>#[doc = "Finds docs for reexported items and duplicates them"];

import std::map;
import rustc::syntax::ast;
import rustc::syntax::ast_util;
import rustc::util::common;

export mk_pass;

fn mk_pass() -> pass {
    run
}

type def_set = map::set<ast::def_id>;
type def_map = map::hashmap<ast::def_id, doc::itemtag>;
type path_map = map::hashmap<str, [(str, doc::itemtag)]>;

fn run(srv: astsrv::srv, doc: doc::cratedoc) -> doc::cratedoc {

    // First gather the set of defs that are used as reexports
    let def_set = build_reexport_def_set(srv);

    // Now find the docs that go with those defs
    let def_map = build_reexport_def_map(srv, doc, def_set);

    // Now create a map that tells us where to insert the duplicated
    // docs into the existing doc tree
    let path_map = build_reexport_path_map(srv, def_map);

    // Finally update the doc tree
    merge_reexports(doc, path_map)
}

// Hash maps are not sendable so converting them back and forth
// to association lists. Yuck.
fn to_assoc_list<K:copy, V:copy>(
    map: map::hashmap<K, V>
) -> [(K, V)] {

    let vec = [];
    map.items {|k, v|
        vec += [(k, v)];
    }
    ret vec;
}

fn from_assoc_list<K:copy, V:copy>(
    list: [(K, V)],
    new_hash: fn() -> map::hashmap<K, V>
) -> map::hashmap<K, V> {

    let map = new_hash();
    vec::iter(list) {|elt|
        let (k, v) = elt;
        map.insert(k, v);
    }
    ret map;
}

fn from_def_assoc_list<V:copy>(
    list: [(ast::def_id, V)]
) -> map::hashmap<ast::def_id, V> {
    from_assoc_list(list, bind common::new_def_hash())
}

fn from_str_assoc_list<V:copy>(
    list: [(str, V)]
) -> map::hashmap<str, V> {
    from_assoc_list(list, bind map::new_str_hash())
}

fn build_reexport_def_set(srv: astsrv::srv) -> def_set {
    let assoc_list = astsrv::exec(srv) {|ctxt|
        let def_set = common::new_def_hash();
        ctxt.exp_map.items {|_path, defs|
            for def in *defs {
                let def_id = ast_util::def_id_of_def(def);
                def_set.insert(def_id, ());
            }
        }
        to_assoc_list(def_set)
    };

    from_def_assoc_list(assoc_list)
}

fn build_reexport_def_map(
    srv: astsrv::srv,
    doc: doc::cratedoc,
    def_set: def_set
) -> def_map {

    type ctxt = {
        srv: astsrv::srv,
        def_set: def_set,
        def_map: def_map
    };

    let ctxt = {
        srv: srv,
        def_set: def_set,
        def_map: common::new_def_hash()
    };

    // FIXME: Do a parallel fold
    let fold = fold::fold({
        fold_mod: fold_mod,
        fold_nmod: fold_nmod
        with *fold::default_seq_fold(ctxt)
    });

    fold.fold_crate(fold, doc);

    ret ctxt.def_map;

    fn fold_mod(fold: fold::fold<ctxt>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        for item in doc.items {
            let def_id = ast_util::local_def(item.id());
            if fold.ctxt.def_set.contains_key(def_id) {
                fold.ctxt.def_map.insert(def_id, item);
            }
        }

        ret doc;
    }

    fn fold_nmod(fold: fold::fold<ctxt>, doc: doc::nmoddoc) -> doc::nmoddoc {
        let doc = fold::default_seq_fold_nmod(fold, doc);

        for fndoc in doc.fns {
            let def_id = ast_util::local_def(fndoc.id());
            if fold.ctxt.def_set.contains_key(def_id) {
                fold.ctxt.def_map.insert(def_id, doc::fntag(fndoc));
            }
        }

        ret doc;
    }
}

fn build_reexport_path_map(srv: astsrv::srv, -def_map: def_map) -> path_map {

    // This is real unfortunate. Lots of copying going on here
    let def_assoc_list = to_assoc_list(def_map);
    #debug("def_map: %?", def_assoc_list);

    let assoc_list = astsrv::exec(srv) {|ctxt|

        let def_map = from_def_assoc_list(def_assoc_list);
        let path_map = map::new_str_hash();

        ctxt.exp_map.items {|path, defs|

            let path = str::split_str(path, "::");
            let modpath = str::connect(vec::init(path), "::");
            let name = option::get(vec::last(path));

            let reexportdocs = [];
            for def in *defs {
                let def_id = ast_util::def_id_of_def(def);
                alt def_map.find(def_id) {
                  some(itemtag) {
                    reexportdocs += [(name, itemtag)];
                  }
                  none { }
                }
            }

            if vec::is_not_empty(reexportdocs) {
                let prevdocs = alt path_map.find(modpath) {
                  some(docs) { docs }
                  none { [] }
                };
                let reexportdocs = prevdocs + reexportdocs;
                path_map.insert(modpath, reexportdocs);
                #debug("path_map entry: %? - %?",
                       modpath, (name, reexportdocs));
            }
        }

        to_assoc_list(path_map)
    };

    from_str_assoc_list(assoc_list)
}

fn merge_reexports(
    doc: doc::cratedoc,
    path_map: path_map
) -> doc::cratedoc {

    let fold = fold::fold({
        fold_mod: fold_mod
        with *fold::default_seq_fold(path_map)
    });

    ret fold.fold_crate(fold, doc);

    fn fold_mod(fold: fold::fold<path_map>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        let is_topmod = doc.id() == rustc::syntax::ast::crate_node_id;

        // In the case of the top mod, it really doesn't have a name;
        // the name we have here is actually the crate name
        let path = if is_topmod {
            doc.path()
        } else {
            doc.path() + [doc.name()]
        };

        let new_items = get_new_items(path, fold.ctxt);
        #debug("merging into %?: %?", path, new_items);

        {
            items: (doc.items + new_items)
            with doc
        }
    }

    fn get_new_items(path: [str], path_map: path_map) -> [doc::itemtag] {
        #debug("looking for reexports in path %?", path);
        alt path_map.find(str::connect(path, "::")) {
          some(name_docs) {
            vec::foldl([], name_docs) {|v, name_doc|
                let (name, doc) = name_doc;
                v + [reexport_doc(doc, name)]
            }
          }
          none { [] }
        }
    }

    fn reexport_doc(doc: doc::itemtag, name: str) -> doc::itemtag {
        alt doc {
          doc::modtag(doc @ {item, _}) {
            doc::modtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::nmodtag(_) { fail }
          doc::consttag(doc @ {item, _}) {
            doc::consttag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::fntag(doc @ {item, _}) {
            doc::fntag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::enumtag(doc @ {item, _}) {
            doc::enumtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::restag(doc @ {item, _}) {
            doc::restag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::ifacetag(doc @ {item, _}) {
            doc::ifacetag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::impltag(doc @ {item, _}) {
            doc::impltag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::tytag(doc @ {item, _}) {
            doc::tytag({
                item: reexport(item, name)
                with doc
            })
          }
        }
    }

    fn reexport(doc: doc::itemdoc, name: str) -> doc::itemdoc {
        {
            name: name,
            reexport: true
            with doc
        }
    }
}

#[test]
fn should_duplicate_reexported_items() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.topmod.mods()[1].fns()[0].name() == "b";
}

#[test]
fn should_mark_reepxorts_as_such() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.topmod.mods()[1].fns()[0].item.reexport == true;
}

#[test]
fn should_duplicate_reexported_native_fns() {
    let source = "native mod a { fn b(); } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.topmod.mods()[0].fns()[0].name() == "b";
}

#[test]
fn should_duplicate_multiple_reexported_items() {
    let source = "mod a { \
                  export b; export c; \
                  fn b() { } fn c() { } \
                  } \
                  mod d { \
                  import a::b; import a::c; \
                  export b; export c; \
                  }";
    astsrv::from_str(source) {|srv|
        let doc = extract::from_srv(srv, "");
        let doc = path_pass::mk_pass()(srv, doc);
        let doc = run(srv, doc);
        // Reexports may not be in any specific order
        let doc = sort_item_name_pass::mk_pass()(srv, doc);
        assert doc.topmod.mods()[1].fns()[0].name() == "b";
        assert doc.topmod.mods()[1].fns()[1].name() == "c";
    }
}

#[test]
fn should_rename_items_reexported_with_different_names() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import x = a::b; export x; }";
    let doc = test::mk_doc(source);
    assert doc.topmod.mods()[1].fns()[0].name() == "x";
}

#[test]
fn should_reexport_in_topmod() {
    fn mk_doc(source: str) -> doc::cratedoc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "core");
            let doc = path_pass::mk_pass()(srv, doc);
            run(srv, doc)
        }
    }
    let source = "import option::{some, none}; \
                  import option = option::t; \
                  export option, some, none; \
                  mod option { \
                  enum t { some, none } \
                  }";
    let doc = mk_doc(source);
    assert doc.topmod.enums()[0].name() == "option";
}

#[test]
fn should_not_reexport_multiple_times() {
    let source = "import option = option::t; \
                  export option; \
                  export option; \
                  mod option { \
                  enum t { none, some } \
                  }";
    let doc = test::mk_doc(source);
    assert vec::len(doc.topmod.enums()) == 1u;
}

#[cfg(test)]
mod test {
    fn mk_doc(source: str) -> doc::cratedoc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "");
            let doc = path_pass::mk_pass()(srv, doc);
            run(srv, doc)
        }
    }
}
<|endoftext|>