<commit_before>use select::SelectType;
use join::Join;
use order_by::OrderBy;
use where_cl::{ToSQL, Where};

pub struct Query<'a, T: 'a + ToSQL> {
    pub select: SelectType<'a>,
    pub from: &'a str,
    pub joins: &'a [Join<'a>],
    pub group_by: &'a [&'a str],
    pub order_by: &'a [OrderBy<'a>],
    pub where_cl: Option<T>
}

impl<'a, T: 'a + ToSQL> Query<'a, T> {
    fn to_sql(&self) -> String {
        let mut rv = String::new();
        rv.push_str("SELECT");
        rv.push(' ');
        rv.push_str(&self.select.to_sql());
        rv.push(' ');
        rv.push_str("FROM");
        rv.push(' ');
        rv.push_str(self.from);

        match *self.where_cl {
            Some(ref clause) => {
                rv.push(' ');
                rv.push_str("WHERE");
                rv.push(' ');
                rv.push_str(&clause.to_sql());
            },
            None => ()
        }

        for join in self.joins {
            rv.push(' ');
            rv.push_str(&join.to_sql());
        }

        if !self.group_by.is_empty() {
            rv.push(' ');
            rv.push_str("GROUP BY");
            rv.push(' ');
            rv.push_str(&self.group_by.connect(", "));
        }

        if !self.order_by.is_empty() {
            rv.push(' ');
            rv.push_str("ORDER BY");
            rv.push(' ');
            rv.push_str(&self.order_by
                        .into_iter()
                        .map(|x| x.to_sql())
                        .collect::<Vec<String>>()
                        .connect(", "));
        }

        rv.push(';');
        rv
    }
}

#[cfg(test)]
mod tests {
    use super::Query;
    use select::SelectType;
    use join::{JoinType, Join};
    use order_by::{Ordering, OrderBy};
    use where_cl::{ToSQL, Where};

    #[test]
    fn select_all() {
        let query: Query<ToSQL> = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_bar() {
        let clauses = &["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(clauses),
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_bar_with_vec_params() {
        let clauses = vec!["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(&clauses),
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_bar_with_vec_params_and_strings() {
        let clauses = vec!["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(&clauses),
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_join_bar() {
        let join = Join {
            join_type: JoinType::Inner,
            target: "target_table",
            clause: "2 == 2"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[join],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };

        assert_eq!(query.to_sql(), "SELECT * FROM test_table INNER JOIN target_table ON 2 == 2;");
    }

    #[test]
    fn select_foo_and_join_bar_and_bazz() {
        let bar_join = Join {
            join_type: JoinType::Inner,
            target: "bar_table",
            clause: "1 == 1"
        };

        let bazz_join = Join {
            join_type: JoinType::Left,
            target: "bazz_table",
            clause: "2 == 2"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[bar_join, bazz_join],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };

        assert_eq!(query.to_sql(), "SELECT * FROM test_table INNER JOIN bar_table ON 1 == 1 LEFT JOIN bazz_table ON 2 == 2;");
    }

    #[test]
    fn select_all_and_group_by_foo() {
        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &["foo"],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table GROUP BY foo;".to_string());
    }

    #[test]
    fn select_all_and_group_by_foo_and_bar() {
        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &["foo", "bar"],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table GROUP BY foo, bar;".to_string());
    }

    #[test]
    fn select_all_and_order_by() {
        let order_by_foo_asc = OrderBy {
            ordering: Ordering::Ascending,
            order_by: "foo"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[order_by_foo_asc],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table ORDER BY foo ASC;".to_string());
    }

    #[test]
    fn select_all_and_multi_order_by() {
        let order_by_foo_asc = OrderBy {
            ordering: Ordering::Ascending,
            order_by: "foo"
        };

        let order_by_bar_desc = OrderBy {
            ordering: Ordering::Descending,
            order_by: "bar"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[order_by_foo_asc, order_by_bar_desc],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table ORDER BY foo ASC, bar DESC;".to_string());
    }

    #[test]
    fn test_complex() {
        let bar_join = Join {
            join_type: JoinType::Inner,
            target: "bar_table",
            clause: "1 == 1"
        };

        let bazz_join = Join {
            join_type: JoinType::Left,
            target: "bazz_table",
            clause: "2 == 2"
        };

        let clauses = ["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(&clauses),
            from: "test_table",
            joins: &[bar_join, bazz_join],
            group_by: &["foo", "bar"],
            order_by: &[],
            where_cl: None
        };

        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table INNER JOIN bar_table ON 1 == 1 LEFT JOIN bazz_table ON 2 == 2 GROUP BY foo, bar;");

    }
}
<commit_msg>Fixed some errors. Still need to figure out what to do with Query type param.<commit_after>use select::SelectType;
use join::Join;
use order_by::OrderBy;
use where_cl::{ToSQL, Where};

pub struct Query<'a, T: ToSQL> {
    pub select: SelectType<'a>,
    pub from: &'a str,
    pub joins: &'a [Join<'a>],
    pub group_by: &'a [&'a str],
    pub order_by: &'a [OrderBy<'a>],
    pub where_cl: Option<T>
}

impl<'a, T: ToSQL> Query<'a, T> {
    fn to_sql(&self) -> String {
        let mut rv = String::new();
        rv.push_str("SELECT");
        rv.push(' ');
        rv.push_str(&self.select.to_sql());
        rv.push(' ');
        rv.push_str("FROM");
        rv.push(' ');
        rv.push_str(self.from);

        match self.where_cl {
            Some(ref clause) => {
                rv.push(' ');
                rv.push_str("WHERE");
                rv.push(' ');
                rv.push_str(&clause.to_sql());
            },
            None => ()
        }

        for join in self.joins {
            rv.push(' ');
            rv.push_str(&join.to_sql());
        }

        if !self.group_by.is_empty() {
            rv.push(' ');
            rv.push_str("GROUP BY");
            rv.push(' ');
            rv.push_str(&self.group_by.connect(", "));
        }

        if !self.order_by.is_empty() {
            rv.push(' ');
            rv.push_str("ORDER BY");
            rv.push(' ');
            rv.push_str(&self.order_by
                        .into_iter()
                        .map(|x| x.to_sql())
                        .collect::<Vec<String>>()
                        .connect(", "));
        }

        rv.push(';');
        rv
    }
}

#[cfg(test)]
mod tests {
    use super::Query;
    use select::SelectType;
    use join::{JoinType, Join};
    use order_by::{Ordering, OrderBy};
    use where_cl::{ToSQL, Where};

    #[test]
    fn select_all() {
        let query: Query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_bar() {
        let clauses = &["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(clauses),
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_bar_with_vec_params() {
        let clauses = vec!["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(&clauses),
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_bar_with_vec_params_and_strings() {
        let clauses = vec!["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(&clauses),
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table;".to_string());
    }

    #[test]
    fn select_foo_and_join_bar() {
        let join = Join {
            join_type: JoinType::Inner,
            target: "target_table",
            clause: "2 == 2"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[join],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };

        assert_eq!(query.to_sql(), "SELECT * FROM test_table INNER JOIN target_table ON 2 == 2;");
    }

    #[test]
    fn select_foo_and_join_bar_and_bazz() {
        let bar_join = Join {
            join_type: JoinType::Inner,
            target: "bar_table",
            clause: "1 == 1"
        };

        let bazz_join = Join {
            join_type: JoinType::Left,
            target: "bazz_table",
            clause: "2 == 2"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[bar_join, bazz_join],
            group_by: &[],
            order_by: &[],
            where_cl: None
        };

        assert_eq!(query.to_sql(), "SELECT * FROM test_table INNER JOIN bar_table ON 1 == 1 LEFT JOIN bazz_table ON 2 == 2;");
    }

    #[test]
    fn select_all_and_group_by_foo() {
        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &["foo"],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table GROUP BY foo;".to_string());
    }

    #[test]
    fn select_all_and_group_by_foo_and_bar() {
        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &["foo", "bar"],
            order_by: &[],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table GROUP BY foo, bar;".to_string());
    }

    #[test]
    fn select_all_and_order_by() {
        let order_by_foo_asc = OrderBy {
            ordering: Ordering::Ascending,
            order_by: "foo"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[order_by_foo_asc],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table ORDER BY foo ASC;".to_string());
    }

    #[test]
    fn select_all_and_multi_order_by() {
        let order_by_foo_asc = OrderBy {
            ordering: Ordering::Ascending,
            order_by: "foo"
        };

        let order_by_bar_desc = OrderBy {
            ordering: Ordering::Descending,
            order_by: "bar"
        };

        let query = Query {
            select: SelectType::All,
            from: "test_table",
            joins: &[],
            group_by: &[],
            order_by: &[order_by_foo_asc, order_by_bar_desc],
            where_cl: None
        };
        assert_eq!(query.to_sql(), "SELECT * FROM test_table ORDER BY foo ASC, bar DESC;".to_string());
    }

    #[test]
    fn test_complex() {
        let bar_join = Join {
            join_type: JoinType::Inner,
            target: "bar_table",
            clause: "1 == 1"
        };

        let bazz_join = Join {
            join_type: JoinType::Left,
            target: "bazz_table",
            clause: "2 == 2"
        };

        let clauses = ["foo", "bar"];
        let query = Query {
            select: SelectType::Specific(&clauses),
            from: "test_table",
            joins: &[bar_join, bazz_join],
            group_by: &["foo", "bar"],
            order_by: &[],
            where_cl: None
        };

        assert_eq!(query.to_sql(), "SELECT foo, bar FROM test_table INNER JOIN bar_table ON 1 == 1 LEFT JOIN bazz_table ON 2 == 2 GROUP BY foo, bar;");

    }
}
<|endoftext|>


<commit_before>use rustc_serialize::json::Json;

use super::Document;

#[derive(Debug, PartialEq)]
pub enum FilterParseError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownFilterType(String),
    NoFilter,
}

#[derive(Debug, PartialEq)]
pub enum Filter {
    Term{field: String, value: Json},
    Terms{field: String, values: Vec<Json>},
    Prefix{field: String, value: String},
    Missing{field: String},
    And{children: Vec<Filter>},
    Or{children: Vec<Filter>},
    Not{child: Box<Filter>},
}

#[derive(Debug, PartialEq)]
pub enum QueryParseError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownQueryType(String),
    NoQuery,
    FilteredNoFilter,
    FilteredNoQuery,
    FilterParseError(FilterParseError),
    MissingQueryString,
    MultiMatchMissingFields,
}

#[derive(Debug, PartialEq)]
pub enum Query {
    MatchAll,
    Match{field: String, query: String},
    MultiMatch{fields: Vec<String>, query: String},
    Filtered{query: Box<Query>, filter: Box<Filter>},
}

impl Filter {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Filter::Term{ref field, ref value} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    return field_value == value;
                }

                false
            }
            Filter::Terms{ref field, ref values} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    for value in values.iter() {
                        if field_value == value {
                            return true;
                        }
                    }
                }

                false
            }
            Filter::Prefix{ref field, ref value} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    if let Json::String(ref field_value) = *field_value {
                        return field_value.starts_with(value);
                    }
                }

                false
            }
            Filter::Missing{ref field} => {
                let obj = doc.data.as_object().unwrap();

                match obj.get(field) {
                    Some(&Json::Null) => true,
                    None => true,
                    _ => false,
                }
            }
            Filter::And{ref children} => {
                for child in children.iter() {
                    if !child.matches(doc) {
                        return false;
                    }
                }

                true
            }
            Filter::Or{ref children} => {
                for child in children.iter() {
                    if child.matches(doc) {
                        return true;
                    }
                }

                false
            }
            Filter::Not{ref child} => !child.matches(doc),
        }
    }
}

pub fn parse_filter(json: &Json) -> Result<Filter, FilterParseError> {
    let filter_json = try!(json.as_object().ok_or(FilterParseError::ExpectedObject));
    let first_key = try!(filter_json.keys().nth(0).ok_or(FilterParseError::NoFilter));

    if first_key == "term" {
        let filter_json = filter_json.get("term").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Term{
            field: first_key.clone(),
            value: filter_json.get(first_key).unwrap().clone()
        })
    } else if first_key == "terms" {
        let filter_json = filter_json.get("terms").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Terms{
            field: first_key.clone(),
            values: filter_json.get(first_key).unwrap().as_array().unwrap().clone()
        })
    } else if first_key == "prefix" {
        let filter_json = filter_json.get("prefix").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();
        let value = filter_json.get(first_key).unwrap().as_string().unwrap();

        Ok(Filter::Prefix{
            field: first_key.clone(),
            value: value.to_owned()
        })
    } else if first_key == "missing" {
        let filter_json = filter_json.get("missing").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Missing{
            field: first_key.clone()
        })
    } else if first_key == "and" {
        Ok(Filter::And{
            children: filter_json.get("and").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "or" {
        Ok(Filter::Or{
            children: filter_json.get("or").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "not" {
        Ok(Filter::Not{
            child: Box::new(parse_filter(filter_json.get("not").unwrap()).unwrap())
        })
    } else {
        Err(FilterParseError::UnknownFilterType(first_key.clone()))
    }
}

impl Query {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Query::MatchAll => true,
            Query::Match{ref field, ref query} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    let mut field_value = field_value.as_string().unwrap().to_lowercase();
                    let mut query = query.to_lowercase();

                    return field_value.contains(&query);
                }

                false
            }
            Query::MultiMatch{ref fields, ref query} => {
                let obj = doc.data.as_object().unwrap();

                for field in fields.iter() {
                    if let Some(field_value) = obj.get(field) {
                        let mut field_value = field_value.as_string().unwrap().to_lowercase();
                        let mut query = query.to_lowercase();

                        if field_value.contains(&query) {
                            return true;
                        }
                    }
                }

                false
            }
            Query::Filtered{ref query, ref filter} => {
                if filter.matches(doc) {
                    query.matches(doc)
                } else {
                    false
                }
            }
        }
    }
}

pub fn parse_match_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QueryParseError::NoQuery));

    match json_object.get(first_key).unwrap() {
        &Json::String(ref query) => {
            Ok(Query::Match {
                field: first_key.clone(),
                query: query.to_owned(),
            })
        }
        &Json::Object(ref object) => {
            Ok(Query::Match {
                field: first_key.clone(),
                query: object.get("query").unwrap().as_string().unwrap().to_owned(),
            })
        }
        // TODO: We actually expect string or object
        _ => Err(QueryParseError::ExpectedString)
    }
}

