<commit_before>use std::collections::HashMap;

use ast::{Function, Expr, Statement};
use rich_number::RichNumber;
use prelude::setup_prelude;

const LOWEST_PRECEDENCE: u8 = 0;
const HIGHEST_PRECEDENCE: u8 = 9;
const FUNCTION_PRECEDENCE: u8 = HIGHEST_PRECEDENCE;

#[derive(Eq, PartialEq, Debug, Clone)]
pub enum Fixity {
    Prefix,
    Infix,
    Postfix,
}

#[derive(PartialEq, Debug, Clone)]
pub enum ContextItem {
    Number(RichNumber),
    Definition {
        // number from 0 to 9 where 9 is the highest precedence
        precedence: u8,
        fixity: Fixity,
        function: Function,
    },
    Constant(String),
    Boolean(bool),
    Nothing,
}

impl ContextItem {
    pub fn unwrap_number(self) -> RichNumber {
        match self {
            ContextItem::Number(num) => num,
            _ => panic!("Expected to unwrap a Number"),
        }
    }

    pub fn unwrap_boolean(self) -> bool {
        match self {
            ContextItem::Boolean(value) => value,
            _ => panic!("Expected to unwrap a Boolean"),
        }
    }
}

pub struct EvalContext {
    symbol_table: HashMap<String, ContextItem>,
}

#[derive(PartialEq, Debug)]
pub enum EvalError {
    NotFound(String),
}

pub type EvalResult = Result<ContextItem, EvalError>;

impl EvalContext {
    pub fn new() -> EvalContext {
        EvalContext {
            symbol_table: HashMap::new(),
        }
    }

    pub fn prelude() -> EvalContext {
        let mut context = EvalContext::new();
        setup_prelude(&mut context);
        context
    }

    /// Gets a single value without evaluating it from the current context
    pub fn get(&self, name: &str) -> Option<ContextItem> {
        self.symbol_table.get(name).map(|item| (*item).clone())
    }

    /// Convenience method for setting a number in the context
    pub fn set_number(&mut self, name: &str, value: RichNumber) {
        self.set(name, ContextItem::Number(value))
    }

    /// Defines a function in the context
    pub fn define(&mut self, name: &str, fixity: Fixity, precedence: u8, function: Function) {
        debug_assert!(precedence >= LOWEST_PRECEDENCE && precedence <= HIGHEST_PRECEDENCE);

        self.set(name, ContextItem::Definition {
            fixity: fixity,
            precedence: precedence,
            function: function,
        })
    }

    /// Creates a constant value
    pub fn set_constant(&mut self, name: &str, value: String) {
        self.set(name, ContextItem::Constant(value));
    }

    /// Creates a boolean value
    pub fn set_boolean(&mut self, name: &str, value: bool) {
        self.set(name, ContextItem::Boolean(value));
    }

    /// Adds to the current context (silently replaces if already present)
    pub fn set(&mut self, name: &str, value: ContextItem) {
        self.symbol_table.insert(name.to_owned(), value);
    }

    pub fn apply(&mut self, statement: Statement) -> EvalResult {
        match statement {
            Statement::NamedFunction {name, definition} => {
                self.define(&name, Fixity::Prefix, FUNCTION_PRECEDENCE, definition);
                Ok(ContextItem::Nothing)
            },
            Statement::AnonymousFunction(_) => Ok(ContextItem::Nothing),
            Statement::Assignment {name, value} => {
                let value = try!(self.evaluate(value));
                self.set(&name, value);
                Ok(ContextItem::Nothing)
            },
            Statement::Expression(value) => self.evaluate(value),
        }
    }

    fn evaluate(&mut self, expr: Expr) -> EvalResult {
        unimplemented!();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ast::*;
    use api::parse;
    use rich_number::RichNumber;

    #[test]
    fn get_set_in_context() {
        let mut context = EvalContext::new();
        assert_eq!(context.get("a"), None);

        let value = RichNumber::from(2f64);
        context.set_number("a", value.clone());
        assert_eq!(context.get("a").map(|n| n.unwrap_number()), Some(value));
    }

    #[test]
    fn define_variable() {
        let mut context = EvalContext::new();
        apply_string(&mut context, "x = 3");

        assert_eq!(context.get("x").map(|n| n.unwrap_number()), Some(RichNumber::from(3f64)));
    }

    #[test]
    fn basic_operators_and_precedence() {
        let result = apply_single("2 * 3 ** 2 - (100e-1 + -77.2) / 13").unwrap();
        assert_eq!(result.unwrap_number(), RichNumber::from(23.1692307692f64));
    }

    fn test_single_boolean(string: &str, expected: bool) {
        let result = apply_single(string).unwrap();
        assert_eq!(result.unwrap_boolean(), expected);
    }

    fn apply_single(string: &str) -> EvalResult {
        let mut context = EvalContext::new();
        let statement = parse_statement(string);

        context.apply(statement)
    }

    fn parse_statement(string: &str) -> Statement {
        let mut parsed = parse(string).unwrap();
        assert!(parsed.len() == 1, "parse_statement() was passed more than one statement");

        parsed.pop().unwrap()
    }

    fn apply_string(mut context: &mut EvalContext, string: &str) {
        apply_program(context, parse(string).unwrap());
    }

    fn apply_program(mut context: &mut EvalContext, program: Program) {
        for statement in program {
            context.apply(statement).expect("Error while applying statement");
        }
    }
}

<commit_msg>Added a BuiltInMethod context item<commit_after>use std::collections::HashMap;

use ast::{Function, Expr, Statement};
use rich_number::RichNumber;
use prelude::setup_prelude;

const LOWEST_PRECEDENCE: u8 = 0;
const HIGHEST_PRECEDENCE: u8 = 9;
const FUNCTION_PRECEDENCE: u8 = HIGHEST_PRECEDENCE;

#[derive(PartialEq, Debug)]
pub enum EvalError {
    UnknownSymbol(String),
}

pub type EvalResult = Result<ContextItem, EvalError>;

#[derive(Eq, PartialEq, Debug, Clone)]
pub enum Fixity {
    Prefix,
    Infix,
    Postfix,
}

#[derive(PartialEq, Debug, Clone)]
pub enum ContextItem {
    Number(RichNumber),
    Definition {
        // number from 0 to 9 where 9 is the highest precedence
        precedence: u8,
        fixity: Fixity,
        function: Function,
    },
    BuiltInMethod(fn(Vec<ContextItem>) -> EvalResult),
    Constant(String),
    Boolean(bool),
    Nothing,
}

impl ContextItem {
    pub fn unwrap_number(self) -> RichNumber {
        match self {
            ContextItem::Number(num) => num,
            _ => panic!("Expected to unwrap a Number"),
        }
    }

    pub fn unwrap_boolean(self) -> bool {
        match self {
            ContextItem::Boolean(value) => value,
            _ => panic!("Expected to unwrap a Boolean"),
        }
    }
}

pub struct EvalContext {
    symbol_table: HashMap<String, ContextItem>,
}

impl EvalContext {
    pub fn new() -> EvalContext {
        EvalContext {
            symbol_table: HashMap::new(),
        }
    }

    pub fn prelude() -> EvalContext {
        let mut context = EvalContext::new();
        setup_prelude(&mut context);
        context
    }

    /// Gets a single value without evaluating it from the current context
    pub fn get(&self, name: &str) -> Option<ContextItem> {
        self.symbol_table.get(name).map(|item| (*item).clone())
    }

    /// Convenience method for setting a number in the context
    pub fn set_number(&mut self, name: &str, value: RichNumber) {
        self.set(name, ContextItem::Number(value))
    }

    /// Defines a function in the context
    pub fn define(&mut self, name: &str, fixity: Fixity, precedence: u8, function: Function) {
        debug_assert!(precedence >= LOWEST_PRECEDENCE && precedence <= HIGHEST_PRECEDENCE);

        self.set(name, ContextItem::Definition {
            fixity: fixity,
            precedence: precedence,
            function: function,
        })
    }

    /// Creates a constant value
    pub fn set_constant(&mut self, name: &str, value: String) {
        self.set(name, ContextItem::Constant(value));
    }

    /// Creates a boolean value
    pub fn set_boolean(&mut self, name: &str, value: bool) {
        self.set(name, ContextItem::Boolean(value));
    }

    /// Adds to the current context (silently replaces if already present)
    pub fn set(&mut self, name: &str, value: ContextItem) {
        self.symbol_table.insert(name.to_owned(), value);
    }

    pub fn apply(&mut self, statement: Statement) -> EvalResult {
        match statement {
            Statement::NamedFunction {name, definition} => {
                self.define(&name, Fixity::Prefix, FUNCTION_PRECEDENCE, definition);
                Ok(ContextItem::Nothing)
            },
            Statement::AnonymousFunction(_) => Ok(ContextItem::Nothing),
            Statement::Assignment {name, value} => {
                let value = try!(self.evaluate(value));
                self.set(&name, value);
                Ok(ContextItem::Nothing)
            },
            Statement::Expression(value) => self.evaluate(value),
        }
    }

