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

use search::term::Term;
use search::index::reader::IndexReader;
use search::query::Query;
use search::query::term_matcher::TermMatcher;

pub enum QuerySetIterator<'a, T: IndexReader<'a>> {
    None,
    All {
        iter: T::AllDocRefIterator,
    },
    Term {
        iter: T::TermDocRefIterator,
    },
    Conjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Disjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Exclusion {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
}

impl <'a, T: IndexReader<'a>> Iterator for QuerySetIterator<'a, T> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        match *self {
            QuerySetIterator::None => None,
            QuerySetIterator::All{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Term{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Conjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();

                                return Some(doc_id_a);
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            return None;
                        }
                    } else {
                        return None;
                    }
                }
            }
            QuerySetIterator::Disjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                if let Some(doc_id_a) = *current_doc_a {
                    if let Some(doc_id_b) = *current_doc_b {
                        if doc_id_a == doc_id_b {
                            *current_doc_a = iter_a.next();
                            *current_doc_b = iter_a.next();
                            Some(doc_id_a)
                        } else if doc_id_a < doc_id_b {
                            *current_doc_a = iter_a.next();
                            Some(doc_id_a)
                        } else {
                            *current_doc_b = iter_b.next();
                            Some(doc_id_b)
                        }
                    } else {
                        *current_doc_a = iter_a.next();
                        Some(doc_id_a)
                    }
                } else {
                    if let Some(doc_id_b) = *current_doc_b {
                        *current_doc_b = iter_b.next();
                        Some(doc_id_b)
                    } else {
                        None
                    }
                }
            }
            QuerySetIterator::Exclusion{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                                return Some(doc_id_a);
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            *current_doc_a = iter_a.next();
                            return Some(doc_id_a);
                        }
                    } else {
                        return None;
                    }
                }
            }
        }
    }
}

