<commit_before>// Copyright 2015 The tiny-http Contributors
// Copyright 2015 The rust-chunked-transfer Contributors
//
// 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.

use std::error::Error;
use std::fmt;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::io::Read;
use std::io::Result as IoResult;

/// Reads HTTP chunks and sends back real data.
///
/// # Example
///
/// ```
/// use chunked_transfer::Decoder;
/// use std::io::Read;
///
/// let encoded = b"3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n";
/// let mut decoded = String::new();
///
/// let mut decoder = Decoder::new(encoded as &[u8]);
/// decoder.read_to_string(&mut decoded);
///
/// assert_eq!(decoded, "hello world!!!");
/// ```
pub struct Decoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,
}

impl<R> Decoder<R>
{
    pub fn new(source: R) -> Decoder<R> {
        Decoder {
            source,
            remaining_chunks_size: None,
        }
    }

    /// Unwraps the Decoder into its inner `Read` source.
    pub fn into_inner(self) -> R {
        self.source
    }
}

impl<R> Decoder<R>
where
    R: Read,
{
    fn read_chunk_size(&mut self) -> IoResult<usize> {
        let mut chunk_size_bytes = Vec::new();
        let mut has_ext = false;

        loop {
            let byte = match self.source.by_ref().bytes().next() {
                Some(b) => b?,
                None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            };

            if byte == b'\r' {
                break;
            }

            if byte == b';' {
                has_ext = true;
                break;
            }

            chunk_size_bytes.push(byte);
        }

        // Ignore extensions for now
        if has_ext {
            loop {
                let byte = match self.source.by_ref().bytes().next() {
                    Some(b) => b?,
                    None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                };
                if byte == b'\r' {
                    break;
                }
            }
        }

        self.read_line_feed()?;

        let chunk_size = String::from_utf8(chunk_size_bytes)
            .ok()
            .and_then(|c| usize::from_str_radix(c.trim(), 16).ok())
            .ok_or(IoError::new(ErrorKind::InvalidInput, DecoderError))?;

        Ok(chunk_size)
    }

    fn read_carriage_return(&mut self) -> IoResult<()> {
        match self.source.by_ref().bytes().next() {
            Some(Ok(b'\r')) => Ok(()),
            _ => Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
        }
    }

    fn read_line_feed(&mut self) -> IoResult<()> {
        match self.source.by_ref().bytes().next() {
            Some(Ok(b'\n')) => Ok(()),
            _ => Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
        }
    }
}

impl<R> Read for Decoder<R>
where
    R: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let remaining_chunks_size = match self.remaining_chunks_size {
            Some(c) => c,
            None => {
                // first possibility: we are not in a chunk, so we'll attempt to determine
                // the chunks size
                let chunk_size = self.read_chunk_size()?;

                // if the chunk size is 0, we are at EOF
                if chunk_size == 0 {
                    self.read_carriage_return()?;
                    self.read_line_feed()?;
                    return Ok(0);
                }

                chunk_size
            }
        };

        // second possibility: we continue reading from a chunk
        if buf.len() < remaining_chunks_size {
            let read = self.source.read(buf)?;
            self.remaining_chunks_size = Some(remaining_chunks_size - read);
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= remaining_chunks_size);

        let buf = &mut buf[..remaining_chunks_size];
        let read = self.source.read(buf)?;

        self.remaining_chunks_size = if read == remaining_chunks_size {
            self.read_carriage_return()?;
            self.read_line_feed()?;
            None
        } else {
            Some(remaining_chunks_size - read)
        };

        Ok(read)
    }
}

#[derive(Debug, Copy, Clone)]
struct DecoderError;

impl fmt::Display for DecoderError {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for DecoderError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

#[cfg(test)]
mod test {
    use super::Decoder;
    use std::io;
    use std::io::Read;

    /// This unit test is taken from from Hyper
    /// https://github.com/hyperium/hyper
    /// Copyright (c) 2014 Sean McArthur
    #[test]
    fn test_read_chunk_size() {
        fn read(s: &str, expected: usize) {
            let mut decoded = Decoder::new(s.as_bytes());
            let actual = decoded.read_chunk_size().unwrap();
            assert_eq!(expected, actual);
        }

        fn read_err(s: &str) {
            let mut decoded = Decoder::new(s.as_bytes());
            let err_kind = decoded.read_chunk_size().unwrap_err().kind();
            assert_eq!(err_kind, io::ErrorKind::InvalidInput);
        }

        read("1\r\n", 1);
        read("01\r\n", 1);
        read("0\r\n", 0);
        read("00\r\n", 0);
        read("A\r\n", 10);
        read("a\r\n", 10);
        read("Ff\r\n", 255);
        read("Ff   \r\n", 255);
        // Missing LF or CRLF
        read_err("F\rF");
        read_err("F");
        // Invalid hex digit
        read_err("X\r\n");
        read_err("1X\r\n");
        read_err("-\r\n");
        read_err("-1\r\n");
        // Acceptable (if not fully valid) extensions do not influence the size
        read("1;extension\r\n", 1);
        read("a;ext name=value\r\n", 10);
        read("1;extension;extension2\r\n", 1);
        read("1;;;  ;\r\n", 1);
        read("2; extension...\r\n", 2);
        read("3   ; extension=123\r\n", 3);
        read("3   ;\r\n", 3);
        read("3   ;   \r\n", 3);
        // Invalid extensions cause an error
        read_err("1 invalid extension\r\n");
        read_err("1 A\r\n");
        read_err("1;no CRLF");
    }

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new(
            "3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n"
                .to_string()
                .into_bytes(),
        );
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    fn test_decode_zero_length() {
        let mut decoder = Decoder::new(b"0\r\n\r\n" as &[u8]);

        let mut decoded = String::new();
        decoder.read_to_string(&mut decoded).unwrap();

        assert_eq!(decoded, "");
    }

    #[test]
    fn test_decode_invalid_chunk_length() {
        let mut decoder = Decoder::new(b"m\r\n\r\n" as &[u8]);

        let mut decoded = String::new();
        assert!(decoder.read_to_string(&mut decoded).is_err());
    }

    #[test]
    fn invalid_input1() {
        let source = io::Cursor::new(
            "2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n"
                .to_string()
                .into_bytes(),
        );
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }

    #[test]
    fn invalid_input2() {
        let source = io::Cursor::new(
            "3\rhel\r\nb\r\nlo world!!!\r\n0\r\n"
                .to_string()
                .into_bytes(),
        );
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }
}
<commit_msg>Move Decoder methods back under R: Read bound.<commit_after>// Copyright 2015 The tiny-http Contributors
// Copyright 2015 The rust-chunked-transfer Contributors
//
// 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.

use std::error::Error;
use std::fmt;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::io::Read;
use std::io::Result as IoResult;

/// Reads HTTP chunks and sends back real data.
///
/// # Example
///
/// ```
/// use chunked_transfer::Decoder;
/// use std::io::Read;
///
/// let encoded = b"3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n";
/// let mut decoded = String::new();
///
/// let mut decoder = Decoder::new(encoded as &[u8]);
/// decoder.read_to_string(&mut decoded);
///
/// assert_eq!(decoded, "hello world!!!");
/// ```
pub struct Decoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,
}

