<commit_before>use super::namepart::{Location, NamePart};
use super::suffix;
use super::surname;
use super::title;
use super::utils::is_mixed_case;
use smallvec::SmallVec;

struct ParseOp<'a> {
    words: SmallVec<[NamePart<'a>; 7]>,
    surname_index: usize,
    generation_from_suffix: Option<usize>,
    possible_false_prefix: Option<NamePart<'a>>,
    possible_false_postfix: Option<NamePart<'a>>,
    use_capitalization: bool,
}

pub fn parse(name: &str) -> Option<(SmallVec<[NamePart; 7]>, usize, Option<usize>)> {
    let mut op = ParseOp {
        words: SmallVec::new(),
        surname_index: 0,
        generation_from_suffix: None,
        possible_false_prefix: None,
        possible_false_postfix: None,
        use_capitalization: is_mixed_case(name),
    };

    if op.run(name) {
        Some((op.words, op.surname_index, op.generation_from_suffix))
    } else {
        None
    }
}

impl<'a> ParseOp<'a> {
    fn run(&mut self, name: &'a str) -> bool {
        // Separate comma-separated titles and suffixes, then flip remaining words
        // around remaining comma, if any
        if name.contains(',') {
            for part in name.split(',') {
                if self.words.is_empty() {
                    // We're in the surname part (if the format is "Smith, John"),
                    // or the only actual name part (if the format is "John Smith,
                    // esq." or just "John Smith")
                    self.handle_before_comma(part, true);
                } else if self.surname_index == 0 {
                    // We already processed one comma-separated part, but we think
                    // it was just the surname, so this might be the given name or
                    // initials
                    self.handle_after_comma(part);
                } else {
                    // We already found the full name, so this is a comma-separated
                    // postfix title or suffix
                    self.handle_after_surname(part);
                }
            }
        } else {
            self.handle_before_comma(name, false);
        }

        // If there are two or fewer words, e.g. "JOHN MA", we treat
        // ambiguous strings like "MA" as surnames rather than titles
        // (not initials, which should only be at the end of the input
        // if they are comma-separated, and we already handled that case)
        if !self.valid() {
            if let Some(ref postfix) = self.possible_false_postfix {
                self.words.push(postfix.clone());
            } else if let Some(ref prefix) = self.possible_false_prefix {
                self.words.insert(0, prefix.clone());
            }
        }

        // Anything trailing that looks like initials is probably a stray postfix
        while self.words.last().filter(|w| !w.is_namelike()).is_some() {
            let removed = self.words.pop().unwrap();

            // If we guessed the surname previously, our guess is no longer valid
            self.surname_index = 0;

            // If the alternative is having less than two words and giving up,
            // check if we can treat the last word as a name if we just ignore
            // case; this handles the not-quite-rare-enough case of an all-caps
            // last name (e.g.Neto John SMITH), among others
            if self.use_capitalization && !self.valid() {
                let word = NamePart::from_word(removed.word, false, Location::End);
                if word.is_namelike() {
                    self.words.push(word);
                    break;
                }
            }
        }

        // Handle case where we thought the whole before-comma part was a surname,
        // but we never found a plausible given name or initial afterwards,
        // as well as the reset just above
        if self.surname_index == 0 && self.words.len() > 1 {
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        }

        // Check the plausibility of what we've found
        self.valid()
    }

    fn valid(&self) -> bool {
        self.words.len() >= 2
            && self
                .words
                .iter()
                .all(|w| w.is_namelike() || w.is_initials())
            && self.words[self.surname_index..]
                .iter()
                .any(|w| w.is_namelike())
    }

    // Called only until any words are found
    fn handle_before_comma(&mut self, part: &'a str, any_after: bool) {
        debug_assert!(
            self.words.is_empty()
                && self.surname_index == 0
                && self.possible_false_prefix.is_none(),
            "Invalid state for handle_before_comma!"
        );

        self.words.extend(NamePart::all_from_text(
            part,
            self.use_capitalization,
            if any_after {
                Location::End
            } else {
                Location::Start
            },
        ));

        if self.words.is_empty() {
            return;
        }

        // Check for title as prefix (e.g. "Dr. John Smith" or "Right Hon.
        // John Smith")
        let prefix_title_len = if any_after || self.words.len() > 2 {
            title::find_prefix_len(&self.words)
        } else {
            0
        };
        for i in (0..prefix_title_len).rev() {
            let word = self.words.remove(i);
            self.found_prefix(word);
        }

        // Strip non-comma-separated titles & suffixes (e.g. "John Smith Jr.")
        let first_postfix_index =
            if any_after || self.words.len() + self.possible_false_prefix.iter().count() > 2 {
                title::find_postfix_index(&self.words[1..], false) + 1
            } else {
                self.words.len()
            };
        if first_postfix_index < self.words.len() {
            let postfix = self.words.swap_remove(first_postfix_index);
            self.found_suffix_or_postfix(postfix, false);
            self.words.truncate(first_postfix_index);
        }

        if prefix_title_len > 0 || !any_after {
            // Finding a prefix title means the next word is a first name or
            // initial (we don't support "Dr. Smith, John")
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        } else {
            // Have to guess whether this is just the surname (as in "Smith, John")
            // or the full name (as in "John Smith")
            //
            // Note we might be wrong, and have to go back, if we think the given
            // name is coming after a comma, but it never does
            self.surname_index = surname::find_surname_index(&self.words);
        }
    }

