<commit_before>// pest. Elegant, efficient grammars
// Copyright (C) 2016  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/// A `macro` useful for implementing the `Parser` `trait` as a recursive descent parser. It only
/// accepts `grammar!` and `process!` calls that get implemented on `self`.
///
/// # Rule
///
/// It also implements an `enum` called `Rule` that has a value for all
/// [non-silent](macro.grammar!#silent-rules-_) rules, but also for
/// [`any` and `eoi`](macro.grammar!). These `Rule`s are used within `Token`s to specify the type
/// of rule that matched.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate pest;
/// # use pest::prelude::*;
/// # fn main() {
/// impl_rdp! {
///     grammar! {
///         expression = _{ paren ~ expression? }
///         paren      =  { ["("] ~ expression? ~ [")"] }
///     }
/// }
///
/// let mut parser = Rdp::new(StringInput::new("(())((())())()"));
///
/// assert!(parser.expression());
/// assert!(parser.end());
/// # }
/// ```
#[macro_export]
macro_rules! impl_rdp {
    // implement rules
    ( @rules $( $name:ident )* ) => {
        #[allow(dead_code, non_camel_case_types)]
        #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
        pub enum Rule {
            any,
            eoi,
            $( $name ),*
        }
    };

    // filter out silent rules
    ( @filter [  ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@rules $( $rules )*);
    };
    ( @filter [ $name:ident = { { $( $_primary:tt )* } $( $ts:tt )* } $( $tail:tt )* ]
      [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* $( $ts )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $name:ident = @{ { $( $_primary:tt )* } $( $ts:tt )* } $( $tail:tt )* ]
      [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* $( $ts )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $name:ident = _{ { $( $_primary:tt )* } $( $ts:tt )* } $( $tail:tt )* ]
      [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* $( $ts )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $name:ident = { $( $_ts:tt )* } $( $tail:tt )* ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $name:ident = @{ $( $_ts:tt )* } $( $tail:tt )* ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $name:ident = _{ $( $_ts:tt )* } $( $tail:tt )* ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* ] [ $( $rules )* ]);
    };

    // implement empty whitespace rule
    ( @ws ) => {
        #[allow(dead_code)]
        #[inline]
        pub fn whitespace(&mut self) -> bool {
            false
        }
    };
    ( @ws whitespace = $( $_ts:tt )* ) => ();
    ( @ws $_name:ident = { $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@ws $( $tail )*);
    };
    ( @ws $_name:ident = @{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@ws $( $tail )*);
    };
    ( @ws $_name:ident = _{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@ws $( $tail )*);
    };

    // implement empty comment rule
    ( @com ) => {
        #[allow(dead_code)]
        #[inline]
        pub fn comment(&mut self) -> bool {
            false
        }
    };
    ( @com comment = $( $_ts:tt )* ) => ();
    ( @com $_name:ident = { $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@com $( $tail )*);
    };
    ( @com $_name:ident = @{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@com $( $tail )*);
    };
    ( @com $_name:ident = _{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@com $( $tail )*);
    };

    ( grammar! { $( $ts:tt )* } $( $mac:ident! { $( $rest:tt )* } )* ) => {
        use std::cell::Cell;
        use std::cmp;

        pub struct Rdp<T> {
            input:       T,
            queue:       Vec<Token<Rule>>,
            queue_index: Cell<usize>,
            failures:    Vec<Rule>,
            fail_pos:    usize,
            atomic:      bool,
            comment:     bool,
            eoi_matched: bool
        }

        impl_rdp!(@filter [ $( $ts )* ] []);

        impl<'input, T: Input<'input>> Rdp<T> {
            pub fn new(input: T) -> Rdp<T> {
                Rdp {
                    input:       input,
                    queue:       vec![],
                    queue_index: Cell::new(0),
                    failures:    vec![],
                    fail_pos:    0,
                    atomic:      false,
                    comment:     false,
                    eoi_matched: false
                }
            }

            impl_rdp!(@ws $( $ts )*);
            impl_rdp!(@com $( $ts )*);

            #[allow(dead_code)]
            #[inline]
            pub fn any(&mut self) -> bool {
                if self.end() {
                    let pos = self.input.pos();

                    self.track(Rule::any, pos);

                    false
                } else {
                    let next = self.input.pos() + 1;
                    self.input.set_pos(next);

                    true
                }
            }

            #[allow(dead_code)]
            #[inline]
            pub fn eoi(&mut self) -> bool {
                let result = self.end();

                if !result {
                    let pos = self.input.pos();

                    self.track(Rule::eoi, pos);
                } else {
                    self.eoi_matched = true;
                }

                result
            }

            grammar! {
                $( $ts )*
            }

            $(
                $mac! {
                    $( $rest )*
                }
            )*
        }

        impl<'input, T: Input<'input>> Parser<'input, T> for Rdp<T> {
            type Rule = Rule;
            type Token = Token<Rule>;

            #[inline]
            fn input(&self) -> &T {
                &self.input
            }

            #[inline]
            fn input_mut(&mut self) -> &mut T {
                &mut self.input
            }

            #[inline]
            fn try<F>(&mut self, revert: bool, rule: F) -> bool
                where F: FnOnce(&mut Self) -> bool {

                let pos = self.input.pos();
                let len = self.queue.len();

                let result = rule(self);

                if revert || !result {
                    self.input.set_pos(pos);
                }

                if !result {
                    self.queue.truncate(len);
                }

                result
            }

            fn prec_climb<F, G>(&mut self, pos: usize, left: usize, min_prec: u8,
                                last_op: Option<(Option<Rule>, u8, bool)>, primary: &mut F,
                                climb: &mut G) -> (Option<(Option<Rule>, u8, bool)>, Option<usize>)
                where F: FnMut(&mut Self) -> bool,
                      G: FnMut(&mut Self) -> Option<(Option<Rule>, u8, bool)> {

                let mut op = if last_op.is_some() {
                    last_op
                } else {
                    climb(self)
                };
                let mut last_right = None;

                while let Some((rule, prec, _)) = op {
                    if prec >= min_prec {
                        let mut new_pos = self.input.pos();
                        let mut right = self.input.pos();
                        let queue_pos = self.queue.len();

                        primary(self);

                        if let Some(token) = self.queue.get(queue_pos) {
                            new_pos = token.start;
                            right   = token.end;
                        }

                        op = climb(self);

                        while let Some((_, new_prec, right_assoc)) = op {
                            if new_prec > prec || right_assoc && new_prec == prec {
                                let (new_op, new_lr) = self.prec_climb(queue_pos, new_pos,
                                                                       new_prec, op, primary,
                                                                       climb);

                                op = new_op;
                                last_right = new_lr;
                            } else {
                                break
                            }
                        }

                        if let Some(pos) = last_right {
                            right = cmp::max(pos, right);
                        } else {
                            last_right = Some(right);
                        }

                        if let Some(rule) = rule {
                            let token = Token {
                                rule:  rule,
                                start: left,
                                end:   right
                            };

                            self.queue.insert(pos, token);
                        }
                    } else {
                        return (op, last_right)
                    }
                }

                (op, last_right)
            }

            #[inline]
            fn end(&self) -> bool {
                self.input.len() == self.input.pos()
            }

            #[inline]
            fn eoi_matched(&self) -> bool {
                self.eoi_matched
            }

            #[inline]
            fn reset(&mut self) {
                self.input.set_pos(0);
                self.queue.clear();
                self.queue_index.set(0);
                self.failures.clear();
                self.fail_pos = 0;
            }

            #[inline]
            fn queue(&self) -> &Vec<Token<Rule>>{
                &self.queue
            }

            #[inline]
            fn queue_index(&self) -> usize {
                self.queue_index.get()
            }

            #[inline]
            fn inc_queue_index(&self) {
                self.queue_index.set(self.queue_index.get() + 1);
            }

            #[inline]
            fn set_queue_index(&self, index: usize) {
                self.queue_index.set(index);
            }

            #[inline]
            fn queue_mut(&mut self) -> &mut Vec<Token<Rule>>{
                &mut self.queue
            }

            #[inline]
            fn skip_ws(&mut self) {
                if self.atomic {
                    return
                }

                loop {
                    if !self.whitespace() {
                        break
                    }
                }
            }

            #[inline]
            fn skip_com(&mut self) {
                if self.atomic {
                    return
                }

                if !self.comment {
                    self.comment = true;

                    loop {
                        if !self.comment() {
                            break
                        }
                    }

                    self.comment = false;
                }
            }

            #[inline]
            fn is_atomic(&self) -> bool {
                self.atomic
            }

            #[inline]
            fn set_atomic(&mut self, value: bool) {
                self.atomic = value;
            }

            #[inline]
            fn track(&mut self, failed: Rule, pos: usize) {
                if self.atomic {
                    return
                }

                if self.failures.is_empty() {
                    self.failures.push(failed);

                    self.fail_pos = pos;
                } else {
                    if pos == self.fail_pos {
                        self.failures.push(failed);
                    } else if pos > self.fail_pos {
                        self.failures.clear();
                        self.failures.push(failed);

                        self.fail_pos = pos;
                    }
                }
            }

            fn tracked_len(&self) -> usize {
                self.failures.len()
            }

            fn expected(&mut self) -> (Vec<Rule>, usize) {
                self.failures.sort();
                self.failures.dedup();

                (self.failures.iter().cloned().collect(), self.fail_pos)
            }
        }
    };
}

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

    impl_rdp! {
        grammar! {
            expression = _{ paren ~ expression? }
            paren = { ["("] ~ expression? ~ [")"] }
            zero = { ["a"]* }
            one = { ["a"]+ }
            comment = _{ ["//"] ~ (!["\n"] ~ any)* ~ ["\n"] }
            whitespace = _{ [" "] }
        }
    }

    #[test]
    fn try() {
        let input = StringInput::new("asdasdf");
        let mut parser = Rdp::new(input);

        assert!(parser.input_mut().match_string("asd"));

        assert!(!parser.try(false, |parser| {
            parser.input_mut().match_string("as") && parser.input_mut().match_string("dd")
        }));

        assert!(parser.try(false, |parser| {
            parser.input_mut().match_string("as") && parser.input_mut().match_string("df")
        }));
    }

    #[test]
    fn end() {
        let input = StringInput::new("asdasdf");
        let mut parser = Rdp::new(input);

        assert!(parser.input_mut().match_string("asdasdf"));
        assert!(parser.end());
    }

    #[test]
    fn reset() {
        let input = StringInput::new("asdasdf");
        let mut parser = Rdp::new(input);

        assert!(parser.input_mut().match_string("asdasdf"));

        parser.reset();

        assert!(parser.input_mut().match_string("asdasdf"));
    }

    #[test]
    fn whitespace_seq() {
        let mut parser = Rdp::new(StringInput::new("  (  ( ))(( () )() )() "));

        assert!(parser.expression());
        assert!(!parser.end());

        let queue = vec![
            Token::new(Rule::paren, 2, 9),
            Token::new(Rule::paren, 5, 8),
            Token::new(Rule::paren, 9, 20),
            Token::new(Rule::paren, 10, 16),
            Token::new(Rule::paren, 12, 14),
            Token::new(Rule::paren, 16, 18),
            Token::new(Rule::paren, 20, 22)
        ];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn whitespace_zero() {
        let mut parser = Rdp::new(StringInput::new("  a a aa aaaa a  "));

        assert!(parser.zero());
        assert!(!parser.end());

        let queue = vec![
            Token::new(Rule::zero, 2, 15)
        ];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn whitespace_one() {
        let mut parser = Rdp::new(StringInput::new("  a a aa aaaa a  "));

        assert!(parser.one());
        assert!(!parser.end());

        let queue = vec![Token {
                             rule: Rule::one,
                             start: 2,
                             end: 15,
                         }];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn comment() {
        let mut parser = Rdp::new(StringInput::new("// hi\n(())"));

        assert!(parser.expression());
        assert!(parser.end());

        let queue = vec![
            Token::new(Rule::paren, 6, 10),
            Token::new(Rule::paren, 7, 9)
        ];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn comment_whitespace() {
        let mut parser = Rdp::new(StringInput::new("   // hi\n  (())"));

        assert!(parser.expression());
        assert!(parser.end());

        let queue = vec![
            Token::new(Rule::paren, 11, 15),
            Token::new(Rule::paren, 12, 14)
        ];

        assert_eq!(parser.queue(), &queue);
    }
}
<commit_msg>Removed useless silent rule enumeration.<commit_after>// pest. Elegant, efficient grammars
// Copyright (C) 2016  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/// A `macro` useful for implementing the `Parser` `trait` as a recursive descent parser. It only
/// accepts `grammar!` and `process!` calls that get implemented on `self`.
///
/// # Rule
///
/// It also implements an `enum` called `Rule` that has a value for all
/// [non-silent](macro.grammar!#silent-rules-_) rules, but also for
/// [`any` and `eoi`](macro.grammar!). These `Rule`s are used within `Token`s to specify the type
/// of rule that matched.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate pest;
/// # use pest::prelude::*;
/// # fn main() {
/// impl_rdp! {
///     grammar! {
///         expression = _{ paren ~ expression? }
///         paren      =  { ["("] ~ expression? ~ [")"] }
///     }
/// }
///
/// let mut parser = Rdp::new(StringInput::new("(())((())())()"));
///
/// assert!(parser.expression());
/// assert!(parser.end());
/// # }
/// ```
#[macro_export]
macro_rules! impl_rdp {
    // implement rules
    ( @rules $( $name:ident )* ) => {
        #[allow(dead_code, non_camel_case_types)]
        #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
        pub enum Rule {
            any,
            eoi,
            $( $name ),*
        }
    };

    // filter out silent rules
    ( @filter [  ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@rules $( $rules )*);
    };
    ( @filter [ $name:ident = { { $( $_primary:tt )* } $( $ts:tt )* } $( $tail:tt )* ]
      [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* $( $ts )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $name:ident = @{ { $( $_primary:tt )* } $( $ts:tt )* } $( $tail:tt )* ]
      [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* $( $ts )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $_name:ident = _{ { $( $_primary:tt )* } $( $ts:tt )* } $( $tail:tt )* ]
      [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* $( $ts )* ] [ $( $rules )* ]);
    };
    ( @filter [ $name:ident = { $( $_ts:tt )* } $( $tail:tt )* ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $name:ident = @{ $( $_ts:tt )* } $( $tail:tt )* ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* ] [ $name $( $rules )* ]);
    };
    ( @filter [ $_name:ident = _{ $( $_ts:tt )* } $( $tail:tt )* ] [ $( $rules:tt )* ] ) => {
        impl_rdp!(@filter [ $( $tail )* ] [ $( $rules )* ]);
    };

    // implement empty whitespace rule
    ( @ws ) => {
        #[allow(dead_code)]
        #[inline]
        pub fn whitespace(&mut self) -> bool {
            false
        }
    };
    ( @ws whitespace = $( $_ts:tt )* ) => ();
    ( @ws $_name:ident = { $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@ws $( $tail )*);
    };
    ( @ws $_name:ident = @{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@ws $( $tail )*);
    };
    ( @ws $_name:ident = _{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@ws $( $tail )*);
    };

    // implement empty comment rule
    ( @com ) => {
        #[allow(dead_code)]
        #[inline]
        pub fn comment(&mut self) -> bool {
            false
        }
    };
    ( @com comment = $( $_ts:tt )* ) => ();
    ( @com $_name:ident = { $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@com $( $tail )*);
    };
    ( @com $_name:ident = @{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@com $( $tail )*);
    };
    ( @com $_name:ident = _{ $( $_ts:tt )* } $( $tail:tt )* ) => {
        impl_rdp!(@com $( $tail )*);
    };

    ( grammar! { $( $ts:tt )* } $( $mac:ident! { $( $rest:tt )* } )* ) => {
        use std::cell::Cell;
        use std::cmp;

        pub struct Rdp<T> {
            input:       T,
            queue:       Vec<Token<Rule>>,
            queue_index: Cell<usize>,
            failures:    Vec<Rule>,
            fail_pos:    usize,
            atomic:      bool,
            comment:     bool,
            eoi_matched: bool
        }

        impl_rdp!(@filter [ $( $ts )* ] []);

        impl<'input, T: Input<'input>> Rdp<T> {
            pub fn new(input: T) -> Rdp<T> {
                Rdp {
                    input:       input,
                    queue:       vec![],
                    queue_index: Cell::new(0),
                    failures:    vec![],
                    fail_pos:    0,
                    atomic:      false,
                    comment:     false,
                    eoi_matched: false
                }
            }

            impl_rdp!(@ws $( $ts )*);
            impl_rdp!(@com $( $ts )*);

            #[allow(dead_code)]
            #[inline]
            pub fn any(&mut self) -> bool {
                if self.end() {
                    let pos = self.input.pos();

                    self.track(Rule::any, pos);

                    false
                } else {
                    let next = self.input.pos() + 1;
                    self.input.set_pos(next);

                    true
                }
            }

            #[allow(dead_code)]
            #[inline]
            pub fn eoi(&mut self) -> bool {
                let result = self.end();

                if !result {
                    let pos = self.input.pos();

                    self.track(Rule::eoi, pos);
                } else {
                    self.eoi_matched = true;
                }

                result
            }

            grammar! {
                $( $ts )*
            }

            $(
                $mac! {
                    $( $rest )*
                }
            )*
        }

        impl<'input, T: Input<'input>> Parser<'input, T> for Rdp<T> {
            type Rule = Rule;
            type Token = Token<Rule>;

            #[inline]
            fn input(&self) -> &T {
                &self.input
            }

            #[inline]
            fn input_mut(&mut self) -> &mut T {
                &mut self.input
            }

            #[inline]
            fn try<F>(&mut self, revert: bool, rule: F) -> bool
                where F: FnOnce(&mut Self) -> bool {

                let pos = self.input.pos();
                let len = self.queue.len();

                let result = rule(self);

                if revert || !result {
                    self.input.set_pos(pos);
                }

                if !result {
                    self.queue.truncate(len);
                }

                result
            }

            fn prec_climb<F, G>(&mut self, pos: usize, left: usize, min_prec: u8,
                                last_op: Option<(Option<Rule>, u8, bool)>, primary: &mut F,
                                climb: &mut G) -> (Option<(Option<Rule>, u8, bool)>, Option<usize>)
                where F: FnMut(&mut Self) -> bool,
                      G: FnMut(&mut Self) -> Option<(Option<Rule>, u8, bool)> {

                let mut op = if last_op.is_some() {
                    last_op
                } else {
                    climb(self)
                };
                let mut last_right = None;

                while let Some((rule, prec, _)) = op {
                    if prec >= min_prec {
                        let mut new_pos = self.input.pos();
                        let mut right = self.input.pos();
                        let queue_pos = self.queue.len();

                        primary(self);

                        if let Some(token) = self.queue.get(queue_pos) {
                            new_pos = token.start;
                            right   = token.end;
                        }

                        op = climb(self);

                        while let Some((_, new_prec, right_assoc)) = op {
                            if new_prec > prec || right_assoc && new_prec == prec {
                                let (new_op, new_lr) = self.prec_climb(queue_pos, new_pos,
                                                                       new_prec, op, primary,
                                                                       climb);

                                op = new_op;
                                last_right = new_lr;
                            } else {
                                break
                            }
                        }

                        if let Some(pos) = last_right {
                            right = cmp::max(pos, right);
                        } else {
                            last_right = Some(right);
                        }

                        if let Some(rule) = rule {
                            let token = Token {
                                rule:  rule,
                                start: left,
                                end:   right
                            };

                            self.queue.insert(pos, token);
                        }
                    } else {
                        return (op, last_right)
                    }
                }

                (op, last_right)
            }

            #[inline]
            fn end(&self) -> bool {
                self.input.len() == self.input.pos()
            }

            #[inline]
            fn eoi_matched(&self) -> bool {
                self.eoi_matched
            }

            #[inline]
            fn reset(&mut self) {
                self.input.set_pos(0);
                self.queue.clear();
                self.queue_index.set(0);
                self.failures.clear();
                self.fail_pos = 0;
            }

            #[inline]
            fn queue(&self) -> &Vec<Token<Rule>>{
                &self.queue
            }

            #[inline]
            fn queue_index(&self) -> usize {
                self.queue_index.get()
            }

            #[inline]
            fn inc_queue_index(&self) {
                self.queue_index.set(self.queue_index.get() + 1);
            }

            #[inline]
            fn set_queue_index(&self, index: usize) {
                self.queue_index.set(index);
            }

            #[inline]
            fn queue_mut(&mut self) -> &mut Vec<Token<Rule>>{
                &mut self.queue
            }

            #[inline]
            fn skip_ws(&mut self) {
                if self.atomic {
                    return
                }

                loop {
                    if !self.whitespace() {
                        break
                    }
                }
            }

            #[inline]
            fn skip_com(&mut self) {
                if self.atomic {
                    return
                }

                if !self.comment {
                    self.comment = true;

                    loop {
                        if !self.comment() {
                            break
                        }
                    }

                    self.comment = false;
                }
            }

            #[inline]
            fn is_atomic(&self) -> bool {
                self.atomic
            }

            #[inline]
            fn set_atomic(&mut self, value: bool) {
                self.atomic = value;
            }

            #[inline]
            fn track(&mut self, failed: Rule, pos: usize) {
                if self.atomic {
                    return
                }

                if self.failures.is_empty() {
                    self.failures.push(failed);

                    self.fail_pos = pos;
                } else {
                    if pos == self.fail_pos {
                        self.failures.push(failed);
                    } else if pos > self.fail_pos {
                        self.failures.clear();
                        self.failures.push(failed);

                        self.fail_pos = pos;
                    }
                }
            }

            fn tracked_len(&self) -> usize {
                self.failures.len()
            }

            fn expected(&mut self) -> (Vec<Rule>, usize) {
                self.failures.sort();
                self.failures.dedup();

                (self.failures.iter().cloned().collect(), self.fail_pos)
            }
        }
    };
}

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

    impl_rdp! {
        grammar! {
            expression = _{ paren ~ expression? }
            paren = { ["("] ~ expression? ~ [")"] }
            zero = { ["a"]* }
            one = { ["a"]+ }
            comment = _{ ["//"] ~ (!["\n"] ~ any)* ~ ["\n"] }
            whitespace = _{ [" "] }
        }
    }

    #[test]
    fn try() {
        let input = StringInput::new("asdasdf");
        let mut parser = Rdp::new(input);

        assert!(parser.input_mut().match_string("asd"));

        assert!(!parser.try(false, |parser| {
            parser.input_mut().match_string("as") && parser.input_mut().match_string("dd")
        }));

        assert!(parser.try(false, |parser| {
            parser.input_mut().match_string("as") && parser.input_mut().match_string("df")
        }));
    }

    #[test]
    fn end() {
        let input = StringInput::new("asdasdf");
        let mut parser = Rdp::new(input);

        assert!(parser.input_mut().match_string("asdasdf"));
        assert!(parser.end());
    }

    #[test]
    fn reset() {
        let input = StringInput::new("asdasdf");
        let mut parser = Rdp::new(input);

        assert!(parser.input_mut().match_string("asdasdf"));

        parser.reset();

        assert!(parser.input_mut().match_string("asdasdf"));
    }

    #[test]
    fn whitespace_seq() {
        let mut parser = Rdp::new(StringInput::new("  (  ( ))(( () )() )() "));

        assert!(parser.expression());
        assert!(!parser.end());

        let queue = vec![
            Token::new(Rule::paren, 2, 9),
            Token::new(Rule::paren, 5, 8),
            Token::new(Rule::paren, 9, 20),
            Token::new(Rule::paren, 10, 16),
            Token::new(Rule::paren, 12, 14),
            Token::new(Rule::paren, 16, 18),
            Token::new(Rule::paren, 20, 22)
        ];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn whitespace_zero() {
        let mut parser = Rdp::new(StringInput::new("  a a aa aaaa a  "));

        assert!(parser.zero());
        assert!(!parser.end());

        let queue = vec![
            Token::new(Rule::zero, 2, 15)
        ];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn whitespace_one() {
        let mut parser = Rdp::new(StringInput::new("  a a aa aaaa a  "));

        assert!(parser.one());
        assert!(!parser.end());

        let queue = vec![Token {
                             rule: Rule::one,
                             start: 2,
                             end: 15,
                         }];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn comment() {
        let mut parser = Rdp::new(StringInput::new("// hi\n(())"));

        assert!(parser.expression());
        assert!(parser.end());

        let queue = vec![
            Token::new(Rule::paren, 6, 10),
            Token::new(Rule::paren, 7, 9)
        ];

        assert_eq!(parser.queue(), &queue);
    }

    #[test]
    fn comment_whitespace() {
        let mut parser = Rdp::new(StringInput::new("   // hi\n  (())"));

        assert!(parser.expression());
        assert!(parser.end());

        let queue = vec![
            Token::new(Rule::paren, 11, 15),
            Token::new(Rule::paren, 12, 14)
        ];

        assert_eq!(parser.queue(), &queue);
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::fmt;

use log::trace;

use rustc_span::DUMMY_SP;

use crate::*;

/// Details of premature program termination.
pub enum TerminationInfo {
    Exit(i64),
    Abort(Option<String>),
    UnsupportedInIsolation(String),
    ExperimentalUb { msg: String, url: String },
    Deadlock,
}

impl fmt::Debug for TerminationInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use TerminationInfo::*;
        match self {
            Exit(code) =>
                write!(f, "the evaluated program completed with exit code {}", code),
            Abort(None) =>
                write!(f, "the evaluated program aborted execution"),
            Abort(Some(msg)) =>
                write!(f, "the evaluated program aborted execution: {}", msg),
            UnsupportedInIsolation(msg) =>
                write!(f, "{}", msg),
            ExperimentalUb { msg, .. } =>
                write!(f, "{}", msg),
            Deadlock =>
                write!(f, "the evaluated program deadlocked"),
        }
    }
}

impl MachineStopType for TerminationInfo {}

/// Miri specific diagnostics
pub enum NonHaltingDiagnostic {
    PoppedTrackedPointerTag(Item),
    CreatedAlloc(AllocId),
}

/// Emit a custom diagnostic without going through the miri-engine machinery
pub fn report_error<'tcx, 'mir>(
    ecx: &InterpCx<'mir, 'tcx, Evaluator<'tcx>>,
    mut e: InterpErrorInfo<'tcx>,
) -> Option<i64> {
    use InterpError::*;

    let (title, helps) = match &e.kind {
        MachineStop(info) => {
            let info = info.downcast_ref::<TerminationInfo>().expect("invalid MachineStop payload");
            use TerminationInfo::*;
            let title = match info {
                Exit(code) => return Some(*code),
                Abort(_) =>
                    "abnormal termination",
                UnsupportedInIsolation(_) =>
                    "unsupported operation",
                ExperimentalUb { .. } =>
                    "Undefined Behavior",
                Deadlock => "deadlock",
            };
            let helps = match info {
                UnsupportedInIsolation(_) =>
                    vec![format!("pass the flag `-Zmiri-disable-isolation` to disable isolation")],
                ExperimentalUb { url, .. } =>
                    vec![
                        format!("this indicates a potential bug in the program: it performed an invalid operation, but the rules it violated are still experimental"),
                        format!("see {} for further information", url),
                    ],
                _ => vec![],
            };
            (title, helps)
        }
        _ => {
            let title = match e.kind {
                Unsupported(_) =>
                    "unsupported operation",
                UndefinedBehavior(_) =>
                    "Undefined Behavior",
                ResourceExhaustion(_) =>
                    "resource exhaustion",
                _ =>
                    bug!("This error should be impossible in Miri: {}", e),
            };
            let helps = match e.kind {
                Unsupported(UnsupportedOpInfo::NoMirFor(..)) =>
                    vec![format!("make sure to use a Miri sysroot, which you can prepare with `cargo miri setup`")],
                Unsupported(_) =>
                    vec![format!("this is likely not a bug in the program; it indicates that the program performed an operation that the interpreter does not support")],
                UndefinedBehavior(_) =>
                    vec![
                        format!("this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior"),
                        format!("see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information"),
                    ],
                _ => vec![],
            };
            (title, helps)
        }
    };

    e.print_backtrace();
    let msg = e.to_string();
    report_msg(ecx, &format!("{}: {}", title, msg), msg, helps, true)
}

/// Report an error or note (depending on the `error` argument) at the current frame's current statement.
/// Also emits a full stacktrace of the interpreter stack.
fn report_msg<'tcx, 'mir>(
    ecx: &InterpCx<'mir, 'tcx, Evaluator<'tcx>>,
    title: &str,
    span_msg: String,
    mut helps: Vec<String>,
    error: bool,
) -> Option<i64> {
    let span = if let Some(frame) = ecx.stack().last() {
        frame.current_source_info().unwrap().span
    } else {
        DUMMY_SP
    };
    let mut err = if error {
        ecx.tcx.sess.struct_span_err(span, title)
    } else {
        ecx.tcx.sess.diagnostic().span_note_diag(span, title)
    };
    err.span_label(span, span_msg);
    if !helps.is_empty() {
        // Add visual separator before backtrace.
        helps.last_mut().unwrap().push_str("\n");
        for help in helps {
            err.help(&help);
        }
    }
    // Add backtrace
    let frames = ecx.generate_stacktrace();
    for (idx, frame_info) in frames.iter().enumerate() {
        let is_local = frame_info.instance.def_id().is_local();
        // No span for non-local frames and the first frame (which is the error site).
        if is_local && idx > 0 {
            err.span_note(frame_info.span, &frame_info.to_string());
        } else {
            err.note(&frame_info.to_string());
        }
    }

    err.emit();

    for (i, frame) in ecx.stack().iter().enumerate() {
        trace!("-------------------");
        trace!("Frame {}", i);
        trace!("    return: {:?}", frame.return_place.map(|p| *p));
        for (i, local) in frame.locals.iter().enumerate() {
            trace!("    local {}: {:?}", i, local.value);
        }
    }
    // Let the reported error determine the return code.
    return None;
}

thread_local! {
    static DIAGNOSTICS: RefCell<Vec<NonHaltingDiagnostic>> = RefCell::new(Vec::new());
}

/// Schedule a diagnostic for emitting. This function works even if you have no `InterpCx` available.
/// The diagnostic will be emitted after the current interpreter step is finished.
pub fn register_diagnostic(e: NonHaltingDiagnostic) {
    DIAGNOSTICS.with(|diagnostics| diagnostics.borrow_mut().push(e));
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    /// Emit all diagnostics that were registed with `register_diagnostics`
    fn process_diagnostics(&self) {
        let this = self.eval_context_ref();
        DIAGNOSTICS.with(|diagnostics| {
            for e in diagnostics.borrow_mut().drain(..) {
                use NonHaltingDiagnostic::*;
                let msg = match e {
                    PoppedTrackedPointerTag(item) =>
                        format!("popped tracked tag for item {:?}", item),
                    CreatedAlloc(AllocId(id)) =>
                        format!("created allocation with id {}", id),
                };
                report_msg(this, "tracking was triggered", msg, vec![], false);
            }
        });
    }
}
<commit_msg>for alignment errors, note that there might be false positives<commit_after>use std::cell::RefCell;
use std::fmt;

use log::trace;

use rustc_span::DUMMY_SP;

use crate::*;

/// Details of premature program termination.
pub enum TerminationInfo {
    Exit(i64),
    Abort(Option<String>),
    UnsupportedInIsolation(String),
    ExperimentalUb { msg: String, url: String },
    Deadlock,
}

impl fmt::Debug for TerminationInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use TerminationInfo::*;
        match self {
            Exit(code) =>
                write!(f, "the evaluated program completed with exit code {}", code),
            Abort(None) =>
                write!(f, "the evaluated program aborted execution"),
            Abort(Some(msg)) =>
                write!(f, "the evaluated program aborted execution: {}", msg),
            UnsupportedInIsolation(msg) =>
                write!(f, "{}", msg),
            ExperimentalUb { msg, .. } =>
                write!(f, "{}", msg),
            Deadlock =>
                write!(f, "the evaluated program deadlocked"),
        }
    }
}

