<commit_before>#[deriving(Show, Eq, PartialEq, Clone)]
pub enum BencodeEvent {
    NumberValue(i64),
    ByteStringValue(Vec<u8>),
    ListStart,
    ListEnd,
    DictStart,
    DictKey(Vec<u8>),
    DictEnd,
    ParseError(Error),
}

#[deriving(Show, Eq, PartialEq, Clone)]
pub struct Error {
    pub pos: u32,
    pub msg: String,
}

#[deriving(Show, Eq, PartialEq, Clone)]
pub enum BencodePosition {
    ListPosition,
    KeyPosition,
    ValuePosition
}

pub struct StreamingParser<T> {
    reader: T,
    pos: u32,
    decoded: u32,
    end: bool,
    curr: Option<u8>,
    stack: Vec<BencodePosition>,
}

macro_rules! expect(($slf:expr, $ch:pat, $ex:expr) => (
    match $slf.curr_char() {
        Some($ch) => $slf.next_byte(),
        _ => return $slf.error($ex)
    }
))

macro_rules! check_nesting(($slf:expr) => (
    if $slf.decoded > 0 && $slf.stack.len() == 0 {
        return $slf.error_msg("Only one value allowed outside of containers".to_string())
    }
))

impl<T: Iterator<u8>> StreamingParser<T> {
    pub fn new(mut reader: T) -> StreamingParser<T> {
        let next = reader.next();
        StreamingParser {
            reader: reader,
            pos: 0,
            decoded: 0,
            end: false,
            curr: next,
            stack: Vec::new()
        }
    }

    fn next_byte(&mut self) {
        self.pos += 1;
        self.curr = self.reader.next();
        self.end = self.curr.is_none();
    }

    fn next_bytes(&mut self, len: uint) -> Result<Vec<u8>, Error> {
        let mut bytes = Vec::with_capacity(len);
        for _ in range(0, len) {
            match self.curr {
                Some(x) =>  bytes.push(x),
                None => {
                    let msg = format!("Expecting {} bytes but only got {}", len, bytes.len());
                    return self.error_msg(msg)
                }
            }
            self.next_byte();
        }
        Ok(bytes)
    }

    #[inline(always)]
    fn curr_char(&mut self) -> Option<char> {
        self.curr.map(|v| v as char)
    }

    fn error<T>(&mut self, expected: String) -> Result<T, Error> {
        let got = self.curr_char();
        let got_char = match got {
            Some(x) => alphanum_to_str(x),
            None => "EOF".to_string()
        };
        self.error_msg(format!("Expecting '{}' but got '{}'", expected, got_char))
    }

    fn error_msg<T>(&mut self, msg: String) -> Result<T, Error> {
        Err(Error {
            pos: self.pos,
            msg: msg
        })
    }

    #[inline(always)]
    fn parse_number(&mut self) -> Result<i64, Error> {
        let sign = match self.curr_char() {
            Some('-') => {
                self.next_byte();
                -1
            }
            _ => 1
        };
        let num = try!(self.parse_number_unsigned());
        expect!(self, 'e', "e".to_string());
        Ok(sign * num)
    }

    #[inline(always)]
    fn parse_number_unsigned(&mut self) -> Result<i64, Error> {
        let mut num = 0;
        match self.curr_char() {
            Some('0') => self.next_byte(),
            Some('1' .. '9') => {
                loop {
                    match self.curr_char() {
                        Some(ch @ '0' .. '9') => self.parse_digit(ch, &mut num),
                        _ => break
                    }
                    self.next_byte();
                }
            }
            _ => return self.error("0-9".to_string())
        };
        Ok(num)
    }

    #[inline(always)]
    fn parse_digit(&self, ch: char, num: &mut i64) {
        *num *= 10;
        *num += ch.to_digit(10).unwrap() as i64;
    }

    #[inline(always)]
    fn parse_bytestring(&mut self) -> Result<Vec<u8>, Error> {
        let len = try!(self.parse_number_unsigned());
        expect!(self, ':', ":".to_string());
        let bytes = try!(self.next_bytes(len as uint));
        Ok(bytes)
    }

    #[inline(always)]
    fn parse_end(&mut self) -> Result<BencodeEvent, Error> {
        self.next_byte();
        match self.stack.pop() {
            Some(ListPosition) => Ok(ListEnd),
            Some(ValuePosition) => {
                Ok(DictEnd)
            }
            _ => return self.error_msg("Unmatched value ending".to_string())
        }
    }

    #[inline(always)]
    fn parse_key(&mut self) -> Result<BencodeEvent, Error> {
        check_nesting!(self);
        match self.curr_char() {
            Some('0' .. '9') => {
                self.decoded += 1;
                let res = try!(self.parse_bytestring());
                Ok(DictKey(res))
            }
            Some('e') => self.parse_end(),
            _ => self.error("0-9 or e".to_string())
        }
    }

    #[inline(always)]
    fn parse_event(&mut self) -> Result<BencodeEvent, Error> {
        match self.curr_char() {
            Some('i') => {
                check_nesting!(self);
                self.next_byte();
                self.decoded += 1;
                let res = try!(self.parse_number());
                Ok(NumberValue(res))
            }
            Some('0' .. '9') => {
                check_nesting!(self);
                self.decoded += 1;
                let res = try!(self.parse_bytestring());
                Ok(ByteStringValue(res))
            }
            Some('l') => {
                check_nesting!(self);
                self.next_byte();
                self.decoded += 1;
                self.stack.push(ListPosition);
                Ok(ListStart)
            }
            Some('d') => {
                check_nesting!(self);
                self.next_byte();
                self.decoded += 1;
                self.stack.push(KeyPosition);
                Ok(DictStart)
            }
            Some('e') => self.parse_end(),
            _ => self.error("i or 0-9 or l or d or e".to_string())
        }
    }

    fn empty_input(&self) -> bool {
        self.curr == None && self.decoded == 0
    }
}

impl<T: Iterator<u8>> Iterator<BencodeEvent> for StreamingParser<T> {
    fn next(&mut self) -> Option<BencodeEvent> {
        if self.end || self.empty_input() {
            return None
        }
        let result = match self.stack.pop() {
            Some(KeyPosition) => {
                self.stack.push(ValuePosition);
                self.parse_key()
            }
            pos => {
                match pos {
                    Some(ValuePosition) => self.stack.push(KeyPosition),
                    Some(_) => self.stack.push(pos.unwrap()),
                    None => {}
                }
                self.parse_event()
            }
        };
        match result {
            Ok(ev) => Some(ev),
            Err(err) => {
                self.end = true;
                Some(ParseError(err))
            }
        }
    }
}

fn alphanum_to_str(ch: char) -> String {
    if ch.is_alphanumeric() {
        ch.to_string()
    } else {
        (ch as u8).to_string()
    }
}

#[cfg(test)]
mod test {
    use std::str::raw;

    use super::{StreamingParser, Error};
    use super::{BencodeEvent, NumberValue, ByteStringValue, ListStart,
                    ListEnd, DictStart, DictKey, DictEnd, ParseError};
    use super::alphanum_to_str;

    fn bytes(s: &str) -> Vec<u8> {
        s.as_bytes().to_vec()
    }

    fn assert_stream_eq(encoded: &str, expected: &[BencodeEvent]) {
        let mut parser = StreamingParser::new(encoded.bytes());
        let result = parser.collect::<Vec<_>>();
        assert_eq!(expected, result.as_slice());
    }

    #[test]
    fn parse_error_on_invalid_first_character() {
        for n in range(::std::u8::MIN, ::std::u8::MAX) {
            match n as char {
                'i' | '0' .. '9' | 'l' | 'd' | 'e' => continue,
                _ => {}
            };
            let msg = format!("Expecting 'i or 0-9 or l or d or e' but got '{}'", alphanum_to_str(n as char));
            assert_stream_eq(unsafe { raw::from_utf8([n]) },
                            [ParseError(Error{
                                pos: 0,
                                msg: msg })]);
        }
    }

    #[test]
    fn parses_empty_input() {
        assert_stream_eq("", []);
    }

    #[test]
    fn parses_number_zero() {
        assert_stream_eq("i0e", [NumberValue(0)]);
    }

    #[test]
    fn parses_positive_numbers() {
        assert_stream_eq("i5e", [NumberValue(5)]);
        assert_stream_eq(format!("i{}e", ::std::i64::MAX).as_slice(), [NumberValue(::std::i64::MAX)]);
    }

    #[test]
    fn parses_negative_numbers() {
        assert_stream_eq("i-5e", [NumberValue(-5)]);
        assert_stream_eq(format!("i{}e", ::std::i64::MIN).as_slice(), [NumberValue(::std::i64::MIN)]);
    }

    #[test]
    fn parse_error_on_number_without_ending() {
        assert_stream_eq("i10", [ParseError(Error{ pos: 3, msg: "Expecting 'e' but got 'EOF'".to_string() })]);
    }

    #[test]
    fn parse_error_on_number_with_leading_zero() {
        assert_stream_eq("i0215e", [ParseError(Error{ pos: 2, msg: "Expecting 'e' but got '2'".to_string() })]);
    }

    #[test]
    fn parse_error_on_empty_number() {
        assert_stream_eq("ie", [ParseError(Error{ pos: 1, msg: "Expecting '0-9' but got 'e'".to_string() })]);
    }

    #[test]
    fn parse_error_on_more_than_one_value_outside_of_containers() {
        let msg = "Only one value allowed outside of containers".to_string();
        assert_stream_eq("i1ei2e", [NumberValue(1), ParseError(Error{ pos: 3, msg: msg.clone() })]);
        assert_stream_eq("i10eli2ee", [NumberValue(10), ParseError(Error{ pos: 4, msg: msg.clone() })]);
        assert_stream_eq("1:ade", [ByteStringValue(bytes("a")), ParseError(Error{ pos: 3, msg: msg.clone() })]);
        assert_stream_eq("3:foo3:bar", [ByteStringValue(bytes("foo")),
                                        ParseError(Error{ pos: 5, msg: msg.clone() })]);
    }

    #[test]
    fn parses_empty_bytestring() {
        assert_stream_eq("0:", [ByteStringValue(Vec::new())]);
    }

    #[test]
    fn parses_short_bytestring() {
        assert_stream_eq("6:abcdef", [ByteStringValue(bytes("abcdef"))]);
    }

    #[test]
    fn parses_long_bytestring() {
        let long = "baz".repeat(10);
        assert_stream_eq(format!("{}:{}", long.len(), long).as_slice(), [ByteStringValue(bytes(long.as_slice()))]);
    }

    #[test]
    fn parse_error_on_too_short_data() {
        assert_stream_eq("5:abcd", [ParseError(Error { pos: 6, msg: "Expecting 5 bytes but only got 4".to_string() })]);
    }

    #[test]
    fn parse_error_on_malformed_bytestring() {
        assert_stream_eq("3abc", [ParseError(Error { pos: 1, msg: "Expecting ':' but got 'a'".to_string() })]);
    }

    #[test]
    fn parse_empty_list() {
        assert_stream_eq("le", [ListStart, ListEnd]);
    }