fn build_conjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T>  {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Conjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_conjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

fn build_disjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T> {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Disjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_disjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

pub fn build_iterator_from_query<'a, T: IndexReader<'a>>(reader: &'a T, query: &Query) -> QuerySetIterator<'a, T> {
    match *query {
        Query::MatchAll => {
            QuerySetIterator::All {
                iter: reader.iter_docids_all(),
            }
        }
        Query::MatchNone => {
            QuerySetIterator::None
        }
        Query::MatchTerm{ref field, ref term, ref matcher} => {
            match *matcher {
                TermMatcher::Exact => {
                    QuerySetIterator::Term {
                        iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                    }
                }
                TermMatcher::Prefix => {
                    // Find all terms in the index that match the prefix
                    let terms = match *term {
                         Term::String(ref term) => {
                             reader.iter_terms().filter_map(|k| {
                                 if let Term::String(ref k) = *k {
                                     if k.starts_with(term) {
                                         return Some(Term::String(k.clone()));
                                     }
                                 }

                                 None
                             }).collect::<Vec<Term>>()
                         }
                         _ => return QuerySetIterator::None,
                    };

                    match terms.len() {
                        0 => QuerySetIterator::None,
                        1 => {
                            QuerySetIterator::Term {
                                iter: reader.iter_docids_with_term(terms[0].clone(), field.clone()),
                            }
                        }
                        _ => {
                            // Produce a disjunction iterator for all the terms
                            let mut iters = VecDeque::new();
                            for term in terms.iter() {
                                iters.push_back(QuerySetIterator::Term {
                                    iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                                });
                            }

                            build_disjunction_iterator(iters)
                        }
                    }
                }
            }
        }
        Query::Conjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_conjunction_iterator(iters)
        }
        Query::Disjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::NDisjunction{ref queries, minimum_should_match} => {
            // TODO
            QuerySetIterator::None
        }
        Query::DisjunctionMax{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::Exclude{ref query, ref exclude} => {
            QuerySetIterator::Exclusion {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, exclude)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Filter{ref query, ref filter} => {
            QuerySetIterator::Conjunction {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, filter)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Score{ref query, mul, add} => {
            build_iterator_from_query(reader, query)
        }
    }
}

#[cfg(test)]
mod benches {
    use test::Bencher;

    use search::term::Term;
    use search::token::Token;
    use search::document::Document;
    use search::index::reader::IndexReader;
    use search::index::store::memory::MemoryIndexStore;
    use search::query_set::QuerySetIterator;

    fn make_test_store() -> MemoryIndexStore {
        let mut store = MemoryIndexStore::new();

        for i in 0..10000 {
            let mut tokens = Vec::new();

            if i % 3 == 0 {
                let position = tokens.len() as u32 + 1;
                tokens.push(Token {
                    term: Term::String("fizz".to_string()),
                    position: position,
                });
            }

            if i % 5 == 0 {
                let position = tokens.len() as u32 + 1;
                tokens.push(Token {
                    term: Term::String("buzz".to_string()),
                    position: position,
                });
            }

            store.insert_or_update_document(Document {
                key: i.to_string(),
                fields: btreemap! {
                    "body".to_string() => tokens
                }
            });
        }

        store
    }

    #[bench]
    fn bench_all(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::All {
                iter: store.iter_docids_all(),
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizz_term(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("fizz".to_string()), "body".to_string()),
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_buzz_term(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizzbuzz_conjunction(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut fizz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut buzz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Conjunction {
                iter_a: Box::new(fizz_iterator),
                iter_b: Box::new(buzz_iterator),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizzbuzz_disjunction(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut fizz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut buzz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Disjunction {
                iter_a: Box::new(fizz_iterator),
                iter_b: Box::new(buzz_iterator),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizzbuzz_exclusion(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut fizz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut buzz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Exclusion {
                iter_a: Box::new(fizz_iterator),
                iter_b: Box::new(buzz_iterator),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            };

            for doc_id in iterator {}
        });
    }
}
<commit_msg>Fixed typos in queryset benches<commit_after>use std::collections::VecDeque;

use search::term::Term;
use search::index::reader::IndexReader;
use search::query::Query;
use search::query::term_matcher::TermMatcher;

pub enum QuerySetIterator<'a, T: IndexReader<'a>> {
    None,
    All {
        iter: T::AllDocRefIterator,
    },
    Term {
        iter: T::TermDocRefIterator,
    },
    Conjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Disjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Exclusion {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
}

impl <'a, T: IndexReader<'a>> Iterator for QuerySetIterator<'a, T> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        match *self {
            QuerySetIterator::None => None,
            QuerySetIterator::All{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Term{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Conjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();

                                return Some(doc_id_a);
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            return None;
                        }
                    } else {
                        return None;
                    }
                }
            }
            QuerySetIterator::Disjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                if let Some(doc_id_a) = *current_doc_a {
                    if let Some(doc_id_b) = *current_doc_b {
                        if doc_id_a == doc_id_b {
                            *current_doc_a = iter_a.next();
                            *current_doc_b = iter_a.next();
                            Some(doc_id_a)
                        } else if doc_id_a < doc_id_b {
                            *current_doc_a = iter_a.next();
                            Some(doc_id_a)
                        } else {
                            *current_doc_b = iter_b.next();
                            Some(doc_id_b)
                        }
                    } else {
                        *current_doc_a = iter_a.next();
                        Some(doc_id_a)
                    }
                } else {
                    if let Some(doc_id_b) = *current_doc_b {
                        *current_doc_b = iter_b.next();
                        Some(doc_id_b)
                    } else {
                        None
                    }
                }
            }
            QuerySetIterator::Exclusion{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                                return Some(doc_id_a);
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            *current_doc_a = iter_a.next();
                            return Some(doc_id_a);
                        }
                    } else {
                        return None;
                    }
                }
            }
        }
    }
}

fn build_conjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T>  {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Conjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_conjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

fn build_disjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T> {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Disjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_disjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

pub fn build_iterator_from_query<'a, T: IndexReader<'a>>(reader: &'a T, query: &Query) -> QuerySetIterator<'a, T> {
    match *query {
        Query::MatchAll => {
            QuerySetIterator::All {
                iter: reader.iter_docids_all(),
            }
        }
        Query::MatchNone => {
            QuerySetIterator::None
        }
        Query::MatchTerm{ref field, ref term, ref matcher} => {
            match *matcher {
                TermMatcher::Exact => {
                    QuerySetIterator::Term {
                        iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                    }
                }
                TermMatcher::Prefix => {
                    // Find all terms in the index that match the prefix
                    let terms = match *term {
                         Term::String(ref term) => {
                             reader.iter_terms().filter_map(|k| {
                                 if let Term::String(ref k) = *k {
                                     if k.starts_with(term) {
                                         return Some(Term::String(k.clone()));
                                     }
                                 }

                                 None
                             }).collect::<Vec<Term>>()
                         }
                         _ => return QuerySetIterator::None,
                    };

                    match terms.len() {
                        0 => QuerySetIterator::None,
                        1 => {
                            QuerySetIterator::Term {
                                iter: reader.iter_docids_with_term(terms[0].clone(), field.clone()),
                            }
                        }
                        _ => {
                            // Produce a disjunction iterator for all the terms
                            let mut iters = VecDeque::new();
                            for term in terms.iter() {
                                iters.push_back(QuerySetIterator::Term {
                                    iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                                });
                            }

                            build_disjunction_iterator(iters)
                        }
                    }
                }
            }
        }
        Query::Conjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_conjunction_iterator(iters)
        }
        Query::Disjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::NDisjunction{ref queries, minimum_should_match} => {
            // TODO
            QuerySetIterator::None
        }
        Query::DisjunctionMax{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::Exclude{ref query, ref exclude} => {
            QuerySetIterator::Exclusion {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, exclude)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Filter{ref query, ref filter} => {
            QuerySetIterator::Conjunction {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, filter)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Score{ref query, mul, add} => {
            build_iterator_from_query(reader, query)
        }
    }
}

#[cfg(test)]
mod benches {
    use test::Bencher;

    use search::term::Term;
    use search::token::Token;
    use search::document::Document;
    use search::index::reader::IndexReader;
    use search::index::store::memory::MemoryIndexStore;
    use search::query_set::QuerySetIterator;

    fn make_test_store() -> MemoryIndexStore {
        let mut store = MemoryIndexStore::new();

        for i in 0..10000 {
            let mut tokens = Vec::new();

            if i % 3 == 0 {
                let position = tokens.len() as u32 + 1;
                tokens.push(Token {
                    term: Term::String("fizz".to_string()),
                    position: position,
                });
            }

            if i % 5 == 0 {
                let position = tokens.len() as u32 + 1;
                tokens.push(Token {
                    term: Term::String("buzz".to_string()),
                    position: position,
                });
            }

            store.insert_or_update_document(Document {
                key: i.to_string(),
                fields: btreemap! {
                    "body".to_string() => tokens
                }
            });
        }

        store
    }

    #[bench]
    fn bench_all(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::All {
                iter: store.iter_docids_all(),
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizz_term(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("fizz".to_string()), "body".to_string()),
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_buzz_term(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizzbuzz_conjunction(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut fizz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("fizz".to_string()), "body".to_string()),
            };
            let mut buzz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Conjunction {
                iter_a: Box::new(fizz_iterator),
                iter_b: Box::new(buzz_iterator),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizzbuzz_disjunction(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut fizz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("fizz".to_string()), "body".to_string()),
            };
            let mut buzz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Disjunction {
                iter_a: Box::new(fizz_iterator),
                iter_b: Box::new(buzz_iterator),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            };

            for doc_id in iterator {}
        });
    }

    #[bench]
    fn bench_fizzbuzz_exclusion(b: &mut Bencher) {
        let store = make_test_store();

        b.iter(|| {
            let mut fizz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("fizz".to_string()), "body".to_string()),
            };
            let mut buzz_iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Term {
                iter: store.iter_docids_with_term(Term::String("buzz".to_string()), "body".to_string()),
            };
            let mut iterator: QuerySetIterator<MemoryIndexStore> = QuerySetIterator::Exclusion {
                iter_a: Box::new(fizz_iterator),
                iter_b: Box::new(buzz_iterator),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            };

            for doc_id in iterator {}
        });
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::option::Option;
use std::sync::{Arc, RwLock, Mutex};

use raftserver::store::{Msg as StoreMsg, Transport, SendCh as StoreSendCh, Callback};
use raftserver::{Result as RaftResult, other as raft_other};
use kvproto::raft_serverpb::RaftMessage;
use kvproto::msgpb::{Message, MessageType};
use kvproto::raft_cmdpb::RaftCmdRequest;
use pd::PdClient;
use util::HandyRwLock;
use super::{SendCh, Msg, ConnData};

pub struct ServerTransport<T: PdClient> {
    cluster_id: u64,
    stores: HashMap<u64, StoreSendCh>,

    pd_client: Arc<RwLock<T>>,
    ch: SendCh,
    msg_id: Mutex<u64>,
}

impl<T: PdClient> ServerTransport<T> {
    pub fn new(cluster_id: u64, ch: SendCh, pd_client: Arc<RwLock<T>>) -> ServerTransport<T> {
        ServerTransport {
            cluster_id: cluster_id,
            stores: HashMap::new(),
            pd_client: pd_client.clone(),
            ch: ch,
            msg_id: Mutex::new(0),
        }
    }

    fn get_sendch(&self, store_id: u64) -> RaftResult<&StoreSendCh> {
        match self.stores.get(&store_id) {
            None => {
                Err(raft_other(format!("send message to invalid store {}, missing send \
                                   channel",
                                       store_id)))
            }

            Some(ch) => Ok(ch),
        }
    }
}

impl<T: PdClient> Transport for ServerTransport<T> {
    fn add_sendch(&mut self, store_id: u64, ch: StoreSendCh) {
        self.stores.insert(store_id, ch);
    }

    fn remove_sendch(&mut self, store_id: u64) -> Option<StoreSendCh> {
        self.stores.remove(&store_id)
    }

    fn send(&self, msg: RaftMessage) -> RaftResult<()> {
        let to_store_id = msg.get_to_peer().get_store_id();
        if let Some(ch) = self.stores.get(&to_store_id) {
            // use store send channel directly.
            return ch.send(StoreMsg::RaftMessage(msg));
        }

        let to_node_id = msg.get_to_peer().get_node_id();
        let node = try!(self.pd_client.rl().get_node(self.cluster_id, to_node_id));

        let mut req = Message::new();
        req.set_msg_type(MessageType::Raft);
        req.set_raft(msg);

        let mut id = self.msg_id.lock().unwrap();
        *id += 1;
        if let Err(e) = self.ch.send(Msg::SendPeer {
            addr: node.get_address().to_owned(),
            data: ConnData::new(*id, req),
        }) {
            return Err(raft_other(format!("send peer to {} err {:?}", node.get_address(), e)));
        }

        Ok(())
    }

    // Send RaftMessage to specified store, the store must exist in current node.
    // Unlike Transport trait Send, this function can only send message to local store.
    fn send_raft_msg(&self, msg: RaftMessage) -> RaftResult<()> {
        let to_store_id = msg.get_to_peer().get_store_id();
        let ch = try!(self.get_sendch(to_store_id));

        try!(ch.send(StoreMsg::RaftMessage(msg)));

        Ok(())
    }

    // Send RaftCmdRequest to specified store, the store must exist in current node.
    // Unlike Transport trait Send, this function can only send message to local store.
    fn send_command(&self, req: RaftCmdRequest, cb: Callback) -> RaftResult<()> {
        let to_store_id = req.get_header().get_peer().get_store_id();
        let ch = try!(self.get_sendch(to_store_id));

        try!(ch.send(StoreMsg::RaftCmd {
            request: req,
            callback: cb,
        }));

        Ok(())
    }
}
<commit_msg>server: Address comment.<commit_after>use std::collections::HashMap;
use std::option::Option;
use std::sync::{Arc, RwLock, Mutex};

use raftserver::store::{Msg as StoreMsg, Transport, SendCh as StoreSendCh, Callback};
use raftserver::{Result as RaftResult, other as raft_other};
use kvproto::raft_serverpb::RaftMessage;
use kvproto::msgpb::{Message, MessageType};
use kvproto::raft_cmdpb::RaftCmdRequest;
use pd::PdClient;
use util::HandyRwLock;
use super::{SendCh, Msg, ConnData};

pub struct ServerTransport<T: PdClient> {
    cluster_id: u64,
    stores: HashMap<u64, StoreSendCh>,

    pd_client: Arc<RwLock<T>>,
    ch: SendCh,
    msg_id: Mutex<u64>,
}

impl<T: PdClient> ServerTransport<T> {
    pub fn new(cluster_id: u64, ch: SendCh, pd_client: Arc<RwLock<T>>) -> ServerTransport<T> {
        ServerTransport {
            cluster_id: cluster_id,
            stores: HashMap::new(),
            pd_client: pd_client.clone(),
            ch: ch,
            msg_id: Mutex::new(0),
        }
    }

    fn get_sendch(&self, store_id: u64) -> RaftResult<&StoreSendCh> {
        match self.stores.get(&store_id) {
            None => {
                Err(raft_other(format!("send message to invalid store {}, missing send \
                                   channel",
                                       store_id)))
            }

            Some(ch) => Ok(ch),
        }
    }
}

impl<T: PdClient> Transport for ServerTransport<T> {
    fn add_sendch(&mut self, store_id: u64, ch: StoreSendCh) {
        self.stores.insert(store_id, ch);
    }

    fn remove_sendch(&mut self, store_id: u64) -> Option<StoreSendCh> {
        self.stores.remove(&store_id)
    }

    fn send(&self, msg: RaftMessage) -> RaftResult<()> {
        let to_store_id = msg.get_to_peer().get_store_id();
        if let Some(ch) = self.stores.get(&to_store_id) {
            // use store send channel directly.
            return ch.send(StoreMsg::RaftMessage(msg));
        }

        let to_node_id = msg.get_to_peer().get_node_id();
        let node = try!(self.pd_client.rl().get_node(self.cluster_id, to_node_id));

        let mut req = Message::new();
        req.set_msg_type(MessageType::Raft);
        req.set_raft(msg);

        let mut id = self.msg_id.lock().unwrap();
        *id += 1;

        self.ch
            .send(Msg::SendPeer {
                addr: node.get_address().to_owned(),
                data: ConnData::new(*id, req),
            })
            .map_err(|e| raft_other(format!("send peer to {} err {:?}", node.get_address(), e)))
    }

    // Send RaftMessage to specified store, the store must exist in current node.
    // Unlike Transport trait Send, this function can only send message to local store.
    fn send_raft_msg(&self, msg: RaftMessage) -> RaftResult<()> {
        let to_store_id = msg.get_to_peer().get_store_id();
        let ch = try!(self.get_sendch(to_store_id));

        try!(ch.send(StoreMsg::RaftMessage(msg)));

        Ok(())
    }

    // Send RaftCmdRequest to specified store, the store must exist in current node.
    // Unlike Transport trait Send, this function can only send message to local store.
    fn send_command(&self, req: RaftCmdRequest, cb: Callback) -> RaftResult<()> {
        let to_store_id = req.get_header().get_peer().get_store_id();
        let ch = try!(self.get_sendch(to_store_id));

        try!(ch.send(StoreMsg::RaftCmd {
            request: req,
            callback: cb,
        }));

        Ok(())
    }
}
<|endoftext|>


<commit_before>use super::Blob;
use failure::{Error, Fail};
use postgres::{transaction::Transaction, Connection};

#[derive(Debug, Fail)]
#[fail(display = "the path is not present in the database")]
struct PathNotFoundError;

pub(crate) struct DatabaseBackend<'a> {
    conn: &'a Connection,
}

impl<'a> DatabaseBackend<'a> {
    pub(crate) fn new(conn: &'a Connection) -> Self {
        Self { conn }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let rows = self.conn.query(
            "SELECT path, mime, date_updated, content FROM files WHERE path = $1;",
            &[&path],
        )?;

        if rows.is_empty() {
            Err(PathNotFoundError.into())
        } else {
            let row = rows.get(0);
            Ok(Blob {
                path: row.get("path"),
                mime: row.get("mime"),
                date_updated: row.get("date_updated"),
                content: row.get("content"),
            })
        }
    }
    pub(super) fn store_batch(&self, batch: &[Blob], trans: &Transaction) -> Result<(), Error> {
        for blob in batch {
            // check if file already exists in database
            let rows = self
                .conn
                .query("SELECT COUNT(*) FROM files WHERE path = $1", &[&blob.path])?;

            if rows.get(0).get::<usize, i64>(0) == 0 {
                trans.query(
                    "INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                    &[&blob.path, &blob.mime, &blob.content],
                )?;
            } else {
                trans.query(
                    "UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                                WHERE path = $1",
                    &[&blob.path, &blob.mime, &blob.content],
                )?;
            }
        }
        Ok(())
    }
}

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

    #[test]
    fn test_path_get() {
        crate::test::wrapper(|env| {
            let conn = env.db().conn();
            let backend = DatabaseBackend::new(&conn);

            // Add a test file to the database
            conn.execute(
                "INSERT INTO files (path, mime, date_updated, content) VALUES ($1, $2, $3, $4);",
                &[
                    &"dir/foo.txt",
                    &"text/plain",
                    &Timespec::new(42, 0),
                    &"Hello world!".as_bytes(),
                ],
            )?;

            // Test that the proper file was returned
            assert_eq!(
                Blob {
                    path: "dir/foo.txt".into(),
                    mime: "text/plain".into(),
                    date_updated: Timespec::new(42, 0),
                    content: "Hello world!".bytes().collect(),
                },
                backend.get("dir/foo.txt")?
            );

            // Test that other files are not returned
            assert!(backend
                .get("dir/bar.txt")
                .unwrap_err()
                .downcast_ref::<PathNotFoundError>()
                .is_some());
            assert!(backend
                .get("foo.txt")
                .unwrap_err()
                .downcast_ref::<PathNotFoundError>()
                .is_some());

            Ok(())
        });
    }
}
<commit_msg>Use `ON CONFLICT` instead of rewriting it manually<commit_after>use super::Blob;
use failure::{Error, Fail};
use postgres::{transaction::Transaction, Connection};

#[derive(Debug, Fail)]
#[fail(display = "the path is not present in the database")]
struct PathNotFoundError;

pub(crate) struct DatabaseBackend<'a> {
    conn: &'a Connection,
}

impl<'a> DatabaseBackend<'a> {
    pub(crate) fn new(conn: &'a Connection) -> Self {
        Self { conn }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let rows = self.conn.query(
            "SELECT path, mime, date_updated, content FROM files WHERE path = $1;",
            &[&path],
        )?;

        if rows.is_empty() {
            Err(PathNotFoundError.into())
        } else {
            let row = rows.get(0);
            Ok(Blob {
                path: row.get("path"),
                mime: row.get("mime"),
                date_updated: row.get("date_updated"),
                content: row.get("content"),
            })
        }
    }
    pub(super) fn store_batch(&self, batch: &[Blob], trans: &Transaction) -> Result<(), Error> {
        for blob in batch {
            trans.query(
                "INSERT INTO files (path, mime, content)
                 VALUES ($1, $2, $3)
                 ON CONFLICT (path) DO UPDATE
                    SET mime = EXCLUDED.mime, content = EXCLUDED.content",
                &[&blob.path, &blob.mime, &blob.content],
            )?;
        }
        Ok(())
    }
}

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

    #[test]
    fn test_path_get() {
        crate::test::wrapper(|env| {
            let conn = env.db().conn();
            let backend = DatabaseBackend::new(&conn);

            // Add a test file to the database
            conn.execute(
                "INSERT INTO files (path, mime, date_updated, content) VALUES ($1, $2, $3, $4);",
                &[
                    &"dir/foo.txt",
                    &"text/plain",
                    &Timespec::new(42, 0),
                    &"Hello world!".as_bytes(),
                ],
            )?;

            // Test that the proper file was returned
            assert_eq!(
                Blob {
                    path: "dir/foo.txt".into(),
                    mime: "text/plain".into(),
                    date_updated: Timespec::new(42, 0),
                    content: "Hello world!".bytes().collect(),
                },
                backend.get("dir/foo.txt")?
            );

            // Test that other files are not returned
            assert!(backend
                .get("dir/bar.txt")
                .unwrap_err()
                .downcast_ref::<PathNotFoundError>()
                .is_some());
            assert!(backend
                .get("foo.txt")
                .unwrap_err()
                .downcast_ref::<PathNotFoundError>()
                .is_some());

            Ok(())
        });
    }
}
<|endoftext|>


<commit_before>use backend::Backend;
use query_builder::*;
use result::QueryResult;
use super::Expression;
use types::BigInt;

/// Creates a SQL `COUNT` expression
///
/// As with most bare functions, this is not exported by default. You can import
/// it specifically as `diesel::expression::count`, or glob import
/// `diesel::dsl::*`
///
/// # Examples
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// # include!("../doctest_setup.rs");
/// # use diesel::dsl::*;
/// #
/// # table! {
/// #     users {
/// #         id -> Integer,
/// #         name -> VarChar,
/// #     }
/// # }
/// #
/// # fn main() {
/// #     use self::animals::dsl::*;
/// #     let connection = establish_connection();
/// assert_eq!(Ok(1), animals.select(count(name)).first(&connection));
/// # }
/// ```
pub fn count<T: Expression>(t: T) -> Count<T> {
    Count { target: t }
}

/// Creates a SQL `COUNT(*)` expression
///
/// For selecting the count of a query, and nothing else, you can just call
/// [`count`](https://docs.diesel.rs/diesel/prelude/trait.CountDsl.html#method.count)
/// on the query instead.
///
/// As with most bare functions, this is not exported by default. You can import
/// it specifically as `diesel::expression::count_star`, or glob import
/// `diesel::dsl::*`
///
/// # Examples
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// # include!("../doctest_setup.rs");
/// # use diesel::dsl::*;
/// #
/// # table! {
/// #     users {
/// #         id -> Integer,
/// #         name -> VarChar,
/// #     }
/// # }
/// #
/// # fn main() {
/// #     use self::users::dsl::*;
/// #     let connection = establish_connection();
/// assert_eq!(Ok(2), users.select(count_star()).first(&connection));
/// # }
/// ```
pub fn count_star() -> CountStar {
    CountStar
}

#[derive(Debug, Clone, Copy)]
#[doc(hidden)]
pub struct Count<T> {
    target: T,
}

impl<T: Expression> Expression for Count<T> {
    type SqlType = BigInt;
}

impl<T: QueryFragment<DB>, DB: Backend> QueryFragment<DB> for Count<T> {
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.push_sql("COUNT(");
        self.target.walk_ast(out.reborrow())?;
        out.push_sql(")");
        Ok(())
    }
}

impl_query_id!(Count<T>);
impl_selectable_expression!(Count<T>);

#[derive(Debug, Clone, Copy)]
#[doc(hidden)]
pub struct CountStar;

impl Expression for CountStar {
    type SqlType = BigInt;
}

impl<DB: Backend> QueryFragment<DB> for CountStar {
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.push_sql("COUNT(*)");
        Ok(())
    }
}