    // Called after the first comma, until we find a given name or first initial
    fn handle_after_comma(&mut self, part: &'a str) {
        debug_assert!(
            !self.words.is_empty() && self.surname_index == 0,
            "Invalid state for handle_after_comma!"
        );

        let mut given_middle_or_postfix_words: SmallVec<[NamePart<'a>; 5]> =
            NamePart::all_from_text(part, self.use_capitalization, Location::Start).collect();

        if given_middle_or_postfix_words.is_empty() {
            return;
        }

        // Handle (unusual) formats like "Smith, Dr. John M."
        if given_middle_or_postfix_words.len() > 1 {
            let prefix_title_len = title::find_prefix_len(&given_middle_or_postfix_words);
            for i in (0..prefix_title_len).rev() {
                self.found_prefix(given_middle_or_postfix_words.remove(i));
            }
        }

        // Handle isolated suffixes or titles as well as (unusual) formats like
        // "Smith, John Jr."
        let first_postfix_index = title::find_postfix_index(&given_middle_or_postfix_words, true);
        if first_postfix_index < given_middle_or_postfix_words.len() {
            let postfix = given_middle_or_postfix_words.swap_remove(first_postfix_index);
            self.found_suffix_or_postfix(postfix, true);
            given_middle_or_postfix_words.truncate(first_postfix_index);
        }

        // Now if there are any words left, they include the given name or first
        // initial (in a format like "Smith, John" or "Smith, J. M."), so we put
        // them in front
        if !given_middle_or_postfix_words.is_empty() {
            self.surname_index = given_middle_or_postfix_words.len();

            self.words.reserve(given_middle_or_postfix_words.len());
            for word in given_middle_or_postfix_words.into_iter().rev() {
                self.words.insert(0, word);
            }
        }
    }

    // Called on any parts remaining after full name is found
    fn handle_after_surname(&mut self, part: &'a str) {
        debug_assert!(
            self.surname_index > 0,
            "Invalid state for handle_after_surname!"
        );

        if self.possible_false_postfix.is_some() && self.generation_from_suffix.is_some() {
            return;
        }

        let mut postfix_words =
            NamePart::all_from_text(part, self.use_capitalization, Location::End);
        while self.possible_false_postfix.is_none() || self.generation_from_suffix.is_none() {
            if let Some(word) = postfix_words.next() {
                self.found_suffix_or_postfix(word, false);
            } else {
                break;
            }
        }
    }

    fn found_prefix(&mut self, prefix: NamePart<'a>) {
        // We drop prefixes, but keep the last word that's namelike,
        // just in case we make a mistake and it turns out by process of
        // elimination that this must actually be a given name
        if self.possible_false_prefix.is_none() && (prefix.is_namelike() || prefix.is_initials()) {
            self.possible_false_prefix = Some(prefix);
        }
    }

    fn found_suffix_or_postfix(&mut self, postfix: NamePart<'a>, expect_initials: bool) {
        if self.generation_from_suffix.is_none() {
            self.generation_from_suffix = suffix::generation_from_suffix(&postfix, expect_initials);
        }

        // We throw away most postfix titles, but keep the first one that's namelike,
        // just in case we make a mistake and it turns out by process of elimination
        // that this must actually be a surname
        if self.possible_false_postfix.is_none() && (postfix.is_namelike() || postfix.is_initials())
        {
            self.possible_false_postfix = Some(postfix);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test::{black_box, Bencher};

    #[test]
    fn first_last() {
        let (parts, surname_index, generation) = parse("John Doe").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn initial_last() {
        let (parts, surname_index, generation) = parse("J. Doe").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_first() {
        let (parts, surname_index, generation) = parse("Doe, John").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_initial() {
        let (parts, surname_index, generation) = parse("Doe, J.").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn suffix() {
        let (parts, surname_index, generation) = parse("John Doe III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[test]
    fn suffix_comma() {
        let (parts, surname_index, generation) = parse("Doe, John III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[bench]
    fn parse_simple(b: &mut Bencher) {
        b.iter(|| black_box(parse("John Doe").is_some()))
    }

    #[bench]
    fn parse_nonascii(b: &mut Bencher) {
        b.iter(|| black_box(parse("이용희").is_some()))
    }

    #[bench]
    fn parse_comma(b: &mut Bencher) {
        b.iter(|| black_box(parse("Doe, John").is_some()))
    }

    #[bench]
    fn parse_all_caps(b: &mut Bencher) {
        b.iter(|| black_box(parse("JOHN DOE").is_some()))
    }

    #[bench]
    fn parse_complex(b: &mut Bencher) {
        b.iter(|| black_box(parse("James S. Brown MD, FRCS, FDSRCS").is_some()))
    }
}
<commit_msg>Further parse cleanup and optimization<commit_after>use super::namepart::{Location, NamePart};
use super::suffix;
use super::surname;
use super::title;
use super::utils::is_mixed_case;
use smallvec::SmallVec;

#[derive(Debug)]
struct ParseOp<'a> {
    words: SmallVec<[NamePart<'a>; 7]>,
    surname_index: usize,
    generation_from_suffix: Option<usize>,
    possible_false_prefix: Option<NamePart<'a>>,
    possible_false_postfix: Option<NamePart<'a>>,
    use_capitalization: bool,
}

pub fn parse(name: &str) -> Option<(SmallVec<[NamePart; 7]>, usize, Option<usize>)> {
    let mut op = ParseOp {
        words: SmallVec::new(),
        surname_index: 0,
        generation_from_suffix: None,
        possible_false_prefix: None,
        possible_false_postfix: None,
        use_capitalization: is_mixed_case(name),
    };

    if op.run(name) {
        Some((op.words, op.surname_index, op.generation_from_suffix))
    } else {
        None
    }
}

impl<'a> ParseOp<'a> {
    fn run(&mut self, name: &'a str) -> bool {
        // Separate comma-separated titles and suffixes, then flip remaining words
        // around remaining comma, if any
        let mut parts = name.split(',').peekable();
        while let Some(part) = parts.next() {
            let first_part = self.words.is_empty();
            let last_part = parts.peek().is_none();

            if first_part && last_part {
                // Simple case
                self.handle_no_comma(part);
            } else if first_part {
                // We're in the surname part if the format is like "Smith, John",
                // or the only actual name part if the format is "John Smith,
                // esq."
                self.handle_before_comma(part);
            } else if self.surname_index == 0 {
                // We already processed one comma-separated part, but we think
                // it was just the surname, so this might be the given name or
                // initials
                let must_include_given = last_part && self.words.len() == 1;
                self.handle_after_comma(part, must_include_given);
            } else {
                // We already found the full name, so this is a comma-separated
                // postfix title or suffix
                self.handle_after_surname(part);
            }
        }

        // If there are two or fewer words, e.g. "JOHN MA", we treat
        // ambiguous strings like "MA" as surnames rather than titles
        // (not initials, which should only be at the end of the input
        // if they are comma-separated, and we already handled that case)
        if !self.valid() {
            if let Some(ref postfix) = self.possible_false_postfix {
                self.words.push(postfix.clone());
            } else if let Some(ref prefix) = self.possible_false_prefix {
                self.words.insert(0, prefix.clone());
            }
        }

        // Anything trailing that looks like initials is probably a stray postfix
        while self.words.last().filter(|w| !w.is_namelike()).is_some() {
            let removed = self.words.pop().unwrap();

            // If we guessed the surname previously, our guess is no longer valid
            self.surname_index = 0;

            // If the alternative is having less than two words and giving up,
            // check if we can treat the last word as a name if we just ignore
            // case; this handles the not-quite-rare-enough case of an all-caps
            // last name (e.g. Neto John SMITH), among others
            if self.use_capitalization && !self.valid() {
                let word = NamePart::from_word(removed.word, false, Location::End);
                if word.is_namelike() {
                    self.words.push(word);
                    break;
                }
            }
        }

        // Handle case where we thought the whole before-comma part was a surname,
        // but we never found a plausible given name or initial afterwards,
        // as well as the reset just above
        if self.surname_index == 0 && self.words.len() > 1 {
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        }

        // Check the plausibility of what we've found
        self.valid()
    }

    fn valid(&self) -> bool {
        self.words.len() >= 2
            && self
                .words
                .iter()
                .all(|w| w.is_namelike() || w.is_initials())
            && self.words[self.surname_index..]
                .iter()
                .any(|w| w.is_namelike())
    }

    fn handle_no_comma(&mut self, name: &'a str) {
        debug_assert!(
            self.words.is_empty()
                && self.surname_index == 0
                && self.possible_false_prefix.is_none()
                && self.possible_false_postfix.is_none(),
            "Invalid state for handle_no_comma!"
        );

        self.words.extend(
            NamePart::all_from_text(name, self.use_capitalization, Location::Start)
                .skip_while(|word| !word.is_namelike() && !word.is_initials()),
        );

        if self.words.is_empty() {
            return;
        }

        // Check for title as prefix (e.g. "Dr. John Smith" or "Right Hon.
        // John Smith")
        let prefix_title_len = if self.words.len() > 2 {
            title::find_prefix_len(&self.words)
        } else {
            0
        };
        self.strip_prefix(prefix_title_len);

        // Strip non-comma-separated titles & suffixes (e.g. "John Smith Jr.")
        let first_postfix_index =
            if self.words.len() + self.possible_false_prefix.iter().count() > 2 {
                title::find_postfix_index(&self.words[1..], false) + 1
            } else {
                self.words.len()
            };
        self.strip_postfix(first_postfix_index);

        self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
    }

    // Called only until any words are found
    fn handle_before_comma(&mut self, part: &'a str) {
        debug_assert!(
            self.words.is_empty()
                && self.surname_index == 0
                && self.possible_false_prefix.is_none()
                && self.possible_false_postfix.is_none(),
            "Invalid state for handle_before_comma!"
        );

        self.words.extend(NamePart::all_from_text(
            part,
            self.use_capitalization,
            Location::End,
        ));

        if self.words.is_empty() {
            return;
        }

        // Check for title as prefix (e.g. "Dr. John Smith, Jr.")
        let prefix_title_len = title::find_prefix_len(&self.words);
        self.strip_prefix(prefix_title_len);

        // Strip non-comma-separated titles & suffixes (e.g. "John Smith Jr., MD")
        let first_postfix_index = title::find_postfix_index(&self.words[1..], false) + 1;
        self.strip_postfix(first_postfix_index);

        if prefix_title_len > 0 {
            // Finding a prefix title means the next word is a first name or
            // initial (we don't support "Dr. Smith, John")
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        } else {
            // Have to guess whether this is just the surname (as in "Smith, John")
            // or the full name (as in "John Smith")
            //
            // Note we might be wrong, and have to go back, if we think the given
            // name is coming after a comma, but it never does
            self.surname_index = surname::find_surname_index(&self.words);
        }
    }

    // Called after the first comma, until we find a given name or first initial
    fn handle_after_comma(&mut self, part: &'a str, must_include_given: bool) {
        debug_assert!(
            !self.words.is_empty() && self.surname_index == 0,
            "Invalid state for handle_after_comma!"
        );

        let mut given_middle_or_postfix_words: SmallVec<[NamePart<'a>; 5]> =
            NamePart::all_from_text(part, self.use_capitalization, Location::Start).collect();

        if given_middle_or_postfix_words.is_empty() {
            return;
        }

        // Handle (unusual) formats like "Smith, Dr. John M."
        if given_middle_or_postfix_words.len() > 1 {
            let prefix_len = title::find_prefix_len(&given_middle_or_postfix_words);
            self.strip_unsaved_prefix(&mut given_middle_or_postfix_words, prefix_len);
        }

        // Handle isolated suffixes or titles as well as (unusual) formats like
        // "Smith, John Jr." and "Smith, Jr., John"
        let first_postfix_index = if must_include_given {
            title::find_postfix_index(&given_middle_or_postfix_words[1..], true) + 1
        } else {
            title::find_postfix_index(&given_middle_or_postfix_words, true)
        };
        self.strip_unsaved_postfix(&mut given_middle_or_postfix_words, first_postfix_index);

        // Now if there are any words left, they include the given name or first
        // initial (in a format like "Smith, John" or "Smith, J. M."), so we put
        // them in front
        if !given_middle_or_postfix_words.is_empty() {
            self.surname_index = given_middle_or_postfix_words.len();

            self.words.reserve(given_middle_or_postfix_words.len());
            self.words
                .insert_many(0, given_middle_or_postfix_words.into_iter());
        }
    }

    // Called on any parts remaining after full name is found
    fn handle_after_surname(&mut self, part: &'a str) {
        debug_assert!(
            self.surname_index > 0,
            "Invalid state for handle_after_surname!"
        );

        if self.possible_false_postfix.is_some() && self.generation_from_suffix.is_some() {
            return;
        }

        let mut postfix_words =
            NamePart::all_from_text(part, self.use_capitalization, Location::End);
        while self.possible_false_postfix.is_none() || self.generation_from_suffix.is_none() {
            if let Some(word) = postfix_words.next() {
                self.found_suffix_or_postfix(word, false);
            } else {
                break;
            }
        }
    }

    fn strip_prefix(&mut self, len: usize) {
        for i in (0..len).rev() {
            let word = self.words.remove(i);
            self.found_prefix(word);
        }
    }

    fn strip_unsaved_prefix(&mut self, words: &mut SmallVec<[NamePart<'a>; 5]>, len: usize) {
        for i in (0..len).rev() {
            self.found_prefix(words.remove(i));
        }
    }

    fn found_prefix(&mut self, prefix: NamePart<'a>) {
        // We drop prefixes, but keep the last word that's namelike,
        // just in case we make a mistake and it turns out by process of
        // elimination that this must actually be a given name
        if self.possible_false_prefix.is_none() && (prefix.is_namelike() || prefix.is_initials()) {
            self.possible_false_prefix = Some(prefix);
        }
    }

    fn strip_postfix(&mut self, index: usize) {
        if index < self.words.len() {
            let postfix = self.words.swap_remove(index);
            self.found_suffix_or_postfix(postfix, false);
            self.words.truncate(index);
        }
    }

    fn strip_unsaved_postfix(&mut self, words: &mut SmallVec<[NamePart<'a>; 5]>, index: usize) {
        if index < words.len() {
            let postfix = words.swap_remove(index);
            self.found_suffix_or_postfix(postfix, false);
            words.truncate(index);
        }
    }

    fn found_suffix_or_postfix(&mut self, postfix: NamePart<'a>, expect_initials: bool) {
        if self.generation_from_suffix.is_none() {
            self.generation_from_suffix = suffix::generation_from_suffix(&postfix, expect_initials);
        }

        // We throw away most postfix titles, but keep the first one that's namelike,
        // just in case we make a mistake and it turns out by process of elimination
        // that this must actually be a surname
        if self.possible_false_postfix.is_none() && (postfix.is_namelike() || postfix.is_initials())
        {
            self.possible_false_postfix = Some(postfix);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test::{black_box, Bencher};

    #[test]
    fn first_last() {
        let (parts, surname_index, generation) = parse("John Doe").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn initial_last() {
        let (parts, surname_index, generation) = parse("J. Doe").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_first() {
        let (parts, surname_index, generation) = parse("Doe, John").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_initial() {
        let (parts, surname_index, generation) = parse("Doe, J.").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn suffix() {
        let (parts, surname_index, generation) = parse("John Doe III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[test]
    fn suffix_comma() {
        let (parts, surname_index, generation) = parse("Doe, John III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[test]
    fn intermediate_suffix() {
        let (parts, surname_index, generation) = parse("Doe, II, John").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(2), generation);

        let (parts, surname_index, generation) = parse("Griffey, Jr., Ken").unwrap();
        assert_eq!("Ken", parts[0].word);
        assert_eq!("Griffey", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(2), generation);
    }

    #[bench]
    fn parse_simple(b: &mut Bencher) {
        b.iter(|| black_box(parse("John Doe").is_some()))
    }

    #[bench]
    fn parse_nonascii(b: &mut Bencher) {
        b.iter(|| black_box(parse("이용희").is_some()))
    }

    #[bench]
    fn parse_comma(b: &mut Bencher) {
        b.iter(|| black_box(parse("Doe, John").is_some()))
    }

    #[bench]
    fn parse_all_caps(b: &mut Bencher) {
        b.iter(|| black_box(parse("JOHN DOE").is_some()))
    }

    #[bench]
    fn parse_complex(b: &mut Bencher) {
        b.iter(|| black_box(parse("James S. Brown MD, FRCS, FDSRCS").is_some()))
    }
}
<|endoftext|>


<commit_before>//! Relay for TCP implementation

use std::io::{self, Read, BufRead};
use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
use std::rc::Rc;
use std::mem;
use std::time::Duration;
use std::net::IpAddr;

use crypto::CipherCategory;
use relay::socks5::Address;
use relay::{BoxIoFuture, boxed_future};
use relay::dns_resolver::resolve;
use relay::Context;
use config::{ServerConfig, ServerAddr};

use tokio_core::net::TcpStream;
use tokio_core::reactor::{Handle, Timeout};
use tokio_io::AsyncRead;
use tokio_io::io::{read_exact, write_all};
use tokio_io::io::{ReadHalf, WriteHalf};

use futures::{self, Future, Poll};

use bytes::{BufMut, BytesMut};

pub use self::crypto_io::{DecryptedRead, EncryptedWrite};

use self::stream::{DecryptedReader as StreamDecryptedReader, EncryptedWriter as StreamEncryptedWriter};
use self::aead::{DecryptedReader as AeadDecryptedReader, EncryptedWriter as AeadEncryptedWriter};

pub mod local;
mod socks5_local;
pub mod server;
mod stream;
pub mod client;
mod crypto_io;
mod aead;
mod utils;

const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer

/// Directions in the tunnel
#[derive(Debug, Copy, Clone)]
pub enum TunnelDirection {
    /// Client -> Server
    Client2Server,
    /// Client <- Server
    Server2Client,
}

type TcpReadHalf = ReadHalf<TcpStream>;
type TcpWriteHalf = WriteHalf<TcpStream>;

/// `ReadHalf `of `TcpStream` with decryption
pub enum DecryptedHalf {
    Stream(StreamDecryptedReader<TcpReadHalf>),
    Aead(AeadDecryptedReader<TcpReadHalf>),
}

impl Read for DecryptedHalf {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            DecryptedHalf::Stream(ref mut d) => d.read(buf),
            DecryptedHalf::Aead(ref mut d) => d.read(buf),
        }
    }
}

impl DecryptedRead for DecryptedHalf {
    fn buffer_size(&self, data: &[u8]) -> usize {
        match *self {
            DecryptedHalf::Stream(ref e) => e.buffer_size(data),
            DecryptedHalf::Aead(ref e) => e.buffer_size(data),
        }
    }
}

impl AsyncRead for DecryptedHalf {}

impl BufRead for DecryptedHalf {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        match *self {
            DecryptedHalf::Stream(ref mut d) => d.fill_buf(),
            DecryptedHalf::Aead(ref mut d) => d.fill_buf(),
        }
    }

    fn consume(&mut self, amt: usize) {
        match *self {
            DecryptedHalf::Stream(ref mut d) => d.consume(amt),
            DecryptedHalf::Aead(ref mut d) => d.consume(amt),
        }
    }
}

impl From<StreamDecryptedReader<TcpReadHalf>> for DecryptedHalf {
    fn from(r: StreamDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
        DecryptedHalf::Stream(r)
    }
}

impl From<AeadDecryptedReader<TcpReadHalf>> for DecryptedHalf {
    fn from(r: AeadDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
        DecryptedHalf::Aead(r)
    }
}

/// `WriteHalf` of `TcpStream` with encryption
pub enum EncryptedHalf {
    Stream(StreamEncryptedWriter<TcpWriteHalf>),
    Aead(AeadEncryptedWriter<TcpWriteHalf>),
}

impl EncryptedWrite for EncryptedHalf {
    fn write_raw(&mut self, data: &[u8]) -> io::Result<usize> {
        match *self {
            EncryptedHalf::Stream(ref mut e) => e.write_raw(data),
            EncryptedHalf::Aead(ref mut e) => e.write_raw(data),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match *self {
            EncryptedHalf::Stream(ref mut e) => e.flush(),
            EncryptedHalf::Aead(ref mut e) => e.flush(),
        }
    }

    fn encrypt<B: BufMut>(&mut self, data: &[u8], buf: &mut B) -> io::Result<()> {
        match *self {
            EncryptedHalf::Stream(ref mut e) => e.encrypt(data, buf),
            EncryptedHalf::Aead(ref mut e) => e.encrypt(data, buf),
        }
    }

    fn buffer_size(&self, data: &[u8]) -> usize {
        match *self {
            EncryptedHalf::Stream(ref e) => e.buffer_size(data),
            EncryptedHalf::Aead(ref e) => e.buffer_size(data),
        }
    }
}

impl From<StreamEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
    fn from(d: StreamEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
        EncryptedHalf::Stream(d)
    }
}

impl From<AeadEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
    fn from(d: AeadEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
        EncryptedHalf::Aead(d)
    }
}

/// Boxed future of `DecryptedHalf`
pub type DecryptedHalfFut = BoxIoFuture<DecryptedHalf>;
/// Boxed future of `EncryptedHalf`
pub type EncryptedHalfFut = BoxIoFuture<EncryptedHalf>;

fn connect_proxy_server(svr_cfg: Rc<ServerConfig>) -> BoxIoFuture<TcpStream> {
    let timeout = *svr_cfg.timeout();
    trace!("Connecting to proxy {:?}, timeout: {:?}",
           svr_cfg.addr(),
           timeout);
    match *svr_cfg.addr() {
        ServerAddr::SocketAddr(ref addr) => {
            Context::with(|ctx| {
                              let handle = ctx.handle();
                              try_timeout(TcpStream::connect(addr, handle), timeout, handle)
                          })
        }
        ServerAddr::DomainName(ref domain, port) => {
            let fut = Context::with(|ctx| {
                let handle = ctx.handle();
                try_timeout(resolve(&domain[..], &handle), timeout, &handle).and_then(move |sockaddr| {
                    let sockaddr = match sockaddr {
                        IpAddr::V4(v4) => SocketAddr::V4(SocketAddrV4::new(v4, port)),
                        IpAddr::V6(v6) => SocketAddr::V6(SocketAddrV6::new(v6, port, 0, 0)),
                    };

                    Context::with(|ctx| {
                                      let handle = ctx.handle();
                                      try_timeout(TcpStream::connect(&sockaddr, &handle), timeout, &handle)
                                  })
                })
            });
            boxed_future(fut)
        }
    }
}

/// Handshake logic for ShadowSocks Client
pub fn proxy_server_handshake(remote_stream: TcpStream, svr_cfg: Rc<ServerConfig>, relay_addr: Address) -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let timeout = *svr_cfg.timeout();
    let fut = proxy_handshake(remote_stream, svr_cfg).and_then(move |(r_fut, w_fut)| {;
        let w_fut = w_fut.and_then(move |enc_w| {
            trace!("Got encrypt stream and going to send addr: {:?}",
                   relay_addr);

            // Send relay address to remote
            let mut buf = BytesMut::with_capacity(relay_addr.len());
            relay_addr.write_to_buf(&mut buf);

            Context::with(|ctx| try_timeout(enc_w.write_all(buf), timeout, ctx.handle()).map(|(w, _)| w))
        });

        Ok((r_fut, boxed_future(w_fut)))
    });
    boxed_future(fut)
}

/// ShadowSocks Client-Server handshake protocol
/// Exchange cipher IV and creates stream wrapper
pub fn proxy_handshake(remote_stream: TcpStream, svr_cfg: Rc<ServerConfig>) -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let fut = futures::lazy(|| Ok(remote_stream.split())).and_then(move |(r, w)| {

        let timeout = svr_cfg.timeout().clone();

        let svr_cfg_cloned = svr_cfg.clone();

        let enc = {
            // Encrypt data to remote server

            // Send initialize vector to remote and create encryptor

            let method = svr_cfg.method();
            let prev_buf = match method.category() {
                CipherCategory::Stream => {
                    let local_iv = method.gen_init_vec();
                    trace!("Going to send initialize vector: {:?}", local_iv);
                    local_iv
                }
                CipherCategory::Aead => {
                    let local_salt = method.gen_salt();
                    trace!("Going to send salt: {:?}", local_salt);
                    local_salt
                }
            };

            Context::with(|ctx| {
                try_timeout(write_all(w, prev_buf), timeout, ctx.handle()).and_then(move |(w, prev_buf)| match svr_cfg.method().category() {
                                                                                CipherCategory::Stream => {
                                                                                let local_iv = prev_buf;
                                                                                Ok(From::from(StreamEncryptedWriter::new(w, svr_cfg.method(), svr_cfg.key(), &local_iv)))
                                                                                }
                                                                                CipherCategory::Aead => {
                                                                                let local_salt = prev_buf;
                                                                                let wtr = AeadEncryptedWriter::new(w, svr_cfg.method(), svr_cfg.key(), &local_salt[..]);
                                                                                Ok(From::from(wtr))
                                                                                }
                                                                                })
            })
        };

        let dec = {
            let svr_cfg = svr_cfg_cloned;

            // Decrypt data from remote server

            let method = svr_cfg.method();
            let prev_len = match method.category() {
                CipherCategory::Stream => method.iv_size(),
                CipherCategory::Aead => method.salt_size(),
            };

            Context::with(|ctx| {
                try_timeout(read_exact(r, vec![0u8; prev_len]), timeout, ctx.handle()).and_then(move |(r, remote_iv)| match svr_cfg.method().category() {
                                                                                CipherCategory::Stream => {
                                                                                trace!("Got initialize vector {:?}", remote_iv);
                                                                                let decrypt_stream = StreamDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv);
                                                                                Ok(From::from(decrypt_stream))
                                                                                }
                                                                                CipherCategory::Aead => {
                                                                                trace!("Got salt {:?}", remote_iv);
                                                                                let dr = AeadDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv);
                                                                                Ok(From::from(dr))
                                                                                }
                                                                                })
            })
        };

        Ok((boxed_future(dec), boxed_future(enc)))
    });

    boxed_future(fut)
}

/// Establish tunnel between server and client
pub fn tunnel<CF, CFI, SF, SFI>(addr: Address, c2s: CF, s2c: SF) -> BoxIoFuture<()>
    where CF: Future<Item = CFI, Error = io::Error> + 'static,
          SF: Future<Item = SFI, Error = io::Error> + 'static
{
    let addr = Rc::new(addr);

    let cloned_addr = addr.clone();
    let c2s = c2s.then(move |res| {
        match res {
            Ok(..) => {
                // Continue reading response from remote server
                trace!("Relay {} client -> server is finished", cloned_addr);

                Ok(TunnelDirection::Client2Server)
            }
            Err(err) => {
                error!("Relay {} client -> server aborted: {}", cloned_addr, err);
                Err(err)
            }
        }
    });

    let cloned_addr = addr.clone();
    let s2c = s2c.then(move |res| match res {
                           Ok(..) => {
                               trace!("Relay {} client <- server is finished", cloned_addr);

                               Ok(TunnelDirection::Server2Client)
                           }
                           Err(err) => {
                               error!("Relay {} client <- server aborted: {}", cloned_addr, err);
                               Err(err)
                           }
                       });

    let fut = c2s.select(s2c)
        .and_then(move |(dir, _)| {
            match dir {
                TunnelDirection::Server2Client => {
                    trace!("Relay {} client <- server is closed, abort connection",
                           addr)
                }
                TunnelDirection::Client2Server => {
                    trace!("Relay {} server -> client is closed, abort connection",
                           addr)
                }
            }

            Ok(())
        })
        .map_err(|(err, _)| err);

    boxed_future(fut)
}

/// Read until EOF, and ignore
pub enum IgnoreUntilEnd<R: Read> {
    Pending { r: R, amt: u64 },
    Empty,
}

impl<R: Read> Future for IgnoreUntilEnd<R> {
    type Item = u64;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match *self {
            IgnoreUntilEnd::Empty => panic!("poll IgnoreUntilEnd after it is finished"),
            IgnoreUntilEnd::Pending {
                ref mut r,
                ref mut amt,
            } => {
                let mut buf = [0u8; 4096];
                loop {
                    let n = try_nb!(r.read(&mut buf));
                    *amt += n as u64;

                    if n == 0 {
                        break;
                    }
                }
            }
        }

        match mem::replace(self, IgnoreUntilEnd::Empty) {
            IgnoreUntilEnd::Pending { amt, .. } => Ok(amt.into()),
            IgnoreUntilEnd::Empty => unreachable!(),
        }
    }
}

/// Ignore all data from the reader
pub fn ignore_until_end<R: Read>(r: R) -> IgnoreUntilEnd<R> {
    IgnoreUntilEnd::Pending { r: r, amt: 0 }
}

fn try_timeout<T, F>(fut: F, dur: Option<Duration>, handle: &Handle) -> BoxIoFuture<T>
    where F: Future<Item = T, Error = io::Error> + 'static,
          T: 'static
{
    match dur {
        Some(dur) => io_timeout(fut, dur, handle),
        None => boxed_future(fut),
    }
}

fn io_timeout<T, F>(fut: F, dur: Duration, handle: &Handle) -> BoxIoFuture<T>
    where F: Future<Item = T, Error = io::Error> + 'static,
          T: 'static
{
    let fut = fut.select(Timeout::new(dur, handle)
            .unwrap() // It must be succeeded!
            .and_then(|_| Err(io::Error::new(io::ErrorKind::TimedOut, "timeout"))))
        .then(|res| match res {
                  Ok((t, _)) => Ok(t),
                  Err((err, _)) => Err(err),
              });
    boxed_future(fut)
}
<commit_msg>use macro to simplify code<commit_after>//! Relay for TCP implementation

use std::io::{self, Read, BufRead};
use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
use std::rc::Rc;
use std::mem;
use std::time::Duration;
use std::net::IpAddr;

use crypto::CipherCategory;
use relay::socks5::Address;
use relay::{BoxIoFuture, boxed_future};
use relay::dns_resolver::resolve;
use relay::Context;
use config::{ServerConfig, ServerAddr};

use tokio_core::net::TcpStream;
use tokio_core::reactor::{Handle, Timeout};
use tokio_io::AsyncRead;
use tokio_io::io::{read_exact, write_all};
use tokio_io::io::{ReadHalf, WriteHalf};

use futures::{self, Future, Poll};

use bytes::{BufMut, BytesMut};

pub use self::crypto_io::{DecryptedRead, EncryptedWrite};

use self::stream::{DecryptedReader as StreamDecryptedReader, EncryptedWriter as StreamEncryptedWriter};
use self::aead::{DecryptedReader as AeadDecryptedReader, EncryptedWriter as AeadEncryptedWriter};

pub mod local;
mod socks5_local;
pub mod server;
mod stream;
pub mod client;
mod crypto_io;
mod aead;
mod utils;

const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer

/// Directions in the tunnel
#[derive(Debug, Copy, Clone)]
pub enum TunnelDirection {
    /// Client -> Server
    Client2Server,
    /// Client <- Server
    Server2Client,
}

type TcpReadHalf = ReadHalf<TcpStream>;
type TcpWriteHalf = WriteHalf<TcpStream>;

/// `ReadHalf `of `TcpStream` with decryption
pub enum DecryptedHalf {
    Stream(StreamDecryptedReader<TcpReadHalf>),
    Aead(AeadDecryptedReader<TcpReadHalf>),
}

macro_rules! ref_half_do {
    ($self:expr,$t:ident,$m:ident$(,$p:expr)*) => {
        match *$self {
            $t::Stream(ref d) => d.$m($($p),*),
            $t::Aead(ref d) => d.$m($($p),*),
        }
    }
}

macro_rules! mut_half_do {
    ($self:expr,$t:ident,$m:ident$(,$p:expr)*) => {
        match *$self {
            $t::Stream(ref mut  d) => d.$m($($p),*),
            $t::Aead(ref mut d) => d.$m($($p),*),
        }
    }
}

impl Read for DecryptedHalf {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        mut_half_do!(self, DecryptedHalf, read, buf)
    }
}

impl DecryptedRead for DecryptedHalf {
    fn buffer_size(&self, data: &[u8]) -> usize {
        ref_half_do!(self, DecryptedHalf, buffer_size, data)
    }
}

impl AsyncRead for DecryptedHalf {}

impl BufRead for DecryptedHalf {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        mut_half_do!(self, DecryptedHalf, fill_buf)
    }

    fn consume(&mut self, amt: usize) {
        mut_half_do!(self, DecryptedHalf, consume, amt)
    }
}

impl From<StreamDecryptedReader<TcpReadHalf>> for DecryptedHalf {
    fn from(r: StreamDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
        DecryptedHalf::Stream(r)
    }
}

impl From<AeadDecryptedReader<TcpReadHalf>> for DecryptedHalf {
    fn from(r: AeadDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
        DecryptedHalf::Aead(r)
    }
}

/// `WriteHalf` of `TcpStream` with encryption
pub enum EncryptedHalf {
    Stream(StreamEncryptedWriter<TcpWriteHalf>),
    Aead(AeadEncryptedWriter<TcpWriteHalf>),
}

impl EncryptedWrite for EncryptedHalf {
    fn write_raw(&mut self, data: &[u8]) -> io::Result<usize> {
        mut_half_do!(self, EncryptedHalf, write_raw, data)
    }

    fn flush(&mut self) -> io::Result<()> {
        mut_half_do!(self, EncryptedHalf, flush)
    }

    fn encrypt<B: BufMut>(&mut self, data: &[u8], buf: &mut B) -> io::Result<()> {
        mut_half_do!(self, EncryptedHalf, encrypt, data, buf)
    }

    fn buffer_size(&self, data: &[u8]) -> usize {
        ref_half_do!(self,EncryptedHalf, buffer_size, data)
    }
}

impl From<StreamEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
    fn from(d: StreamEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
        EncryptedHalf::Stream(d)
    }
}

impl From<AeadEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
    fn from(d: AeadEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
        EncryptedHalf::Aead(d)
    }
}

/// Boxed future of `DecryptedHalf`
pub type DecryptedHalfFut = BoxIoFuture<DecryptedHalf>;
/// Boxed future of `EncryptedHalf`
pub type EncryptedHalfFut = BoxIoFuture<EncryptedHalf>;

fn connect_proxy_server(svr_cfg: Rc<ServerConfig>) -> BoxIoFuture<TcpStream> {
    let timeout = *svr_cfg.timeout();
    trace!("Connecting to proxy {:?}, timeout: {:?}",
           svr_cfg.addr(),
           timeout);
    match *svr_cfg.addr() {
        ServerAddr::SocketAddr(ref addr) => {
            Context::with(|ctx| {
                              let handle = ctx.handle();
                              try_timeout(TcpStream::connect(addr, handle), timeout, handle)
                          })
        }
        ServerAddr::DomainName(ref domain, port) => {
            let fut = Context::with(|ctx| {
                let handle = ctx.handle();
                try_timeout(resolve(&domain[..], &handle), timeout, &handle).and_then(move |sockaddr| {
                    let sockaddr = match sockaddr {
                        IpAddr::V4(v4) => SocketAddr::V4(SocketAddrV4::new(v4, port)),
                        IpAddr::V6(v6) => SocketAddr::V6(SocketAddrV6::new(v6, port, 0, 0)),
                    };

                    Context::with(|ctx| {
                                      let handle = ctx.handle();
                                      try_timeout(TcpStream::connect(&sockaddr, &handle), timeout, &handle)
                                  })
                })
            });
            boxed_future(fut)
        }
    }
}

/// Handshake logic for ShadowSocks Client
pub fn proxy_server_handshake(remote_stream: TcpStream,
                              svr_cfg: Rc<ServerConfig>,
                              relay_addr: Address)
                              -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let timeout = *svr_cfg.timeout();
    let fut = proxy_handshake(remote_stream, svr_cfg).and_then(move |(r_fut, w_fut)| {;
        let w_fut = w_fut.and_then(move |enc_w| {
            trace!("Got encrypt stream and going to send addr: {:?}",
                   relay_addr);

            // Send relay address to remote
            let mut buf = BytesMut::with_capacity(relay_addr.len());
            relay_addr.write_to_buf(&mut buf);

            Context::with(|ctx| try_timeout(enc_w.write_all(buf), timeout, ctx.handle()).map(|(w, _)| w))
        });

        Ok((r_fut, boxed_future(w_fut)))
    });
    boxed_future(fut)
}

/// ShadowSocks Client-Server handshake protocol
/// Exchange cipher IV and creates stream wrapper
pub fn proxy_handshake(remote_stream: TcpStream,
                       svr_cfg: Rc<ServerConfig>)
                       -> BoxIoFuture<(DecryptedHalfFut, EncryptedHalfFut)> {
    let fut = futures::lazy(|| Ok(remote_stream.split())).and_then(move |(r, w)| {

        let timeout = svr_cfg.timeout().clone();

        let svr_cfg_cloned = svr_cfg.clone();

        let enc = {
            // Encrypt data to remote server

            // Send initialize vector to remote and create encryptor

            let method = svr_cfg.method();
            let prev_buf = match method.category() {
                CipherCategory::Stream => {
                    let local_iv = method.gen_init_vec();
                    trace!("Going to send initialize vector: {:?}", local_iv);
                    local_iv
                }
                CipherCategory::Aead => {
                    let local_salt = method.gen_salt();
                    trace!("Going to send salt: {:?}", local_salt);
                    local_salt
                }
            };

            Context::with(|ctx| {
                try_timeout(write_all(w, prev_buf), timeout, ctx.handle())
                    .and_then(move |(w, prev_buf)| match svr_cfg.method().category() {
                                  CipherCategory::Stream => {
                                      let local_iv = prev_buf;
                                      Ok(From::from(StreamEncryptedWriter::new(w,
                                                                               svr_cfg.method(),
                                                                               svr_cfg.key(),
                                                                               &local_iv)))
                                  }
                                  CipherCategory::Aead => {
                                      let local_salt = prev_buf;
                                      let wtr =
                                          AeadEncryptedWriter::new(w, svr_cfg.method(), svr_cfg.key(), &local_salt[..]);
                                      Ok(From::from(wtr))
                                  }
                              })
            })
        };

        let dec = {
            let svr_cfg = svr_cfg_cloned;

            // Decrypt data from remote server

            let method = svr_cfg.method();
            let prev_len = match method.category() {
                CipherCategory::Stream => method.iv_size(),
                CipherCategory::Aead => method.salt_size(),
            };

            Context::with(|ctx| {
                try_timeout(read_exact(r, vec![0u8; prev_len]), timeout, ctx.handle())
                    .and_then(move |(r, remote_iv)| match svr_cfg.method().category() {
                                  CipherCategory::Stream => {
                                      trace!("Got initialize vector {:?}", remote_iv);
                                      let decrypt_stream =
                                          StreamDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv);
                                      Ok(From::from(decrypt_stream))
                                  }
                                  CipherCategory::Aead => {
                                      trace!("Got salt {:?}", remote_iv);
                                      let dr = AeadDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv);
                                      Ok(From::from(dr))
                                  }
                              })
            })
        };

        Ok((boxed_future(dec), boxed_future(enc)))
    });

    boxed_future(fut)
}

/// Establish tunnel between server and client
pub fn tunnel<CF, CFI, SF, SFI>(addr: Address, c2s: CF, s2c: SF) -> BoxIoFuture<()>
    where CF: Future<Item = CFI, Error = io::Error> + 'static,
          SF: Future<Item = SFI, Error = io::Error> + 'static
{
    let addr = Rc::new(addr);

    let cloned_addr = addr.clone();
    let c2s = c2s.then(move |res| {
        match res {
            Ok(..) => {
                // Continue reading response from remote server
                trace!("Relay {} client -> server is finished", cloned_addr);

                Ok(TunnelDirection::Client2Server)
            }
            Err(err) => {
                error!("Relay {} client -> server aborted: {}", cloned_addr, err);
                Err(err)
            }
        }
    });

    let cloned_addr = addr.clone();
    let s2c = s2c.then(move |res| match res {
                           Ok(..) => {
                               trace!("Relay {} client <- server is finished", cloned_addr);

                               Ok(TunnelDirection::Server2Client)
                           }
                           Err(err) => {
                               error!("Relay {} client <- server aborted: {}", cloned_addr, err);
                               Err(err)
                           }
                       });

    let fut = c2s.select(s2c)
        .and_then(move |(dir, _)| {
            match dir {
                TunnelDirection::Server2Client => {
                    trace!("Relay {} client <- server is closed, abort connection",
                           addr)
                }
                TunnelDirection::Client2Server => {
                    trace!("Relay {} server -> client is closed, abort connection",
                           addr)
                }
            }

            Ok(())
        })
        .map_err(|(err, _)| err);

    boxed_future(fut)
}

/// Read until EOF, and ignore
pub enum IgnoreUntilEnd<R: Read> {
    Pending { r: R, amt: u64 },
    Empty,
}

impl<R: Read> Future for IgnoreUntilEnd<R> {
    type Item = u64;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match *self {
            IgnoreUntilEnd::Empty => panic!("poll IgnoreUntilEnd after it is finished"),
            IgnoreUntilEnd::Pending {
                ref mut r,
                ref mut amt,
            } => {
                let mut buf = [0u8; 4096];
                loop {
                    let n = try_nb!(r.read(&mut buf));
                    *amt += n as u64;

                    if n == 0 {
                        break;
                    }
                }
            }
        }

        match mem::replace(self, IgnoreUntilEnd::Empty) {
            IgnoreUntilEnd::Pending { amt, .. } => Ok(amt.into()),
            IgnoreUntilEnd::Empty => unreachable!(),
        }
    }
}

/// Ignore all data from the reader
pub fn ignore_until_end<R: Read>(r: R) -> IgnoreUntilEnd<R> {
    IgnoreUntilEnd::Pending { r: r, amt: 0 }
}

fn try_timeout<T, F>(fut: F, dur: Option<Duration>, handle: &Handle) -> BoxIoFuture<T>
    where F: Future<Item = T, Error = io::Error> + 'static,
          T: 'static
{
    match dur {
        Some(dur) => io_timeout(fut, dur, handle),
        None => boxed_future(fut),
    }
}

fn io_timeout<T, F>(fut: F, dur: Duration, handle: &Handle) -> BoxIoFuture<T>
    where F: Future<Item = T, Error = io::Error> + 'static,
          T: 'static
{
    let fut = fut.select(Timeout::new(dur, handle)
                         .unwrap() // It must be succeeded!
                         .and_then(|_| Err(io::Error::new(io::ErrorKind::TimedOut, "timeout"))))
        .then(|res| match res {
                  Ok((t, _)) => Ok(t),
                  Err((err, _)) => Err(err),
              });
    boxed_future(fut)
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
#![allow(unused_variables)]
extern crate capnp;

use std::{error, fmt, str};

use index::Index;
use key_builder::KeyBuilder;
use stems::{StemmedWord, Stems};

// TODO vmx 2016-11-02: Make it import "rocksdb" properly instead of needing to import the individual tihngs
use rocksdb::{DBIterator, SeekKey};
use rocksdb::rocksdb::Snapshot;
use records_capnp::payload;

#[derive(Debug)]
enum Error {
    Parse(String),
    Capnp(capnp::Error),
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Parse(ref description) => description,
            Error::Capnp(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::Parse(_) => None,
            Error::Capnp(ref err) => Some(err as &error::Error),
        }
    }
}

impl From<capnp::Error> for Error {
    fn from(err: capnp::Error) -> Error {
        Error::Capnp(err)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Parse(ref err) => write!(f, "Parse error: {}", err),
            Error::Capnp(ref err) => write!(f, "Capnproto error: {}", err),
        }
    }
}

struct DocResult {
    seq: u64,
    array_paths: Vec<Vec<u64>>,
}

impl DocResult {
    fn new() -> DocResult {
        DocResult {
            seq: 0,
            array_paths: Vec::new(),
        }
    }

    fn truncate_array_paths(&mut self, array_depth: usize) {
        for array_path in self.array_paths.iter_mut() {
            debug_assert!(array_path.len() >= array_depth);
            array_path.resize(array_depth, 0);
        }
    }

    fn intersect_array_paths(aa: &DocResult, bb: &DocResult) -> Option<DocResult> {
        let mut doc_result = DocResult::new();
        debug_assert_eq!(aa.seq, bb.seq);
        doc_result.seq = aa.seq;
        for array_path_a in &aa.array_paths {
            for array_path_b in &bb.array_paths {
                if array_path_a == array_path_b {
                    doc_result.array_paths.push(array_path_a.clone());
                }
            }
        }
        if doc_result.array_paths.is_empty() {
            None
        } else {
            Some(doc_result)
        }
    }
}

//trait QueryRuntimeFilter {
//struct QueryRuntimeFilter {}

trait QueryRuntimeFilter {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error>;
    fn next_result(&mut self) -> Result<Option<DocResult>, Error>;
}

pub struct Query {}

pub struct QueryResults {
}

//struct SnapshotIteratorCreator {
//    snapshot: rocksdb::Snapshot,
//}
//
//impl SnapshotIteratorCreator {
//    fn new(db: &rocksdb::DB) {
//        let snapshot = rocksdb::Snapshot::new(db);
//        SnapshotIteratorCreator{
//            snapshot: snapshot,
//        }
//    }
//
//    fn new_iterator(&self) {
//        self.snapshot.iter()
//    }
//}

struct ExactMatchFilter<'a> {
    iter: DBIterator<'a>,
    kb: KeyBuilder,
    stemmed_offset: u64,
    suffix: String,
    suffix_offset: u64,
}

impl<'a> ExactMatchFilter<'a> {
    fn new(iter: DBIterator<'a>, stemmed_word: &StemmedWord, mut kb: KeyBuilder) -> ExactMatchFilter<'a> {
        kb.push_word(&stemmed_word.stemmed);
        ExactMatchFilter{
            iter: iter,
            kb: kb,
            stemmed_offset: stemmed_word.stemmed_offset as u64,
            suffix: stemmed_word.suffix.clone(),
            suffix_offset: stemmed_word.suffix_offset as u64,
        }
    }
}

impl<'a> QueryRuntimeFilter for ExactMatchFilter<'a> {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error> {
        // Build the full key
        self.kb.push_doc_seq(start_id);

        // Seek in index to >= entry
        self.iter.seek(SeekKey::from(self.kb.key().as_bytes()));

        // Revert
        self.kb.pop_doc_seq();

        self.next_result()
    }

    fn next_result(&mut self) -> Result<Option<DocResult>, Error> {
        let mut doc_result = DocResult::new();

        loop {
            if !self.iter.valid() {
                return Ok(None)
            }

            // New scope needed as the iter.next() below invalidates the
            // current key and value
            {
                let key = self.iter.key();
                if !key.starts_with(self.kb.key().as_bytes()) {
                    return Ok(None)
                }
                let seq = &key[self.kb.key().len()..];

                let value = self.iter.value();
                // NOTE vmx 2016-10-13: I'm not really sure why the dereferencing is needed
                // and why we pass on mutable reference of it to `read_message()`
                let mut ref_value = &*value;
                let message_reader = ::capnp::serialize_packed::read_message(
                    &mut ref_value, ::capnp::message::ReaderOptions::new()).unwrap();
                let payload = message_reader.get_root::<payload::Reader>().unwrap();

                for aos_wis in try!(payload.get_arrayoffsets_to_wordinfos()).iter() {
                    for wi in try!(aos_wis.get_wordinfos()).iter() {
                        if self.stemmed_offset == wi.get_stemmed_offset() &&
                            self.suffix_offset == wi.get_suffix_offset() &&
                            self.suffix == try!(wi.get_suffix_text()) {
                                // We have a candidate document to return
                                let arrayoffsets = try!(aos_wis.get_arrayoffsets());
                                doc_result.array_paths.push(arrayoffsets.iter().collect::<>());
                                doc_result.seq = str::from_utf8(&seq).unwrap().parse().unwrap();
                                break;
                            }
                    }
                }
            }
            self.iter.next();

            if doc_result.seq > 0 {
                return Ok(Some(doc_result));
            }
        }
    }
}

struct DummyFilter {}

impl QueryRuntimeFilter for DummyFilter {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error> {
        Ok(None)
    }
    fn next_result(&mut self) -> Result<Option<DocResult>, Error> {
        Ok(None)
    }
}

struct AndFilter<'a> {
    filters: Vec<Box<QueryRuntimeFilter + 'a>>,
    current_filter: usize,
    array_depth: usize,
}

impl<'a> AndFilter<'a> {
    fn new(filters: Vec<Box<QueryRuntimeFilter + 'a>>, array_depth: usize) -> AndFilter<'a> {
        AndFilter {
            filters: filters,
            current_filter: 0,
            array_depth: array_depth,
        }
   }

    fn result(&mut self, base: Option<DocResult>) -> Result<Option<DocResult>, Error> {
        let mut matches_count = self.filters.len();
        // TODO vmx 2016-11-04: Make it nicer
        let mut base_result = match base {
            Some(base_result) => base_result,
            None => return Ok(None),
        };
        loop {
            base_result.truncate_array_paths(self.array_depth);

            self.current_filter += 1;
            if self.current_filter > self.filters.len() {
                self.current_filter = 0;
            }

            let next = try!(self.filters[self.current_filter].first_result(base_result.seq));
            let mut next_result = match next {
                Some(next_result) => next_result,
                None => return Ok(None),
            };
            next_result.truncate_array_paths(self.array_depth);

            if base_result.seq == next_result.seq {
                match DocResult::intersect_array_paths(&base_result, &next_result) {
                    Some(new_result) => {
                        base_result = new_result;
                        matches_count -= 1;
                    },
                    None => {
                        let new_result = try!(self.filters[self.current_filter].first_result(base_result.seq));
                        base_result = match new_result {
                            Some(base_result) => base_result,
                            None => return Ok(None),
                        };
                        matches_count = self.filters.len()
                    }
                }
            } else {
                base_result = next_result;
                matches_count = self.filters.len();
            }
        }
    }
}

impl<'a> QueryRuntimeFilter for AndFilter<'a> {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error> {
        let base_result = try!(self.filters[self.current_filter].first_result(start_id));
        self.result(base_result)
        //Ok(None)
    }

    fn next_result(&mut self) -> Result<Option<DocResult>, Error> {
        //let base_result = try!(self.filters[self.current_filter].next_result());
        //self.result(base_result)
        Ok(None)
    }
}

struct Parser<'a> {
    query: &'a str,
    offset: usize,
    kb: KeyBuilder,
    snapshot: &'a Snapshot<'a>,
}

impl<'a> Parser<'a> {
    fn new(query: &'a str, snapshot: &'a Snapshot<'a>) -> Parser<'a> {
        Parser{
            query: query,
            offset: 0,
            kb: KeyBuilder::new(),
            snapshot: &snapshot,
        }
    }

    fn whitespace(&mut self) {
        loop {
            if let Some(char) = self.query[self.offset..].chars().next() {
                // Stop when the character isn't a whitespace
                if !char.is_whitespace() {
                    break;
                }
                self.offset += char.len_utf8();
            }
        }
    }

    fn consume(&mut self, token: &str) -> bool {
        if self.could_consume(token) {
            self.offset += token.len();
            true
        } else {
            false
        }
    }

    fn could_consume(&mut self, token: &str) -> bool {
        self.query[self.offset..].starts_with(token)
    }

    fn consume_field(&mut self) -> Option<String> {
        let mut result = String::new();
        for char in self.query[self.offset..].chars() {
            if char.is_alphanumeric() {
                result.push(char);
            } else {
                break;
            }
        }
        if result.len() > 0 {
            self.offset += result.len();
            self.whitespace();
            Some(result)
        } else {
            None
        }
    }

    fn consume_string_literal(&mut self) -> Result<Option<String>, String> {
        // Does not unescape yet
        let mut lit = String::new();
        if self.consume("\"") {
            for char in self.query[self.offset..].chars() {
                if char != '"' {
                    lit.push(char);
                }
                self.offset += char.len_utf8();
            }
            if self.consume("\"") {
                self.whitespace();
                Ok(Some(lit))
            } else {
                Err("Expected \"".to_string())
            }
        } else {
            Ok(None)
        }
    }

    fn array() {
        //XXX GO ON HERE 2016-11-04
    }

    fn compare<'b>(&'b mut self) -> Result<Box<QueryRuntimeFilter + 'a>, String> {
        match self.consume_field() {
            Some(field) => {
                if self.consume(".") {
                    self.kb.push_object_key(field);
                    let ret = self.compare();
                    self.kb.pop_object_key();
                    ret
                } else if self.consume("=") {
                    match self.consume_string_literal() {
                        Ok(Some(literal)) => {
                            self.kb.push_object_key(field);

                            let stems = Stems::new(&literal);
                            let mut filters: Vec<Box<QueryRuntimeFilter + 'a>> = Vec::new();
                            for stem in stems {
                                let iter = self.snapshot.iter();
                                let filter = Box::new(ExactMatchFilter::new(
                                   iter, &stem, self.kb.clone()));
                                filters.push(filter);
                            }

                            self.kb.pop_object_key();

                            match filters.len() {
                                0 => panic!("Cannot create a ExactMatchFilter"),
                                1 => Ok(filters.pop().unwrap()),
                                _ => Ok(Box::new(AndFilter::<'a>::new(
                                    filters, self.kb.array_depth))),
                            }
                        },
                        // Empty literal
                        Ok(None) => {Err("Expected string".to_string())},
                        Err(error) => {
                            Err(error)
                        }
                    }
                } else if self.could_consume("[") {
                    self.kb.push_object_key(field);
                    //let ret = self.array();
                    self.kb.pop_object_key();
                    //ret
                    Err("Not yet implemented".to_string())
                } else {
                    Err("Expected comparison or array operator".to_string())
                }
            },
            None => {
                Err("Expected comparison or array operator".to_string())
            }
        }
    }

    fn build_filter(&mut self) -> Result<QueryResults, String> {
        self.whitespace();
        Ok(QueryResults{})
    }
}

impl Query {
    //pub fn get_matches<'a>(query: &str, index: &'a Index) -> Result<QueryResults, String> {
    pub fn get_matches<'a>(query: &str, snapshot: &'a Snapshot) -> Result<QueryResults, String> {
    //        match &index.rocks {
//            &Some(ref rocks) => {
//                let snapshot = Snapshot::new(rocks);
//                let parser = Parser::new(query, &snapshot);
//                Ok(QueryResults{})
//            },
//            &None => {
//                Err("You must open the index first".to_string())
//            },
//        }
        //let rocks = &index.rocks.unwrap();
        // This one would work as well
        //let rocks = index.rocks.as_ref().unwrap();
        //let snapshot = Snapshot::new(rocks);
        let parser = Parser::new(query, &snapshot);
        Ok(QueryResults{})
    }
}

#[cfg(test)]
mod tests {
    use super::Parser;