impl MachineStopType for TerminationInfo {}

/// Miri specific diagnostics
pub enum NonHaltingDiagnostic {
    PoppedTrackedPointerTag(Item),
    CreatedAlloc(AllocId),
}

/// Emit a custom diagnostic without going through the miri-engine machinery
pub fn report_error<'tcx, 'mir>(
    ecx: &InterpCx<'mir, 'tcx, Evaluator<'tcx>>,
    mut e: InterpErrorInfo<'tcx>,
) -> Option<i64> {
    use InterpError::*;

    let (title, helps) = match &e.kind {
        MachineStop(info) => {
            let info = info.downcast_ref::<TerminationInfo>().expect("invalid MachineStop payload");
            use TerminationInfo::*;
            let title = match info {
                Exit(code) => return Some(*code),
                Abort(_) =>
                    "abnormal termination",
                UnsupportedInIsolation(_) =>
                    "unsupported operation",
                ExperimentalUb { .. } =>
                    "Undefined Behavior",
                Deadlock => "deadlock",
            };
            let helps = match info {
                UnsupportedInIsolation(_) =>
                    vec![format!("pass the flag `-Zmiri-disable-isolation` to disable isolation")],
                ExperimentalUb { url, .. } =>
                    vec![
                        format!("this indicates a potential bug in the program: it performed an invalid operation, but the rules it violated are still experimental"),
                        format!("see {} for further information", url),
                    ],
                _ => vec![],
            };
            (title, helps)
        }
        _ => {
            let title = match e.kind {
                Unsupported(_) =>
                    "unsupported operation",
                UndefinedBehavior(_) =>
                    "Undefined Behavior",
                ResourceExhaustion(_) =>
                    "resource exhaustion",
                _ =>
                    bug!("This error should be impossible in Miri: {}", e),
            };
            let helps = match e.kind {
                Unsupported(UnsupportedOpInfo::NoMirFor(..)) =>
                    vec![format!("make sure to use a Miri sysroot, which you can prepare with `cargo miri setup`")],
                Unsupported(_) =>
                    vec![format!("this is likely not a bug in the program; it indicates that the program performed an operation that the interpreter does not support")],
                UndefinedBehavior(UndefinedBehaviorInfo::AlignmentCheckFailed { .. }) =>
                    vec![
                        format!("this usually indicates that your program performed an invalid operation and caused Undefined Behavior"),
                        format!("but alignment errors can also be false positives, see https://github.com/rust-lang/miri/issues/1074"),
                    ],
                UndefinedBehavior(_) =>
                    vec![
                        format!("this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior"),
                        format!("see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information"),
                    ],
                _ => vec![],
            };
            (title, helps)
        }
    };

    e.print_backtrace();
    let msg = e.to_string();
    report_msg(ecx, &format!("{}: {}", title, msg), msg, helps, true)
}