pub fn parse_multi_match_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));

    // Convert "fields" into a Vec<String>
    let fields_json = try!(json_object.get("fields").ok_or(QueryParseError::MultiMatchMissingFields));
    let fields = try!(fields_json.as_array().ok_or(QueryParseError::ExpectedArray))
                      .iter().map(|s| s.as_string().unwrap().to_owned())
                      .collect::<Vec<_>>();

    let query_json = try!(json_object.get("query").ok_or(QueryParseError::MissingQueryString));
    let query = try!(query_json.as_string().ok_or(QueryParseError::ExpectedString)).to_owned();

    Ok(Query::MultiMatch {
        fields: fields,
        query: query,
    })
}

pub fn parse_filtered_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));

    let filter_json = try!(json_object.get("filter").ok_or(QueryParseError::FilteredNoFilter));
    let filter = match parse_filter(filter_json) {
        Ok(filter) => filter,
        Err(err) => return Err(QueryParseError::FilterParseError(err)),
    };

    let query_json = try!(json_object.get("query").ok_or(QueryParseError::FilteredNoQuery));
    let query = try!(parse_query(query_json));

    Ok(Query::Filtered {
        filter: Box::new(filter),
        query: Box::new(query),
    })
}

pub fn parse_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QueryParseError::NoQuery));

    if first_key == "match_all" {
        Ok(Query::MatchAll)
    } else if first_key == "match" {
        let inner_query = json_object.get("match").unwrap();
        Ok(try!(parse_match_query(inner_query)))
    } else if first_key == "multi_match" {
        let inner_query = json_object.get("multi_match").unwrap();
        Ok(try!(parse_multi_match_query(inner_query)))
    } else if first_key == "filtered" {
        let inner_query = json_object.get("filtered").unwrap();
        Ok(try!(parse_filtered_query(inner_query)))
    } else {
        Err(QueryParseError::UnknownQueryType(first_key.clone()))
    }
}

#[cfg(test)]
mod tests {
    use rustc_serialize::json::Json;
    use super::{Query, Filter, QueryParseError, parse_query};

    #[test]
    fn test_match_all_query() {
        let query = parse_query(&Json::from_str("
            {
                \"match_all\": {}
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::MatchAll))
    }

    #[test]
    fn test_match_query() {
        let query = parse_query(&Json::from_str("
            {
                \"match\": {
                    \"title\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::Match{
            field: "title".to_owned(),
            query: "Hello world!".to_owned(),
        }))
    }

    #[test]
    fn test_match_query_without_field() {
        let query = parse_query(&Json::from_str("
            {
                \"match\": {}
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::NoQuery))
    }

    #[test]
    fn test_multi_match_query() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"fields\": [\"title\", \"body\"],
                    \"query\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::MultiMatch{
            fields: vec!["title".to_owned(), "body".to_owned()],
            query: "Hello world!".to_owned(),
        }))
    }

    #[test]
    fn test_multi_match_query_without_fields() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"query\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::MultiMatchMissingFields))
    }

    #[test]
    fn test_multi_match_query_without_query() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"fields\": [\"title\", \"body\"]
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::MissingQueryString))
    }

    #[test]
    fn test_filtered_query() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"query\": {
                        \"match\": {
                            \"title\": \"Hello world!\"
                        }
                    },
                    \"filter\": {
                        \"term\": {
                            \"date\": \"2016-01-25\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::Filtered{
            query: Box::new(Query::Match{
                field: "title".to_owned(),
                query: "Hello world!".to_owned(),
            }),
            filter: Box::new(Filter::Term{
                field: "date".to_owned(),
                value: Json::from_str("\"2016-01-25\"").unwrap(),
            })
        }))
    }

    #[test]
    fn test_filtered_query_without_query() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"filter\": {
                        \"term\": {
                            \"date\": \"2016-01-25\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::FilteredNoQuery))
    }

    #[test]
    fn test_filtered_query_without_filter() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"query\": {
                        \"match\": {
                            \"title\": \"Hello world!\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::FilteredNoFilter))
    }
}
<commit_msg>Added query operator parsing<commit_after>use rustc_serialize::json::Json;