    fn evaluate(&mut self, expr: Expr) -> EvalResult {
        unimplemented!();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ast::*;
    use api::parse;
    use rich_number::RichNumber;

    #[test]
    fn get_set_in_context() {
        let mut context = EvalContext::new();
        assert_eq!(context.get("a"), None);

        let value = RichNumber::from(2f64);
        context.set_number("a", value.clone());
        assert_eq!(context.get("a").map(|n| n.unwrap_number()), Some(value));
    }

    #[test]
    fn define_variable() {
        let mut context = EvalContext::new();
        apply_string(&mut context, "x = 3");

        assert_eq!(context.get("x").map(|n| n.unwrap_number()), Some(RichNumber::from(3f64)));
    }

    #[test]
    fn basic_operators_and_precedence() {
        let result = apply_single("2 * 3 ** 2 - (100e-1 + -77.2) / 13").unwrap();
        assert_eq!(result.unwrap_number(), RichNumber::from(23.1692307692f64));
    }

    fn test_single_boolean(string: &str, expected: bool) {
        let result = apply_single(string).unwrap();
        assert_eq!(result.unwrap_boolean(), expected);
    }

    fn apply_single(string: &str) -> EvalResult {
        let mut context = EvalContext::new();
        let statement = parse_statement(string);

        context.apply(statement)
    }

    fn parse_statement(string: &str) -> Statement {
        let mut parsed = parse(string).unwrap();
        assert!(parsed.len() == 1, "parse_statement() was passed more than one statement");

        parsed.pop().unwrap()
    }

    fn apply_string(mut context: &mut EvalContext, string: &str) {
        apply_program(context, parse(string).unwrap());
    }

    fn apply_program(mut context: &mut EvalContext, program: Program) {
        for statement in program {
            context.apply(statement).expect("Error while applying statement");
        }
    }
}

<|endoftext|>


<commit_before>use regex::Regex;
use chrono::NaiveDate;
use reqwest;
use std::io::Read;
use select::document::Document;
use select::predicate::{Name, Attr};
use rustc_serialize::json::Json;
use errors::*;

// Some formats borrowed from https://github.com/amir/article-date-extractor
static FMTS: &'static [&str] = &["%A, %B %e, %Y", "%Y-%m-%dT%H:%M:%S%:z", "/%Y/%m/%d/", "/%Y/%d/%m/", "%Y-%m-%d", "%B %e, %Y", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%B %k, %Y, %H:%M %p", "%Y-%m-%d %H:%M:%S.000000"];

// Use lazy_static to ensure we only compile the regex once
lazy_static! {
    // Regex by Newspaper3k  - https://github.com/codelucas/newspaper/blob/master/newspaper/urls.py
    static ref RE: Regex = Regex::new(r"([\./\-_]{0,1}(19|20)\d{2})[\./\-_]{0,1}(([0-3]{0,1}[0-9][\./\-_])|(\w{3,5}[\./\-_]))([0-3]{0,1}[0-9][\./\-]{0,1})").unwrap();
}

// Parse the date, trying out each format
fn parse_date(input: &str) -> Result<NaiveDate> {
    let mut result: Result<NaiveDate> = Err("None of the formats matched the date".into());

    'outer: for fmt in FMTS {
        if let Ok(v) = NaiveDate::parse_from_str(input, fmt) {
            { result = Ok(v);
              break 'outer; }
        }
    }

    result
}

// Extract date from a URL
fn extract_from_url(url: &str) -> Option<String> {
    if let Some(val) = RE.find(url) {
        return Some(val.as_str().to_string())
    } else {
        return None
    }
}

// Extract date from JSON-LD
fn extract_from_ldjson<'a>(html: &'a Document) -> Option<String> {
    let mut json_date: Option<String> = None;
    let mut ldjson = String::new();
    if let Some(ldj) = html.find(Attr("type", "application/ld+json")).next() {
        ldjson = ldj.text();
    }

    let mut _decoded_ldjson = Json::from_str("{}").unwrap();

    match Json::from_str(ldjson.as_str()) {
        Ok(v) => _decoded_ldjson = v,
        _ => return None,
    }

    if let Some(date_published) = _decoded_ldjson.search("datePublished") {
        if let Some(date) = date_published.as_string() {
            json_date = Some(date.to_string())
        }
    } else if let Some(date_created) = _decoded_ldjson.search("dateCreated") {
        if let Some(date) = date_created.as_string() {
            json_date = Some(date.to_string())
        }
    }

    json_date
}

// Extract date from meta tags
fn extract_from_meta<'a>(html: &'a Document) -> Option<String> {
    let mut meta_date: Option<String> = None;

    'outer: for meta in html.find(Name("meta")) {
        let meta_name     = meta.attr("name").unwrap_or("").to_lowercase();
        let item_prop     = meta.attr("itemprop").unwrap_or("").to_lowercase();
        let http_equiv    = meta.attr("http-equiv").unwrap_or("").to_lowercase();
        let meta_property = meta.attr("property").unwrap_or("").to_lowercase();

        match meta_name.as_ref() {
            "pubdate"               | "publishdate"                  | "timestamp"       |
            "dc.date.issued"        | "date"                         | "sailthru.date"   |
            "article.published"     | "published-date"               | "article.created" |
            "article_date_original" | "cxenseparse:recs:publishtime" | "date_published"  => { if let Some(ct) = meta.attr("content") {
                                                                                meta_date = Some(ct.trim().to_string())
                                                                                }
                                                                                break 'outer; },
            _ => {},
        }

        match item_prop.as_ref() {
            "datepublished" | "datecreated" => { if let Some(ct) = meta.attr("content") {
                                                   meta_date = Some(ct.trim().to_string())
                                                 }
                                                 break 'outer; },
            _ => {},
        }

        match http_equiv.as_ref() {
            "date" =>  { if let Some(ct) = meta.attr("content") {
                           meta_date = Some(ct.trim().to_string())
                         }
                         break 'outer; },
            _ => {},
        }

        match meta_property.as_ref() {
            "article:published_time" | "bt:pubdate" => { if let Some(ct) = meta.attr("content") {
                                                           meta_date = Some(ct.trim().to_string())
                                                         }
                                                         break 'outer; },
            "og:image"                              => { if let Some(url) = meta.attr("content") {
                                                           meta_date = extract_from_url(url.trim())
                                                         }
                                                         break 'outer; },

            _ => {},
        }

    }

    meta_date
}

// Extract from html tags
fn extract_from_html_tag<'a>(html: &'a Document) -> Option<String> {
    lazy_static! {
        static ref TAG_RE: Regex = Regex::new(r"(?i)publishdate|pubdate|timestamp|article_date|articledate|date").unwrap();
    }

    let mut date: Option<String> = None;

    'initial: for time in html.find(Name("time")) {
        if let Some(dt) = time.attr("datetime") {
            { date = Some(dt.to_string());
              break 'initial; }
        } else if let Some("timestamp") = time.attr("class") {
            { date = Some(time.text().trim_matches('\n').to_string());
              break 'initial; }
        }
    }

    if date.is_none() {
        'outer: for tag in html.find(Name("span")) {
            if let Some("datePublished") = tag.attr("itemprop") {
                if let Some(v) = tag.attr("content") {
                    { date = Some(v.to_string());
                      break 'outer; }
                } else if tag.text() != "" {
                    { date = Some(tag.text().trim_matches('\n').to_string());
                      break 'outer; }
                }
            }
        }
    }

    // These next three loops are due to the lack of `find_all` method for select.rs library
    // See also: https://github.com/Webhose/article-date-extractor/blob/master/articleDateExtractor/__init__.py#L191
    if date.is_none() {
        'outer_first: for tag in html.find(Name("span")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_first; }
            }
        }
    }

    if date.is_none() {
        'outer_second: for tag in html.find(Name("p")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_second; }
            }
        }
    }

    if date.is_none() {
        'outer_third: for tag in html.find(Name("div")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_third; }
            }
        }
    }

    date
}

// Try to extract the date by using each function one by one
pub fn extract_article_published_date(link: &str, html: Option<String>) -> Result<NaiveDate> {
    let mut body: String = String::new();
    let mut _parsed_body: Option<Document> = None;

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    }

    if html.is_none() {
        if let Ok(mut response) = reqwest::get(link) {
            response.read_to_string(&mut body).unwrap();
            let doc = Document::from(body.as_str());
            _parsed_body = Some(doc);
        } else {
            return Err("Couldn't open the link".into())
        }
    } else {
        _parsed_body = Some(Document::from(html.unwrap().as_str()))
    }

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_ldjson(_parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_meta(_parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_html_tag(_parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else {
        return Err("Couldn't find the date to parse".into())
    }
}

// Unit tests
#[cfg(test)]
mod test {
    use super::extract_from_url;
    use super::parse_date;
    use super::extract_from_meta;
    use super::extract_from_ldjson;
    use super::extract_from_html_tag;
    use chrono::NaiveDate;
    use reqwest;
    use std::io::Read;
    use select::document::Document;

    #[test]
    fn parsing_date() {
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/11/30/").unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/30/11/").unwrap());

        assert!(parse_date("bad_format").is_err());
    }

    #[test]
    fn extracting_from_url() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        assert_eq!(Some("/2015/11/28/".to_string()), extract_from_url(link));

        let link = "";
        assert_eq!(None, extract_from_url(link));
    }

    #[test]
    fn extracting_from_ldjson() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-12-01T07:50:48Z".to_string()), extract_from_ldjson(&document));
    }

    #[test]
    fn extracting_from_meta() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some(("2015-11-30 23:50:48".to_string())), extract_from_meta(&document));
    }

    #[test]
    fn extracting_from_html_tag() {
        let mut response = reqwest::get("https://research.googleblog.com/2017/03/announcing-guetzli-new-open-source-jpeg.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("Thursday, March 16, 2017".to_string()), extract_from_html_tag(&document));
    }
}
<commit_msg>Add types to all variables<commit_after>use regex::Regex;
use chrono::NaiveDate;
use reqwest;
use std::io::Read;
use select::document::Document;
use select::predicate::{Name, Attr};
use rustc_serialize::json::Json;
use errors::*;

// Some formats borrowed from https://github.com/amir/article-date-extractor
static FMTS: &'static [&str] = &["%A, %B %e, %Y", "%Y-%m-%dT%H:%M:%S%:z", "/%Y/%m/%d/", "/%Y/%d/%m/", "%Y-%m-%d", "%B %e, %Y", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%B %k, %Y, %H:%M %p", "%Y-%m-%d %H:%M:%S.000000"];

// Use lazy_static to ensure we only compile the regex once
lazy_static! {
    // Regex by Newspaper3k  - https://github.com/codelucas/newspaper/blob/master/newspaper/urls.py
    static ref RE: Regex = Regex::new(r"([\./\-_]{0,1}(19|20)\d{2})[\./\-_]{0,1}(([0-3]{0,1}[0-9][\./\-_])|(\w{3,5}[\./\-_]))([0-3]{0,1}[0-9][\./\-]{0,1})").unwrap();
}

