<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! A tokenizer is a pattern matcher that is intended to turn a stream of symbols into another stream of symbols based on the patterns
//! that are matched. Every pattern can produce a different output symbol. If two input strings can ndfa in two different output
//! symbols, then the output symbol that is ordered lower is the one that's produced (ie, if the output symbols are numbers, then '0' will
//! be produced instead of '1' in the event of a clash)
//!

use super::countable::*;
use super::symbol_range::*;
use super::regular_pattern::*;
use super::state_machine::*;
use super::ndfa::*;
use super::prepare::*;
use super::symbol_range_dfa::*;
use super::symbol_reader::*;
use super::pattern_matcher::*;
use super::matches::*;
use super::tape::*;

///
/// Used for generating tokenizing pattern matchers
///
pub struct TokenMatcher<InputSymbol: Clone+Ord+Countable, OutputSymbol: Clone+Ord> {
    patterns: Vec<(Pattern<InputSymbol>, OutputSymbol)>
}

impl<InputSymbol: Clone+Ord+Countable+'static, OutputSymbol: Clone+Ord+'static> TokenMatcher<InputSymbol, OutputSymbol> {
    ///
    /// Creates a new TokenMatcher
    ///
    pub fn new() -> TokenMatcher<InputSymbol, OutputSymbol> {
        TokenMatcher { patterns: vec![] }
    }

    ///
    /// Adds a new pattern that will generate the specified output symbol
    ///
    pub fn add_pattern<TPattern: ToPattern<InputSymbol>>(&mut self, pattern: TPattern, output: OutputSymbol) {
        self.patterns.push((pattern.to_pattern(), output));
    }

    ///
    /// Compiles an NDFA from this TokenMatcher
    ///
    pub fn to_ndfa(&self) -> Box<StateMachine<SymbolRange<InputSymbol>, OutputSymbol>> {
        let mut ndfa = Ndfa::new();

        for &(ref pattern, ref output) in &self.patterns {
            // Compile each pattern starting at state 0
            let end_state = pattern.compile(&mut ndfa, 0);

            // Set the output for this pattern
            ndfa.set_output_symbol(end_state, output.clone());
        }

        // Clear out any overlapping ranges so we can build a valid DFA
        ndfa.fix_overlapping_ranges();

        Box::new(ndfa)
    }
}

impl<'a, InputSymbol: Clone+Ord+Countable+'static, OutputSymbol: Clone+Ord+'static> PrepareToMatch<SymbolRangeDfa<InputSymbol, OutputSymbol>>
for &'a TokenMatcher<InputSymbol, OutputSymbol> {
    #[inline]
    fn prepare_to_match(self) -> SymbolRangeDfa<InputSymbol, OutputSymbol> {
        let ndfa = self.to_ndfa();

        ndfa.prepare_to_match()
    }
}

///
/// A tokenizer is a type of symbol stream that uses a pattern matcher to convert a symbol stream into a stream of tokens
///
pub struct Tokenizer<'a, InputSymbol: Clone+Ord+Countable+'a, OutputSymbol: Clone+Ord+'a, Reader: SymbolReader<InputSymbol>> {
    /// The pattern matcher for this tokenizer
    dfa: &'a SymbolRangeDfa<InputSymbol, OutputSymbol>,

    /// Tape of input symbols that will be used to generate the result
    tape: Tape<InputSymbol, Reader>,
}

impl<'a, InputSymbol: Clone+Ord+Countable, OutputSymbol: Clone+Ord, Reader: SymbolReader<InputSymbol>> Tokenizer<'a, InputSymbol, OutputSymbol, Reader> {
    ///
    /// Creates a new tokenizer from a pattern (usually a TokenMatcher)
    ///
    pub fn new<'b, Prepare: PrepareToMatch<SymbolRangeDfa<InputSymbol, OutputSymbol>>>(source: Reader, pattern: Prepare) -> Tokenizer<'b, InputSymbol, OutputSymbol, Reader> {
        /*
        let generated_dfa   = RefCell::new(Some(pattern.prepare_to_match())).borrow();
        let dfa             = generated_dfa;
        let tape            = Tape::new(source);

        Tokenizer { generated_dfa: generated_dfa, dfa: dfa, tape: tape }
        */
        unimplemented!()
    }

    ///
    /// Creates a new tokenizer from a prepared pattern
    ///
    pub fn new_prepared<'b>(source: Reader, pattern: &'b SymbolRangeDfa<InputSymbol, OutputSymbol>) -> Tokenizer<'b, InputSymbol, OutputSymbol, Reader> {
        Tokenizer { dfa: pattern, tape: Tape::new(source) }
    }

    ///
    /// Returns the current position in the source (the position after the last matched symbol)
    ///
    pub fn get_source_position(&self) -> usize {
        self.tape.get_source_position()
    }

    ///
    /// Skips an input symbol (returning the symbol that was skipped)
    ///
    pub fn skip_input(&mut self) -> Option<InputSymbol> {
        self.tape.next_symbol()
    }

    ///
    /// True if we've reached the end of the source reader
    ///
    /// If `next_symbol` returns `None` and `at_end_of_reader` is false, then the input stream does not contain a symbol matching the DFA
    ///
    pub fn at_end_of_reader(&self) -> bool {
        self.tape.at_end_of_reader()
    }
}