/// Report an error or note (depending on the `error` argument) at the current frame's current statement.
/// Also emits a full stacktrace of the interpreter stack.
fn report_msg<'tcx, 'mir>(
    ecx: &InterpCx<'mir, 'tcx, Evaluator<'tcx>>,
    title: &str,
    span_msg: String,
    mut helps: Vec<String>,
    error: bool,
) -> Option<i64> {
    let span = if let Some(frame) = ecx.stack().last() {
        frame.current_source_info().unwrap().span
    } else {
        DUMMY_SP
    };
    let mut err = if error {
        ecx.tcx.sess.struct_span_err(span, title)
    } else {
        ecx.tcx.sess.diagnostic().span_note_diag(span, title)
    };
    err.span_label(span, span_msg);
    if !helps.is_empty() {
        // Add visual separator before backtrace.
        helps.last_mut().unwrap().push_str("\n");
        for help in helps {
            err.help(&help);
        }
    }
    // Add backtrace
    let frames = ecx.generate_stacktrace();
    for (idx, frame_info) in frames.iter().enumerate() {
        let is_local = frame_info.instance.def_id().is_local();
        // No span for non-local frames and the first frame (which is the error site).
        if is_local && idx > 0 {
            err.span_note(frame_info.span, &frame_info.to_string());
        } else {
            err.note(&frame_info.to_string());
        }
    }

    err.emit();

    for (i, frame) in ecx.stack().iter().enumerate() {
        trace!("-------------------");
        trace!("Frame {}", i);
        trace!("    return: {:?}", frame.return_place.map(|p| *p));
        for (i, local) in frame.locals.iter().enumerate() {
            trace!("    local {}: {:?}", i, local.value);
        }
    }
    // Let the reported error determine the return code.
    return None;
}