// Parse the date, trying out each format
fn parse_date(input: &str) -> Result<NaiveDate> {
    let mut result: Result<NaiveDate> = Err("None of the formats matched the date".into());

    'outer: for fmt in FMTS {
        if let Ok(v) = NaiveDate::parse_from_str(input, fmt) {
            { result = Ok(v);
              break 'outer; }
        }
    }

    result
}

// Extract date from a URL
fn extract_from_url(url: &str) -> Option<String> {
    if let Some(val) = RE.find(url) {
        return Some(val.as_str().to_string())
    } else {
        return None
    }
}

// Extract date from JSON-LD
fn extract_from_ldjson<'a>(html: &'a Document) -> Option<String> {
    let mut json_date: Option<String> = None;
    let mut ldjson: String = String::new();
    if let Some(ldj) = html.find(Attr("type", "application/ld+json")).next() {
        ldjson = ldj.text();
    }

    let mut _decoded_ldjson: Json = Json::from_str("{}").unwrap();

    match Json::from_str(ldjson.as_str()) {
        Ok(v) => _decoded_ldjson = v,
        _ => return None,
    }

    if let Some(date_published) = _decoded_ldjson.search("datePublished") {
        if let Some(date) = date_published.as_string() {
            json_date = Some(date.to_string())
        }
    } else if let Some(date_created) = _decoded_ldjson.search("dateCreated") {
        if let Some(date) = date_created.as_string() {
            json_date = Some(date.to_string())
        }
    }

    json_date
}

// Extract date from meta tags
fn extract_from_meta<'a>(html: &'a Document) -> Option<String> {
    let mut meta_date: Option<String> = None;

    'outer: for meta in html.find(Name("meta")) {
        let meta_name: String     = meta.attr("name").unwrap_or("").to_lowercase();
        let item_prop: String     = meta.attr("itemprop").unwrap_or("").to_lowercase();
        let http_equiv: String    = meta.attr("http-equiv").unwrap_or("").to_lowercase();
        let meta_property: String = meta.attr("property").unwrap_or("").to_lowercase();

        match meta_name.as_ref() {
            "pubdate"               | "publishdate"                  | "timestamp"       |
            "dc.date.issued"        | "date"                         | "sailthru.date"   |
            "article.published"     | "published-date"               | "article.created" |
            "article_date_original" | "cxenseparse:recs:publishtime" | "date_published"  => { if let Some(ct) = meta.attr("content") {
                                                                                meta_date = Some(ct.trim().to_string())
                                                                                }
                                                                                break 'outer; },
            _ => {},
        }

        match item_prop.as_ref() {
            "datepublished" | "datecreated" => { if let Some(ct) = meta.attr("content") {
                                                   meta_date = Some(ct.trim().to_string())
                                                 }
                                                 break 'outer; },
            _ => {},
        }

        match http_equiv.as_ref() {
            "date" =>  { if let Some(ct) = meta.attr("content") {
                           meta_date = Some(ct.trim().to_string())
                         }
                         break 'outer; },
            _ => {},
        }

        match meta_property.as_ref() {
            "article:published_time" | "bt:pubdate" => { if let Some(ct) = meta.attr("content") {
                                                           meta_date = Some(ct.trim().to_string())
                                                         }
                                                         break 'outer; },
            "og:image"                              => { if let Some(url) = meta.attr("content") {
                                                           meta_date = extract_from_url(url.trim())
                                                         }
                                                         break 'outer; },

            _ => {},
        }

    }

    meta_date
}

// Extract from html tags
fn extract_from_html_tag<'a>(html: &'a Document) -> Option<String> {
    lazy_static! {
        static ref TAG_RE: Regex = Regex::new(r"(?i)publishdate|pubdate|timestamp|article_date|articledate|date").unwrap();
    }

    let mut date: Option<String> = None;

    'initial: for time in html.find(Name("time")) {
        if let Some(dt) = time.attr("datetime") {
            { date = Some(dt.to_string());
              break 'initial; }
        } else if let Some("timestamp") = time.attr("class") {
            { date = Some(time.text().trim_matches('\n').to_string());
              break 'initial; }
        }
    }

    if date.is_none() {
        'outer: for tag in html.find(Name("span")) {
            if let Some("datePublished") = tag.attr("itemprop") {
                if let Some(v) = tag.attr("content") {
                    { date = Some(v.to_string());
                      break 'outer; }
                } else if tag.text() != "" {
                    { date = Some(tag.text().trim_matches('\n').to_string());
                      break 'outer; }
                }
            }
        }
    }

    // These next three loops are due to the lack of `find_all` method for select.rs library
    // See also: https://github.com/Webhose/article-date-extractor/blob/master/articleDateExtractor/__init__.py#L191
    if date.is_none() {
        'outer_first: for tag in html.find(Name("span")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_first; }
            }
        }
    }

    if date.is_none() {
        'outer_second: for tag in html.find(Name("p")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_second; }
            }
        }
    }

    if date.is_none() {
        'outer_third: for tag in html.find(Name("div")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_third; }
            }
        }
    }

    date
}

// Try to extract the date by using each function one by one
pub fn extract_article_published_date(link: &str, html: Option<String>) -> Result<NaiveDate> {
    let mut body: String = String::new();
    let mut _parsed_body: Option<Document> = None;

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    }

    if html.is_none() {
        if let Ok(mut response) = reqwest::get(link) {
            response.read_to_string(&mut body).unwrap();
            let doc = Document::from(body.as_str());
            _parsed_body = Some(doc);
        } else {
            return Err("Couldn't open the link".into())
        }
    } else {
        _parsed_body = Some(Document::from(html.unwrap().as_str()))
    }

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_ldjson(_parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_meta(_parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_html_tag(_parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else {
        return Err("Couldn't find the date to parse".into())
    }
}

// Unit tests
#[cfg(test)]
mod test {
    use super::extract_from_url;
    use super::parse_date;
    use super::extract_from_meta;
    use super::extract_from_ldjson;
    use super::extract_from_html_tag;
    use chrono::NaiveDate;
    use reqwest;
    use reqwest::Response;
    use std::io::Read;
    use select::document::Document;

    #[test]
    fn parsing_date() {
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/11/30/").unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/30/11/").unwrap());

        assert!(parse_date("bad_format").is_err());
    }

    #[test]
    fn extracting_from_url() {
        let link: &str = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        assert_eq!(Some("/2015/11/28/".to_string()), extract_from_url(link));

        let link: &str = "";
        assert_eq!(None, extract_from_url(link));
    }

    #[test]
    fn extracting_from_ldjson() {
        let mut response: Response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body: String = String::new();
        response.read_to_string(&mut body).unwrap();
        let document: Document = Document::from(body.as_str());

        assert_eq!(Some("2015-12-01T07:50:48Z".to_string()), extract_from_ldjson(&document));
    }

    #[test]
    fn extracting_from_meta() {
        let mut response: Response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body: String = String::new();
        response.read_to_string(&mut body).unwrap();
        let document: Document = Document::from(body.as_str());

        assert_eq!(Some(("2015-11-30 23:50:48".to_string())), extract_from_meta(&document));
    }

    #[test]
    fn extracting_from_html_tag() {
        let mut response: Response = reqwest::get("https://research.googleblog.com/2017/03/announcing-guetzli-new-open-source-jpeg.html").unwrap();
        let mut body: String = String::new();
        response.read_to_string(&mut body).unwrap();
        let document: Document = Document::from(body.as_str());

        assert_eq!(Some("Thursday, March 16, 2017".to_string()), extract_from_html_tag(&document));
    }
}
<|endoftext|>


<commit_before>#![feature(phase, macro_rules, if_let)]

extern crate serialize;
extern crate libc;
#[phase(plugin, link)] extern crate log;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
extern crate time;
extern crate debug;

extern crate argparse;
extern crate quire;
#[phase(plugin, link)] extern crate lithos;

use std::os::args;
use std::io::stderr;
use std::io::fs::readdir;
use std::os::{set_exit_status, self_exe_path};
use std::io::fs::PathExtensions;
use std::default::Default;

use argparse::{ArgumentParser, Store, StoreOption};
use quire::parse_config;

use lithos::tree_config::TreeConfig;
use lithos::container_config::ContainerConfig;
use lithos::child_config::ChildConfig;
use lithos::signal;

fn check_config(cfg: &TreeConfig) -> Result<(), String> {
    if !Path::new(cfg.devfs_dir.as_slice()).exists() {
        return Err(format!(
            "Devfs dir ({}) must exist and contain device nodes",
            cfg.devfs_dir));
    }
    return Ok(());
}

fn check(config_file: Path, config_dir: Option<Path>) -> Result<(), String> {
    let cfg: TreeConfig = try_str!(parse_config(&config_file,
        &*TreeConfig::validator(), Default::default()));

    try!(check_config(&cfg));
    let config_dir = config_dir.unwrap_or(cfg.config_dir);

    debug!("Checking child dir {}", config_dir.display());
    let dirlist = try_str!(readdir(&config_dir));
    for child_fn in dirlist.into_iter() {
        match (child_fn.filestem_str(), child_fn.extension_str()) {
            (Some(""), _) => continue,  // Hidden files
            (_, Some("yaml")) => {}
            _ => continue,  // Non-yaml, old, whatever, files
        }
        debug!("Checking {}", child_fn.display());
        let child_cfg: ChildConfig = try_str!(parse_config(&child_fn,
            &*ChildConfig::validator(), Default::default()));
        debug!("Opening config {}", child_fn.display());
        let config: ContainerConfig = try_str!(parse_config(
            &cfg.image_dir
                .join(child_cfg.image)
                .join(child_cfg.config.path_relative_from(
                    &Path::new("/")).unwrap()),
            &*ContainerConfig::validator(), Default::default()));
    }

    return Ok(());
}

fn check_binaries() -> bool {
    let dir = match self_exe_path() {
        Some(dir) => dir,
        None => return false,
    };
    if !dir.join("lithos_tree").exists() {
        error!("Can't find lithos_tree binary");
        return false;
    }
    if !dir.join("lithos_knot").exists() {
        error!("Can't find lithos_knot binary");
        return false;
    }
    return true;
}

