<commit_before>use std::rc::Rc;

pub struct Node<T> {
    elem: T,
    next: Link<T>,
}

pub struct List<T> {
    head: Link<T>,
}

type Link<T> = Option<Rc<Node<T>>>;

impl<T> List<T> {
    fn new() -> Self {
        List {
            head: None
        }
    }

    fn append(&self, elem: T) -> Self {
        let new_node = Node {
            elem: elem,
            next: self.head.clone(),
        };
        List {
            head: Some(Rc::new(new_node)),
        }
    }

    fn tail(&self) -> Self {
        List {
            head: self.head.as_ref().and_then(|x| x.next.clone())
        }
    }

    fn head(&self) -> Option<&T> {
        self.head.as_ref().map(|x| &x.elem)
    }

}

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

    #[test]
    fn basics() {
        let list = List::new();
        assert_eq!(list.head(), None);

        let list = list.append(1).append(2).append(3);
        assert_eq!(list.head(), Some(&3));

        let list = list.tail();
        assert_eq!(list.head(), Some(&2));

        let list = list.tail();
        assert_eq!(list.head(), Some(&1));

        let list = list.tail();
        assert_eq!(list.head(), None);
    }
}

<commit_msg>iterator implementation for persistent linked list<commit_after>use std::rc::Rc;

pub struct Node<T> {
    elem: T,
    next: Link<T>,
}

pub struct List<T> {
    head: Link<T>,
}

type Link<T> = Option<Rc<Node<T>>>;

impl<T> List<T> {
    fn new() -> Self {
        List {
            head: None
        }
    }

    fn append(&self, elem: T) -> Self {
        let new_node = Node {
            elem: elem,
            next: self.head.clone(),
        };
        List {
            head: Some(Rc::new(new_node)),
        }
    }

    fn tail(&self) -> Self {
        List {
            head: self.head.as_ref().and_then(|x| x.next.clone())
        }
    }

    fn head(&self) -> Option<&T> {
        self.head.as_ref().map(|x| &x.elem)
    }

}

pub struct Iter<'a, T: 'a> {
    next: Option<&'a Node<T>>,
}

impl<T> List<T> {
    fn iter<'a>(&'a self) -> Iter<'a, T> {
        Iter {
            next: self.head.as_ref().map(|node| &**node)
        }
    }
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        self.next.map(|item| {
            self.next = item.next.as_ref().map(|node| {
                &**node
            });
            &item.elem
        })
    }
}

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

    #[test]
    fn basics() {
        let list = List::new();
        assert_eq!(list.head(), None);

        let list = list.append(1).append(2).append(3);
        assert_eq!(list.head(), Some(&3));

        let list = list.tail();
        assert_eq!(list.head(), Some(&2));

        let list = list.tail();
        assert_eq!(list.head(), Some(&1));

        let list = list.tail();
        assert_eq!(list.head(), None);
    }

    #[test]
    fn iter() {
        let list = List::new().append(1).append(2);
        let mut iter = list.iter();

        assert_eq!(iter.next(), Some(&2));
        assert_eq!(iter.next(), Some(&1));
        assert_eq!(iter.next(), None);
        assert_eq!(list.head(), Some(&2));
    }
}

<|endoftext|>


<commit_before>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use czmq::{ZMsg, ZSock};
use error::{Error, Result};
use std::error::Error as StdError;

pub struct Msg;

impl Msg {
    /// Receive a variable number of args from sock. We expect the number
    /// of args from `sock` to be between `min` and `max`.
    /// If max = None then we allow a variable number of args.
    pub fn expect_recv(sock: &ZSock, min: usize, max: Option<usize>, block: bool) -> Result<ZMsg> {
        // Avoid blocking unless it is explicitly required.
        let zmsg = if sock.rcvmore() || block {
            try!(ZMsg::recv(sock))
        } else {
            ZMsg::new()
        };

        if min > zmsg.size() || (max.is_some() && max.unwrap() < zmsg.size()) {
            if try!(Self::sock_is_writeable(sock)) {
                let zmsg = ZMsg::new();
                try!(zmsg.addstr("Err"));
                try!(zmsg.addstr("Invalid args count"));
                try!(zmsg.send(sock));
            }

            Err(Error::InvalidArgsCount)
        } else {
            Ok(zmsg)
        }
    }

    pub fn send(sock: &ZSock, frames: Vec<&str>) -> Result<()> {
        let zmsg = ZMsg::new();
        for f in frames {
            try!(zmsg.addstr(&f));
        }
        try!(zmsg.send(sock));
        Ok(())
    }

    pub fn send_ok(sock: &ZSock, mut frames: Vec<&str>) -> Result<()> {
        frames.insert(0, "Ok");
        Self::send(sock, frames)
    }

    pub fn send_err(sock: &ZSock, err: Error) -> Result<()> {
        Self::send(sock, vec!["Err", err.description()])
    }

    fn sock_is_writeable(sock: &ZSock) -> Result<bool> {
        match try!(sock.type_str()).unwrap() {
            "PAIR" |
            "PUB" |
            "REQ" |
            "REP" |
            "DEALER" |
            "ROUTER" |
            "PUSH" |
            "XSUB" => Ok(true),

            "SUB" |
            "PULL" |
            "XPUB" => Ok(false),

            _ => unimplemented!(),
        }
    }
}

#[cfg(test)]
mod tests {
    use czmq::{ZMsg, ZSock, zsys_init};
    use super::*;

    /// Test providing 1 and only 1 arg
    #[test]
    fn test_recv_args_ok_eq() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_ok_eq").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_ok_eq").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 1, Some(1), true).unwrap();
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "0");
    }

    /// Test providing 1 or more args
    #[test]
    fn test_recv_args_ok_range() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_ok_range").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_ok_range").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 1, Some(2), true).unwrap();
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "0");
    }

    /// Test providing 1+ args
    #[test]
    fn test_recv_args_ok_variable() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_ok_variable").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_ok_variable").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 2, None, true).unwrap();
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "0");
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "1");
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "2");
    }

    /// Test failing less than 3 args
    #[test]
    fn test_recv_args_err_min() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_err_min").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_err_min").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.addstr("1").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 3, None, true);
        assert!(rcv_msg.is_err());
    }

    /// Test failing more than 1 arg
    #[test]
    fn test_recv_args_err_max() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_err_max").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_err_max").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 0, Some(1), true);
        assert!(rcv_msg.is_err());
    }
}
<commit_msg>Fixed bug with ZMsg where rcvmore() is not reliable for detecting single frame msgs<commit_after>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use czmq::{ZMsg, ZSock};
use error::{Error, Result};
use std::error::Error as StdError;

pub struct Msg;

impl Msg {
    /// Receive a variable number of args from sock. We expect the number
    /// of args from `sock` to be between `min` and `max`.
    /// If max = None then we allow a variable number of args.
    pub fn expect_recv(sock: &ZSock, min: usize, max: Option<usize>, block: bool) -> Result<ZMsg> {
        let zmsg = if block {
            try!(ZMsg::recv(sock))
        } else {
            let rcvtimeo = sock.rcvtimeo();
            sock.set_rcvtimeo(Some(0));
            let result = ZMsg::recv(sock);
            sock.set_rcvtimeo(rcvtimeo);
            result.unwrap_or(ZMsg::new())
        };

        if min > zmsg.size() || (max.is_some() && max.unwrap() < zmsg.size()) {
            if try!(Self::sock_is_writeable(sock)) {
                let zmsg = ZMsg::new();
                try!(zmsg.addstr("Err"));
                try!(zmsg.addstr("Invalid args count"));
                try!(zmsg.send(sock));
            }

            Err(Error::InvalidArgsCount)
        } else {
            Ok(zmsg)
        }
    }

    pub fn send(sock: &ZSock, frames: Vec<&str>) -> Result<()> {
        let zmsg = ZMsg::new();
        for f in frames {
            try!(zmsg.addstr(&f));
        }
        try!(zmsg.send(sock));
        Ok(())
    }

    pub fn send_ok(sock: &ZSock, mut frames: Vec<&str>) -> Result<()> {
        frames.insert(0, "Ok");
        Self::send(sock, frames)
    }

    pub fn send_err(sock: &ZSock, err: Error) -> Result<()> {
        Self::send(sock, vec!["Err", err.description()])
    }

    fn sock_is_writeable(sock: &ZSock) -> Result<bool> {
        match try!(sock.type_str()).unwrap() {
            "PAIR" |
            "PUB" |
            "REQ" |
            "REP" |
            "DEALER" |
            "ROUTER" |
            "PUSH" |
            "XSUB" => Ok(true),

            "SUB" |
            "PULL" |
            "XPUB" => Ok(false),

            _ => unimplemented!(),
        }
    }
}

#[cfg(test)]
mod tests {
    use czmq::{ZMsg, ZSock, zsys_init};
    use super::*;

    /// Test providing 1 and only 1 arg
    #[test]
    fn test_recv_args_ok_eq() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_ok_eq").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_ok_eq").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 1, Some(1), true).unwrap();
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "0");
    }

    /// Test providing 1 or more args
    #[test]
    fn test_recv_args_ok_range() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_ok_range").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_ok_range").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 1, Some(2), true).unwrap();
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "0");
    }

    /// Test providing 1+ args
    #[test]
    fn test_recv_args_ok_variable() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_ok_variable").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_ok_variable").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 2, None, true).unwrap();
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "0");
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "1");
        assert_eq!(rcv_msg.popstr().unwrap().unwrap(), "2");
    }

    /// Test failing less than 3 args
    #[test]
    fn test_recv_args_err_min() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_err_min").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_err_min").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.addstr("1").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 3, None, true);
        assert!(rcv_msg.is_err());
    }

    /// Test failing more than 1 arg
    #[test]
    fn test_recv_args_err_max() {
        zsys_init();

        let server = ZSock::new_rep("inproc://msg_recv_args_err_max").unwrap();
        let client = ZSock::new_req("inproc://msg_recv_args_err_max").unwrap();

        let msg = ZMsg::new();
        msg.addstr("0").unwrap();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.send(&client).unwrap();

        let rcv_msg = Msg::expect_recv(&server, 0, Some(1), true);
        assert!(rcv_msg.is_err());
    }
}
<|endoftext|>


<commit_before>use std::error;
use std::fmt;
use std::io;
use std::result;

use ByteOrder;

/// A short-hand for `result::Result<T, byteorder::Error>`.
pub type Result<T> = result::Result<T, Error>;

/// An error type for reading bytes.
///
/// This is a thin wrapper over the standard `io::Error` type. Namely, it
/// adds one additional error case: an unexpected EOF.
///
/// Note that this error is also used for the `write` methods to keep things
/// consistent.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Error {
    /// An unexpected EOF.
    ///
    /// This occurs when a call to the underlying reader returns `0` bytes,
    /// but more bytes are required to decode a meaningful value.
    UnexpectedEOF,
    /// Any underlying IO error that occurs while reading bytes.
    Io(io::Error),
}

impl error::FromError<io::Error> for Error {
    fn from_error(err: io::Error) -> Error { Error::Io(err) }
}