use super::Document;

#[derive(Debug, PartialEq)]
pub enum FilterParseError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownFilterType(String),
    NoFilter,
}

#[derive(Debug, PartialEq)]
pub enum Filter {
    Term{field: String, value: Json},
    Terms{field: String, values: Vec<Json>},
    Prefix{field: String, value: String},
    Missing{field: String},
    And{children: Vec<Filter>},
    Or{children: Vec<Filter>},
    Not{child: Box<Filter>},
}

#[derive(Debug, PartialEq)]
pub enum QueryParseError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownQueryType(String),
    NoQuery,
    FilteredNoFilter,
    FilteredNoQuery,
    FilterParseError(FilterParseError),
    MissingQueryString,
    MultiMatchMissingFields,
    InvalidQueryOperator,
}

#[derive(Debug, PartialEq)]
pub enum QueryOperator {
    Or,
    And,
}

impl Default for QueryOperator {
    fn default() -> QueryOperator {
        QueryOperator::Or
    }
}

#[derive(Debug, PartialEq)]
pub enum Query {
    MatchAll,
    Match{field: String, query: String, operator: QueryOperator},
    MultiMatch{fields: Vec<String>, query: String},
    Filtered{query: Box<Query>, filter: Box<Filter>},
}

impl Filter {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Filter::Term{ref field, ref value} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    return field_value == value;
                }

                false
            }
            Filter::Terms{ref field, ref values} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    for value in values.iter() {
                        if field_value == value {
                            return true;
                        }
                    }
                }

                false
            }
            Filter::Prefix{ref field, ref value} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    if let Json::String(ref field_value) = *field_value {
                        return field_value.starts_with(value);
                    }
                }

                false
            }
            Filter::Missing{ref field} => {
                let obj = doc.data.as_object().unwrap();

                match obj.get(field) {
                    Some(&Json::Null) => true,
                    None => true,
                    _ => false,
                }
            }
            Filter::And{ref children} => {
                for child in children.iter() {
                    if !child.matches(doc) {
                        return false;
                    }
                }

                true
            }
            Filter::Or{ref children} => {
                for child in children.iter() {
                    if child.matches(doc) {
                        return true;
                    }
                }

                false
            }
            Filter::Not{ref child} => !child.matches(doc),
        }
    }
}

pub fn parse_filter(json: &Json) -> Result<Filter, FilterParseError> {
    let filter_json = try!(json.as_object().ok_or(FilterParseError::ExpectedObject));
    let first_key = try!(filter_json.keys().nth(0).ok_or(FilterParseError::NoFilter));

    if first_key == "term" {
        let filter_json = filter_json.get("term").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Term{
            field: first_key.clone(),
            value: filter_json.get(first_key).unwrap().clone()
        })
    } else if first_key == "terms" {
        let filter_json = filter_json.get("terms").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Terms{
            field: first_key.clone(),
            values: filter_json.get(first_key).unwrap().as_array().unwrap().clone()
        })
    } else if first_key == "prefix" {
        let filter_json = filter_json.get("prefix").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();
        let value = filter_json.get(first_key).unwrap().as_string().unwrap();

        Ok(Filter::Prefix{
            field: first_key.clone(),
            value: value.to_owned()
        })
    } else if first_key == "missing" {
        let filter_json = filter_json.get("missing").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Missing{
            field: first_key.clone()
        })
    } else if first_key == "and" {
        Ok(Filter::And{
            children: filter_json.get("and").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "or" {
        Ok(Filter::Or{
            children: filter_json.get("or").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "not" {
        Ok(Filter::Not{
            child: Box::new(parse_filter(filter_json.get("not").unwrap()).unwrap())
        })
    } else {
        Err(FilterParseError::UnknownFilterType(first_key.clone()))
    }
}

impl Query {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Query::MatchAll => true,
            Query::Match{ref field, ref query, ref operator} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    let mut field_value = field_value.as_string().unwrap().to_lowercase();
                    let mut query = query.to_lowercase();

                    return field_value.contains(&query);
                }

                false
            }
            Query::MultiMatch{ref fields, ref query} => {
                let obj = doc.data.as_object().unwrap();

                for field in fields.iter() {
                    if let Some(field_value) = obj.get(field) {
                        let mut field_value = field_value.as_string().unwrap().to_lowercase();
                        let mut query = query.to_lowercase();

                        if field_value.contains(&query) {
                            return true;
                        }
                    }
                }

                false
            }
            Query::Filtered{ref query, ref filter} => {
                if filter.matches(doc) {
                    query.matches(doc)
                } else {
                    false
                }
            }
        }
    }
}

pub fn parse_query_operator(json: Option<&Json>) -> Result<QueryOperator, QueryParseError> {
    match json {
        Some(json) => {
            match *json {
                Json::String(ref value) => {
                    match value.as_ref() {
                        "or" => Ok(QueryOperator::Or),
                        "and" => Ok(QueryOperator::And),
                        _ => return Err(QueryParseError::InvalidQueryOperator),
                    }
                }
                _ => return Err(QueryParseError::InvalidQueryOperator),
            }
        }
        None => Ok(QueryOperator::default())
    }
}

pub fn parse_match_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QueryParseError::NoQuery));

    match json_object.get(first_key).unwrap() {
        &Json::String(ref query) => {
            Ok(Query::Match {
                field: first_key.clone(),
                query: query.to_owned(),
                operator: QueryOperator::default(),
            })
        }
        &Json::Object(ref object) => {
            Ok(Query::Match {
                field: first_key.clone(),
                query: object.get("query").unwrap().as_string().unwrap().to_owned(),
                operator: try!(parse_query_operator(object.get("operator"))),
            })
        }
        // TODO: We actually expect string or object
        _ => Err(QueryParseError::ExpectedString)
    }
}

pub fn parse_multi_match_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));

    // Convert "fields" into a Vec<String>
    let fields_json = try!(json_object.get("fields").ok_or(QueryParseError::MultiMatchMissingFields));
    let fields = try!(fields_json.as_array().ok_or(QueryParseError::ExpectedArray))
                      .iter().map(|s| s.as_string().unwrap().to_owned())
                      .collect::<Vec<_>>();

    let query_json = try!(json_object.get("query").ok_or(QueryParseError::MissingQueryString));
    let query = try!(query_json.as_string().ok_or(QueryParseError::ExpectedString)).to_owned();

    Ok(Query::MultiMatch {
        fields: fields,
        query: query,
    })
}

pub fn parse_filtered_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));

    let filter_json = try!(json_object.get("filter").ok_or(QueryParseError::FilteredNoFilter));
    let filter = match parse_filter(filter_json) {
        Ok(filter) => filter,
        Err(err) => return Err(QueryParseError::FilterParseError(err)),
    };

    let query_json = try!(json_object.get("query").ok_or(QueryParseError::FilteredNoQuery));
    let query = try!(parse_query(query_json));

    Ok(Query::Filtered {
        filter: Box::new(filter),
        query: Box::new(query),
    })
}

pub fn parse_query(json: &Json) -> Result<Query, QueryParseError> {
    let json_object = try!(json.as_object().ok_or(QueryParseError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QueryParseError::NoQuery));

    if first_key == "match_all" {
        Ok(Query::MatchAll)
    } else if first_key == "match" {
        let inner_query = json_object.get("match").unwrap();
        Ok(try!(parse_match_query(inner_query)))
    } else if first_key == "multi_match" {
        let inner_query = json_object.get("multi_match").unwrap();
        Ok(try!(parse_multi_match_query(inner_query)))
    } else if first_key == "filtered" {
        let inner_query = json_object.get("filtered").unwrap();
        Ok(try!(parse_filtered_query(inner_query)))
    } else {
        Err(QueryParseError::UnknownQueryType(first_key.clone()))
    }
}

#[cfg(test)]
mod tests {
    use rustc_serialize::json::Json;
    use super::{Query, Filter, QueryParseError, QueryOperator, parse_query};