fn main() {

    signal::block_all();

    if !check_binaries() {
        set_exit_status(127);
        return;
    }
    let mut config_file = Path::new("/etc/lithos.yaml");
    let mut config_dir = None;
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Checks if lithos configuration is ok");
        ap.refer(&mut config_file)
          .add_option(["-C", "--config"], box Store::<Path>,
            "Name of the global configuration file (default /etc/lithos.yaml)")
          .metavar("FILE");
        ap.refer(&mut config_dir)
          .add_option(["-D", "--config-dir"], box StoreOption::<Path>,
            concat!("Name of the alterate directory with configs. ",
                    "Useful to test configuration directory before ",
                    "switching it to be primary one"))
          .metavar("DIR");
        match ap.parse_args() {
            Ok(()) => {}
            Err(x) => {
                set_exit_status(x);
                return;
            }
        }
    }
    match check(config_file, config_dir) {
        Ok(()) => {
            set_exit_status(0);
        }
        Err(e) => {
            (write!(stderr(), "Fatal error: {}\n", e)).ok();
            set_exit_status(1);
        }
    }
}
<commit_msg>Implement hostname check in lithos_check<commit_after>#![feature(phase, macro_rules, if_let)]

extern crate serialize;
extern crate libc;
#[phase(plugin, link)] extern crate log;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
extern crate time;
extern crate debug;

extern crate argparse;
extern crate quire;
#[phase(plugin, link)] extern crate lithos;

use std::os::args;
use std::io::stderr;
use std::io::fs::readdir;
use std::os::{set_exit_status, self_exe_path};
use std::io::fs::PathExtensions;
use std::default::Default;

use argparse::{ArgumentParser, Store, StoreOption};
use quire::parse_config;

use lithos::tree_config::TreeConfig;
use lithos::container_config::ContainerConfig;
use lithos::child_config::ChildConfig;
use lithos::signal;
use lithos::network::{get_host_name, get_host_ip};

fn check_config(cfg: &TreeConfig) -> Result<(), String> {
    // TODO(tailhook) maybe check host only if we need it for hosts file
    try_str!(get_host_name());
    try_str!(get_host_ip());

    if !Path::new(cfg.devfs_dir.as_slice()).exists() {
        return Err(format!(
            "Devfs dir ({}) must exist and contain device nodes",
            cfg.devfs_dir));
    }
    return Ok(());
}

fn check(config_file: Path, config_dir: Option<Path>) -> Result<(), String> {
    let cfg: TreeConfig = try_str!(parse_config(&config_file,
        &*TreeConfig::validator(), Default::default()));

    try!(check_config(&cfg));
    let config_dir = config_dir.unwrap_or(cfg.config_dir);

    debug!("Checking child dir {}", config_dir.display());
    let dirlist = try_str!(readdir(&config_dir));
    for child_fn in dirlist.into_iter() {
        match (child_fn.filestem_str(), child_fn.extension_str()) {
            (Some(""), _) => continue,  // Hidden files
            (_, Some("yaml")) => {}
            _ => continue,  // Non-yaml, old, whatever, files
        }
        debug!("Checking {}", child_fn.display());
        let child_cfg: ChildConfig = try_str!(parse_config(&child_fn,
            &*ChildConfig::validator(), Default::default()));
        debug!("Opening config {}", child_fn.display());
        let config: ContainerConfig = try_str!(parse_config(
            &cfg.image_dir
                .join(child_cfg.image)
                .join(child_cfg.config.path_relative_from(
                    &Path::new("/")).unwrap()),
            &*ContainerConfig::validator(), Default::default()));
    }

    return Ok(());
}

fn check_binaries() -> bool {
    let dir = match self_exe_path() {
        Some(dir) => dir,
        None => return false,
    };
    if !dir.join("lithos_tree").exists() {
        error!("Can't find lithos_tree binary");
        return false;
    }
    if !dir.join("lithos_knot").exists() {
        error!("Can't find lithos_knot binary");
        return false;
    }
    return true;
}

fn main() {

    signal::block_all();

    if !check_binaries() {
        set_exit_status(127);
        return;
    }
    let mut config_file = Path::new("/etc/lithos.yaml");
    let mut config_dir = None;
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Checks if lithos configuration is ok");
        ap.refer(&mut config_file)
          .add_option(["-C", "--config"], box Store::<Path>,
            "Name of the global configuration file (default /etc/lithos.yaml)")
          .metavar("FILE");
        ap.refer(&mut config_dir)
          .add_option(["-D", "--config-dir"], box StoreOption::<Path>,
            concat!("Name of the alterate directory with configs. ",
                    "Useful to test configuration directory before ",
                    "switching it to be primary one"))
          .metavar("DIR");
        match ap.parse_args() {
            Ok(()) => {}
            Err(x) => {
                set_exit_status(x);
                return;
            }
        }
    }
    match check(config_file, config_dir) {
        Ok(()) => {
            set_exit_status(0);
        }
        Err(e) => {
            (write!(stderr(), "Fatal error: {}\n", e)).ok();
            set_exit_status(1);
        }
    }
}
<|endoftext|>


<commit_before>extern crate sha2;

use std::path::PathBuf;
use std::io::Read;
use std::fs;
use std::fs::File;

use self::sha2::Sha256;
use self::sha2::Digest;

use super::super::routing::EffectDesc;

/// Resource manager. Where to search for various file types (e.g. Effects).
/// Uses a 'dumb' implementation - doesn't try to auto-configure paths (/usr/bin/share/[...],
/// ~/.friendship, etc). Instead, designed to be configured by the host.
struct ResMan {
    dirs: Vec<PathBuf>,
}

impl ResMan {
    pub fn new() -> Self {
        Self {
            dirs: Vec::new(),
        }
    }
    pub fn add_dir(&mut self, dir: PathBuf) {
        self.dirs.push(dir);
    }
    pub fn find_effect<'a>(&'a self, desc: &'a EffectDesc) -> impl Iterator<Item=PathBuf> + 'a{
        self.iter_all_files().filter(move |f| {
            match desc.sha256() {
                &None => true,
                &Some(hash) => hash == self.file_sha256_hash(f),
            }
        })
    }
    fn iter_all_files<'a>(&'a self) -> impl Iterator<Item=PathBuf> + 'a{
        // dirs as PathBuf -> valid ReadDir objects
        self.dirs.iter().filter_map(|dir_path| {
            fs::read_dir(dir_path).ok()
        })
        // ReadDir objects -> flat list of Result<DirEntry>
        .flat_map(|read_dir| {
            read_dir
        })
        // Result<DirEntry> -> DirEntry
        .filter_map(|dir_entry| {
            dir_entry.ok()
        })
        // keep only the files
        .filter(|dir_entry| {
            if let Ok(file_type)=dir_entry.file_type() {
                file_type.is_file()
            } else {
                false
            }
        })
        // DirEntry -> Path
        .map(|dir_entry| {
            dir_entry.path()
        })
    }

    /// Returns the sha256 hash of the file's contents.
    fn file_sha256_hash(&self, path: &PathBuf) -> [u8; 32] {
        // TODO: Rewrite surroundings to avoid the possibility that the file doesn't exist / has
        // been deleted since the time the directory was enumerated.
        let mut hasher = Sha256::new();
        let mut file_contents = Vec::new();
        hasher.input({
            File::open(path).unwrap().read_to_end(&mut file_contents);
            &file_contents
        });
        let mut res: [u8; 32] = Default::default();
        res.copy_from_slice(hasher.result().as_slice());
        res
    }
}
<commit_msg>Implement `find_effect` in a way that is generic over the underlying transit/file system<commit_after>extern crate sha2;

use std::path::PathBuf;
use std::io::Read;
use std::fs;
use std::fs::File;

use self::sha2::Sha256;
use self::sha2::Digest;

use super::super::routing::EffectDesc;

/// Resource manager. Where to search for various file types (e.g. Effects).
/// Uses a 'dumb' implementation - doesn't try to auto-configure paths (/usr/bin/share/[...],
/// ~/.friendship, etc). Instead, designed to be configured by the host.
struct ResMan {
    dirs: Vec<PathBuf>,
}

impl ResMan {
    pub fn new() -> Self {
        Self {
            dirs: Vec::new(),
        }
    }
    pub fn add_dir(&mut self, dir: PathBuf) {
        self.dirs.push(dir);
    }
    /// Returns all definitions of the given effect in the form of an iterator
    ///   over boxed objects implementing io::Read.
    pub fn find_effect<'a>(&'a self, desc: &'a EffectDesc) -> impl Iterator<Item=Box<Read>> + 'a {
        self.iter_effect_files(desc).map(|path| {
            Box::new(File::open(path).unwrap()) as Box<Read>
        })
    }
    fn iter_effect_files<'a>(&'a self, desc: &'a EffectDesc) -> impl Iterator<Item=PathBuf> + 'a{
        self.iter_all_files().filter(move |f| {
            match desc.sha256() {
                &None => true,
                &Some(hash) => hash == self.file_sha256_hash(f),
            }
        })
    }
    fn iter_all_files<'a>(&'a self) -> impl Iterator<Item=PathBuf> + 'a{
        // dirs as PathBuf -> valid ReadDir objects
        self.dirs.iter().filter_map(|dir_path| {
            fs::read_dir(dir_path).ok()
        })
        // ReadDir objects -> flat list of Result<DirEntry>
        .flat_map(|read_dir| {
            read_dir
        })
        // Result<DirEntry> -> DirEntry
        .filter_map(|dir_entry| {
            dir_entry.ok()
        })
        // keep only the files
        .filter(|dir_entry| {
            if let Ok(file_type)=dir_entry.file_type() {
                file_type.is_file()
            } else {
                false
            }
        })
        // DirEntry -> Path
        .map(|dir_entry| {
            dir_entry.path()
        })
    }

    /// Returns the sha256 hash of the file's contents.
    fn file_sha256_hash(&self, path: &PathBuf) -> [u8; 32] {
        // TODO: Rewrite surroundings to avoid the possibility that the file doesn't exist / has
        // been deleted since the time the directory was enumerated.
        let mut hasher = Sha256::new();
        let mut file_contents = Vec::new();
        hasher.input({
            File::open(path).unwrap().read_to_end(&mut file_contents);
            &file_contents
        });
        let mut res: [u8; 32] = Default::default();
        res.copy_from_slice(hasher.result().as_slice());
        res
    }
}
<|endoftext|>


