<commit_before>use regex::Regex;
use reqwest;
use serde_json;

use web::data::SolrTermSummary;

pub struct Search {
    solr_url: String,
}

#[derive(Deserialize, Debug)]
struct SolrResponse {
    pub docs: Vec<SolrTermSummary>,
}

#[derive(Deserialize, Debug)]
struct SolrResponseContainer {
    pub response: SolrResponse,
}

impl Search {
    pub fn new(solr_url: String) -> Search {
        Search {
            solr_url: solr_url,
        }
    }

    fn get_query_part(&self, words: &Vec<String>, weight: f32) -> String {
        let mut ret = String::new();

        let words_length = words.len();

        for (i, word) in words.iter().enumerate() {
            if i == words_length - 1 {
                ret += &format!("{}^{} {}*", word, weight, word);
            } else {
                ret += &format!("{}^{} {}~0.8 ", word, weight, word);
            }
        }

        ret
    }

    pub fn term_complete(&self, cv_name: &str, q: &str)
                         -> Result<Vec<SolrTermSummary>, reqwest::Error>
    {
        let mut terms_url =
            self.solr_url.to_owned() + "/terms/select?wt=json&q=";

        let termid_re_string = r"(?P<prefix>[\w_]+):(?P<accession>\d+)";
        let termid_re = Regex::new(termid_re_string).unwrap();

        let parent_re_string = r"^\[".to_owned() + termid_re_string + r"\]$";
        let parent_re = Regex::new(&parent_re_string).unwrap();

        let maybe_captures = parent_re.captures(cv_name);

        if let Some(captures) = maybe_captures {
            let prefix = captures.name("prefix").unwrap();
            let accession = captures.name("accession").unwrap();
            terms_url += &format!("(interesting_parents:{}\\:{} OR id:{}\\:{})",
                                  prefix, accession, prefix, accession);
        } else {
            terms_url += "cv_name:";
            terms_url += cv_name;
        }

        if let Some(captures) = termid_re.captures(q) {
            let prefix = captures.name("prefix").unwrap();
            let accession = captures.name("accession").unwrap();
            terms_url += " AND id:";
            terms_url += prefix;
            terms_url += r"\:";
            terms_url += accession;
        } else {
            terms_url += " AND (name:(";

            let clean_words: Vec<String> =
                Regex::new(r"(\w+)").unwrap().captures_iter(q)
                .map(|cap| cap.at(1).unwrap().to_owned()).collect();

            if clean_words.len() == 0 {
                return Ok(vec![]);
            }

            let clean_words_length = clean_words.len();

            for (i, word) in clean_words.iter().enumerate() {
                if i == clean_words_length - 1 {
                    terms_url += &format!("{}^4 {}*", word, word);
                } else {
                    terms_url += &format!("{}^4 {}~0.8 ", word, word);
                }
            }

            terms_url += ") OR close_synonyms:(";
            terms_url += &self.get_query_part(&clean_words, 0.5);
            terms_url += ") OR distant_synonyms:(";
            terms_url += &self.get_query_part(&clean_words, 0.1);
            terms_url += "))";
        }
        print!("{:?}\n", terms_url);

        let res = reqwest::get(&terms_url)?;

        println!("Status: {}", res.status());
        println!("Headers:\n{}", res.headers());

        match serde_json::from_reader(res) {
            Ok(solr_response_container) => {
                let container: SolrResponseContainer = solr_response_container;
                Ok(container.response.docs)
            },
            Err(err) => {
                panic!(format!("{:?}", err));
            }
        }
    }
}
<commit_msg>Fix fuzzy search of last word in completion string<commit_after>use regex::Regex;
use reqwest;
use serde_json;

use web::data::SolrTermSummary;

pub struct Search {
    solr_url: String,
}

#[derive(Deserialize, Debug)]
struct SolrResponse {
    pub docs: Vec<SolrTermSummary>,
}

#[derive(Deserialize, Debug)]
struct SolrResponseContainer {
    pub response: SolrResponse,
}

impl Search {
    pub fn new(solr_url: String) -> Search {
        Search {
            solr_url: solr_url,
        }
    }

    fn get_query_part(&self, words: &Vec<String>, weight: f32) -> String {
        let mut ret = String::new();

        let words_length = words.len();

        for (i, word) in words.iter().enumerate() {
            if i == words_length - 1 {
                ret += &format!("{}^{} {}*", word, weight, word);
            } else {
                ret += &format!("{}^{} {}~0.8 ", word, weight, word);
            }
        }

        ret
    }

    pub fn term_complete(&self, cv_name: &str, q: &str)
                         -> Result<Vec<SolrTermSummary>, reqwest::Error>
    {
        let mut terms_url =
            self.solr_url.to_owned() + "/terms/select?wt=json&q=";

        let termid_re_string = r"(?P<prefix>[\w_]+):(?P<accession>\d+)";
        let termid_re = Regex::new(termid_re_string).unwrap();

        let parent_re_string = r"^\[".to_owned() + termid_re_string + r"\]$";
        let parent_re = Regex::new(&parent_re_string).unwrap();

        let maybe_captures = parent_re.captures(cv_name);

        if let Some(captures) = maybe_captures {
            let prefix = captures.name("prefix").unwrap();
            let accession = captures.name("accession").unwrap();
            terms_url += &format!("(interesting_parents:{}\\:{} OR id:{}\\:{})",
                                  prefix, accession, prefix, accession);
        } else {
            terms_url += "cv_name:";
            terms_url += cv_name;
        }

        if let Some(captures) = termid_re.captures(q) {
            let prefix = captures.name("prefix").unwrap();
            let accession = captures.name("accession").unwrap();
            terms_url += " AND id:";
            terms_url += prefix;
            terms_url += r"\:";
            terms_url += accession;
        } else {
            terms_url += " AND (name:(";

            let clean_words: Vec<String> =
                Regex::new(r"(\w+)").unwrap().captures_iter(q)
                .map(|cap| cap.at(1).unwrap().to_owned()).collect();

            if clean_words.len() == 0 {
                return Ok(vec![]);
            }

            let clean_words_length = clean_words.len();

            for (i, word) in clean_words.iter().enumerate() {
                if i == clean_words_length - 1 {
                    terms_url += &format!("{}^4 {}~0.8 {}*", word, word, word);
                } else {
                    terms_url += &format!("{}^4 {}~0.8 ", word, word);
                }
            }

            terms_url += ") OR close_synonyms:(";
            terms_url += &self.get_query_part(&clean_words, 0.5);
            terms_url += ") OR distant_synonyms:(";
            terms_url += &self.get_query_part(&clean_words, 0.1);
            terms_url += "))";
        }
        print!("{:?}\n", terms_url);

        let res = reqwest::get(&terms_url)?;

        println!("Status: {}", res.status());
        println!("Headers:\n{}", res.headers());

        match serde_json::from_reader(res) {
            Ok(solr_response_container) => {
                let container: SolrResponseContainer = solr_response_container;
                Ok(container.response.docs)
            },
            Err(err) => {
                panic!(format!("{:?}", err));
            }
        }
    }
}
<|endoftext|>