thread_local! {
    static DIAGNOSTICS: RefCell<Vec<NonHaltingDiagnostic>> = RefCell::new(Vec::new());
}

/// Schedule a diagnostic for emitting. This function works even if you have no `InterpCx` available.
/// The diagnostic will be emitted after the current interpreter step is finished.
pub fn register_diagnostic(e: NonHaltingDiagnostic) {
    DIAGNOSTICS.with(|diagnostics| diagnostics.borrow_mut().push(e));
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    /// Emit all diagnostics that were registed with `register_diagnostics`
    fn process_diagnostics(&self) {
        let this = self.eval_context_ref();
        DIAGNOSTICS.with(|diagnostics| {
            for e in diagnostics.borrow_mut().drain(..) {
                use NonHaltingDiagnostic::*;
                let msg = match e {
                    PoppedTrackedPointerTag(item) =>
                        format!("popped tracked tag for item {:?}", item),
                    CreatedAlloc(AllocId(id)) =>
                        format!("created allocation with id {}", id),
                };
                report_msg(this, "tracking was triggered", msg, vec![], false);
            }
        });
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::io::BufReader;
use std::io::Error as IoError;
use std::ascii::AsciiExt;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(read: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(read);
        let mut magic: [u8; 2] = [0, 0];
        try!(buf.read_exact(&mut magic[..])); // Skip magic constant
        if magic[0] != b'P' || (magic[1] != b'3' && magic[1] != b'6') {
            return Err(ImageError::FormatError("Expected magic constant for ppm, P3 or P6".to_string()));
        }

        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));

        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
        })
    }

    /// Reads a string as well as a single whitespace after it, ignoring comments
    fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
        let mut bytes = Vec::new();

        let mark_comments = reader
            .bytes()
            .scan(true, |partof, read| {
                let byte = match read {
                    Err(err) => return Some((*partof, Err(err))),
                    Ok(byte) => byte,
                };
                let cur_enabled = *partof && byte != b'#';
                let next_enabled = cur_enabled || (byte == b'\r' || byte == b'\n');
                *partof = next_enabled;
                return Some((cur_enabled, Ok(byte)));
            });

        for (_, byte) in mark_comments.filter(|ref e| e.0) {
            match byte {
                Ok(b'\t') | Ok(b'\n') | Ok(b'\x0b') | Ok(b'\x0c') | Ok(b'\r') | Ok(b' ') => {
                    if !bytes.is_empty() {
                        break // We're done as we already have some content
                    }
                },
                Ok(byte) => {
                    bytes.push(byte);
                },
                Err(_) => break,
            }
        }

        if bytes.is_empty() {
            return Err(ImageError::FormatError("Unexpected eof".to_string()))
        }

        if !bytes.as_slice().is_ascii() {
            return Err(ImageError::FormatError("Non ascii character in preamble"))
        }

        String::from_utf8(bytes).map_err(|_| ImageError::FormatError("Couldn't read preamble".to_string()))
    }

    fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
        let s = try!(PPMDecoder::read_next_string(reader));
        s.parse::<u32>().map_err(|_| ImageError::FormatError("Invalid number in preamble".to_string()))
    }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.bytewidth() {
            1 => Ok(ColorType::RGB(8)),
            2 => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*3*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let opt_size = self.width.checked_mul(self.height)
            .map_or(None, |v| v.checked_mul(3))
            .map_or(None, |v| v.checked_mul(self.bytewidth()));

        let size = match opt_size {
            Some(v) => v,
            None => return Err(ImageError::DimensionError),
        };

        let mut data = vec![0 as u8; size as usize];

        match self.reader.read_exact(&mut data) {
            Ok(_) => {},
            Err(e) => return Err(ImageError::IoError(e)),
        };

        if self.bytewidth() == 1 {
            Ok(DecodingResult::U8(data))
        } else {
            let mut out = vec![0 as u16; (self.width*self.height*3) as usize];
            for (o, i) in out.chunks_mut(1).zip(data.chunks(2)) {
                o[0] = BigEndian::read_u16(i);
            }
            Ok(DecodingResult::U16(out))
        }
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }
}