impl<'a, InputSymbol: Clone+Ord+Countable, OutputSymbol: Clone+Ord+'static, Reader: SymbolReader<InputSymbol>> SymbolReader<OutputSymbol> for Tokenizer<'a, InputSymbol, OutputSymbol, Reader> {
    fn next_symbol(&mut self) -> Option<OutputSymbol> {
        // Start of the next symbol
        let start_pos = self.tape.get_source_position();

        // Match against it
        let match_result = match_pattern(self.dfa.start(), &mut self.tape);

        let end_pos = self.tape.get_source_position();
        match match_result {
            Accept(length, outputsymbol) => {
                if length > 0 {
                    // Rewind the tape to after the accepted symbol
                    self.tape.rewind(end_pos-start_pos - length);

                    // Won't try to match anything before this position
                    self.tape.cut();

                    // Result is the oputput symbol
                    Some(outputsymbol.clone())
                } else {
                    // Zero-length match
                    // If we accepted matches of length 0 we'd get an infinite stream when we hit a symbol that doesn't match
                    self.tape.rewind(end_pos-start_pos);

                    // Return no match
                    None
                }
            },

            Reject => {
                // Rewind back to the start position
                self.tape.rewind(end_pos-start_pos);

                // No match
                None
            },

            _ => {
                panic!("Unexpected output state of ");
            }
        }
    }
}

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

    #[test]
    fn can_match_tokens_like_any_other_pattern() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            AllAs,
            AllBs
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(1), TestToken::AllAs);
        token_matcher.add_pattern("b".repeat_forever(1), TestToken::AllBs);

        assert!(matches("aaaa", &token_matcher) == Some(4));
        assert!(matches("bbbbb", &token_matcher) == Some(5));
        assert!(matches("abbb", &token_matcher) == Some(1));
        assert!(matches("bbaaa", &token_matcher) == Some(2));
    }

    #[test]
    fn can_distinguish_simple_tokens() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            AllAs,
            AllBs
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(1), TestToken::AllAs);
        token_matcher.add_pattern("b".repeat_forever(1), TestToken::AllBs);

        let matcher = token_matcher.prepare_to_match();

        assert!(match_pattern(matcher.start(), &mut "aaaaa".read_symbols()).is_accepted(&TestToken::AllAs));
        assert!(match_pattern(matcher.start(), &mut "bbbb".read_symbols()).is_accepted(&TestToken::AllBs));
    }

    #[test]
    fn can_distinguish_simple_tokens_with_zero_repeats() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            AllAs,
            AllBs
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(0), TestToken::AllAs);
        token_matcher.add_pattern("b".repeat_forever(0), TestToken::AllBs);

        let matcher = token_matcher.prepare_to_match();

        assert!(match_pattern(matcher.start(), &mut "aaaaa".read_symbols()).is_accepted(&TestToken::AllAs));
        assert!(match_pattern(matcher.start(), &mut "bbbb".read_symbols()).is_accepted(&TestToken::AllBs));
    }

    #[test]
    fn clashes_producer_lower_tokens() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Abbb,
            Aaab
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(1).append("b"), TestToken::Aaab);
        token_matcher.add_pattern("a".append("b".repeat_forever(1)), TestToken::Abbb);

        let matcher = token_matcher.prepare_to_match();

        assert!(match_pattern(matcher.start(), &mut "aaab".read_symbols()).is_accepted(&TestToken::Aaab));
        assert!(match_pattern(matcher.start(), &mut "ab".read_symbols()).is_accepted(&TestToken::Abbb));
        assert!(match_pattern(matcher.start(), &mut "abbbb".read_symbols()).is_accepted(&TestToken::Abbb));
    }

    #[test]
    fn can_match_number_stream() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat_forever(1), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat_forever(1), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 390  32 ".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.get_source_position() == 2);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.get_source_position() == 3);
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.get_source_position() == 6);
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.get_source_position() == 8);
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.get_source_position() == 10);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == None);
    }

    #[test]
    fn can_skip_bad_symbols() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat_forever(1), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat_forever(1), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 ab 12".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('a'));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.skip_input() == Some('b'));
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.at_end_of_reader());
    }

    #[test]
    fn wont_match_zero_length() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat_forever(0), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat_forever(0), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 ab 12".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('a'));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('b'));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.at_end_of_reader());
    }

    #[test]
    fn wont_match_zero_length_2() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat(0..4), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat(0..4), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 ab 12".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('a'));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('b'));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.at_end_of_reader());
    }
}
<commit_msg>Need to declare an enum and do a bunch of crap to work around rust<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! A tokenizer is a pattern matcher that is intended to turn a stream of symbols into another stream of symbols based on the patterns
//! that are matched. Every pattern can produce a different output symbol. If two input strings can ndfa in two different output
//! symbols, then the output symbol that is ordered lower is the one that's produced (ie, if the output symbols are numbers, then '0' will
//! be produced instead of '1' in the event of a clash)
//!

use super::countable::*;
use super::symbol_range::*;
use super::regular_pattern::*;
use super::state_machine::*;
use super::ndfa::*;
use super::prepare::*;
use super::symbol_range_dfa::*;
use super::symbol_reader::*;
use super::pattern_matcher::*;
use super::matches::*;
use super::tape::*;

///
/// Used for generating tokenizing pattern matchers
///
pub struct TokenMatcher<InputSymbol: Clone+Ord+Countable, OutputSymbol: Clone+Ord> {
    patterns: Vec<(Pattern<InputSymbol>, OutputSymbol)>
}

impl<InputSymbol: Clone+Ord+Countable+'static, OutputSymbol: Clone+Ord+'static> TokenMatcher<InputSymbol, OutputSymbol> {
    ///
    /// Creates a new TokenMatcher
    ///
    pub fn new() -> TokenMatcher<InputSymbol, OutputSymbol> {
        TokenMatcher { patterns: vec![] }
    }

    ///
    /// Adds a new pattern that will generate the specified output symbol
    ///
    pub fn add_pattern<TPattern: ToPattern<InputSymbol>>(&mut self, pattern: TPattern, output: OutputSymbol) {
        self.patterns.push((pattern.to_pattern(), output));
    }

    ///
    /// Compiles an NDFA from this TokenMatcher
    ///
    pub fn to_ndfa(&self) -> Box<StateMachine<SymbolRange<InputSymbol>, OutputSymbol>> {
        let mut ndfa = Ndfa::new();

        for &(ref pattern, ref output) in &self.patterns {
            // Compile each pattern starting at state 0
            let end_state = pattern.compile(&mut ndfa, 0);

            // Set the output for this pattern
            ndfa.set_output_symbol(end_state, output.clone());
        }

        // Clear out any overlapping ranges so we can build a valid DFA
        ndfa.fix_overlapping_ranges();

        Box::new(ndfa)
    }
}

impl<'a, InputSymbol: Clone+Ord+Countable+'static, OutputSymbol: Clone+Ord+'static> PrepareToMatch<SymbolRangeDfa<InputSymbol, OutputSymbol>>
for &'a TokenMatcher<InputSymbol, OutputSymbol> {
    #[inline]
    fn prepare_to_match(self) -> SymbolRangeDfa<InputSymbol, OutputSymbol> {
        let ndfa = self.to_ndfa();

        ndfa.prepare_to_match()
    }
}

///
/// Thing that can be a reference or the actual thing
///
/// An elegant way of dealing with the 'sometimes we store a reference, but sometimes we also need the actual thing' problem
/// that we have here would be for rust to support self-references in structures (or perhaps some other way - via RefCell or Box
/// perhaps - not possible in any obvious way at present due to borrowing order problems).
///
enum ReferenceOrOwned<'a, Thing: 'a> {
    Reference(&'a Thing),
    Owned(Thing)
}

use self::ReferenceOrOwned::*;

impl<'a, Thing> ReferenceOrOwned<'a, Thing> {
    #[inline]
    fn get(&'a self) -> &'a Thing {
        match self {
            &Reference(ref thing)    => *thing,
            &Owned(ref thing)        => thing
        }
    }
}

///
/// A tokenizer is a type of symbol stream that uses a pattern matcher to convert a symbol stream into a stream of tokens
///
pub struct Tokenizer<'a, InputSymbol: Clone+Ord+Countable+'a, OutputSymbol: Clone+Ord+'a, Reader: SymbolReader<InputSymbol>> {
    /// The pattern matcher for this tokenizer
    dfa: ReferenceOrOwned<'a, SymbolRangeDfa<InputSymbol, OutputSymbol>>,

    /// Tape of input symbols that will be used to generate the result
    tape: Tape<InputSymbol, Reader>,
}

impl<'a, InputSymbol: Clone+Ord+Countable, OutputSymbol: Clone+Ord, Reader: SymbolReader<InputSymbol>> Tokenizer<'a, InputSymbol, OutputSymbol, Reader> {
    ///
    /// Creates a new tokenizer from a pattern (usually a TokenMatcher)
    ///
    pub fn new<'b, Prepare: PrepareToMatch<SymbolRangeDfa<InputSymbol, OutputSymbol>>>(source: Reader, pattern: Prepare) -> Tokenizer<'b, InputSymbol, OutputSymbol, Reader> {
        Tokenizer { dfa: Owned(pattern.prepare_to_match()), tape: Tape::new(source) }
    }