impl error::FromError<Error> for io::Error {
    fn from_error(err: Error) -> io::Error {
        match err {
            Error::Io(err) => err,
            Error::UnexpectedEOF => io::Error::new(io::ErrorKind::Other,
                                                   "unexpected EOF", None)
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::UnexpectedEOF => write!(f, "Unexpected end of file."),
            Error::Io(ref err) => err.fmt(f),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::UnexpectedEOF => "Unexpected end of file.",
            Error::Io(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::UnexpectedEOF => None,
            Error::Io(ref err) => err.cause(),
        }
    }
}

/// Extends `Read` with methods for reading numbers. (For `std::io`.)
///
/// Most of the methods defined here have an unconstrained type parameter that
/// must be explicitly instantiated. Typically, it is instantiated with either
/// the `BigEndian` or `LittleEndian` types defined in this crate.
///
/// # Examples
///
/// Read unsigned 16 bit big-endian integers from a `Read`:
///
/// ```rust
/// use std::io::Cursor;
/// use byteorder::{BigEndian, ReadBytesExt};
///
/// let mut rdr = Cursor::new(vec![2, 5, 3, 0]);
/// assert_eq!(517, rdr.read_u16::<BigEndian>().unwrap());
/// assert_eq!(768, rdr.read_u16::<BigEndian>().unwrap());
/// ```
pub trait ReadBytesExt: io::Read {
    /// Reads an unsigned 8 bit integer from the underlying reader.
    ///
    /// Note that since this reads a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn read_u8(&mut self) -> Result<u8> {
        let mut buf = [0; 1];
        try!(read_full(self, &mut buf));
        Ok(buf[0])
    }

    /// Reads a signed 8 bit integer from the underlying reader.
    ///
    /// Note that since this reads a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn read_i8(&mut self) -> Result<i8> {
        let mut buf = [0; 1];
        try!(read_full(self, &mut buf));
        Ok(buf[0] as i8)
    }

    /// Reads an unsigned 16 bit integer from the underlying reader.
    fn read_u16<T: ByteOrder>(&mut self) -> Result<u16> {
        let mut buf = [0; 2];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_u16(&buf))
    }

    /// Reads a signed 16 bit integer from the underlying reader.
    fn read_i16<T: ByteOrder>(&mut self) -> Result<i16> {
        let mut buf = [0; 2];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_i16(&buf))
    }

    /// Reads an unsigned 32 bit integer from the underlying reader.
    fn read_u32<T: ByteOrder>(&mut self) -> Result<u32> {
        let mut buf = [0; 4];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_u32(&buf))
    }

    /// Reads a signed 32 bit integer from the underlying reader.
    fn read_i32<T: ByteOrder>(&mut self) -> Result<i32> {
        let mut buf = [0; 4];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_i32(&buf))
    }

    /// Reads an unsigned 64 bit integer from the underlying reader.
    fn read_u64<T: ByteOrder>(&mut self) -> Result<u64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_u64(&buf))
    }

    /// Reads a signed 64 bit integer from the underlying reader.
    fn read_i64<T: ByteOrder>(&mut self) -> Result<i64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_i64(&buf))
    }

    /// Reads an unsigned n-bytes integer from the underlying reader.
    fn read_uint<T: ByteOrder>(&mut self, nbytes: usize) -> Result<u64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf[0..nbytes]));
        Ok(<T as ByteOrder>::read_uint(&buf, nbytes))
    }

    /// Reads a signed n-bytes integer from the underlying reader.
    fn read_int<T: ByteOrder>(&mut self, nbytes: usize) -> Result<i64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf[0..nbytes]));
        Ok(<T as ByteOrder>::read_int(&buf, nbytes))
    }

    /// Reads a IEEE754 single-precision (4 bytes) floating point number from
    /// the underlying reader.
    fn read_f32<T: ByteOrder>(&mut self) -> Result<f32> {
        let mut buf = [0; 4];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_f32(&buf))
    }

    /// Reads a IEEE754 double-precision (8 bytes) floating point number from
    /// the underlying reader.
    fn read_f64<T: ByteOrder>(&mut self) -> Result<f64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_f64(&buf))
    }
}

/// All types that implement `Read` get methods defined in `ReadBytesExt`
/// for free.
impl<R: io::Read> ReadBytesExt for R {}

fn read_full<R: io::Read + ?Sized>(rdr: &mut R, buf: &mut [u8]) -> Result<()> {
    let mut nread = 0usize;
    while nread < buf.len() {
        match try!(rdr.read(&mut buf[nread..])) {
            0 => return Err(Error::UnexpectedEOF),
            n => nread += n,
        }
    }
    Ok(())
}

fn write_all<W: io::Write + ?Sized>(wtr: &mut W, buf: &[u8]) -> Result<()> {
    wtr.write_all(buf).map_err(error::FromError::from_error)
}

/// Extends `Write` with methods for writing numbers. (For `std::io`.)
///
/// Most of the methods defined here have an unconstrained type parameter that
/// must be explicitly instantiated. Typically, it is instantiated with either
/// the `BigEndian` or `LittleEndian` types defined in this crate.
///
/// # Examples
///
/// Write unsigned 16 bit big-endian integers to a `Write`:
///
/// ```rust
/// use byteorder::{BigEndian, WriteBytesExt};
///
/// let mut wtr = vec![];
/// wtr.write_u16::<BigEndian>(517).unwrap();
/// wtr.write_u16::<BigEndian>(768).unwrap();
/// assert_eq!(wtr, vec![2, 5, 3, 0]);
/// ```
pub trait WriteBytesExt: io::Write {
    /// Writes an unsigned 8 bit integer to the underlying writer.
    ///
    /// Note that since this writes a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn write_u8(&mut self, n: u8) -> Result<()> {
        write_all(self, &[n])
    }

    /// Writes a signed 8 bit integer to the underlying writer.
    ///
    /// Note that since this writes a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn write_i8(&mut self, n: i8) -> Result<()> {
        write_all(self, &[n as u8])
    }

    /// Writes an unsigned 16 bit integer to the underlying writer.
    fn write_u16<T: ByteOrder>(&mut self, n: u16) -> Result<()> {
        let mut buf = [0; 2];
        <T as ByteOrder>::write_u16(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a signed 16 bit integer to the underlying writer.
    fn write_i16<T: ByteOrder>(&mut self, n: i16) -> Result<()> {
        let mut buf = [0; 2];
        <T as ByteOrder>::write_i16(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes an unsigned 32 bit integer to the underlying writer.
    fn write_u32<T: ByteOrder>(&mut self, n: u32) -> Result<()> {
        let mut buf = [0; 4];
        <T as ByteOrder>::write_u32(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a signed 32 bit integer to the underlying writer.
    fn write_i32<T: ByteOrder>(&mut self, n: i32) -> Result<()> {
        let mut buf = [0; 4];
        <T as ByteOrder>::write_i32(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes an unsigned 64 bit integer to the underlying writer.
    fn write_u64<T: ByteOrder>(&mut self, n: u64) -> Result<()> {
        let mut buf = [0; 8];
        <T as ByteOrder>::write_u64(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a signed 64 bit integer to the underlying writer.
    fn write_i64<T: ByteOrder>(&mut self, n: i64) -> Result<()> {
        let mut buf = [0; 8];
        <T as ByteOrder>::write_i64(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a IEEE754 single-precision (4 bytes) floating point number to
    /// the underlying writer.
    fn write_f32<T: ByteOrder>(&mut self, n: f32) -> Result<()> {
        let mut buf = [0; 4];
        <T as ByteOrder>::write_f32(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a IEEE754 double-precision (8 bytes) floating point number to
    /// the underlying writer.
    fn write_f64<T: ByteOrder>(&mut self, n: f64) -> Result<()> {
        let mut buf = [0; 8];
        <T as ByteOrder>::write_f64(&mut buf, n);
        write_all(self, &buf)
    }
}

/// All types that implement `Write` get methods defined in `WriteBytesExt`
/// for free.
impl<W: io::Write> WriteBytesExt for W {}
<commit_msg>Retry on interrupt.<commit_after>use std::error;
use std::fmt;
use std::io;
use std::result;

use ByteOrder;

/// A short-hand for `result::Result<T, byteorder::Error>`.
pub type Result<T> = result::Result<T, Error>;

/// An error type for reading bytes.
///
/// This is a thin wrapper over the standard `io::Error` type. Namely, it
/// adds one additional error case: an unexpected EOF.
///
/// Note that this error is also used for the `write` methods to keep things
/// consistent.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Error {
    /// An unexpected EOF.
    ///
    /// This occurs when a call to the underlying reader returns `0` bytes,
    /// but more bytes are required to decode a meaningful value.
    UnexpectedEOF,
    /// Any underlying IO error that occurs while reading bytes.
    Io(io::Error),
}

impl error::FromError<io::Error> for Error {
    fn from_error(err: io::Error) -> Error { Error::Io(err) }
}

impl error::FromError<Error> for io::Error {
    fn from_error(err: Error) -> io::Error {
        match err {
            Error::Io(err) => err,
            Error::UnexpectedEOF => io::Error::new(io::ErrorKind::Other,
                                                   "unexpected EOF", None)
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::UnexpectedEOF => write!(f, "Unexpected end of file."),
            Error::Io(ref err) => err.fmt(f),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::UnexpectedEOF => "Unexpected end of file.",
            Error::Io(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::UnexpectedEOF => None,
            Error::Io(ref err) => err.cause(),
        }
    }
}

/// Extends `Read` with methods for reading numbers. (For `std::io`.)
///
/// Most of the methods defined here have an unconstrained type parameter that
/// must be explicitly instantiated. Typically, it is instantiated with either
/// the `BigEndian` or `LittleEndian` types defined in this crate.
///
/// # Examples
///
/// Read unsigned 16 bit big-endian integers from a `Read`:
///
/// ```rust
/// use std::io::Cursor;
/// use byteorder::{BigEndian, ReadBytesExt};
///
/// let mut rdr = Cursor::new(vec![2, 5, 3, 0]);
/// assert_eq!(517, rdr.read_u16::<BigEndian>().unwrap());
/// assert_eq!(768, rdr.read_u16::<BigEndian>().unwrap());
/// ```
pub trait ReadBytesExt: io::Read {
    /// Reads an unsigned 8 bit integer from the underlying reader.
    ///
    /// Note that since this reads a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn read_u8(&mut self) -> Result<u8> {
        let mut buf = [0; 1];
        try!(read_full(self, &mut buf));
        Ok(buf[0])
    }

    /// Reads a signed 8 bit integer from the underlying reader.
    ///
    /// Note that since this reads a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn read_i8(&mut self) -> Result<i8> {
        let mut buf = [0; 1];
        try!(read_full(self, &mut buf));
        Ok(buf[0] as i8)
    }

    /// Reads an unsigned 16 bit integer from the underlying reader.
    fn read_u16<T: ByteOrder>(&mut self) -> Result<u16> {
        let mut buf = [0; 2];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_u16(&buf))
    }

    /// Reads a signed 16 bit integer from the underlying reader.
    fn read_i16<T: ByteOrder>(&mut self) -> Result<i16> {
        let mut buf = [0; 2];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_i16(&buf))
    }

    /// Reads an unsigned 32 bit integer from the underlying reader.
    fn read_u32<T: ByteOrder>(&mut self) -> Result<u32> {
        let mut buf = [0; 4];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_u32(&buf))
    }

    /// Reads a signed 32 bit integer from the underlying reader.
    fn read_i32<T: ByteOrder>(&mut self) -> Result<i32> {
        let mut buf = [0; 4];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_i32(&buf))
    }

    /// Reads an unsigned 64 bit integer from the underlying reader.
    fn read_u64<T: ByteOrder>(&mut self) -> Result<u64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_u64(&buf))
    }

    /// Reads a signed 64 bit integer from the underlying reader.
    fn read_i64<T: ByteOrder>(&mut self) -> Result<i64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_i64(&buf))
    }

    /// Reads an unsigned n-bytes integer from the underlying reader.
    fn read_uint<T: ByteOrder>(&mut self, nbytes: usize) -> Result<u64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf[0..nbytes]));
        Ok(<T as ByteOrder>::read_uint(&buf, nbytes))
    }

    /// Reads a signed n-bytes integer from the underlying reader.
    fn read_int<T: ByteOrder>(&mut self, nbytes: usize) -> Result<i64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf[0..nbytes]));
        Ok(<T as ByteOrder>::read_int(&buf, nbytes))
    }

    /// Reads a IEEE754 single-precision (4 bytes) floating point number from
    /// the underlying reader.
    fn read_f32<T: ByteOrder>(&mut self) -> Result<f32> {
        let mut buf = [0; 4];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_f32(&buf))
    }

    /// Reads a IEEE754 double-precision (8 bytes) floating point number from
    /// the underlying reader.
    fn read_f64<T: ByteOrder>(&mut self) -> Result<f64> {
        let mut buf = [0; 8];
        try!(read_full(self, &mut buf));
        Ok(<T as ByteOrder>::read_f64(&buf))
    }
}