impl<R> Decoder<R>
where
    R: Read,
{
    pub fn new(source: R) -> Decoder<R> {
        Decoder {
            source,
            remaining_chunks_size: None,
        }
    }

    /// Unwraps the Decoder into its inner `Read` source.
    pub fn into_inner(self) -> R {
        self.source
    }

    fn read_chunk_size(&mut self) -> IoResult<usize> {
        let mut chunk_size_bytes = Vec::new();
        let mut has_ext = false;

        loop {
            let byte = match self.source.by_ref().bytes().next() {
                Some(b) => b?,
                None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            };

            if byte == b'\r' {
                break;
            }

            if byte == b';' {
                has_ext = true;
                break;
            }

            chunk_size_bytes.push(byte);
        }

        // Ignore extensions for now
        if has_ext {
            loop {
                let byte = match self.source.by_ref().bytes().next() {
                    Some(b) => b?,
                    None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                };
                if byte == b'\r' {
                    break;
                }
            }
        }

        self.read_line_feed()?;

        let chunk_size = String::from_utf8(chunk_size_bytes)
            .ok()
            .and_then(|c| usize::from_str_radix(c.trim(), 16).ok())
            .ok_or(IoError::new(ErrorKind::InvalidInput, DecoderError))?;

        Ok(chunk_size)
    }

    fn read_carriage_return(&mut self) -> IoResult<()> {
        match self.source.by_ref().bytes().next() {
            Some(Ok(b'\r')) => Ok(()),
            _ => Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
        }
    }

    fn read_line_feed(&mut self) -> IoResult<()> {
        match self.source.by_ref().bytes().next() {
            Some(Ok(b'\n')) => Ok(()),
            _ => Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
        }
    }
}

impl<R> Read for Decoder<R>
where
    R: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let remaining_chunks_size = match self.remaining_chunks_size {
            Some(c) => c,
            None => {
                // first possibility: we are not in a chunk, so we'll attempt to determine
                // the chunks size
                let chunk_size = self.read_chunk_size()?;

                // if the chunk size is 0, we are at EOF
                if chunk_size == 0 {
                    self.read_carriage_return()?;
                    self.read_line_feed()?;
                    return Ok(0);
                }

                chunk_size
            }
        };

        // second possibility: we continue reading from a chunk
        if buf.len() < remaining_chunks_size {
            let read = self.source.read(buf)?;
            self.remaining_chunks_size = Some(remaining_chunks_size - read);
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= remaining_chunks_size);

        let buf = &mut buf[..remaining_chunks_size];
        let read = self.source.read(buf)?;

        self.remaining_chunks_size = if read == remaining_chunks_size {
            self.read_carriage_return()?;
            self.read_line_feed()?;
            None
        } else {
            Some(remaining_chunks_size - read)
        };

        Ok(read)
    }
}

#[derive(Debug, Copy, Clone)]
struct DecoderError;

impl fmt::Display for DecoderError {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for DecoderError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

#[cfg(test)]
mod test {
    use super::Decoder;
    use std::io;
    use std::io::Read;

    /// This unit test is taken from from Hyper
    /// https://github.com/hyperium/hyper
    /// Copyright (c) 2014 Sean McArthur
    #[test]
    fn test_read_chunk_size() {
        fn read(s: &str, expected: usize) {
            let mut decoded = Decoder::new(s.as_bytes());
            let actual = decoded.read_chunk_size().unwrap();
            assert_eq!(expected, actual);
        }

        fn read_err(s: &str) {
            let mut decoded = Decoder::new(s.as_bytes());
            let err_kind = decoded.read_chunk_size().unwrap_err().kind();
            assert_eq!(err_kind, io::ErrorKind::InvalidInput);
        }

        read("1\r\n", 1);
        read("01\r\n", 1);
        read("0\r\n", 0);
        read("00\r\n", 0);
        read("A\r\n", 10);
        read("a\r\n", 10);
        read("Ff\r\n", 255);
        read("Ff   \r\n", 255);
        // Missing LF or CRLF
        read_err("F\rF");
        read_err("F");
        // Invalid hex digit
        read_err("X\r\n");
        read_err("1X\r\n");
        read_err("-\r\n");
        read_err("-1\r\n");
        // Acceptable (if not fully valid) extensions do not influence the size
        read("1;extension\r\n", 1);
        read("a;ext name=value\r\n", 10);
        read("1;extension;extension2\r\n", 1);
        read("1;;;  ;\r\n", 1);
        read("2; extension...\r\n", 2);
        read("3   ; extension=123\r\n", 3);
        read("3   ;\r\n", 3);
        read("3   ;   \r\n", 3);
        // Invalid extensions cause an error
        read_err("1 invalid extension\r\n");
        read_err("1 A\r\n");
        read_err("1;no CRLF");
    }

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new(
            "3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n"
                .to_string()
                .into_bytes(),
        );
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    fn test_decode_zero_length() {
        let mut decoder = Decoder::new(b"0\r\n\r\n" as &[u8]);

        let mut decoded = String::new();
        decoder.read_to_string(&mut decoded).unwrap();

        assert_eq!(decoded, "");
    }

    #[test]
    fn test_decode_invalid_chunk_length() {
        let mut decoder = Decoder::new(b"m\r\n\r\n" as &[u8]);

        let mut decoded = String::new();
        assert!(decoder.read_to_string(&mut decoded).is_err());
    }

    #[test]
    fn invalid_input1() {
        let source = io::Cursor::new(
            "2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n"
                .to_string()
                .into_bytes(),
        );
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }

    #[test]
    fn invalid_input2() {
        let source = io::Cursor::new(
            "3\rhel\r\nb\r\nlo world!!!\r\n0\r\n"
                .to_string()
                .into_bytes(),
        );
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2016 Antoine Wacheux

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
associated documentation files (the "Software"), to deal in the Software without restriction,
including without limitation the rights to use, copy, modify, merge, publish, distribute,
sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or
substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT
OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

pub struct LexicalUnit {
    name : String,
    regex : String,
}

pub struct ProductionRule {
    name: String,
    rule: Vec<RuleComponent>
}

pub enum RuleComponent {
    ProductionRule(String),
    LexicalUnit(String)
}

pub struct Grammar {
    lexical_units : Vec<LexicalUnit>,
    production_rules: Vec<ProductionRule>,
    axiom : String
}

impl Grammar {
    fn new() -> Grammar {
        Grammar {
            lexical_units: Vec::new(),
            production_rules: Vec::new(),
            axiom: String::new(),
        }
    }
}

//Methods
//get_rule(name: String&) -> Vec<ProductionRule>&

struct ParserContext {
    src: String,
    row: u32,
    col: u32,
    str_pos: usize,
}

impl ParserContext {
    fn new(src : String) -> ParserContext {
        ParserContext {
            src: src,
            row: 1,
            col: 0,
            str_pos: 0
        }
    }

    fn next_word(&mut self) -> Option<String> {
        let mut iter = self.src.chars();
        let mut ch = iter.nth(self.str_pos);
        let mut c : char;
        let mut read_word = String::new();

        match ch {
            None => { return None },
            Some(chr) => { c = chr}
        }

        macro_rules! next_char {
            ($err:stmt) => ({
                self.str_pos += 1;
                ch = iter.next();

                match ch {
                    None => { $err },
                    Some(chr) => { c = chr }
                }

                match c { //TODO : Implement a better line/column tracking and forward this information in the word
                    ' ' | '\t' => self.col += 1,
                    '\n' => self.row += 1,
                    _ => {}
                }
            });
        }

        while (c == ' ') || (c == '\n') || (c == '\t') { //We skip leading whitespaces
            next_char! { return None }
        }

        while !((c == ' ') || (c == '\n') || (c == '\t')) {
            read_word.push(c);
            next_char! { return Some(read_word) }
        }

        Some(read_word)
    }

    pub fn read_line(&mut self) -> Option<String> {
        let mut iter = self.src.chars();
        let mut ch = iter.nth(self.str_pos);
        let mut c : char;
        let mut read_word = String::new();

        match ch {
            None => { return None },
            Some(chr) => { c = chr}
        }

        macro_rules! next_char {
            ($err:stmt) => ({
                self.str_pos += 1;
                ch = iter.next();

                match ch {
                    None => { $err },
                    Some(chr) => { c = chr }
                }

                match c {
                    ' ' | '\t' => self.col += 1,
                    '\n' => self.row += 1,
                    _ => {}
                }
            });
        }

        while (c == ' ') || (c == '\n') || (c == '\t') { //We skip leading whitespaces
            next_char! { return None }
        }

        while c != '\n' {
            read_word.push(c);
            next_char! { return Some(read_word) }
        }

        Some(read_word)
    }
}

struct GrammarParser {
    ctx: ParserContext,
    grammar: Grammar
}

impl GrammarParser {
    fn new(src: String) -> GrammarParser {
        GrammarParser {
            ctx : ParserContext::new(src),
            grammar: Grammar::new(),
        }
    }

    fn get_resulting_grammar(self) -> Grammar {
        self.grammar
    }

    fn parse(&mut self) -> Result<(), String> {
        while let Some(str_word) = self.ctx.next_word() {
            match str_word.chars().next().unwrap() {
                'A'...'Z' => {
                    match self.parse_grammar_rule(str_word) {
                        Ok(()) => {},
                        Err(msg) => {return Err(msg)},
                    }
                },
                'a'...'z' => {
                    match self.parse_lexical_unit(str_word) {
                        Ok(()) => {},
                        Err(msg) => {return Err(msg)},
                    }
                }
                _ => { return Err(String::from("Unknown token : ") + str_word.as_str())}
            }
        }

        Ok(())
    }

    fn parse_grammar_rule(&mut self, name : String) -> Result<(), String> {
        match self.ctx.next_word() {
            None => {return Err(format!("line {} : Expected '->'", self.ctx.row))},
            Some(token) => {
                if token != "->" {
                    return Err(format!("line {} : Expected '->' found '{}'", self.ctx.row, token));
                }
            }
        }

        if let Some(rule_str) = self.ctx.read_line() {
            let mut production_rule = ProductionRule{
                name: name,
                rule: Vec::new()
            };

            for comp_str in rule_str.split_whitespace() {
                let component_str = String::from(comp_str);
                let component = match component_str.chars().next().unwrap() {
                    'a'...'z' => {RuleComponent::LexicalUnit(component_str.to_owned())},
                    'A'...'Z' => {RuleComponent::ProductionRule(component_str.to_owned())},
                    '\'' => {
                        let s = component_str.len();
                        if s < 3 {
                            return Err(format!("line {} : The regex in a production rule cannot be empty.", self.ctx.row));
                        }

                        let regex = unsafe {component_str.slice_unchecked(1, s-1) }; //TODO : Remove this unsafe
                        let name = format!("auto_{}", regex);

                        if self.grammar.lexical_units.iter().find(|rule| rule.name == name).is_none() {
                            self.grammar.lexical_units.push(LexicalUnit{name:name.clone(), regex:String::from(regex)});
                        }

                        RuleComponent::LexicalUnit(name)
                    },
                    _ => {return Err(format!("line {}: '{}' is unexpected", self.ctx.row, component_str))}
                };

                production_rule.rule.push(component);
            }
            self.grammar.production_rules.push(production_rule);
            Ok(())
        }
        else {
            Err(format!("A production rule cannot be empty"))
        }
    }

    fn parse_lexical_unit(&mut self, name : String) -> Result<(), String> {
        //Let's read the ':'
        match self.ctx.next_word() {
            None => {return Err(format!("line {} : Expected ':'", self.ctx.row))},
            Some(token) => {
                if token != ":" {
                    return Err(format!("line {} : Expected ':' found '{}'", self.ctx.row, token));
                }
            }
        }

        //Ok, now the regex
        if let Some(regex) = self.ctx.read_line() {
            let lexical_rule = LexicalUnit {
                name: name,
                regex: regex,
            };

            if self.grammar.lexical_units.iter().find(|rule| rule.name == lexical_rule.name).is_some() {
                return Err(format!("line {} : the lexical unit {} has already been declared", self.ctx.row, lexical_rule.name));
            }

            //TODO : Implement a check by regex too
            //TODO : Check that the regex is syntacticly corret (https://crates.io/crates/regex-syntax)
            self.grammar.lexical_units.push(lexical_rule);
            Ok(())
        }
        else {
            Err(format!("line {} : The regex is expected after ':' but nothing was found", self.ctx.row))
        }
    }

    fn check_consistency(& self) -> Result<(), String> {
        for production_rule in &self.grammar.production_rules {
            for rule_component in &production_rule.rule {
                match rule_component {
                    &RuleComponent::ProductionRule(ref name) => {
                        if self.grammar.production_rules.iter().find(|r| *name == r.name).is_none() {
                            return Err(format!("The production rule '{}' does not exist", name));
                        }
                    },
                    &RuleComponent::LexicalUnit(ref name) => {
                        if self.grammar.lexical_units.iter().find(|lu| *name  == lu.name).is_none() {
                            return Err(format!("The lexical unit '{}' does not exist", name));
                        }
                    },
                }
            }
        }

        Ok(())
    }
}

pub fn parse_grammar(src : String) -> Result<Grammar, String> {
    let mut parser = GrammarParser::new(src);

    if let Err(msg) = parser.parse() {
        return Err(msg);
    }

    if let Err(msg) = parser.check_consistency() {
        return Err(msg);
    }

    Ok(parser.get_resulting_grammar())
}

#[test]
fn test_next_word() {
    let mut ctx = ParserContext::new("test and run".to_owned());

    assert_eq!(ctx.next_word().unwrap(), "test");
    assert_eq!(ctx.next_word().unwrap(), "and");
    assert_eq!(ctx.next_word().unwrap(), "run");
}

#[test]
fn test_read_line() {
    let mut ctx = ParserContext::new("line 1\nline2\nline3\nline4".to_owned());
    ctx.next_word();

    assert_eq!(ctx.read_line().unwrap(), "1");
    assert_eq!(ctx.read_line().unwrap(), "line2");
    assert_eq!(ctx.read_line().unwrap(), "line3");
    assert_eq!(ctx.read_line().unwrap(), "line4");
}
<commit_msg>Avoid left recursive rules<commit_after>/*
Copyright (c) 2016 Antoine Wacheux

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
associated documentation files (the "Software"), to deal in the Software without restriction,
including without limitation the rights to use, copy, modify, merge, publish, distribute,
sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or
substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT
OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

pub struct LexicalUnit {
    name : String,
    regex : String,
}

pub struct ProductionRule {
    name: String,
    rule: Vec<RuleComponent>
}

pub enum RuleComponent {
    ProductionRule(String),
    LexicalUnit(String)
}

pub struct Grammar {
    lexical_units : Vec<LexicalUnit>,
    production_rules: Vec<ProductionRule>,
    axiom : String
}

impl Grammar {
    fn new() -> Grammar {
        Grammar {
            lexical_units: Vec::new(),
            production_rules: Vec::new(),
            axiom: String::new(),
        }
    }

    fn get_production_rule(&self, name: &String) -> Vec<&ProductionRule> {
        self.production_rules.iter().filter(|x| x.name == *name).collect()
    }
}

//Methods
//get_rule(name: String&) -> Vec<ProductionRule>&

struct ParserContext {
    src: String,
    row: u32,
    col: u32,
    str_pos: usize,
}

impl ParserContext {
    fn new(src : String) -> ParserContext {
        ParserContext {
            src: src,
            row: 1,
            col: 0,
            str_pos: 0
        }
    }

    fn next_word(&mut self) -> Option<String> {
        let mut iter = self.src.chars();
        let mut ch = iter.nth(self.str_pos);
        let mut c : char;
        let mut read_word = String::new();

        match ch {
            None => { return None },
            Some(chr) => { c = chr}
        }

        macro_rules! next_char {
            ($err:stmt) => ({
                self.str_pos += 1;
                ch = iter.next();

                match ch {
                    None => { $err },
                    Some(chr) => { c = chr }
                }

                match c { //TODO : Implement a better line/column tracking and forward this information in the word
                    ' ' | '\t' => self.col += 1,
                    '\n' => self.row += 1,
                    _ => {}
                }
            });
        }

        while (c == ' ') || (c == '\n') || (c == '\t') { //We skip leading whitespaces
            next_char! { return None }
        }

        while !((c == ' ') || (c == '\n') || (c == '\t')) {
            read_word.push(c);
            next_char! { return Some(read_word) }
        }

        Some(read_word)
    }

    pub fn read_line(&mut self) -> Option<String> {
        let mut iter = self.src.chars();
        let mut ch = iter.nth(self.str_pos);
        let mut c : char;
        let mut read_word = String::new();

        match ch {
            None => { return None },
            Some(chr) => { c = chr}
        }

        macro_rules! next_char {
            ($err:stmt) => ({
                self.str_pos += 1;
                ch = iter.next();

                match ch {
                    None => { $err },
                    Some(chr) => { c = chr }
                }

                match c {
                    ' ' | '\t' => self.col += 1,
                    '\n' => self.row += 1,
                    _ => {}
                }
            });
        }

        while (c == ' ') || (c == '\n') || (c == '\t') { //We skip leading whitespaces
            next_char! { return None }
        }

        while c != '\n' {
            read_word.push(c);
            next_char! { return Some(read_word) }
        }

        Some(read_word)
    }
}

struct GrammarParser {
    ctx: ParserContext,
    grammar: Grammar
}

impl GrammarParser {
    fn new(src: String) -> GrammarParser {
        GrammarParser {
            ctx : ParserContext::new(src),
            grammar: Grammar::new(),
        }
    }

    fn get_resulting_grammar(self) -> Grammar {
        self.grammar
    }

    fn parse(&mut self) -> Result<(), String> {
        while let Some(str_word) = self.ctx.next_word() {
            match str_word.chars().next().unwrap() {
                'A'...'Z' => {
                    match self.parse_grammar_rule(str_word) {
                        Ok(()) => {},
                        Err(msg) => {return Err(msg)},
                    }
                },
                'a'...'z' => {
                    match self.parse_lexical_unit(str_word) {
                        Ok(()) => {},
                        Err(msg) => {return Err(msg)},
                    }
                }
                _ => { return Err(String::from("Unknown token : ") + str_word.as_str())}
            }
        }

        Ok(())
    }

    fn parse_grammar_rule(&mut self, name : String) -> Result<(), String> {
        match self.ctx.next_word() {
            None => {return Err(format!("line {} : Expected '->'", self.ctx.row))},
            Some(token) => {
                if token != "->" {
                    return Err(format!("line {} : Expected '->' found '{}'", self.ctx.row, token));
                }
            }
        }

        if let Some(rule_str) = self.ctx.read_line() {
            let mut production_rule = ProductionRule{
                name: name,
                rule: Vec::new()
            };

            for comp_str in rule_str.split_whitespace() {
                let component_str = String::from(comp_str);
                let component = match component_str.chars().next().unwrap() {
                    'a'...'z' => {RuleComponent::LexicalUnit(component_str.to_owned())},
                    'A'...'Z' => {RuleComponent::ProductionRule(component_str.to_owned())},
                    '\'' => {
                        let s = component_str.len();
                        if s < 3 {
                            return Err(format!("line {} : The regex in a production rule cannot be empty.", self.ctx.row));
                        }

                        let regex = unsafe {component_str.slice_unchecked(1, s-1) }; //TODO : Remove this unsafe
                        let name = format!("auto_{}", regex);

                        if self.grammar.lexical_units.iter().find(|rule| rule.name == name).is_none() {
                            self.grammar.lexical_units.push(LexicalUnit{name:name.clone(), regex:String::from(regex)});
                        }

                        RuleComponent::LexicalUnit(name)
                    },
                    _ => {return Err(format!("line {}: '{}' is unexpected", self.ctx.row, component_str))}
                };

                production_rule.rule.push(component);
            }
            self.grammar.production_rules.push(production_rule);
            Ok(())
        }
        else {
            Err(format!("A production rule cannot be empty"))
        }
    }

    fn parse_lexical_unit(&mut self, name : String) -> Result<(), String> {
        //Let's read the ':'
        match self.ctx.next_word() {
            None => {return Err(format!("line {} : Expected ':'", self.ctx.row))},
            Some(token) => {
                if token != ":" {
                    return Err(format!("line {} : Expected ':' found '{}'", self.ctx.row, token));
                }
            }
        }

        //Ok, now the regex
        if let Some(regex) = self.ctx.read_line() {
            let lexical_rule = LexicalUnit {
                name: name,
                regex: regex,
            };

            if self.grammar.lexical_units.iter().find(|rule| rule.name == lexical_rule.name).is_some() {
                return Err(format!("line {} : the lexical unit {} has already been declared", self.ctx.row, lexical_rule.name));
            }

            //TODO : Implement a check by regex too
            //TODO : Check that the regex is syntacticly corret (https://crates.io/crates/regex-syntax)
            self.grammar.lexical_units.push(lexical_rule);
            Ok(())
        }
        else {
            Err(format!("line {} : The regex is expected after ':' but nothing was found", self.ctx.row))
        }
    }

    fn check_consistency(& self) -> Result<(), String> {
        for production_rule in &self.grammar.production_rules {
            for rule_component in &production_rule.rule {
                match rule_component {
                    &RuleComponent::ProductionRule(ref name) => {
                        if self.grammar.production_rules.iter().find(|r| *name == r.name).is_none() {
                            return Err(format!("The production rule '{}' does not exist", name));
                        }
                    },
                    &RuleComponent::LexicalUnit(ref name) => {
                        if self.grammar.lexical_units.iter().find(|lu| *name  == lu.name).is_none() {
                            return Err(format!("The lexical unit '{}' does not exist", name));
                        }
                    },
                }
            }
        }

        Ok(())
    }

    fn check_for_left_recursion(& self) -> Result<(), String> {
        self.grammar.production_rules.iter()
            .map(|p| self.begin_with_production_rule(&p.name, &p.name))
            .fold(Ok(()), |acc, x| acc.and(x))
    }

    fn begin_with_production_rule(& self, name: &String, production_to_look_in: &String) -> Result<(), String> {
        for p in self.grammar.get_production_rule(production_to_look_in) {
            let first_component = &p.rule.iter().next().unwrap();
            if let &RuleComponent::ProductionRule(ref p_name) = *first_component {
                if name == p_name {
                    return Err(format!("Left recursion for the production rule {}", name));
                }
                else {
                    return self.begin_with_production_rule(name, p_name);
                }
            }
        }

        Ok(())
    }
}

pub fn parse_grammar(src : String) -> Result<Grammar, String> {
    let mut parser = GrammarParser::new(src);

    if let Err(msg) = parser.parse() {
        return Err(msg);
    }

    if let Err(msg) = parser.check_consistency() {
        return Err(msg);
    }

    if let Err(msg) = parser.check_for_left_recursion() {
        return Err(msg);
    }

    Ok(parser.get_resulting_grammar())
}

#[test]
fn test_next_word() {
    let mut ctx = ParserContext::new("test and run".to_owned());

    assert_eq!(ctx.next_word().unwrap(), "test");
    assert_eq!(ctx.next_word().unwrap(), "and");
    assert_eq!(ctx.next_word().unwrap(), "run");
}

#[test]
fn test_read_line() {
    let mut ctx = ParserContext::new("line 1\nline2\nline3\nline4".to_owned());
    ctx.next_word();

    assert_eq!(ctx.read_line().unwrap(), "1");
    assert_eq!(ctx.read_line().unwrap(), "line2");
    assert_eq!(ctx.read_line().unwrap(), "line3");
    assert_eq!(ctx.read_line().unwrap(), "line4");
}
<|endoftext|>


<commit_before>//! Defines a generic 2D image type.

use ndarray;
use ndarray::prelude::*;
use num_traits::{Zero};

use std::iter::{IntoIterator};

use errors::*;
use traits::Pixel;

/// 2-dimensional image type.
pub struct Image2D

    where P: Pixel
{
    buffer: Array2

}

/// Type of immutable `Image2D` views.
pub type Image2DView<'a, P> = ndarray::ArrayView<'a, P, Ix2>;

impl<'a, P> Image2D

    where P: Pixel
{
    /// Create a new image of specified dimensions from a `Vec`.
    ///
    /// **Error**: `InvalidDimensions` if the dimensions do not match the length of `v`.
    pub fn from_vec(w: usize, h: usize, v: Vec
) -> Result<Image2D
> {
        let buf = try!(Array2::from_shape_vec((w, h), v));
        Ok(Image2D{ buffer: buf })
    }

    /// Return the pixel at the specified coordinates.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn get_pixel(&self, x: usize, y: usize) -> P {
        self.buffer[[x, y]].clone()
    }

    /// Set the pixel at the specified coordinates to the specified value.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn put_pixel(&mut self, x: usize, y: usize, pixel: P) {
        self.buffer[[x, y]] = pixel;
    }

    /// Return the width of the image.
    pub fn width(&self) -> usize { self.buffer.cols() }
    /// Return the height of the image.
    pub fn height(&self) -> usize { self.buffer.rows() }
    /// Return the dimensions of the image as a `(width, height)` tuple.
    pub fn dimensions(&self) -> (usize, usize) { self.buffer.dim() }

    /// Return a view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image(&'a self, x: usize, y: usize, w: usize, h: usize) -> Image2DView<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice(s![x..x + w, y..y + h])
    }
}

impl<'a, P> IntoIterator for &'a Image2D

    where P: Pixel + 'a
{
    type Item = &'a P;
    type IntoIter = ndarray::iter::Iter<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.into_iter()
    }
}

impl<'a, P> IntoIterator for &'a mut Image2D

    where P: Pixel + 'a
{
    type Item = &'a mut P;
    type IntoIter = ndarray::iter::IterMut<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.iter_mut()
    }
}

impl
 Image2D

    where P: Pixel + Zero
{
    /// Create a new image of specified dimensions filled with zeros.
    pub fn zeros(width: usize, height: usize) -> Image2D
 {
        Image2D { buffer: Array2::zeros((width, height)) }
    }
}

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

    use std::iter::FromIterator;
    use std::fmt::Debug;

    #[test]
    fn test_from_vec() {
        let v1 = Vec::from_iter((1u8..10u8).map(|n| Luma::new([n])));
        let v2 = Vec::from_iter((1u8..7u8).map(|n| Luma::new([n])));

        assert!(Image2D::from_vec(3, 3, v1.clone()).is_ok());
        assert!(Image2D::from_vec(2, 3, v2.clone()).is_ok());
        assert!(Image2D::from_vec(3, 2, v2.clone()).is_ok());

        assert!(Image2D::from_vec(3, 3, v2.clone()).is_err());
        assert!(Image2D::from_vec(4, 2, v2.clone()).is_err());
    }

    #[test]
    fn test_zeros() {
        fn test_zeros_helper
(w: usize, h: usize)
            where P: Pixel + Zero + Debug
        {
            let img = Image2D::
::zeros(w, h);
            for pixel in &img {
                assert!(pixel.is_zero());
            }
        };
        test_zeros_helper::<Luma<u8>>(100, 100);
        test_zeros_helper::<Luma<f32>>(100, 100);
    }

    #[test]
    fn test_into_iter() {
        let v: Vec<Luma<u8>> = (1..10).map(|i| Luma::from(i)).collect();
        let img = Image2D::from_vec(3, 3, v.clone()).unwrap();

        for (p, i) in img.into_iter().zip(v.into_iter()) {
            assert!(&i == p);
        }
    }

    #[test]
    fn test_sub_image() {
        let v: Vec<Luma<usize>> = (1..10).map(|n| Luma::new([n])).collect();
        let img = Image2D::from_vec(3, 3, v).unwrap();
        let subimg1 = img.sub_image(1, 1, 2, 2);

        fn subimg_vec_eq<'a>(subimg: Image2DView<'a, Luma<usize>>, v: &Vec<usize>) -> bool {
            let img_iter = subimg.into_iter();
            let v_iter   = v.into_iter();
            match img_iter.len() == v_iter.len() {
                true => !img_iter.zip(v_iter).any(|(p, i)| p.data[0] != *i),
                false => false
            }
        }

        let subimg1_vec: Vec<usize> = vec![5, 6, 8, 9];

        assert!(subimg_vec_eq(subimg1, &subimg1_vec));
    }
}
<commit_msg>Switch image dimensions interface from usize to u32<commit_after>//! Defines a generic 2D image type.

use ndarray;
use ndarray::prelude::*;
use num_traits::{Zero};

use std::iter::{IntoIterator};

use errors::*;
use traits::Pixel;

/// 2-dimensional image type.
pub struct Image2D

    where P: Pixel
{
    buffer: Array2

}

/// Type of immutable `Image2D` views.
pub type Image2DView<'a, P> = ndarray::ArrayView<'a, P, Ix2>;

impl<'a, P> Image2D

    where P: Pixel
{
    /// Create a new image of specified dimensions from a `Vec`.
    ///
    /// **Error**: `InvalidDimensions` if the dimensions do not match the length of `v`.
    pub fn from_vec(w: u32, h: u32, v: Vec
) -> Result<Image2D
> {
        let buf = try!(Array2::from_shape_vec((w as usize, h as usize), v));
        Ok(Image2D{ buffer: buf })
    }

    /// Return the pixel at the specified coordinates.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn get_pixel(&self, x: u32, y: u32) -> P {
        self.buffer[[x as usize, y as usize]].clone()
    }

    /// Set the pixel at the specified coordinates to the specified value.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn put_pixel(&mut self, x: u32, y: u32, pixel: P) {
        self.buffer[[x as usize, y as usize]] = pixel;
    }

    /// Return the width of the image.
    pub fn width(&self) -> u32 { self.buffer.cols() as u32 }
    /// Return the height of the image.
    pub fn height(&self) -> u32 { self.buffer.rows() as u32 }
    /// Return the dimensions of the image as a `(width, height)` tuple.
    pub fn dimensions(&self) -> (u32, u32) { (self.width(), self.height()) }

    /// Return a view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image(&'a self, x: u32, y: u32, w: u32, h: u32) -> Image2DView<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice(s![x..x + w, y..y + h])
    }
}

impl<'a, P> IntoIterator for &'a Image2D

    where P: Pixel + 'a
{
    type Item = &'a P;
    type IntoIter = ndarray::iter::Iter<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.into_iter()
    }
}

impl<'a, P> IntoIterator for &'a mut Image2D

    where P: Pixel + 'a
{
    type Item = &'a mut P;
    type IntoIter = ndarray::iter::IterMut<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.iter_mut()
    }
}

impl
 Image2D

    where P: Pixel + Zero
{
    /// Create a new image of specified dimensions filled with zeros.
    pub fn zeros(width: u32, height: u32) -> Image2D
 {
        Image2D { buffer: Array2::zeros((width as usize, height as usize)) }
    }
}

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

    use std::iter::FromIterator;
    use std::fmt::Debug;

    #[test]
    fn test_from_vec() {
        let v1 = Vec::from_iter((1u8..10u8).map(|n| Luma::new([n])));
        let v2 = Vec::from_iter((1u8..7u8).map(|n| Luma::new([n])));

        assert!(Image2D::from_vec(3, 3, v1.clone()).is_ok());
        assert!(Image2D::from_vec(2, 3, v2.clone()).is_ok());
        assert!(Image2D::from_vec(3, 2, v2.clone()).is_ok());

        assert!(Image2D::from_vec(3, 3, v2.clone()).is_err());
        assert!(Image2D::from_vec(4, 2, v2.clone()).is_err());
    }

    #[test]
    fn test_zeros() {
        fn test_zeros_helper
(w: u32, h: u32)
            where P: Pixel + Zero + Debug
        {
            let img = Image2D::
::zeros(w, h);
            for pixel in &img {
                assert!(pixel.is_zero());
            }
        };
        test_zeros_helper::<Luma<u8>>(100, 100);
        test_zeros_helper::<Luma<f32>>(100, 100);
    }

    #[test]
    fn test_into_iter() {
        let v: Vec<Luma<u8>> = (1..10).map(|i| Luma::from(i)).collect();
        let img = Image2D::from_vec(3, 3, v.clone()).unwrap();

        for (p, i) in img.into_iter().zip(v.into_iter()) {
            assert!(&i == p);
        }
    }

    #[test]
    fn test_sub_image() {
        let v: Vec<Luma<usize>> = (1..10).map(|n| Luma::new([n])).collect();
        let img = Image2D::from_vec(3, 3, v).unwrap();
        let subimg1 = img.sub_image(1, 1, 2, 2);

        fn subimg_vec_eq<'a>(subimg: Image2DView<'a, Luma<usize>>, v: &Vec<usize>) -> bool {
            let img_iter = subimg.into_iter();
            let v_iter   = v.into_iter();
            match img_iter.len() == v_iter.len() {
                true => !img_iter.zip(v_iter).any(|(p, i)| p.data[0] != *i),
                false => false
            }
        }

        let subimg1_vec: Vec<usize> = vec![5, 6, 8, 9];

        assert!(subimg_vec_eq(subimg1, &subimg1_vec));
    }
}
<|endoftext|>


<commit_before>//! Defines a generic 2D image type.

use ndarray;
use ndarray::prelude::*;
use num_traits::{Zero};

use std::iter::{IntoIterator};

use errors::*;
use traits::Pixel;

/// 2-dimensional image type.
pub struct Image2D

    where P: Pixel
{
    buffer: Array2

}

/// Type of immutable `Image2D` views.
pub type Image2DView<'a, P> = ndarray::ArrayView<'a, P, Ix2>;
pub type Image2DViewMut<'a, P> = ndarray::ArrayViewMut<'a, P, Ix2>;

impl<'a, P> Image2D

    where P: Pixel
{
    /// Create a new image of specified dimensions filled with zeros.
    pub fn new(width: u32, height: u32) -> Image2D

        where P: Pixel + Zero
    {
        Image2D { buffer: Array2::zeros((width as usize, height as usize)) }
    }

    /// Create a new image of specified dimensions from a `Vec`.
    ///
    /// **Error**: `InvalidDimensions` if the dimensions do not match the length of `v`.
    pub fn from_vec(w: u32, h: u32, v: Vec
) -> Result<Image2D
> {
        let buf = try!(Array2::from_shape_vec((w as usize, h as usize), v));
        Ok(Image2D{ buffer: buf })
    }

    /// Return the pixel at the specified coordinates.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn get_pixel(&self, x: u32, y: u32) -> P {
        self.buffer[[x as usize, y as usize]].clone()
    }

    /// Set the pixel at the specified coordinates to the specified value.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn put_pixel(&mut self, x: u32, y: u32, pixel: P) {
        self.buffer[[x as usize, y as usize]] = pixel;
    }

    /// Return the width of the image.
    pub fn width(&self) -> u32 { self.buffer.cols() as u32 }
    /// Return the height of the image.
    pub fn height(&self) -> u32 { self.buffer.rows() as u32 }
    /// Return the dimensions of the image as a `(width, height)` tuple.
    pub fn dimensions(&self) -> (u32, u32) { (self.width(), self.height()) }

    /// Return a view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image(&'a self, x: u32, y: u32, w: u32, h: u32) -> Image2DView<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice(s![x..x + w, y..y + h])
    }

    /// Return a mutable view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image_mut(&'a mut self, x: u32, y: u32, w: u32, h: u32) -> Image2DViewMut<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice_mut(s![x..x + w, y..y + h])
    }
}

impl<'a, P> IntoIterator for &'a Image2D

    where P: Pixel + 'a
{
    type Item = &'a P;
    type IntoIter = ndarray::iter::Iter<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.into_iter()
    }
}

impl<'a, P> IntoIterator for &'a mut Image2D

    where P: Pixel + 'a
{
    type Item = &'a mut P;
    type IntoIter = ndarray::iter::IterMut<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.iter_mut()
    }
}

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

    use std::iter::FromIterator;
    use std::fmt::Debug;

    #[test]
    fn test_from_vec() {
        let v1 = Vec::from_iter((1u8..10u8).map(|n| Luma::new([n])));
        let v2 = Vec::from_iter((1u8..7u8).map(|n| Luma::new([n])));

        assert!(Image2D::from_vec(3, 3, v1.clone()).is_ok());
        assert!(Image2D::from_vec(2, 3, v2.clone()).is_ok());
        assert!(Image2D::from_vec(3, 2, v2.clone()).is_ok());

        assert!(Image2D::from_vec(3, 3, v2.clone()).is_err());
        assert!(Image2D::from_vec(4, 2, v2.clone()).is_err());
    }

    #[test]
    fn test_new() {
        fn test_zeros_helper
(w: u32, h: u32)
            where P: Pixel + Zero + Debug
        {
            let img = Image2D::
::new(w, h);
            for pixel in &img {
                assert!(pixel.is_zero());
            }
        };
        test_zeros_helper::<Luma<u8>>(100, 100);
        test_zeros_helper::<Luma<f32>>(100, 100);
    }

    #[test]
    fn test_into_iter() {
        let v: Vec<Luma<u8>> = (1..10).map(|i| Luma::from(i)).collect();
        let img = Image2D::from_vec(3, 3, v.clone()).unwrap();

        for (p, i) in img.into_iter().zip(v.into_iter()) {
            assert!(&i == p);
        }
    }

    #[test]
    fn test_sub_image() {
        let v: Vec<Luma<usize>> = (1..10).map(|n| Luma::new([n])).collect();
        let img = Image2D::from_vec(3, 3, v).unwrap();
        let subimg1 = img.sub_image(1, 1, 2, 2);

        fn subimg_vec_eq<'a>(subimg: Image2DView<'a, Luma<usize>>, v: &Vec<usize>) -> bool {
            let img_iter = subimg.into_iter();
            let v_iter   = v.into_iter();
            match img_iter.len() == v_iter.len() {
                true => !img_iter.zip(v_iter).any(|(p, i)| p.data[0] != *i),
                false => false
            }
        }

        let subimg1_vec: Vec<usize> = vec![5, 6, 8, 9];

        assert!(subimg_vec_eq(subimg1, &subimg1_vec));
    }
}
<commit_msg>Add enumerate_pixels and enumerate_pixels_mut methods<commit_after>//! Defines a generic 2D image type.

use ndarray;
use ndarray::prelude::*;
use num_traits::{Zero};

use std::iter::{IntoIterator};

use errors::*;
use traits::Pixel;

/// 2-dimensional image type.
pub struct Image2D

    where P: Pixel
{
    buffer: Array2

}

/// Type of immutable `Image2D` views.
pub type Image2DView<'a, P> = ndarray::ArrayView<'a, P, Ix2>;
pub type Image2DViewMut<'a, P> = ndarray::ArrayViewMut<'a, P, Ix2>;

impl<'a, P> Image2D

    where P: Pixel
{
    /// Create a new image of specified dimensions filled with zeros.
    pub fn new(width: u32, height: u32) -> Image2D

        where P: Pixel + Zero
    {
        Image2D { buffer: Array2::zeros((width as usize, height as usize)) }
    }

    /// Create a new image of specified dimensions from a `Vec`.
    ///
    /// **Error**: `InvalidDimensions` if the dimensions do not match the length of `v`.
    pub fn from_vec(w: u32, h: u32, v: Vec
) -> Result<Image2D
> {
        let buf = try!(Array2::from_shape_vec((w as usize, h as usize), v));
        Ok(Image2D{ buffer: buf })
    }

    /// Return the pixel at the specified coordinates.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn get_pixel(&self, x: u32, y: u32) -> P {
        self.buffer[[x as usize, y as usize]].clone()
    }

    /// Set the pixel at the specified coordinates to the specified value.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn put_pixel(&mut self, x: u32, y: u32, pixel: P) {
        self.buffer[[x as usize, y as usize]] = pixel;
    }

    /// Return the width of the image.
    pub fn width(&self) -> u32 { self.buffer.cols() as u32 }
    /// Return the height of the image.
    pub fn height(&self) -> u32 { self.buffer.rows() as u32 }
    /// Return the dimensions of the image as a `(width, height)` tuple.
    pub fn dimensions(&self) -> (u32, u32) { (self.width(), self.height()) }

    // TODO: map to u32's for coherence
    /// Return an iterator to the pixels and their indices. The type of the iterator is ((usize, usize), &P)
    pub fn enumerate_pixels(&self) -> ndarray::iter::IndexedIter
 {
        self.buffer.indexed_iter()
    }

    // TODO: map to u32's for coherence
    /// Return an iterator to the pixels and their indices. The type of the iterator is ((usize, usize), &mut P)
    pub fn enumerate_pixels_mut(&mut self) -> ndarray::iter::IndexedIterMut
 {
        self.buffer.indexed_iter_mut()
    }

    /// Return a view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image(&'a self, x: u32, y: u32, w: u32, h: u32) -> Image2DView<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice(s![x..x + w, y..y + h])
    }

    /// Return a mutable view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image_mut(&'a mut self, x: u32, y: u32, w: u32, h: u32) -> Image2DViewMut<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice_mut(s![x..x + w, y..y + h])
    }
}

impl<'a, P> IntoIterator for &'a Image2D

    where P: Pixel + 'a
{
    type Item = &'a P;
    type IntoIter = ndarray::iter::Iter<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.into_iter()
    }
}

impl<'a, P> IntoIterator for &'a mut Image2D

    where P: Pixel + 'a
{
    type Item = &'a mut P;
    type IntoIter = ndarray::iter::IterMut<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.iter_mut()
    }
}

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

    use std::iter::FromIterator;
    use std::fmt::Debug;

    #[test]
    fn test_from_vec() {
        let v1 = Vec::from_iter((1u8..10u8).map(|n| Luma::new([n])));
        let v2 = Vec::from_iter((1u8..7u8).map(|n| Luma::new([n])));

        assert!(Image2D::from_vec(3, 3, v1.clone()).is_ok());
        assert!(Image2D::from_vec(2, 3, v2.clone()).is_ok());
        assert!(Image2D::from_vec(3, 2, v2.clone()).is_ok());

        assert!(Image2D::from_vec(3, 3, v2.clone()).is_err());
        assert!(Image2D::from_vec(4, 2, v2.clone()).is_err());
    }

    #[test]
    fn test_new() {
        fn test_zeros_helper
(w: u32, h: u32)
            where P: Pixel + Zero + Debug
        {
            let img = Image2D::
::new(w, h);
            for pixel in &img {
                assert!(pixel.is_zero());
            }
        };
        test_zeros_helper::<Luma<u8>>(100, 100);
        test_zeros_helper::<Luma<f32>>(100, 100);
    }

    #[test]
    fn test_into_iter() {
        let v: Vec<Luma<u8>> = (1..10).map(|i| Luma::from(i)).collect();
        let img = Image2D::from_vec(3, 3, v.clone()).unwrap();

        for (p, i) in img.into_iter().zip(v.into_iter()) {
            assert!(&i == p);
        }
    }

    #[test]
    fn test_enumerate_pixels() {
        let mut v: Vec<Luma<u8>> = vec![];
        for x in 0..3 {
            for y in 0..3 {
                v.push(Luma::from((2*x + 3*y) as u8));
            }
        }
        let img = Image2D::from_vec(3, 3, v.clone()).unwrap();

        for ((x, y), p) in img.enumerate_pixels().map(|((x, y), p)| ((x, y), p.channels()[0])) {
            assert!((2*x + 3*y) as u8 == p);
        }
    }

    #[test]
    fn test_sub_image() {
        let v: Vec<Luma<usize>> = (1..10).map(|n| Luma::new([n])).collect();
        let img = Image2D::from_vec(3, 3, v).unwrap();
        let subimg1 = img.sub_image(1, 1, 2, 2);

        fn subimg_vec_eq<'a>(subimg: Image2DView<'a, Luma<usize>>, v: &Vec<usize>) -> bool {
            let img_iter = subimg.into_iter();
            let v_iter   = v.into_iter();
            match img_iter.len() == v_iter.len() {
                true => !img_iter.zip(v_iter).any(|(p, i)| p.data[0] != *i),
                false => false
            }
        }

        let subimg1_vec: Vec<usize> = vec![5, 6, 8, 9];

        assert!(subimg_vec_eq(subimg1, &subimg1_vec));
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    reactor::{Reactor, ReactorHandle},
    socket_state::SocketState,
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Poll, Token};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

pub(crate) const SOCKET: Token = Token(1);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: ReactorHandle,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    socket: T,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        mut socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = Reactor::new(poll, heartbeat.clone());
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);

        reactor.register(SOCKET, socket_state.handle(), &mut socket, registered)?;
        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            socket,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn ensure_setup(&mut self) {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
            }
            self.status = Status::Setup;
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable()
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.reactor.shutdown();
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup();
        self.poll_socket_events()?;
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.reactor.shutdown();
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream();
            if let Err(e) = self.socket_state.handle_write_result(res) {
                error!("error writing: {:?}", e);
                self.critical_error(e)?;
            }
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            if let Err(e) = self.socket_state.handle_read_result(res) {
                error!("error reading: {:?}", e);
                self.critical_error(e)?;
            }
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.socket_state
                .handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<commit_msg>IoLoop: refactor handle_write_result<commit_after>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    reactor::{Reactor, ReactorHandle},
    socket_state::SocketState,
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Poll, Token};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

pub(crate) const SOCKET: Token = Token(1);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: ReactorHandle,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    socket: T,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        mut socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = Reactor::new(poll, heartbeat.clone());
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);

        reactor.register(SOCKET, socket_state.handle(), &mut socket, registered)?;
        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            socket,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn ensure_setup(&mut self) {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
            }
            self.status = Status::Setup;
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable()
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.reactor.shutdown();
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup();
        self.poll_socket_events()?;
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.reactor.shutdown();
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self.socket_state.handle_write_result(result) {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            if let Err(e) = self.socket_state.handle_read_result(res) {
                error!("error reading: {:?}", e);
                self.critical_error(e)?;
            }
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.socket_state
                .handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>