    #[test]
    fn parses_list_with_number() {
        assert_stream_eq("li2006ee", [ListStart, NumberValue(2006), ListEnd]);
    }

    #[test]
    fn parses_list_with_bytestring() {
        assert_stream_eq("l9:foobarbaze", [ListStart, ByteStringValue(bytes("foobarbaz")), ListEnd]);
    }

    #[test]
    fn parses_list_with_mixed_elements() {
        assert_stream_eq("l4:rusti2006ee",
                         [ListStart,
                          ByteStringValue(bytes("rust")),
                          NumberValue(2006),
                          ListEnd]);
    }

    #[test]
    fn parses_nested_lists() {
        assert_stream_eq("li1983el3:c++e3:oope",
                         [ListStart,
                          NumberValue(1983),
                          ListStart,
                          ByteStringValue(bytes("c++")),
                          ListEnd,
                          ByteStringValue(bytes("oop")),
                          ListEnd]);
    }

    #[test]
    fn parses_empty_dict() {
        assert_stream_eq("de", [DictStart, DictEnd]);
    }

    #[test]
    fn parses_dict_with_number_value() {
        assert_stream_eq("d3:fooi2006ee",
                         [DictStart,
                          DictKey(bytes("foo")),
                          NumberValue(2006),
                          DictEnd])
    }

    #[test]
    fn parses_dict_with_bytestring_value() {
        assert_stream_eq("d3:foo4:2006e",
                         [DictStart,
                          DictKey(bytes("foo")),
                          ByteStringValue(bytes("2006")),
                          DictEnd])
    }

    #[test]
    fn parses_dict_with_list_value() {
        assert_stream_eq("d3:fooli2006eee",
                         [DictStart,
                          DictKey(bytes("foo")),
                          ListStart,
                          NumberValue(2006),
                          ListEnd,
                          DictEnd]);
        assert_stream_eq("d1:ai123e1:bl3:foo4:baaree",
                        [DictStart,
                         DictKey(bytes("a")),
                         NumberValue(123),
                         DictKey(bytes("b")),
                         ListStart,
                         ByteStringValue(bytes("foo")),
                         ByteStringValue(bytes("baar")),
                         ListEnd,
                         DictEnd])
    }

    #[test]
    fn parses_nested_dicts() {
        assert_stream_eq("d3:food3:bari2006eee",
                         [DictStart,
                          DictKey(bytes("foo")),
                          DictStart,
                          DictKey(bytes("bar")),
                          NumberValue(2006),
                          DictEnd,
                          DictEnd])
    }

    #[test]
    fn parse_error_on_key_of_wrong_type() {
        assert_stream_eq("di2006ei1ee",
                         [DictStart,
                          ParseError(Error{ pos: 1, msg: "Expecting '0-9 or e' but got 'i'".to_string() })]);
        assert_stream_eq("dleei1ee",
                         [DictStart,
                          ParseError(Error{ pos: 1, msg: "Expecting '0-9 or e' but got 'l'".to_string() })]);
        assert_stream_eq("ddei1ee",
                         [DictStart,
                          ParseError(Error{ pos: 1, msg: "Expecting '0-9 or e' but got 'd'".to_string() })]);
    }

    #[test]
    fn parse_error_on_unmatched_value_ending() {
        let msg = "Unmatched value ending".to_string();
        assert_stream_eq("e", [ParseError(Error{ pos: 1, msg: msg.clone() })]);
        assert_stream_eq("lee", [ListStart, ListEnd, ParseError(Error{ pos: 3, msg: msg.clone() })]);
    }

}
<commit_msg>Adjust to new match range syntax<commit_after>#[deriving(Show, Eq, PartialEq, Clone)]
pub enum BencodeEvent {
    NumberValue(i64),
    ByteStringValue(Vec<u8>),
    ListStart,
    ListEnd,
    DictStart,
    DictKey(Vec<u8>),
    DictEnd,
    ParseError(Error),
}

#[deriving(Show, Eq, PartialEq, Clone)]
pub struct Error {
    pub pos: u32,
    pub msg: String,
}

#[deriving(Show, Eq, PartialEq, Clone)]
pub enum BencodePosition {
    ListPosition,
    KeyPosition,
    ValuePosition
}

pub struct StreamingParser<T> {
    reader: T,
    pos: u32,
    decoded: u32,
    end: bool,
    curr: Option<u8>,
    stack: Vec<BencodePosition>,
}

macro_rules! expect(($slf:expr, $ch:pat, $ex:expr) => (
    match $slf.curr_char() {
        Some($ch) => $slf.next_byte(),
        _ => return $slf.error($ex)
    }
))

macro_rules! check_nesting(($slf:expr) => (
    if $slf.decoded > 0 && $slf.stack.len() == 0 {
        return $slf.error_msg("Only one value allowed outside of containers".to_string())
    }
))

impl<T: Iterator<u8>> StreamingParser<T> {
    pub fn new(mut reader: T) -> StreamingParser<T> {
        let next = reader.next();
        StreamingParser {
            reader: reader,
            pos: 0,
            decoded: 0,
            end: false,
            curr: next,
            stack: Vec::new()
        }
    }

    fn next_byte(&mut self) {
        self.pos += 1;
        self.curr = self.reader.next();
        self.end = self.curr.is_none();
    }

    fn next_bytes(&mut self, len: uint) -> Result<Vec<u8>, Error> {
        let mut bytes = Vec::with_capacity(len);
        for _ in range(0, len) {
            match self.curr {
                Some(x) =>  bytes.push(x),
                None => {
                    let msg = format!("Expecting {} bytes but only got {}", len, bytes.len());
                    return self.error_msg(msg)
                }
            }
            self.next_byte();
        }
        Ok(bytes)
    }

    #[inline(always)]
    fn curr_char(&mut self) -> Option<char> {
        self.curr.map(|v| v as char)
    }

    fn error<T>(&mut self, expected: String) -> Result<T, Error> {
        let got = self.curr_char();
        let got_char = match got {
            Some(x) => alphanum_to_str(x),
            None => "EOF".to_string()
        };
        self.error_msg(format!("Expecting '{}' but got '{}'", expected, got_char))
    }

    fn error_msg<T>(&mut self, msg: String) -> Result<T, Error> {
        Err(Error {
            pos: self.pos,
            msg: msg
        })
    }

    #[inline(always)]
    fn parse_number(&mut self) -> Result<i64, Error> {
        let sign = match self.curr_char() {
            Some('-') => {
                self.next_byte();
                -1
            }
            _ => 1
        };
        let num = try!(self.parse_number_unsigned());
        expect!(self, 'e', "e".to_string());
        Ok(sign * num)
    }

    #[inline(always)]
    fn parse_number_unsigned(&mut self) -> Result<i64, Error> {
        let mut num = 0;
        match self.curr_char() {
            Some('0') => self.next_byte(),
            Some('1' ... '9') => {
                loop {
                    match self.curr_char() {
                        Some(ch @ '0' ... '9') => self.parse_digit(ch, &mut num),
                        _ => break
                    }
                    self.next_byte();
                }
            }
            _ => return self.error("0-9".to_string())
        };
        Ok(num)
    }

    #[inline(always)]
    fn parse_digit(&self, ch: char, num: &mut i64) {
        *num *= 10;
        *num += ch.to_digit(10).unwrap() as i64;
    }

    #[inline(always)]
    fn parse_bytestring(&mut self) -> Result<Vec<u8>, Error> {
        let len = try!(self.parse_number_unsigned());
        expect!(self, ':', ":".to_string());
        let bytes = try!(self.next_bytes(len as uint));
        Ok(bytes)
    }

    #[inline(always)]
    fn parse_end(&mut self) -> Result<BencodeEvent, Error> {
        self.next_byte();
        match self.stack.pop() {
            Some(ListPosition) => Ok(ListEnd),
            Some(ValuePosition) => {
                Ok(DictEnd)
            }
            _ => return self.error_msg("Unmatched value ending".to_string())
        }
    }

    #[inline(always)]
    fn parse_key(&mut self) -> Result<BencodeEvent, Error> {
        check_nesting!(self);
        match self.curr_char() {
            Some('0' ... '9') => {
                self.decoded += 1;
                let res = try!(self.parse_bytestring());
                Ok(DictKey(res))
            }
            Some('e') => self.parse_end(),
            _ => self.error("0-9 or e".to_string())
        }
    }

    #[inline(always)]
    fn parse_event(&mut self) -> Result<BencodeEvent, Error> {
        match self.curr_char() {
            Some('i') => {
                check_nesting!(self);
                self.next_byte();
                self.decoded += 1;
                let res = try!(self.parse_number());
                Ok(NumberValue(res))
            }
            Some('0' ... '9') => {
                check_nesting!(self);
                self.decoded += 1;
                let res = try!(self.parse_bytestring());
                Ok(ByteStringValue(res))
            }
            Some('l') => {
                check_nesting!(self);
                self.next_byte();
                self.decoded += 1;
                self.stack.push(ListPosition);
                Ok(ListStart)
            }
            Some('d') => {
                check_nesting!(self);
                self.next_byte();
                self.decoded += 1;
                self.stack.push(KeyPosition);
                Ok(DictStart)
            }
            Some('e') => self.parse_end(),
            _ => self.error("i or 0-9 or l or d or e".to_string())
        }
    }

    fn empty_input(&self) -> bool {
        self.curr == None && self.decoded == 0
    }
}

impl<T: Iterator<u8>> Iterator<BencodeEvent> for StreamingParser<T> {
    fn next(&mut self) -> Option<BencodeEvent> {
        if self.end || self.empty_input() {
            return None
        }
        let result = match self.stack.pop() {
            Some(KeyPosition) => {
                self.stack.push(ValuePosition);
                self.parse_key()
            }
            pos => {
                match pos {
                    Some(ValuePosition) => self.stack.push(KeyPosition),
                    Some(_) => self.stack.push(pos.unwrap()),
                    None => {}
                }
                self.parse_event()
            }
        };
        match result {
            Ok(ev) => Some(ev),
            Err(err) => {
                self.end = true;
                Some(ParseError(err))
            }
        }
    }
}

fn alphanum_to_str(ch: char) -> String {
    if ch.is_alphanumeric() {
        ch.to_string()
    } else {
        (ch as u8).to_string()
    }
}

#[cfg(test)]
mod test {
    use std::str::raw;

    use super::{StreamingParser, Error};
    use super::{BencodeEvent, NumberValue, ByteStringValue, ListStart,
                    ListEnd, DictStart, DictKey, DictEnd, ParseError};
    use super::alphanum_to_str;

    fn bytes(s: &str) -> Vec<u8> {
        s.as_bytes().to_vec()
    }

    fn assert_stream_eq(encoded: &str, expected: &[BencodeEvent]) {
        let mut parser = StreamingParser::new(encoded.bytes());
        let result = parser.collect::<Vec<_>>();
        assert_eq!(expected, result.as_slice());
    }