impl_query_id!(CountStar);
impl_selectable_expression!(CountStar);
<commit_msg>Docs: Fixes for dsl::count and dsl::count_star<commit_after>use backend::Backend;
use query_builder::*;
use result::QueryResult;
use super::Expression;
use types::BigInt;

/// Creates a SQL `COUNT` expression
///
/// As with most bare functions, this is not exported by default. You can import
/// it specifically as `diesel::dsl::count`, or glob import
/// `diesel::dsl::*`
///
/// # Examples
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// # include!("../doctest_setup.rs");
/// # use diesel::dsl::*;
/// #
/// # table! {
/// #     users {
/// #         id -> Integer,
/// #         name -> VarChar,
/// #     }
/// # }
/// #
/// # fn main() {
/// #     use self::animals::dsl::*;
/// #     let connection = establish_connection();
/// assert_eq!(Ok(1), animals.select(count(name)).first(&connection));
/// # }
/// ```
pub fn count<T: Expression>(t: T) -> Count<T> {
    Count { target: t }
}

/// Creates a SQL `COUNT(*)` expression
///
/// For selecting the count of a query, and nothing else, you can just call
/// [`count`](https://docs.diesel.rs/diesel/prelude/trait.CountDsl.html#method.count)
/// on the query instead.
///
/// As with most bare functions, this is not exported by default. You can import
/// it specifically as `diesel::dsl::count_star`, or glob import
/// `diesel::dsl::*`
///
/// # Examples
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// # include!("../doctest_setup.rs");
/// # use diesel::dsl::*;
/// #
/// # table! {
/// #     users {
/// #         id -> Integer,
/// #         name -> VarChar,
/// #     }
/// # }
/// #
/// # fn main() {
/// #     use self::users::dsl::*;
/// #     let connection = establish_connection();
/// assert_eq!(Ok(2), users.select(count_star()).first(&connection));
/// # }
/// ```
pub fn count_star() -> CountStar {
    CountStar
}