    ///
    /// Creates a new tokenizer from a prepared pattern
    ///
    pub fn new_prepared<'b>(source: Reader, pattern: &'b SymbolRangeDfa<InputSymbol, OutputSymbol>) -> Tokenizer<'b, InputSymbol, OutputSymbol, Reader> {
        Tokenizer { dfa: Reference(pattern), tape: Tape::new(source) }
    }

    ///
    /// Returns the current position in the source (the position after the last matched symbol)
    ///
    pub fn get_source_position(&self) -> usize {
        self.tape.get_source_position()
    }

    ///
    /// Skips an input symbol (returning the symbol that was skipped)
    ///
    pub fn skip_input(&mut self) -> Option<InputSymbol> {
        self.tape.next_symbol()
    }

    ///
    /// True if we've reached the end of the source reader
    ///
    /// If `next_symbol` returns `None` and `at_end_of_reader` is false, then the input stream does not contain a symbol matching the DFA
    ///
    pub fn at_end_of_reader(&self) -> bool {
        self.tape.at_end_of_reader()
    }
}

impl<'a, InputSymbol: Clone+Ord+Countable, OutputSymbol: Clone+Ord+'static, Reader: SymbolReader<InputSymbol>> SymbolReader<OutputSymbol> for Tokenizer<'a, InputSymbol, OutputSymbol, Reader> {
    fn next_symbol(&mut self) -> Option<OutputSymbol> {
        // Start of the next symbol
        let start_pos = self.tape.get_source_position();

        // Match against it
        let match_result = match_pattern(self.dfa.get().start(), &mut self.tape);

        let end_pos = self.tape.get_source_position();
        match match_result {
            Accept(length, outputsymbol) => {
                if length > 0 {
                    // Rewind the tape to after the accepted symbol
                    self.tape.rewind(end_pos-start_pos - length);

                    // Won't try to match anything before this position
                    self.tape.cut();

                    // Result is the oputput symbol
                    Some(outputsymbol.clone())
                } else {
                    // Zero-length match
                    // If we accepted matches of length 0 we'd get an infinite stream when we hit a symbol that doesn't match
                    self.tape.rewind(end_pos-start_pos);

                    // Return no match
                    None
                }
            },

            Reject => {
                // Rewind back to the start position
                self.tape.rewind(end_pos-start_pos);

                // No match
                None
            },

            _ => {
                panic!("Unexpected output state of ");
            }
        }
    }
}

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

    #[test]
    fn can_match_tokens_like_any_other_pattern() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            AllAs,
            AllBs
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(1), TestToken::AllAs);
        token_matcher.add_pattern("b".repeat_forever(1), TestToken::AllBs);

        assert!(matches("aaaa", &token_matcher) == Some(4));
        assert!(matches("bbbbb", &token_matcher) == Some(5));
        assert!(matches("abbb", &token_matcher) == Some(1));
        assert!(matches("bbaaa", &token_matcher) == Some(2));
    }

    #[test]
    fn can_distinguish_simple_tokens() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            AllAs,
            AllBs
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(1), TestToken::AllAs);
        token_matcher.add_pattern("b".repeat_forever(1), TestToken::AllBs);

        let matcher = token_matcher.prepare_to_match();

        assert!(match_pattern(matcher.start(), &mut "aaaaa".read_symbols()).is_accepted(&TestToken::AllAs));
        assert!(match_pattern(matcher.start(), &mut "bbbb".read_symbols()).is_accepted(&TestToken::AllBs));
    }

    #[test]
    fn can_distinguish_simple_tokens_with_zero_repeats() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            AllAs,
            AllBs
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(0), TestToken::AllAs);
        token_matcher.add_pattern("b".repeat_forever(0), TestToken::AllBs);

        let matcher = token_matcher.prepare_to_match();

        assert!(match_pattern(matcher.start(), &mut "aaaaa".read_symbols()).is_accepted(&TestToken::AllAs));
        assert!(match_pattern(matcher.start(), &mut "bbbb".read_symbols()).is_accepted(&TestToken::AllBs));
    }

    #[test]
    fn clashes_producer_lower_tokens() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Abbb,
            Aaab
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern("a".repeat_forever(1).append("b"), TestToken::Aaab);
        token_matcher.add_pattern("a".append("b".repeat_forever(1)), TestToken::Abbb);

        let matcher = token_matcher.prepare_to_match();

        assert!(match_pattern(matcher.start(), &mut "aaab".read_symbols()).is_accepted(&TestToken::Aaab));
        assert!(match_pattern(matcher.start(), &mut "ab".read_symbols()).is_accepted(&TestToken::Abbb));
        assert!(match_pattern(matcher.start(), &mut "abbbb".read_symbols()).is_accepted(&TestToken::Abbb));
    }

    #[test]
    fn can_match_number_stream() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat_forever(1), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat_forever(1), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 390  32 ".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.get_source_position() == 2);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.get_source_position() == 3);
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.get_source_position() == 6);
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.get_source_position() == 8);
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.get_source_position() == 10);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == None);
    }

    #[test]
    fn can_skip_bad_symbols() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat_forever(1), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat_forever(1), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 ab 12".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('a'));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.skip_input() == Some('b'));
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.at_end_of_reader());
    }

    #[test]
    fn wont_match_zero_length() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat_forever(0), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat_forever(0), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 ab 12".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('a'));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('b'));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.at_end_of_reader());
    }

    #[test]
    fn wont_match_zero_length_2() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Digit,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat(0..4), TestToken::Digit);
        token_matcher.add_pattern(" ".repeat(0..4), TestToken::Whitespace);

        let mut tokenizer = Tokenizer::new("12 ab 12".read_symbols(), &token_matcher);

        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('a'));
        assert!(tokenizer.next_symbol() == None);
        assert!(!tokenizer.at_end_of_reader());
        assert!(tokenizer.skip_input() == Some('b'));
        assert!(tokenizer.next_symbol() == Some(TestToken::Whitespace));
        assert!(tokenizer.next_symbol() == Some(TestToken::Digit));
        assert!(tokenizer.next_symbol() == None);
        assert!(tokenizer.at_end_of_reader());
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;

use internal;

use packet::Packet;
use packet::ip::IpNextHeaderProtocol;
use packet::ipv4::Ipv4Packet;
use packet::udp::UdpPacket;
use packet::icmp::IcmpPacket;
use self::TransportChannelType::{Layer3, Layer4};
use self::TransportProtocol::{Ipv4, Ipv6};
use sockets;

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::mem;
use std::net::{self, IpAddr};
use std::sync::Arc;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<internal::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for receiving at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<internal::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) |
            Layer3(IpNextHeaderProtocol(proto)) => {
                sockets::socket(sockets::AF_INET, sockets::SOCK_RAW, proto as libc::c_int)
            }
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) => {
                sockets::socket(sockets::AF_INET6, sockets::SOCK_RAW, proto as libc::c_int)
            }
        }
    };
    if socket == sockets::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            sockets::setsockopt(socket,
                                sockets::IPPROTO_IP,
                                sockets::IP_HDRINCL,
                                (&hincl as *const libc::c_int) as sockets::Buf,
                                mem::size_of::<libc::c_int>() as sockets::SockLen)
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                sockets::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(internal::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: repeat(0u8).take(buffer_size).collect(),
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = internal::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const sockets::SockAddrStorage) as *const sockets::SockAddr;

        internal::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided destination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use packet::MutablePacket;
        use packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: sockets::SockAddrStorage = unsafe { mem::zeroed() };
                let res = internal::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = internal::sockaddr_to_addr(
                                        &caddr,
                                        mem::size_of::<sockets::SockAddrStorage>()
                                   );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet, Ipv4TransportChannelIterator, ipv4_packet_iter);