<commit_before>extern crate bytes;
use std::io::{Error};
use std::io;
use std::thread;

use mio::{EventLoop, EventSet, PollOpt, Handler, Token, TryWrite, TryRead};
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;
use self::bytes::{Buf, ByteBuf, MutByteBuf, SliceBuf};

const SERVER_CLIENTS: Token = Token(0);
const SERVER_PEERS: Token = Token(1);

pub struct Server {
    peers: Vec<String>,
    cli_handler: ClientHandler,
    peer_handler: PeerHandler,
}

impl Server {
    pub fn new(
        peer_port: u16,
        cli_port: u16,
        peers: Vec<String>,
    ) -> Result<Server, Error> {

        let cli_addr =
            format!("0.0.0.0:{}", cli_port).parse().unwrap();
        let cli_srv_sock =
            try!(TcpListener::bind(&cli_addr));

        let peer_addr =
            format!("0.0.0.0:{}", peer_port).parse().unwrap();
        let peer_srv_sock =
            try!(TcpListener::bind(&peer_addr));

        Ok(Server {
            peers: peers,
            cli_handler: ClientHandler {
                conns: Slab::new_starting_at(Token(129), 4096),
                srv_sock: cli_srv_sock,
            },
            peer_handler: PeerHandler {
                srv_sock: peer_srv_sock,
                conns: Slab::new_starting_at(Token(2), 128),
            },
        })
    }

    pub fn start(&mut self) -> Result<(), String> {

        let mut event_loop = EventLoop::new().unwrap();

        event_loop.register_opt(
            &self.cli_handler.srv_sock,
            SERVER_CLIENTS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.register_opt(
            &self.peer_handler.srv_sock,
            SERVER_PEERS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.run(self).unwrap();

        Err("event_loop should not have returned.".to_string())
    }
}

impl Handler for Server {
    type Timeout = ();
    type Message = ();

    fn ready(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        token: Token,
        events: EventSet,
    ) {
       if events.is_readable() {
            match token {
                SERVER_PEERS => {
                    info!("got SERVER_PEERS accept");
                    self.peer_handler.accept(event_loop).unwrap();
                },
                SERVER_CLIENTS => {
                    info!("got SERVER_CLIENTS accept");
                    self.cli_handler.accept(event_loop).unwrap();
                },
                peer if peer.as_usize() > 1 && peer.as_usize() <= 128 =>
                    self.peer_handler.conn_readable(event_loop, peer).unwrap(),
                cli if cli.as_usize() > 128 && cli.as_usize() <= 4096 =>
                    self.cli_handler.conn_readable(event_loop, cli).unwrap(),
                t => panic!("unknown token: {}", t.as_usize()),
            }
        }

        if events.is_writable() {
            match token {
                SERVER_PEERS =>
                    panic!("received writable for SERVER_PEERS"),
                SERVER_CLIENTS =>
                    panic!("received writable for token SERVER_CLIENTS"),
                peer if peer.as_usize() > 1 && peer.as_usize() <= 128 =>
                    self.peer_handler.conn_writable(event_loop, peer),
                cli if cli.as_usize() > 128 && cli.as_usize() <= 4096 =>
                    self.cli_handler.conn_writable(event_loop, cli),
                t => panic!("received writable for out-of-range token: {}", t.as_usize()),
            };
        }
    }
}

struct ServerConn {
    sock: TcpStream,
    buf: Option<ByteBuf>,
    mut_buf: Option<MutByteBuf>,
    token: Option<Token>,
    interest: EventSet
}

impl ServerConn {
    fn new(sock: TcpStream) -> ServerConn {
        ServerConn {
            sock: sock,
            buf: None,
            mut_buf: Some(ByteBuf::mut_with_capacity(2048)),
            token: None,
            interest: EventSet::hup()
        }
    }

    fn writable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        let mut buf = self.buf.take().unwrap();

        match self.sock.try_write_buf(&mut buf) {
            Ok(None) => {
                info!("client flushing buf; WOULDBLOCK");

                self.buf = Some(buf);
                self.interest.insert(EventSet::writable());
            }
            Ok(Some(r)) => {
                info!("CONN : we wrote {} bytes!", r);

                self.mut_buf = Some(buf.flip());

                self.interest.insert(EventSet::readable());
                self.interest.remove(EventSet::writable());
            }
            Err(e) => {
                match e.raw_os_error() {
                    Some(32) => {
                        info!("client disconnected");
                    },
                    Some(e) =>
                        info!("not implemented; client os err={:?}", e),
                    _ =>
                        info!("not implemented; client err={:?}", e),
                };
                // Don't reregister.
                return Err(e);
            },
        }

        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge() | PollOpt::oneshot(),
        )
    }

    fn readable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        let mut buf = self.mut_buf.take().unwrap();

        match self.sock.try_read_buf(&mut buf) {
            Ok(None) => {
                panic!("We just got readable, but were unable to read from the socket?");
            }
            Ok(Some(r)) => {
                info!("CONN : we read {} bytes!", r);
                self.interest.remove(EventSet::readable());
                self.interest.insert(EventSet::writable());
            }
            Err(e) => {
                info!("not implemented; client err={:?}", e);
                self.interest.remove(EventSet::readable());
            }

        };

        // prepare to provide this to writable
        self.buf = Some(buf.flip());
        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge(),
        )
    }
}
pub struct PeerHandler {
    conns: Slab<ServerConn>,
    srv_sock: TcpListener,
}

impl PeerHandler {
    fn accept(
        &mut self,
        event_loop: &mut EventLoop<Server>
    ) -> io::Result<()> {

        info!("peer server accepting socket");

        let sock = self.srv_sock.accept().unwrap().unwrap();
        let conn = ServerConn::new(sock,);
        let tok = self.conns.insert(conn)
            .ok().expect("could not add connection to slab");

        // Re-register accepting socket
        event_loop.reregister(
            &self.srv_sock,
            SERVER_PEERS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        );

        // Register the connection
        self.conns[tok].token = Some(tok);
        event_loop.register_opt(
            &self.conns[tok].sock,
            tok,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).ok().expect("could not register socket with event loop");

        Ok(())
    }

    fn conn_readable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("peer server conn readable; tok={:?}", tok);

        self.conn(tok).readable(event_loop)
    }

    fn conn_writable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("peer server conn writable; tok={:?}", tok);
        self.conn(tok).writable(event_loop)
    }

    fn conn<'a>(&'a mut self, tok: Token) -> &'a mut ServerConn {
        &mut self.conns[tok]
    }
}

pub struct ClientHandler {
    conns: Slab<ServerConn>,
    srv_sock: TcpListener,
}

impl ClientHandler {
    fn accept(
        &mut self,
        event_loop: &mut EventLoop<Server>,
    ) -> io::Result<()> {

        info!("cli server accepting socket");

        let sock = self.srv_sock.accept().unwrap().unwrap();
        let conn = ServerConn::new(sock,);
        let tok = self.conns.insert(conn)
            .ok().expect("could not add connection to slab");

        // Re-register accepting socket
        event_loop.reregister(
            &self.srv_sock,
            SERVER_CLIENTS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        );

        // Register the connection
        self.conns[tok].token = Some(tok);
        event_loop.register_opt(
            &self.conns[tok].sock,
            tok,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot()
        ).ok().expect("could not register socket with event loop");

        Ok(())
    }

    fn conn_readable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("cli server conn readable; tok={:?}", tok);
        self.conn(tok).readable(event_loop)
    }

    fn conn_writable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("cli server conn writable; tok={:?}", tok);
        self.conn(tok).writable(event_loop)
    }

    fn conn<'a>(&'a mut self, tok: Token) -> &'a mut ServerConn {
        &mut self.conns[tok]
    }
}
<commit_msg>unify some handling logic<commit_after>extern crate bytes;
use std::io::{Error};
use std::io;
use std::thread;

use mio::{EventLoop, EventSet, PollOpt, Handler, Token, TryWrite, TryRead};
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;
use self::bytes::{Buf, ByteBuf, MutByteBuf, SliceBuf};

const SERVER_CLIENTS: Token = Token(0);
const SERVER_PEERS: Token = Token(1);

pub struct Server {
    peers: Vec<String>,
    cli_handler: ConnSet,
    peer_handler: ConnSet,
}

impl Server {
    pub fn new(
        peer_port: u16,
        cli_port: u16,
        peers: Vec<String>,
    ) -> Result<Server, Error> {

        let cli_addr =
            format!("0.0.0.0:{}", cli_port).parse().unwrap();
        let cli_srv_sock =
            try!(TcpListener::bind(&cli_addr));

        let peer_addr =
            format!("0.0.0.0:{}", peer_port).parse().unwrap();
        let peer_srv_sock =
            try!(TcpListener::bind(&peer_addr));

        Ok(Server {
            peers: peers,
            cli_handler: ConnSet {
                srv_sock: cli_srv_sock,
                srv_token: SERVER_CLIENTS,
                conns: Slab::new_starting_at(Token(129), 4096),
            },
            peer_handler: ConnSet {
                srv_sock: peer_srv_sock,
                srv_token: SERVER_PEERS,
                conns: Slab::new_starting_at(Token(2), 128),
            },
        })
    }