    use index::{Index, OpenOptions};

    use rocksdb::Snapshot;

    #[test]
    fn test_whitespace() {
        let mut index = Index::new();
        index.open("test_whitespace", Some(OpenOptions::Create)).unwrap();
        let rocks = &index.rocks.unwrap();
        let snapshot = Snapshot::new(rocks);

        let mut query = " \n \t test";
        let mut parser = Parser::new(query, &snapshot);
        parser.whitespace();
        assert_eq!(parser.offset, 5);

        query = "test";
        parser = Parser::new(query, &snapshot);
        parser.whitespace();
        assert_eq!(parser.offset, 0);
    }
}
<commit_msg>Adding array and bool filter<commit_after>#![allow(dead_code)]
#![allow(unused_variables)]
extern crate capnp;

use std::{error, fmt, str};

use index::Index;
use key_builder::KeyBuilder;
use stems::{StemmedWord, Stems};

// TODO vmx 2016-11-02: Make it import "rocksdb" properly instead of needing to import the individual tihngs
use rocksdb::{DBIterator, SeekKey};
use rocksdb::rocksdb::Snapshot;
use records_capnp::payload;

#[derive(Debug)]
enum Error {
    Parse(String),
    Capnp(capnp::Error),
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Parse(ref description) => description,
            Error::Capnp(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::Parse(_) => None,
            Error::Capnp(ref err) => Some(err as &error::Error),
        }
    }
}