    #[test]
    fn test_match_all_query() {
        let query = parse_query(&Json::from_str("
            {
                \"match_all\": {}
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::MatchAll))
    }

    #[test]
    fn test_match_query() {
        let query = parse_query(&Json::from_str("
            {
                \"match\": {
                    \"title\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::Match{
            field: "title".to_owned(),
            query: "Hello world!".to_owned(),
            operator: QueryOperator::Or,
        }))
    }

    #[test]
    fn test_match_query_without_field() {
        let query = parse_query(&Json::from_str("
            {
                \"match\": {}
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::NoQuery))
    }

    #[test]
    fn test_multi_match_query() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"fields\": [\"title\", \"body\"],
                    \"query\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::MultiMatch{
            fields: vec!["title".to_owned(), "body".to_owned()],
            query: "Hello world!".to_owned(),
        }))
    }

    #[test]
    fn test_multi_match_query_without_fields() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"query\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::MultiMatchMissingFields))
    }

    #[test]
    fn test_multi_match_query_without_query() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"fields\": [\"title\", \"body\"]
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::MissingQueryString))
    }

    #[test]
    fn test_filtered_query() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"query\": {
                        \"match\": {
                            \"title\": \"Hello world!\"
                        }
                    },
                    \"filter\": {
                        \"term\": {
                            \"date\": \"2016-01-25\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::Filtered{
            query: Box::new(Query::Match{
                field: "title".to_owned(),
                query: "Hello world!".to_owned(),
                operator: QueryOperator::Or,
            }),
            filter: Box::new(Filter::Term{
                field: "date".to_owned(),
                value: Json::from_str("\"2016-01-25\"").unwrap(),
            })
        }))
    }

    #[test]
    fn test_filtered_query_without_query() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"filter\": {
                        \"term\": {
                            \"date\": \"2016-01-25\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::FilteredNoQuery))
    }

    #[test]
    fn test_filtered_query_without_filter() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"query\": {
                        \"match\": {
                            \"title\": \"Hello world!\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QueryParseError::FilteredNoFilter))
    }
}
<|endoftext|>


<commit_before>use crate::api::Encodable;
use crate::api::Reader;
use crate::api::WriteError;
use crate::api::Writer;
use crate::core::ReadResult;
use crate::core::TryIter;
use crossbeam_utils::Backoff;
use parking_lot::Mutex;
use std::iter::FusedIterator;
use std::iter::Iterator;
use std::sync::Arc;

/// A nonblocking iterator over messages in the channel, which tries multiple times to read
/// a message from a channel.
#[repr(transparent)]
pub struct RetryIter<'a, R: Reader> {
    inner: TryIter<'a, R>,
}

impl<'a, R: Reader> From<TryIter<'a, R>> for RetryIter<'a, R> {
    fn from(try_iter: TryIter<'a, R>) -> RetryIter<'a, R> {
        RetryIter { inner: try_iter }
    }
}

impl<'a, R: Reader> Iterator for RetryIter<'a, R> {
    type Item = ReadResult<'a>;
    /// Tries multiple times to read a message from channel.
    ///
    /// # Errors
    ///
    /// If the ReadResult is a Failure all subsequent call will return None.
    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let backoff = Backoff::new();
        loop {
            let res = self.inner.next();
            match res {
                Some(ReadResult::Nothing) => {
                    if backoff.is_completed() {
                        return res;
                    } else {
                        backoff.snooze();
                    }
                }
                Some(_) => return res,
                None => return None,
            }
        }
    }
    ///Returns (0, None) if records may be still available in the channel or (0, Some(0)) if
    ///the channel is exhausted.
    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

impl<'a, R: Reader> FusedIterator for RetryIter<'a, R> {}

/// Writer which decorates another writer in order to make it available to multiple threads.
/// This writer is non-blocking but will try multiple times before it give up.
#[repr(transparent)]
pub struct RetryWriter<W: Writer> {
    mx_writer: Arc<Mutex<W>>,
}

impl<W: Writer> RetryWriter<W> {
    #[inline]
    pub fn new(mx_writer: Arc<Mutex<W>>) -> RetryWriter<W> {
        RetryWriter { mx_writer }
    }
}

impl<W: Writer> Writer for RetryWriter<W> {
    ///Tries to acquire the inner writer than writes the given message into channel.
    ///
    ///	# Errors
    ///
    /// Any error returned by the decorated writer will be passed on.
    /// WriteError::Wait will be returned if the inner writer cannot be acquired.
    #[inline]
    fn write<E: Encodable>(&mut self, data: &E) -> Result<u32, WriteError> {
        let backoff = Backoff::new();
        loop {
            let try_write = self.mx_writer.try_lock();
            match try_write {
                Some(mut writer) => {
                    return writer.write(data);
                }
                None => {
                    if backoff.is_completed() {
                        return Err(WriteError::Wait);
                    } else {
                        backoff.snooze();
                    }
                }
            }
        }
    }
}
#[cfg(test)]
mod test {
    use super::*;
    use crate::api::*;
    use crate::core::*;
    use assert_matches::*;
    use tempdir::TempDir;
    #[test]
    fn retry_iter() {
        let metadata = Metadata::new(100, 1000, 10000, 1000, 1000, TickUnit::Millis);
        let test_tmp_dir = TempDir::new("kektest").unwrap();
        let txt = "There are 10 kinds of people";
        let mut msgs = txt.split_whitespace();
        let mut writer = shm_writer(&test_tmp_dir.path(), &metadata, EncoderHandler::default()).unwrap();
        let mut reader = shm_reader(&test_tmp_dir.path(), 1000).unwrap();
        let mut retry_iter: RetryIter<ShmReader> = reader.try_iter().into();
        assert_matches!(retry_iter.size_hint(), (0, None));
        assert_matches!(retry_iter.next(), Some(ReadResult::Nothing));
        for _i in 0..3 {
            let to_wr = msgs.next().unwrap().as_bytes();
            writer.write(&to_wr).unwrap();
        }
        for _i in 0..3 {
            assert_matches!(retry_iter.next(), Some(ReadResult::Record(_)));
        }
        assert_matches!(retry_iter.next(), Some(ReadResult::Nothing));
        std::mem::drop(writer);
        assert_matches!(retry_iter.next(), Some(ReadResult::Failed(ReadError::Closed)));
        assert_matches!(retry_iter.next(), None);
        assert_matches!(retry_iter.size_hint(), (0, Some(0)));
    }

    #[test]
    fn retry_write() {
        let metadata = Metadata::new(100, 1000, 10000, 1000, 1000, TickUnit::Millis);
        let test_tmp_dir = TempDir::new("kektest").unwrap();
        let writer = shm_writer(&test_tmp_dir.path(), &metadata, EncoderHandler::default()).unwrap();
        let arc_mx = Arc::new(Mutex::new(writer));
        let handles: Vec<std::thread::JoinHandle<()>> = (0..5)
            .map(|i| (i, arc_mx.clone()))
            .map(|(i, arc_mx)| {
                std::thread::spawn(move || {
                    let to_wr = format!("Hello {}", &i);
                    let mut retry_w = RetryWriter::new(arc_mx);
                    for _i in 0..3 {
                        loop {
                            match retry_w.write(&to_wr) {
                                Ok(_) => break,
                                Err(WriteError::Wait) => (),
                                Err(_) => panic!("Write failure"),
                            };
                        }
                    }
                })
            })
            .collect();
        for handle in handles {
            handle.join().unwrap();
        }
        let mut read_results = std::collections::HashMap::<&str, i32>::new();
        let mut shm_reader = shm_reader(&test_tmp_dir.path(), 1000).unwrap();
        let reader_iter = shm_reader.try_iter();
        for msg in reader_iter {
            match msg {
                ReadResult::Record(data) => {
                    *read_results.entry(std::str::from_utf8(data).unwrap()).or_default() += 1;
                }
                _ => break,
            }
        }
        assert_eq!(read_results.len(), 5);
        for (key, value) in read_results {
            dbg!(key, value);
            assert_eq!(value, 3);
        }
    }
}
<commit_msg>Minor clippy warning fixed.<commit_after>use crate::api::Encodable;
use crate::api::Reader;
use crate::api::WriteError;
use crate::api::Writer;
use crate::core::ReadResult;
use crate::core::TryIter;
use crossbeam_utils::Backoff;
use parking_lot::Mutex;
use std::iter::FusedIterator;
use std::iter::Iterator;
use std::sync::Arc;

/// A nonblocking iterator over messages in the channel, which tries multiple times to read
/// a message from a channel.
#[repr(transparent)]
pub struct RetryIter<'a, R: Reader> {
    inner: TryIter<'a, R>,
}

impl<'a, R: Reader> From<TryIter<'a, R>> for RetryIter<'a, R> {
    fn from(try_iter: TryIter<'a, R>) -> RetryIter<'a, R> {
        RetryIter { inner: try_iter }
    }
}