    #[test]
    fn parse_error_on_invalid_first_character() {
        for n in range(::std::u8::MIN, ::std::u8::MAX) {
            match n as char {
                'i' | '0' ... '9' | 'l' | 'd' | 'e' => continue,
                _ => {}
            };
            let msg = format!("Expecting 'i or 0-9 or l or d or e' but got '{}'", alphanum_to_str(n as char));
            assert_stream_eq(unsafe { raw::from_utf8([n]) },
                            [ParseError(Error{
                                pos: 0,
                                msg: msg })]);
        }
    }

    #[test]
    fn parses_empty_input() {
        assert_stream_eq("", []);
    }

    #[test]
    fn parses_number_zero() {
        assert_stream_eq("i0e", [NumberValue(0)]);
    }

    #[test]
    fn parses_positive_numbers() {
        assert_stream_eq("i5e", [NumberValue(5)]);
        assert_stream_eq(format!("i{}e", ::std::i64::MAX).as_slice(), [NumberValue(::std::i64::MAX)]);
    }

    #[test]
    fn parses_negative_numbers() {
        assert_stream_eq("i-5e", [NumberValue(-5)]);
        assert_stream_eq(format!("i{}e", ::std::i64::MIN).as_slice(), [NumberValue(::std::i64::MIN)]);
    }

    #[test]
    fn parse_error_on_number_without_ending() {
        assert_stream_eq("i10", [ParseError(Error{ pos: 3, msg: "Expecting 'e' but got 'EOF'".to_string() })]);
    }

    #[test]
    fn parse_error_on_number_with_leading_zero() {
        assert_stream_eq("i0215e", [ParseError(Error{ pos: 2, msg: "Expecting 'e' but got '2'".to_string() })]);
    }

    #[test]
    fn parse_error_on_empty_number() {
        assert_stream_eq("ie", [ParseError(Error{ pos: 1, msg: "Expecting '0-9' but got 'e'".to_string() })]);
    }

    #[test]
    fn parse_error_on_more_than_one_value_outside_of_containers() {
        let msg = "Only one value allowed outside of containers".to_string();
        assert_stream_eq("i1ei2e", [NumberValue(1), ParseError(Error{ pos: 3, msg: msg.clone() })]);
        assert_stream_eq("i10eli2ee", [NumberValue(10), ParseError(Error{ pos: 4, msg: msg.clone() })]);
        assert_stream_eq("1:ade", [ByteStringValue(bytes("a")), ParseError(Error{ pos: 3, msg: msg.clone() })]);
        assert_stream_eq("3:foo3:bar", [ByteStringValue(bytes("foo")),
                                        ParseError(Error{ pos: 5, msg: msg.clone() })]);
    }

    #[test]
    fn parses_empty_bytestring() {
        assert_stream_eq("0:", [ByteStringValue(Vec::new())]);
    }

    #[test]
    fn parses_short_bytestring() {
        assert_stream_eq("6:abcdef", [ByteStringValue(bytes("abcdef"))]);
    }

    #[test]
    fn parses_long_bytestring() {
        let long = "baz".repeat(10);
        assert_stream_eq(format!("{}:{}", long.len(), long).as_slice(), [ByteStringValue(bytes(long.as_slice()))]);
    }

    #[test]
    fn parse_error_on_too_short_data() {
        assert_stream_eq("5:abcd", [ParseError(Error { pos: 6, msg: "Expecting 5 bytes but only got 4".to_string() })]);
    }

    #[test]
    fn parse_error_on_malformed_bytestring() {
        assert_stream_eq("3abc", [ParseError(Error { pos: 1, msg: "Expecting ':' but got 'a'".to_string() })]);
    }

    #[test]
    fn parse_empty_list() {
        assert_stream_eq("le", [ListStart, ListEnd]);
    }

    #[test]
    fn parses_list_with_number() {
        assert_stream_eq("li2006ee", [ListStart, NumberValue(2006), ListEnd]);
    }

    #[test]
    fn parses_list_with_bytestring() {
        assert_stream_eq("l9:foobarbaze", [ListStart, ByteStringValue(bytes("foobarbaz")), ListEnd]);
    }

    #[test]
    fn parses_list_with_mixed_elements() {
        assert_stream_eq("l4:rusti2006ee",
                         [ListStart,
                          ByteStringValue(bytes("rust")),
                          NumberValue(2006),
                          ListEnd]);
    }

    #[test]
    fn parses_nested_lists() {
        assert_stream_eq("li1983el3:c++e3:oope",
                         [ListStart,
                          NumberValue(1983),
                          ListStart,
                          ByteStringValue(bytes("c++")),
                          ListEnd,
                          ByteStringValue(bytes("oop")),
                          ListEnd]);
    }

    #[test]
    fn parses_empty_dict() {
        assert_stream_eq("de", [DictStart, DictEnd]);
    }

    #[test]
    fn parses_dict_with_number_value() {
        assert_stream_eq("d3:fooi2006ee",
                         [DictStart,
                          DictKey(bytes("foo")),
                          NumberValue(2006),
                          DictEnd])
    }

    #[test]
    fn parses_dict_with_bytestring_value() {
        assert_stream_eq("d3:foo4:2006e",
                         [DictStart,
                          DictKey(bytes("foo")),
                          ByteStringValue(bytes("2006")),
                          DictEnd])
    }

    #[test]
    fn parses_dict_with_list_value() {
        assert_stream_eq("d3:fooli2006eee",
                         [DictStart,
                          DictKey(bytes("foo")),
                          ListStart,
                          NumberValue(2006),
                          ListEnd,
                          DictEnd]);
        assert_stream_eq("d1:ai123e1:bl3:foo4:baaree",
                        [DictStart,
                         DictKey(bytes("a")),
                         NumberValue(123),
                         DictKey(bytes("b")),
                         ListStart,
                         ByteStringValue(bytes("foo")),
                         ByteStringValue(bytes("baar")),
                         ListEnd,
                         DictEnd])
    }

    #[test]
    fn parses_nested_dicts() {
        assert_stream_eq("d3:food3:bari2006eee",
                         [DictStart,
                          DictKey(bytes("foo")),
                          DictStart,
                          DictKey(bytes("bar")),
                          NumberValue(2006),
                          DictEnd,
                          DictEnd])
    }

    #[test]
    fn parse_error_on_key_of_wrong_type() {
        assert_stream_eq("di2006ei1ee",
                         [DictStart,
                          ParseError(Error{ pos: 1, msg: "Expecting '0-9 or e' but got 'i'".to_string() })]);
        assert_stream_eq("dleei1ee",
                         [DictStart,
                          ParseError(Error{ pos: 1, msg: "Expecting '0-9 or e' but got 'l'".to_string() })]);
        assert_stream_eq("ddei1ee",
                         [DictStart,
                          ParseError(Error{ pos: 1, msg: "Expecting '0-9 or e' but got 'd'".to_string() })]);
    }

    #[test]
    fn parse_error_on_unmatched_value_ending() {
        let msg = "Unmatched value ending".to_string();
        assert_stream_eq("e", [ParseError(Error{ pos: 1, msg: msg.clone() })]);
        assert_stream_eq("lee", [ListStart, ListEnd, ParseError(Error{ pos: 3, msg: msg.clone() })]);
    }

}
<|endoftext|>