impl From<capnp::Error> for Error {
    fn from(err: capnp::Error) -> Error {
        Error::Capnp(err)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Parse(ref err) => write!(f, "Parse error: {}", err),
            Error::Capnp(ref err) => write!(f, "Capnproto error: {}", err),
        }
    }
}

struct DocResult {
    seq: u64,
    array_paths: Vec<Vec<u64>>,
}

impl DocResult {
    fn new() -> DocResult {
        DocResult {
            seq: 0,
            array_paths: Vec::new(),
        }
    }

    fn truncate_array_paths(&mut self, array_depth: usize) {
        for array_path in self.array_paths.iter_mut() {
            debug_assert!(array_path.len() >= array_depth);
            array_path.resize(array_depth, 0);
        }
    }

    fn intersect_array_paths(aa: &DocResult, bb: &DocResult) -> Option<DocResult> {
        let mut doc_result = DocResult::new();
        debug_assert_eq!(aa.seq, bb.seq);
        doc_result.seq = aa.seq;
        for array_path_a in &aa.array_paths {
            for array_path_b in &bb.array_paths {
                if array_path_a == array_path_b {
                    doc_result.array_paths.push(array_path_a.clone());
                }
            }
        }
        if doc_result.array_paths.is_empty() {
            None
        } else {
            Some(doc_result)
        }
    }
}