    pub fn start(&mut self) -> Result<(), String> {

        let mut event_loop = EventLoop::new().unwrap();

        event_loop.register_opt(
            &self.cli_handler.srv_sock,
            SERVER_CLIENTS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.register_opt(
            &self.peer_handler.srv_sock,
            SERVER_PEERS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.run(self).unwrap();

        Err("event_loop should not have returned.".to_string())
    }
}

impl Handler for Server {
    type Timeout = ();
    type Message = ();

    fn ready(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        token: Token,
        events: EventSet,
    ) {
       if events.is_readable() {
            match token {
                SERVER_PEERS => {
                    info!("got SERVER_PEERS accept");
                    self.peer_handler.accept(event_loop).unwrap();
                },
                SERVER_CLIENTS => {
                    info!("got SERVER_CLIENTS accept");
                    self.cli_handler.accept(event_loop).unwrap();
                },
                peer if peer.as_usize() > 1 && peer.as_usize() <= 128 =>
                    self.peer_handler.conn_readable(event_loop, peer).unwrap(),
                cli if cli.as_usize() > 128 && cli.as_usize() <= 4096 =>
                    self.cli_handler.conn_readable(event_loop, cli).unwrap(),
                t => panic!("unknown token: {}", t.as_usize()),
            }
        }

        if events.is_writable() {
            match token {
                SERVER_PEERS =>
                    panic!("received writable for SERVER_PEERS"),
                SERVER_CLIENTS =>
                    panic!("received writable for token SERVER_CLIENTS"),
                peer if peer.as_usize() > 1 && peer.as_usize() <= 128 =>
                    self.peer_handler.conn_writable(event_loop, peer),
                cli if cli.as_usize() > 128 && cli.as_usize() <= 4096 =>
                    self.cli_handler.conn_writable(event_loop, cli),
                t => panic!("received writable for out-of-range token: {}", t.as_usize()),
            };
        }
    }
}

struct ServerConn {
    sock: TcpStream,
    buf: Option<ByteBuf>,
    mut_buf: Option<MutByteBuf>,
    token: Option<Token>,
    interest: EventSet
}

impl ServerConn {
    fn new(sock: TcpStream) -> ServerConn {
        ServerConn {
            sock: sock,
            buf: None,
            mut_buf: Some(ByteBuf::mut_with_capacity(2048)),
            token: None,
            interest: EventSet::hup()
        }
    }

    fn writable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        let mut buf = self.buf.take().unwrap();

        match self.sock.try_write_buf(&mut buf) {
            Ok(None) => {
                info!("client flushing buf; WOULDBLOCK");

                self.buf = Some(buf);
                self.interest.insert(EventSet::writable());
            }
            Ok(Some(r)) => {
                info!("CONN : we wrote {} bytes!", r);

                self.mut_buf = Some(buf.flip());

                self.interest.insert(EventSet::readable());
                self.interest.remove(EventSet::writable());
            }
            Err(e) => {
                match e.raw_os_error() {
                    Some(32) => {
                        info!("client disconnected");
                    },
                    Some(e) =>
                        info!("not implemented; client os err={:?}", e),
                    _ =>
                        info!("not implemented; client err={:?}", e),
                };
                // Don't reregister.
                return Err(e);
            },
        }

        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge() | PollOpt::oneshot(),
        )
    }

    fn readable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        let mut buf = self.mut_buf.take().unwrap();

        match self.sock.try_read_buf(&mut buf) {
            Ok(None) => {
                panic!("We just got readable, but were unable to read from the socket?");
            }
            Ok(Some(r)) => {
                info!("CONN : we read {} bytes!", r);
                self.interest.remove(EventSet::readable());
                self.interest.insert(EventSet::writable());
            }
            Err(e) => {
                info!("not implemented; client err={:?}", e);
                self.interest.remove(EventSet::readable());
            }

        };

        // prepare to provide this to writable
        self.buf = Some(buf.flip());
        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge(),
        )
    }
}

pub struct ConnSet {
    srv_sock: TcpListener,
    srv_token: Token,
    conns: Slab<ServerConn>,
}

impl ConnSet {
    fn accept(
        &mut self,
        event_loop: &mut EventLoop<Server>,
    ) -> io::Result<()> {

        info!("cli server accepting socket");

        let sock = self.srv_sock.accept().unwrap().unwrap();
        let conn = ServerConn::new(sock,);
        let tok = self.conns.insert(conn)
            .ok().expect("could not add connection to slab");

        // Re-register accepting socket
        event_loop.reregister(
            &self.srv_sock,
            self.srv_token,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        );

        // Register the connection
        self.conns[tok].token = Some(tok);
        event_loop.register_opt(
            &self.conns[tok].sock,
            tok,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot()
        ).ok().expect("could not register socket with event loop");

        Ok(())
    }

    fn conn_readable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("cli server conn readable; tok={:?}", tok);
        self.conn(tok).readable(event_loop)
    }

    fn conn_writable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("cli server conn writable; tok={:?}", tok);
        self.conn(tok).writable(event_loop)
    }

    fn conn<'a>(&'a mut self, tok: Token) -> &'a mut ServerConn {
        &mut self.conns[tok]
    }
}
<|endoftext|>


<commit_before>extern crate discord;

use std::env;
use discord::{Discord, State};
use discord::model::Event;

// A simple DJ bot example.
// Use by issuing the command "!dj <youtube-link>" in PM or a visible text channel.
// The bot will join the voice channel of the person issuing the command.
// "!dj stop" will stop playing, and "!dj quit" will quit the voice channel.
// The bot will quit any voice channel it is the last user in.

pub fn main() {
	// log in to the API
	let discord = Discord::from_bot_token(&env::var("DISCORD_TOKEN").expect("Bad DISCORD_TOKEN")).expect("Login failed");

	// establish websocket and voice connection
	let (mut connection, ready) = discord.connect().expect("connect failed");
	println!("[Ready] {} is serving {} servers", ready.user.username, ready.servers.len());
	let mut state = State::new(ready);

	// receive events forever
	loop {
		let event = match connection.recv_event() {
			Ok(event) => event,
			Err(err) => {
				println!("[Warning] Receive error: {:?}", err);
				if let discord::Error::WebSocket(..) = err {
					// Handle the websocket connection being dropped
					let (new_connection, ready) = discord.connect().expect("connect failed");
					connection = new_connection;
					state = State::new(ready);
					println!("[Ready] Reconnected successfully.");
				}
				if let discord::Error::Closed(..) = err {
					break
				}
				continue
			},
		};
		state.update(&event);

		match event {
			Event::MessageCreate(message) => {
				use std::ascii::AsciiExt;
				// safeguard: stop if the message is from us
				if message.author.id == state.user().id {
					continue
				}

				// reply to a command if there was one
				let mut split = message.content.split(" ");
				let first_word = split.next().unwrap_or("");
				let argument = split.next().unwrap_or("");

				if first_word.eq_ignore_ascii_case("!dj") {
					let vchan = state.find_voice_user(message.author.id);
					if argument.eq_ignore_ascii_case("stop") {
						vchan.map(|(sid, _)| connection.voice(sid).stop());
					} else if argument.eq_ignore_ascii_case("quit") {
						vchan.map(|(sid, _)| connection.drop_voice(sid));
					} else {
						let output = if let Some((server_id, channel_id)) = vchan {
							match discord::voice::open_ytdl_stream(&argument) {
								Ok(stream) => {
									let voice = connection.voice(server_id);
									voice.set_deaf(true);
									voice.connect(channel_id);
									voice.play(stream);
									String::new()
								},
								Err(error) => format!("Error: {}", error),
							}
						} else {
							"You must be in a voice channel to DJ".to_owned()
						};
						if output.len() > 0 {
							warn(discord.send_message(&message.channel_id, &output, "", false));
						}
					}
				}
			}
			Event::VoiceStateUpdate(server_id, _) => {
				// If someone moves/hangs up, and we are in a voice channel,
				if let Some(cur_channel) = connection.voice(server_id).current_channel() {
					// and our current voice channel is empty, disconnect from voice
					if let Some(srv) = state.servers().iter().find(|srv| srv.id == server_id) {
						if srv.voice_states.iter().filter(|vs| vs.channel_id == Some(cur_channel)).count() <= 1 {
							connection.voice(server_id).disconnect();
						}
					}
				}
			}
			_ => {}, // discard other events
		}
	}
}

fn warn<T, E: ::std::fmt::Debug>(result: Result<T, E>) {
	match result {
		Ok(_) => {},
		Err(err) => println!("[Warning] {:?}", err)
	}
}
<commit_msg>Improve speed of the dj.rs example<commit_after>extern crate discord;

use std::env;
use discord::{Discord, State};
use discord::model::Event;

// A simple DJ bot example.
// Use by issuing the command "!dj <youtube-link>" in PM or a visible text channel.
// The bot will join the voice channel of the person issuing the command.
// "!dj stop" will stop playing, and "!dj quit" will quit the voice channel.
// The bot will quit any voice channel it is the last user in.

pub fn main() {
	// log in to the API
	let discord = Discord::from_bot_token(&env::var("DISCORD_TOKEN").expect("Bad DISCORD_TOKEN")).expect("Login failed");

	// establish websocket and voice connection
	let (mut connection, ready) = discord.connect().expect("connect failed");
	println!("[Ready] {} is serving {} servers", ready.user.username, ready.servers.len());
	let mut state = State::new(ready);

	// receive events forever
	loop {
		let event = match connection.recv_event() {
			Ok(event) => event,
			Err(err) => {
				println!("[Warning] Receive error: {:?}", err);
				if let discord::Error::WebSocket(..) = err {
					// Handle the websocket connection being dropped
					let (new_connection, ready) = discord.connect().expect("connect failed");
					connection = new_connection;
					state = State::new(ready);
					println!("[Ready] Reconnected successfully.");
				}
				if let discord::Error::Closed(..) = err {
					break
				}
				continue
			},
		};
		state.update(&event);

		match event {
			Event::MessageCreate(message) => {
				use std::ascii::AsciiExt;
				// safeguard: stop if the message is from us
				if message.author.id == state.user().id {
					continue
				}

				// reply to a command if there was one
				let mut split = message.content.split(' ');
				let first_word = split.next().unwrap_or("");
				let argument = split.next().unwrap_or("");

				if first_word.eq_ignore_ascii_case("!dj") {
					let vchan = state.find_voice_user(message.author.id);
					if argument.eq_ignore_ascii_case("stop") {
						vchan.map(|(sid, _)| connection.voice(sid).stop());
					} else if argument.eq_ignore_ascii_case("quit") {
						vchan.map(|(sid, _)| connection.drop_voice(sid));
					} else {
						let output = if let Some((server_id, channel_id)) = vchan {
							match discord::voice::open_ytdl_stream(argument) {
								Ok(stream) => {
									let voice = connection.voice(server_id);
									voice.set_deaf(true);
									voice.connect(channel_id);
									voice.play(stream);
									String::new()
								},
								Err(error) => format!("Error: {}", error),
							}
						} else {
							"You must be in a voice channel to DJ".to_owned()
						};
						if output.is_empty() {
							warn(discord.send_message(&message.channel_id, &output, "", false));
						}
					}
				}
			}
			Event::VoiceStateUpdate(server_id, _) => {
				// If someone moves/hangs up, and we are in a voice channel,
				if let Some(cur_channel) = connection.voice(server_id).current_channel() {
					// and our current voice channel is empty, disconnect from voice
					if let Some(srv) = state.servers().iter().find(|srv| srv.id == server_id) {
						if srv.voice_states.iter().filter(|vs| vs.channel_id == Some(cur_channel)).count() <= 1 {
							connection.voice(server_id).disconnect();
						}
					}
				}
			}
			_ => {}, // discard other events
		}
	}
}