#[derive(Debug, Clone, Copy)]
#[doc(hidden)]
pub struct Count<T> {
    target: T,
}

impl<T: Expression> Expression for Count<T> {
    type SqlType = BigInt;
}

impl<T: QueryFragment<DB>, DB: Backend> QueryFragment<DB> for Count<T> {
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.push_sql("COUNT(");
        self.target.walk_ast(out.reborrow())?;
        out.push_sql(")");
        Ok(())
    }
}

impl_query_id!(Count<T>);
impl_selectable_expression!(Count<T>);

#[derive(Debug, Clone, Copy)]
#[doc(hidden)]
pub struct CountStar;

impl Expression for CountStar {
    type SqlType = BigInt;
}

impl<DB: Backend> QueryFragment<DB> for CountStar {
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.push_sql("COUNT(*)");
        Ok(())
    }
}

impl_query_id!(CountStar);
impl_selectable_expression!(CountStar);
<|endoftext|>


<commit_before>
use Renderable;
use context::Context;
use token::Token;
use LiquidOptions;
use template::Template;
use parser;
use lexer;
use error::{Result, Error};

use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};

struct Include {
    partial: Template,
}

impl Renderable for Include {
    fn render(&self, mut context: &mut Context) -> Result<Option<String>> {
        self.partial.render(&mut context)
    }
}

fn parse_partial<P: AsRef<Path>>(path: P, options: &LiquidOptions) -> Result<Template> {
    let file_system = options.file_system.clone().unwrap_or(PathBuf::new());
    let path = file_system.join(path);

    // check if file exists
    if !path.exists() {
        return Err(Error::from(&*format!("{:?} does not exist", path)));
    }

    let mut file = try!(File::open(path));

    let mut content = String::new();
    try!(file.read_to_string(&mut content));

    let tokens = try!(lexer::tokenize(&content));
    parser::parse(&tokens, &options).map(Template::new)
}

pub fn include_tag(_tag_name: &str,
                   arguments: &[Token],
                   options: &LiquidOptions)
                   -> Result<Box<Renderable>> {
    let mut args = arguments.iter();

    let path = match args.next() {
        Some(&Token::StringLiteral(ref path)) => path,

        Some(&Token::Identifier(_)) => {
            return Err(Error::from("This should be a string literal")); // TODO: better error message
        }

        arg => return Error::parser("Identifier", arg),
    };

    Ok(Box::new(Include { partial: try!(parse_partial(&path, &options)) }))
}

#[cfg(test)]
mod test {
    use context::Context;
    use Renderable;
    use parse;
    use error::Error;
    use LiquidOptions;
    use std::path::PathBuf;

    fn options() -> LiquidOptions {
        LiquidOptions {
            file_system: Some(PathBuf::from("tests/fixtures/input")),
            ..Default::default()
        }
    }

    #[test]
    fn include_tag() {
        let text = "{% include 'example.txt' %}";
        let template = parse(text, options()).unwrap();

        let mut context = Context::new();
        assert_eq!(template.render(&mut context).unwrap(),
                   Some("5 wot wot\n".to_owned()));
    }