//trait QueryRuntimeFilter {
//struct QueryRuntimeFilter {}

trait QueryRuntimeFilter {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error>;
    fn next_result(&mut self) -> Result<Option<DocResult>, Error>;
}

pub struct Query {}

pub struct QueryResults {
}

//struct SnapshotIteratorCreator {
//    snapshot: rocksdb::Snapshot,
//}
//
//impl SnapshotIteratorCreator {
//    fn new(db: &rocksdb::DB) {
//        let snapshot = rocksdb::Snapshot::new(db);
//        SnapshotIteratorCreator{
//            snapshot: snapshot,
//        }
//    }
//
//    fn new_iterator(&self) {
//        self.snapshot.iter()
//    }
//}

struct ExactMatchFilter<'a> {
    iter: DBIterator<'a>,
    kb: KeyBuilder,
    stemmed_offset: u64,
    suffix: String,
    suffix_offset: u64,
}

impl<'a> ExactMatchFilter<'a> {
    fn new(iter: DBIterator<'a>, stemmed_word: &StemmedWord, mut kb: KeyBuilder) -> ExactMatchFilter<'a> {
        kb.push_word(&stemmed_word.stemmed);
        ExactMatchFilter{
            iter: iter,
            kb: kb,
            stemmed_offset: stemmed_word.stemmed_offset as u64,
            suffix: stemmed_word.suffix.clone(),
            suffix_offset: stemmed_word.suffix_offset as u64,
        }
    }
}

impl<'a> QueryRuntimeFilter for ExactMatchFilter<'a> {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error> {
        // Build the full key
        self.kb.push_doc_seq(start_id);