fn warn<T, E: ::std::fmt::Debug>(result: Result<T, E>) {
	match result {
		Ok(_) => {},
		Err(err) => println!("[Warning] {:?}", err)
	}
}
<|endoftext|>


<commit_before>//! Provides the functionality for memory management accross devices.
//!
//! A SharedMemory tracks the memory copies accross the devices of the Backend and manages
//!
//! * the location of these memory copies
//! * the location of the latest memory copy and
//! * the synchronisation of memory copies between devices
//!
//! A [memory copy][mem] represents one logical unit of data, which might me located at the host. The
//! SharedMemory, tracks the location of the data blob accross the various devices that the backend might
//! consist of. This allows us to run operations on various backends with the same data blob.
//!
//! [frameworks]: ../frameworks/index.html
//! [mem]: ../memory/index.html
//!
//! ## Examples
//!
//! Create SharedMemory and fill it with some numbers:
//!
//! ```
//! #![feature(clone_from_slice)]
//! # extern crate collenchyma;
//! use collenchyma::framework::IFramework;
//! use collenchyma::frameworks::Native;
//! use collenchyma::shared_memory::SharedMemory;
//! # fn main() {
//! // allocate memory
//! let native = Native::new();
//! let device = native.new_device(native.hardwares()).unwrap();
//! let shared_data = &mut SharedMemory::<i32>::new(&device, 5);
//! // fill memory with some numbers
//! let local_data = [0, 1, 2, 3, 4];
//! let data = shared_data.get_mut(&device).unwrap().as_mut_native().unwrap();
//! data.as_mut_slice().clone_from_slice(&local_data);
//! # }
//! ```

use std::collections::HashMap;
use device::{IDevice, DeviceType};
use memory::MemoryType;
use std::marker::PhantomData;
use std::{fmt, mem, error};

#[derive(Debug)]
/// Container that handles synchronization of [Memory][1] of type `T`.
/// [1]: ../memory/index.html
pub struct SharedMemory<T> {
    latest_location: DeviceType,
    copies: HashMap<DeviceType, MemoryType>,
    cap: usize,
    phantom: PhantomData<T>,
}

impl<T> SharedMemory<T> {
    /// Create new SharedMemory by allocating [Memory][1] on a Device.
    /// [1]: ../memory/index.html
    pub fn new(dev: &DeviceType, capacity: usize) -> SharedMemory<T> {
        let mut copies = HashMap::<DeviceType, MemoryType>::new();
        let copy: MemoryType;
        let alloc_size = mem::size_of::<T>() * capacity;
        match *dev {
            DeviceType::Native(ref cpu) => copy = MemoryType::Native(cpu.alloc_memory(alloc_size)),
            DeviceType::OpenCL(ref context) => copy = MemoryType::OpenCL(context.alloc_memory(alloc_size)),
        }
        copies.insert(dev.clone(), copy);
        SharedMemory {
            latest_location: dev.clone(),
            copies: copies,
            cap: capacity,
            phantom: PhantomData,
        }
    }

    /// Synchronize memory from latest location to `destination`.
    pub fn sync(&mut self, destination: &DeviceType) -> Result<(), Error> {
        if &self.latest_location != destination {
            let latest = self.latest_location.clone();
            try!(self.sync_from_to(&latest, &destination));
            self.latest_location = destination.clone();
        }
        Ok(())
    }

    /// Get a reference to the memory copy on the provided `device`.
    ///
    /// Returns `None` if there is no memory copy on the device.
    pub fn get(&self, device: &DeviceType) -> Option<&MemoryType> {
        self.copies.get(device)
    }

    /// Get a mutable reference to the memory copy on the provided `device`.
    ///
    /// Returns `None` if there is no memory copy on the device.
    pub fn get_mut(&mut self, device: &DeviceType) -> Option<&mut MemoryType> {
        self.copies.get_mut(device)
    }

    /// Synchronize memory from `source` device to `destination` device.
    fn sync_from_to(&mut self, source: &DeviceType, destination: &DeviceType) -> Result<(), Error> {
        if source != destination {
            match self.aquire_copies(source, destination) {
                Ok((source_copy, mut destination_copy)) => {
                    match source.clone() {
                        DeviceType::Native(cpu) => {
                            if let MemoryType::Native(ref src) = source_copy {
                                cpu.sync_memory_to(&src, &mut destination_copy, destination)
                            }
                        },
                        DeviceType::OpenCL(context) => {
                            if let MemoryType::OpenCL(ref src) = source_copy {
                                context.sync_memory_to(&src,&mut destination_copy, destination)
                            }
                        },
                    }
                    self.return_copies(source, source_copy, destination, destination_copy);
                    Ok(())
                },
                Err(err) => Err(err),
            }
        } else {
            Ok(())
        }
    }

    /// Aquire ownership over the copies for synchronizing.
    fn aquire_copies(&mut self, source: &DeviceType, destination: &DeviceType) -> Result<(MemoryType, MemoryType), Error> {
        let source_copy: MemoryType;
        let destination_copy: MemoryType;
        match self.copies.remove(source) {
            Some(source_cpy) => source_copy = source_cpy,
            None => return Err(Error::MissingSource(format!("SharedMemory does not hold a copy on source device {:?}.", source)))
        }
        match self.copies.remove(destination) {
            Some(destination_cpy) => destination_copy = destination_cpy,
            None => return Err(Error::MissingDestination(format!("SharedMemory does not hold a copy on destination device {:?}.", destination)))
        }

        Ok((source_copy, destination_copy))
    }

    /// Return ownership over the copies after synchronizing.
    fn return_copies(&mut self, src: &DeviceType, src_mem: MemoryType, dest: &DeviceType, dest_mem: MemoryType) {
        self.copies.insert(src.clone(), src_mem);
        self.copies.insert(dest.clone(), dest_mem);
    }

    /// Track a new `device` and allocate memory on it.
    ///
    /// Returns an error if the SharedMemory is already tracking the `device`.
    pub fn add_device(&mut self, device: &DeviceType) -> Result<&mut Self, Error> {
        match self.copies.get(device) {
            Some(_) => Err(Error::InvalidMemoryAllocation(format!("SharedMemory already tracks memory for this device. No memory allocation."))),
            None => {
                let copy: MemoryType;
                match *device {
                    DeviceType::Native(ref cpu) => copy = MemoryType::Native(cpu.alloc_memory(mem::size_of::<T>())),
                    DeviceType::OpenCL(ref context) => copy = MemoryType::OpenCL(context.alloc_memory(mem::size_of::<T>())),
                };
                self.copies.insert(device.clone(), copy);
                Ok(self)
            }
        }
    }

    /// Returns the device that contains the up-to-date memory copy.
    pub fn latest_device(&self) -> &DeviceType {
        &self.latest_location
    }

    /// Returns the number of elements for which the SharedMemory has been allocated.
    pub fn capacity(&self) -> usize {
        self.cap
    }
}

/// Errors than can occur when synchronizing memory.
#[derive(Debug)]
pub enum Error {
    /// No copy on source device.
    MissingSource(String),
    /// No copy on destination device.
    MissingDestination(String),
    /// No memory allocation on specified device happened.
    InvalidMemoryAllocation(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::MissingSource(ref err) => write!(f, "{:?}", err),
            Error::MissingDestination(ref err) => write!(f, "{:?}", err),
            Error::InvalidMemoryAllocation(ref err) => write!(f, "{:?}", err),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::MissingSource(ref err) => err,
            Error::MissingDestination(ref err) => err,
            Error::InvalidMemoryAllocation(ref err) => err,
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::MissingSource(_) => None,
            Error::MissingDestination(_) => None,
            Error::InvalidMemoryAllocation(_) => None,
        }
    }
}

impl From<Error> for ::error::Error {
    fn from(err: Error) -> ::error::Error {
        ::error::Error::SharedMemory(err)
    }
}
<commit_msg>Auto merge of #9 - skade:patch-1, r=hobofan<commit_after>//! Provides the functionality for memory management across devices.
//!
//! A SharedMemory tracks the memory copies across the devices of the Backend and manages
//!
//! * the location of these memory copies
//! * the location of the latest memory copy and
//! * the synchronisation of memory copies between devices
//!
//! A [memory copy][mem] represents one logical unit of data, which might me located at the host. The
//! SharedMemory, tracks the location of the data blob across the various devices that the backend might
//! consist of. This allows us to run operations on various backends with the same data blob.
//!
//! [frameworks]: ../frameworks/index.html
//! [mem]: ../memory/index.html
//!
//! ## Examples
//!
//! Create SharedMemory and fill it with some numbers:
//!
//! ```
//! #![feature(clone_from_slice)]
//! # extern crate collenchyma;
//! use collenchyma::framework::IFramework;
//! use collenchyma::frameworks::Native;
//! use collenchyma::shared_memory::SharedMemory;
//! # fn main() {
//! // allocate memory
//! let native = Native::new();
//! let device = native.new_device(native.hardwares()).unwrap();
//! let shared_data = &mut SharedMemory::<i32>::new(&device, 5);
//! // fill memory with some numbers
//! let local_data = [0, 1, 2, 3, 4];
//! let data = shared_data.get_mut(&device).unwrap().as_mut_native().unwrap();
//! data.as_mut_slice().clone_from_slice(&local_data);
//! # }
//! ```