<commit_before>/// Solve least square problem `|b - Ax|` with multi-column `b`
use approx::AbsDiffEq;
use ndarray::*;
use ndarray_linalg::*;

/// A is square. `x = A^{-1} b`, `|b - Ax| = 0`
fn test_exact<T: Scalar + Lapack>(a: Array2<T>, b: Array2<T>) {
    assert_eq!(a.layout().unwrap().size(), (3, 3));
    assert_eq!(b.layout().unwrap().size(), (3, 2));

    let result = a.least_squares(&b).unwrap();
    dbg!(&result);
    // unpack result
    let x: Array2<T> = result.solution;
    let residual_l2_square: Array1<T::Real> = result.residual_sum_of_squares.unwrap();

    // must be full-rank
    assert_eq!(result.rank, 3);

    // |b - Ax| == 0
    for &residual in &residual_l2_square {
        assert!(residual < T::real(1.0e-4));
    }

    // b == Ax
    let ax = a.dot(&x);
    assert_close_l2!(&b, &ax, T::real(1.0e-4));
}

macro_rules! impl_exact {
    ($scalar:ty) => {
        paste::item! {
            #[test]
            fn [<least_squares_ $scalar _exact_ac_bc>]() {
                let a: Array2<$scalar> = random((3, 3));
                let b: Array2<$scalar> = random((3, 2));
                test_exact(a, b)
            }

            /* Unsupported currently. See https://github.com/rust-ndarray/ndarray-linalg/issues/234

            #[test]
            fn [<least_squares_ $scalar _exact_ac_bf>]() {
                let a: Array2<$scalar> = random((3, 3));
                let b: Array2<$scalar> = random((3, 2).f());
                test_exact(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _exact_af_bc>]() {
                let a: Array2<$scalar> = random((3, 3).f());
                let b: Array2<$scalar> = random((3, 2));
                test_exact(a, b)
            }

            */

            #[test]
            fn [<least_squares_ $scalar _exact_af_bf>]() {
                let a: Array2<$scalar> = random((3, 3).f());
                let b: Array2<$scalar> = random((3, 2).f());
                test_exact(a, b)
            }
        }
    };
}

impl_exact!(f32);
impl_exact!(f64);
impl_exact!(c32);
impl_exact!(c64);

/// #column < #row case.
/// Linear problem is overdetermined, `|b - Ax| > 0`.
fn test_overdetermined<T: Scalar + Lapack>(a: Array2<T>, bs: Array2<T>)
where
    T::Real: AbsDiffEq<Epsilon = T::Real>,
{
    assert_eq!(a.layout().unwrap().size(), (4, 3));
    assert_eq!(bs.layout().unwrap().size(), (4, 2));

    let result = a.least_squares(&bs).unwrap();
    // unpack result
    let xs = result.solution;
    let residual_l2_square = result.residual_sum_of_squares.unwrap();

    // Must be full-rank
    assert_eq!(result.rank, 3);

    for j in 0..2 {
        let b = bs.index_axis(Axis(1), j);
        let x = xs.index_axis(Axis(1), j);
        let residual = &b - &a.dot(&x);
        let residual_l2_sq = residual_l2_square[j];
        assert!(residual_l2_sq.abs_diff_eq(&residual.norm_l2().powi(2), T::real(1.0e-4)));

        // `|residual| < |b|`
        assert!(residual.norm_l2() < b.norm_l2());
    }
}

macro_rules! impl_overdetermined {
    ($scalar:ty) => {
        paste::item! {
            #[test]
            fn [<least_squares_ $scalar _overdetermined_ac_bc>]() {
                let a: Array2<$scalar> = random((4, 3));
                let b: Array2<$scalar> = random((4, 2));
                test_overdetermined(a, b)
            }

            /* Unsupported currently. See https://github.com/rust-ndarray/ndarray-linalg/issues/234

            #[test]
            fn [<least_squares_ $scalar _overdetermined_af_bc>]() {
                let a: Array2<$scalar> = random((4, 3).f());
                let b: Array2<$scalar> = random((4, 2));
                test_overdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _overdetermined_ac_bf>]() {
                let a: Array2<$scalar> = random((4, 3));
                let b: Array2<$scalar> = random((4, 2).f());
                test_overdetermined(a, b)
            }

            */

            #[test]
            fn [<least_squares_ $scalar _overdetermined_af_bf>]() {
                let a: Array2<$scalar> = random((4, 3).f());
                let b: Array2<$scalar> = random((4, 2).f());
                test_overdetermined(a, b)
            }
        }
    };
}

impl_overdetermined!(f32);
impl_overdetermined!(f64);
impl_overdetermined!(c32);
impl_overdetermined!(c64);

/// #column > #row case.
/// Linear problem is underdetermined, `|b - Ax| = 0` and `x` is not unique
fn test_underdetermined<T: Scalar + Lapack>(a: Array2<T>, b: Array2<T>) {
    assert_eq!(a.layout().unwrap().size(), (3, 4));
    assert_eq!(b.layout().unwrap().size(), (3, 2));

    let result = a.least_squares(&b).unwrap();
    assert_eq!(result.rank, 3);
    assert!(result.residual_sum_of_squares.is_none());

    // b == Ax
    let x = result.solution;
    let ax = a.dot(&x);
    assert_close_l2!(&b, &ax, T::real(1.0e-4));
}

macro_rules! impl_underdetermined {
    ($scalar:ty) => {
        paste::item! {
            #[test]
            fn [<least_squares_ $scalar _underdetermined_ac_bc>]() {
                let a: Array2<$scalar> = random((3, 4));
                let b: Array2<$scalar> = random((3, 2));
                test_underdetermined(a, b)
            }

            /* Unsupported currently. See https://github.com/rust-ndarray/ndarray-linalg/issues/234

            #[test]
            fn [<least_squares_ $scalar _underdetermined_af_bc>]() {
                let a: Array2<$scalar> = random((3, 4).f());
                let b: Array2<$scalar> = random((3, 2));
                test_underdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _underdetermined_ac_bf>]() {
                let a: Array2<$scalar> = random((3, 4));
                let b: Array2<$scalar> = random((3, 2).f());
                test_underdetermined(a, b)
            }

            */

            #[test]
            fn [<least_squares_ $scalar _underdetermined_af_bf>]() {
                let a: Array2<$scalar> = random((3, 4).f());
                let b: Array2<$scalar> = random((3, 2).f());
                test_underdetermined(a, b)
            }
        }
    };
}

impl_underdetermined!(f32);
impl_underdetermined!(f64);
impl_underdetermined!(c32);
impl_underdetermined!(c64);
<commit_msg>Enable tests for C/F mixed cases #234<commit_after>/// Solve least square problem `|b - Ax|` with multi-column `b`
use approx::AbsDiffEq;
use ndarray::*;
use ndarray_linalg::*;