/// Tests parsing binary buffers were written based on and validated against `identify` from
/// ImageMagick (https://www.imagemagick.org/script/identify.php).
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn minimal_form() {
        // Violates current specification (October 2016 ) but accepted by both netpbm and ImageMagick
        decode_minimal_image(&b"P61 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123\xFF"[..]); // Too long should not be an issue
    }

    #[test]
    fn comment_in_token() {
        decode_minimal_image(&b"P6 1 1 2#comment\n55 123"[..]); // Terminating LF
        decode_minimal_image(&b"P6 1 1 2#comment\r55 123"[..]); // Terminating CR
        decode_minimal_image(&b"P6 1 1#comment\n 255 123"[..]); // Comment after token
        decode_minimal_image(&b"P6 1 1 #comment\n255 123"[..]); // Comment before token
        decode_minimal_image(&b"P6#comment\n 1 1 255 123"[..]); // Begin of header
        decode_minimal_image(&b"P6 1 1 255#comment\n 123"[..]); // End of header
    }

    #[test]
    fn whitespace() {
        decode_minimal_image(&b"P6\x091\x091\x09255\x09123"[..]); // TAB
        decode_minimal_image(&b"P6\x0a1\x0a1\x0a255\x0a123"[..]); // LF
        decode_minimal_image(&b"P6\x0b1\x0b1\x0b255\x0b123"[..]); // VT
        decode_minimal_image(&b"P6\x0c1\x0c1\x0c255\x0c123"[..]); // FF
        decode_minimal_image(&b"P6\x0d1\x0d1\x0d255\x0d123"[..]); // CR
        // Spaces tested before
        decode_minimal_image(&b"P61\x09\x0a\x0b\x0c\x0d1 255 123"[..]); // All whitespace, combined
    }

    /// Tests for decoding error, assuming `encoded` is ppm encoding for the very simplistic image
    /// containing a single pixel with one byte values (1, 2, 3).
    fn decode_minimal_image(encoded: &[u8]) {
        let content = vec![49 as u8, 50, 51];
        let mut decoder = super::PPMDecoder::new(encoded).unwrap();

        assert_eq!(decoder.dimensions().unwrap(), (1, 1));
        assert_eq!(decoder.colortype().unwrap(), ColorType::RGB(8));
        assert_eq!(decoder.row_len().unwrap(), 3);
        assert_eq!(decoder.bytewidth(), 1);

        match decoder.read_image().unwrap() {
            DecodingResult::U8(image) => assert_eq!(image, content),
            _ => assert!(false),
        }
    }
}
<commit_msg>Add tests for failure<commit_after>use std::io::Read;
use std::io::BufReader;
use std::io::Error as IoError;
use std::ascii::AsciiExt;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(read: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(read);
        let mut magic: [u8; 2] = [0, 0];
        try!(buf.read_exact(&mut magic[..])); // Skip magic constant
        if magic[0] != b'P' || (magic[1] != b'3' && magic[1] != b'6') {
            return Err(ImageError::FormatError("Expected magic constant for ppm, P3 or P6".to_string()));
        }

        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));

        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
        })
    }

    /// Reads a string as well as a single whitespace after it, ignoring comments
    fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
        let mut bytes = Vec::new();

        // pair input bytes with a bool mask to remove comments
        let mark_comments = reader
            .bytes()
            .scan(true, |partof, read| {
                let byte = match read {
                    Err(err) => return Some((*partof, Err(err))),
                    Ok(byte) => byte,
                };
                let cur_enabled = *partof && byte != b'#';
                let next_enabled = cur_enabled || (byte == b'\r' || byte == b'\n');
                *partof = next_enabled;
                return Some((cur_enabled, Ok(byte)));
            });

        for (_, byte) in mark_comments.filter(|ref e| e.0) {
            match byte {
                Ok(b'\t') | Ok(b'\n') | Ok(b'\x0b') | Ok(b'\x0c') | Ok(b'\r') | Ok(b' ') => {
                    if !bytes.is_empty() {
                        break // We're done as we already have some content
                    }
                },
                Ok(byte) => {
                    bytes.push(byte);
                },
                Err(_) => break,
            }
        }

        if bytes.is_empty() {
            return Err(ImageError::FormatError("Unexpected eof".to_string()))
        }

        if !bytes.as_slice().is_ascii() {
            return Err(ImageError::FormatError("Non ascii character in preamble".to_string()))
        }

        String::from_utf8(bytes).map_err(|_| ImageError::FormatError("Couldn't read preamble".to_string()))
    }

    fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
        let s = try!(PPMDecoder::read_next_string(reader));
        s.parse::<u32>().map_err(|_| ImageError::FormatError("Invalid number in preamble".to_string()))
    }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.bytewidth() {
            1 => Ok(ColorType::RGB(8)),
            2 => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*3*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let opt_size = self.width.checked_mul(self.height)
            .map_or(None, |v| v.checked_mul(3))
            .map_or(None, |v| v.checked_mul(self.bytewidth()));

        let size = match opt_size {
            Some(v) => v,
            None => return Err(ImageError::DimensionError),
        };

        let mut data = vec![0 as u8; size as usize];

        match self.reader.read_exact(&mut data) {
            Ok(_) => {},
            Err(e) => return Err(ImageError::IoError(e)),
        };

        if self.bytewidth() == 1 {
            Ok(DecodingResult::U8(data))
        } else {
            let mut out = vec![0 as u16; (self.width*self.height*3) as usize];
            for (o, i) in out.chunks_mut(1).zip(data.chunks(2)) {
                o[0] = BigEndian::read_u16(i);
            }
            Ok(DecodingResult::U16(out))
        }
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }
}