use std::collections::HashMap;
use device::{IDevice, DeviceType};
use memory::MemoryType;
use std::marker::PhantomData;
use std::{fmt, mem, error};

#[derive(Debug)]
/// Container that handles synchronization of [Memory][1] of type `T`.
/// [1]: ../memory/index.html
pub struct SharedMemory<T> {
    latest_location: DeviceType,
    copies: HashMap<DeviceType, MemoryType>,
    cap: usize,
    phantom: PhantomData<T>,
}

impl<T> SharedMemory<T> {
    /// Create new SharedMemory by allocating [Memory][1] on a Device.
    /// [1]: ../memory/index.html
    pub fn new(dev: &DeviceType, capacity: usize) -> SharedMemory<T> {
        let mut copies = HashMap::<DeviceType, MemoryType>::new();
        let copy: MemoryType;
        let alloc_size = mem::size_of::<T>() * capacity;
        match *dev {
            DeviceType::Native(ref cpu) => copy = MemoryType::Native(cpu.alloc_memory(alloc_size)),
            DeviceType::OpenCL(ref context) => copy = MemoryType::OpenCL(context.alloc_memory(alloc_size)),
        }
        copies.insert(dev.clone(), copy);
        SharedMemory {
            latest_location: dev.clone(),
            copies: copies,
            cap: capacity,
            phantom: PhantomData,
        }
    }

    /// Synchronize memory from latest location to `destination`.
    pub fn sync(&mut self, destination: &DeviceType) -> Result<(), Error> {
        if &self.latest_location != destination {
            let latest = self.latest_location.clone();
            try!(self.sync_from_to(&latest, &destination));
            self.latest_location = destination.clone();
        }
        Ok(())
    }

    /// Get a reference to the memory copy on the provided `device`.
    ///
    /// Returns `None` if there is no memory copy on the device.
    pub fn get(&self, device: &DeviceType) -> Option<&MemoryType> {
        self.copies.get(device)
    }

    /// Get a mutable reference to the memory copy on the provided `device`.
    ///
    /// Returns `None` if there is no memory copy on the device.
    pub fn get_mut(&mut self, device: &DeviceType) -> Option<&mut MemoryType> {
        self.copies.get_mut(device)
    }

    /// Synchronize memory from `source` device to `destination` device.
    fn sync_from_to(&mut self, source: &DeviceType, destination: &DeviceType) -> Result<(), Error> {
        if source != destination {
            match self.aquire_copies(source, destination) {
                Ok((source_copy, mut destination_copy)) => {
                    match source.clone() {
                        DeviceType::Native(cpu) => {
                            if let MemoryType::Native(ref src) = source_copy {
                                cpu.sync_memory_to(&src, &mut destination_copy, destination)
                            }
                        },
                        DeviceType::OpenCL(context) => {
                            if let MemoryType::OpenCL(ref src) = source_copy {
                                context.sync_memory_to(&src,&mut destination_copy, destination)
                            }
                        },
                    }
                    self.return_copies(source, source_copy, destination, destination_copy);
                    Ok(())
                },
                Err(err) => Err(err),
            }
        } else {
            Ok(())
        }
    }

    /// Aquire ownership over the copies for synchronizing.
    fn aquire_copies(&mut self, source: &DeviceType, destination: &DeviceType) -> Result<(MemoryType, MemoryType), Error> {
        let source_copy: MemoryType;
        let destination_copy: MemoryType;
        match self.copies.remove(source) {
            Some(source_cpy) => source_copy = source_cpy,
            None => return Err(Error::MissingSource(format!("SharedMemory does not hold a copy on source device {:?}.", source)))
        }
        match self.copies.remove(destination) {
            Some(destination_cpy) => destination_copy = destination_cpy,
            None => return Err(Error::MissingDestination(format!("SharedMemory does not hold a copy on destination device {:?}.", destination)))
        }

        Ok((source_copy, destination_copy))
    }

    /// Return ownership over the copies after synchronizing.
    fn return_copies(&mut self, src: &DeviceType, src_mem: MemoryType, dest: &DeviceType, dest_mem: MemoryType) {
        self.copies.insert(src.clone(), src_mem);
        self.copies.insert(dest.clone(), dest_mem);
    }

    /// Track a new `device` and allocate memory on it.
    ///
    /// Returns an error if the SharedMemory is already tracking the `device`.
    pub fn add_device(&mut self, device: &DeviceType) -> Result<&mut Self, Error> {
        match self.copies.get(device) {
            Some(_) => Err(Error::InvalidMemoryAllocation(format!("SharedMemory already tracks memory for this device. No memory allocation."))),
            None => {
                let copy: MemoryType;
                match *device {
                    DeviceType::Native(ref cpu) => copy = MemoryType::Native(cpu.alloc_memory(mem::size_of::<T>())),
                    DeviceType::OpenCL(ref context) => copy = MemoryType::OpenCL(context.alloc_memory(mem::size_of::<T>())),
                };
                self.copies.insert(device.clone(), copy);
                Ok(self)
            }
        }
    }

    /// Returns the device that contains the up-to-date memory copy.
    pub fn latest_device(&self) -> &DeviceType {
        &self.latest_location
    }

    /// Returns the number of elements for which the SharedMemory has been allocated.
    pub fn capacity(&self) -> usize {
        self.cap
    }
}

/// Errors than can occur when synchronizing memory.
#[derive(Debug)]
pub enum Error {
    /// No copy on source device.
    MissingSource(String),
    /// No copy on destination device.
    MissingDestination(String),
    /// No memory allocation on specified device happened.
    InvalidMemoryAllocation(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::MissingSource(ref err) => write!(f, "{:?}", err),
            Error::MissingDestination(ref err) => write!(f, "{:?}", err),
            Error::InvalidMemoryAllocation(ref err) => write!(f, "{:?}", err),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::MissingSource(ref err) => err,
            Error::MissingDestination(ref err) => err,
            Error::InvalidMemoryAllocation(ref err) => err,
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::MissingSource(_) => None,
            Error::MissingDestination(_) => None,
            Error::InvalidMemoryAllocation(_) => None,
        }
    }
}

impl From<Error> for ::error::Error {
    fn from(err: Error) -> ::error::Error {
        ::error::Error::SharedMemory(err)
    }
}
<|endoftext|>


<commit_before>#![allow(non_upper_case_globals)]
#![feature(std_misc)]
/* workaround: the rust compiler doesn't recognize
   the feature std_misc yet, still it warns
   about using it */
#![allow(unused_features)]
#![feature(std_misc)]

extern crate libc;
use libc::{c_int, c_ulonglong, c_char, size_t};

include!("src/core.rs");
include!("src/crypto_aead_chacha20poly1305.rs");

include!("src/crypto_auth.rs");
include!("src/crypto_auth_hmacsha256.rs");
include!("src/crypto_auth_hmacsha512.rs");
include!("src/crypto_auth_hmacsha512256.rs");

include!("src/crypto_box_curve25519xsalsa20poly1305.rs");

include!("src/crypto_hash.rs");
include!("src/crypto_hash_sha256.rs");
include!("src/crypto_hash_sha512.rs");

include!("src/crypto_onetimeauth.rs");
include!("src/crypto_onetimeauth_poly1305.rs");

include!("src/crypto_pwhash_scryptsalsa208sha256.rs");

include!("src/crypto_scalarmult_curve25519.rs");
include!("src/crypto_secretbox_xsalsa20poly1305.rs");
include!("src/crypto_shorthash_siphash24.rs");
include!("src/crypto_sign_ed25519.rs");
include!("src/crypto_sign_edwards25519sha512batch.rs");

include!("src/crypto_stream.rs");
include!("src/crypto_stream_aes128ctr.rs");
include!("src/crypto_stream_chacha20.rs");
include!("src/crypto_stream_salsa20.rs");
include!("src/crypto_stream_salsa2012.rs");
include!("src/crypto_stream_salsa208.rs");
include!("src/crypto_stream_xsalsa20.rs");

include!("src/crypto_verify_16.rs");
include!("src/crypto_verify_32.rs");

include!("src/randombytes.rs");
<commit_msg>Removed unnecessary #![feature(..)] directives<commit_after>#![allow(non_upper_case_globals)]

extern crate libc;
use libc::{c_int, c_ulonglong, c_char, size_t};

include!("src/core.rs");
include!("src/crypto_aead_chacha20poly1305.rs");

include!("src/crypto_auth.rs");
include!("src/crypto_auth_hmacsha256.rs");
include!("src/crypto_auth_hmacsha512.rs");
include!("src/crypto_auth_hmacsha512256.rs");

include!("src/crypto_box_curve25519xsalsa20poly1305.rs");

include!("src/crypto_hash.rs");
include!("src/crypto_hash_sha256.rs");
include!("src/crypto_hash_sha512.rs");

include!("src/crypto_onetimeauth.rs");
include!("src/crypto_onetimeauth_poly1305.rs");

include!("src/crypto_pwhash_scryptsalsa208sha256.rs");

include!("src/crypto_scalarmult_curve25519.rs");
include!("src/crypto_secretbox_xsalsa20poly1305.rs");
include!("src/crypto_shorthash_siphash24.rs");
include!("src/crypto_sign_ed25519.rs");
include!("src/crypto_sign_edwards25519sha512batch.rs");

include!("src/crypto_stream.rs");
include!("src/crypto_stream_aes128ctr.rs");
include!("src/crypto_stream_chacha20.rs");
include!("src/crypto_stream_salsa20.rs");
include!("src/crypto_stream_salsa2012.rs");
include!("src/crypto_stream_salsa208.rs");
include!("src/crypto_stream_xsalsa20.rs");

include!("src/crypto_verify_16.rs");
include!("src/crypto_verify_32.rs");

include!("src/randombytes.rs");
<|endoftext|>