<commit_before>//  Copyright (C) 2020 Éloïs SANCHEZ.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#[macro_export]
macro_rules! db_schema {
    ($db_name:ident, [ $([$col_path:literal, $col_name:ident, $K:ty, $V:ty]),*, ]) => {
        paste::paste! {
            $(
                // Define each collection event type
                #[derive(Debug, PartialEq)]
                pub enum [<$col_name Event>] {
                    Upsert { key: $K, value: $V },
                    Remove { key: $K },
                    RemoveAll,
                }
                impl kv_typed::prelude::EventTrait for [<$col_name Event>] {
                    type K = $K;
                    type V = $V;

                    fn clear() -> Self { Self::RemoveAll }
                    fn upsert(k: Self::K, v: Self::V) -> Self { Self::Upsert { key: k, value: v, } }
                    fn remove(k: Self::K) -> Self { Self::Remove { key: k } }
                }
            )*
            // Inner module used to hide internals types that must not be exposed on public api
            pub use __inner::{[<$db_name Db>], [<$db_name DbRo>], [<$db_name DbWritable>], [<$db_name DbReadable>]};
            mod __inner {
                use super::*;
                use kv_typed::prelude::*;
                // DbCollections
                #[derive(Clone, Debug)]
                pub struct [<$db_name ColsRo>]<BC: BackendCol> {
                    $([<$col_name:snake>]: ColRo<BC, [<$col_name Event>]>,)*
                }
                #[derive(Clone, Debug)]
                pub struct [<$db_name ColsRw>]<BC: BackendCol> {
                    $([<$col_name:snake>]: ColRw<BC, [<$col_name Event>]>,)*
                }
                impl<BC: BackendCol> [<$db_name ColsRw>]<BC> {
                    fn to_ro(&self) -> [<$db_name ColsRo>]<BC> {
                        [<$db_name ColsRo>] {
                            $([<$col_name:snake>]: self.[<$col_name:snake>].to_ro().clone(),)*
                        }
                    }
                }
                // Db
                #[derive(Debug)]
                pub struct [<$db_name Db>]<B: Backend> {
                    collections: [<$db_name ColsRw>]<B::Col>,
                }
                impl<B: Backend> [<$db_name Db>]<B> {
                    pub const NAME: &'static str = stringify!([<$db_name:snake>]);
                }
                impl<B: Backend> Clone for [<$db_name Db>]<B> {
                    fn clone(&self) -> Self {
                        [<$db_name Db>] {
                            collections: self.collections.clone(),
                        }
                    }
                }
                #[cfg(feature = "explorer")]
                impl<B: Backend> kv_typed::explorer::DbExplorable for [<$db_name Db>]<B> {
                    fn explore<'a>(
                        &self,
                        collection_name: &str,
                        action: kv_typed::explorer::ExplorerAction<'a>,
                        stringify_json_value: fn(serde_json::Value) -> serde_json::Value,
                    ) -> KvResult<std::result::Result<kv_typed::explorer::ExplorerActionResponse, ExplorerActionErr>> {
                        $( if stringify!([<$col_name:snake>]) == collection_name {
                            return action.exec(&self.collections.[<$col_name:snake>], stringify_json_value);
                        } )*
                        Ok(Err(ExplorerActionErr(format!("collection '{}' not exist in database '{}'.", collection_name, stringify!([<$db_name Db>])).into())))
                    }
                    fn list_collections() -> Vec<(&'static str, &'static str, &'static str)> {
                        vec![
                            $((stringify!([<$col_name:snake>]), stringify!($K), stringify!($V)),)*
                        ]
                    }
                }
                // Batch
                pub struct [<$db_name DbBatch>]<B: Backend> {
                    $([<$col_name:snake>]: Batch<B::Col, ColRw<B::Col, [<$col_name Event>]>>,)*
                }
                impl<B: Backend> Default for [<$db_name DbBatch>]<B> {
                    fn default() -> Self {
                        [<$db_name DbBatch>] {
                            $([<$col_name:snake>]: Batch::default(),)*
                        }
                    }
                }
                impl<B: Backend> [<$db_name DbBatch>]<B> {
                    $(pub fn [<$col_name:snake>](&mut self) -> &mut Batch<B::Col, ColRw<B::Col, [<$col_name Event>]>> { &mut self.[<$col_name:snake>] })*
                }
                // impl TransactionalWrite for Db
                #[derive(Debug)]
                pub struct [<$db_name DbTxRw>]<'tx, BC: BackendCol> {
                    $(pub [<$col_name:snake>]: TxColRw<'tx, BC, [<$col_name Event>]>,)*
                }
                impl<'tx, B: Backend> TransactionalWrite<'tx, B::Col> for &'tx [<$db_name Db>]<B> {
                    type TxCols = [<$db_name DbTxRw>]<'tx, B::Col>;

                    fn write<D, F: FnOnce(Self::TxCols) -> KvResult<D>>(&'tx self, f: F) -> KvResult<D> {
                        $(let [<$col_name:snake _upgradable_guard>] = self.collections.[<$col_name:snake>].upgradable_read();)*

                        $(let mut [<$col_name:snake _batch>] =  Batch::<B::Col, ColRw<B::Col, [<$col_name Event>]>>::default();)*

                        let db_tx = [<$db_name DbTxRw>] {
                            $([<$col_name:snake>]: TxColRw::new(
                                unsafe { std::mem::transmute(&mut [<$col_name:snake _batch>]) },
                                unsafe { std::mem::transmute(&[<$col_name:snake _upgradable_guard>]) },
                            ),)*
                        };

                        let data = f(db_tx)?;

                        // Prepare commit
                        $(let ([<$col_name:snake _backend_batch>], [<$col_name:snake _events>]) = [<$col_name:snake _batch>].into_backend_batch_and_events();)*

                        // Acquire exclusive lock
                        $(let mut [<$col_name:snake _write_guard>] = parking_lot::RwLockUpgradableReadGuard::upgrade([<$col_name:snake _upgradable_guard>]);)*;

                        // Commit
                        $(self.collections.[<$col_name:snake>].write_backend_batch(
                            [<$col_name:snake _backend_batch>],
                            [<$col_name:snake _events>],
                            &mut [<$col_name:snake _write_guard>],
                        )?;)*

                        Ok(data)
                    }
                }
                // DbRo
                #[derive(Debug)]
                pub struct [<$db_name DbRo>]<B: Backend> {
                    collections: [<$db_name ColsRo>]<B::Col>,
                }
                impl<B: Backend> [<$db_name DbRo>]<B> {
                    pub const NAME: &'static str = stringify!([<$db_name:snake>]);
                }
                impl<B: Backend> Clone for [<$db_name DbRo>]<B> {
                    fn clone(&self) -> Self {
                        [<$db_name DbRo>] {
                            collections: self.collections.clone(),
                        }
                    }
                }
                // Read operations
                pub trait [<$db_name DbReadable>]: Sized {
                    type Backend: Backend;

                    $(fn [<$col_name:snake>](&self) -> &ColRo<<Self::Backend as Backend>::Col, [<$col_name Event>]>;)*
                }
                impl<B: Backend> [<$db_name DbReadable>] for [<$db_name Db>]<B> {
                    type Backend = B;

                    $(fn [<$col_name:snake>](&self) -> &ColRo<B::Col, [<$col_name Event>]> { &self.collections.[<$col_name:snake>].to_ro() })*
                }
                impl<B: Backend> [<$db_name DbReadable>] for [<$db_name DbRo>]<B>{
                    type Backend = B;

                    $(fn [<$col_name:snake>](&self) -> &ColRo<B::Col, [<$col_name Event>]> { &self.collections.[<$col_name:snake>] })*
                }
                // Write operations
                pub trait [<$db_name DbWritable>]: [<$db_name DbReadable>] {
                    type Backend: Backend;
                    type Batch;
                    $(type [<$col_name ColRw>]: DbCollectionRw;)*
                    type DbRo: Sized;

                    fn clear(&self) -> KvResult<()>;
                    fn get_ro_handler(&self) -> Self::DbRo;
                    fn open(
                        backend_conf: <<Self as [<$db_name DbWritable>]>::Backend as kv_typed::backend::Backend>::Conf,
                    ) -> KvResult <Self>;
                    fn new_batch(&self) -> Self::Batch;
                    fn save(&self) -> KvResult<()>;
                    fn write_batch(&self, batch: Self::Batch) -> KvResult<()>;
                    $(fn [<$col_name:snake _write>](&self) -> &Self::[<$col_name ColRw>];)*
                }
                impl<B: Backend> [<$db_name DbWritable>] for [<$db_name Db>]<B> {
                    type Backend = B;
                    type Batch = [<$db_name DbBatch>]<B>;
                    $(type [<$col_name ColRw>] = ColRw<B::Col, [<$col_name Event>]>;)*
                    type DbRo = [<$db_name DbRo>]<B>;

                    #[inline(always)]
                    fn clear(&self) -> KvResult<()> {
                        $(self.collections.[<$col_name:snake>].clear()?;)*
                        Ok(())
                    }
                    #[inline(always)]
                    fn get_ro_handler(&self) -> Self::DbRo {
                        [<$db_name DbRo>] {
                            collections: self.collections.to_ro(),
                        }
                    }
                    #[inline(always)]
                    fn new_batch(&self) -> Self::Batch {
                        <[<$db_name DbBatch>]::<B>>::default()
                    }
                    fn write_batch(&self, batch: Self::Batch) -> KvResult<()> {
                        $(self.collections.[<$col_name:snake>].write_batch(batch.[<$col_name:snake>])?;)*
                        Ok(())
                    }
                    fn open(
                        backend_conf: <<Self as [<$db_name DbWritable>]>::Backend as kv_typed::backend::Backend>::Conf,
                    ) -> KvResult <Self> {
                        let mut db = B::open(&backend_conf)?;
                        Ok([<$db_name Db>] {
                            collections: [<$db_name ColsRw>] {
                                $([<$col_name:snake>]: <ColRw<B::Col, [<$col_name Event>]>>::new(
                                    db.open_col(&backend_conf, $col_path)?
                                ),)*
                            },
                        })
                    }
                    #[inline(always)]
                    fn save(&self) -> KvResult<()> {
                        $(self.collections.[<$col_name:snake>].save()?;)*
                        Ok(())
                    }
                    $(
                        #[inline(always)]
                        fn [<$col_name:snake _write>](&self) -> &ColRw<B::Col, [<$col_name Event>]> { &self.collections.[<$col_name:snake>] }
                    )*
                }
            }
        }
    };
}
<commit_msg>[fix] kv_typed: *DbTxRw must be exposed<commit_after>//  Copyright (C) 2020 Éloïs SANCHEZ.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#[macro_export]
macro_rules! db_schema {
    ($db_name:ident, [ $([$col_path:literal, $col_name:ident, $K:ty, $V:ty]),*, ]) => {
        paste::paste! {
            $(
                // Define each collection event type
                #[derive(Debug, PartialEq)]
                pub enum [<$col_name Event>] {
                    Upsert { key: $K, value: $V },
                    Remove { key: $K },
                    RemoveAll,
                }
                impl kv_typed::prelude::EventTrait for [<$col_name Event>] {
                    type K = $K;
                    type V = $V;

                    fn clear() -> Self { Self::RemoveAll }
                    fn upsert(k: Self::K, v: Self::V) -> Self { Self::Upsert { key: k, value: v, } }
                    fn remove(k: Self::K) -> Self { Self::Remove { key: k } }
                }
            )*
            // Inner module used to hide internals types that must not be exposed on public api
            pub use __inner::{[<$db_name Db>], [<$db_name DbRo>], [<$db_name DbWritable>], [<$db_name DbReadable>], [<$db_name DbTxRw>]};
            mod __inner {
                use super::*;
                use kv_typed::prelude::*;
                // DbCollections
                #[derive(Clone, Debug)]
                pub struct [<$db_name ColsRo>]<BC: BackendCol> {
                    $([<$col_name:snake>]: ColRo<BC, [<$col_name Event>]>,)*
                }
                #[derive(Clone, Debug)]
                pub struct [<$db_name ColsRw>]<BC: BackendCol> {
                    $([<$col_name:snake>]: ColRw<BC, [<$col_name Event>]>,)*
                }
                impl<BC: BackendCol> [<$db_name ColsRw>]<BC> {
                    fn to_ro(&self) -> [<$db_name ColsRo>]<BC> {
                        [<$db_name ColsRo>] {
                            $([<$col_name:snake>]: self.[<$col_name:snake>].to_ro().clone(),)*
                        }
                    }
                }
                // Db
                #[derive(Debug)]
                pub struct [<$db_name Db>]<B: Backend> {
                    collections: [<$db_name ColsRw>]<B::Col>,
                }
                impl<B: Backend> [<$db_name Db>]<B> {
                    pub const NAME: &'static str = stringify!([<$db_name:snake>]);
                }
                impl<B: Backend> Clone for [<$db_name Db>]<B> {
                    fn clone(&self) -> Self {
                        [<$db_name Db>] {
                            collections: self.collections.clone(),
                        }
                    }
                }
                #[cfg(feature = "explorer")]
                impl<B: Backend> kv_typed::explorer::DbExplorable for [<$db_name Db>]<B> {
                    fn explore<'a>(
                        &self,
                        collection_name: &str,
                        action: kv_typed::explorer::ExplorerAction<'a>,
                        stringify_json_value: fn(serde_json::Value) -> serde_json::Value,
                    ) -> KvResult<std::result::Result<kv_typed::explorer::ExplorerActionResponse, ExplorerActionErr>> {
                        $( if stringify!([<$col_name:snake>]) == collection_name {
                            return action.exec(&self.collections.[<$col_name:snake>], stringify_json_value);
                        } )*
                        Ok(Err(ExplorerActionErr(format!("collection '{}' not exist in database '{}'.", collection_name, stringify!([<$db_name Db>])).into())))
                    }
                    fn list_collections() -> Vec<(&'static str, &'static str, &'static str)> {
                        vec![
                            $((stringify!([<$col_name:snake>]), stringify!($K), stringify!($V)),)*
                        ]
                    }
                }
                // Batch
                pub struct [<$db_name DbBatch>]<B: Backend> {
                    $([<$col_name:snake>]: Batch<B::Col, ColRw<B::Col, [<$col_name Event>]>>,)*
                }
                impl<B: Backend> Default for [<$db_name DbBatch>]<B> {
                    fn default() -> Self {
                        [<$db_name DbBatch>] {
                            $([<$col_name:snake>]: Batch::default(),)*
                        }
                    }
                }
                impl<B: Backend> [<$db_name DbBatch>]<B> {
                    $(pub fn [<$col_name:snake>](&mut self) -> &mut Batch<B::Col, ColRw<B::Col, [<$col_name Event>]>> { &mut self.[<$col_name:snake>] })*
                }
                // impl TransactionalWrite for Db
                #[derive(Debug)]
                pub struct [<$db_name DbTxRw>]<'tx, BC: BackendCol> {
                    $(pub [<$col_name:snake>]: TxColRw<'tx, BC, [<$col_name Event>]>,)*
                }
                impl<'tx, B: Backend> TransactionalWrite<'tx, B::Col> for &'tx [<$db_name Db>]<B> {
                    type TxCols = [<$db_name DbTxRw>]<'tx, B::Col>;

                    fn write<D, F: FnOnce(Self::TxCols) -> KvResult<D>>(&'tx self, f: F) -> KvResult<D> {
                        $(let [<$col_name:snake _upgradable_guard>] = self.collections.[<$col_name:snake>].upgradable_read();)*

                        $(let mut [<$col_name:snake _batch>] =  Batch::<B::Col, ColRw<B::Col, [<$col_name Event>]>>::default();)*

                        let db_tx = [<$db_name DbTxRw>] {
                            $([<$col_name:snake>]: TxColRw::new(
                                unsafe { std::mem::transmute(&mut [<$col_name:snake _batch>]) },
                                unsafe { std::mem::transmute(&[<$col_name:snake _upgradable_guard>]) },
                            ),)*
                        };

                        let data = f(db_tx)?;

                        // Prepare commit
                        $(let ([<$col_name:snake _backend_batch>], [<$col_name:snake _events>]) = [<$col_name:snake _batch>].into_backend_batch_and_events();)*

                        // Acquire exclusive lock
                        $(let mut [<$col_name:snake _write_guard>] = parking_lot::RwLockUpgradableReadGuard::upgrade([<$col_name:snake _upgradable_guard>]);)*;

                        // Commit
                        $(self.collections.[<$col_name:snake>].write_backend_batch(
                            [<$col_name:snake _backend_batch>],
                            [<$col_name:snake _events>],
                            &mut [<$col_name:snake _write_guard>],
                        )?;)*

                        Ok(data)
                    }
                }
                // DbRo
                #[derive(Debug)]
                pub struct [<$db_name DbRo>]<B: Backend> {
                    collections: [<$db_name ColsRo>]<B::Col>,
                }
                impl<B: Backend> [<$db_name DbRo>]<B> {
                    pub const NAME: &'static str = stringify!([<$db_name:snake>]);
                }
                impl<B: Backend> Clone for [<$db_name DbRo>]<B> {
                    fn clone(&self) -> Self {
                        [<$db_name DbRo>] {
                            collections: self.collections.clone(),
                        }
                    }
                }
                // Read operations
                pub trait [<$db_name DbReadable>]: Sized {
                    type Backend: Backend;

                    $(fn [<$col_name:snake>](&self) -> &ColRo<<Self::Backend as Backend>::Col, [<$col_name Event>]>;)*
                }
                impl<B: Backend> [<$db_name DbReadable>] for [<$db_name Db>]<B> {
                    type Backend = B;

                    $(fn [<$col_name:snake>](&self) -> &ColRo<B::Col, [<$col_name Event>]> { &self.collections.[<$col_name:snake>].to_ro() })*
                }
                impl<B: Backend> [<$db_name DbReadable>] for [<$db_name DbRo>]<B>{
                    type Backend = B;

                    $(fn [<$col_name:snake>](&self) -> &ColRo<B::Col, [<$col_name Event>]> { &self.collections.[<$col_name:snake>] })*
                }
                // Write operations
                pub trait [<$db_name DbWritable>]: [<$db_name DbReadable>] {
                    type Backend: Backend;
                    type Batch;
                    $(type [<$col_name ColRw>]: DbCollectionRw;)*
                    type DbRo: Sized;

                    fn clear(&self) -> KvResult<()>;
                    fn get_ro_handler(&self) -> Self::DbRo;
                    fn open(
                        backend_conf: <<Self as [<$db_name DbWritable>]>::Backend as kv_typed::backend::Backend>::Conf,
                    ) -> KvResult <Self>;
                    fn new_batch(&self) -> Self::Batch;
                    fn save(&self) -> KvResult<()>;
                    fn write_batch(&self, batch: Self::Batch) -> KvResult<()>;
                    $(fn [<$col_name:snake _write>](&self) -> &Self::[<$col_name ColRw>];)*
                }
                impl<B: Backend> [<$db_name DbWritable>] for [<$db_name Db>]<B> {
                    type Backend = B;
                    type Batch = [<$db_name DbBatch>]<B>;
                    $(type [<$col_name ColRw>] = ColRw<B::Col, [<$col_name Event>]>;)*
                    type DbRo = [<$db_name DbRo>]<B>;

                    #[inline(always)]
                    fn clear(&self) -> KvResult<()> {
                        $(self.collections.[<$col_name:snake>].clear()?;)*
                        Ok(())
                    }
                    #[inline(always)]
                    fn get_ro_handler(&self) -> Self::DbRo {
                        [<$db_name DbRo>] {
                            collections: self.collections.to_ro(),
                        }
                    }
                    #[inline(always)]
                    fn new_batch(&self) -> Self::Batch {
                        <[<$db_name DbBatch>]::<B>>::default()
                    }
                    fn write_batch(&self, batch: Self::Batch) -> KvResult<()> {
                        $(self.collections.[<$col_name:snake>].write_batch(batch.[<$col_name:snake>])?;)*
                        Ok(())
                    }
                    fn open(
                        backend_conf: <<Self as [<$db_name DbWritable>]>::Backend as kv_typed::backend::Backend>::Conf,
                    ) -> KvResult <Self> {
                        let mut db = B::open(&backend_conf)?;
                        Ok([<$db_name Db>] {
                            collections: [<$db_name ColsRw>] {
                                $([<$col_name:snake>]: <ColRw<B::Col, [<$col_name Event>]>>::new(
                                    db.open_col(&backend_conf, $col_path)?
                                ),)*
                            },
                        })
                    }
                    #[inline(always)]
                    fn save(&self) -> KvResult<()> {
                        $(self.collections.[<$col_name:snake>].save()?;)*
                        Ok(())
                    }
                    $(
                        #[inline(always)]
                        fn [<$col_name:snake _write>](&self) -> &ColRw<B::Col, [<$col_name Event>]> { &self.collections.[<$col_name:snake>] }
                    )*
                }
            }
        }
    };
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Defines the SORT plan