/// Tests parsing binary buffers were written based on and validated against `pamfile` from
/// netpbm (http://netpbm.sourceforge.net/).
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn minimal_form() {
        // Violates current specification (October 2016 ) but accepted by both netpbm and ImageMagick
        decode_minimal_image(&b"P61 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123\xFF"[..]); // Too long should not be an issue
    }

    #[test]
    fn comment_in_token() {
        decode_minimal_image(&b"P6 1 1 2#comment\n55 123"[..]); // Terminating LF
        decode_minimal_image(&b"P6 1 1 2#comment\r55 123"[..]); // Terminating CR
        decode_minimal_image(&b"P6 1 1#comment\n 255 123"[..]); // Comment after token
        decode_minimal_image(&b"P6 1 1 #comment\n255 123"[..]); // Comment before token
        decode_minimal_image(&b"P6#comment\n 1 1 255 123"[..]); // Begin of header
        decode_minimal_image(&b"P6 1 1 255#comment\n 123"[..]); // End of header
    }

    #[test]
    fn whitespace() {
        decode_minimal_image(&b"P6\x091\x091\x09255\x09123"[..]); // TAB
        decode_minimal_image(&b"P6\x0a1\x0a1\x0a255\x0a123"[..]); // LF
        decode_minimal_image(&b"P6\x0b1\x0b1\x0b255\x0b123"[..]); // VT
        decode_minimal_image(&b"P6\x0c1\x0c1\x0c255\x0c123"[..]); // FF
        decode_minimal_image(&b"P6\x0d1\x0d1\x0d255\x0d123"[..]); // CR
        // Spaces tested before
        decode_minimal_image(&b"P61\x09\x0a\x0b\x0c\x0d1 255 123"[..]); // All whitespace, combined
    }

    /// Tests for decoding error, assuming `encoded` is ppm encoding for the very simplistic image
    /// containing a single pixel with one byte values (1, 2, 3).
    fn decode_minimal_image(encoded: &[u8]) {
        let content = vec![49 as u8, 50, 51];
        let mut decoder = PPMDecoder::new(encoded).unwrap();

        assert_eq!(decoder.dimensions().unwrap(), (1, 1));
        assert_eq!(decoder.colortype().unwrap(), ColorType::RGB(8));
        assert_eq!(decoder.row_len().unwrap(), 3);
        assert_eq!(decoder.bytewidth(), 1);

        match decoder.read_image().unwrap() {
            DecodingResult::U8(image) => assert_eq!(image, content),
            _ => assert!(false),
        }
    }

    #[test]
    fn wrong_tag() {
        assert!(PPMDecoder::new(&b"P5 1 1 255 1"[..]).is_err());
    }

    #[test]
    fn invalid_characters() {
        assert!(PPMDecoder::new(&b"P6 1chars1 255 1"[..]).is_err()); // No text outside of comments
        assert!(PPMDecoder::new(&b"P6 1\xFF1 255 1"[..]).is_err()); // No invalid ascii chars
        assert!(PPMDecoder::new(&b"P6 0x01 1 255 1"[..]).is_err()); // Numbers only as decimal
    }

    /// These violate the narrow specification of ppm but are commonly supported in other programs.
    /// Fail fast and concise is important here as these might be received as input files.
    #[test]
    fn unsupported_extensions() {
        assert!(PPMDecoder::new(&b"P6 1 1 65536 1"[..]).is_err()); // No bitwidth above 16
    }
}
<|endoftext|>


<commit_before>extern crate linked_list_allocator;

use self::linked_list_allocator::Heap;
use core::alloc::Alloc;
use core::alloc::GlobalAlloc;
use core::alloc::Layout;
use core::mem;
use core::ptr;
use core::ptr::NonNull;
use syscalls;

const HEAP_SIZE: usize = 0x400;

// None-threaded heap wrapper based on `r9` register instead of global variable
pub(crate) struct TockAllocator;

impl TockAllocator {
    unsafe fn heap(&self) -> &mut Heap {
        let heap: *mut Heap;
        asm!("mov $0, r9" : "=r"(heap) : : : "volatile");
        &mut *heap
    }

    /// Initializes an empty heap
    ///
    /// # Unsafety
    ///
    /// This function must be called at most once. The memory between [`heap_location`] and [`heap_top`] must not overlap with any other memory section.
    #[inline(never)]
    unsafe fn init(&mut self, heap_bottom: usize, heap_top: usize) {
        asm!("mov r9, $0" : : "r"(heap_bottom) : : "volatile");

        let effective_heap_bottom = heap_bottom + mem::size_of::<Heap>();

        let heap = heap_bottom as *mut Heap;
        *heap = Heap::new(effective_heap_bottom, heap_top - effective_heap_bottom);
    }
}

unsafe impl GlobalAlloc for TockAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.heap().alloc(layout).unwrap().as_ptr()
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.heap().dealloc(NonNull::new_unchecked(ptr), layout)
    }
}

/// Tock programs' entry point
#[doc(hidden)]
#[no_mangle]
#[naked]
#[link_section = ".start"]
pub unsafe extern "C" fn _start(
    text_start: usize,
    mem_start: usize,
    _memory_len: usize,
    _app_heap_break: usize,
) -> ! {
    extern "C" {
        // This function is created internally by`rustc`. See `src/lang_items.rs` for more details.
        fn main(argc: isize, argv: *const *const u8) -> isize;
    }

    let stack_size = *(text_start as *const usize).offset(9);
    let real_stack_top = mem_start + stack_size;
    // Set the effective stack top below the real stack stop and use the space in between for the heap
    let effective_stack_top = real_stack_top - HEAP_SIZE;
    TockAllocator.init(effective_stack_top, real_stack_top);

    asm!("mov sp, $0" : : "r"(effective_stack_top) : "memory" :  "volatile" );

    main(0, ptr::null());

    loop {
        syscalls::yieldk();
    }
}
<commit_msg>Tell the kernel where the app put its stack<commit_after>extern crate linked_list_allocator;

use self::linked_list_allocator::Heap;
use core::alloc::Alloc;
use core::alloc::GlobalAlloc;
use core::alloc::Layout;
use core::mem;
use core::ptr;
use core::ptr::NonNull;
use syscalls;

const HEAP_SIZE: usize = 0x400;

// None-threaded heap wrapper based on `r9` register instead of global variable
pub(crate) struct TockAllocator;

impl TockAllocator {
    unsafe fn heap(&self) -> &mut Heap {
        let heap: *mut Heap;
        asm!("mov $0, r9" : "=r"(heap) : : : "volatile");
        &mut *heap
    }

    /// Initializes an empty heap
    ///
    /// # Unsafety
    ///
    /// This function must be called at most once. The memory between [`heap_location`] and [`heap_top`] must not overlap with any other memory section.
    #[inline(never)]
    unsafe fn init(&mut self, heap_bottom: usize, heap_top: usize) {
        asm!("mov r9, $0" : : "r"(heap_bottom) : : "volatile");

        let effective_heap_bottom = heap_bottom + mem::size_of::<Heap>();

        let heap = heap_bottom as *mut Heap;
        *heap = Heap::new(effective_heap_bottom, heap_top - effective_heap_bottom);
    }
}

unsafe impl GlobalAlloc for TockAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.heap().alloc(layout).unwrap().as_ptr()
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.heap().dealloc(NonNull::new_unchecked(ptr), layout)
    }
}

/// Tock programs' entry point
#[doc(hidden)]
#[no_mangle]
#[naked]
#[link_section = ".start"]
pub unsafe extern "C" fn _start(
    text_start: usize,
    mem_start: usize,
    _memory_len: usize,
    _app_heap_break: usize,
) -> ! {
    extern "C" {
        // This function is created internally by`rustc`. See `src/lang_items.rs` for more details.
        fn main(argc: isize, argv: *const *const u8) -> isize;
    }

    let stack_size = *(text_start as *const usize).offset(9);
    let real_stack_top = mem_start + stack_size;
    // Set the effective stack top below the real stack stop and use the space in between for the heap
    let effective_stack_top = real_stack_top - HEAP_SIZE;
    TockAllocator.init(effective_stack_top, real_stack_top);

    asm!("mov sp, $0" : : "r"(effective_stack_top) : "memory" :  "volatile" );

    syscalls::memop(0, effective_stack_top + HEAP_SIZE);
    syscalls::memop(11, effective_stack_top + HEAP_SIZE);
    syscalls::memop(10, effective_stack_top);

    main(0, ptr::null());

    loop {
        syscalls::yieldk();
    }
}
<|endoftext|>


<commit_before>extern crate linked_list_allocator;

use self::linked_list_allocator::Heap;
use alloc::allocator::Alloc;
use alloc::allocator::AllocErr;
use alloc::allocator::Layout;
use alloc::string::String;
use console::Console;
use core::mem;
use core::ptr;
use fmt;
use syscalls;

// None-threaded heap wrapper based on `r9` register instead of global variable
pub(crate) struct StackOwnedHeap;

impl StackOwnedHeap {
    unsafe fn heap(&self) -> &mut Heap {
        let mut heap: *mut Heap;
        asm!("mov $0, r9" : "=r"(heap) : : : "volatile");
        heap = ((heap as u32) - HEAP_SIZE as u32 + HEAP_OFFSET as u32) as *mut Heap;
        &mut *heap
    }

    /// Initializes an empty heap
    ///
    /// # Unsafety
    ///
    /// This function must be called at most once. The heap_buffer must remain valid until the end of the process.
    #[inline(never)]
    unsafe fn init(&mut self, heap_buffer: &mut [u8]) {
        let heap_location = heap_buffer.as_ptr() as usize;

        let heap_bottom = heap_location + mem::size_of::<Heap>();
        let heap_top = heap_location + heap_buffer.len();
        *self.heap() = Heap::new(heap_bottom, heap_top - heap_bottom);
    }
}