transport_channel_iterator!(UdpPacket, UdpTransportChannelIterator, udp_packet_iter);

transport_channel_iterator!(IcmpPacket, IcmpTransportChannelIterator, icmp_packet_iter);<commit_msg>Add tcp_packet_iter function.<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;

use internal;

use packet::Packet;
use packet::ip::IpNextHeaderProtocol;
use packet::ipv4::Ipv4Packet;
use packet::udp::UdpPacket;
use packet::icmp::IcmpPacket;
use packet::tcp::TcpPacket;
use self::TransportChannelType::{Layer3, Layer4};
use self::TransportProtocol::{Ipv4, Ipv6};
use sockets;

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::mem;
use std::net::{self, IpAddr};
use std::sync::Arc;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<internal::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for receiving at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<internal::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) |
            Layer3(IpNextHeaderProtocol(proto)) => {
                sockets::socket(sockets::AF_INET, sockets::SOCK_RAW, proto as libc::c_int)
            }
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) => {
                sockets::socket(sockets::AF_INET6, sockets::SOCK_RAW, proto as libc::c_int)
            }
        }
    };
    if socket == sockets::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            sockets::setsockopt(socket,
                                sockets::IPPROTO_IP,
                                sockets::IP_HDRINCL,
                                (&hincl as *const libc::c_int) as sockets::Buf,
                                mem::size_of::<libc::c_int>() as sockets::SockLen)
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                sockets::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(internal::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: repeat(0u8).take(buffer_size).collect(),
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = internal::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const sockets::SockAddrStorage) as *const sockets::SockAddr;

        internal::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided destination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use packet::MutablePacket;
        use packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: sockets::SockAddrStorage = unsafe { mem::zeroed() };
                let res = internal::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = internal::sockaddr_to_addr(
                                        &caddr,
                                        mem::size_of::<sockets::SockAddrStorage>()
                                   );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet, Ipv4TransportChannelIterator, ipv4_packet_iter);

transport_channel_iterator!(UdpPacket, UdpTransportChannelIterator, udp_packet_iter);

transport_channel_iterator!(IcmpPacket, IcmpTransportChannelIterator, icmp_packet_iter);

transport_channel_iterator!(TcpPacket, TcpTransportChannelIterator, tcp_packet_iter);

<|endoftext|>


<commit_before>extern crate time;

use std::fmt;
use std::io;
use std::io::{Error, Write};
use std::cmp::{Ord, Ordering};

pub enum LogLevel {
    ERROR,
    WARN,
    INFO,
    VERBOSE
}

// ERROR > WARN > INFO > VERBOSE
impl Ord for LogLevel {
    fn cmp(&self, other: &Self) -> Ordering {
        match *self {
            LogLevel::ERROR => match *other {
                        LogLevel::ERROR => Ordering::Equal,
                        _ => Ordering::Greater,
                    },
            LogLevel::WARN => match *other {
                        LogLevel::ERROR => Ordering::Less,
                        LogLevel::WARN => Ordering::Equal,
                        _ => Ordering::Greater,
                    },
            LogLevel::INFO => match *other {
                        LogLevel::VERBOSE => Ordering::Greater,
                        LogLevel::INFO => Ordering::Equal,
                        _ => Ordering::Less,
                    },
            LogLevel::VERBOSE => match *other {
                        LogLevel::VERBOSE => Ordering::Equal,
                        _ => Ordering::Less,
                    },
        }
    }
}

impl Eq for LogLevel {}

impl PartialEq for LogLevel {

    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }

    fn ne(&self, other: &Self) -> bool {
        !self.eq(other)
    }

}

impl PartialOrd for LogLevel {

    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }

    fn lt(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Less
    }

    fn le(&self, other: &Self) -> bool {
        self.cmp(other) != Ordering::Greater
    }

    fn gt(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Greater
    }

    fn ge(&self, other: &Self) -> bool {
        self.cmp(other) != Ordering::Less
    }
}

/// Set global log level here
static LOG_LEVEL: LogLevel = LogLevel::INFO;

/// the log function returning a result
///
/// # Examples
/// '''
/// log(LogLevel::ERROR, format_args!("This is bad!")).unwrap();
/// '''
pub fn log(level: LogLevel, args: fmt::Arguments) -> Result<(), Error> {

    if level < LOG_LEVEL {
        return Ok(());
    }

    let stderr = io::stderr();
    let stdout = io::stdout();

    let mut stderr_lock = stderr.lock();
    let mut stdout_lock = stdout.lock();

    let cur_time = time::now();

    match level {
        LogLevel::ERROR => writeln!(stderr_lock, "[ERROR] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
        LogLevel::WARN => writeln!(stdout_lock, "[WARN] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
        LogLevel::INFO => writeln!(stdout_lock, "[INFO] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
        LogLevel::VERBOSE => writeln!(stdout_lock, "[VERBOSE] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
    }

}

/// a log function ignoring any failure
/// "log something, if possible"
///
/// # Examples
/// '''
/// try_log(LogLevel::ERROR, format_args!("This is bad!"));
/// '''
pub fn try_log(level: LogLevel, args: fmt::Arguments) {
    let _ = log(level, args);
}

// Helper Macro

/// Error Macro
///
/// # Examples
/// '''
/// log_error!("This is bad!");
/// '''
macro_rules! log_error {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::ERROR, format_args!($($arg)*)));
}

/// Warning Macro
///
/// # Examples
/// '''
/// log_warn!("This is maybe bad?!");
/// '''
macro_rules! log_warn {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::WARN, format_args!($($arg)*)));
}

/// Info Macro
///
/// # Examples
/// '''
/// log_info!("This is just to let you know!");
/// '''
macro_rules! log_info {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::INFO, format_args!($($arg)*)));
}

/// Verbose Macro
///
/// # Examples
/// '''
/// log_verbose!("You must really like spam!");
/// '''
macro_rules! log_verbose {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::VERBOSE, format_args!($($arg)*)));
}

#[test]
fn test() {
    log_info!("Test!");
}

//To-Do add a way to check the result automatically
#[test]
fn multi_thread_test() {
    use std::thread;

    let mut joins = Vec::new();
    for _ in 0..10 {
        joins.push(thread::spawn(move || {
            log_info!("No corruption should happen in this ultra-extra-super-long string output even if run in multiple threads.");
        }));
    }
    for handle in joins {
        let _ = handle.join();
    }
}
<commit_msg>Put Logger-Tests into test module<commit_after>extern crate time;

use std::fmt;
use std::io;
use std::io::{Error, Write};
use std::cmp::{Ord, Ordering};

pub enum LogLevel {
    ERROR,
    WARN,
    INFO,
    VERBOSE
}