use std::any::Any;
use std::sync::Arc;

use arrow::array::ArrayRef;
pub use arrow::compute::SortOptions;
use arrow::compute::{concat, lexsort_to_indices, take, SortColumn, TakeOptions};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;

use super::SendableRecordBatchStream;
use crate::error::{DataFusionError, Result};
use crate::physical_plan::common::SizedRecordBatchStream;
use crate::physical_plan::expressions::PhysicalSortExpr;
use crate::physical_plan::{common, Distribution, ExecutionPlan, Partitioning};

use async_trait::async_trait;

/// Sort execution plan
#[derive(Debug)]
pub struct SortExec {
    /// Input schema
    input: Arc<dyn ExecutionPlan>,
    /// Sort expressions
    expr: Vec<PhysicalSortExpr>,
    /// Number of threads to execute input partitions on before combining into a single partition
    concurrency: usize,
}

impl SortExec {
    /// Create a new sort execution plan
    pub fn try_new(
        expr: Vec<PhysicalSortExpr>,
        input: Arc<dyn ExecutionPlan>,
        concurrency: usize,
    ) -> Result<Self> {
        Ok(Self {
            expr,
            input,
            concurrency,
        })
    }
}

#[async_trait]
impl ExecutionPlan for SortExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.input.schema().clone()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        Partitioning::UnknownPartitioning(1)
    }

    fn required_child_distribution(&self) -> Distribution {
        Distribution::SinglePartition
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(SortExec::try_new(
                self.expr.clone(),
                children[0].clone(),
                self.concurrency,
            )?)),
            _ => Err(DataFusionError::Internal(
                "SortExec wrong number of children".to_string(),
            )),
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        if 0 != partition {
            return Err(DataFusionError::Internal(format!(
                "SortExec invalid partition {}",
                partition
            )));
        }

        // sort needs to operate on a single partition currently
        if 1 != self.input.output_partitioning().partition_count() {
            return Err(DataFusionError::Internal(
                "SortExec requires a single input partition".to_owned(),
            ));
        }
        let it = self.input.execute(0).await?;
        let batches = common::collect(it).await?;

        // combine all record batches into one for each column
        let combined_batch = RecordBatch::try_new(
            self.schema(),
            self.schema()
                .fields()
                .iter()
                .enumerate()
                .map(|(i, _)| -> Result<ArrayRef> {
                    Ok(concat(
                        &batches
                            .iter()
                            .map(|batch| batch.columns()[i].clone())
                            .collect::<Vec<ArrayRef>>(),
                    )?)
                })
                .collect::<Result<Vec<ArrayRef>>>()?,
        )?;

        // sort combined record batch
        let indices = lexsort_to_indices(
            &self
                .expr
                .iter()
                .map(|e| e.evaluate_to_sort_column(&combined_batch))
                .collect::<Result<Vec<SortColumn>>>()?,
        )?;

        // reorder all rows based on sorted indices
        let sorted_batch = RecordBatch::try_new(
            self.schema(),
            combined_batch
                .columns()
                .iter()
                .map(|column| -> Result<ArrayRef> {
                    Ok(take(
                        column,
                        &indices,
                        // disable bound check overhead since indices are already generated from
                        // the same record batch
                        Some(TakeOptions {
                            check_bounds: false,
                        }),
                    )?)
                })
                .collect::<Result<Vec<ArrayRef>>>()?,
        )?;

        Ok(Box::pin(SizedRecordBatchStream::new(
            self.schema(),
            vec![Arc::new(sorted_batch)],
        )))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::physical_plan::csv::{CsvExec, CsvReadOptions};
    use crate::physical_plan::expressions::col;
    use crate::physical_plan::memory::MemoryExec;
    use crate::physical_plan::merge::MergeExec;
    use crate::test;
    use arrow::array::*;
    use arrow::datatypes::*;

    #[tokio::test]
    async fn test_sort() -> Result<()> {
        let schema = test::aggr_test_schema();
        let partitions = 4;
        let path = test::create_partitioned_csv("aggregate_test_100.csv", partitions)?;
        let csv =
            CsvExec::try_new(&path, CsvReadOptions::new().schema(&schema), None, 1024)?;

        let sort_exec = Arc::new(SortExec::try_new(
            vec![
                // c1 string column
                PhysicalSortExpr {
                    expr: col("c1"),
                    options: SortOptions::default(),
                },
                // c2 uin32 column
                PhysicalSortExpr {
                    expr: col("c2"),
                    options: SortOptions::default(),
                },
                // c7 uin8 column
                PhysicalSortExpr {
                    expr: col("c7"),
                    options: SortOptions::default(),
                },
            ],
            Arc::new(MergeExec::new(Arc::new(csv))),
            2,
        )?);

        let result: Vec<RecordBatch> = test::execute(sort_exec).await?;
        assert_eq!(result.len(), 1);

        let columns = result[0].columns();

        let c1 = as_string_array(&columns[0]);
        assert_eq!(c1.value(0), "a");
        assert_eq!(c1.value(c1.len() - 1), "e");

        let c2 = as_primitive_array::<UInt32Type>(&columns[1]);
        assert_eq!(c2.value(0), 1);
        assert_eq!(c2.value(c2.len() - 1), 5,);

        let c7 = as_primitive_array::<UInt8Type>(&columns[6]);
        assert_eq!(c7.value(0), 15);
        assert_eq!(c7.value(c7.len() - 1), 254,);

        Ok(())
    }

    #[tokio::test]
    async fn test_lex_sort_by_float() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("a", DataType::Float32, true),
            Field::new("b", DataType::Float64, true),
        ]));

        // define data.
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Float32Array::from(vec![
                    Some(f32::NAN),
                    None,
                    None,
                    Some(f32::NAN),
                    Some(1.0_f32),
                    Some(1.0_f32),
                    Some(2.0_f32),
                    Some(3.0_f32),
                ])),
                Arc::new(Float64Array::from(vec![
                    Some(200.0_f64),
                    Some(20.0_f64),
                    Some(10.0_f64),
                    Some(100.0_f64),
                    Some(f64::NAN),
                    None,
                    None,
                    Some(f64::NAN),
                ])),
            ],
        )?;

        let sort_exec = Arc::new(SortExec::try_new(
            vec![
                PhysicalSortExpr {
                    expr: col("a"),
                    options: SortOptions {
                        descending: true,
                        nulls_first: true,
                    },
                },
                PhysicalSortExpr {
                    expr: col("b"),
                    options: SortOptions {
                        descending: false,
                        nulls_first: false,
                    },
                },
            ],
            Arc::new(MemoryExec::try_new(&vec![vec![batch]], schema, None)?),
            2,
        )?);

        assert_eq!(DataType::Float32, *sort_exec.schema().field(0).data_type());
        assert_eq!(DataType::Float64, *sort_exec.schema().field(1).data_type());

        let result: Vec<RecordBatch> = test::execute(sort_exec).await?;
        assert_eq!(result.len(), 1);

        let columns = result[0].columns();

        assert_eq!(DataType::Float32, *columns[0].data_type());
        assert_eq!(DataType::Float64, *columns[1].data_type());

        let a = as_primitive_array::<Float32Type>(&columns[0]);
        let b = as_primitive_array::<Float64Type>(&columns[1]);

        // convert result to strings to allow comparing to expected result containing NaN
        let result: Vec<(Option<String>, Option<String>)> = (0..result[0].num_rows())
            .map(|i| {
                let aval = if a.is_valid(i) {
                    Some(a.value(i).to_string())
                } else {
                    None
                };
                let bval = if b.is_valid(i) {
                    Some(b.value(i).to_string())
                } else {
                    None
                };
                (aval, bval)
            })
            .collect();

        let expected: Vec<(Option<String>, Option<String>)> = vec![
            (None, Some("10".to_owned())),
            (None, Some("20".to_owned())),
            (Some("NaN".to_owned()), Some("100".to_owned())),
            (Some("NaN".to_owned()), Some("200".to_owned())),
            (Some("3".to_owned()), Some("NaN".to_owned())),
            (Some("2".to_owned()), None),
            (Some("1".to_owned()), Some("NaN".to_owned())),
            (Some("1".to_owned()), None),
        ];

        assert_eq!(expected, result);

        Ok(())
    }
}
<commit_msg>ARROW-10673: [Rust] [DataFusion] Made sort not collect on `execute`.<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Defines the SORT plan