unsafe impl<'a> Alloc for &'a StackOwnedHeap {
    unsafe fn alloc(&mut self, layout: Layout) -> Result<*mut u8, AllocErr> {
        self.heap().alloc(layout)
    }

    unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
        self.heap().dealloc(ptr, layout)
    }
}

const HEAP_SIZE: usize = 0x200;
const HEAP_OFFSET: usize = 0x1c;
const STACK_SIZE: usize = 0x400;

/// Tock programs' entry point
#[doc(hidden)]
#[no_mangle]
#[naked]
#[link_section = ".start"]
pub extern "C" fn _start(
    _text_start: usize,
    mem_start: usize,
    _memory_len: usize,
    _app_heap_break: usize,
) -> ! {
    unsafe {
        asm!("mov sp, $0" : : "r"(mem_start+HEAP_SIZE+STACK_SIZE) : "memory" :  "volatile" );
        asm!("mov r9, $0" : : "r"(mem_start+HEAP_SIZE+STACK_SIZE) : : "volatile");
        run_with_new_stack()
    }
}

#[inline(never)]
unsafe fn run_with_new_stack() -> ! {
    extern "C" {
        // This function is created internally by`rustc`. See `src/lang_items.rs` for more details.
        fn main(argc: isize, argv: *const *const u8) -> isize;
    }
    let mut heap: [u8; HEAP_SIZE as usize] = [0; HEAP_SIZE as usize];

    StackOwnedHeap.init(&mut heap);

    let mut console = Console::new();
    console.write(String::from(
        "\nProcess started\n===============\nHeap position: \n",
    ));
    console.write(fmt::u32_as_hex(heap.as_ptr() as u32));
    console.write(String::from("\n"));

    main(0, ptr::null());

    loop {
        syscalls::yieldk();
    }
}
<commit_msg>Remove magical calculations from entry_point.rs<commit_after>extern crate linked_list_allocator;

use self::linked_list_allocator::Heap;
use alloc::allocator::Alloc;
use alloc::allocator::AllocErr;
use alloc::allocator::Layout;
use alloc::string::String;
use console::Console;
use core::mem;
use core::ptr;
use fmt;
use syscalls;

const HEAP_SIZE: usize = 0x200;
const STACK_SIZE: usize = 0x400;

// None-threaded heap wrapper based on `r9` register instead of global variable
pub(crate) struct StackOwnedHeap;

impl StackOwnedHeap {
    unsafe fn heap(&self) -> &mut Heap {
        let heap: *mut Heap;
        asm!("mov $0, r9" : "=r"(heap) : : : "volatile");
        &mut *heap
    }

    /// Initializes an empty heap
    ///
    /// # Unsafety
    ///
    /// This function must be called at most once. The heap_buffer must remain valid until the end of the process.
    #[inline(never)]
    unsafe fn init(&mut self, heap_buffer: &mut [u8]) {
        let heap_location = heap_buffer.as_ptr() as usize;
        asm!("mov r9, $0" : : "r"(heap_location) : : "volatile");
        let heap = heap_location as *mut Heap;

        let heap_bottom = heap_location + mem::size_of::<Heap>();
        let heap_top = heap_location + heap_buffer.len();
        *heap = Heap::new(heap_bottom, heap_top - heap_bottom);
    }
}

unsafe impl<'a> Alloc for &'a StackOwnedHeap {
    unsafe fn alloc(&mut self, layout: Layout) -> Result<*mut u8, AllocErr> {
        self.heap().alloc(layout)
    }

    unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
        self.heap().dealloc(ptr, layout)
    }
}

/// Tock programs' entry point
#[doc(hidden)]
#[no_mangle]
#[naked]
#[link_section = ".start"]
pub extern "C" fn _start(
    _text_start: usize,
    mem_start: usize,
    _memory_len: usize,
    _app_heap_break: usize,
) -> ! {
    unsafe {
        asm!("mov sp, $0" : : "r"(mem_start + HEAP_SIZE + STACK_SIZE) : "memory" :  "volatile" );
        run_with_new_stack()
    }
}

#[naked]
#[inline(never)]
unsafe fn run_with_new_stack() -> ! {
    extern "C" {
        // This function is created internally by`rustc`. See `src/lang_items.rs` for more details.
        fn main(argc: isize, argv: *const *const u8) -> isize;
    }
    let mut heap: [u8; HEAP_SIZE as usize] = [0; HEAP_SIZE as usize];

    StackOwnedHeap.init(&mut heap);

    let mut console = Console::new();
    console.write(String::from(
        "\nProcess started\n===============\nHeap position: ",
    ));
    console.write(fmt::u32_as_hex(heap.as_ptr() as u32));
    console.write(String::from("\n\n"));

    main(0, ptr::null());

    loop {
        syscalls::yieldk();
    }
}
<|endoftext|>


<commit_before>use std::io::net::ip::SocketAddr;
use std::io::net::tcp::TcpStream;
use std::io::{BufferedStream, Stream};
use std::vec::bytes::push_bytes;
use std::vec;
use std::str::from_utf8;

pub enum Result {
  Nil,
  Int(i64),
  Data(~[u8]),
  List(~[Result]),
  Error(~str),
  Status(~str),
  ProtocolError(&'static str)
}

fn read_char<T: Stream>(io: &mut BufferedStream<T>) -> Option<char> {
  match io.read_byte() {
    Some(ch) => Some(ch as char),
    None     => None
  }
}

fn parse_data<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let res =
    if (len > 0) {
      let bytes = io.read_bytes(len);
      if bytes.len() != len {
        return ProtocolError("Invalid number of bytes")
      } else {
        Data(bytes)
      }
    } else {
      Data(~[])
    };

  if (read_char(io) != Some('\r')) {
    return ProtocolError("Carriage return expected"); // TODO: ignore
  }

  if (read_char(io) != Some('\n')) {
    return ProtocolError("Newline expected");
  }

  return res;
}

fn parse_list<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let mut list: ~[Result] = vec::with_capacity(len);

  for _ in range(0, len) {
    match parse(io) {
      ProtocolError(err) => {
        return ProtocolError(err);
      }
      other => {
        list.push(other);
      }
    }
  }

  List(list)
}

fn parse_int_line<T: Stream>(io: &mut BufferedStream<T>) -> Option<i64> {
  let mut i: i64 = 0;
  let mut digits: uint = 0;
  let mut negative: bool = false;

  loop {
    match read_char(io) {
      None => { return None }
      Some(ch) => {
        match ch {
          '0' .. '9' => {
            digits += 1;
            i = (i * 10) + (ch as i64 - '0' as i64);
          }
          '-' => {
            if negative { return None }
            negative = true
          }
          '\r' => {
            if read_char(io) != Some('\n') { return None }
            break
          }
          '\n' => { break }
          _ => { return None }
        }
      }
    }
  }

  if digits == 0 { return None }

  if negative { Some(-i) }
  else { Some(i) }
}

fn parse_n<T: Stream>(io: &mut BufferedStream<T>, f: |uint, &mut BufferedStream<T>| -> Result) -> Result {
  match parse_int_line(io) {
    Some(-1) => Nil,
    Some(len) if len >= 0 => f(len as uint, io), // XXX: i64 might be larger than uint
    _ => ProtocolError("Invalid number")
  }
}

fn parse_status<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Status(line),
    None       => ProtocolError("Invalid status line")
  }
}

fn parse_error<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Error(line),
    None       => ProtocolError("Invalid error line")
  }
}

pub fn parse<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match read_char(io) {
    Some('$') => parse_n(io, parse_data),
    Some('*') => parse_n(io, parse_list),
    Some('+') => parse_status(io),
    Some('-') => parse_error(io),
    Some(':') => {
      match parse_int_line(io) {
        Some(i) => Int(i),
        None => ProtocolError("Invalid number")
      }
    },
    Some(_)   => ProtocolError("Invalid character"),
    None      => ProtocolError("Invalid EOF")
  }
}

pub struct CommandWriter {
  buf: ~[u8]
}

impl CommandWriter {
  pub fn new() -> CommandWriter {
    CommandWriter { buf: ~[] }
  }

  pub fn args<'a>(&'a mut self, n: uint) -> &'a mut CommandWriter {
    self.write_char('*');
    self.write_uint(n);
    self.write_crnl();
    self
  }

  pub fn arg_bin<'a>(&'a mut self, arg: &[u8]) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write(arg);
    self.write_crnl();
    self
  }

  pub fn nil(&mut self) {
    self.write_str("$-1");
    self.write_crnl();
  }

  pub fn arg_str<'a>(&'a mut self, arg: &str) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write_str(arg);
    self.write_crnl();
    self
  }

  pub fn error(&mut self, err: &str) {
    self.write_char('-');
    self.write_str(err);
    self.write_crnl();
  }

  pub fn status(&mut self, status: &str) {
    self.write_char('+');
    self.write_str(status);
    self.write_crnl();
  }

  fn write_crnl(&mut self) {
    self.write_char('\r');
    self.write_char('\n');
  }

  fn write_uint(&mut self, n: uint) {
    if n < 10 {
      self.write_byte('0' as u8 + (n as u8));
    }
    else {
      push_bytes(&mut self.buf, n.to_str().into_bytes()); // XXX: Optimize
    }
  }

  fn write_str(&mut self, s: &str) {
    push_bytes(&mut self.buf, s.as_bytes());
  }

  fn write(&mut self, s: &[u8]) {
    push_bytes(&mut self.buf, s);
  }

  fn write_char(&mut self, s: char) {
    self.buf.push(s as u8);
  }

  fn write_byte(&mut self, b: u8) {
    self.buf.push(b);
  }

  pub fn with_buf<T>(&self, f: |&[u8]| -> T) -> T {
    f(self.buf.as_slice())
  }
}