        // Seek in index to >= entry
        self.iter.seek(SeekKey::from(self.kb.key().as_bytes()));

        // Revert
        self.kb.pop_doc_seq();

        self.next_result()
    }

    fn next_result(&mut self) -> Result<Option<DocResult>, Error> {
        let mut doc_result = DocResult::new();

        loop {
            if !self.iter.valid() {
                return Ok(None)
            }

            // New scope needed as the iter.next() below invalidates the
            // current key and value
            {
                let key = self.iter.key();
                if !key.starts_with(self.kb.key().as_bytes()) {
                    return Ok(None)
                }
                let seq = &key[self.kb.key().len()..];

                let value = self.iter.value();
                // NOTE vmx 2016-10-13: I'm not really sure why the dereferencing is needed
                // and why we pass on mutable reference of it to `read_message()`
                let mut ref_value = &*value;
                let message_reader = ::capnp::serialize_packed::read_message(
                    &mut ref_value, ::capnp::message::ReaderOptions::new()).unwrap();
                let payload = message_reader.get_root::<payload::Reader>().unwrap();

                for aos_wis in try!(payload.get_arrayoffsets_to_wordinfos()).iter() {
                    for wi in try!(aos_wis.get_wordinfos()).iter() {
                        if self.stemmed_offset == wi.get_stemmed_offset() &&
                            self.suffix_offset == wi.get_suffix_offset() &&
                            self.suffix == try!(wi.get_suffix_text()) {
                                // We have a candidate document to return
                                let arrayoffsets = try!(aos_wis.get_arrayoffsets());
                                doc_result.array_paths.push(arrayoffsets.iter().collect::<>());
                                doc_result.seq = str::from_utf8(&seq).unwrap().parse().unwrap();
                                break;
                            }
                    }
                }
            }
            self.iter.next();

            if doc_result.seq > 0 {
                return Ok(Some(doc_result));
            }
        }
    }
}

struct DummyFilter {}

impl QueryRuntimeFilter for DummyFilter {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error> {
        Ok(None)
    }
    fn next_result(&mut self) -> Result<Option<DocResult>, Error> {
        Ok(None)
    }
}

struct AndFilter<'a> {
    filters: Vec<Box<QueryRuntimeFilter + 'a>>,
    current_filter: usize,
    array_depth: usize,
}