    #[test]
    fn no_file() {
        let text = "{% include 'file_does_not_exist.liquid' %}";
        let output = parse(text, options());

        assert!(output.is_err());
        if let Err(Error::Other(val)) = output {
            assert_eq!(format!("{}", val),
                       "\"tests/fixtures/input/file_does_not_exist.liquid\" does not exist"
                           .to_owned());
        } else {
            assert!(false);
        }
    }
}
<commit_msg>removed identifier match arm<commit_after>
use Renderable;
use context::Context;
use token::Token;
use LiquidOptions;
use template::Template;
use parser;
use lexer;
use error::{Result, Error};

use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};

struct Include {
    partial: Template,
}

impl Renderable for Include {
    fn render(&self, mut context: &mut Context) -> Result<Option<String>> {
        self.partial.render(&mut context)
    }
}

fn parse_partial<P: AsRef<Path>>(path: P, options: &LiquidOptions) -> Result<Template> {
    let file_system = options.file_system.clone().unwrap_or(PathBuf::new());
    let path = file_system.join(path);

    // check if file exists
    if !path.exists() {
        return Err(Error::from(&*format!("{:?} does not exist", path)));
    }

    let mut file = try!(File::open(path));

    let mut content = String::new();
    try!(file.read_to_string(&mut content));

    let tokens = try!(lexer::tokenize(&content));
    parser::parse(&tokens, &options).map(Template::new)
}

pub fn include_tag(_tag_name: &str,
                   arguments: &[Token],
                   options: &LiquidOptions)
                   -> Result<Box<Renderable>> {
    let mut args = arguments.iter();

    let path = match args.next() {
        Some(&Token::StringLiteral(ref path)) => path,
        arg => return Error::parser("String Literal", arg),
    };

    Ok(Box::new(Include { partial: try!(parse_partial(&path, &options)) }))
}

#[cfg(test)]
mod test {
    use context::Context;
    use Renderable;
    use parse;
    use error::Error;
    use LiquidOptions;
    use std::path::PathBuf;

    fn options() -> LiquidOptions {
        LiquidOptions {
            file_system: Some(PathBuf::from("tests/fixtures/input")),
            ..Default::default()
        }
    }

    #[test]
    fn include_tag() {
        let text = "{% include 'example.txt' %}";
        let template = parse(text, options()).unwrap();

        let mut context = Context::new();
        assert_eq!(template.render(&mut context).unwrap(),
                   Some("5 wot wot\n".to_owned()));
    }

    #[test]
    fn no_file() {
        let text = "{% include 'file_does_not_exist.liquid' %}";
        let output = parse(text, options());

        assert!(output.is_err());
        if let Err(Error::Other(val)) = output {
            assert_eq!(format!("{}", val),
                       "\"tests/fixtures/input/file_does_not_exist.liquid\" does not exist"
                           .to_owned());
        } else {
            assert!(false);
        }
    }
}
<|endoftext|>


<commit_before>//! Checks for needless boolean results of if-else expressions
//!
//! This lint is **warn** by default

use crate::utils::sugg::Sugg;
use crate::utils::{in_macro, span_lint, span_lint_and_sugg};
use rustc::hir::*;
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::{declare_tool_lint, lint_array};
use rustc_errors::Applicability;
use syntax::ast::LitKind;
use syntax::source_map::Spanned;

/// **What it does:** Checks for expressions of the form `if c { true } else {
/// false }`
/// (or vice versa) and suggest using the condition directly.
///
/// **Why is this bad?** Redundant code.
///
/// **Known problems:** Maybe false positives: Sometimes, the two branches are
/// painstakingly documented (which we of course do not detect), so they *may*
/// have some value. Even then, the documentation can be rewritten to match the
/// shorter code.
///
/// **Example:**
/// ```rust
/// if x {
///     false
/// } else {
///     true
/// }
/// ```
declare_clippy_lint! {
    pub NEEDLESS_BOOL,
    complexity,
    "if-statements with plain booleans in the then- and else-clause, e.g. `if p { true } else { false }`"
}

/// **What it does:** Checks for expressions of the form `x == true`,
/// `x != true` and order comparisons such as `x < true` (or vice versa) and
/// suggest using the variable directly.
///
/// **Why is this bad?** Unnecessary code.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// if x == true {} // could be `if x { }`
/// ```
declare_clippy_lint! {
    pub BOOL_COMPARISON,
    complexity,
    "comparing a variable to a boolean, e.g. `if x == true` or `if x != true`"
}

#[derive(Copy, Clone)]
pub struct NeedlessBool;

impl LintPass for NeedlessBool {
    fn get_lints(&self) -> LintArray {
        lint_array!(NEEDLESS_BOOL)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for NeedlessBool {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, e: &'tcx Expr) {
        use self::Expression::*;
        if let ExprKind::If(ref pred, ref then_block, Some(ref else_expr)) = e.node {

            let reduce = |ret, not| {
                let mut applicability = Applicability::MachineApplicable;
                let snip = Sugg::hir_with_applicability(cx, pred, "<predicate>", &mut applicability);
                let snip = if not { !snip } else { snip };

                let mut hint = if ret {
                    format!("return {}", snip)
                } else {
                    snip.to_string()
                };

                if parent_node_is_if_expr(&e, &cx) {
                    hint = format!("{{ {} }}", hint);
                }

                span_lint_and_sugg(
                    cx,
                    NEEDLESS_BOOL,
                    e.span,
                    "this if-then-else expression returns a bool literal",
                    "you can reduce it to",
                    hint,
                    applicability,
                );
            };
            if let ExprKind::Block(ref then_block, _) = then_block.node {
                match (fetch_bool_block(then_block), fetch_bool_expr(else_expr)) {
                    (RetBool(true), RetBool(true)) | (Bool(true), Bool(true)) => {
                        span_lint(
                            cx,
                            NEEDLESS_BOOL,
                            e.span,
                            "this if-then-else expression will always return true",
                        );
                    },
                    (RetBool(false), RetBool(false)) | (Bool(false), Bool(false)) => {
                        span_lint(
                            cx,
                            NEEDLESS_BOOL,
                            e.span,
                            "this if-then-else expression will always return false",
                        );
                    },
                    (RetBool(true), RetBool(false)) => reduce(true, false),
                    (Bool(true), Bool(false)) => reduce(false, false),
                    (RetBool(false), RetBool(true)) => reduce(true, true),
                    (Bool(false), Bool(true)) => reduce(false, true),
                    _ => (),
                }
            } else {
                panic!("IfExpr 'then' node is not an ExprKind::Block");
            }
        }
    }
}

fn parent_node_is_if_expr<'a, 'b>(expr: &Expr, cx: &LateContext<'a, 'b>) -> bool {
    let parent_id = cx.tcx.hir().get_parent_node(expr.id);
    let parent_node = cx.tcx.hir().get(parent_id);

    if let rustc::hir::Node::Expr(e) = parent_node {
        if let ExprKind::If(_,_,_) = e.node {
            return true;
        }
    }

    false
}

#[derive(Copy, Clone)]
pub struct BoolComparison;

impl LintPass for BoolComparison {
    fn get_lints(&self) -> LintArray {
        lint_array!(BOOL_COMPARISON)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for BoolComparison {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, e: &'tcx Expr) {
        if in_macro(e.span) {
            return;
        }

        if let ExprKind::Binary(Spanned { node, .. }, ..) = e.node {
            let ignore_case = None::<(fn(_) -> _, &str)>;
            let ignore_no_literal = None::<(fn(_, _) -> _, &str)>;
            match node {
                BinOpKind::Eq => {
                    let true_case = Some((|h| h, "equality checks against true are unnecessary"));
                    let false_case = Some((
                        |h: Sugg<'_>| !h,
                        "equality checks against false can be replaced by a negation",
                    ));
                    check_comparison(cx, e, true_case, false_case, true_case, false_case, ignore_no_literal)
                },
                BinOpKind::Ne => {
                    let true_case = Some((
                        |h: Sugg<'_>| !h,
                        "inequality checks against true can be replaced by a negation",
                    ));
                    let false_case = Some((|h| h, "inequality checks against false are unnecessary"));
                    check_comparison(cx, e, true_case, false_case, true_case, false_case, ignore_no_literal)
                },
                BinOpKind::Lt => check_comparison(
                    cx,
                    e,
                    ignore_case,
                    Some((|h| h, "greater than checks against false are unnecessary")),
                    Some((
                        |h: Sugg<'_>| !h,
                        "less than comparison against true can be replaced by a negation",
                    )),
                    ignore_case,
                    Some((
                        |l: Sugg<'_>, r: Sugg<'_>| (!l).bit_and(&r),
                        "order comparisons between booleans can be simplified",
                    )),
                ),
                BinOpKind::Gt => check_comparison(
                    cx,
                    e,
                    Some((
                        |h: Sugg<'_>| !h,
                        "less than comparison against true can be replaced by a negation",
                    )),
                    ignore_case,
                    ignore_case,
                    Some((|h| h, "greater than checks against false are unnecessary")),
                    Some((
                        |l: Sugg<'_>, r: Sugg<'_>| l.bit_and(&(!r)),
                        "order comparisons between booleans can be simplified",
                    )),
                ),
                _ => (),
            }
        }
    }
}

fn check_comparison<'a, 'tcx>(
    cx: &LateContext<'a, 'tcx>,
    e: &'tcx Expr,
    left_true: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    left_false: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    right_true: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    right_false: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    no_literal: Option<(impl FnOnce(Sugg<'a>, Sugg<'a>) -> Sugg<'a>, &str)>,
) {
    use self::Expression::*;

    if let ExprKind::Binary(_, ref left_side, ref right_side) = e.node {
        let mut applicability = Applicability::MachineApplicable;
        match (fetch_bool_expr(left_side), fetch_bool_expr(right_side)) {
            (Bool(true), Other) => left_true.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, right_side, applicability, m, h)
            }),
            (Other, Bool(true)) => right_true.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, left_side, applicability, m, h)
            }),
            (Bool(false), Other) => left_false.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, right_side, applicability, m, h)
            }),
            (Other, Bool(false)) => right_false.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, left_side, applicability, m, h)
            }),
            (Other, Other) => no_literal.map_or((), |(h, m)| {
                let (l_ty, r_ty) = (cx.tables.expr_ty(left_side), cx.tables.expr_ty(right_side));
                if l_ty.is_bool() && r_ty.is_bool() {
                    let left_side = Sugg::hir_with_applicability(cx, left_side, "..", &mut applicability);
                    let right_side = Sugg::hir_with_applicability(cx, right_side, "..", &mut applicability);
                    span_lint_and_sugg(
                        cx,
                        BOOL_COMPARISON,
                        e.span,
                        m,
                        "try simplifying it as shown",
                        h(left_side, right_side).to_string(),
                        applicability,
                    )
                }
            }),
            _ => (),
        }
    }
}