// ERROR > WARN > INFO > VERBOSE
impl Ord for LogLevel {
    fn cmp(&self, other: &Self) -> Ordering {
        match *self {
            LogLevel::ERROR => match *other {
                        LogLevel::ERROR => Ordering::Equal,
                        _ => Ordering::Greater,
                    },
            LogLevel::WARN => match *other {
                        LogLevel::ERROR => Ordering::Less,
                        LogLevel::WARN => Ordering::Equal,
                        _ => Ordering::Greater,
                    },
            LogLevel::INFO => match *other {
                        LogLevel::VERBOSE => Ordering::Greater,
                        LogLevel::INFO => Ordering::Equal,
                        _ => Ordering::Less,
                    },
            LogLevel::VERBOSE => match *other {
                        LogLevel::VERBOSE => Ordering::Equal,
                        _ => Ordering::Less,
                    },
        }
    }
}

impl Eq for LogLevel {}

impl PartialEq for LogLevel {

    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }

    fn ne(&self, other: &Self) -> bool {
        !self.eq(other)
    }

}

impl PartialOrd for LogLevel {

    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }

    fn lt(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Less
    }

    fn le(&self, other: &Self) -> bool {
        self.cmp(other) != Ordering::Greater
    }

    fn gt(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Greater
    }

    fn ge(&self, other: &Self) -> bool {
        self.cmp(other) != Ordering::Less
    }
}

/// Set global log level here
static LOG_LEVEL: LogLevel = LogLevel::INFO;

/// the log function returning a result
///
/// # Examples
/// '''
/// log(LogLevel::ERROR, format_args!("This is bad!")).unwrap();
/// '''
pub fn log(level: LogLevel, args: fmt::Arguments) -> Result<(), Error> {

    if level < LOG_LEVEL {
        return Ok(());
    }

    let stderr = io::stderr();
    let stdout = io::stdout();

    let mut stderr_lock = stderr.lock();
    let mut stdout_lock = stdout.lock();

    let cur_time = time::now();

    match level {
        LogLevel::ERROR => writeln!(stderr_lock, "[ERROR] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
        LogLevel::WARN => writeln!(stdout_lock, "[WARN] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
        LogLevel::INFO => writeln!(stdout_lock, "[INFO] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
        LogLevel::VERBOSE => writeln!(stdout_lock, "[VERBOSE] {}:{}:{} - {}", cur_time.tm_hour, cur_time.tm_min, cur_time.tm_sec, args),
    }

}

/// a log function ignoring any failure
/// "log something, if possible"
///
/// # Examples
/// '''
/// try_log(LogLevel::ERROR, format_args!("This is bad!"));
/// '''
pub fn try_log(level: LogLevel, args: fmt::Arguments) {
    let _ = log(level, args);
}

// Helper Macro

/// Error Macro
///
/// # Examples
/// '''
/// log_error!("This is bad!");
/// '''
macro_rules! log_error {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::ERROR, format_args!($($arg)*)));
}

/// Warning Macro
///
/// # Examples
/// '''
/// log_warn!("This is maybe bad?!");
/// '''
macro_rules! log_warn {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::WARN, format_args!($($arg)*)));
}

/// Info Macro
///
/// # Examples
/// '''
/// log_info!("This is just to let you know!");
/// '''
macro_rules! log_info {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::INFO, format_args!($($arg)*)));
}

/// Verbose Macro
///
/// # Examples
/// '''
/// log_verbose!("You must really like spam!");
/// '''
macro_rules! log_verbose {
    ($($arg:tt)*) => (::utils::log::try_log(::utils::log::LogLevel::VERBOSE, format_args!($($arg)*)));
}

#[cfg(test)]
mod test {
    use std::thread;
    // use super::*;

    #[test]
    fn test() {
        log_info!("Test!");
    }

    //To-Do add a way to check the result automatically
    #[test]
    fn multi_thread_test() {
        let mut joins = Vec::new();
        for _ in 0..10 {
            joins.push(thread::spawn(move || {
                log_info!("No corruption should happen in this ultra-extra-super-long string output even if run in multiple threads.");
            }));
        }
        for handle in joins {
            let _ = handle.join();
        }
    }
}
<|endoftext|>


<commit_before>//! This modules provides an implementation of the Lempel–Ziv–Welch Compression Algorithm

// FIXME remove as soon as gif encoder is done
#![allow(dead_code)]

use std::old_io;

use utils::bitstream::{BitReader, BitWriter};

const MAX_CODESIZE: u8 = 12;
const MAX_ENTRIES: usize = 1 << MAX_CODESIZE as usize;

/// Alias for a LZW code point
type Code = u16;

/// Decoding dictionary
/// It is not generic due to current limitations of Rust
/// Inspired by http://www.cplusplus.com/articles/iL18T05o/
struct DecodingDict {
    min_size: u8,
    table: Vec<(Option<Code>, u8)>,
    buffer: Vec<u8>,
}

impl DecodingDict {
    /// Creates a new dict
    fn new(min_size: u8) -> DecodingDict {
        DecodingDict {
            min_size: min_size,
            table: Vec::with_capacity(512),
            buffer: Vec::with_capacity((1 << MAX_CODESIZE as usize) - 1)
        }
    }

    /// Resets the dictionary
    fn reset(&mut self) {
        self.table.clear();
        for i in (0..(1u16 << self.min_size as usize)) {
            self.table.push((None, i as u8));
        }
    }

    /// Inserts a value into the dict
    #[inline(always)]
    fn push(&mut self, key: Option<Code>, value: u8) {
        self.table.push((key, value))
    }

    /// Reconstructs the data for the corresponding code
    fn reconstruct(&mut self, code: Option<Code>) -> old_io::IoResult<&[u8]> {
        self.buffer.clear();
        let mut code = code;
        let mut cha;
        // Check the first access more thoroughly since a bad code
        // could occur if the data is malformed
        if let Some(k) = code {
            match self.table.get(k as usize) {
                Some(&(code_, cha_)) => {
                    code = code_;
                    cha = cha_;
                }
                None => return Err(old_io::IoError {
                    kind: old_io::InvalidInput,
                    desc: "invalid code occured",
                    detail: Some(format!("{} < {} expected", k, self.table.len()))
                })
            }
            self.buffer.push(cha);
        }
        while let Some(k) = code {
            //(code, cha) = self.table[k as usize];
            // Node this could possibly replaced by unsafe access because this
            // struct had been contructed by this algorithm correctly
            let entry = self.table[k as usize]; code = entry.0; cha = entry.1;
            self.buffer.push(cha);
        }
        self.buffer.reverse();
        Ok(&self.buffer[])
    }

    /// Returns the buffer constructed by the last reconstruction
    #[inline(always)]
    fn buffer(&self) -> &[u8] {
        &self.buffer[]
    }

    /// Number of entries in the dictionary
    #[inline(always)]
    fn next_code(&self) -> u16 {
        self.table.len() as u16
    }
}