impl<'a> AndFilter<'a> {
    fn new(filters: Vec<Box<QueryRuntimeFilter + 'a>>, array_depth: usize) -> AndFilter<'a> {
        AndFilter {
            filters: filters,
            current_filter: 0,
            array_depth: array_depth,
        }
   }

    fn result(&mut self, base: Option<DocResult>) -> Result<Option<DocResult>, Error> {
        let mut matches_count = self.filters.len();
        // TODO vmx 2016-11-04: Make it nicer
        let mut base_result = match base {
            Some(base_result) => base_result,
            None => return Ok(None),
        };
        loop {
            base_result.truncate_array_paths(self.array_depth);

            self.current_filter += 1;
            if self.current_filter > self.filters.len() {
                self.current_filter = 0;
            }

            let next = try!(self.filters[self.current_filter].first_result(base_result.seq));
            let mut next_result = match next {
                Some(next_result) => next_result,
                None => return Ok(None),
            };
            next_result.truncate_array_paths(self.array_depth);

            if base_result.seq == next_result.seq {
                match DocResult::intersect_array_paths(&base_result, &next_result) {
                    Some(new_result) => {
                        base_result = new_result;
                        matches_count -= 1;
                    },
                    None => {
                        let new_result = try!(self.filters[self.current_filter].first_result(base_result.seq));
                        base_result = match new_result {
                            Some(base_result) => base_result,
                            None => return Ok(None),
                        };
                        matches_count = self.filters.len()
                    }
                }
            } else {
                base_result = next_result;
                matches_count = self.filters.len();
            }
        }
    }
}

impl<'a> QueryRuntimeFilter for AndFilter<'a> {
    fn first_result(&mut self, start_id: u64) -> Result<Option<DocResult>, Error> {
        let base_result = try!(self.filters[self.current_filter].first_result(start_id));
        self.result(base_result)
        //Ok(None)
    }