fn suggest_bool_comparison<'a, 'tcx>(
    cx: &LateContext<'a, 'tcx>,
    e: &'tcx Expr,
    expr: &Expr,
    mut applicability: Applicability,
    message: &str,
    conv_hint: impl FnOnce(Sugg<'a>) -> Sugg<'a>,
) {
    let hint = Sugg::hir_with_applicability(cx, expr, "..", &mut applicability);
    span_lint_and_sugg(
        cx,
        BOOL_COMPARISON,
        e.span,
        message,
        "try simplifying it as shown",
        conv_hint(hint).to_string(),
        applicability,
    );
}

enum Expression {
    Bool(bool),
    RetBool(bool),
    Other,
}

fn fetch_bool_block(block: &Block) -> Expression {
    match (&*block.stmts, block.expr.as_ref()) {
        (&[], Some(e)) => fetch_bool_expr(&**e),
        (&[ref e], None) => {
            if let StmtKind::Semi(ref e, _) = e.node {
                if let ExprKind::Ret(_) = e.node {
                    fetch_bool_expr(&**e)
                } else {
                    Expression::Other
                }
            } else {
                Expression::Other
            }
        },
        _ => Expression::Other,
    }
}

fn fetch_bool_expr(expr: &Expr) -> Expression {
    match expr.node {
        ExprKind::Block(ref block, _) => fetch_bool_block(block),
        ExprKind::Lit(ref lit_ptr) => {
            if let LitKind::Bool(value) = lit_ptr.node {
                Expression::Bool(value)
            } else {
                Expression::Other
            }
        },
        ExprKind::Ret(Some(ref expr)) => match fetch_bool_expr(expr) {
            Expression::Bool(value) => Expression::RetBool(value),
            _ => Expression::Other,
        },
        _ => Expression::Other,
    }
}
<commit_msg>Update clippy_lints/src/needless_bool.rs<commit_after>//! Checks for needless boolean results of if-else expressions
//!
//! This lint is **warn** by default

use crate::utils::sugg::Sugg;
use crate::utils::{in_macro, span_lint, span_lint_and_sugg};
use rustc::hir::*;
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::{declare_tool_lint, lint_array};
use rustc_errors::Applicability;
use syntax::ast::LitKind;
use syntax::source_map::Spanned;

/// **What it does:** Checks for expressions of the form `if c { true } else {
/// false }`
/// (or vice versa) and suggest using the condition directly.
///
/// **Why is this bad?** Redundant code.
///
/// **Known problems:** Maybe false positives: Sometimes, the two branches are
/// painstakingly documented (which we of course do not detect), so they *may*
/// have some value. Even then, the documentation can be rewritten to match the
/// shorter code.
///
/// **Example:**
/// ```rust
/// if x {
///     false
/// } else {
///     true
/// }
/// ```
declare_clippy_lint! {
    pub NEEDLESS_BOOL,
    complexity,
    "if-statements with plain booleans in the then- and else-clause, e.g. `if p { true } else { false }`"
}

/// **What it does:** Checks for expressions of the form `x == true`,
/// `x != true` and order comparisons such as `x < true` (or vice versa) and
/// suggest using the variable directly.
///
/// **Why is this bad?** Unnecessary code.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// if x == true {} // could be `if x { }`
/// ```
declare_clippy_lint! {
    pub BOOL_COMPARISON,
    complexity,
    "comparing a variable to a boolean, e.g. `if x == true` or `if x != true`"
}

#[derive(Copy, Clone)]
pub struct NeedlessBool;

impl LintPass for NeedlessBool {
    fn get_lints(&self) -> LintArray {
        lint_array!(NEEDLESS_BOOL)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for NeedlessBool {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, e: &'tcx Expr) {
        use self::Expression::*;
        if let ExprKind::If(ref pred, ref then_block, Some(ref else_expr)) = e.node {

            let reduce = |ret, not| {
                let mut applicability = Applicability::MachineApplicable;
                let snip = Sugg::hir_with_applicability(cx, pred, "<predicate>", &mut applicability);
                let snip = if not { !snip } else { snip };

                let mut hint = if ret {
                    format!("return {}", snip)
                } else {
                    snip.to_string()
                };

                if parent_node_is_if_expr(&e, &cx) {
                    hint = format!("{{ {} }}", hint);
                }

                span_lint_and_sugg(
                    cx,
                    NEEDLESS_BOOL,
                    e.span,
                    "this if-then-else expression returns a bool literal",
                    "you can reduce it to",
                    hint,
                    applicability,
                );
            };
            if let ExprKind::Block(ref then_block, _) = then_block.node {
                match (fetch_bool_block(then_block), fetch_bool_expr(else_expr)) {
                    (RetBool(true), RetBool(true)) | (Bool(true), Bool(true)) => {
                        span_lint(
                            cx,
                            NEEDLESS_BOOL,
                            e.span,
                            "this if-then-else expression will always return true",
                        );
                    },
                    (RetBool(false), RetBool(false)) | (Bool(false), Bool(false)) => {
                        span_lint(
                            cx,
                            NEEDLESS_BOOL,
                            e.span,
                            "this if-then-else expression will always return false",
                        );
                    },
                    (RetBool(true), RetBool(false)) => reduce(true, false),
                    (Bool(true), Bool(false)) => reduce(false, false),
                    (RetBool(false), RetBool(true)) => reduce(true, true),
                    (Bool(false), Bool(true)) => reduce(false, true),
                    _ => (),
                }
            } else {
                panic!("IfExpr 'then' node is not an ExprKind::Block");
            }
        }
    }
}

fn parent_node_is_if_expr<'a, 'b>(expr: &Expr, cx: &LateContext<'a, 'b>) -> bool {
    let parent_id = cx.tcx.hir().get_parent_node(expr.id);
    let parent_node = cx.tcx.hir().get(parent_id);

    if let rustc::hir::Node::Expr(e) = parent_node {
        if let ExprKind::If(_, _, _) = e.node {
            return true;
        }
    }

    false
}

#[derive(Copy, Clone)]
pub struct BoolComparison;

impl LintPass for BoolComparison {
    fn get_lints(&self) -> LintArray {
        lint_array!(BOOL_COMPARISON)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for BoolComparison {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, e: &'tcx Expr) {
        if in_macro(e.span) {
            return;
        }

        if let ExprKind::Binary(Spanned { node, .. }, ..) = e.node {
            let ignore_case = None::<(fn(_) -> _, &str)>;
            let ignore_no_literal = None::<(fn(_, _) -> _, &str)>;
            match node {
                BinOpKind::Eq => {
                    let true_case = Some((|h| h, "equality checks against true are unnecessary"));
                    let false_case = Some((
                        |h: Sugg<'_>| !h,
                        "equality checks against false can be replaced by a negation",
                    ));
                    check_comparison(cx, e, true_case, false_case, true_case, false_case, ignore_no_literal)
                },
                BinOpKind::Ne => {
                    let true_case = Some((
                        |h: Sugg<'_>| !h,
                        "inequality checks against true can be replaced by a negation",
                    ));
                    let false_case = Some((|h| h, "inequality checks against false are unnecessary"));
                    check_comparison(cx, e, true_case, false_case, true_case, false_case, ignore_no_literal)
                },
                BinOpKind::Lt => check_comparison(
                    cx,
                    e,
                    ignore_case,
                    Some((|h| h, "greater than checks against false are unnecessary")),
                    Some((
                        |h: Sugg<'_>| !h,
                        "less than comparison against true can be replaced by a negation",
                    )),
                    ignore_case,
                    Some((
                        |l: Sugg<'_>, r: Sugg<'_>| (!l).bit_and(&r),
                        "order comparisons between booleans can be simplified",
                    )),
                ),
                BinOpKind::Gt => check_comparison(
                    cx,
                    e,
                    Some((
                        |h: Sugg<'_>| !h,
                        "less than comparison against true can be replaced by a negation",
                    )),
                    ignore_case,
                    ignore_case,
                    Some((|h| h, "greater than checks against false are unnecessary")),
                    Some((
                        |l: Sugg<'_>, r: Sugg<'_>| l.bit_and(&(!r)),
                        "order comparisons between booleans can be simplified",
                    )),
                ),
                _ => (),
            }
        }
    }
}

fn check_comparison<'a, 'tcx>(
    cx: &LateContext<'a, 'tcx>,
    e: &'tcx Expr,
    left_true: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    left_false: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    right_true: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    right_false: Option<(impl FnOnce(Sugg<'a>) -> Sugg<'a>, &str)>,
    no_literal: Option<(impl FnOnce(Sugg<'a>, Sugg<'a>) -> Sugg<'a>, &str)>,
) {
    use self::Expression::*;

    if let ExprKind::Binary(_, ref left_side, ref right_side) = e.node {
        let mut applicability = Applicability::MachineApplicable;
        match (fetch_bool_expr(left_side), fetch_bool_expr(right_side)) {
            (Bool(true), Other) => left_true.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, right_side, applicability, m, h)
            }),
            (Other, Bool(true)) => right_true.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, left_side, applicability, m, h)
            }),
            (Bool(false), Other) => left_false.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, right_side, applicability, m, h)
            }),
            (Other, Bool(false)) => right_false.map_or((), |(h, m)| {
                suggest_bool_comparison(cx, e, left_side, applicability, m, h)
            }),
            (Other, Other) => no_literal.map_or((), |(h, m)| {
                let (l_ty, r_ty) = (cx.tables.expr_ty(left_side), cx.tables.expr_ty(right_side));
                if l_ty.is_bool() && r_ty.is_bool() {
                    let left_side = Sugg::hir_with_applicability(cx, left_side, "..", &mut applicability);
                    let right_side = Sugg::hir_with_applicability(cx, right_side, "..", &mut applicability);
                    span_lint_and_sugg(
                        cx,
                        BOOL_COMPARISON,
                        e.span,
                        m,
                        "try simplifying it as shown",
                        h(left_side, right_side).to_string(),
                        applicability,
                    )
                }
            }),
            _ => (),
        }
    }
}