macro_rules! define_decoder_function {
    {$(
        $name:ident, $offset:expr, #[$doc:meta];
    )*} => {

$( // START function definition

#[$doc]
pub fn $name<R, W>(mut r: R, w: &mut W, min_code_size: u8) -> old_io::IoResult<()>
where R: BitReader, W: Writer {
    let mut prev = None;
    let clear_code = 1 << min_code_size as usize;
    let end_code = clear_code + 1;
    let mut table = DecodingDict::new(min_code_size);
    let mut code_size = min_code_size + 1;
    loop {
        let code = try!(r.read_bits(code_size));
        println!("{:02} {:?}", code_size, code);
        if code == clear_code {
            table.reset();
            table.push(None, 0); // clear code
            table.push(None, 0); // end code
            code_size = min_code_size + 1;
            prev = None;
        } else if code == end_code {
            return Ok(())
        } else {
            let next_code = table.next_code();
            if prev.is_none() {
                try!(w.write_u8(code as u8));
            } else {
                let data = if code == next_code {
                    let cha = try!(table.reconstruct(prev))[0];
                    table.push(prev, cha);
                    try!(table.reconstruct(Some(code)))
                } else if code < next_code {
                    let cha = try!(table.reconstruct(Some(code)))[0];
                    table.push(prev, cha);
                    table.buffer()
                } else {
                    return Err(old_io::IoError {
                        kind: old_io::InvalidInput,
                        desc: "Invalid code",
                        detail: Some(format!("expected {} <= {}",
                                     code,
                                     next_code)
                                )
                    })
                };
                try!(w.write(data));
            }
            if next_code == (1 << code_size as usize) - 1 - $offset
               && code_size < MAX_CODESIZE {
                code_size += 1;
            }
            prev = Some(code);
        }
    }
}

)* // END function definition

    }
}

define_decoder_function!{
    decode, 0, #[doc = "Decodes a lzw compressed stream."];
    decode_early_change, 1, #[doc = "Decodes a lzw compressed stream using an “early change” algorithm."];
}

struct Node {
    prefix: Option<Code>,
    c: u8,
    left: Option<Code>,
    right: Option<Code>,
}

impl Node {
    #[inline(always)]
    fn new(c: u8) -> Node {
        Node {
            prefix: None,
            c: c,
            left: None,
            right: None
        }
    }
}

struct EncodingDict {
    table: Vec<Node>,
    min_size: u8,

}

/// Encoding dictionary based on a binary tree
impl EncodingDict {
    fn new(min_size: u8) -> EncodingDict {
        let mut this = EncodingDict {
            table: Vec::with_capacity(MAX_ENTRIES),
            min_size: min_size
        };
        this.reset();
        this
    }

    fn reset(&mut self) {
        self.table.clear();
        for i in range(0, (1u16 << self.min_size as usize)) {
            self.push_node(Node::new(i as u8));
        }
    }

    #[inline(always)]
    fn push_node(&mut self, node: Node) {
        self.table.push(node)
    }

    #[inline(always)]
    fn clear_code(&self) -> Code {
        1u16 << self.min_size
    }

    #[inline(always)]
    fn end_code(&self) -> Code {
        self.clear_code() + 1
    }

    // Searches for a new prefix
    fn search_and_insert(&mut self, i: Option<Code>, c: u8) -> Option<Code> {
        if let Some(i) = i.map(|v| v as usize) {
            let table_size = self.table.len() as Code;
            if let Some(mut j) = self.table[i].prefix {
                loop {
                    let entry = &mut self.table[j as usize];
                    if c < entry.c {
                        if let Some(k) = entry.left {
                            j = k
                        } else {
                            entry.left = Some(table_size);
                            break
                        }
                    } else if c > entry.c {
                        if let Some(k) = entry.right {
                            j = k
                        } else {
                            entry.right = Some(table_size);
                            break
                        }
                    } else {
                        return Some(j)
                    }
                }
            } else {
                self.table[i].prefix = Some(table_size);
            }
            self.table.push(Node::new(c));
            None
        } else {
            Some(self.search_initials(c as Code))
        }
    }

    fn next_code(&self) -> usize {
        self.table.len()
    }

    fn search_initials(&self, i: Code) -> Code {
        self.table[i as usize].c as Code
    }
}

pub fn encode<R, W>(mut r: R, w: &mut W, min_code_size: u8) -> old_io::IoResult<()>
where R: Reader, W: BitWriter {
    let mut dict = EncodingDict::new(min_code_size);
    dict.push_node(Node::new(0)); // clear code
    dict.push_node(Node::new(0)); // end code
    let mut code_size = min_code_size + 1;
    let mut i = None;
    // gif spec: first clear code
    try!(w.write_bits(dict.clear_code(), code_size));
    while let Ok(c) = r.read_byte() {
        let prev = i;
        i = dict.search_and_insert(prev, c);
        if i.is_none() {
            if let Some(code) = prev {
                try!(w.write_bits(code, code_size));
            }
            i = Some(dict.search_initials(c as Code))
        }
        // There is a hit: do not write out code but continue
        let next_code = dict.next_code();
        if next_code > (1 << code_size as usize)
           && code_size < MAX_CODESIZE {
            code_size += 1;
        }
        if next_code > MAX_ENTRIES {
            dict.reset();
            dict.push_node(Node::new(0)); // clear code
            dict.push_node(Node::new(0)); // end code
            try!(w.write_bits(dict.clear_code(), code_size));
            code_size = min_code_size + 1;
        }

    }
    if let Some(code) = i {
        try!(w.write_bits(code, code_size));
    }
    try!(w.write_bits(dict.end_code(), code_size));
    try!(w.flush());
    Ok(())
}
<commit_msg>Attribution for LZW algorithm.<commit_after>//! This modules provides an implementation of the Lempel–Ziv–Welch Compression Algorithm

// FIXME remove as soon as gif encoder is done
#![allow(dead_code)]

// Note: This implementation borrows heavily from the work of Julius Pettersson
// See http://www.cplusplus.com/articles/iL18T05o/ for his extensive explanations
// and a C++ implementatation

use std::old_io;

use utils::bitstream::{BitReader, BitWriter};

const MAX_CODESIZE: u8 = 12;
const MAX_ENTRIES: usize = 1 << MAX_CODESIZE as usize;

/// Alias for a LZW code point
type Code = u16;

/// Decoding dictionary
/// It is not generic due to current limitations of Rust
/// Inspired by http://www.cplusplus.com/articles/iL18T05o/
struct DecodingDict {
    min_size: u8,
    table: Vec<(Option<Code>, u8)>,
    buffer: Vec<u8>,
}

impl DecodingDict {
    /// Creates a new dict
    fn new(min_size: u8) -> DecodingDict {
        DecodingDict {
            min_size: min_size,
            table: Vec::with_capacity(512),
            buffer: Vec::with_capacity((1 << MAX_CODESIZE as usize) - 1)
        }
    }

    /// Resets the dictionary
    fn reset(&mut self) {
        self.table.clear();
        for i in (0..(1u16 << self.min_size as usize)) {
            self.table.push((None, i as u8));
        }
    }

    /// Inserts a value into the dict
    #[inline(always)]
    fn push(&mut self, key: Option<Code>, value: u8) {
        self.table.push((key, value))
    }

    /// Reconstructs the data for the corresponding code
    fn reconstruct(&mut self, code: Option<Code>) -> old_io::IoResult<&[u8]> {
        self.buffer.clear();
        let mut code = code;
        let mut cha;
        // Check the first access more thoroughly since a bad code
        // could occur if the data is malformed
        if let Some(k) = code {
            match self.table.get(k as usize) {
                Some(&(code_, cha_)) => {
                    code = code_;
                    cha = cha_;
                }
                None => return Err(old_io::IoError {
                    kind: old_io::InvalidInput,
                    desc: "invalid code occured",
                    detail: Some(format!("{} < {} expected", k, self.table.len()))
                })
            }
            self.buffer.push(cha);
        }
        while let Some(k) = code {
            //(code, cha) = self.table[k as usize];
            // Node this could possibly replaced by unsafe access because this
            // struct had been contructed by this algorithm correctly
            let entry = self.table[k as usize]; code = entry.0; cha = entry.1;
            self.buffer.push(cha);
        }
        self.buffer.reverse();
        Ok(&self.buffer[])
    }