impl<'a, R: Reader> Iterator for RetryIter<'a, R> {
    type Item = ReadResult<'a>;
    /// Tries multiple times to read a message from channel.
    ///
    /// # Errors
    ///
    /// If the ReadResult is a Failure all subsequent call will return None.
    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let backoff = Backoff::new();
        loop {
            let res = self.inner.next();
            match res {
                Some(ReadResult::Nothing) => {
                    if backoff.is_completed() {
                        return res;
                    } else {
                        backoff.snooze();
                    }
                }
                Some(_) => return res,
                None => return None,
            }
        }
    }
    ///Returns (0, None) if records may be still available in the channel or (0, Some(0)) if
    ///the channel is exhausted.
    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

impl<'a, R: Reader> FusedIterator for RetryIter<'a, R> {}

/// Writer which decorates another writer in order to make it available to multiple threads.
/// This writer is non-blocking but will try multiple times before it give up.
#[repr(transparent)]
pub struct RetryWriter<W: Writer> {
    mx_writer: Arc<Mutex<W>>,
}

impl<W: Writer> RetryWriter<W> {
    #[inline]
    pub fn new(mx_writer: Arc<Mutex<W>>) -> RetryWriter<W> {
        RetryWriter { mx_writer }
    }
}

impl<W: Writer> Writer for RetryWriter<W> {
    ///Tries to acquire the inner writer than writes the given message into channel.
    ///
    /// # Errors
    ///
    /// Any error returned by the decorated writer will be passed on.
    /// WriteError::Wait will be returned if the inner writer cannot be acquired.
    #[inline]
    fn write<E: Encodable>(&mut self, data: &E) -> Result<u32, WriteError> {
        let backoff = Backoff::new();
        loop {
            let try_write = self.mx_writer.try_lock();
            match try_write {
                Some(mut writer) => {
                    return writer.write(data);
                }
                None => {
                    if backoff.is_completed() {
                        return Err(WriteError::Wait);
                    } else {
                        backoff.snooze();
                    }
                }
            }
        }
    }
}
#[cfg(test)]
mod test {
    use super::*;
    use crate::api::*;
    use crate::core::*;
    use assert_matches::*;
    use tempdir::TempDir;
    #[test]
    fn retry_iter() {
        let metadata = Metadata::new(100, 1000, 10000, 1000, 1000, TickUnit::Millis);
        let test_tmp_dir = TempDir::new("kektest").unwrap();
        let txt = "There are 10 kinds of people";
        let mut msgs = txt.split_whitespace();
        let mut writer = shm_writer(&test_tmp_dir.path(), &metadata, EncoderHandler::default()).unwrap();
        let mut reader = shm_reader(&test_tmp_dir.path(), 1000).unwrap();
        let mut retry_iter: RetryIter<ShmReader> = reader.try_iter().into();
        assert_matches!(retry_iter.size_hint(), (0, None));
        assert_matches!(retry_iter.next(), Some(ReadResult::Nothing));
        for _i in 0..3 {
            let to_wr = msgs.next().unwrap().as_bytes();
            writer.write(&to_wr).unwrap();
        }
        for _i in 0..3 {
            assert_matches!(retry_iter.next(), Some(ReadResult::Record(_)));
        }
        assert_matches!(retry_iter.next(), Some(ReadResult::Nothing));
        std::mem::drop(writer);
        assert_matches!(retry_iter.next(), Some(ReadResult::Failed(ReadError::Closed)));
        assert_matches!(retry_iter.next(), None);
        assert_matches!(retry_iter.size_hint(), (0, Some(0)));
    }

    #[test]
    fn retry_write() {
        let metadata = Metadata::new(100, 1000, 10000, 1000, 1000, TickUnit::Millis);
        let test_tmp_dir = TempDir::new("kektest").unwrap();
        let writer = shm_writer(&test_tmp_dir.path(), &metadata, EncoderHandler::default()).unwrap();
        let arc_mx = Arc::new(Mutex::new(writer));
        let handles: Vec<std::thread::JoinHandle<()>> = (0..5)
            .map(|i| (i, arc_mx.clone()))
            .map(|(i, arc_mx)| {
                std::thread::spawn(move || {
                    let to_wr = format!("Hello {}", &i);
                    let mut retry_w = RetryWriter::new(arc_mx);
                    for _i in 0..3 {
                        loop {
                            match retry_w.write(&to_wr) {
                                Ok(_) => break,
                                Err(WriteError::Wait) => (),
                                Err(_) => panic!("Write failure"),
                            };
                        }
                    }
                })
            })
            .collect();
        for handle in handles {
            handle.join().unwrap();
        }
        let mut read_results = std::collections::HashMap::<&str, i32>::new();
        let mut shm_reader = shm_reader(&test_tmp_dir.path(), 1000).unwrap();
        let reader_iter = shm_reader.try_iter();
        for msg in reader_iter {
            match msg {
                ReadResult::Record(data) => {
                    *read_results.entry(std::str::from_utf8(data).unwrap()).or_default() += 1;
                }
                _ => break,
            }
        }
        assert_eq!(read_results.len(), 5);
        for (key, value) in read_results {
            dbg!(key, value);
            assert_eq!(value, 3);
        }
    }
}
<|endoftext|>


<commit_before>#![crate_name = "uu_rm"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Arcterus <arcterus@mail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

#[macro_use]
extern crate uucore;

use std::collections::VecDeque;
use std::fs;
use std::io::{stdin, stderr, BufRead, Write};
use std::ops::BitOr;
use std::path::{Path, PathBuf};

#[derive(Eq, PartialEq, Clone, Copy)]
enum InteractiveMode {
    InteractiveNone,
    InteractiveOnce,
    InteractiveAlways
}

static NAME: &'static str = "rm";
static VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn uumain(args: Vec<String>) -> i32 {
    // TODO: make getopts support -R in addition to -r
    let mut opts = getopts::Options::new();

    opts.optflag("f", "force", "ignore nonexistent files and arguments, never prompt");
    opts.optflag("i", "", "prompt before every removal");
    opts.optflag("I", "", "prompt once before removing more than three files, or when removing recursively.  Less intrusive than -i, while still giving some protection against most mistakes");
    opts.optflagopt("", "interactive", "prompt according to WHEN: never, once (-I), or always (-i).  Without WHEN, prompts always", "WHEN");
    opts.optflag("", "one-file-system", "when removing a hierarchy recursively, skip any directory that is on a file system different from that of the corresponding command line argument (NOT IMPLEMENTED)");
    opts.optflag("", "no-preserve-root", "do not treat '/' specially");
    opts.optflag("", "preserve-root", "do not remove '/' (default)");
    opts.optflag("r", "recursive", "remove directories and their contents recursively");
    opts.optflag("d", "dir", "remove empty directories");
    opts.optflag("v", "verbose", "explain what is being done");
    opts.optflag("h", "help", "display this help and exit");
    opts.optflag("V", "version", "output version information and exit");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => crash!(1, "{}", f)
    };
    if matches.opt_present("help") {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {0} [OPTION]... [FILE]...", NAME);
        println!("");
        println!("{}", opts.usage("Remove (unlink) the FILE(s)."));
        println!("By default, rm does not remove directories.  Use the --recursive (-r)");
        println!("option to remove each listed directory, too, along with all of its contents");
        println!("");
        println!("To remove a file whose name starts with a '-', for example '-foo',");
        println!("use one of these commands:");
        println!("rm -- -foo");
        println!("");
        println!("rm ./-foo");
        println!("");
        println!("Note that if you use rm to remove a file, it might be possible to recover");
        println!("some of its contents, given sufficient expertise and/or time.  For greater");
        println!("assurance that the contents are truly unrecoverable, consider using shred.");
    } else if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
    } else if matches.free.is_empty() {
        show_error!("missing an argument");
        show_error!("for help, try '{0} --help'", NAME);
        return 1;
    } else {
        let force = matches.opt_present("force");
        let interactive =
            if matches.opt_present("i") {
                InteractiveMode::InteractiveAlways
            } else if matches.opt_present("I") {
                InteractiveMode::InteractiveOnce
            } else if matches.opt_present("interactive") {
                match &matches.opt_str("interactive").unwrap()[..] {
                    "none" => InteractiveMode::InteractiveNone,
                    "once" => InteractiveMode::InteractiveOnce,
                    "always" => InteractiveMode::InteractiveAlways,
                    val => {
                        crash!(1, "Invalid argument to interactive ({})", val)
                    }
                }
            } else {
                InteractiveMode::InteractiveNone
            };
        let one_fs = matches.opt_present("one-file-system");
        let preserve_root = !matches.opt_present("no-preserve-root");
        let recursive = matches.opt_present("recursive");
        let dir = matches.opt_present("dir");
        let verbose = matches.opt_present("verbose");
        if interactive == InteractiveMode::InteractiveOnce && (recursive || matches.free.len() > 3) {
            let msg =
                if recursive {
                    "Remove all arguments recursively? "
                } else {
                    "Remove all arguments? "
                };
            if !prompt(msg) {
                return 0;
            }
        }

        if remove(matches.free, force, interactive, one_fs, preserve_root, recursive, dir, verbose) {
            return 1;
        }
    }

    0
}