fn suggest_bool_comparison<'a, 'tcx>(
    cx: &LateContext<'a, 'tcx>,
    e: &'tcx Expr,
    expr: &Expr,
    mut applicability: Applicability,
    message: &str,
    conv_hint: impl FnOnce(Sugg<'a>) -> Sugg<'a>,
) {
    let hint = Sugg::hir_with_applicability(cx, expr, "..", &mut applicability);
    span_lint_and_sugg(
        cx,
        BOOL_COMPARISON,
        e.span,
        message,
        "try simplifying it as shown",
        conv_hint(hint).to_string(),
        applicability,
    );
}

enum Expression {
    Bool(bool),
    RetBool(bool),
    Other,
}

fn fetch_bool_block(block: &Block) -> Expression {
    match (&*block.stmts, block.expr.as_ref()) {
        (&[], Some(e)) => fetch_bool_expr(&**e),
        (&[ref e], None) => {
            if let StmtKind::Semi(ref e, _) = e.node {
                if let ExprKind::Ret(_) = e.node {
                    fetch_bool_expr(&**e)
                } else {
                    Expression::Other
                }
            } else {
                Expression::Other
            }
        },
        _ => Expression::Other,
    }
}

fn fetch_bool_expr(expr: &Expr) -> Expression {
    match expr.node {
        ExprKind::Block(ref block, _) => fetch_bool_block(block),
        ExprKind::Lit(ref lit_ptr) => {
            if let LitKind::Bool(value) = lit_ptr.node {
                Expression::Bool(value)
            } else {
                Expression::Other
            }
        },
        ExprKind::Ret(Some(ref expr)) => match fetch_bool_expr(expr) {
            Expression::Bool(value) => Expression::RetBool(value),
            _ => Expression::Other,
        },
        _ => Expression::Other,
    }
}
<|endoftext|>


<commit_before>mod terminal;
mod ansi;

use std::io::{self, Write, BufWriter};
use self::terminal::{Terminal, Event, Key};
use self::ansi::{clear, color, cursor, style};
use rayon::prelude::*;
use rff::choice::Choice;

#[derive(Debug)]
pub enum Error {
    Exit,
    Write(io::Error),
    Reset(terminal::Error)
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Write(err)
    }
}

impl From<terminal::Error> for Error {
    fn from(err: terminal::Error) -> Error {
        Error::Reset(err)
    }
}

pub struct Options {
    pub choices: Vec<String>,
    pub initial: String,
}

pub struct Interface {
    choices: Vec<String>,
    matching: Vec<Choice>,
    selected: usize,
    search: String,
    width: usize,
    terminal: Terminal,
}

impl Interface {
    /// Creates a new Interface from the provided options.
    pub fn from_opts(opts: Options) -> Interface {
        let mut term = Terminal::from("/dev/tty").unwrap();
        term.set_raw_mode().unwrap();

        Interface {
            choices: opts.choices,
            matching: Vec::new(),
            selected: 0,
            search: opts.initial,
            width: term.max_width,
            terminal: term
        }
    }