    /// Returns the buffer constructed by the last reconstruction
    #[inline(always)]
    fn buffer(&self) -> &[u8] {
        &self.buffer[]
    }

    /// Number of entries in the dictionary
    #[inline(always)]
    fn next_code(&self) -> u16 {
        self.table.len() as u16
    }
}

macro_rules! define_decoder_function {
    {$(
        $name:ident, $offset:expr, #[$doc:meta];
    )*} => {

$( // START function definition

#[$doc]
pub fn $name<R, W>(mut r: R, w: &mut W, min_code_size: u8) -> old_io::IoResult<()>
where R: BitReader, W: Writer {
    let mut prev = None;
    let clear_code = 1 << min_code_size as usize;
    let end_code = clear_code + 1;
    let mut table = DecodingDict::new(min_code_size);
    let mut code_size = min_code_size + 1;
    loop {
        let code = try!(r.read_bits(code_size));
        println!("{:02} {:?}", code_size, code);
        if code == clear_code {
            table.reset();
            table.push(None, 0); // clear code
            table.push(None, 0); // end code
            code_size = min_code_size + 1;
            prev = None;
        } else if code == end_code {
            return Ok(())
        } else {
            let next_code = table.next_code();
            if prev.is_none() {
                try!(w.write_u8(code as u8));
            } else {
                let data = if code == next_code {
                    let cha = try!(table.reconstruct(prev))[0];
                    table.push(prev, cha);
                    try!(table.reconstruct(Some(code)))
                } else if code < next_code {
                    let cha = try!(table.reconstruct(Some(code)))[0];
                    table.push(prev, cha);
                    table.buffer()
                } else {
                    return Err(old_io::IoError {
                        kind: old_io::InvalidInput,
                        desc: "Invalid code",
                        detail: Some(format!("expected {} <= {}",
                                     code,
                                     next_code)
                                )
                    })
                };
                try!(w.write(data));
            }
            if next_code == (1 << code_size as usize) - 1 - $offset
               && code_size < MAX_CODESIZE {
                code_size += 1;
            }
            prev = Some(code);
        }
    }
}

)* // END function definition

    }
}

define_decoder_function!{
    decode, 0, #[doc = "Decodes a lzw compressed stream."];
    decode_early_change, 1, #[doc = "Decodes a lzw compressed stream using an “early change” algorithm."];
}

struct Node {
    prefix: Option<Code>,
    c: u8,
    left: Option<Code>,
    right: Option<Code>,
}

impl Node {
    #[inline(always)]
    fn new(c: u8) -> Node {
        Node {
            prefix: None,
            c: c,
            left: None,
            right: None
        }
    }
}

struct EncodingDict {
    table: Vec<Node>,
    min_size: u8,

}

/// Encoding dictionary based on a binary tree
impl EncodingDict {
    fn new(min_size: u8) -> EncodingDict {
        let mut this = EncodingDict {
            table: Vec::with_capacity(MAX_ENTRIES),
            min_size: min_size
        };
        this.reset();
        this
    }

    fn reset(&mut self) {
        self.table.clear();
        for i in range(0, (1u16 << self.min_size as usize)) {
            self.push_node(Node::new(i as u8));
        }
    }

    #[inline(always)]
    fn push_node(&mut self, node: Node) {
        self.table.push(node)
    }

    #[inline(always)]
    fn clear_code(&self) -> Code {
        1u16 << self.min_size
    }

    #[inline(always)]
    fn end_code(&self) -> Code {
        self.clear_code() + 1
    }

    // Searches for a new prefix
    fn search_and_insert(&mut self, i: Option<Code>, c: u8) -> Option<Code> {
        if let Some(i) = i.map(|v| v as usize) {
            let table_size = self.table.len() as Code;
            if let Some(mut j) = self.table[i].prefix {
                loop {
                    let entry = &mut self.table[j as usize];
                    if c < entry.c {
                        if let Some(k) = entry.left {
                            j = k
                        } else {
                            entry.left = Some(table_size);
                            break
                        }
                    } else if c > entry.c {
                        if let Some(k) = entry.right {
                            j = k
                        } else {
                            entry.right = Some(table_size);
                            break
                        }
                    } else {
                        return Some(j)
                    }
                }
            } else {
                self.table[i].prefix = Some(table_size);
            }
            self.table.push(Node::new(c));
            None
        } else {
            Some(self.search_initials(c as Code))
        }
    }

    fn next_code(&self) -> usize {
        self.table.len()
    }

    fn search_initials(&self, i: Code) -> Code {
        self.table[i as usize].c as Code
    }
}

pub fn encode<R, W>(mut r: R, w: &mut W, min_code_size: u8) -> old_io::IoResult<()>
where R: Reader, W: BitWriter {
    let mut dict = EncodingDict::new(min_code_size);
    dict.push_node(Node::new(0)); // clear code
    dict.push_node(Node::new(0)); // end code
    let mut code_size = min_code_size + 1;
    let mut i = None;
    // gif spec: first clear code
    try!(w.write_bits(dict.clear_code(), code_size));
    while let Ok(c) = r.read_byte() {
        let prev = i;
        i = dict.search_and_insert(prev, c);
        if i.is_none() {
            if let Some(code) = prev {
                try!(w.write_bits(code, code_size));
            }
            i = Some(dict.search_initials(c as Code))
        }
        // There is a hit: do not write out code but continue
        let next_code = dict.next_code();
        if next_code > (1 << code_size as usize)
           && code_size < MAX_CODESIZE {
            code_size += 1;
        }
        if next_code > MAX_ENTRIES {
            dict.reset();
            dict.push_node(Node::new(0)); // clear code
            dict.push_node(Node::new(0)); // end code
            try!(w.write_bits(dict.clear_code(), code_size));
            code_size = min_code_size + 1;
        }

    }
    if let Some(code) = i {
        try!(w.write_bits(code, code_size));
    }
    try!(w.write_bits(dict.end_code(), code_size));
    try!(w.flush());
    Ok(())
}
<|endoftext|>


<commit_before>use std::io;
use {Rng, SeedableRng, Rand};

#[derive(Debug, Clone, Copy)]
pub struct Xoroshiro128Rng {
  state: [u64; 2]
}

impl Xoroshiro128Rng {
  /// Creates a new `Xoroshiro128Rng` instance which is randomly seeded.
  ///
  /// # Errors
  ///
  /// When created this way the initial values of this `Rng` are seeded using an `OsRng`. If this
  /// fails then this method will pass through the error from `OsRng`.
  ///
  /// # Examples
  ///
  /// ```rust,no_run
  /// use xoroshiro128::{Rng, Xoroshiro128Rng};
  ///
  /// let rng = Xoroshiro128Rng::new();
  /// let x: u32 = rng.unwrap().gen();
  /// println!("{}", x);
  /// ```
  pub fn new() -> io::Result<Self> {
    ::rand::OsRng::new().map(|mut x: ::rand::OsRng| x.gen::<Self>())
  }