    fn next_result(&mut self) -> Result<Option<DocResult>, Error> {
        //let base_result = try!(self.filters[self.current_filter].next_result());
        //self.result(base_result)
        Ok(None)
    }
}

struct Parser<'a> {
    query: &'a str,
    offset: usize,
    kb: KeyBuilder,
    snapshot: &'a Snapshot<'a>,
}

impl<'a> Parser<'a> {
    fn new(query: &'a str, snapshot: &'a Snapshot<'a>) -> Parser<'a> {
        Parser{
            query: query,
            offset: 0,
            kb: KeyBuilder::new(),
            snapshot: &snapshot,
        }
    }

    fn whitespace(&mut self) {
        loop {
            if let Some(char) = self.query[self.offset..].chars().next() {
                // Stop when the character isn't a whitespace
                if !char.is_whitespace() {
                    break;
                }
                self.offset += char.len_utf8();
            }
        }
    }

    fn consume(&mut self, token: &str) -> bool {
        if self.could_consume(token) {
            self.offset += token.len();
            true
        } else {
            false
        }
    }

    fn could_consume(&mut self, token: &str) -> bool {
        self.query[self.offset..].starts_with(token)
    }

    fn consume_field(&mut self) -> Option<String> {
        let mut result = String::new();
        for char in self.query[self.offset..].chars() {
            if char.is_alphanumeric() {
                result.push(char);
            } else {
                break;
            }
        }
        if result.len() > 0 {
            self.offset += result.len();
            self.whitespace();
            Some(result)
        } else {
            None
        }
    }

    fn consume_string_literal(&mut self) -> Result<Option<String>, Error> {
        // Does not unescape yet
        let mut lit = String::new();
        if self.consume("\"") {
            for char in self.query[self.offset..].chars() {
                if char != '"' {
                    lit.push(char);
                }
                self.offset += char.len_utf8();
            }
            if self.consume("\"") {
                self.whitespace();
                Ok(Some(lit))
            } else {
                Err(Error::Parse("Expected \"".to_string()))
            }
        } else {
            Ok(None)
        }
    }

    fn bool(&mut self) -> Result<Box<QueryRuntimeFilter + 'a>, Error> {
        let left = try!(self.compare());
        let mut filters = vec![left];
        loop {
            if !self.consume("&") {
                break;
            }

            let right = try!(self.compare());
            filters.push(right);
        }
        if filters.len() == 1 {
            Ok(filters.pop().unwrap())
        } else {
            Ok(Box::new(AndFilter::new(filters, self.kb.array_depth)))
        }
    }

    fn array(&mut self) -> Result<Box<QueryRuntimeFilter + 'a>, Error> {
        if !self.consume("[") {
            return Err(Error::Parse("Expected '['".to_string()));
        }
        self.kb.push_array();
        let filter = try!(self.bool());
        self.kb.pop_array();
        if self.consume("]") {
            return Err(Error::Parse("Expected ']'".to_string()));
        }
        Ok(filter)
    }

    fn compare<'b>(&'b mut self) -> Result<Box<QueryRuntimeFilter + 'a>, Error> {
        match self.consume_field() {
            Some(field) => {
                if self.consume(".") {
                    self.kb.push_object_key(field);
                    let ret = self.compare();
                    self.kb.pop_object_key();
                    ret
                } else if self.consume("=") {
                    match self.consume_string_literal() {
                        Ok(Some(literal)) => {
                            self.kb.push_object_key(field);

                            let stems = Stems::new(&literal);
                            let mut filters: Vec<Box<QueryRuntimeFilter + 'a>> = Vec::new();
                            for stem in stems {
                                let iter = self.snapshot.iter();
                                let filter = Box::new(ExactMatchFilter::new(
                                   iter, &stem, self.kb.clone()));
                                filters.push(filter);
                            }

                            self.kb.pop_object_key();

                            match filters.len() {
                                0 => panic!("Cannot create a ExactMatchFilter"),
                                1 => Ok(filters.pop().unwrap()),
                                _ => Ok(Box::new(AndFilter::<'a>::new(
                                    filters, self.kb.array_depth))),
                            }
                        },
                        // Empty literal
                        Ok(None) => {Err(Error::Parse("Expected string".to_string()))},
                        Err(error) => {
                            Err(error)
                        }
                    }
                } else if self.could_consume("[") {
                    self.kb.push_object_key(field);
                    let ret = self.array();
                    self.kb.pop_object_key();
                    ret
                } else {
                    Err(Error::Parse("Expected comparison or array operator".to_string()))
                }
            },
            None => {
                Err(Error::Parse("Expected comparison or array operator".to_string()))
            }
        }
    }

    fn build_filter(&mut self) -> Result<QueryResults, String> {
        self.whitespace();
        Ok(QueryResults{})
    }
}

impl Query {
    //pub fn get_matches<'a>(query: &str, index: &'a Index) -> Result<QueryResults, String> {
    pub fn get_matches<'a>(query: &str, snapshot: &'a Snapshot) -> Result<QueryResults, String> {
    //        match &index.rocks {
//            &Some(ref rocks) => {
//                let snapshot = Snapshot::new(rocks);
//                let parser = Parser::new(query, &snapshot);
//                Ok(QueryResults{})
//            },
//            &None => {
//                Err("You must open the index first".to_string())
//            },
//        }
        //let rocks = &index.rocks.unwrap();
        // This one would work as well
        //let rocks = index.rocks.as_ref().unwrap();
        //let snapshot = Snapshot::new(rocks);
        let parser = Parser::new(query, &snapshot);
        Ok(QueryResults{})
    }
}

#[cfg(test)]
mod tests {
    use super::Parser;

    use index::{Index, OpenOptions};

    use rocksdb::Snapshot;

    #[test]
    fn test_whitespace() {
        let mut index = Index::new();
        index.open("test_whitespace", Some(OpenOptions::Create)).unwrap();
        let rocks = &index.rocks.unwrap();
        let snapshot = Snapshot::new(rocks);

        let mut query = " \n \t test";
        let mut parser = Parser::new(query, &snapshot);
        parser.whitespace();
        assert_eq!(parser.offset, 5);

        query = "test";
        parser = Parser::new(query, &snapshot);
        parser.whitespace();
        assert_eq!(parser.offset, 0);
    }
}
<|endoftext|>


<commit_before>/// A Query organizes a search of an inverted index.
/// It is recursively hierarchical, allowing flexibility
/// in exactly how a search is specified
#[derive(Debug)]
pub enum Query<'a> {
    /// The simplest query, represents a search using
    /// the given string
    Match(&'a str),

    /// A query requesting the intersection of the documents
    /// returned in each sub-query
    And(&'a [Query<'a>]),

    /// A query requesting the union of the documents returned
    /// in each sub-query
    Or(&'a [Query<'a>]),

    /// An exact-match query. The given phrase must appear in all documents returned.
    /// False positives may occur.
    Phrase(&'a str),

    /// A prefix query that returns all documents containing terms with the given prefix.
    /// Note that, unlike `Match` and `Phrase`, this query is not tokenized before searching
    /// the index. Thus, Prefix("hi bob") is likely to match zero documents, since indexed
    /// documents typically have their content tokenized upon spaces.
    Prefix(&'a str),
}
<commit_msg>Derive all possible traits for Query<commit_after>/// A Query organizes a search of an inverted index.
/// It is recursively hierarchical, allowing flexibility
/// in exactly how a search is specified
#[derive(Copy, Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, RustcEncodable)]
pub enum Query<'a> {
    /// The simplest query, represents a search using
    /// the given string
    Match(&'a str),

    /// A query requesting the intersection of the documents
    /// returned in each sub-query
    And(&'a [Query<'a>]),

    /// A query requesting the union of the documents returned
    /// in each sub-query
    Or(&'a [Query<'a>]),

    /// An exact-match query. The given phrase must appear in all documents returned.
    /// False positives may occur.
    Phrase(&'a str),

    /// A prefix query that returns all documents containing terms with the given prefix.
    /// Note that, unlike `Match` and `Phrase`, this query is not tokenized before searching
    /// the index. Thus, Prefix("hi bob") is likely to match zero documents, since indexed
    /// documents typically have their content tokenized upon spaces.
    Prefix(&'a str),
}
<|endoftext|>