/// All types that implement `Read` get methods defined in `ReadBytesExt`
/// for free.
impl<R: io::Read> ReadBytesExt for R {}

fn read_full<R: io::Read + ?Sized>(rdr: &mut R, buf: &mut [u8]) -> Result<()> {
    let mut nread = 0usize;
    while nread < buf.len() {
        match rdr.read(&mut buf[nread..]) {
            Ok(0) => return Err(Error::UnexpectedEOF),
            Ok(n) => nread += n,
            Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {},
            Err(e) => return Err(error::FromError::from_error(e))
        }
    }
    Ok(())
}

fn write_all<W: io::Write + ?Sized>(wtr: &mut W, buf: &[u8]) -> Result<()> {
    wtr.write_all(buf).map_err(error::FromError::from_error)
}

/// Extends `Write` with methods for writing numbers. (For `std::io`.)
///
/// Most of the methods defined here have an unconstrained type parameter that
/// must be explicitly instantiated. Typically, it is instantiated with either
/// the `BigEndian` or `LittleEndian` types defined in this crate.
///
/// # Examples
///
/// Write unsigned 16 bit big-endian integers to a `Write`:
///
/// ```rust
/// use byteorder::{BigEndian, WriteBytesExt};
///
/// let mut wtr = vec![];
/// wtr.write_u16::<BigEndian>(517).unwrap();
/// wtr.write_u16::<BigEndian>(768).unwrap();
/// assert_eq!(wtr, vec![2, 5, 3, 0]);
/// ```
pub trait WriteBytesExt: io::Write {
    /// Writes an unsigned 8 bit integer to the underlying writer.
    ///
    /// Note that since this writes a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn write_u8(&mut self, n: u8) -> Result<()> {
        write_all(self, &[n])
    }

    /// Writes a signed 8 bit integer to the underlying writer.
    ///
    /// Note that since this writes a single byte, no byte order conversions
    /// are used. It is included for completeness.
    fn write_i8(&mut self, n: i8) -> Result<()> {
        write_all(self, &[n as u8])
    }

    /// Writes an unsigned 16 bit integer to the underlying writer.
    fn write_u16<T: ByteOrder>(&mut self, n: u16) -> Result<()> {
        let mut buf = [0; 2];
        <T as ByteOrder>::write_u16(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a signed 16 bit integer to the underlying writer.
    fn write_i16<T: ByteOrder>(&mut self, n: i16) -> Result<()> {
        let mut buf = [0; 2];
        <T as ByteOrder>::write_i16(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes an unsigned 32 bit integer to the underlying writer.
    fn write_u32<T: ByteOrder>(&mut self, n: u32) -> Result<()> {
        let mut buf = [0; 4];
        <T as ByteOrder>::write_u32(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a signed 32 bit integer to the underlying writer.
    fn write_i32<T: ByteOrder>(&mut self, n: i32) -> Result<()> {
        let mut buf = [0; 4];
        <T as ByteOrder>::write_i32(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes an unsigned 64 bit integer to the underlying writer.
    fn write_u64<T: ByteOrder>(&mut self, n: u64) -> Result<()> {
        let mut buf = [0; 8];
        <T as ByteOrder>::write_u64(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a signed 64 bit integer to the underlying writer.
    fn write_i64<T: ByteOrder>(&mut self, n: i64) -> Result<()> {
        let mut buf = [0; 8];
        <T as ByteOrder>::write_i64(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a IEEE754 single-precision (4 bytes) floating point number to
    /// the underlying writer.
    fn write_f32<T: ByteOrder>(&mut self, n: f32) -> Result<()> {
        let mut buf = [0; 4];
        <T as ByteOrder>::write_f32(&mut buf, n);
        write_all(self, &buf)
    }

    /// Writes a IEEE754 double-precision (8 bytes) floating point number to
    /// the underlying writer.
    fn write_f64<T: ByteOrder>(&mut self, n: f64) -> Result<()> {
        let mut buf = [0; 8];
        <T as ByteOrder>::write_f64(&mut buf, n);
        write_all(self, &buf)
    }
}

/// All types that implement `Write` get methods defined in `WriteBytesExt`
/// for free.
impl<W: io::Write> WriteBytesExt for W {}
<|endoftext|>


<commit_before>//! Decoding and Encoding of PNG Images
//!
//! PNG (Portable Network Graphics) is an image format that supports lossless compression.
//!
//! # Related Links
//! * <http://www.w3.org/TR/PNG/> - The PNG Specification
//!

extern crate png;

use std::io::{self, Read, Write};

use color::ColorType;
use image::{ImageDecoder, ImageError, ImageResult};

/// PNG Reader
///
/// This reader will try to read the png one row at a time,
/// however for interlaced png files this is not posible and
/// these are therefore readed at once.
pub struct PNGReader<R: Read> {
    reader: png::Reader<R>,
    buffer: Vec<u8>,
    index: usize,
}

impl<R: Read> PNGReader<R> {
    fn new(mut reader: png::Reader<R>) -> ImageResult<PNGReader<R>> {
        let len = reader.output_buffer_size();
        // Since interlaced images do not come in
        // scanline order it is almost impossible to
        // read them in a streaming fashion, however
        // this shouldn't be a too big of a problem
        // as most interlaced images should fit in memory.
        let buffer = if reader.info().interlaced {
            let mut buffer = vec![0; len];
            reader.next_frame(&mut buffer)?;
            buffer
        } else {
            Vec::new()
        };

        Ok(PNGReader {
            reader,
            buffer,
            index: 0,
        })
    }
}

impl<R: Read> Read for PNGReader<R> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        // io::Write::write for slice cannot fail
        let readed = buf.write(&self.buffer[self.index..]).unwrap();

        let mut bytes = readed;
        self.index += readed;

        while self.index >= self.buffer.len() {
            match self.reader.next_row()? {
                Some(row) => {
                    // Faster to copy directly to external buffer
                    let readed  = buf.write(row).unwrap();
                    bytes += readed;

                    self.buffer = (&row[readed..]).to_owned();
                    self.index = 0;
                }
                None => return Ok(bytes)
            }
        }

        Ok(bytes)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut bytes = self.buffer.len();
        buf.extend_from_slice(&self.buffer);
        self.buffer = Vec::new();
        self.index = 0;

        while let Some(row) = self.reader.next_row()? {
            buf.extend_from_slice(row);
            bytes += row.len();
        }

        Ok(bytes)
    }
}

/// PNG decoder
pub struct PNGDecoder<R: Read> {
    colortype: ColorType,
    reader: png::Reader<R>,
}

impl<R: Read> PNGDecoder<R> {
    /// Creates a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PNGDecoder<R>> {
        let limits = png::Limits {
            bytes: usize::max_value(),
        };
        let decoder = png::Decoder::new_with_limits(r, limits);
        let (_, mut reader) = decoder.read_info()?;
        let colortype = reader.output_color_type().into();

        Ok(PNGDecoder { colortype, reader })
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for PNGDecoder<R> {
    type Reader = PNGReader<R>;

    fn dimensions(&self) -> (u64, u64) {
        let (w, h) = self.reader.info().size();
        (u64::from(w), u64::from(h))
    }

    fn colortype(&self) -> ColorType {
        self.colortype
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        PNGReader::new(self.reader)
    }

    fn read_image(mut self) -> ImageResult<Vec<u8>> {
        // This should be slightly faster than the default implementation
        let mut data = vec![0; self.reader.output_buffer_size()];
        self.reader.next_frame(&mut data)?;
        Ok(data)
    }

    fn scanline_bytes(&self) -> u64 {
        let width = self.reader.info().width;
        self.reader.output_line_size(width) as u64
    }
}

/// PNG encoder
pub struct PNGEncoder<W: Write> {
    w: W,
}

impl<W: Write> PNGEncoder<W> {
    /// Create a new encoder that writes its output to ```w```
    pub fn new(w: W) -> PNGEncoder<W> {
        PNGEncoder { w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```
    pub fn encode(self, data: &[u8], width: u32, height: u32, color: ColorType) -> io::Result<()> {
        let (ct, bits) = match color {
            ColorType::Gray(bits) => Some((png::ColorType::Grayscale, bits)),
            ColorType::RGB(bits) => Some((png::ColorType::RGB, bits)),
            ColorType::Palette(bits) => Some((png::ColorType::Indexed, bits)),
            ColorType::GrayA(bits) => Some((png::ColorType::GrayscaleAlpha, bits)),
            ColorType::RGBA(bits) => Some((png::ColorType::RGBA, bits)),
            _ => None,
        }
        .and_then(|(ct, bits)| Some((ct, png::BitDepth::from_u8(bits)?)))
        .ok_or(io::Error::new(io::ErrorKind::InvalidInput, ImageError::UnsupportedColor(color)))?;

        let mut encoder = png::Encoder::new(self.w, width, height);
        encoder.set_color(ct);
        encoder.set_depth(bits);
        let mut writer = encoder.write_header()?;
        writer.write_image_data(data).map_err(|e| e.into())
    }
}

impl From<(png::ColorType, png::BitDepth)> for ColorType {
    fn from((ct, bits): (png::ColorType, png::BitDepth)) -> ColorType {
        use self::png::ColorType::*;
        let bits = bits as u8;
        match ct {
            Grayscale => ColorType::Gray(bits),
            RGB => ColorType::RGB(bits),
            Indexed => ColorType::Palette(bits),
            GrayscaleAlpha => ColorType::GrayA(bits),
            RGBA => ColorType::RGBA(bits),
        }
    }
}

impl From<ColorType> for (png::ColorType, png::BitDepth) {
    fn from(ct: ColorType) -> (png::ColorType, png::BitDepth) {
        use self::png::ColorType::*;
        let (ct, bits) = match ct {
            ColorType::Gray(bits) => (Grayscale, bits),
            ColorType::RGB(bits) => (RGB, bits),
            ColorType::Palette(bits) => (Indexed, bits),
            ColorType::GrayA(bits) => (GrayscaleAlpha, bits),
            ColorType::RGBA(bits) => (RGBA, bits),
            ColorType::BGRA(bits) => (RGBA, bits),
            ColorType::BGR(bits) => (RGB, bits),
        };
        (ct, png::BitDepth::from_u8(bits).unwrap())
    }
}

impl From<png::DecodingError> for ImageError {
    fn from(err: png::DecodingError) -> ImageError {
        use self::png::DecodingError::*;
        match err {
            IoError(err) => ImageError::IoError(err),
            Format(desc) => ImageError::FormatError(desc.into_owned()),
            InvalidSignature => ImageError::FormatError("invalid signature".into()),
            CrcMismatch { .. } => ImageError::FormatError("CRC error".into()),
            Other(desc) => ImageError::FormatError(desc.into_owned()),
            CorruptFlateStream => {
                ImageError::FormatError("compressed data stream corrupted".into())
            }
            LimitsExceeded => ImageError::InsufficientMemory,
        }
    }
}

#[cfg(test)]
mod tests {
    use image::ImageDecoder;
    use std::io::Read;
    use super::*;

    #[test]
    fn ensure_no_decoder_off_by_one() {
        let dec = PNGDecoder::new(std::fs::File::open("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap())
            .expect("Unable to read PNG file (does it exist?)");

        assert_eq![(2000, 1000), dec.dimensions()];

        assert_eq![
            ColorType::RGB(8),
            dec.colortype(),
            "Image MUST have the RGB(8) format"
        ];

        let correct_bytes = dec
            .into_reader()
            .expect("Unable to read file")
            .bytes()
            .map(|x| x.expect("Unable to read byte"))
            .collect::<Vec<u8>>();

        assert_eq![6_000_000, correct_bytes.len()];
    }
}
<commit_msg>Add a FIXME to change the error type<commit_after>//! Decoding and Encoding of PNG Images
//!
//! PNG (Portable Network Graphics) is an image format that supports lossless compression.
//!
//! # Related Links
//! * <http://www.w3.org/TR/PNG/> - The PNG Specification
//!

extern crate png;

use std::io::{self, Read, Write};

use color::ColorType;
use image::{ImageDecoder, ImageError, ImageResult};

/// PNG Reader
///
/// This reader will try to read the png one row at a time,
/// however for interlaced png files this is not posible and
/// these are therefore readed at once.
pub struct PNGReader<R: Read> {
    reader: png::Reader<R>,
    buffer: Vec<u8>,
    index: usize,
}

impl<R: Read> PNGReader<R> {
    fn new(mut reader: png::Reader<R>) -> ImageResult<PNGReader<R>> {
        let len = reader.output_buffer_size();
        // Since interlaced images do not come in
        // scanline order it is almost impossible to
        // read them in a streaming fashion, however
        // this shouldn't be a too big of a problem
        // as most interlaced images should fit in memory.
        let buffer = if reader.info().interlaced {
            let mut buffer = vec![0; len];
            reader.next_frame(&mut buffer)?;
            buffer
        } else {
            Vec::new()
        };

        Ok(PNGReader {
            reader,
            buffer,
            index: 0,
        })
    }
}

impl<R: Read> Read for PNGReader<R> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        // io::Write::write for slice cannot fail
        let readed = buf.write(&self.buffer[self.index..]).unwrap();

        let mut bytes = readed;
        self.index += readed;

        while self.index >= self.buffer.len() {
            match self.reader.next_row()? {
                Some(row) => {
                    // Faster to copy directly to external buffer
                    let readed  = buf.write(row).unwrap();
                    bytes += readed;

                    self.buffer = (&row[readed..]).to_owned();
                    self.index = 0;
                }
                None => return Ok(bytes)
            }
        }

        Ok(bytes)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut bytes = self.buffer.len();
        buf.extend_from_slice(&self.buffer);
        self.buffer = Vec::new();
        self.index = 0;

        while let Some(row) = self.reader.next_row()? {
            buf.extend_from_slice(row);
            bytes += row.len();
        }

        Ok(bytes)
    }
}

/// PNG decoder
pub struct PNGDecoder<R: Read> {
    colortype: ColorType,
    reader: png::Reader<R>,
}

impl<R: Read> PNGDecoder<R> {
    /// Creates a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PNGDecoder<R>> {
        let limits = png::Limits {
            bytes: usize::max_value(),
        };
        let decoder = png::Decoder::new_with_limits(r, limits);
        let (_, mut reader) = decoder.read_info()?;
        let colortype = reader.output_color_type().into();

        Ok(PNGDecoder { colortype, reader })
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for PNGDecoder<R> {
    type Reader = PNGReader<R>;

    fn dimensions(&self) -> (u64, u64) {
        let (w, h) = self.reader.info().size();
        (u64::from(w), u64::from(h))
    }

    fn colortype(&self) -> ColorType {
        self.colortype
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        PNGReader::new(self.reader)
    }

    fn read_image(mut self) -> ImageResult<Vec<u8>> {
        // This should be slightly faster than the default implementation
        let mut data = vec![0; self.reader.output_buffer_size()];
        self.reader.next_frame(&mut data)?;
        Ok(data)
    }

    fn scanline_bytes(&self) -> u64 {
        let width = self.reader.info().width;
        self.reader.output_line_size(width) as u64
    }
}

/// PNG encoder
pub struct PNGEncoder<W: Write> {
    w: W,
}

impl<W: Write> PNGEncoder<W> {
    /// Create a new encoder that writes its output to ```w```
    pub fn new(w: W) -> PNGEncoder<W> {
        PNGEncoder { w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```
    pub fn encode(self, data: &[u8], width: u32, height: u32, color: ColorType) -> io::Result<()> {
        let (ct, bits) = match color {
            ColorType::Gray(bits) => Some((png::ColorType::Grayscale, bits)),
            ColorType::RGB(bits) => Some((png::ColorType::RGB, bits)),
            ColorType::Palette(bits) => Some((png::ColorType::Indexed, bits)),
            ColorType::GrayA(bits) => Some((png::ColorType::GrayscaleAlpha, bits)),
            ColorType::RGBA(bits) => Some((png::ColorType::RGBA, bits)),
            _ => None,
        }
        .and_then(|(ct, bits)| Some((ct, png::BitDepth::from_u8(bits)?)))
        // FIXME: After #1066 lands, return the ImageError directly.
        .ok_or(io::Error::new(io::ErrorKind::InvalidInput, ImageError::UnsupportedColor(color)))?;

        let mut encoder = png::Encoder::new(self.w, width, height);
        encoder.set_color(ct);
        encoder.set_depth(bits);
        let mut writer = encoder.write_header()?;
        writer.write_image_data(data).map_err(|e| e.into())
    }
}

impl From<(png::ColorType, png::BitDepth)> for ColorType {
    fn from((ct, bits): (png::ColorType, png::BitDepth)) -> ColorType {
        use self::png::ColorType::*;
        let bits = bits as u8;
        match ct {
            Grayscale => ColorType::Gray(bits),
            RGB => ColorType::RGB(bits),
            Indexed => ColorType::Palette(bits),
            GrayscaleAlpha => ColorType::GrayA(bits),
            RGBA => ColorType::RGBA(bits),
        }
    }
}

impl From<ColorType> for (png::ColorType, png::BitDepth) {
    fn from(ct: ColorType) -> (png::ColorType, png::BitDepth) {
        use self::png::ColorType::*;
        let (ct, bits) = match ct {
            ColorType::Gray(bits) => (Grayscale, bits),
            ColorType::RGB(bits) => (RGB, bits),
            ColorType::Palette(bits) => (Indexed, bits),
            ColorType::GrayA(bits) => (GrayscaleAlpha, bits),
            ColorType::RGBA(bits) => (RGBA, bits),
            ColorType::BGRA(bits) => (RGBA, bits),
            ColorType::BGR(bits) => (RGB, bits),
        };
        (ct, png::BitDepth::from_u8(bits).unwrap())
    }
}

impl From<png::DecodingError> for ImageError {
    fn from(err: png::DecodingError) -> ImageError {
        use self::png::DecodingError::*;
        match err {
            IoError(err) => ImageError::IoError(err),
            Format(desc) => ImageError::FormatError(desc.into_owned()),
            InvalidSignature => ImageError::FormatError("invalid signature".into()),
            CrcMismatch { .. } => ImageError::FormatError("CRC error".into()),
            Other(desc) => ImageError::FormatError(desc.into_owned()),
            CorruptFlateStream => {
                ImageError::FormatError("compressed data stream corrupted".into())
            }
            LimitsExceeded => ImageError::InsufficientMemory,
        }
    }
}

#[cfg(test)]
mod tests {
    use image::ImageDecoder;
    use std::io::Read;
    use super::*;

    #[test]
    fn ensure_no_decoder_off_by_one() {
        let dec = PNGDecoder::new(std::fs::File::open("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap())
            .expect("Unable to read PNG file (does it exist?)");

        assert_eq![(2000, 1000), dec.dimensions()];

        assert_eq![
            ColorType::RGB(8),
            dec.colortype(),
            "Image MUST have the RGB(8) format"
        ];

        let correct_bytes = dec
            .into_reader()
            .expect("Unable to read file")
            .bytes()
            .map(|x| x.expect("Unable to read byte"))
            .collect::<Vec<u8>>();

        assert_eq![6_000_000, correct_bytes.len()];
    }
}
<|endoftext|>


<commit_before>use std::str;
use std::convert;
use std::ops;
use regex::Regex;

use {
    Result,
    Error,
    ErrorType
};

#[derive(Debug)]
pub enum Pos {
    Line(usize),
    Current,
    End
}

pub trait Converter<F, T> {
    fn convert(&self, from: F) -> T;
}

static POS_RE: &'static str = r"^((?P<current>\.)|(?P<end>\$)|(?P<line>\d+))$";
static RANGE_RE: &'static str = r"^((?P<all>%{1})|(?P<first>[.$0-9]+)(,(?P<second>[.$0-9]+))?)$";

impl str::FromStr for Pos {
    type Err = Error;
    fn from_str(s: &str) -> Result<Pos> {
        let re = Regex::new(POS_RE).unwrap();

        if let Some(captures) = re.captures(s) {

            if let Some(_) = captures.name("current") {
                return Ok(Pos::Current);
            }

            if let Some(_) = captures.name("end") {
                return Ok(Pos::End);
            }

            if let Some(line) = captures.name("line") {
                return  match line.parse() {
                    Ok(n) => Ok(Pos::Line(n)),
                    Err(_) => Err(Error::new(ErrorType::ParseError))
                };
            }
        }

        Err(Error::detailed(ErrorType::ParseError, s.to_string()))
    }
}

impl convert::From<Range> for Pos {
    fn from(r: Range) -> Pos {
        match r {
            Range::Line(p) => p,
            Range::Range(_, p) => p
        }
    }
}

#[derive(Debug)]
pub enum Range {
    Line(Pos),
    Range(Pos, Pos)
}

impl str::FromStr for Range {
    type Err = Error;

    fn from_str(s: &str) -> Result<Range> {
        let re = Regex::new(RANGE_RE).unwrap();

        if let Some(captures) = re.captures(s) {
            if let Some(_) = captures.name("all") {
                return Ok( Range::Range( Pos::Line(1), Pos::End ) );
            }

            if let Some(first) = captures.name("first") {

                let first_pos = try!(first.parse::<Pos>());

                if let Some(second) = captures.name("second") {
                    let second_pos = try!(second.parse::<Pos>());

                    return Ok(Range::Range( first_pos, second_pos ));
                }

                return Ok(Range::Line( first_pos ));
            }
        }

        Err(Error::detailed(ErrorType::ParseError, s.to_string()))
    }
}

impl Range {

    pub fn to_range<'a, C>(&'a self, conv: &C) -> ops::Range<usize>
        where C: Converter<&'a Pos, usize> {
        match *self {
            Range::Line(ref p) => {
                let pos = conv.convert(&p);
                ops::Range {
                    start: pos - 1,
                    end: pos
                }
            },

            Range::Range(ref f, ref t) => ops::Range {
                start: conv.convert(&f) - 1,
                end: conv.convert(&t)
            }
        }
    }

    pub fn current_line() -> Range {
        Range::Line(Pos::Current)
    }

}

<commit_msg>fixed line handling<commit_after>use std::str;
use std::convert;
use std::ops;
use regex::Regex;

use {
    Result,
    Error,
    ErrorType
};

#[derive(Debug)]
pub enum Pos {
    Line(usize),
    Current,
    End
}

pub trait Converter<F, T> {
    fn convert(&self, from: F) -> T;
}

static POS_RE: &'static str = r"^((?P<current>\.)|(?P<end>\$)|(?P<line>\d+))$";
static RANGE_RE: &'static str = r"^((?P<all>%{1})|(?P<first>[.$0-9]+)(,(?P<second>[.$0-9]+))?)$";

impl str::FromStr for Pos {
    type Err = Error;
    fn from_str(s: &str) -> Result<Pos> {
        let re = Regex::new(POS_RE).unwrap();

        if let Some(captures) = re.captures(s) {

            if let Some(_) = captures.name("current") {
                return Ok(Pos::Current);
            }

            if let Some(_) = captures.name("end") {
                return Ok(Pos::End);
            }

            if let Some(line) = captures.name("line") {
                return  match line.parse() {
                    Ok(n) => Ok(Pos::Line(n)),
                    Err(_) => Err(Error::new(ErrorType::ParseError))
                };
            }
        }

        Err(Error::detailed(ErrorType::ParseError, s.to_string()))
    }
}

impl convert::From<Range> for Pos {
    fn from(r: Range) -> Pos {
        match r {
            Range::Line(p) => p,
            Range::Range(_, p) => p
        }
    }
}

#[derive(Debug)]
pub enum Range {
    Line(Pos),
    Range(Pos, Pos)
}

impl str::FromStr for Range {
    type Err = Error;

    fn from_str(s: &str) -> Result<Range> {
        let re = Regex::new(RANGE_RE).unwrap();

        if let Some(captures) = re.captures(s) {
            if let Some(_) = captures.name("all") {
                return Ok( Range::Range( Pos::Line(1), Pos::End ) );
            }

            if let Some(first) = captures.name("first") {

                let first_pos = try!(first.parse::<Pos>());

                if let Some(second) = captures.name("second") {
                    let second_pos = try!(second.parse::<Pos>());

                    return Ok(Range::Range( first_pos, second_pos ));
                }

                return Ok(Range::Line( first_pos ));
            }
        }

        Err(Error::detailed(ErrorType::ParseError, s.to_string()))
    }
}

impl Range {

    pub fn to_range<'a, C>(&'a self, conv: &C) -> ops::Range<usize>
        where C: Converter<&'a Pos, usize> {
        match *self {
            Range::Line(ref p) => {
                let pos = conv.convert(&p);
                if pos == 0 {
                    ops::Range {
                        start: 0,
                        end: 0
                    }
                } else {
                    ops::Range {
                        start: pos - 1,
                        end: pos
                    }
                }
            },

            Range::Range(ref f, ref t) => ops::Range {
                start: conv.convert(&f) - 1,
                end: conv.convert(&t)
            }
        }
    }

    pub fn current_line() -> Range {
        Range::Line(Pos::Current)
    }

}

<|endoftext|>


<commit_before>//! Serialize a Rust data structure into bencode data.

mod string;

use std::str;
use std::mem;
use serde::ser;
use crate::error::{Error, Result};

/// A structure for serializing Rust values into bencode.
#[derive(Debug)]
pub struct Serializer {
    buf: Vec<u8>,
}

impl Serializer {
    /// Create a new serializer.
    pub fn new() -> Serializer {
        Serializer { buf: Vec::new() }
    }

    /// Consume the serializer and return the contents as a byte vector.
    pub fn into_vec(self) -> Vec<u8> {
        self.buf
    }

    fn push<T: AsRef<[u8]>>(&mut self, token: T) {
        self.buf.extend_from_slice(token.as_ref());
    }
}

impl AsRef<[u8]> for Serializer {
    fn as_ref(&self) -> &[u8] {
        self.buf.as_ref()
    }
}

impl<'a> ser::SerializeSeq for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_element<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        self.push("e");
        Ok(())
    }
}

impl<'a> ser::SerializeTuple for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_element<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        ser::SerializeSeq::end(self)
    }
}

impl<'a> ser::SerializeTupleStruct for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        ser::SerializeSeq::end(self)
    }
}

impl<'a> ser::SerializeTupleVariant for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        self.push("ee");
        Ok(())
    }
}

#[doc(hidden)]
// TODO: This should be pub(crate).
pub struct SerializeMap<'a> {
    ser: &'a mut Serializer,
    entries: Vec<(Vec<u8>, Vec<u8>)>,
    cur_key: Option<Vec<u8>>,
}

impl<'a> SerializeMap<'a> {
    pub fn new(ser: &'a mut Serializer, len: usize) -> SerializeMap {
        SerializeMap {
            ser: ser,
            entries: Vec::with_capacity(len),
            cur_key: None,
        }
    }

    fn end_map(&mut self) -> Result<()> {
        if self.cur_key.is_some() {
            return Err(Error::InvalidValue("`serialize_key` called without calling  `serialize_value`".to_string()));
        }
        let mut entries = mem::replace(&mut self.entries, Vec::new());
        entries.sort_by(|&(ref a, _), &(ref b, _)| a.cmp(b));
        self.ser.push("d");
        for (k, v) in entries {
            ser::Serializer::serialize_bytes(&mut *self.ser, k.as_ref())?;
            self.ser.push(v);
        }
        self.ser.push("e");
        Ok(())
    }
}

impl<'a> ser::SerializeMap for SerializeMap<'a> {
    type Ok = ();
    type Error = Error;
    fn serialize_key<T: ?Sized + ser::Serialize>(&mut self, key: &T) -> Result<()> {
        if self.cur_key.is_some() {
            return Err(Error::InvalidValue("`serialize_key` called multiple times without calling  `serialize_value`".to_string()));
        }
        self.cur_key = Some(key.serialize(&mut string::StringSerializer)?);
        Ok(())
    }
    fn serialize_value<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        let key = self.cur_key
            .take()
            .ok_or(Error::InvalidValue("`serialize_value` called without calling `serialize_key`"
                                           .to_string()))?;
        let mut ser = Serializer::new();
        value.serialize(&mut ser)?;
        let value = ser.into_vec();
        if !value.is_empty() {
            self.entries.push((key, value));
        }
        Ok(())
    }
    fn serialize_entry<K, V>(&mut self, key: &K, value: &V) -> Result<()>
        where K: ?Sized + ser::Serialize,
              V: ?Sized + ser::Serialize
    {
        if self.cur_key.is_some() {
            return Err(Error::InvalidValue("`serialize_key` called multiple times without calling  `serialize_value`".to_string()));
        }
        let key = key.serialize(&mut string::StringSerializer)?;
        let mut ser = Serializer::new();
        value.serialize(&mut ser)?;
        let value = ser.into_vec();
        if !value.is_empty() {
            self.entries.push((key, value));
        }
        Ok(())
    }
    fn end(mut self) -> Result<()> {
        self.end_map()
    }
}