/// A is square. `x = A^{-1} b`, `|b - Ax| = 0`
fn test_exact<T: Scalar + Lapack>(a: Array2<T>, b: Array2<T>) {
    assert_eq!(a.layout().unwrap().size(), (3, 3));
    assert_eq!(b.layout().unwrap().size(), (3, 2));

    let result = a.least_squares(&b).unwrap();
    dbg!(&result);
    // unpack result
    let x: Array2<T> = result.solution;
    let residual_l2_square: Array1<T::Real> = result.residual_sum_of_squares.unwrap();

    // must be full-rank
    assert_eq!(result.rank, 3);

    // |b - Ax| == 0
    for &residual in &residual_l2_square {
        assert!(residual < T::real(1.0e-4));
    }

    // b == Ax
    let ax = a.dot(&x);
    assert_close_l2!(&b, &ax, T::real(1.0e-4));
}

macro_rules! impl_exact {
    ($scalar:ty) => {
        paste::item! {
            #[test]
            fn [<least_squares_ $scalar _exact_ac_bc>]() {
                let a: Array2<$scalar> = random((3, 3));
                let b: Array2<$scalar> = random((3, 2));
                test_exact(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _exact_ac_bf>]() {
                let a: Array2<$scalar> = random((3, 3));
                let b: Array2<$scalar> = random((3, 2).f());
                test_exact(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _exact_af_bc>]() {
                let a: Array2<$scalar> = random((3, 3).f());
                let b: Array2<$scalar> = random((3, 2));
                test_exact(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _exact_af_bf>]() {
                let a: Array2<$scalar> = random((3, 3).f());
                let b: Array2<$scalar> = random((3, 2).f());
                test_exact(a, b)
            }
        }
    };
}

impl_exact!(f32);
impl_exact!(f64);
impl_exact!(c32);
impl_exact!(c64);

/// #column < #row case.
/// Linear problem is overdetermined, `|b - Ax| > 0`.
fn test_overdetermined<T: Scalar + Lapack>(a: Array2<T>, bs: Array2<T>)
where
    T::Real: AbsDiffEq<Epsilon = T::Real>,
{
    assert_eq!(a.layout().unwrap().size(), (4, 3));
    assert_eq!(bs.layout().unwrap().size(), (4, 2));

    let result = a.least_squares(&bs).unwrap();
    // unpack result
    let xs = result.solution;
    let residual_l2_square = result.residual_sum_of_squares.unwrap();

    // Must be full-rank
    assert_eq!(result.rank, 3);

    for j in 0..2 {
        let b = bs.index_axis(Axis(1), j);
        let x = xs.index_axis(Axis(1), j);
        let residual = &b - &a.dot(&x);
        let residual_l2_sq = residual_l2_square[j];
        assert!(residual_l2_sq.abs_diff_eq(&residual.norm_l2().powi(2), T::real(1.0e-4)));

        // `|residual| < |b|`
        assert!(residual.norm_l2() < b.norm_l2());
    }
}

macro_rules! impl_overdetermined {
    ($scalar:ty) => {
        paste::item! {
            #[test]
            fn [<least_squares_ $scalar _overdetermined_ac_bc>]() {
                let a: Array2<$scalar> = random((4, 3));
                let b: Array2<$scalar> = random((4, 2));
                test_overdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _overdetermined_af_bc>]() {
                let a: Array2<$scalar> = random((4, 3).f());
                let b: Array2<$scalar> = random((4, 2));
                test_overdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _overdetermined_ac_bf>]() {
                let a: Array2<$scalar> = random((4, 3));
                let b: Array2<$scalar> = random((4, 2).f());
                test_overdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _overdetermined_af_bf>]() {
                let a: Array2<$scalar> = random((4, 3).f());
                let b: Array2<$scalar> = random((4, 2).f());
                test_overdetermined(a, b)
            }
        }
    };
}

impl_overdetermined!(f32);
impl_overdetermined!(f64);
impl_overdetermined!(c32);
impl_overdetermined!(c64);

/// #column > #row case.
/// Linear problem is underdetermined, `|b - Ax| = 0` and `x` is not unique
fn test_underdetermined<T: Scalar + Lapack>(a: Array2<T>, b: Array2<T>) {
    assert_eq!(a.layout().unwrap().size(), (3, 4));
    assert_eq!(b.layout().unwrap().size(), (3, 2));

    let result = a.least_squares(&b).unwrap();
    assert_eq!(result.rank, 3);
    assert!(result.residual_sum_of_squares.is_none());

    // b == Ax
    let x = result.solution;
    let ax = a.dot(&x);
    assert_close_l2!(&b, &ax, T::real(1.0e-4));
}

macro_rules! impl_underdetermined {
    ($scalar:ty) => {
        paste::item! {
            #[test]
            fn [<least_squares_ $scalar _underdetermined_ac_bc>]() {
                let a: Array2<$scalar> = random((3, 4));
                let b: Array2<$scalar> = random((3, 2));
                test_underdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _underdetermined_af_bc>]() {
                let a: Array2<$scalar> = random((3, 4).f());
                let b: Array2<$scalar> = random((3, 2));
                test_underdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _underdetermined_ac_bf>]() {
                let a: Array2<$scalar> = random((3, 4));
                let b: Array2<$scalar> = random((3, 2).f());
                test_underdetermined(a, b)
            }

            #[test]
            fn [<least_squares_ $scalar _underdetermined_af_bf>]() {
                let a: Array2<$scalar> = random((3, 4).f());
                let b: Array2<$scalar> = random((3, 2).f());
                test_underdetermined(a, b)
            }
        }
    };
}

impl_underdetermined!(f32);
impl_underdetermined!(f64);
impl_underdetermined!(c32);
impl_underdetermined!(c64);
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

async fn acl_lookup<Remote>(
    acl: &Option<AccessControl>,
    local: &upstream::LocalUpstream,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut, decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut, if remote_response.is_none() => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider, if !use_remote => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
        }
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let remote_upstream = remote_upstream.clone();
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), context.local_dns(), remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Add else to tokio::select<commit_after>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

async fn acl_lookup<Remote>(
    acl: &Option<AccessControl>,
    local: &upstream::LocalUpstream,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut, decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut, if remote_response.is_none() => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider, if !use_remote => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
            else => unreachable!(),
        }
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let remote_upstream = remote_upstream.clone();
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), context.local_dns(), remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>//! Main interface for a stream of events the Supervisor can send out
//! in the course of its operations.
//!
//! Currently, the Supervisor is able to send events to a [NATS
//! Streaming][1] server. The `init_stream` function must be called
//! before sending events to initialize the publishing thread in the
//! background. Thereafter, you can pass "event" structs to the
//! `event` function, which will publish the event to the stream.
//!
//! All events are published under the "habitat" subject.
//!
//! [1]:https://github.com/nats-io/nats-streaming-server

mod error;

#[cfg(feature = "nitox_stream")]
#[path = "event/nitox.rs"]
mod stream_impl;

// This is the default, but setting it as a default in Cargo.toml
// makes it trickier to use nitox instead.
#[cfg(any(feature = "ratsio_stream", not(feature = "nitox_stream")))]
#[path = "event/ratsio.rs"]
mod stream_impl;
mod types;

pub(crate) use self::types::ServiceMetadata;
use self::types::{EventMessage,
                  EventMetadata,
                  HealthCheckEvent,
                  ServiceStartedEvent,
                  ServiceStoppedEvent};
use crate::manager::{service::{HealthCheckResult,
                               Service},
                     sys::Sys};
use clap::ArgMatches;
pub use error::{Error,
                Result};
use futures::sync::mpsc::UnboundedSender;
use habitat_common::types::{AutomateAuthToken,
                            EventStreamMetadata};
use habitat_core::env::Config as EnvConfig;
use state::Container;
use std::{net::SocketAddr,
          num::ParseIntError,
          str::FromStr,
          sync::Once,
          time::Duration};

static INIT: Once = Once::new();
lazy_static! {
    // TODO (CM): When const fn support lands in stable, we can ditch
    // this lazy_static call.

    /// Reference to the event stream.
    static ref EVENT_STREAM: Container = Container::new();
    /// Core information that is shared between all events.
    static ref EVENT_CORE: Container = Container::new();
}

/// Starts a new thread for sending events to a NATS Streaming
/// server. Stashes the handle to the stream, as well as the core
/// event information that will be a part of all events, in a global
/// static reference for access later.
pub fn init_stream(config: EventStreamConfig, event_core: EventCore) -> Result<()> {
    // call_once can't return a Result (or anything), so we'll fake it
    // by hanging onto any error we might receive.
    let mut return_value: Result<()> = Ok(());

    INIT.call_once(|| {
            let conn_info = EventConnectionInfo::new(config.token, config.url);
            match stream_impl::init_stream(conn_info) {
                Ok(event_stream) => {
                    EVENT_STREAM.set(event_stream);
                    EVENT_CORE.set(event_core);
                }
                Err(e) => return_value = Err(e),
            }
        });

    return_value
}

/// Captures all event stream-related configuration options that would
/// be passed in by a user
#[derive(Clone, Debug)]
pub struct EventStreamConfig {
    environment: String,
    application: String,
    meta:        EventStreamMetadata,
    token:       AutomateAuthToken,
    url:         String,
}

impl<'a> From<&'a ArgMatches<'a>> for EventStreamConfig {
    fn from(m: &ArgMatches) -> Self {
        EventStreamConfig { environment: m.value_of("EVENT_STREAM_ENVIRONMENT")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            application: m.value_of("EVENT_STREAM_APPLICATION")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            meta:        EventStreamMetadata::from(m),
                            token:       AutomateAuthToken::from(m),
                            url:         m.value_of("EVENT_STREAM_URL")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"), }
    }
}

/// All the information needed to establish a connection to a NATS
/// Streaming server.
// TODO: This will change as we firm up what the interaction between
// Habitat and A2 looks like.
pub struct EventConnectionInfo {
    pub name:        String,
    pub verbose:     bool,
    pub cluster_uri: String,
    pub cluster_id:  String,
    pub auth_token:  AutomateAuthToken,
}

impl EventConnectionInfo {
    pub fn new(auth_token: AutomateAuthToken, cluster_uri: String) -> Self {
        EventConnectionInfo { name: "hab_client".to_string(),
                              verbose: true,
                              cluster_uri,
                              cluster_id: "event-service".to_string(),
                              auth_token }
    }
}

/// A collection of data that will be present in all events. Rather
/// than baking this into the structure of each event, we represent it
/// once and merge the information into the final rendered form of the
/// event.
///
/// This prevents us from having to thread information throughout the
/// system, just to get it to the places where the events are
/// generated (e.g., not all code has direct access to the
/// Supervisor's ID).
#[derive(Clone, Debug)]
pub struct EventCore {
    /// The unique identifier of the Supervisor sending the event.
    supervisor_id: String,
    ip_address: SocketAddr,
    // TODO (CM): could add application and environment to the meta
    // map directly... hrmm
    application: String,
    environment: String,
    meta:        EventStreamMetadata,
}

impl EventCore {
    pub fn new(config: &EventStreamConfig, sys: &Sys) -> Self {
        EventCore { supervisor_id: sys.member_id.clone(),
                    ip_address:    sys.gossip_listen(),
                    environment:   config.environment.clone(),
                    application:   config.application.clone(),
                    meta:          config.meta.clone(), }
    }
}

/// Send an event for the start of a Service.
pub fn service_started(service: &Service) {
    if stream_initialized() {
        publish(ServiceStartedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

/// Send an event for the stop of a Service.
pub fn service_stopped(service: &Service) {
    if stream_initialized() {
        publish(ServiceStoppedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

// Takes metadata directly, rather than a `&Service` like other event
// functions, because of how the asynchronous health checking
// currently works. Revisit when async/await + Pin is all stabilized.
pub fn health_check(metadata: ServiceMetadata,
                    check_result: HealthCheckResult,
                    duration: Option<Duration>) {
    if stream_initialized() {
        let check_result: types::HealthCheck = check_result.into();
        publish(HealthCheckEvent { service_metadata: Some(metadata),
                                   event_metadata:   None,
                                   result:           i32::from(check_result),
                                   duration:         duration.map(Duration::into), });
    }
}

////////////////////////////////////////////////////////////////////////

/// Internal helper function to know whether or not to go to the trouble of
/// creating event structures. If the event stream hasn't been
/// initialized, then we shouldn't need to do anything.
fn stream_initialized() -> bool { EVENT_STREAM.try_get::<EventStream>().is_some() }

/// Publish an event. This is the main interface that client code will
/// use.
///
/// If `init_stream` has not been called already, this function will
/// be a no-op.
fn publish(mut event: impl EventMessage) {
    if let Some(e) = EVENT_STREAM.try_get::<EventStream>() {
        // TODO (CM): Yeah... this is looking pretty gross. The
        // intention is to be able to timestamp the events right as
        // they go out.
        //
        // We *could* set the time when we convert the EventCore to a
        // EventMetadata struct, but that seems odd.
        //
        // It probably means that this structure just isn't the right
        // one.
        //
        // The ugliness is at least contained, though.
        event.event_metadata(EventMetadata { timestamp:
                                                 Some(std::time::SystemTime::now().into()),
                                             ..EVENT_CORE.get::<EventCore>().to_event_metadata() });

        e.send(event.to_bytes());
    }
}

/// A lightweight handle for the event stream. All events get to the
/// event stream through this.
struct EventStream(UnboundedSender<Vec<u8>>);

impl EventStream {
    /// Queues an event to be sent out.
    fn send(&self, event: Vec<u8>) {
        trace!("About to queue an event: {:?}", event);
        if let Err(e) = self.0.unbounded_send(event) {
            error!("Failed to queue event: {:?}", e);
        }
    }
}

////////////////////////////////////////////////////////////////////////

/// How long should we for the event thread to start up before
/// abandoning it and shutting down?
#[derive(Clone, Debug)]
struct EventThreadStartupWait {
    secs: u64,
}

impl Default for EventThreadStartupWait {
    fn default() -> Self { Self { secs: 5 } }
}

impl FromStr for EventThreadStartupWait {
    type Err = ParseIntError;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> { Ok(Self { secs: s.parse()? }) }
}

impl EnvConfig for EventThreadStartupWait {
    const ENVVAR: &'static str = "HAB_EVENT_THREAD_STARTUP_WAIT_SEC";
}

impl Into<Duration> for EventThreadStartupWait {
    fn into(self) -> Duration { Duration::from_secs(self.secs) }
}
<commit_msg>Remove some unnecessary TODOs<commit_after>//! Main interface for a stream of events the Supervisor can send out
//! in the course of its operations.
//!
//! Currently, the Supervisor is able to send events to a [NATS
//! Streaming][1] server. The `init_stream` function must be called
//! before sending events to initialize the publishing thread in the
//! background. Thereafter, you can pass "event" structs to the
//! `event` function, which will publish the event to the stream.
//!
//! All events are published under the "habitat" subject.
//!
//! [1]:https://github.com/nats-io/nats-streaming-server

mod error;

#[cfg(feature = "nitox_stream")]
#[path = "event/nitox.rs"]
mod stream_impl;

// This is the default, but setting it as a default in Cargo.toml
// makes it trickier to use nitox instead.
#[cfg(any(feature = "ratsio_stream", not(feature = "nitox_stream")))]
#[path = "event/ratsio.rs"]
mod stream_impl;
mod types;

pub(crate) use self::types::ServiceMetadata;
use self::types::{EventMessage,
                  EventMetadata,
                  HealthCheckEvent,
                  ServiceStartedEvent,
                  ServiceStoppedEvent};
use crate::manager::{service::{HealthCheckResult,
                               Service},
                     sys::Sys};
use clap::ArgMatches;
pub use error::{Error,
                Result};
use futures::sync::mpsc::UnboundedSender;
use habitat_common::types::{AutomateAuthToken,
                            EventStreamMetadata};
use habitat_core::env::Config as EnvConfig;
use state::Container;
use std::{net::SocketAddr,
          num::ParseIntError,
          str::FromStr,
          sync::Once,
          time::Duration};

static INIT: Once = Once::new();
lazy_static! {
    // TODO (CM): When const fn support lands in stable, we can ditch
    // this lazy_static call.

    /// Reference to the event stream.
    static ref EVENT_STREAM: Container = Container::new();
    /// Core information that is shared between all events.
    static ref EVENT_CORE: Container = Container::new();
}

/// Starts a new thread for sending events to a NATS Streaming
/// server. Stashes the handle to the stream, as well as the core
/// event information that will be a part of all events, in a global
/// static reference for access later.
pub fn init_stream(config: EventStreamConfig, event_core: EventCore) -> Result<()> {
    // call_once can't return a Result (or anything), so we'll fake it
    // by hanging onto any error we might receive.
    let mut return_value: Result<()> = Ok(());

    INIT.call_once(|| {
            let conn_info = EventConnectionInfo::new(config.token, config.url);
            match stream_impl::init_stream(conn_info) {
                Ok(event_stream) => {
                    EVENT_STREAM.set(event_stream);
                    EVENT_CORE.set(event_core);
                }
                Err(e) => return_value = Err(e),
            }
        });

    return_value
}

/// Captures all event stream-related configuration options that would
/// be passed in by a user
#[derive(Clone, Debug)]
pub struct EventStreamConfig {
    environment: String,
    application: String,
    meta:        EventStreamMetadata,
    token:       AutomateAuthToken,
    url:         String,
}

impl<'a> From<&'a ArgMatches<'a>> for EventStreamConfig {
    fn from(m: &ArgMatches) -> Self {
        EventStreamConfig { environment: m.value_of("EVENT_STREAM_ENVIRONMENT")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            application: m.value_of("EVENT_STREAM_APPLICATION")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            meta:        EventStreamMetadata::from(m),
                            token:       AutomateAuthToken::from(m),
                            url:         m.value_of("EVENT_STREAM_URL")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"), }
    }
}

/// All the information needed to establish a connection to a NATS
/// Streaming server.
pub struct EventConnectionInfo {
    pub name:        String,
    pub verbose:     bool,
    pub cluster_uri: String,
    pub cluster_id:  String,
    pub auth_token:  AutomateAuthToken,
}

impl EventConnectionInfo {
    pub fn new(auth_token: AutomateAuthToken, cluster_uri: String) -> Self {
        EventConnectionInfo { name: "hab_client".to_string(),
                              verbose: true,
                              cluster_uri,
                              cluster_id: "event-service".to_string(),
                              auth_token }
    }
}

/// A collection of data that will be present in all events. Rather
/// than baking this into the structure of each event, we represent it
/// once and merge the information into the final rendered form of the
/// event.
///
/// This prevents us from having to thread information throughout the
/// system, just to get it to the places where the events are
/// generated (e.g., not all code has direct access to the
/// Supervisor's ID).
#[derive(Clone, Debug)]
pub struct EventCore {
    /// The unique identifier of the Supervisor sending the event.
    supervisor_id: String,
    ip_address: SocketAddr,
    application: String,
    environment: String,
    meta: EventStreamMetadata,
}

impl EventCore {
    pub fn new(config: &EventStreamConfig, sys: &Sys) -> Self {
        EventCore { supervisor_id: sys.member_id.clone(),
                    ip_address:    sys.gossip_listen(),
                    environment:   config.environment.clone(),
                    application:   config.application.clone(),
                    meta:          config.meta.clone(), }
    }
}

/// Send an event for the start of a Service.
pub fn service_started(service: &Service) {
    if stream_initialized() {
        publish(ServiceStartedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

/// Send an event for the stop of a Service.
pub fn service_stopped(service: &Service) {
    if stream_initialized() {
        publish(ServiceStoppedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

// Takes metadata directly, rather than a `&Service` like other event
// functions, because of how the asynchronous health checking
// currently works. Revisit when async/await + Pin is all stabilized.
pub fn health_check(metadata: ServiceMetadata,
                    check_result: HealthCheckResult,
                    duration: Option<Duration>) {
    if stream_initialized() {
        let check_result: types::HealthCheck = check_result.into();
        publish(HealthCheckEvent { service_metadata: Some(metadata),
                                   event_metadata:   None,
                                   result:           i32::from(check_result),
                                   duration:         duration.map(Duration::into), });
    }
}

////////////////////////////////////////////////////////////////////////

/// Internal helper function to know whether or not to go to the trouble of
/// creating event structures. If the event stream hasn't been
/// initialized, then we shouldn't need to do anything.
fn stream_initialized() -> bool { EVENT_STREAM.try_get::<EventStream>().is_some() }

/// Publish an event. This is the main interface that client code will
/// use.
///
/// If `init_stream` has not been called already, this function will
/// be a no-op.
fn publish(mut event: impl EventMessage) {
    if let Some(e) = EVENT_STREAM.try_get::<EventStream>() {
        // TODO (CM): Yeah... this is looking pretty gross. The
        // intention is to be able to timestamp the events right as
        // they go out.
        //
        // We *could* set the time when we convert the EventCore to a
        // EventMetadata struct, but that seems odd.
        //
        // It probably means that this structure just isn't the right
        // one.
        //
        // The ugliness is at least contained, though.
        event.event_metadata(EventMetadata { timestamp:
                                                 Some(std::time::SystemTime::now().into()),
                                             ..EVENT_CORE.get::<EventCore>().to_event_metadata() });

        e.send(event.to_bytes());
    }
}

/// A lightweight handle for the event stream. All events get to the
/// event stream through this.
struct EventStream(UnboundedSender<Vec<u8>>);

impl EventStream {
    /// Queues an event to be sent out.
    fn send(&self, event: Vec<u8>) {
        trace!("About to queue an event: {:?}", event);
        if let Err(e) = self.0.unbounded_send(event) {
            error!("Failed to queue event: {:?}", e);
        }
    }
}

////////////////////////////////////////////////////////////////////////

/// How long should we for the event thread to start up before
/// abandoning it and shutting down?
#[derive(Clone, Debug)]
struct EventThreadStartupWait {
    secs: u64,
}

impl Default for EventThreadStartupWait {
    fn default() -> Self { Self { secs: 5 } }
}

impl FromStr for EventThreadStartupWait {
    type Err = ParseIntError;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> { Ok(Self { secs: s.parse()? }) }
}

impl EnvConfig for EventThreadStartupWait {
    const ENVVAR: &'static str = "HAB_EVENT_THREAD_STARTUP_WAIT_SEC";
}

impl Into<Duration> for EventThreadStartupWait {
    fn into(self) -> Duration { Duration::from_secs(self.secs) }
}
<|endoftext|>


<commit_before>//! # Channel Identification via Sliding Windows
//!
//! This strategy aims to locate price channels over arbitrary time periods by analyzing the
//! minimums and maximums of prices over various time periods.  By determining new lows of a
//! macro trend and watching for new highs in smaller time periods, subtrends and the channels
//! that they make up can be located and analyzed for trade opportunities.

#![feature(test)]

extern crate futures;
extern crate algobot_util;
extern crate test;

mod window_manager;

use futures::Complete;

use algobot_util::strategies::Strategy;
use algobot_util::tick::SymbolTick;

struct SldingWindows{}

impl Strategy for SldingWindows {
    fn process(&mut self, t: SymbolTick) {

    }

    fn exit_now(&mut self, ready: Complete<()>) {

    }
}
<commit_msg>Boilerplate for strategy initialization<commit_after>//! # Channel Identification via Sliding Windows
//!
//! This strategy aims to locate price channels over arbitrary time periods by analyzing the
//! minimums and maximums of prices over various time periods.  By determining new lows of a
//! macro trend and watching for new highs in smaller time periods, subtrends and the channels
//! that they make up can be located and analyzed for trade opportunities.

#![feature(test, conservative_impl_trait)]

extern crate futures;
extern crate algobot_util;
extern crate test;

mod window_manager;

use futures::Complete;
use algobot_util::strategies::Strategy;
use algobot_util::tick::SymbolTick;
use algobot_util::transport::command_server::CommandServer;
use algobot_util::transport::query_server::QueryServer;

use window_manager::WindowManager;

struct SlidingWindows {
    wm: WindowManager
}

impl SlidingWindows {
    fn new() -> SlidingWindows {
        SlidingWindows {
            wm: WindowManager::new()
        }
    }
}

impl Strategy for SlidingWindows {
    /// This should initialize the strategy and start it running.  It should be asynchronous so
    /// the strategy doesn't block the main thread, so it's likely that the strategy should run
    /// on a separate thread.
    fn init(cs: CommandServer, qs: QueryServer) {
        let strat = SlidingWindows::new();
    }

    fn process(&mut self, t: SymbolTick) {

    }

    fn exit_now(&mut self, ready: Complete<()>) {

    }
}
<|endoftext|>


<commit_before>use std::iter::Empty;
use std::ops::Range;
use std::sync::Arc;

use command_buffer::Submission;
use device::Device;
use format::FormatDesc;
use format::FormatTy;
use image::sys::Dimensions;
use image::sys::Layout;
use image::sys::UnsafeImage;
use image::sys::UnsafeImageView;
use image::sys::Usage;
use image::traits::AccessRange;
use image::traits::Image;
use image::traits::ImageContent;
use image::traits::ImageView;
use memory::DeviceMemory;
use sync::Sharing;

use OomError;

pub struct AttachmentImage<F> {
    image: UnsafeImage,
    view: UnsafeImageView,
    memory: DeviceMemory,
    format: F,

    // Layout to use when the image is used as a framebuffer attachment.
    // Should be either "depth-stencil optimal" or "color optimal".
    attachment_layout: Layout,
}

impl<F> AttachmentImage<F> {
    pub fn new(device: &Arc<Device>, dimensions: [u32; 2], format: F)
               -> Result<Arc<AttachmentImage<F>>, OomError>
        where F: FormatDesc
    {
        let is_depth = match format.format().ty() {
            FormatTy::Depth => true,
            FormatTy::DepthStencil => true,
            FormatTy::Compressed => panic!(),
            _ => false
        };

        let usage = Usage {
            transfer_source: true,
            sampled: true,
            color_attachment: !is_depth,
            depth_stencil_attachment: is_depth,
            .. Usage::none()
        };

        let (image, mem_reqs) = unsafe {
            try!(UnsafeImage::new(device, &usage, format.format(),
                                  Dimensions::Dim2d { width: dimensions[0], height: dimensions[1] },
                                  1, 1, Sharing::Exclusive::<Empty<u32>>, false, false))
        };

        let mem_ty = {
            let device_local = device.physical_device().memory_types()
                                     .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0)
                                     .filter(|t| t.is_device_local());
            let any = device.physical_device().memory_types()
                            .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0);
            device_local.chain(any).next().unwrap()
        };

        // note: alignment doesn't need to be checked because allocating memory is guaranteed to
        //       fulfill any alignment requirement

        let mem = try!(DeviceMemory::alloc(device, &mem_ty, mem_reqs.size));
        unsafe { try!(image.bind_memory(&mem, 0 .. mem_reqs.size)); }

        let view = unsafe {
            try!(UnsafeImageView::new(&image))
        };

        Ok(Arc::new(AttachmentImage {
            image: image,
            view: view,
            memory: mem,
            format: format,
            attachment_layout: if is_depth { Layout::DepthStencilAttachmentOptimal }
                               else { Layout::ColorAttachmentOptimal },
        }))
    }

    #[inline]
    pub fn dimensions(&self) -> Dimensions {
        self.image.dimensions()
    }
}

unsafe impl<F> Image for AttachmentImage<F> {
    #[inline]
    fn inner_image(&self) -> &UnsafeImage {
        &self.image
    }

    #[inline]
    fn blocks(&self, _: Range<u32>, _: Range<u32>) -> Vec<(u32, u32)> {
        vec![(0, 0)]
    }

    #[inline]
    fn initial_layout(&self, _: (u32, u32), _: Layout) -> Layout {
        self.attachment_layout
    }

    #[inline]
    fn final_layout(&self, _: (u32, u32), _: Layout) -> Layout {
        self.attachment_layout
    }

    fn needs_fence(&self, access: &mut Iterator<Item = AccessRange>) -> Option<bool> {
        Some(false)
    }

    unsafe fn gpu_access(&self, access: &mut Iterator<Item = AccessRange>,
                         submission: &Arc<Submission>) -> Vec<Arc<Submission>>
    {
        vec![]
    }
}

unsafe impl<P, F> ImageContent<P> for AttachmentImage<F> {
    #[inline]
    fn matches_format(&self) -> bool {
        true        // FIXME:
    }
}

unsafe impl<F: 'static> ImageView for AttachmentImage<F> {
    #[inline]
    fn parent(&self) -> &Image {
        self
    }

    #[inline]
    fn parent_arc(me: &Arc<Self>) -> Arc<Image> where Self: Sized {
        me.clone() as Arc<_>
    }

    #[inline]
    fn inner_view(&self) -> &UnsafeImageView {
        &self.view
    }

    #[inline]
    fn descriptor_set_storage_image_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn descriptor_set_combined_image_sampler_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn descriptor_set_sampled_image_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn descriptor_set_input_attachment_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn identity_swizzle(&self) -> bool {
        true
    }
}
<commit_msg>Fix AttachmentImage not being usable as input attachment<commit_after>use std::iter::Empty;
use std::ops::Range;
use std::sync::Arc;

use command_buffer::Submission;
use device::Device;
use format::FormatDesc;
use format::FormatTy;
use image::sys::Dimensions;
use image::sys::Layout;
use image::sys::UnsafeImage;
use image::sys::UnsafeImageView;
use image::sys::Usage;
use image::traits::AccessRange;
use image::traits::Image;
use image::traits::ImageContent;
use image::traits::ImageView;
use memory::DeviceMemory;
use sync::Sharing;

use OomError;

pub struct AttachmentImage<F> {
    image: UnsafeImage,
    view: UnsafeImageView,
    memory: DeviceMemory,
    format: F,

    // Layout to use when the image is used as a framebuffer attachment.
    // Should be either "depth-stencil optimal" or "color optimal".
    attachment_layout: Layout,
}

impl<F> AttachmentImage<F> {
    pub fn new(device: &Arc<Device>, dimensions: [u32; 2], format: F)
               -> Result<Arc<AttachmentImage<F>>, OomError>
        where F: FormatDesc
    {
        let is_depth = match format.format().ty() {
            FormatTy::Depth => true,
            FormatTy::DepthStencil => true,
            FormatTy::Compressed => panic!(),
            _ => false
        };

        let usage = Usage {
            transfer_source: true,
            sampled: true,
            color_attachment: !is_depth,
            depth_stencil_attachment: is_depth,
            input_attachment: true,
            .. Usage::none()
        };

        let (image, mem_reqs) = unsafe {
            try!(UnsafeImage::new(device, &usage, format.format(),
                                  Dimensions::Dim2d { width: dimensions[0], height: dimensions[1] },
                                  1, 1, Sharing::Exclusive::<Empty<u32>>, false, false))
        };

        let mem_ty = {
            let device_local = device.physical_device().memory_types()
                                     .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0)
                                     .filter(|t| t.is_device_local());
            let any = device.physical_device().memory_types()
                            .filter(|t| (mem_reqs.memory_type_bits & (1 << t.id())) != 0);
            device_local.chain(any).next().unwrap()
        };

        // note: alignment doesn't need to be checked because allocating memory is guaranteed to
        //       fulfill any alignment requirement

        let mem = try!(DeviceMemory::alloc(device, &mem_ty, mem_reqs.size));
        unsafe { try!(image.bind_memory(&mem, 0 .. mem_reqs.size)); }

        let view = unsafe {
            try!(UnsafeImageView::new(&image))
        };

        Ok(Arc::new(AttachmentImage {
            image: image,
            view: view,
            memory: mem,
            format: format,
            attachment_layout: if is_depth { Layout::DepthStencilAttachmentOptimal }
                               else { Layout::ColorAttachmentOptimal },
        }))
    }

    #[inline]
    pub fn dimensions(&self) -> Dimensions {
        self.image.dimensions()
    }
}

unsafe impl<F> Image for AttachmentImage<F> {
    #[inline]
    fn inner_image(&self) -> &UnsafeImage {
        &self.image
    }

    #[inline]
    fn blocks(&self, _: Range<u32>, _: Range<u32>) -> Vec<(u32, u32)> {
        vec![(0, 0)]
    }

    #[inline]
    fn initial_layout(&self, _: (u32, u32), _: Layout) -> Layout {
        self.attachment_layout
    }

    #[inline]
    fn final_layout(&self, _: (u32, u32), _: Layout) -> Layout {
        self.attachment_layout
    }

    fn needs_fence(&self, access: &mut Iterator<Item = AccessRange>) -> Option<bool> {
        Some(false)
    }

    unsafe fn gpu_access(&self, access: &mut Iterator<Item = AccessRange>,
                         submission: &Arc<Submission>) -> Vec<Arc<Submission>>
    {
        vec![]
    }
}

unsafe impl<P, F> ImageContent<P> for AttachmentImage<F> {
    #[inline]
    fn matches_format(&self) -> bool {
        true        // FIXME:
    }
}

unsafe impl<F: 'static> ImageView for AttachmentImage<F> {
    #[inline]
    fn parent(&self) -> &Image {
        self
    }

    #[inline]
    fn parent_arc(me: &Arc<Self>) -> Arc<Image> where Self: Sized {
        me.clone() as Arc<_>
    }

    #[inline]
    fn inner_view(&self) -> &UnsafeImageView {
        &self.view
    }

    #[inline]
    fn descriptor_set_storage_image_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn descriptor_set_combined_image_sampler_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn descriptor_set_sampled_image_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn descriptor_set_input_attachment_layout(&self, _: AccessRange) -> Layout {
        Layout::ColorAttachmentOptimal
    }

    #[inline]
    fn identity_swizzle(&self) -> bool {
        true
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::iter;

use adivon::bag::Bag;

/// Weighted directed edge
#[derive(Clone, Copy)]
pub struct DirectedEdge {
    v: usize,
    w: usize,
    weight: f64
}

impl DirectedEdge {
    pub fn new(v: usize, w: usize, weight: f64) -> DirectedEdge {
        assert!(!weight.is_nan(), "weight is NaN");
        DirectedEdge {
            v: v,
            w: w,
            weight: weight
        }
    }

    #[inline]
    pub fn from(&self) -> usize {
        self.v
    }

    #[inline]
    pub fn to(&self) -> usize {
        self.w
    }

    #[inline]
    pub fn weight(&self) -> f64 {
        self.weight
    }
}

impl fmt::Debug for DirectedEdge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} -> {} {:5.2}", self.v, self.w, self.weight)
    }
}

#[test]
fn test_directed_edge() {
    let e = DirectedEdge::new(12, 24, 3.14);
    assert_eq!(format!("{:?}", e), "12 -> 24  3.14");
}

/// Edge-weighted digraph, implemented using adjacency lists
#[derive(Clone)]
pub struct EdgeWeightedDigraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<DirectedEdge>>
}

impl EdgeWeightedDigraph {
    pub fn new(v: usize) -> EdgeWeightedDigraph {
        EdgeWeightedDigraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex must be between 0 and V");
    }

    pub fn add_edge(&mut self, e: DirectedEdge) {
        let v = e.from();
        let w = e.to();
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e);
        self.e += 1;
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<DirectedEdge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
    }

    pub fn outdegree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<DirectedEdge> {
        self.adj.iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
            })
            .collect::<Vec<DirectedEdge>>()
            .into_iter()
    }

    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("digraph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.from();
            let w = e.to();
            dot.push_str(&format!("  {} -> {} [ label=\"{}\" ];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot
    }
}
<commit_msg>shortest_paths: implement Dijkstra shortest path<commit_after>use std::fmt;
use std::iter;
use std::f64;

use adivon::bag::Bag;
use adivon::stack::Stack;
use adivon::priority_queue::IndexMinPQ;

/// Weighted directed edge
#[derive(Clone, Copy)]
pub struct DirectedEdge {
    v: usize,
    w: usize,
    weight: f64
}

impl DirectedEdge {
    pub fn new(v: usize, w: usize, weight: f64) -> DirectedEdge {
        assert!(!weight.is_nan(), "weight is NaN");
        DirectedEdge {
            v: v,
            w: w,
            weight: weight
        }
    }

    #[inline]
    pub fn from(&self) -> usize {
        self.v
    }

    #[inline]
    pub fn to(&self) -> usize {
        self.w
    }

    #[inline]
    pub fn weight(&self) -> f64 {
        self.weight
    }
}

impl fmt::Debug for DirectedEdge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} -> {} {:5.2}", self.v, self.w, self.weight)
    }
}

#[test]
fn test_directed_edge() {
    let e = DirectedEdge::new(12, 24, 3.14);
    assert_eq!(format!("{:?}", e), "12 -> 24  3.14");
}

/// Edge-weighted digraph, implemented using adjacency lists
#[derive(Clone)]
pub struct EdgeWeightedDigraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<DirectedEdge>>
}

impl EdgeWeightedDigraph {
    pub fn new(v: usize) -> EdgeWeightedDigraph {
        EdgeWeightedDigraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex must be between 0 and V");
    }

    pub fn add_edge(&mut self, e: DirectedEdge) {
        let v = e.from();
        let w = e.to();
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e);
        self.e += 1;
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<DirectedEdge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
    }

    pub fn outdegree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<DirectedEdge> {
        self.adj.iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<DirectedEdge>>().into_iter()
            })
            .collect::<Vec<DirectedEdge>>()
            .into_iter()
    }

    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("digraph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.from();
            let w = e.to();
            dot.push_str(&format!("  {} -> {} [ label=\"{}\" ];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot
    }
}

// Single-source shortest paths API
pub struct DijkstraSP<'a> {
    graph: &'a EdgeWeightedDigraph,
    dist_to: Vec<f64>,
    edge_to: Vec<Option<DirectedEdge>>,
    pq: IndexMinPQ<f64>,
    s: usize
}