  /// Creates a new `Xoroshiro128Rng` instance which is not seeded.
  ///
  /// The initial values of this Rng are constants, so all generators created by this function
  /// will yield the same stream of random numbers. It is highly recommended that this is created
  /// through `SeedableRng` instead of this function.
  pub fn new_unseeded() -> Self {
    Xoroshiro128Rng::from_seed([
      0x193a6754a8a7d469,
      0x97830e05113ba7bb
    ])
  }
}

#[inline(always)]
fn rotate_left(x: u64, k: i32) -> u64 {
  (x << k) | (x >> (64 - k))
}

impl Rng for Xoroshiro128Rng {
  fn next_u32(&mut self) -> u32 {
    self.next_u64() as u32
  }

  fn next_u64(&mut self) -> u64 {
    let s0: u64 = self.state[0];
    let mut s1: u64 = self.state[1];
    let result: u64 = s0.wrapping_add(s1);

    s1 ^= s0;
    self.state[0] = rotate_left(s0, 55) ^ s1 ^ (s1 << 14); // a, b
    self.state[1] = rotate_left(s1, 36); // c

    result
  }
}

/// Seed a `Xoroshiro128Rng` with a given seed.
///
/// # Panics
///
/// `Xoroshiro128Rng` is undefined for the seed `[0, 0]` and will panic if this seed is provided.
impl SeedableRng<[u64; 2]> for Xoroshiro128Rng {
  fn reseed(&mut self, seed: [u64; 2]) {
    assert!(seed != [0, 0], "Invalid seed: seed must not be 0.");
    self.state = seed
  }

  fn from_seed(seed: [u64; 2]) -> Self {
    assert!(seed != [0, 0], "Invalid seed: seed must not be 0.");
    Xoroshiro128Rng { state: seed }
  }
}

#[inline(always)]
fn splitmix_seed(seed: u64) -> [u64; 2] {
  ::SplitMix64Rng::from_seed(seed).gen()
}

impl SeedableRng<u64> for Xoroshiro128Rng {
  fn reseed(&mut self, seed: u64) {
    self.reseed(splitmix_seed(seed))
  }

  fn from_seed(seed: u64) -> Self {
    Xoroshiro128Rng::from_seed(splitmix_seed(seed))
  }
}

/// Seed a `Xoroshiro128Rng` based on an `OsRng`.
///
/// # Panics
///
/// If `OsRng` is unavailable then this will panic. A safer option is `Xoroshiro128Rng::new()`
impl SeedableRng<()> for Xoroshiro128Rng {
  fn reseed(&mut self, _: ()) {
    self.reseed(::rand::OsRng::new().unwrap().gen::<[u64; 2]>())
  }

  fn from_seed(_: ()) -> Self {
    ::rand::OsRng::new().unwrap().gen::<Self>()
  }
}

impl Rand for Xoroshiro128Rng {
  fn rand<R: Rng>(rng: &mut R) -> Xoroshiro128Rng {
    let mut seed: [u64; 2] = rng.gen();
    while seed == [0, 0] {
      seed = rng.gen();
    }

    Xoroshiro128Rng::from_seed(seed)
  }
}
<commit_msg>Use u64::rotate_left in Xoroshiro128Rng<commit_after>use std::io;
use {Rng, SeedableRng, Rand};

#[derive(Debug, Clone, Copy)]
pub struct Xoroshiro128Rng {
  state: [u64; 2]
}

impl Xoroshiro128Rng {
  /// Creates a new `Xoroshiro128Rng` instance which is randomly seeded.
  ///
  /// # Errors
  ///
  /// When created this way the initial values of this `Rng` are seeded using an `OsRng`. If this
  /// fails then this method will pass through the error from `OsRng`.
  ///
  /// # Examples
  ///
  /// ```rust,no_run
  /// use xoroshiro128::{Rng, Xoroshiro128Rng};
  ///
  /// let rng = Xoroshiro128Rng::new();
  /// let x: u32 = rng.unwrap().gen();
  /// println!("{}", x);
  /// ```
  pub fn new() -> io::Result<Self> {
    ::rand::OsRng::new().map(|mut x: ::rand::OsRng| x.gen::<Self>())
  }

  /// Creates a new `Xoroshiro128Rng` instance which is not seeded.
  ///
  /// The initial values of this Rng are constants, so all generators created by this function
  /// will yield the same stream of random numbers. It is highly recommended that this is created
  /// through `SeedableRng` instead of this function.
  pub fn new_unseeded() -> Self {
    Xoroshiro128Rng::from_seed([
      0x193a6754a8a7d469,
      0x97830e05113ba7bb
    ])
  }
}

impl Rng for Xoroshiro128Rng {
  fn next_u32(&mut self) -> u32 {
    self.next_u64() as u32
  }

  fn next_u64(&mut self) -> u64 {
    let s0: u64 = self.state[0];
    let mut s1: u64 = self.state[1];
    let result: u64 = s0.wrapping_add(s1);

    s1 ^= s0;
    self.state[0] = s0.rotate_left(55) ^ s1 ^ (s1 << 14); // a, b
    self.state[1] = s1.rotate_left(36); // c

    result
  }
}

/// Seed a `Xoroshiro128Rng` with a given seed.
///
/// # Panics
///
/// `Xoroshiro128Rng` is undefined for the seed `[0, 0]` and will panic if this seed is provided.
impl SeedableRng<[u64; 2]> for Xoroshiro128Rng {
  fn reseed(&mut self, seed: [u64; 2]) {
    assert!(seed != [0, 0], "Invalid seed: seed must not be 0.");
    self.state = seed
  }

  fn from_seed(seed: [u64; 2]) -> Self {
    assert!(seed != [0, 0], "Invalid seed: seed must not be 0.");
    Xoroshiro128Rng { state: seed }
  }
}

#[inline(always)]
fn splitmix_seed(seed: u64) -> [u64; 2] {
  ::SplitMix64Rng::from_seed(seed).gen()
}

impl SeedableRng<u64> for Xoroshiro128Rng {
  fn reseed(&mut self, seed: u64) {
    self.reseed(splitmix_seed(seed))
  }

  fn from_seed(seed: u64) -> Self {
    Xoroshiro128Rng::from_seed(splitmix_seed(seed))
  }
}

/// Seed a `Xoroshiro128Rng` based on an `OsRng`.
///
/// # Panics
///
/// If `OsRng` is unavailable then this will panic. A safer option is `Xoroshiro128Rng::new()`
impl SeedableRng<()> for Xoroshiro128Rng {
  fn reseed(&mut self, _: ()) {
    self.reseed(::rand::OsRng::new().unwrap().gen::<[u64; 2]>())
  }

  fn from_seed(_: ()) -> Self {
    ::rand::OsRng::new().unwrap().gen::<Self>()
  }
}

impl Rand for Xoroshiro128Rng {
  fn rand<R: Rng>(rng: &mut R) -> Xoroshiro128Rng {
    let mut seed: [u64; 2] = rng.gen();
    while seed == [0, 0] {
      seed = rng.gen();
    }

    Xoroshiro128Rng::from_seed(seed)
  }
}
<|endoftext|>