impl<'a> ser::SerializeStruct for SerializeMap<'a> {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self,
                                                   key: &'static str,
                                                   value: &T)
                                                   -> Result<()> {
        ser::SerializeMap::serialize_entry(self, key, value)
    }
    fn end(mut self) -> Result<()> {
        self.end_map()
    }
}

impl<'a> ser::SerializeStructVariant for SerializeMap<'a> {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self,
                                                   key: &'static str,
                                                   value: &T)
                                                   -> Result<()> {
        ser::SerializeMap::serialize_entry(self, key, value)
    }
    fn end(mut self) -> Result<()> {
        self.end_map()?;
        self.ser.push("e");
        Ok(())
    }
}

impl<'a> ser::Serializer for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    type SerializeSeq = Self;
    type SerializeTuple = Self;
    type SerializeTupleStruct = Self;
    type SerializeTupleVariant = Self;
    type SerializeMap = SerializeMap<'a>;
    type SerializeStruct = SerializeMap<'a>;
    type SerializeStructVariant = SerializeMap<'a>;

    fn serialize_bool(self, value: bool) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i8(self, value: i8) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i16(self, value: i16) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i32(self, value: i32) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i64(self, value: i64) -> Result<()> {
        self.push("i");
        self.push(value.to_string());
        self.push("e");
        Ok(())
    }
    fn serialize_u8(self, value: u8) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_u16(self, value: u16) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_u32(self, value: u32) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_u64(self, value: u64) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_f32(self, _value: f32) -> Result<()> {
        Err(Error::InvalidValue("Cannot serialize f32".to_string()))
    }
    fn serialize_f64(self, _value: f64) -> Result<()> {
        Err(Error::InvalidValue("Cannot serialize f64".to_string()))
    }
    fn serialize_char(self, value: char) -> Result<()> {
        let mut buffer = [0; 4];
        self.serialize_bytes(value.encode_utf8(&mut buffer).as_bytes())
    }

    fn serialize_str(self, value: &str) -> Result<()> {
        self.serialize_bytes(value.as_bytes())
    }
    fn serialize_bytes(self, value: &[u8]) -> Result<()> {
        self.push(value.len().to_string());
        self.push(":");
        self.push(value);
        Ok(())
    }
    fn serialize_unit(self) -> Result<()> {
        Ok(())
    }
    fn serialize_unit_struct(self, _name: &'static str) -> Result<()> {
        self.serialize_unit()
    }
    fn serialize_unit_variant(self,
                              _name: &'static str,
                              _variant_index: u32,
                              variant: &'static str)
                              -> Result<()> {
        self.serialize_str(variant)
    }
    fn serialize_newtype_struct<T: ?Sized + ser::Serialize>(self,
                                                            _name: &'static str,
                                                            value: &T)
                                                            -> Result<()> {
        value.serialize(self)
    }
    fn serialize_newtype_variant<T: ?Sized + ser::Serialize>(self,
                                                             _name: &'static str,
                                                             _variant_index: u32,
                                                             variant: &'static str,
                                                             value: &T)
                                                             -> Result<()> {
        self.push("d");
        self.serialize_bytes(variant.as_bytes())?;
        value.serialize(&mut *self)?;
        self.push("e");
        Ok(())
    }
    fn serialize_none(self) -> Result<()> {
        Ok(())
    }
    fn serialize_some<T: ?Sized + ser::Serialize>(self, value: &T) -> Result<()> {
        value.serialize(self)
    }
    fn serialize_seq(self, _len: Option<usize>) -> Result<Self> {
        self.push("l");
        Ok(self)
    }
    fn serialize_tuple(self, size: usize) -> Result<Self> {
        self.serialize_seq(Some(size))
    }
    fn serialize_tuple_struct(self, _name: &'static str, len: usize) -> Result<Self> {
        self.serialize_seq(Some(len))
    }
    fn serialize_tuple_variant(self,
                               _name: &'static str,
                               _variant_index: u32,
                               variant: &'static str,
                               _len: usize)
                               -> Result<Self::SerializeTupleVariant> {
        self.push("d");
        self.serialize_bytes(variant.as_bytes())?;
        self.push("l");
        Ok(self)
    }
    fn serialize_map(self, len: Option<usize>) -> Result<Self::SerializeMap> {
        Ok(SerializeMap::new(self, len.unwrap_or(0)))
    }
    fn serialize_struct(self, _name: &'static str, len: usize) -> Result<Self::SerializeStruct> {
        self.serialize_map(Some(len))
    }
    fn serialize_struct_variant(self,
                                _name: &'static str,
                                _variant_index: u32,
                                variant: &'static str,
                                len: usize)
                                -> Result<Self::SerializeStructVariant> {
        self.push("d");
        self.serialize_bytes(variant.as_bytes())?;
        Ok(SerializeMap::new(self, len))
    }
}