fn execute<T: Stream>(cmd: &[u8], io: &mut BufferedStream<T>) -> Result {
  io.write(cmd);
  io.flush();
  parse(io)
}

pub struct Client<T> {
  priv io: BufferedStream<T>
}

impl Client<TcpStream> {
  pub fn new(sock_addr: &str) -> Client<TcpStream> {
    let addr = from_str::<SocketAddr>(sock_addr).unwrap();
    let tcp_stream = TcpStream::connect(addr).unwrap();
    Client::new_from_stream(tcp_stream)
  }
}

impl<T: Stream> Client<T> {
  pub fn new_from_stream(io: T) -> Client<T> {
    Client { io: BufferedStream::new(io) }
  }

  pub fn get(&mut self, key: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(2).
        arg_str("GET").
        arg_str(key).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }

  pub fn get_str(&mut self, key: &str) -> Option<~str> {
    match self.get(key) {
      Nil => None,
      Int(i) => Some(i.to_str()),
      Data(ref bytes) => Some(from_utf8(*bytes).to_owned()),
      _ => fail!("Invalid result type from Redis")
    }
  }

  pub fn get_int(&mut self, key: &str) -> Option<i64> {
    match self.get(key) {
      Nil => None,
      Data(ref bytes) => from_str(from_utf8(*bytes)), // XXX
      _ => fail!("Invalid result type from Redis")
    }
  }

  pub fn set(&mut self, key: &str, val: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(3).
        arg_str("SET").
        arg_str(key).
        arg_str(val).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }

  pub fn set_int(&mut self, key: &str, val: i64) -> Result {
    self.set(key, val.to_str())
  }

  pub fn incr(&mut self, key: &str) -> i64 {
    let mut cwr = CommandWriter::new();
    let res = cwr.args(2).
        arg_str("INCR").
        arg_str(key).
        with_buf(|cmd| execute(cmd, &mut self.io));
    match res {
      Int(i) => i,
      _ => fail!()
    }
  }
}
<commit_msg>Update for 0.10-pre<commit_after>use std::io::net::ip::SocketAddr;
use std::io::net::tcp::TcpStream;
use std::io::{BufferedStream, Stream};
use std::vec::bytes::push_bytes;
use std::vec;
use std::str::from_utf8;

pub enum Result {
  Nil,
  Int(i64),
  Data(~[u8]),
  List(~[Result]),
  Error(~str),
  Status(~str),
  ProtocolError(&'static str)
}

fn read_char<T: Stream>(io: &mut BufferedStream<T>) -> Option<char> {
  match io.read_byte() {
    Some(ch) => Some(ch as char),
    None     => None
  }
}

fn parse_data<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let res =
    if (len > 0) {
      let bytes = io.read_bytes(len);
      if bytes.len() != len {
        return ProtocolError("Invalid number of bytes")
      } else {
        Data(bytes)
      }
    } else {
      Data(~[])
    };

  if (read_char(io) != Some('\r')) {
    return ProtocolError("Carriage return expected"); // TODO: ignore
  }

  if (read_char(io) != Some('\n')) {
    return ProtocolError("Newline expected");
  }

  return res;
}

fn parse_list<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let mut list: ~[Result] = vec::with_capacity(len);

  for _ in range(0, len) {
    match parse(io) {
      ProtocolError(err) => {
        return ProtocolError(err);
      }
      other => {
        list.push(other);
      }
    }
  }

  List(list)
}

fn parse_int_line<T: Stream>(io: &mut BufferedStream<T>) -> Option<i64> {
  let mut i: i64 = 0;
  let mut digits: uint = 0;
  let mut negative: bool = false;

  loop {
    match read_char(io) {
      None => { return None }
      Some(ch) => {
        match ch {
          '0' .. '9' => {
            digits += 1;
            i = (i * 10) + (ch as i64 - '0' as i64);
          }
          '-' => {
            if negative { return None }
            negative = true
          }
          '\r' => {
            if read_char(io) != Some('\n') { return None }
            break
          }
          '\n' => { break }
          _ => { return None }
        }
      }
    }
  }

  if digits == 0 { return None }

  if negative { Some(-i) }
  else { Some(i) }
}

fn parse_n<T: Stream>(io: &mut BufferedStream<T>, f: |uint, &mut BufferedStream<T>| -> Result) -> Result {
  match parse_int_line(io) {
    Some(-1) => Nil,
    Some(len) if len >= 0 => f(len as uint, io), // XXX: i64 might be larger than uint
    _ => ProtocolError("Invalid number")
  }
}

fn parse_status<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Status(line),
    None       => ProtocolError("Invalid status line")
  }
}

fn parse_error<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Error(line),
    None       => ProtocolError("Invalid error line")
  }
}

pub fn parse<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match read_char(io) {
    Some('$') => parse_n(io, parse_data),
    Some('*') => parse_n(io, parse_list),
    Some('+') => parse_status(io),
    Some('-') => parse_error(io),
    Some(':') => {
      match parse_int_line(io) {
        Some(i) => Int(i),
        None => ProtocolError("Invalid number")
      }
    },
    Some(_)   => ProtocolError("Invalid character"),
    None      => ProtocolError("Invalid EOF")
  }
}

pub struct CommandWriter {
  buf: ~[u8]
}

impl CommandWriter {
  pub fn new() -> CommandWriter {
    CommandWriter { buf: ~[] }
  }

  pub fn args<'a>(&'a mut self, n: uint) -> &'a mut CommandWriter {
    self.write_char('*');
    self.write_uint(n);
    self.write_crnl();
    self
  }

  pub fn arg_bin<'a>(&'a mut self, arg: &[u8]) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write(arg);
    self.write_crnl();
    self
  }

  pub fn nil(&mut self) {
    self.write_str("$-1");
    self.write_crnl();
  }

  pub fn arg_str<'a>(&'a mut self, arg: &str) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write_str(arg);
    self.write_crnl();
    self
  }

  pub fn error(&mut self, err: &str) {
    self.write_char('-');
    self.write_str(err);
    self.write_crnl();
  }

  pub fn status(&mut self, status: &str) {
    self.write_char('+');
    self.write_str(status);
    self.write_crnl();
  }

  fn write_crnl(&mut self) {
    self.write_char('\r');
    self.write_char('\n');
  }

  fn write_uint(&mut self, n: uint) {
    if n < 10 {
      self.write_byte('0' as u8 + (n as u8));
    }
    else {
      push_bytes(&mut self.buf, n.to_str().into_bytes()); // XXX: Optimize
    }
  }

  fn write_str(&mut self, s: &str) {
    push_bytes(&mut self.buf, s.as_bytes());
  }

  fn write(&mut self, s: &[u8]) {
    push_bytes(&mut self.buf, s);
  }

  fn write_char(&mut self, s: char) {
    self.buf.push(s as u8);
  }

  fn write_byte(&mut self, b: u8) {
    self.buf.push(b);
  }

  pub fn with_buf<T>(&self, f: |&[u8]| -> T) -> T {
    f(self.buf.as_slice())
  }
}

fn execute<T: Stream>(cmd: &[u8], io: &mut BufferedStream<T>) -> Result {
  io.write(cmd);
  io.flush();
  parse(io)
}

pub struct Client<T> {
  priv io: BufferedStream<T>
}

impl Client<TcpStream> {
  pub fn new(sock_addr: &str) -> Client<TcpStream> {
    let addr = from_str::<SocketAddr>(sock_addr).unwrap();
    let tcp_stream = TcpStream::connect(addr).unwrap();
    Client::new_from_stream(tcp_stream)
  }
}

impl<T: Stream> Client<T> {
  pub fn new_from_stream(io: T) -> Client<T> {
    Client { io: BufferedStream::new(io) }
  }

  pub fn get(&mut self, key: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(2).
        arg_str("GET").
        arg_str(key).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }

  pub fn get_str(&mut self, key: &str) -> Option<~str> {
    match self.get(key) {
      Nil => None,
      Int(i) => Some(i.to_str()),
      Data(ref bytes) => Some(from_utf8(*bytes).unwrap().to_owned()),
      _ => fail!("Invalid result type from Redis")
    }
  }

  pub fn get_int(&mut self, key: &str) -> Option<i64> {
    match self.get(key) {
      Nil => None,
      Data(ref bytes) => from_str(from_utf8(*bytes).unwrap()), // XXX
      _ => fail!("Invalid result type from Redis")
    }
  }

  pub fn set(&mut self, key: &str, val: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(3).
        arg_str("SET").
        arg_str(key).
        arg_str(val).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }

  pub fn set_int(&mut self, key: &str, val: i64) -> Result {
    self.set(key, val.to_str())
  }

  pub fn incr(&mut self, key: &str) -> i64 {
    let mut cwr = CommandWriter::new();
    let res = cwr.args(2).
        arg_str("INCR").
        arg_str(key).
        with_buf(|cmd| execute(cmd, &mut self.io));
    match res {
      Int(i) => i,
      _ => fail!()
    }
  }
}
<|endoftext|>