use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use futures::stream::Stream;
use futures::Future;

use pin_project_lite::pin_project;

pub use arrow::compute::SortOptions;
use arrow::compute::{concat, lexsort_to_indices, take, SortColumn, TakeOptions};
use arrow::datatypes::SchemaRef;
use arrow::error::Result as ArrowResult;
use arrow::record_batch::RecordBatch;
use arrow::{array::ArrayRef, error::ArrowError};

use super::{RecordBatchStream, SendableRecordBatchStream};
use crate::error::{DataFusionError, Result};
use crate::physical_plan::expressions::PhysicalSortExpr;
use crate::physical_plan::{common, Distribution, ExecutionPlan, Partitioning};

use async_trait::async_trait;

/// Sort execution plan
#[derive(Debug)]
pub struct SortExec {
    /// Input schema
    input: Arc<dyn ExecutionPlan>,
    /// Sort expressions
    expr: Vec<PhysicalSortExpr>,
    /// Number of threads to execute input partitions on before combining into a single partition
    concurrency: usize,
}

impl SortExec {
    /// Create a new sort execution plan
    pub fn try_new(
        expr: Vec<PhysicalSortExpr>,
        input: Arc<dyn ExecutionPlan>,
        concurrency: usize,
    ) -> Result<Self> {
        Ok(Self {
            expr,
            input,
            concurrency,
        })
    }
}

#[async_trait]
impl ExecutionPlan for SortExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.input.schema().clone()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        Partitioning::UnknownPartitioning(1)
    }

    fn required_child_distribution(&self) -> Distribution {
        Distribution::SinglePartition
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(SortExec::try_new(
                self.expr.clone(),
                children[0].clone(),
                self.concurrency,
            )?)),
            _ => Err(DataFusionError::Internal(
                "SortExec wrong number of children".to_string(),
            )),
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        if 0 != partition {
            return Err(DataFusionError::Internal(format!(
                "SortExec invalid partition {}",
                partition
            )));
        }

        // sort needs to operate on a single partition currently
        if 1 != self.input.output_partitioning().partition_count() {
            return Err(DataFusionError::Internal(
                "SortExec requires a single input partition".to_owned(),
            ));
        }
        let input = self.input.execute(0).await?;

        Ok(Box::pin(SortStream::new(input, self.expr.clone())))
    }
}

fn sort_batches(
    batches: &Vec<RecordBatch>,
    schema: &SchemaRef,
    expr: &[PhysicalSortExpr],
) -> ArrowResult<RecordBatch> {
    // combine all record batches into one for each column
    let combined_batch = RecordBatch::try_new(
        schema.clone(),
        schema
            .fields()
            .iter()
            .enumerate()
            .map(|(i, _)| {
                concat(
                    &batches
                        .iter()
                        .map(|batch| batch.columns()[i].clone())
                        .collect::<Vec<ArrayRef>>(),
                )
            })
            .collect::<ArrowResult<Vec<ArrayRef>>>()?,
    )?;

    // sort combined record batch
    let indices = lexsort_to_indices(
        &expr
            .iter()
            .map(|e| e.evaluate_to_sort_column(&combined_batch))
            .collect::<Result<Vec<SortColumn>>>()
            .map_err(DataFusionError::into_arrow_external_error)?,
    )?;

    // reorder all rows based on sorted indices
    RecordBatch::try_new(
        schema.clone(),
        combined_batch
            .columns()
            .iter()
            .map(|column| {
                take(
                    column,
                    &indices,
                    // disable bound check overhead since indices are already generated from
                    // the same record batch
                    Some(TakeOptions {
                        check_bounds: false,
                    }),
                )
            })
            .collect::<ArrowResult<Vec<ArrayRef>>>()?,
    )
}

pin_project! {
    struct SortStream {
        #[pin]
        output: futures::channel::oneshot::Receiver<ArrowResult<RecordBatch>>,
        finished: bool,
        schema: SchemaRef,
    }
}

impl SortStream {
    fn new(input: SendableRecordBatchStream, expr: Vec<PhysicalSortExpr>) -> Self {
        let (tx, rx) = futures::channel::oneshot::channel();

        let schema = input.schema();
        tokio::spawn(async move {
            let schema = input.schema();
            let sorted_batch = common::collect(input)
                .await
                .map_err(DataFusionError::into_arrow_external_error)
                .and_then(move |batches| sort_batches(&batches, &schema, &expr));

            tx.send(sorted_batch)
        });

        Self {
            output: rx,
            finished: false,
            schema,
        }
    }
}

impl Stream for SortStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if self.finished {
            return Poll::Ready(None);
        }

        // is the output ready?
        let this = self.project();
        let output_poll = this.output.poll(cx);

        match output_poll {
            Poll::Ready(result) => {
                *this.finished = true;

                // check for error in receiving channel and unwrap actual result
                let result = match result {
                    Err(e) => Err(ArrowError::ExternalError(Box::new(e))), // error receiving
                    Ok(result) => result,
                };
                Poll::Ready(Some(result))
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl RecordBatchStream for SortStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::physical_plan::csv::{CsvExec, CsvReadOptions};
    use crate::physical_plan::expressions::col;
    use crate::physical_plan::memory::MemoryExec;
    use crate::physical_plan::merge::MergeExec;
    use crate::test;
    use arrow::array::*;
    use arrow::datatypes::*;

    #[tokio::test]
    async fn test_sort() -> Result<()> {
        let schema = test::aggr_test_schema();
        let partitions = 4;
        let path = test::create_partitioned_csv("aggregate_test_100.csv", partitions)?;
        let csv =
            CsvExec::try_new(&path, CsvReadOptions::new().schema(&schema), None, 1024)?;

        let sort_exec = Arc::new(SortExec::try_new(
            vec![
                // c1 string column
                PhysicalSortExpr {
                    expr: col("c1"),
                    options: SortOptions::default(),
                },
                // c2 uin32 column
                PhysicalSortExpr {
                    expr: col("c2"),
                    options: SortOptions::default(),
                },
                // c7 uin8 column
                PhysicalSortExpr {
                    expr: col("c7"),
                    options: SortOptions::default(),
                },
            ],
            Arc::new(MergeExec::new(Arc::new(csv))),
            2,
        )?);

        let result: Vec<RecordBatch> = test::execute(sort_exec).await?;
        assert_eq!(result.len(), 1);

        let columns = result[0].columns();

        let c1 = as_string_array(&columns[0]);
        assert_eq!(c1.value(0), "a");
        assert_eq!(c1.value(c1.len() - 1), "e");

        let c2 = as_primitive_array::<UInt32Type>(&columns[1]);
        assert_eq!(c2.value(0), 1);
        assert_eq!(c2.value(c2.len() - 1), 5,);

        let c7 = as_primitive_array::<UInt8Type>(&columns[6]);
        assert_eq!(c7.value(0), 15);
        assert_eq!(c7.value(c7.len() - 1), 254,);

        Ok(())
    }

    #[tokio::test]
    async fn test_lex_sort_by_float() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("a", DataType::Float32, true),
            Field::new("b", DataType::Float64, true),
        ]));

        // define data.
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Float32Array::from(vec![
                    Some(f32::NAN),
                    None,
                    None,
                    Some(f32::NAN),
                    Some(1.0_f32),
                    Some(1.0_f32),
                    Some(2.0_f32),
                    Some(3.0_f32),
                ])),
                Arc::new(Float64Array::from(vec![
                    Some(200.0_f64),
                    Some(20.0_f64),
                    Some(10.0_f64),
                    Some(100.0_f64),
                    Some(f64::NAN),
                    None,
                    None,
                    Some(f64::NAN),
                ])),
            ],
        )?;

        let sort_exec = Arc::new(SortExec::try_new(
            vec![
                PhysicalSortExpr {
                    expr: col("a"),
                    options: SortOptions {
                        descending: true,
                        nulls_first: true,
                    },
                },
                PhysicalSortExpr {
                    expr: col("b"),
                    options: SortOptions {
                        descending: false,
                        nulls_first: false,
                    },
                },
            ],
            Arc::new(MemoryExec::try_new(&vec![vec![batch]], schema, None)?),
            2,
        )?);

        assert_eq!(DataType::Float32, *sort_exec.schema().field(0).data_type());
        assert_eq!(DataType::Float64, *sort_exec.schema().field(1).data_type());

        let result: Vec<RecordBatch> = test::execute(sort_exec).await?;
        assert_eq!(result.len(), 1);

        let columns = result[0].columns();

        assert_eq!(DataType::Float32, *columns[0].data_type());
        assert_eq!(DataType::Float64, *columns[1].data_type());

        let a = as_primitive_array::<Float32Type>(&columns[0]);
        let b = as_primitive_array::<Float64Type>(&columns[1]);

        // convert result to strings to allow comparing to expected result containing NaN
        let result: Vec<(Option<String>, Option<String>)> = (0..result[0].num_rows())
            .map(|i| {
                let aval = if a.is_valid(i) {
                    Some(a.value(i).to_string())
                } else {
                    None
                };
                let bval = if b.is_valid(i) {
                    Some(b.value(i).to_string())
                } else {
                    None
                };
                (aval, bval)
            })
            .collect();

        let expected: Vec<(Option<String>, Option<String>)> = vec![
            (None, Some("10".to_owned())),
            (None, Some("20".to_owned())),
            (Some("NaN".to_owned()), Some("100".to_owned())),
            (Some("NaN".to_owned()), Some("200".to_owned())),
            (Some("3".to_owned()), Some("NaN".to_owned())),
            (Some("2".to_owned()), None),
            (Some("1".to_owned()), Some("NaN".to_owned())),
            (Some("1".to_owned()), None),
        ];

        assert_eq!(expected, result);

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