// TODO: implement one-file-system
#[allow(unused_variables)]
fn remove(files: Vec<String>, force: bool, interactive: InteractiveMode, one_fs: bool, preserve_root: bool, recursive: bool, dir: bool, verbose: bool) -> bool {
    let mut had_err = false;

    for filename in &files {
        let filename = &filename[..];
        let file = Path::new(filename);
        if file.exists() {
            if file.symlink_metadata().unwrap().is_dir() {
                if recursive && (filename != "/" || !preserve_root) {
                    if interactive != InteractiveMode::InteractiveAlways {
                        if let Err(e) = fs::remove_dir_all(file) {
                            had_err = true;
                            show_error!("could not remove '{}': {}", filename, e);
                        };
                    } else {
                        let mut dirs: VecDeque<PathBuf> = VecDeque::new();
                        let mut files: Vec<PathBuf> = Vec::new();
                        let mut rmdirstack: Vec<PathBuf> = Vec::new();
                        dirs.push_back(file.to_path_buf());

                        while !dirs.is_empty() {
                            let dir = dirs.pop_front().unwrap();
                            if !prompt(&(format!("rm: descend into directory '{}'? ", dir.display()))[..]) {
                                continue;
                            }

                            // iterate over items in this directory, adding to either file or
                            // directory queue
                            match fs::read_dir(dir.as_path()) {
                                Ok(rdir) => {
                                    for ent in rdir {
                                        match ent {
                                            Ok(ref f) => match f.file_type() {
                                                Ok(t) => {
                                                    if t.is_dir() {
                                                        dirs.push_back(f.path());
                                                    } else {
                                                        files.push(f.path());
                                                    }
                                                },
                                                Err(e) => {
                                                    had_err = true;
                                                    show_error!("reading '{}': {}", f.path().display(), e);
                                                },
                                            },
                                            Err(ref e) => {
                                                had_err = true;
                                                show_error!("recursing into '{}': {}", filename, e);
                                            },
                                        };
                                    }
                                },
                                Err(e) => {
                                    had_err = true;
                                    show_error!("could not recurse into '{}': {}", dir.display(), e);
                                    continue;
                                },
                            };

                            for f in &files {
                                had_err = remove_file(f.as_path(), interactive, verbose).bitor(had_err);
                            }

                            files.clear();
                            rmdirstack.push(dir);
                        }

                        for d in rmdirstack.iter().rev() {
                            had_err = remove_dir(d.as_path(), interactive, verbose).bitor(had_err);
                        }
                    }
                } else if dir && (filename != "/" || !preserve_root) {
                    had_err = remove_dir(&file, interactive, verbose).bitor(had_err);
                } else {
                    if recursive {
                        show_error!("could not remove directory '{}'", filename);
                        had_err = true;
                    } else {
                        show_error!("could not remove directory '{}' (did you mean to pass '-r'?)", filename);
                        had_err = true;
                    }
                }
            } else {
                had_err = remove_file(&file, interactive, verbose).bitor(had_err);
            }
        } else if !force {
            show_error!("no such file or directory '{}'", filename);
            had_err = true;
        }
    }

    had_err
}

fn remove_dir(path: &Path, interactive: InteractiveMode, verbose: bool) -> bool {
    let response =
        if interactive == InteractiveMode::InteractiveAlways {
            prompt_file(path, true)
        } else {
            true
        };
    if response {
        match fs::remove_dir(path) {
            Ok(_) => if verbose { println!("removed '{}'", path.display()); },
            Err(e) => {
                show_error!("removing '{}': {}", path.display(), e);
                return true;
            }
        }
    }

    false
}

fn remove_file(path: &Path, interactive: InteractiveMode, verbose: bool) -> bool {
    let response =
        if interactive == InteractiveMode::InteractiveAlways {
            prompt_file(path, false)
        } else {
            true
        };
    if response {
        match fs::remove_file(path) {
            Ok(_) => if verbose { println!("removed '{}'", path.display()); },
            Err(e) => {
                show_error!("removing '{}': {}", path.display(), e);
                return true;
            }
        }
    }

    false
}

fn prompt_file(path: &Path, is_dir: bool) -> bool {
    if is_dir {
        prompt(&(format!("rm: remove directory '{}'? ", path.display()))[..])
    } else {
        prompt(&(format!("rm: remove file '{}'? ", path.display()))[..])
    }
}

fn prompt(msg: &str) -> bool {
    stderr().write_all(msg.as_bytes()).unwrap_or(());
    stderr().flush().unwrap_or(());
    let mut buf = Vec::new();
    let stdin = stdin();
    let mut stdin = stdin.lock();
    match stdin.read_until('\n' as u8, &mut buf) {
        Ok(x) if x > 0 => {
            match buf[0] {
                0x59 | 0x79 => true,
                _ => false,
            }
        }
        _ => false,
    }
}
<commit_msg>rm: Fail when symlink_metadata() fails<commit_after>#![crate_name = "uu_rm"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Arcterus <arcterus@mail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

#[macro_use]
extern crate uucore;

use std::collections::VecDeque;
use std::fs;
use std::io::{stdin, stderr, BufRead, Write};
use std::ops::BitOr;
use std::path::{Path, PathBuf};

#[derive(Eq, PartialEq, Clone, Copy)]
enum InteractiveMode {
    InteractiveNone,
    InteractiveOnce,
    InteractiveAlways
}

static NAME: &'static str = "rm";
static VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn uumain(args: Vec<String>) -> i32 {
    // TODO: make getopts support -R in addition to -r
    let mut opts = getopts::Options::new();

    opts.optflag("f", "force", "ignore nonexistent files and arguments, never prompt");
    opts.optflag("i", "", "prompt before every removal");
    opts.optflag("I", "", "prompt once before removing more than three files, or when removing recursively.  Less intrusive than -i, while still giving some protection against most mistakes");
    opts.optflagopt("", "interactive", "prompt according to WHEN: never, once (-I), or always (-i).  Without WHEN, prompts always", "WHEN");
    opts.optflag("", "one-file-system", "when removing a hierarchy recursively, skip any directory that is on a file system different from that of the corresponding command line argument (NOT IMPLEMENTED)");
    opts.optflag("", "no-preserve-root", "do not treat '/' specially");
    opts.optflag("", "preserve-root", "do not remove '/' (default)");
    opts.optflag("r", "recursive", "remove directories and their contents recursively");
    opts.optflag("d", "dir", "remove empty directories");
    opts.optflag("v", "verbose", "explain what is being done");
    opts.optflag("h", "help", "display this help and exit");
    opts.optflag("V", "version", "output version information and exit");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => crash!(1, "{}", f)
    };
    if matches.opt_present("help") {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {0} [OPTION]... [FILE]...", NAME);
        println!("");
        println!("{}", opts.usage("Remove (unlink) the FILE(s)."));
        println!("By default, rm does not remove directories.  Use the --recursive (-r)");
        println!("option to remove each listed directory, too, along with all of its contents");
        println!("");
        println!("To remove a file whose name starts with a '-', for example '-foo',");
        println!("use one of these commands:");
        println!("rm -- -foo");
        println!("");
        println!("rm ./-foo");
        println!("");
        println!("Note that if you use rm to remove a file, it might be possible to recover");
        println!("some of its contents, given sufficient expertise and/or time.  For greater");
        println!("assurance that the contents are truly unrecoverable, consider using shred.");
    } else if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
    } else if matches.free.is_empty() {
        show_error!("missing an argument");
        show_error!("for help, try '{0} --help'", NAME);
        return 1;
    } else {
        let force = matches.opt_present("force");
        let interactive =
            if matches.opt_present("i") {
                InteractiveMode::InteractiveAlways
            } else if matches.opt_present("I") {
                InteractiveMode::InteractiveOnce
            } else if matches.opt_present("interactive") {
                match &matches.opt_str("interactive").unwrap()[..] {
                    "none" => InteractiveMode::InteractiveNone,
                    "once" => InteractiveMode::InteractiveOnce,
                    "always" => InteractiveMode::InteractiveAlways,
                    val => {
                        crash!(1, "Invalid argument to interactive ({})", val)
                    }
                }
            } else {
                InteractiveMode::InteractiveNone
            };
        let one_fs = matches.opt_present("one-file-system");
        let preserve_root = !matches.opt_present("no-preserve-root");
        let recursive = matches.opt_present("recursive");
        let dir = matches.opt_present("dir");
        let verbose = matches.opt_present("verbose");
        if interactive == InteractiveMode::InteractiveOnce && (recursive || matches.free.len() > 3) {
            let msg =
                if recursive {
                    "Remove all arguments recursively? "
                } else {
                    "Remove all arguments? "
                };
            if !prompt(msg) {
                return 0;
            }
        }

        if remove(matches.free, force, interactive, one_fs, preserve_root, recursive, dir, verbose) {
            return 1;
        }
    }

    0
}