impl<'a> DijkstraSP<'a> {
    fn new<'b>(graph: &'b EdgeWeightedDigraph, s: usize) -> DijkstraSP<'b> {
        let n = graph.v();
        for e in graph.edges() {
            if e.weight() < 0.0 {
                panic!("edge has negative weight in DijkstraSP");
            }
        }
        let dist_to = iter::repeat(f64::INFINITY).take(n).collect();
        let edge_to = iter::repeat(None).take(n).collect();
        let pq = IndexMinPQ::with_capacity(n);
        let mut sp = DijkstraSP {
            graph: graph,
            s: s,
            dist_to: dist_to,
            edge_to: edge_to,
            pq: pq
        };
        // alogrithm
        sp.dist_to[s] = 0.0;

        sp.pq.insert(s, 0.0);
        while !sp.pq.is_empty() {
            let v = sp.pq.del_min().unwrap();
            for e in graph.adj(v) {
                sp.relax(e);
            }
        }

        sp
    }

    fn relax(&mut self, e: DirectedEdge) {
        let v = e.from();
        let w = e.to();
        if self.dist_to[w] > self.dist_to[v] + e.weight() {
            self.dist_to[w] = self.dist_to[v] + e.weight();
            self.edge_to[w] = Some(e);
            if self.pq.contains(w) {
                self.pq.decrease_key(w, self.dist_to[w]);
            } else {
                self.pq.insert(w, self.dist_to[w]);
            }
        }
    }

    // length of shortest path from s to v
    pub fn dist_to(&self, v: usize) -> f64 {
        self.dist_to[v]
    }

    pub fn has_path_to(&self, v: usize) -> bool {
        self.dist_to[v] < f64::INFINITY
    }

    // shortest path from s to v
    pub fn path_to(&self, v: usize) -> ::std::vec::IntoIter<DirectedEdge> {
        if !self.has_path_to(v) {
            vec!().into_iter()
        } else {
            let mut path = Stack::new();
            let mut e = self.edge_to[v];
            while e.is_some() {
                path.push(e.unwrap());
                e = self.edge_to[e.unwrap().from()]
            }
            path.into_iter().collect::<Vec<DirectedEdge>>().into_iter()
        }
    }
}

impl EdgeWeightedDigraph {
    pub fn dijkstra_sp<'a>(&'a self, s: usize) -> DijkstraSP<'a> {
        DijkstraSP::new(self, s)
    }
}

#[test]
fn test_dijkstra_shortest_path() {
    let mut g = EdgeWeightedDigraph::new(6);
    g.add_edge(DirectedEdge::new(0, 1, 7.0));
    g.add_edge(DirectedEdge::new(1, 2, 10.0));
    g.add_edge(DirectedEdge::new(0, 2, 9.0));
    g.add_edge(DirectedEdge::new(0, 5, 14.0));
    g.add_edge(DirectedEdge::new(1, 3, 15.0));
    g.add_edge(DirectedEdge::new(2, 5, 2.0));
    g.add_edge(DirectedEdge::new(2, 3, 11.0));
    g.add_edge(DirectedEdge::new(4, 5, 9.0));
    g.add_edge(DirectedEdge::new(3, 4, 6.0));
    g.add_edge(DirectedEdge::new(2, 2, 1.0));

    assert_eq!(20.0, g.dijkstra_sp(0).dist_to(3));
    assert_eq!(26.0, g.dijkstra_sp(0).path_to(4).map(|e| e.weight()).sum());
}
<|endoftext|>