<commit_before>// Claxon -- A FLAC decoding library in Rust
// Copyright (C) 2014-2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License, version 3,
// as published by the Free Software Foundation.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#![feature(path_ext)]

extern crate claxon;
extern crate hound;

use std::fs;
use std::io;
use std::path;

fn run_metaflac(fname: &path::Path) -> String {
    use std::process::Command;

    // Run metaflac on the specified file and print all streaminfo data.
    let output = Command::new("metaflac")
                         .arg("--show-min-blocksize")
                         .arg("--show-max-blocksize")
                         .arg("--show-min-framesize")
                         .arg("--show-max-framesize")
                         .arg("--show-sample-rate")
                         .arg("--show-channels")
                         .arg("--show-bps")
                         .arg("--show-total-samples")
                         .arg("--show-md5sum")
                         .arg(fname.to_str().expect("unsupported filename"))
                         .output().ok().expect("failed to run metaflac");
    String::from_utf8(output.stdout).ok().expect("metaflac wrote invalid UTF-8")
}

fn decode_file(fname: &path::Path) {
    use std::process::Command;

    // Run the the reference flac decoder on the file.
    let success = Command::new("flac")
                          .arg("--decode")
                          .arg(fname.to_str().expect("unsupported filename"))
                          .status().ok().expect("failed to run flac")
                          .success();
    assert!(success);
}

fn print_hex(seq: &[u8]) -> String {
    let vec: Vec<String> = seq.iter().map(|x| format!("{:0>2x}", *x)).collect();
    vec.concat()
}

fn read_streaminfo(fname: &path::Path) -> String {
    use claxon::FlacStream;

    // Use a buffered reader, this speeds up the test by 120%.
    let file = fs::File::open(fname).unwrap();
    let mut reader = io::BufReader::new(file);
    let stream = FlacStream::new(&mut reader).unwrap();
    let streaminfo = stream.streaminfo();

    // Format the streaminfo in the same way that metaflac prints it.
    format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
            streaminfo.min_block_size,
            streaminfo.max_block_size,
            streaminfo.min_frame_size.unwrap_or(0),
            streaminfo.max_frame_size.unwrap_or(0),
            streaminfo.sample_rate,
            streaminfo.channels,
            streaminfo.bits_per_sample,
            streaminfo.samples.unwrap_or(0),
            print_hex(&streaminfo.md5sum)) // TODO implement LowerHex for &[u8] and submit a PR.
}

fn compare_metaflac(fname: &path::Path) {
    let metaflac = run_metaflac(fname);
    let streaminfo = read_streaminfo(fname);
    let mut mf_lines = metaflac.lines_any();
    let mut si_lines = streaminfo.lines_any();
    while let (Some(mf), Some(si)) = (mf_lines.next(), si_lines.next()) {
        if mf != si {
            println!("metaflac\n--------\n{}", metaflac);
            println!("streaminfo\n----------\n{}", streaminfo);
            panic!("metaflac disagrees on parsed streaminfo");
        }
    };
}

fn compare_decoded_stream(fname: &path::Path) {
    use std::fs::PathExt;

    let ref_fname = fname.with_extension("wav");
    if !ref_fname.exists() {
        // TODO: actually, we might only want to run the test for files that
        // do exist already. It is fine to run streaminfo or decode-only tests
        // on thousands of files, but decoding thousands of files to wav before
        // running the tests will take a lot of space. For now, while the
        // number of samples is closer to 10 than to 100, this is fine.
        decode_file(fname);
    }

    // If the reference file does exist after decoding, we can compare it to
    // how Claxon decodes it, sample by sample.
    if ref_fname.exists() {
        let mut ref_stream = hound::WavReader::open(ref_fname).unwrap();

        let try_file = fs::File::open(fname).unwrap();
        let mut try_reader = io::BufReader::new(try_file);
        // TODO: Claxon should take R by value, not by mut reference.
        let mut try_stream = claxon::FlacStream::new(&mut try_reader).unwrap();

        // The streaminfo test will ensure that things like bit depth and
        // sample rate match, only the actual samples are compared here.

        let mut ref_samples = ref_stream.samples::<i32>();

        let samples = try_stream.streaminfo().samples.unwrap();
        let n_channels = try_stream.streaminfo().channels;
        let mut blocks = try_stream.blocks::<i32>();
        let mut sample = 0u64;
        let mut b = 0u64;

        while sample < samples {
            let block = blocks.read_next().unwrap();
            let mut channels: Vec<_> = (0 .. n_channels)
                                       .map(|i| block.channel(i).iter().cloned())
                                       .collect();
            for i in (0 .. block.len()) {
                for ch in (0 .. n_channels as usize) {
                    let ref_sample = ref_samples.next().map(|r| r.ok().unwrap());
                    let try_sample = channels[ch].next();
                    if ref_sample != try_sample {
                        println!("disagreement at sample {} of block {} in channel {}: reference is {} but decoded is {}",
                                 i, b, ch, ref_sample.unwrap(), try_sample.unwrap());
                        panic!("decoding differs from reference decoder");
                    }
                }
            }
            sample = sample + block.len() as u64;
            b = b + 1;
        }
    }
}

fn for_each_test_sample<F: Fn(&path::Path)>(f: F) {
    use std::ffi::OsStr;
    use std::fs::PathExt;

    // Enumerate all the flac files in the testsamples directory, and execute
    // the function `f` for it.
    let dir = fs::read_dir("testsamples")
                 .ok().expect("failed to enumerate flac files");
    for path in dir {
        let path = path.ok().expect("failed to obtain path info").path();
        if path.is_file() && path.extension() == Some(OsStr::new("flac")) {
            print!("    comparing {} ...", path.to_str()
                                               .expect("unsupported filename"));
            f(&path);
            println!(" ok");
        }
    }
}

#[test]
fn verify_streaminfo() {
    for_each_test_sample(compare_metaflac);
}

#[test]
fn verify_decoded_stream() {
    for_each_test_sample(compare_decoded_stream);
}
<commit_msg>add hard-coded streaminfo tests for populated test samples<commit_after>// Claxon -- A FLAC decoding library in Rust
// Copyright (C) 2014-2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License, version 3,
// as published by the Free Software Foundation.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#![feature(path_ext)]

extern crate claxon;
extern crate hound;

use std::fs;
use std::io;
use std::path;

fn run_metaflac(fname: &path::Path) -> String {
    use std::process::Command;

    // Run metaflac on the specified file and print all streaminfo data.
    let output = Command::new("metaflac")
                         .arg("--show-min-blocksize")
                         .arg("--show-max-blocksize")
                         .arg("--show-min-framesize")
                         .arg("--show-max-framesize")
                         .arg("--show-sample-rate")
                         .arg("--show-channels")
                         .arg("--show-bps")
                         .arg("--show-total-samples")
                         .arg("--show-md5sum")
                         .arg(fname.to_str().expect("unsupported filename"))
                         .output().ok().expect("failed to run metaflac");
    String::from_utf8(output.stdout).ok().expect("metaflac wrote invalid UTF-8")
}

fn decode_file(fname: &path::Path) {
    use std::process::Command;

    // Run the the reference flac decoder on the file.
    let success = Command::new("flac")
                          .arg("--decode")
                          .arg(fname.to_str().expect("unsupported filename"))
                          .status().ok().expect("failed to run flac")
                          .success();
    assert!(success);
}

fn print_hex(seq: &[u8]) -> String {
    let vec: Vec<String> = seq.iter().map(|x| format!("{:0>2x}", *x)).collect();
    vec.concat()
}

fn read_streaminfo(fname: &path::Path) -> String {
    use claxon::FlacStream;

    // Use a buffered reader, this speeds up the test by 120%.
    let file = fs::File::open(fname).unwrap();
    let mut reader = io::BufReader::new(file);
    let stream = FlacStream::new(&mut reader).unwrap();
    let streaminfo = stream.streaminfo();

    // Format the streaminfo in the same way that metaflac prints it.
    format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
            streaminfo.min_block_size,
            streaminfo.max_block_size,
            streaminfo.min_frame_size.unwrap_or(0),
            streaminfo.max_frame_size.unwrap_or(0),
            streaminfo.sample_rate,
            streaminfo.channels,
            streaminfo.bits_per_sample,
            streaminfo.samples.unwrap_or(0),
            print_hex(&streaminfo.md5sum)) // TODO implement LowerHex for &[u8] and submit a PR.
}

fn compare_metaflac(fname: &path::Path) {
    let metaflac = run_metaflac(fname);
    let streaminfo = read_streaminfo(fname);
    let mut mf_lines = metaflac.lines_any();
    let mut si_lines = streaminfo.lines_any();
    while let (Some(mf), Some(si)) = (mf_lines.next(), si_lines.next()) {
        if mf != si {
            println!("metaflac\n--------\n{}", metaflac);
            println!("streaminfo\n----------\n{}", streaminfo);
            panic!("metaflac disagrees on parsed streaminfo");
        }
    };
}