// TODO: implement one-file-system
#[allow(unused_variables)]
fn remove(files: Vec<String>, force: bool, interactive: InteractiveMode, one_fs: bool, preserve_root: bool, recursive: bool, dir: bool, verbose: bool) -> bool {
    let mut had_err = false;

    for filename in &files {
        let filename = &filename[..];
        let file = Path::new(filename);
        if file.exists() {
            let is_dir = match file.symlink_metadata() {
                Ok(metadata) => metadata.is_dir(),
                Err(e) => {
                    had_err = true;
                    show_error!("could not read metadata of '{}': {}", filename, e);
                    continue;
                }
            };
            if is_dir {
                if recursive && (filename != "/" || !preserve_root) {
                    if interactive != InteractiveMode::InteractiveAlways {
                        if let Err(e) = fs::remove_dir_all(file) {
                            had_err = true;
                            show_error!("could not remove '{}': {}", filename, e);
                        };
                    } else {
                        let mut dirs: VecDeque<PathBuf> = VecDeque::new();
                        let mut files: Vec<PathBuf> = Vec::new();
                        let mut rmdirstack: Vec<PathBuf> = Vec::new();
                        dirs.push_back(file.to_path_buf());

                        while !dirs.is_empty() {
                            let dir = dirs.pop_front().unwrap();
                            if !prompt(&(format!("rm: descend into directory '{}'? ", dir.display()))[..]) {
                                continue;
                            }

                            // iterate over items in this directory, adding to either file or
                            // directory queue
                            match fs::read_dir(dir.as_path()) {
                                Ok(rdir) => {
                                    for ent in rdir {
                                        match ent {
                                            Ok(ref f) => match f.file_type() {
                                                Ok(t) => {
                                                    if t.is_dir() {
                                                        dirs.push_back(f.path());
                                                    } else {
                                                        files.push(f.path());
                                                    }
                                                },
                                                Err(e) => {
                                                    had_err = true;
                                                    show_error!("reading '{}': {}", f.path().display(), e);
                                                },
                                            },
                                            Err(ref e) => {
                                                had_err = true;
                                                show_error!("recursing into '{}': {}", filename, e);
                                            },
                                        };
                                    }
                                },
                                Err(e) => {
                                    had_err = true;
                                    show_error!("could not recurse into '{}': {}", dir.display(), e);
                                    continue;
                                },
                            };

                            for f in &files {
                                had_err = remove_file(f.as_path(), interactive, verbose).bitor(had_err);
                            }

                            files.clear();
                            rmdirstack.push(dir);
                        }

                        for d in rmdirstack.iter().rev() {
                            had_err = remove_dir(d.as_path(), interactive, verbose).bitor(had_err);
                        }
                    }
                } else if dir && (filename != "/" || !preserve_root) {
                    had_err = remove_dir(&file, interactive, verbose).bitor(had_err);
                } else {
                    if recursive {
                        show_error!("could not remove directory '{}'", filename);
                        had_err = true;
                    } else {
                        show_error!("could not remove directory '{}' (did you mean to pass '-r'?)", filename);
                        had_err = true;
                    }
                }
            } else {
                had_err = remove_file(&file, interactive, verbose).bitor(had_err);
            }
        } else if !force {
            show_error!("no such file or directory '{}'", filename);
            had_err = true;
        }
    }

    had_err
}

fn remove_dir(path: &Path, interactive: InteractiveMode, verbose: bool) -> bool {
    let response =
        if interactive == InteractiveMode::InteractiveAlways {
            prompt_file(path, true)
        } else {
            true
        };
    if response {
        match fs::remove_dir(path) {
            Ok(_) => if verbose { println!("removed '{}'", path.display()); },
            Err(e) => {
                show_error!("removing '{}': {}", path.display(), e);
                return true;
            }
        }
    }

    false
}

fn remove_file(path: &Path, interactive: InteractiveMode, verbose: bool) -> bool {
    let response =
        if interactive == InteractiveMode::InteractiveAlways {
            prompt_file(path, false)
        } else {
            true
        };
    if response {
        match fs::remove_file(path) {
            Ok(_) => if verbose { println!("removed '{}'", path.display()); },
            Err(e) => {
                show_error!("removing '{}': {}", path.display(), e);
                return true;
            }
        }
    }

    false
}

fn prompt_file(path: &Path, is_dir: bool) -> bool {
    if is_dir {
        prompt(&(format!("rm: remove directory '{}'? ", path.display()))[..])
    } else {
        prompt(&(format!("rm: remove file '{}'? ", path.display()))[..])
    }
}