/// Serialize the given data into a bencode byte vector.
///
/// # Examples
/// ```
/// # fn main() -> Result<(), serde_bencode::Error> {
/// use serde_derive::{Serialize, Deserialize};
///
/// #[derive(Serialize, Deserialize, PartialEq, Eq, Debug)]
/// struct Address {
///     street: String,
///     city: String,
/// }
///
/// let address = Address {
///     street: "1313 Webfoot Walk".to_string(),
///     city: "Duckburg, Calisota".to_string(),
/// };
///
/// let bytes = serde_bencode::to_bytes(&address)?;
/// assert_eq!(
///     String::from_utf8(bytes),
///     Ok("d4:city18:Duckburg, Calisota6:street17:1313 Webfoot Walke".to_string()),
/// );
/// # Ok(())
/// # }
/// ```
///
/// # Errors
///
/// Serialization can fail if `T`'s implementation of `Serialize` decides to fail or `T` contains
/// floating point values, which bencode cannot serialize.
pub fn to_bytes<T: ser::Serialize>(b: &T) -> Result<Vec<u8>> {
    let mut ser = Serializer::new();
    b.serialize(&mut ser)?;
    Ok(ser.into_vec())
}

/// Serialize the given data into a String of bencode.
///
/// # Examples
/// ```
/// # fn main() -> Result<(), serde_bencode::Error> {
/// use serde_derive::{Serialize, Deserialize};
///
/// #[derive(Serialize, Deserialize, PartialEq, Eq, Debug)]
/// struct Address {
///     street: String,
///     city: String,
/// }
///
/// let address = Address {
///     street: "1313 Webfoot Walk".to_string(),
///     city: "Duckburg, Calisota".to_string(),
/// };
///
/// assert_eq!(
///     serde_bencode::to_string(&address)?,
///     "d4:city18:Duckburg, Calisota6:street17:1313 Webfoot Walke".to_string(),
/// );
/// # Ok(())
/// # }
/// ```
///
/// # Errors
///
/// Serialization can fail if `T`'s implementation of `Serialize` decides to fail or `T` contains
/// floating point values, which bencode cannot serialize.
pub fn to_string<T: ser::Serialize>(b: &T) -> Result<String> {
    let mut ser = Serializer::new();
    b.serialize(&mut ser)?;
    str::from_utf8(ser.as_ref())
        .map(|s| s.to_string())
        .map_err(|_| Error::InvalidValue("Not an UTF-8".to_string()))
}
<commit_msg>fix unsigned overflow<commit_after>//! Serialize a Rust data structure into bencode data.