fn compare_decoded_stream(fname: &path::Path) {
    use std::fs::PathExt;

    let ref_fname = fname.with_extension("wav");
    if !ref_fname.exists() {
        // TODO: actually, we might only want to run the test for files that
        // do exist already. It is fine to run streaminfo or decode-only tests
        // on thousands of files, but decoding thousands of files to wav before
        // running the tests will take a lot of space. For now, while the
        // number of samples is closer to 10 than to 100, this is fine.
        decode_file(fname);
    }

    // If the reference file does exist after decoding, we can compare it to
    // how Claxon decodes it, sample by sample.
    if ref_fname.exists() {
        let mut ref_stream = hound::WavReader::open(ref_fname).unwrap();

        let try_file = fs::File::open(fname).unwrap();
        let mut try_reader = io::BufReader::new(try_file);
        // TODO: Claxon should take R by value, not by mut reference.
        let mut try_stream = claxon::FlacStream::new(&mut try_reader).unwrap();

        // The streaminfo test will ensure that things like bit depth and
        // sample rate match, only the actual samples are compared here.

        let mut ref_samples = ref_stream.samples::<i32>();

        let samples = try_stream.streaminfo().samples.unwrap();
        let n_channels = try_stream.streaminfo().channels;
        let mut blocks = try_stream.blocks::<i32>();
        let mut sample = 0u64;
        let mut b = 0u64;

        while sample < samples {
            let block = blocks.read_next().unwrap();
            let mut channels: Vec<_> = (0 .. n_channels)
                                       .map(|i| block.channel(i).iter().cloned())
                                       .collect();
            for i in (0 .. block.len()) {
                for ch in (0 .. n_channels as usize) {
                    let ref_sample = ref_samples.next().map(|r| r.ok().unwrap());
                    let try_sample = channels[ch].next();
                    if ref_sample != try_sample {
                        println!("disagreement at sample {} of block {} in channel {}: reference is {} but decoded is {}",
                                 i, b, ch, ref_sample.unwrap(), try_sample.unwrap());
                        panic!("decoding differs from reference decoder");
                    }
                }
            }
            sample = sample + block.len() as u64;
            b = b + 1;
        }
    }
}

fn for_each_test_sample<F: Fn(&path::Path)>(f: F) {
    use std::ffi::OsStr;
    use std::fs::PathExt;

    // Enumerate all the flac files in the testsamples directory, and execute
    // the function `f` for it.
    let dir = fs::read_dir("testsamples")
                 .ok().expect("failed to enumerate flac files");
    for path in dir {
        let path = path.ok().expect("failed to obtain path info").path();
        if path.is_file() && path.extension() == Some(OsStr::new("flac")) {
            print!("    comparing {} ...", path.to_str()
                                               .expect("unsupported filename"));
            f(&path);
            println!(" ok");
        }
    }
}

#[test]
fn verify_streaminfo_all() {
    for_each_test_sample(compare_metaflac);
}

#[test]
fn verify_decoded_stream_all() {
    for_each_test_sample(compare_decoded_stream);
}

// Hard-coded tests for the test samples added py populate.sh.
// This allows us to run these tests in parallel and follow progress without
// enabling --nocapture.

#[test]
fn verify_streaminfo_p0() {
    compare_metaflac(path::Path::new("testsamples/p0.flac"));
}

#[test]
fn verify_streaminfo_p1() {
    compare_metaflac(path::Path::new("testsamples/p1.flac"));
}

#[test]
fn verify_streaminfo_p2() {
    compare_metaflac(path::Path::new("testsamples/p2.flac"));
}

#[test]
fn verify_streaminfo_p3() {
    compare_metaflac(path::Path::new("testsamples/p3.flac"));
}

#[test]
fn verify_streaminfo_p4() {
    compare_metaflac(path::Path::new("testsamples/p4.flac"));
}
<|endoftext|>


<commit_before>//! Minecraft protocol length-prefixed array data type

use std::io;
use std::io::prelude::*;
use std::iter::{ AdditiveIterator, FromIterator };
use std::marker::PhantomData;
use std::num::{ NumCast, ToPrimitive };

use packet::Protocol;

pub struct Arr<L, T>(PhantomData<(fn() -> L, T)>);

impl<L: Protocol, T: Protocol> Protocol for Arr<L, T> where L::Clean: NumCast {
    type Clean = Vec<T::Clean>;

    fn proto_len(value: &Vec<T::Clean>) -> usize {
        let len_len = <L as Protocol>::proto_len(&(<<L as Protocol>::Clean as NumCast>::from(value.len()).unwrap()));
        let len_values = value.iter().map(|elt| <T as Protocol>::proto_len(elt)).sum();
        len_len + len_values
    }

    fn proto_encode(value: &Vec<T::Clean>, dst: &mut Write) -> io::Result<()> {
        let len = try!(<L::Clean as NumCast>::from(value.len()).ok_or(io::Error::new(io::ErrorKind::InvalidInput, "could not convert length of vector to Array length type", None)));
        try!(<L as Protocol>::proto_encode(&len, dst));
        for elt in value {
            try!(<T as Protocol>::proto_encode(elt, dst));
        }
        Ok(())
    }

    fn proto_decode(src: &mut Read) -> io::Result<Vec<T::Clean>> {
        let len = try!(try!(<L as Protocol>::proto_decode(src)).to_usize().ok_or(io::Error::new(io::ErrorKind::InvalidInput, "could not read length of vector from Array length type", None)));
        <io::Result<Vec<T::Clean>> as FromIterator<_>>::from_iter((0..len).map(|_| <T as Protocol>::proto_decode(src)))
    }
}

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

    use std::io;

    use packet::Protocol;
    use types::Var;

    #[test]
    fn arr_encode_i8_varint() {
        let mut dst = Vec::new();
        let value = vec![0i32, -1i32];
        <Arr<i8, Var<i32>> as Protocol>::proto_encode(&value, &mut dst).unwrap();
        let bytes = vec![2, 0, 0xff, 0xff, 0xff, 0xff, 0xf];
        assert_eq!(&dst, &bytes);
    }

    #[test]
    fn arr_decode_i8_varint() {
        let bytes = vec![2, 0, 0xff, 0xff, 0xff, 0xff, 0xf];
        let arr = vec![0i32, -1i32];
        let mut src = io::Cursor::new(bytes);
        let value = <Arr<i8, Var<i32>> as Protocol>::proto_decode(&mut src).unwrap();
        assert_eq!(arr, value);
    }

    #[test]
    fn arr_encode_i32_i32() {
        let mut dst = Vec::new();
        let value = vec![0i32, -1i32];
        <Arr<i32, i32> as Protocol>::proto_encode(&value, &mut dst).unwrap();
        let bytes = vec![
            0x00, 0x00, 0x00, 0x02,
            0x00, 0x00, 0x00, 0x00,
            0xff, 0xff, 0xff, 0xff
        ];
        assert_eq!(&dst, &bytes);
    }

    #[test]
    fn arr_decode_i32_i32() {
        let bytes = vec![
            0x00, 0x00, 0x00, 0x02,
            0x00, 0x00, 0x00, 0x00,
            0xff, 0xff, 0xff, 0xff
        ];
        let arr = vec![0i32, -1i32];
        let mut src = io::Cursor::new(bytes);
        let value = <Arr<i32, i32> as Protocol>::proto_decode(&mut src).unwrap();
        assert_eq!(arr, value);
    }
}
<commit_msg>Clean up from_iter call<commit_after>//! Minecraft protocol length-prefixed array data type

use std::io;
use std::io::prelude::*;
use std::iter::{ AdditiveIterator, FromIterator };
use std::marker::PhantomData;
use std::num::{ NumCast, ToPrimitive };

use packet::Protocol;

pub struct Arr<L, T>(PhantomData<(fn() -> L, T)>);

impl<L: Protocol, T: Protocol> Protocol for Arr<L, T> where L::Clean: NumCast {
    type Clean = Vec<T::Clean>;

    fn proto_len(value: &Vec<T::Clean>) -> usize {
        let len_len = <L as Protocol>::proto_len(&(<<L as Protocol>::Clean as NumCast>::from(value.len()).unwrap()));
        let len_values = value.iter().map(|elt| <T as Protocol>::proto_len(elt)).sum();
        len_len + len_values
    }

    fn proto_encode(value: &Vec<T::Clean>, dst: &mut Write) -> io::Result<()> {
        let len = try!(<L::Clean as NumCast>::from(value.len()).ok_or(io::Error::new(io::ErrorKind::InvalidInput, "could not convert length of vector to Array length type", None)));
        try!(<L as Protocol>::proto_encode(&len, dst));
        for elt in value {
            try!(<T as Protocol>::proto_encode(elt, dst));
        }
        Ok(())
    }

    fn proto_decode(src: &mut Read) -> io::Result<Vec<T::Clean>> {
        let len = try!(try!(<L as Protocol>::proto_decode(src)).to_usize().ok_or(io::Error::new(io::ErrorKind::InvalidInput, "could not read length of vector from Array length type", None)));
        io::Result::from_iter((0..len).map(|_| <T as Protocol>::proto_decode(src)))
    }
}

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

    use std::io;

    use packet::Protocol;
    use types::Var;

    #[test]
    fn arr_encode_i8_varint() {
        let mut dst = Vec::new();
        let value = vec![0i32, -1i32];
        <Arr<i8, Var<i32>> as Protocol>::proto_encode(&value, &mut dst).unwrap();
        let bytes = vec![2, 0, 0xff, 0xff, 0xff, 0xff, 0xf];
        assert_eq!(&dst, &bytes);
    }

    #[test]
    fn arr_decode_i8_varint() {
        let bytes = vec![2, 0, 0xff, 0xff, 0xff, 0xff, 0xf];
        let arr = vec![0i32, -1i32];
        let mut src = io::Cursor::new(bytes);
        let value = <Arr<i8, Var<i32>> as Protocol>::proto_decode(&mut src).unwrap();
        assert_eq!(arr, value);
    }

    #[test]
    fn arr_encode_i32_i32() {
        let mut dst = Vec::new();
        let value = vec![0i32, -1i32];
        <Arr<i32, i32> as Protocol>::proto_encode(&value, &mut dst).unwrap();
        let bytes = vec![
            0x00, 0x00, 0x00, 0x02,
            0x00, 0x00, 0x00, 0x00,
            0xff, 0xff, 0xff, 0xff
        ];
        assert_eq!(&dst, &bytes);
    }

    #[test]
    fn arr_decode_i32_i32() {
        let bytes = vec![
            0x00, 0x00, 0x00, 0x02,
            0x00, 0x00, 0x00, 0x00,
            0xff, 0xff, 0xff, 0xff
        ];
        let arr = vec![0i32, -1i32];
        let mut src = io::Cursor::new(bytes);
        let value = <Arr<i32, i32> as Protocol>::proto_decode(&mut src).unwrap();
        assert_eq!(arr, value);
    }
}
<|endoftext|>