fn prompt(msg: &str) -> bool {
    stderr().write_all(msg.as_bytes()).unwrap_or(());
    stderr().flush().unwrap_or(());
    let mut buf = Vec::new();
    let stdin = stdin();
    let mut stdin = stdin.lock();
    match stdin.read_until('\n' as u8, &mut buf) {
        Ok(x) if x > 0 => {
            match buf[0] {
                0x59 | 0x79 => true,
                _ => false,
            }
        }
        _ => false,
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::path::Path;
use std::vec::Vec;

use cgmath::prelude::*;
use cgmath::Point3;

use glium::VertexBuffer;
use glium::backend::Facade;

use crate::aabb::AABB;
use crate::bvh::{BVH, SplitMode};
use crate::mesh::{Mesh, GPUMesh};
use crate::material::{Material, GPUMaterial};
use crate::obj_load;
use crate::triangle::{RTTriangle, RTTriangleBuilder};
use crate::vertex::Vertex;

/// Scene containing all the CPU resources
pub struct Scene {
    pub vertices: Vec<Vertex>,
    pub meshes: Vec<Mesh>,
    pub materials: Vec<Material>,
    pub triangles: Vec<RTTriangle>,
    pub lights: Vec<RTTriangle>,
    pub aabb: AABB,
    pub bvh: BVH
}

/// Scene containing minimum resources for GPU rendering
pub struct GPUScene {
    pub meshes: Vec<GPUMesh>,
    pub materials: Vec<GPUMaterial>,
    pub vertex_buffer: VertexBuffer<Vertex>,
}

impl Scene {
    pub fn new(scene_path: &Path) -> Scene {
        let mut scene = Scene {
            vertices: Vec::new(),
            meshes: Vec::new(),
            materials: Vec::new(),
            triangles: Vec::new(),
            lights: Vec::new(),
            aabb: AABB { min: Point3::origin(), max: Point3::origin() },
            bvh: BVH::empty(),
        };
        scene.load_scene(scene_path);
        scene.bvh = BVH::build(&mut scene.triangles, SplitMode::SAH);
        scene
    }

    pub fn without_bvh(scene_path: &Path) -> Scene {
        let mut scene = Scene {
            vertices: Vec::new(),
            meshes: Vec::new(),
            materials: Vec::new(),
            triangles: Vec::new(),
            lights: Vec::new(),
            aabb: AABB { min: Point3::origin(), max: Point3::origin() },
            bvh: BVH::empty(),
        };
        scene.load_scene(scene_path);
        scene
    }

    /// Load a scene from the given path
    fn load_scene(&mut self, scene_path: &Path) {
        let obj = obj_load::load_obj(scene_path).expect("Failed to load scene");

        // Closure to calculate planar normal for a triangle
        let calculate_normal = |triangle: &obj_load::Triangle| -> [f32; 3] {
            let pos_i1 = triangle.index_vertices[0].pos_i;
            let pos_i2 = triangle.index_vertices[1].pos_i;
            let pos_i3 = triangle.index_vertices[2].pos_i;
            let pos_1 = obj.positions[pos_i1];
            let pos_2 = obj.positions[pos_i2];
            let pos_3 = obj.positions[pos_i3];
            let u = [pos_2[0] - pos_1[0],
                    pos_2[1] - pos_1[1],
                    pos_2[2] - pos_1[2]];
            let v = [pos_3[0] - pos_1[0],
                    pos_3[1] - pos_1[1],
                    pos_3[2] - pos_1[2]];
            let normal = [u[1]*v[2] - u[2]*v[1],
                          u[2]*v[0] - u[0]*v[2],
                          u[0]*v[1] - u[1]*v[0]];
            let length = (normal[0].powi(2) + normal[1].powi(2) + normal[2].powi(2)).sqrt();
            [normal[0] / length, normal[1] / length, normal[2] / length]
        };

        // Group the polygons by materials for easy rendering
        let mut vertex_map = HashMap::new();
        let mut material_map = HashMap::new();
        for range in &obj.material_ranges {
            let material_i = match material_map.get(&range.name) {
                Some(&i) => i,
                None => {
                    let obj_mat = obj.materials.get(&range.name)
                        .unwrap_or_else(|| panic!("Couldn't find material {}!", range.name));
                    let material = Material::new(obj_mat);
                    let i = self.materials.len();
                    self.materials.push(material);
                    material_map.insert(&range.name, i);
                    i
                }
            };
            let material = &self.materials[material_i];
            let mut mesh = Mesh::new(material_i);
            for tri in &obj.triangles[range.start_i..range.end_i] {
                let mut tri_builder = RTTriangleBuilder::new();
                let default_tex_coords = [0.0; 2];
                for index_vertex in &tri.index_vertices {
                    match vertex_map.get(index_vertex) {
                        // Vertex has already been added
                        Some(&i) => {
                            mesh.indices.push(i as u32);
                            tri_builder.add_vertex(self.vertices[i]);
                        }
                        None => {
                            let pos = obj.positions[index_vertex.pos_i];
                            self.aabb.add_point(&Point3::from(pos));

                            let tex_coords = match index_vertex.tex_i {
                                Some(tex_i) => obj.tex_coords[tex_i],
                                None => default_tex_coords
                            };
                            let normal = match index_vertex.normal_i {
                                Some(normal_i) => obj.normals[normal_i],
                                None => calculate_normal(tri)
                            };

                            mesh.indices.push(self.vertices.len() as u32);
                            vertex_map.insert(index_vertex, self.vertices.len());
                            self.vertices.push(Vertex { pos, normal, tex_coords });
                            tri_builder.add_vertex(*self.vertices.last().unwrap());
                        }
                    }
                }
                let triangle = tri_builder.build(material_i).expect("Failed to build tri!");
                if material.emissive.is_some() {
                    self.lights.push(triangle.clone());
                }
                self.triangles.push(triangle);
            }
            if !mesh.indices.is_empty() {
                self.meshes.push(mesh);
            }
        }
    }

    /// Load the textures + vertex and index buffers to the GPU
    pub fn upload_data<F: Facade>(&self, facade: &F) -> GPUScene {
        let vertex_buffer = VertexBuffer::new(facade, &self.vertices)
                                  .expect("Failed to create vertex buffer!");
        let mut meshes = Vec::new();
        let mut materials = Vec::new();
        for mesh in &self.meshes {
            meshes.push(mesh.upload_data(facade));
        }
        for material in &self.materials {
            materials.push(material.upload_textures(facade));
        }
        GPUScene {
            meshes,
            materials,
            vertex_buffer,
        }
    }

    /// Get the center of the scene as defined by the bounding box
    pub fn center(&self) -> Point3<f32> {
        self.aabb.center()
    }

    /// Get the approximate size of the scene
    pub fn size(&self) -> f32 {
        self.aabb.longest_edge()
    }
}
<commit_msg>Improve scene load panic message<commit_after>use std::collections::HashMap;
use std::path::Path;
use std::vec::Vec;

use cgmath::prelude::*;
use cgmath::Point3;

use glium::VertexBuffer;
use glium::backend::Facade;

use crate::aabb::AABB;
use crate::bvh::{BVH, SplitMode};
use crate::mesh::{Mesh, GPUMesh};
use crate::material::{Material, GPUMaterial};
use crate::obj_load;
use crate::triangle::{RTTriangle, RTTriangleBuilder};
use crate::vertex::Vertex;

/// Scene containing all the CPU resources
pub struct Scene {
    pub vertices: Vec<Vertex>,
    pub meshes: Vec<Mesh>,
    pub materials: Vec<Material>,
    pub triangles: Vec<RTTriangle>,
    pub lights: Vec<RTTriangle>,
    pub aabb: AABB,
    pub bvh: BVH
}

/// Scene containing minimum resources for GPU rendering
pub struct GPUScene {
    pub meshes: Vec<GPUMesh>,
    pub materials: Vec<GPUMaterial>,
    pub vertex_buffer: VertexBuffer<Vertex>,
}

impl Scene {
    pub fn new(scene_path: &Path) -> Scene {
        let mut scene = Scene {
            vertices: Vec::new(),
            meshes: Vec::new(),
            materials: Vec::new(),
            triangles: Vec::new(),
            lights: Vec::new(),
            aabb: AABB { min: Point3::origin(), max: Point3::origin() },
            bvh: BVH::empty(),
        };
        scene.load_scene(scene_path);
        scene.bvh = BVH::build(&mut scene.triangles, SplitMode::SAH);
        scene
    }

    pub fn without_bvh(scene_path: &Path) -> Scene {
        let mut scene = Scene {
            vertices: Vec::new(),
            meshes: Vec::new(),
            materials: Vec::new(),
            triangles: Vec::new(),
            lights: Vec::new(),
            aabb: AABB { min: Point3::origin(), max: Point3::origin() },
            bvh: BVH::empty(),
        };
        scene.load_scene(scene_path);
        scene
    }

    /// Load a scene from the given path
    fn load_scene(&mut self, scene_path: &Path) {
        let obj = obj_load::load_obj(scene_path)
            .unwrap_or_else(|err| panic!("Failed to load scene {:?}: {}", scene_path, err));

        // Closure to calculate planar normal for a triangle
        let calculate_normal = |triangle: &obj_load::Triangle| -> [f32; 3] {
            let pos_i1 = triangle.index_vertices[0].pos_i;
            let pos_i2 = triangle.index_vertices[1].pos_i;
            let pos_i3 = triangle.index_vertices[2].pos_i;
            let pos_1 = obj.positions[pos_i1];
            let pos_2 = obj.positions[pos_i2];
            let pos_3 = obj.positions[pos_i3];
            let u = [pos_2[0] - pos_1[0],
                    pos_2[1] - pos_1[1],
                    pos_2[2] - pos_1[2]];
            let v = [pos_3[0] - pos_1[0],
                    pos_3[1] - pos_1[1],
                    pos_3[2] - pos_1[2]];
            let normal = [u[1]*v[2] - u[2]*v[1],
                          u[2]*v[0] - u[0]*v[2],
                          u[0]*v[1] - u[1]*v[0]];
            let length = (normal[0].powi(2) + normal[1].powi(2) + normal[2].powi(2)).sqrt();
            [normal[0] / length, normal[1] / length, normal[2] / length]
        };

        // Group the polygons by materials for easy rendering
        let mut vertex_map = HashMap::new();
        let mut material_map = HashMap::new();
        for range in &obj.material_ranges {
            let material_i = match material_map.get(&range.name) {
                Some(&i) => i,
                None => {
                    let obj_mat = obj.materials.get(&range.name)
                        .unwrap_or_else(|| panic!("Couldn't find material {}!", range.name));
                    let material = Material::new(obj_mat);
                    let i = self.materials.len();
                    self.materials.push(material);
                    material_map.insert(&range.name, i);
                    i
                }
            };
            let material = &self.materials[material_i];
            let mut mesh = Mesh::new(material_i);
            for tri in &obj.triangles[range.start_i..range.end_i] {
                let mut tri_builder = RTTriangleBuilder::new();
                let default_tex_coords = [0.0; 2];
                for index_vertex in &tri.index_vertices {
                    match vertex_map.get(index_vertex) {
                        // Vertex has already been added
                        Some(&i) => {
                            mesh.indices.push(i as u32);
                            tri_builder.add_vertex(self.vertices[i]);
                        }
                        None => {
                            let pos = obj.positions[index_vertex.pos_i];
                            self.aabb.add_point(&Point3::from(pos));

                            let tex_coords = match index_vertex.tex_i {
                                Some(tex_i) => obj.tex_coords[tex_i],
                                None => default_tex_coords
                            };
                            let normal = match index_vertex.normal_i {
                                Some(normal_i) => obj.normals[normal_i],
                                None => calculate_normal(tri)
                            };

                            mesh.indices.push(self.vertices.len() as u32);
                            vertex_map.insert(index_vertex, self.vertices.len());
                            self.vertices.push(Vertex { pos, normal, tex_coords });
                            tri_builder.add_vertex(*self.vertices.last().unwrap());
                        }
                    }
                }
                let triangle = tri_builder.build(material_i).expect("Failed to build tri!");
                if material.emissive.is_some() {
                    self.lights.push(triangle.clone());
                }
                self.triangles.push(triangle);
            }
            if !mesh.indices.is_empty() {
                self.meshes.push(mesh);
            }
        }
    }

    /// Load the textures + vertex and index buffers to the GPU
    pub fn upload_data<F: Facade>(&self, facade: &F) -> GPUScene {
        let vertex_buffer = VertexBuffer::new(facade, &self.vertices)
                                  .expect("Failed to create vertex buffer!");
        let mut meshes = Vec::new();
        let mut materials = Vec::new();
        for mesh in &self.meshes {
            meshes.push(mesh.upload_data(facade));
        }
        for material in &self.materials {
            materials.push(material.upload_textures(facade));
        }
        GPUScene {
            meshes,
            materials,
            vertex_buffer,
        }
    }

    /// Get the center of the scene as defined by the bounding box
    pub fn center(&self) -> Point3<f32> {
        self.aabb.center()
    }

    /// Get the approximate size of the scene
    pub fn size(&self) -> f32 {
        self.aabb.longest_edge()
    }
}
<|endoftext|>