mod string;

use std::str;
use std::mem;
use serde::ser;
use crate::error::{Error, Result};

/// A structure for serializing Rust values into bencode.
#[derive(Debug)]
pub struct Serializer {
    buf: Vec<u8>,
}

impl Serializer {
    /// Create a new serializer.
    pub fn new() -> Serializer {
        Serializer { buf: Vec::new() }
    }

    /// Consume the serializer and return the contents as a byte vector.
    pub fn into_vec(self) -> Vec<u8> {
        self.buf
    }

    fn push<T: AsRef<[u8]>>(&mut self, token: T) {
        self.buf.extend_from_slice(token.as_ref());
    }
}

impl AsRef<[u8]> for Serializer {
    fn as_ref(&self) -> &[u8] {
        self.buf.as_ref()
    }
}

impl<'a> ser::SerializeSeq for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_element<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        self.push("e");
        Ok(())
    }
}

impl<'a> ser::SerializeTuple for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_element<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        ser::SerializeSeq::end(self)
    }
}

impl<'a> ser::SerializeTupleStruct for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        ser::SerializeSeq::end(self)
    }
}

impl<'a> ser::SerializeTupleVariant for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        value.serialize(&mut **self)
    }
    fn end(self) -> Result<()> {
        self.push("ee");
        Ok(())
    }
}

#[doc(hidden)]
// TODO: This should be pub(crate).
pub struct SerializeMap<'a> {
    ser: &'a mut Serializer,
    entries: Vec<(Vec<u8>, Vec<u8>)>,
    cur_key: Option<Vec<u8>>,
}

impl<'a> SerializeMap<'a> {
    pub fn new(ser: &'a mut Serializer, len: usize) -> SerializeMap {
        SerializeMap {
            ser: ser,
            entries: Vec::with_capacity(len),
            cur_key: None,
        }
    }

    fn end_map(&mut self) -> Result<()> {
        if self.cur_key.is_some() {
            return Err(Error::InvalidValue("`serialize_key` called without calling  `serialize_value`".to_string()));
        }
        let mut entries = mem::replace(&mut self.entries, Vec::new());
        entries.sort_by(|&(ref a, _), &(ref b, _)| a.cmp(b));
        self.ser.push("d");
        for (k, v) in entries {
            ser::Serializer::serialize_bytes(&mut *self.ser, k.as_ref())?;
            self.ser.push(v);
        }
        self.ser.push("e");
        Ok(())
    }
}

impl<'a> ser::SerializeMap for SerializeMap<'a> {
    type Ok = ();
    type Error = Error;
    fn serialize_key<T: ?Sized + ser::Serialize>(&mut self, key: &T) -> Result<()> {
        if self.cur_key.is_some() {
            return Err(Error::InvalidValue("`serialize_key` called multiple times without calling  `serialize_value`".to_string()));
        }
        self.cur_key = Some(key.serialize(&mut string::StringSerializer)?);
        Ok(())
    }
    fn serialize_value<T: ?Sized + ser::Serialize>(&mut self, value: &T) -> Result<()> {
        let key = self.cur_key
            .take()
            .ok_or(Error::InvalidValue("`serialize_value` called without calling `serialize_key`"
                                           .to_string()))?;
        let mut ser = Serializer::new();
        value.serialize(&mut ser)?;
        let value = ser.into_vec();
        if !value.is_empty() {
            self.entries.push((key, value));
        }
        Ok(())
    }
    fn serialize_entry<K, V>(&mut self, key: &K, value: &V) -> Result<()>
        where K: ?Sized + ser::Serialize,
              V: ?Sized + ser::Serialize
    {
        if self.cur_key.is_some() {
            return Err(Error::InvalidValue("`serialize_key` called multiple times without calling  `serialize_value`".to_string()));
        }
        let key = key.serialize(&mut string::StringSerializer)?;
        let mut ser = Serializer::new();
        value.serialize(&mut ser)?;
        let value = ser.into_vec();
        if !value.is_empty() {
            self.entries.push((key, value));
        }
        Ok(())
    }
    fn end(mut self) -> Result<()> {
        self.end_map()
    }
}

impl<'a> ser::SerializeStruct for SerializeMap<'a> {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self,
                                                   key: &'static str,
                                                   value: &T)
                                                   -> Result<()> {
        ser::SerializeMap::serialize_entry(self, key, value)
    }
    fn end(mut self) -> Result<()> {
        self.end_map()
    }
}

impl<'a> ser::SerializeStructVariant for SerializeMap<'a> {
    type Ok = ();
    type Error = Error;
    fn serialize_field<T: ?Sized + ser::Serialize>(&mut self,
                                                   key: &'static str,
                                                   value: &T)
                                                   -> Result<()> {
        ser::SerializeMap::serialize_entry(self, key, value)
    }
    fn end(mut self) -> Result<()> {
        self.end_map()?;
        self.ser.push("e");
        Ok(())
    }
}

impl<'a> ser::Serializer for &'a mut Serializer {
    type Ok = ();
    type Error = Error;
    type SerializeSeq = Self;
    type SerializeTuple = Self;
    type SerializeTupleStruct = Self;
    type SerializeTupleVariant = Self;
    type SerializeMap = SerializeMap<'a>;
    type SerializeStruct = SerializeMap<'a>;
    type SerializeStructVariant = SerializeMap<'a>;

    fn serialize_bool(self, value: bool) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i8(self, value: i8) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i16(self, value: i16) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i32(self, value: i32) -> Result<()> {
        self.serialize_i64(value as i64)
    }
    fn serialize_i64(self, value: i64) -> Result<()> {
        self.push("i");
        self.push(value.to_string());
        self.push("e");
        Ok(())
    }
    fn serialize_u8(self, value: u8) -> Result<()> {
        self.serialize_u64(value as u64)
    }
    fn serialize_u16(self, value: u16) -> Result<()> {
        self.serialize_u64(value as u64)
    }
    fn serialize_u32(self, value: u32) -> Result<()> {
        self.serialize_u64(value as u64)
    }
    fn serialize_u64(self, value: u64) -> Result<()> {
        self.push("i");
        self.push(value.to_string());
        self.push("e");
        Ok(())
    }
    fn serialize_f32(self, _value: f32) -> Result<()> {
        Err(Error::InvalidValue("Cannot serialize f32".to_string()))
    }
    fn serialize_f64(self, _value: f64) -> Result<()> {
        Err(Error::InvalidValue("Cannot serialize f64".to_string()))
    }
    fn serialize_char(self, value: char) -> Result<()> {
        let mut buffer = [0; 4];
        self.serialize_bytes(value.encode_utf8(&mut buffer).as_bytes())
    }