    // Starts the interface
    pub fn run(&mut self) -> Result<&str, Error> {
        self.filter_choices();
        self.render()?;

        for event in self.terminal.events()? {
            if let Event::Key(key) = event? {
                match key {
                    Key::Ctrl('c') | Key::Ctrl('d') => {
                        self.clear()?;
                        return Err(Error::Exit);
                    }

                    Key::Char('\n') => {
                        break;
                    },

                    Key::Ctrl('n') => {
                        self.selected += 1;
                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Ctrl('p') => {
                        self.selected = self.selected.
                            checked_sub(1).
                            unwrap_or(0);

                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Char(ch) => {
                        self.search.push(ch);
                        self.filter_existing_matches();
                        self.render()?;
                    },

                    Key::Backspace | Key::Ctrl('h') => {
                        self.search.pop();
                        self.filter_choices();
                        self.render()?;
                    }

                    Key::Ctrl('u') => {
                        self.search.clear();
                        self.filter_choices();
                        self.render()?;
                    }

                    _ => {}
                }
            };
        }

        self.clear()?;
        Ok(self.result())
    }

    fn filter_choices(&mut self) {
        let mut matches = self.choices.
            par_iter().
            cloned().
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn filter_existing_matches(&mut self) {
        let mut matches = self.matching.
            par_iter().
            map(|c| c.text().to_string()).
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn prompt(&self) -> String {
        format!("> {}", self.search)
    }

    fn render(&mut self) -> io::Result<()> {
        let prompt = self.prompt();
        let matches = self.matching.iter().take(10);
        let n = matches.len() as u16;

        let mut term = BufWriter::new(&mut self.terminal);

        write!(term, "{}{}{}", cursor::Column(1), clear::Screen, prompt)?;

        for (i, choice) in matches.enumerate() {
            write!(term, "\r\n")?;

            let selected = i == self.selected;

            if selected {
                write!(term, "{}", style::Invert)?;
                print_choice(&mut term, choice, self.width)?;
                write!(term, "{}", style::NoInvert)?;
            } else {
                print_choice(&mut term, choice, self.width)?;
            }
        }

        if n > 0 {
            let col = (prompt.len() + 1) as u16;
            write!(term, "{}{}", cursor::Up(n), cursor::Column(col))?;
        }

        Ok(())
    }

    #[inline(always)]
    fn clamp_selected(&mut self) {
        let mut max = self.matching.len();
        if max > 10 { max = 10; }

        if self.selected >= max {
            self.selected = max - 1;
        }
    }

    fn clear(&mut self) -> Result<(), Error> {
        write!(self.terminal, "{}{}", cursor::Column(1), clear::Screen)?;
        self.terminal.reset()?;
        Ok(())
    }

    fn result(&mut self) -> &str {
        self.matching.iter().
            map(|c| c.text()).
            nth(self.selected).
            unwrap_or(&self.search)
    }
}

fn print_choice(term: &mut BufWriter<&mut Terminal>, choice: &Choice, max_width: usize) -> io::Result<()> {
    let chars = choice.text().chars().take(max_width);

    if let Some(positions) = choice.positions() {
        for (i, ch) in chars.enumerate() {
            let is_match = positions.iter().any(|p| *p == i);

            if is_match {
                let color = color::Fg(color::Colors::Magenta);
                let reset = color::Fg(color::Reset);
                write!(term, "{}{}{}", color, ch, reset)?;
            } else {
                write!(term, "{}", ch)?;
            }
        }
    } else {
        write!(term, "{}", chars.collect::<String>())?;
    }

    Ok(())
}
<commit_msg>Print number of matches in prompt<commit_after>mod terminal;
mod ansi;

use std::io::{self, Write, BufWriter};
use self::terminal::{Terminal, Event, Key};
use self::ansi::{clear, color, cursor, style};
use rayon::prelude::*;
use rff::choice::Choice;

#[derive(Debug)]
pub enum Error {
    Exit,
    Write(io::Error),
    Reset(terminal::Error)
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Write(err)
    }
}

impl From<terminal::Error> for Error {
    fn from(err: terminal::Error) -> Error {
        Error::Reset(err)
    }
}

pub struct Options {
    pub choices: Vec<String>,
    pub initial: String,
}

pub struct Interface {
    choices: Vec<String>,
    choices_width: usize,
    matching: Vec<Choice>,
    selected: usize,
    search: String,
    width: usize,
    terminal: Terminal,
}

impl Interface {
    /// Creates a new Interface from the provided options.
    pub fn from_opts(opts: Options) -> Interface {
        let mut term = Terminal::from("/dev/tty").unwrap();
        term.set_raw_mode().unwrap();

        let choices_width = format!("{}", opts.choices.len()).len();

        Interface {
            choices: opts.choices,
            choices_width: choices_width,
            matching: Vec::new(),
            selected: 0,
            search: opts.initial,
            width: term.max_width,
            terminal: term
        }
    }

    // Starts the interface
    pub fn run(&mut self) -> Result<&str, Error> {
        self.filter_choices();
        self.render()?;

        for event in self.terminal.events()? {
            if let Event::Key(key) = event? {
                match key {
                    Key::Ctrl('c') | Key::Ctrl('d') => {
                        self.clear()?;
                        return Err(Error::Exit);
                    }

                    Key::Char('\n') => {
                        break;
                    },

                    Key::Ctrl('n') => {
                        self.selected += 1;
                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Ctrl('p') => {
                        self.selected = self.selected.
                            checked_sub(1).
                            unwrap_or(0);

                        self.clamp_selected();
                        self.render()?;
                    },

                    Key::Char(ch) => {
                        self.search.push(ch);
                        self.filter_existing_matches();
                        self.render()?;
                    },

                    Key::Backspace | Key::Ctrl('h') => {
                        self.search.pop();
                        self.filter_choices();
                        self.render()?;
                    }

                    Key::Ctrl('u') => {
                        self.search.clear();
                        self.filter_choices();
                        self.render()?;
                    }

                    _ => {}
                }
            };
        }

        self.clear()?;
        Ok(self.result())
    }

    fn filter_choices(&mut self) {
        let mut matches = self.choices.
            par_iter().
            cloned().
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn filter_existing_matches(&mut self) {
        let mut matches = self.matching.
            par_iter().
            map(|c| c.text().to_string()).
            filter_map(|choice| Choice::with_positions(&self.search, choice)).
            collect::<Vec<_>>();

        matches.sort_by(|a, b| b.partial_cmp(a).unwrap());

        self.matching = matches;
    }

    fn prompt(&self) -> String {
        let count = self.matching.len();
        format!("{:width$} > {}", count, self.search, width = self.choices_width)
    }

    fn render(&mut self) -> io::Result<()> {
        let prompt = self.prompt();
        let matches = self.matching.iter().take(10);
        let n = matches.len() as u16;

        let mut term = BufWriter::new(&mut self.terminal);

        write!(term, "{}{}{}", cursor::Column(1), clear::Screen, prompt)?;

        for (i, choice) in matches.enumerate() {
            write!(term, "\r\n")?;

            let selected = i == self.selected;

            if selected {
                write!(term, "{}", style::Invert)?;
                print_choice(&mut term, choice, self.width)?;
                write!(term, "{}", style::NoInvert)?;
            } else {
                print_choice(&mut term, choice, self.width)?;
            }
        }

        if n > 0 {
            let col = (prompt.len() + 1) as u16;
            write!(term, "{}{}", cursor::Up(n), cursor::Column(col))?;
        }

        Ok(())
    }

    #[inline(always)]
    fn clamp_selected(&mut self) {
        let mut max = self.matching.len();
        if max > 10 { max = 10; }

        if self.selected >= max {
            self.selected = max - 1;
        }
    }

    fn clear(&mut self) -> Result<(), Error> {
        write!(self.terminal, "{}{}", cursor::Column(1), clear::Screen)?;
        self.terminal.reset()?;
        Ok(())
    }

    fn result(&mut self) -> &str {
        self.matching.iter().
            map(|c| c.text()).
            nth(self.selected).
            unwrap_or(&self.search)
    }
}

fn print_choice(term: &mut BufWriter<&mut Terminal>, choice: &Choice, max_width: usize) -> io::Result<()> {
    let chars = choice.text().chars().take(max_width);

    if let Some(positions) = choice.positions() {
        for (i, ch) in chars.enumerate() {
            let is_match = positions.iter().any(|p| *p == i);

            if is_match {
                let color = color::Fg(color::Colors::Magenta);
                let reset = color::Fg(color::Reset);
                write!(term, "{}{}{}", color, ch, reset)?;
            } else {
                write!(term, "{}", ch)?;
            }
        }
    } else {
        write!(term, "{}", chars.collect::<String>())?;
    }

    Ok(())
}
<|endoftext|>