    fn serialize_str(self, value: &str) -> Result<()> {
        self.serialize_bytes(value.as_bytes())
    }
    fn serialize_bytes(self, value: &[u8]) -> Result<()> {
        self.push(value.len().to_string());
        self.push(":");
        self.push(value);
        Ok(())
    }
    fn serialize_unit(self) -> Result<()> {
        Ok(())
    }
    fn serialize_unit_struct(self, _name: &'static str) -> Result<()> {
        self.serialize_unit()
    }
    fn serialize_unit_variant(self,
                              _name: &'static str,
                              _variant_index: u32,
                              variant: &'static str)
                              -> Result<()> {
        self.serialize_str(variant)
    }
    fn serialize_newtype_struct<T: ?Sized + ser::Serialize>(self,
                                                            _name: &'static str,
                                                            value: &T)
                                                            -> Result<()> {
        value.serialize(self)
    }
    fn serialize_newtype_variant<T: ?Sized + ser::Serialize>(self,
                                                             _name: &'static str,
                                                             _variant_index: u32,
                                                             variant: &'static str,
                                                             value: &T)
                                                             -> Result<()> {
        self.push("d");
        self.serialize_bytes(variant.as_bytes())?;
        value.serialize(&mut *self)?;
        self.push("e");
        Ok(())
    }
    fn serialize_none(self) -> Result<()> {
        Ok(())
    }
    fn serialize_some<T: ?Sized + ser::Serialize>(self, value: &T) -> Result<()> {
        value.serialize(self)
    }
    fn serialize_seq(self, _len: Option<usize>) -> Result<Self> {
        self.push("l");
        Ok(self)
    }
    fn serialize_tuple(self, size: usize) -> Result<Self> {
        self.serialize_seq(Some(size))
    }
    fn serialize_tuple_struct(self, _name: &'static str, len: usize) -> Result<Self> {
        self.serialize_seq(Some(len))
    }
    fn serialize_tuple_variant(self,
                               _name: &'static str,
                               _variant_index: u32,
                               variant: &'static str,
                               _len: usize)
                               -> Result<Self::SerializeTupleVariant> {
        self.push("d");
        self.serialize_bytes(variant.as_bytes())?;
        self.push("l");
        Ok(self)
    }
    fn serialize_map(self, len: Option<usize>) -> Result<Self::SerializeMap> {
        Ok(SerializeMap::new(self, len.unwrap_or(0)))
    }
    fn serialize_struct(self, _name: &'static str, len: usize) -> Result<Self::SerializeStruct> {
        self.serialize_map(Some(len))
    }
    fn serialize_struct_variant(self,
                                _name: &'static str,
                                _variant_index: u32,
                                variant: &'static str,
                                len: usize)
                                -> Result<Self::SerializeStructVariant> {
        self.push("d");
        self.serialize_bytes(variant.as_bytes())?;
        Ok(SerializeMap::new(self, len))
    }
}

/// Serialize the given data into a bencode byte vector.
///
/// # Examples
/// ```
/// # fn main() -> Result<(), serde_bencode::Error> {
/// use serde_derive::{Serialize, Deserialize};
///
/// #[derive(Serialize, Deserialize, PartialEq, Eq, Debug)]
/// struct Address {
///     street: String,
///     city: String,
/// }
///
/// let address = Address {
///     street: "1313 Webfoot Walk".to_string(),
///     city: "Duckburg, Calisota".to_string(),
/// };
///
/// let bytes = serde_bencode::to_bytes(&address)?;
/// assert_eq!(
///     String::from_utf8(bytes),
///     Ok("d4:city18:Duckburg, Calisota6:street17:1313 Webfoot Walke".to_string()),
/// );
/// # Ok(())
/// # }
/// ```
///
/// # Errors
///
/// Serialization can fail if `T`'s implementation of `Serialize` decides to fail or `T` contains
/// floating point values, which bencode cannot serialize.
pub fn to_bytes<T: ser::Serialize>(b: &T) -> Result<Vec<u8>> {
    let mut ser = Serializer::new();
    b.serialize(&mut ser)?;
    Ok(ser.into_vec())
}

/// Serialize the given data into a String of bencode.
///
/// # Examples
/// ```
/// # fn main() -> Result<(), serde_bencode::Error> {
/// use serde_derive::{Serialize, Deserialize};
///
/// #[derive(Serialize, Deserialize, PartialEq, Eq, Debug)]
/// struct Address {
///     street: String,
///     city: String,
/// }
///
/// let address = Address {
///     street: "1313 Webfoot Walk".to_string(),
///     city: "Duckburg, Calisota".to_string(),
/// };
///
/// assert_eq!(
///     serde_bencode::to_string(&address)?,
///     "d4:city18:Duckburg, Calisota6:street17:1313 Webfoot Walke".to_string(),
/// );
/// # Ok(())
/// # }
/// ```
///
/// # Errors
///
/// Serialization can fail if `T`'s implementation of `Serialize` decides to fail or `T` contains
/// floating point values, which bencode cannot serialize.
pub fn to_string<T: ser::Serialize>(b: &T) -> Result<String> {
    let mut ser = Serializer::new();
    b.serialize(&mut ser)?;
    str::from_utf8(ser.as_ref())
        .map(|s| s.to_string())
        .map_err(|_| Error::InvalidValue("Not an UTF-8".to_string()))
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
 * An implementation of the RC4 (also sometimes called ARC4) stream cipher. THIS IMPLEMENTATION IS
 * NOT A FIXED TIME IMPLEMENTATION.
 */

use buffer::{BufferResult, RefReadBuffer, RefWriteBuffer};
use symmetriccipher::{Encryptor, Decryptor, SynchronousStreamCipher, SymmetricCipherError};
use cryptoutil::symm_enc_or_dec;

#[derive(Copy)]
pub struct Rc4 {
    i: uint,
    j: uint,
    state: [u8; 256]
}

impl Rc4 {
    pub fn new(key: &[u8]) -> Rc4 {
        assert!(key.len() >= 1 && key.len() <= 256);
        let mut rc4 = Rc4 { i: 0, j: 0, state: [0; 256] };
        for (i, x) in rc4.state.iter_mut().enumerate() {
            *x = i as u8;
        }
        let mut j: u8 = 0;
        for i in range(0u, 256) {
            j = j + rc4.state[i] + key[i % key.len()];
            rc4.state.swap(i, j as uint);
        }
        rc4
    }
    fn next(&mut self) -> u8 {
        self.i = (self.i + 1) % 256;
        self.j = (self.j + self.state[self.i] as uint) % 256;
        self.state.swap(self.i, self.j);
        let k = self.state[(self.state[self.i] + self.state[self.j]) as uint];
        k
    }
}

impl SynchronousStreamCipher for Rc4 {
    fn process(&mut self, input: &[u8], output: &mut [u8]) {
        assert!(input.len() == output.len());
        for (x, y) in input.iter().zip(output.iter_mut()) {
            *y = *x ^ self.next();
        }
    }
}

impl Encryptor for Rc4 {
    fn encrypt(&mut self, input: &mut RefReadBuffer, output: &mut RefWriteBuffer, _: bool)
            -> Result<BufferResult, SymmetricCipherError> {
        symm_enc_or_dec(self, input, output)
    }
}

impl Decryptor for Rc4 {
    fn decrypt(&mut self, input: &mut RefReadBuffer, output: &mut RefWriteBuffer, _: bool)
            -> Result<BufferResult, SymmetricCipherError> {
        symm_enc_or_dec(self, input, output)
    }
}

#[cfg(test)]
mod test {
    use std::iter::repeat;

    use symmetriccipher::SynchronousStreamCipher;
    use rc4::Rc4;

    struct Test {
        key: &'static str,
        input: &'static str,
        output: Vec<u8>
    }

    fn tests() -> Vec<Test> {
        vec![
            Test {
                key: "Key",
                input: "Plaintext",
                output: vec![0xBB, 0xF3, 0x16, 0xE8, 0xD9, 0x40, 0xAF, 0x0A, 0xD3]
            },
            Test {
                key: "Wiki",
                input: "pedia",
                output: vec![0x10, 0x21, 0xBF, 0x04, 0x20]
            },
            Test {
                key: "Secret",
                input: "Attack at dawn",
                output: vec![0x45, 0xA0, 0x1F, 0x64, 0x5F, 0xC3, 0x5B,
                          0x38, 0x35, 0x52, 0x54, 0x4B, 0x9B, 0xF5]
            }
        ]
    }

    #[test]
    fn wikipedia_tests() {
        let tests = tests();
        for t in tests.iter() {
            let mut rc4 = Rc4::new(t.key.as_bytes());
            let mut result: Vec<u8> = repeat(0).take(t.output.len()).collect();
            rc4.process(t.input.as_bytes(), result.as_mut_slice());
            assert!(result == t.output);
        }
    }
}

#[cfg(test)]
mod bench {
    use test::Bencher;
    use symmetriccipher::SynchronousStreamCipher;
    use rc4::Rc4;

    #[bench]
    pub fn rc4_10(bh: & mut Bencher) {
        let mut rc4 = Rc4::new("key".as_bytes());
        let input = [1u8; 10];
        let mut output = [0u8; 10];
        bh.iter( || {
            rc4.process(&input, &mut output);
        });
        bh.bytes = input.len() as u64;
    }

    #[bench]
    pub fn rc4_1k(bh: & mut Bencher) {
        let mut rc4 = Rc4::new("key".as_bytes());
        let input = [1u8; 1024];
        let mut output = [0u8; 1024];
        bh.iter( || {
            rc4.process(&input, &mut output);
        });
        bh.bytes = input.len() as u64;
    }

    #[bench]
    pub fn rc4_64k(bh: & mut Bencher) {
        let mut rc4 = Rc4::new("key".as_bytes());
        let input = [1u8; 65536];
        let mut output = [0u8; 65536];
        bh.iter( || {
            rc4.process(&input, &mut output);
        });
        bh.bytes = input.len() as u64;
    }
}
<commit_msg>Pick more appropriate int types for rc4<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
 * An implementation of the RC4 (also sometimes called ARC4) stream cipher. THIS IMPLEMENTATION IS
 * NOT A FIXED TIME IMPLEMENTATION.
 */

use buffer::{BufferResult, RefReadBuffer, RefWriteBuffer};
use symmetriccipher::{Encryptor, Decryptor, SynchronousStreamCipher, SymmetricCipherError};
use cryptoutil::symm_enc_or_dec;

#[derive(Copy)]
pub struct Rc4 {
    i: u8,
    j: u8,
    state: [u8; 256]
}

impl Rc4 {
    pub fn new(key: &[u8]) -> Rc4 {
        assert!(key.len() >= 1 && key.len() <= 256);
        let mut rc4 = Rc4 { i: 0, j: 0, state: [0; 256] };
        for (i, x) in rc4.state.iter_mut().enumerate() {
            *x = i as u8;
        }
        let mut j: u8 = 0;
        for i in range(0, 256) {
            j = j + rc4.state[i] + key[i % key.len()];
            rc4.state.swap(i, j as usize);
        }
        rc4
    }
    fn next(&mut self) -> u8 {
        self.i = self.i + 1; // we want this to wrap after 255
        self.j = self.j + self.state[self.i as usize]; // we want this to wrap after 255
        self.state.swap(self.i as usize, self.j as usize);
        let k = self.state[(self.state[self.i as usize] + self.state[self.j as usize]) as usize];
        k
    }
}

impl SynchronousStreamCipher for Rc4 {
    fn process(&mut self, input: &[u8], output: &mut [u8]) {
        assert!(input.len() == output.len());
        for (x, y) in input.iter().zip(output.iter_mut()) {
            *y = *x ^ self.next();
        }
    }
}

impl Encryptor for Rc4 {
    fn encrypt(&mut self, input: &mut RefReadBuffer, output: &mut RefWriteBuffer, _: bool)
            -> Result<BufferResult, SymmetricCipherError> {
        symm_enc_or_dec(self, input, output)
    }
}

impl Decryptor for Rc4 {
    fn decrypt(&mut self, input: &mut RefReadBuffer, output: &mut RefWriteBuffer, _: bool)
            -> Result<BufferResult, SymmetricCipherError> {
        symm_enc_or_dec(self, input, output)
    }
}

#[cfg(test)]
mod test {
    use std::iter::repeat;

    use symmetriccipher::SynchronousStreamCipher;
    use rc4::Rc4;

    struct Test {
        key: &'static str,
        input: &'static str,
        output: Vec<u8>
    }

    fn tests() -> Vec<Test> {
        vec![
            Test {
                key: "Key",
                input: "Plaintext",
                output: vec![0xBB, 0xF3, 0x16, 0xE8, 0xD9, 0x40, 0xAF, 0x0A, 0xD3]
            },
            Test {
                key: "Wiki",
                input: "pedia",
                output: vec![0x10, 0x21, 0xBF, 0x04, 0x20]
            },
            Test {
                key: "Secret",
                input: "Attack at dawn",
                output: vec![0x45, 0xA0, 0x1F, 0x64, 0x5F, 0xC3, 0x5B,
                          0x38, 0x35, 0x52, 0x54, 0x4B, 0x9B, 0xF5]
            }
        ]
    }

    #[test]
    fn wikipedia_tests() {
        let tests = tests();
        for t in tests.iter() {
            let mut rc4 = Rc4::new(t.key.as_bytes());
            let mut result: Vec<u8> = repeat(0).take(t.output.len()).collect();
            rc4.process(t.input.as_bytes(), result.as_mut_slice());
            assert!(result == t.output);
        }
    }
}

#[cfg(test)]
mod bench {
    use test::Bencher;
    use symmetriccipher::SynchronousStreamCipher;
    use rc4::Rc4;

    #[bench]
    pub fn rc4_10(bh: & mut Bencher) {
        let mut rc4 = Rc4::new("key".as_bytes());
        let input = [1u8; 10];
        let mut output = [0u8; 10];
        bh.iter( || {
            rc4.process(&input, &mut output);
        });
        bh.bytes = input.len() as u64;
    }

    #[bench]
    pub fn rc4_1k(bh: & mut Bencher) {
        let mut rc4 = Rc4::new("key".as_bytes());
        let input = [1u8; 1024];
        let mut output = [0u8; 1024];
        bh.iter( || {
            rc4.process(&input, &mut output);
        });
        bh.bytes = input.len() as u64;
    }

    #[bench]
    pub fn rc4_64k(bh: & mut Bencher) {
        let mut rc4 = Rc4::new("key".as_bytes());
        let input = [1u8; 65536];
        let mut output = [0u8; 65536];
        bh.iter( || {
            rc4.process(&input, &mut output);
        });
        bh.bytes = input.len() as u64;
    }
}
<|endoftext|>


<commit_before>use std::fmt::Display;

pub trait Concatable {
  fn concat2(&self) -> String;
  fn join2(&self, delim: &str) -> String;
}

impl<V: Display> Concatable for [V] {
  fn concat2(&self) -> String {
    let mut buf = String::new();
    for x in self.iter() {
      buf.push_str(&format!("{}", x));
    }
    buf
  }

  fn join2(&self, delim: &str) -> String {
    let mut buf = String::new();
    let mut first = true;

    for s in self.iter() {
      if first {
        first = false;
      } else {
        buf.push_str(delim);
      }

      buf.push_str(&format!("{}", s));
    }

    buf
  }
}

pub trait ConcatableWithMap<V, F: Fn(&V) -> String> {
  fn concat_with_map(&self, f: F) -> String;
  fn join_with_map(&self, delim: &str, f: F) -> String;
}

impl<V, F: Fn(&V) -> String> ConcatableWithMap<V, F> for [V] {

  fn concat_with_map(&self, f: F) -> String {
     let mut buf = String::new();
     for v in self.iter() {
       buf.push_str(&f(v));
     }
     buf
  }

  fn join_with_map(&self, delim: &str, f: F) -> String {
    let mut buf = String::new();
    let mut first = true;

    for v in self.iter() {
      if first {
        first = false;
      } else {
        buf.push_str(delim);
      }

      buf.push_str(&f(v));
    }

    buf
  }
}

pub fn join<S, T>(strs: T, delimiter: &str) -> String
    where S: Display,
          T: IntoIterator<Item=S> {

  let mut buf = String::new();
  let mut first = true;

  for s in strs.into_iter() {

    if first {
      first = false;
    } else {
      buf.push_str(delimiter);
    }

    buf.push_str(&format!("{}", s));
  }

  buf
}

pub fn join_with_map<'a, S, T, F>(strs: T, delimiter: &str, f: F) -> String
    where T: IntoIterator<Item=S>,
          F: Fn(S) -> String {

  let mut buf = String::new();
  let mut first = true;

  for s in strs.into_iter() {

    if first {
      first = false;
    } else {
      buf.push_str(delimiter);
    }

    buf.push_str(&f(s));
  }

  buf
}

#[cfg(test)]
mod tests {
  use std::fmt;
  use super::{Concatable, ConcatableWithMap};
  use super::{join, join_with_map};

  pub struct Test {
    x: i32
  }

  impl fmt::Display for Test {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
      write!(f, "<{}>", self.x)
    }
  }

#[test]
  fn test_concatable() {
    let ts = vec![
      Test {x: 1}, Test {x: 2}
    ];

    assert_eq!("<1><2>".to_string(), ts.concat2());
  }

  #[test]
  fn test_join() {
    let str_vec = vec!["a", "b", "c"];
    let str_arr = ["a", "b", "c"];

    assert_eq!("a,b,c".to_string(), join(&str_vec, ","));
    assert_eq!("a,b,c".to_string(), join(&str_arr, ","));

    assert_eq!("abc".to_string(), str_vec.concat2());
    assert_eq!("abc".to_string(), str_arr.concat2());

    assert_eq!("abc".to_string(),
      str_vec.concat_with_map(|x| format!("{}", x)));
    assert_eq!("abc".to_string(),
      str_arr.concat_with_map(|x| format!("{}", x)));

    assert_eq!("a,b,c".to_string(),
      join_with_map(&str_vec, ",", |x| x.to_string()));
    assert_eq!("a,b,c".to_string(),
      join_with_map(&str_arr, ",", |x| x.to_string()));
  }
}
<commit_msg>Clean up tests in str module.<commit_after>use std::fmt::Display;

pub trait Concatable {
  fn concat2(&self) -> String;
  fn join2(&self, delim: &str) -> String;
}

impl<V: Display> Concatable for [V] {
  fn concat2(&self) -> String {
    let mut buf = String::new();
    for x in self.iter() {
      buf.push_str(&format!("{}", x));
    }
    buf
  }

  fn join2(&self, delim: &str) -> String {
    let mut buf = String::new();
    let mut first = true;

    for s in self.iter() {
      if first {
        first = false;
      } else {
        buf.push_str(delim);
      }

      buf.push_str(&format!("{}", s));
    }

    buf
  }
}

pub trait ConcatableWithMap<V, F: Fn(&V) -> String> {
  fn concat_with_map(&self, f: F) -> String;
  fn join_with_map(&self, delim: &str, f: F) -> String;
}

impl<V, F: Fn(&V) -> String> ConcatableWithMap<V, F> for [V] {

  fn concat_with_map(&self, f: F) -> String {
     let mut buf = String::new();
     for v in self.iter() {
       buf.push_str(&f(v));
     }
     buf
  }

  fn join_with_map(&self, delim: &str, f: F) -> String {
    let mut buf = String::new();
    let mut first = true;

    for v in self.iter() {
      if first {
        first = false;
      } else {
        buf.push_str(delim);
      }

      buf.push_str(&f(v));
    }

    buf
  }
}

pub fn join<S, T>(strs: T, delimiter: &str) -> String
    where S: Display,
          T: IntoIterator<Item=S> {

  let mut buf = String::new();
  let mut first = true;

  for s in strs.into_iter() {

    if first {
      first = false;
    } else {
      buf.push_str(delimiter);
    }

    buf.push_str(&format!("{}", s));
  }

  buf
}

pub fn join_with_map<'a, S, T, F>(strs: T, delimiter: &str, f: F) -> String
    where T: IntoIterator<Item=S>,
          F: Fn(S) -> String {

  let mut buf = String::new();
  let mut first = true;

  for s in strs.into_iter() {

    if first {
      first = false;
    } else {
      buf.push_str(delimiter);
    }

    buf.push_str(&f(s));
  }

  buf
}

#[cfg(test)]
mod tests {
  use std::fmt;
  use super::{Concatable, ConcatableWithMap};
  use super::{join, join_with_map};

  pub struct Test {
    x: i32
  }

  impl fmt::Display for Test {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
      write!(f, "<{}>", self.x)
    }
  }

#[test]
  fn test_concatable() {
    let ts = vec![
      Test {x: 1}, Test {x: 2}
    ];
    assert_eq!("<1><2>".to_string(), ts.concat2());
    assert_eq!("<1>=<2>".to_string(), ts.join2("="));

    let str_vec = vec!["a", "b", "c"];
    let str_arr = ["a", "b", "c"];
    assert_eq!("abc".to_string(), str_vec.concat2());
    assert_eq!("abc".to_string(), str_arr.concat2());
  }

  #[test]
  fn test_concatable_with_map() {
    let str_vec = vec!["a", "b", "c"];
    let str_arr = ["a", "b", "c"];

    assert_eq!("abc".to_string(),
      str_vec.concat_with_map(|x| format!("{}", x)));
    assert_eq!("abc".to_string(),
      str_arr.concat_with_map(|x| format!("{}", x)));

    assert_eq!("a,b,c".to_string(),
      str_vec.join_with_map(",", |x| format!("{}", x)));
    assert_eq!("a,b,c".to_string(),
      str_arr.join_with_map(",", |x| format!("{}", x)));
  }

  #[test]
  fn test_join() {
    let str_vec = vec!["a", "b", "c"];
    let str_arr = ["a", "b", "c"];

    assert_eq!("a,b,c".to_string(), concat(&str_vec, ","));
    assert_eq!("a,b,c".to_string(), concat(&str_arr, ","));

    assert_eq!("a,b,c".to_string(), join(&str_vec, ","));
    assert_eq!("a,b,c".to_string(), join(&str_arr, ","));
  }
}
<|endoftext|>