<commit_before>//! Logic for all supported compression codecs in Avro.
use crate::{
    errors::{AvroResult, Error},
    types::Value,
};
use libflate::deflate::{Decoder, Encoder};
use std::io::{Read, Write};
use strum_macros::{EnumString, IntoStaticStr};

/// The compression codec used to compress blocks.
#[derive(Clone, Copy, Debug, PartialEq, EnumString, IntoStaticStr)]
#[strum(serialize_all = "kebab_case")]
pub enum Codec {
    /// The `Null` codec simply passes through data uncompressed.
    Null,
    /// The `Deflate` codec writes the data block using the deflate algorithm
    /// as specified in RFC 1951, and typically implemented using the zlib library.
    /// Note that this format (unlike the "zlib format" in RFC 1950) does not have a checksum.
    Deflate,
    #[cfg(feature = "snappy")]
    /// The `Snappy` codec uses Google's [Snappy](http://google.github.io/snappy/)
    /// compression library. Each compressed block is followed by the 4-byte, big-endian
    /// CRC32 checksum of the uncompressed data in the block.
    Snappy,
}

impl From<Codec> for Value {
    fn from(value: Codec) -> Self {
        Self::Bytes(<&str>::from(value).as_bytes().to_vec())
    }
}

impl Codec {
    /// Compress a stream of bytes in-place.
    pub fn compress(self, stream: &mut Vec<u8>) -> AvroResult<()> {
        match self {
            Codec::Null => (),
            Codec::Deflate => {
                let mut encoder = Encoder::new(Vec::new());
                encoder.write_all(stream)?;
                // Deflate errors seem to just be io::Error
                *stream = encoder.finish().into_result()?;
            }
            #[cfg(feature = "snappy")]
            Codec::Snappy => {
                use byteorder::ByteOrder;

                let mut encoded: Vec<u8> = vec![0; snap::max_compress_len(stream.len())];
                let compressed_size =
                    snap::Encoder::new().compress(&stream[..], &mut encoded[..])?;

                let crc = crc::crc32::checksum_ieee(&stream[..]);
                byteorder::BigEndian::write_u32(&mut encoded[compressed_size..], crc);
                encoded.truncate(compressed_size + 4);

                *stream = encoded;
            }
        };

        Ok(())
    }

    /// Decompress a stream of bytes in-place.
    pub fn decompress(self, stream: &mut Vec<u8>) -> AvroResult<()> {
        *stream = match self {
            Codec::Null => return Ok(()),
            Codec::Deflate => {
                let mut decoded = Vec::new();
                let mut decoder = Decoder::new(&stream[..]);
                decoder.read_to_end(&mut decoded)?;
                decoded
            }
            #[cfg(feature = "snappy")]
            Codec::Snappy => {
                use byteorder::ByteOrder;

                let decompressed_size = snap::decompress_len(&stream[..stream.len() - 4])?;
                let mut decoded = vec![0; decompressed_size];
                snap::Decoder::new().decompress(&stream[..stream.len() - 4], &mut decoded[..])?;

                let expected_crc = byteorder::BigEndian::read_u32(&stream[stream.len() - 4..]);
                let actual_crc = crc::crc32::checksum_ieee(&decoded);

                if expected_crc != actual_crc {
                    return Err(Error::SnappyCrcError {
                        expected: expected_crc,
                        found: actual_crc,
                    });
                }
                decoded
            }
        };
        Ok(())
    }
}

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

    const INPUT: &[u8] = b"theanswertolifetheuniverseandeverythingis42theanswertolifetheuniverseandeverythingis4theanswertolifetheuniverseandeverythingis2";

    #[test]
    fn null_compress_and_decompress() {
        let codec = Codec::Null;
        let mut stream = INPUT.to_vec();
        codec.compress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
        codec.decompress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
    }

    #[test]
    fn deflate_compress_and_decompress() {
        let codec = Codec::Deflate;
        let mut stream = INPUT.to_vec();
        codec.compress(&mut stream).unwrap();
        assert_ne!(INPUT, stream.as_slice());
        assert!(INPUT.len() > stream.len());
        codec.decompress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
    }

    #[cfg(feature = "snappy")]
    #[test]
    fn snappy_compress_and_decompress() {
        let codec = Codec::Snappy;
        let mut stream = INPUT.to_vec();
        codec.compress(&mut stream).unwrap();
        assert_ne!(INPUT, stream.as_slice());
        assert!(INPUT.len() > stream.len());
        codec.decompress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
    }

    #[test]
    fn codec_to_str() {
        assert_eq!(<&str>::from(Codec::Null), "null");
        assert_eq!(<&str>::from(Codec::Deflate), "deflate");

        #[cfg(feature = "snappy")]
        assert_eq!(<&str>::from(Codec::Snappy), "snappy");
    }

    #[test]
    fn codec_from_str() {
        use std::str::FromStr;

        assert_eq!(Codec::from_str("null").unwrap(), Codec::Null);
        assert_eq!(Codec::from_str("deflate").unwrap(), Codec::Deflate);

        #[cfg(feature = "snappy")]
        assert_eq!(Codec::from_str("snappy").unwrap(), Codec::Snappy);

        assert!(Codec::from_str("not a codec").is_err());
    }
}
<commit_msg>Move codec to use site (#142)<commit_after>//! Logic for all supported compression codecs in Avro.
use crate::{errors::AvroResult, types::Value};
use libflate::deflate::{Decoder, Encoder};
use std::io::{Read, Write};
use strum_macros::{EnumString, IntoStaticStr};

/// The compression codec used to compress blocks.
#[derive(Clone, Copy, Debug, PartialEq, EnumString, IntoStaticStr)]
#[strum(serialize_all = "kebab_case")]
pub enum Codec {
    /// The `Null` codec simply passes through data uncompressed.
    Null,
    /// The `Deflate` codec writes the data block using the deflate algorithm
    /// as specified in RFC 1951, and typically implemented using the zlib library.
    /// Note that this format (unlike the "zlib format" in RFC 1950) does not have a checksum.
    Deflate,
    #[cfg(feature = "snappy")]
    /// The `Snappy` codec uses Google's [Snappy](http://google.github.io/snappy/)
    /// compression library. Each compressed block is followed by the 4-byte, big-endian
    /// CRC32 checksum of the uncompressed data in the block.
    Snappy,
}

impl From<Codec> for Value {
    fn from(value: Codec) -> Self {
        Self::Bytes(<&str>::from(value).as_bytes().to_vec())
    }
}

impl Codec {
    /// Compress a stream of bytes in-place.
    pub fn compress(self, stream: &mut Vec<u8>) -> AvroResult<()> {
        match self {
            Codec::Null => (),
            Codec::Deflate => {
                let mut encoder = Encoder::new(Vec::new());
                encoder.write_all(stream)?;
                // Deflate errors seem to just be io::Error
                *stream = encoder.finish().into_result()?;
            }
            #[cfg(feature = "snappy")]
            Codec::Snappy => {
                use byteorder::ByteOrder;

                let mut encoded: Vec<u8> = vec![0; snap::max_compress_len(stream.len())];
                let compressed_size =
                    snap::Encoder::new().compress(&stream[..], &mut encoded[..])?;

                let crc = crc::crc32::checksum_ieee(&stream[..]);
                byteorder::BigEndian::write_u32(&mut encoded[compressed_size..], crc);
                encoded.truncate(compressed_size + 4);

                *stream = encoded;
            }
        };

        Ok(())
    }

    /// Decompress a stream of bytes in-place.
    pub fn decompress(self, stream: &mut Vec<u8>) -> AvroResult<()> {
        *stream = match self {
            Codec::Null => return Ok(()),
            Codec::Deflate => {
                let mut decoded = Vec::new();
                let mut decoder = Decoder::new(&stream[..]);
                decoder.read_to_end(&mut decoded)?;
                decoded
            }
            #[cfg(feature = "snappy")]
            Codec::Snappy => {
                use crate::Error;
                use byteorder::ByteOrder;

                let decompressed_size = snap::decompress_len(&stream[..stream.len() - 4])?;
                let mut decoded = vec![0; decompressed_size];
                snap::Decoder::new().decompress(&stream[..stream.len() - 4], &mut decoded[..])?;

                let expected_crc = byteorder::BigEndian::read_u32(&stream[stream.len() - 4..]);
                let actual_crc = crc::crc32::checksum_ieee(&decoded);

                if expected_crc != actual_crc {
                    return Err(Error::SnappyCrcError {
                        expected: expected_crc,
                        found: actual_crc,
                    });
                }
                decoded
            }
        };
        Ok(())
    }
}

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

    const INPUT: &[u8] = b"theanswertolifetheuniverseandeverythingis42theanswertolifetheuniverseandeverythingis4theanswertolifetheuniverseandeverythingis2";

    #[test]
    fn null_compress_and_decompress() {
        let codec = Codec::Null;
        let mut stream = INPUT.to_vec();
        codec.compress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
        codec.decompress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
    }

    #[test]
    fn deflate_compress_and_decompress() {
        let codec = Codec::Deflate;
        let mut stream = INPUT.to_vec();
        codec.compress(&mut stream).unwrap();
        assert_ne!(INPUT, stream.as_slice());
        assert!(INPUT.len() > stream.len());
        codec.decompress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
    }

    #[cfg(feature = "snappy")]
    #[test]
    fn snappy_compress_and_decompress() {
        let codec = Codec::Snappy;
        let mut stream = INPUT.to_vec();
        codec.compress(&mut stream).unwrap();
        assert_ne!(INPUT, stream.as_slice());
        assert!(INPUT.len() > stream.len());
        codec.decompress(&mut stream).unwrap();
        assert_eq!(INPUT, stream.as_slice());
    }

    #[test]
    fn codec_to_str() {
        assert_eq!(<&str>::from(Codec::Null), "null");
        assert_eq!(<&str>::from(Codec::Deflate), "deflate");

        #[cfg(feature = "snappy")]
        assert_eq!(<&str>::from(Codec::Snappy), "snappy");
    }

    #[test]
    fn codec_from_str() {
        use std::str::FromStr;

        assert_eq!(Codec::from_str("null").unwrap(), Codec::Null);
        assert_eq!(Codec::from_str("deflate").unwrap(), Codec::Deflate);

        #[cfg(feature = "snappy")]
        assert_eq!(Codec::from_str("snappy").unwrap(), Codec::Snappy);

        assert!(Codec::from_str("not a codec").is_err());
    }
}
<|endoftext|>


<commit_before>use std::io::{self, ErrorKind, Cursor};
use std::error::Error;
use bytes::BytesMut;
use tokio_io::codec::{Encoder, Decoder};

use mqtt3::{self, Packet, MqttWrite, MqttRead};
use slog::{Logger, Drain};
use slog_term;

pub struct MqttCodec {
    logger: Logger,
}

impl MqttCodec {
    pub fn new() -> Self {
        MqttCodec {
            logger: codec_logger(),
        }
    }
}

impl Decoder for MqttCodec {
    type Item = Packet;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<Packet>> {
        // NOTE: `decode` might be called with `buf.len == 0` when prevous
        // decode call read all the bytes in the stream. We should return
        // Ok(None) in those cases or else the `read` call will return
        // Ok(0) => translated to UnexpectedEOF by `byteorder` crate.
        // `read` call Ok(0) happens when buffer specified was 0 bytes in len
        // https://doc.rust-lang.org/std/io/trait.Read.html#tymethod.read
        if buf.len() == 0 {
            return Ok(None);
        }

        let (packet, len) = {
            let mut buf_ref = buf.as_ref();
            match buf_ref.read_packet_with_len() {
                Err(e) => {
                    if let mqtt3::Error::Io(e) = e {
                        match e.kind() {
                            ErrorKind::TimedOut | ErrorKind::WouldBlock => return Ok(None),
                            _ => {
                                error!(self.logger, "mqtt3 io error = {:?}", e);
                                return Err(io::Error::new(e.kind(), e.description()))
                            },
                        }
                    } else {
                        error!(self.logger, "mqtt3 read error = {:?}", e);
                        return Err(io::Error::new(ErrorKind::Other, e.description()));
                    }
                }
                Ok(v) => v,
            }
        };

        // NOTE: It's possible that `decode` got called before `buf` has full bytes
        // necessary to frame raw bytes into a packet. In that case return Ok(None)
        // and the next time decode` gets called, there will be more bytes in `buf`,
        // hopefully enough to frame the packet
        if buf.len() < len {
            return Ok(None);
        }

        // println!("{:?}, {:?}, {:?}", len, packet, buf.len());

        buf.split_to(len);

        Ok(Some(packet))
    }
}

impl Encoder for MqttCodec {
    type Item = Packet;
    type Error = io::Error;

    fn encode(&mut self, msg: Packet, buf: &mut BytesMut) -> io::Result<()> {
        let mut stream = Cursor::new(Vec::new());

        // TODO: Implement `write_packet` for `&mut BytesMut`
        if let Err(_) = stream.write_packet(&msg) {
            return Err(io::Error::new(io::ErrorKind::Other, "Unable to encode!"));
        }

        buf.extend(stream.get_ref());

        Ok(())
    }
}

fn codec_logger() -> Logger {
    use std::sync::Mutex;

    let decorator = slog_term::TermDecorator::new().build();
    let drain = slog_term::FullFormat::new(decorator).build().fuse();
    let drain = Mutex::new(drain).fuse();
    Logger::root(drain, o!("codec" => env!("CARGO_PKG_VERSION")))
}
<commit_msg>fix unexpected eof bug caused due to insufficient input buf size<commit_after>use std::io::{self, ErrorKind, Cursor};
use std::error::Error;
use bytes::BytesMut;
use tokio_io::codec::{Encoder, Decoder};

use mqtt3::{self, Packet, MqttWrite, MqttRead};
use slog::{Logger, Drain};
use slog_term;

pub struct MqttCodec {
    logger: Logger,
}

impl MqttCodec {
    pub fn new() -> Self {
        MqttCodec {
            logger: codec_logger(),
        }
    }
}

impl Decoder for MqttCodec {
    type Item = Packet;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<Packet>> {
        // NOTE: `decode` might be called with `buf.len == 0` when prevous
        // decode call read all the bytes in the stream. We should return
        // Ok(None) in those cases or else the `read` call will return
        // Ok(0) => translated to UnexpectedEOF by `byteorder` crate.
        // `read` call Ok(0) happens when buffer specified was 0 bytes in len
        // https://doc.rust-lang.org/std/io/trait.Read.html#tymethod.read
        if buf.len() < 2 {
            return Ok(None);
        }

        let (packet, len) = {
            let mut buf_ref = buf.as_ref();
            match buf_ref.read_packet_with_len() {
                Err(e) => {
                    if let mqtt3::Error::Io(e) = e {
                        match e.kind() {
                            ErrorKind::TimedOut | ErrorKind::WouldBlock => return Ok(None),
                            ErrorKind::UnexpectedEof => return Ok(None),
                            _ => {
                                error!(self.logger, "mqtt3 io error = {:?}", e);
                                return Err(io::Error::new(e.kind(), e.description()))
                            },
                        }
                    } else {
                        error!(self.logger, "mqtt3 read error = {:?}", e);
                        return Err(io::Error::new(ErrorKind::Other, e.description()));
                    }
                }
                Ok(v) => v,
            }
        };

        // NOTE: It's possible that `decode` got called before `buf` has full bytes
        // necessary to frame raw bytes into a packet. In that case return Ok(None)
        // and the next time decode` gets called, there will be more bytes in `buf`,
        // hopefully enough to frame the packet
        if buf.len() < len {
            return Ok(None);
        }

        // println!("{:?}, {:?}, {:?}", len, packet, buf.len());

        buf.split_to(len);

        Ok(Some(packet))
    }
}

impl Encoder for MqttCodec {
    type Item = Packet;
    type Error = io::Error;

    fn encode(&mut self, msg: Packet, buf: &mut BytesMut) -> io::Result<()> {
        let mut stream = Cursor::new(Vec::new());

        // TODO: Implement `write_packet` for `&mut BytesMut`
        if let Err(_) = stream.write_packet(&msg) {
            return Err(io::Error::new(io::ErrorKind::Other, "Unable to encode!"));
        }

        buf.extend(stream.get_ref());

        Ok(())
    }
}

fn codec_logger() -> Logger {
    use std::sync::Mutex;

    let decorator = slog_term::TermDecorator::new().build();
    let drain = slog_term::FullFormat::new(decorator).build().fuse();
    let drain = Mutex::new(drain).fuse();
    Logger::root(drain, o!("codec" => env!("CARGO_PKG_VERSION")))
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::fmt::Arguments;
use std::io::Read;
use std::io::Write;
use std::io;
use std::net::TcpStream;
use std::os::unix::io::{AsRawFd, RawFd};
use std::str;
use std::time::Duration;

use msg::{Msg, Pfx, Command};
use utils::find_byte;

pub struct Comms {
    /// The TCP connection to the server.
    stream        : TcpStream,

    pub serv_name : String,

    /// _Partial_ messages collected here until they make a complete message.
    msg_buf       : Vec<u8>,
}

#[derive(Debug)]
pub enum CommsRet {
    Disconnected {
        fd        : RawFd,
    },

    Err {
        err_msg   : String,
    },

    IncomingMsg {
        pfx       : Pfx,
        ty        : String,
        args      : Vec<String>,
    },

    /// A message without prefix. From RFC 2812:
    /// > If the prefix is missing from the message, it is assumed to have
    /// > originated from the connection from which it was received from.
    SentMsg {
        ty        : String,
        msg       : String,
    }
}

impl Comms {
    pub fn try_connect(serv_addr : &str, serv_name : &str,
                       nick : &str, hostname : &str, realname : &str)
                       -> io::Result<Comms> {
        let stream = try!(TcpStream::connect(serv_addr));
        try!(stream.set_read_timeout(Some(Duration::from_millis(10))));
        try!(stream.set_write_timeout(None));
        // try!(stream.set_nonblocking(true));

        let mut comms = Comms {
            stream:     stream,
            serv_name:  serv_name.to_owned(),
            msg_buf:    Vec::new(),
        };
        try!(comms.introduce(nick, hostname, realname));
        Ok(comms)
    }

    /// Get the RawFd, to be used with select() or other I/O multiplexer.
    pub fn get_raw_fd(&self) -> RawFd {
        self.stream.as_raw_fd()
    }

    ////////////////////////////////////////////////////////////////////////////
    // Sending messages

    fn introduce(&mut self, nick : &str, hostname : &str, realname : &str) -> io::Result<()> {
        try!(Msg::user(hostname, realname, &mut self.stream));
        Msg::nick(nick, &mut self.stream)
    }

    ////////////////////////////////////////////////////////////////////////////
    // Receiving messages

    pub fn read_incoming_msg(&mut self) -> Vec<CommsRet> {
        let mut read_buf : [u8; 512] = [0; 512];

        // Handle disconnects
        match self.stream.read(&mut read_buf) {
            Err(_) => {
                // TODO: I don't understand why this happens. I'm ``randomly''
                // getting "temporarily unavailable" errors.
                // return vec![CommsRet::ShowErr(format!("error in read(): {:?}", err))];
                return vec![];
            },
            Ok(bytes_read) => {
                if bytes_read == 0 {
                    vec![CommsRet::Disconnected { fd: self.get_raw_fd() }]
                } else {
                    self.add_to_msg_buf(&read_buf[ 0 .. bytes_read ]);
                    self.handle_msgs()
                }
            }
        }
    }

    #[inline]
    fn add_to_msg_buf(&mut self, slice : &[u8]) {
        // Some invisible ASCII characters causing glitches on some terminals,
        // we filter those out here.
        self.msg_buf.extend(slice.iter().filter(|c| **c != 0x1 /* SOH */ ||
                                                    **c != 0x2 /* STX */ ||
                                                    **c != 0x0 /* NUL */ ||
                                                    **c != 0x4 /* EOT */));
    }

    fn handle_msgs(&mut self) -> Vec<CommsRet> {
        let mut ret = Vec::with_capacity(1);

        // Have we read any CRLFs? In that case just process the message and
        // update buffers. Otherwise just leave the partial message in the
        // buffer.
        loop {
            match find_byte(self.msg_buf.borrow(), b'\r') {
                None => { break; },
                Some(cr_idx) => {
                    // We have a CR, however, we don't have any guarantees that
                    // a single read() will read both CR and LF. So if we have a
                    // CR, but that's the last byte, we should just wait until
                    // we read NL too.
                    if cr_idx == self.msg_buf.len() - 1 {
                        break;
                    } else {
                        assert!(self.msg_buf[cr_idx + 1] == b'\n');
                        // Don't include CRLF
                        let msg = Msg::parse(&self.msg_buf[ 0 .. cr_idx ]);
                        Comms::handle_msg(&mut self.stream, msg, &mut ret);
                        // Update the buffer (drop CRLF)
                        self.msg_buf.drain(0 .. cr_idx + 2);
                    }
                }
            }
        }

        ret
    }

    fn handle_msg(stream : &mut TcpStream, msg : Result<Msg, String>,
                  ret : &mut Vec<CommsRet>) {
        match msg {
            Err(err_msg) => {
                ret.push(CommsRet::Err { err_msg: err_msg });
            },
            Ok(Msg { pfx, command, params }) => {
                match command {
                    Command::Str(str) =>
                        Comms::handle_str_command(stream, ret, pfx, str, params),
                    Command::Num(num) =>
                        Comms::handle_num_command(stream, ret, pfx, num, params),
                }
            }
        }
    }

    fn handle_str_command(stream : &mut TcpStream, ret : &mut Vec<CommsRet>,
                          pfx : Option<Pfx>, cmd : String, params : Vec<Vec<u8>>) {
        if cmd == "PING" {
            debug_assert!(params.len() == 1);
            Msg::pong(unsafe {
                        str::from_utf8_unchecked(params.into_iter().nth(0).unwrap().as_ref())
                      }, stream).unwrap();
        } else {
            match pfx {
                None => {
                    ret.push(CommsRet::SentMsg {
                        ty: cmd,
                        msg: params.into_iter().map(|s| unsafe {
                            String::from_utf8_unchecked(s)
                        }).collect::<Vec<_>>().join(" "), // FIXME: intermediate vector
                    });
                },
                Some(pfx) => {
                    ret.push(CommsRet::IncomingMsg {
                        pfx: pfx,
                        ty: cmd,
                        args: params.into_iter().map(|s| unsafe {
                            String::from_utf8_unchecked(s)
                        }).collect(),
                    });
                }
            }
        }
    }

    fn handle_num_command(stream : &mut TcpStream,
                          ret : &mut Vec<CommsRet>,
                          prefix : Option<Pfx>, num : u16, params : Vec<Vec<u8>>) {
        // TODO
        Comms::handle_str_command(stream, ret, prefix, "UNKNOWN".to_owned(), params)
    }
}

impl Write for Comms {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.stream.write(buf)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.stream.write_all(buf)
    }

    #[inline]
    fn write_fmt(&mut self, fmt: Arguments) -> io::Result<()> {
        self.stream.write_fmt(fmt)
    }

    #[inline]
    fn by_ref(&mut self) -> &mut Comms {
        self
    }
}
<commit_msg>Log all messages in debug mode<commit_after>use std::borrow::Borrow;
use std::fmt::Arguments;
use std::fs::File;
use std::fs;
use std::io::Read;
use std::io::Write;
use std::io;
use std::net::TcpStream;
use std::os::unix::io::{AsRawFd, RawFd};
use std::str;
use std::time::Duration;

use msg::{Msg, Pfx, Command};
use utils::find_byte;

pub struct Comms {
    /// The TCP connection to the server.
    stream        : TcpStream,

    pub serv_name : String,

    /// _Partial_ messages collected here until they make a complete message.
    msg_buf       : Vec<u8>,

    /// A file to log incoming messages for debugging purposes. Only available
    /// when `debug_assertions` is available.
    log_file      : Option<File>,
}

#[derive(Debug)]
pub enum CommsRet {
    Disconnected {
        fd        : RawFd,
    },

    Err {
        err_msg   : String,
    },

    IncomingMsg {
        pfx       : Pfx,
        ty        : String,
        args      : Vec<String>,
    },

    /// A message without prefix. From RFC 2812:
    /// > If the prefix is missing from the message, it is assumed to have
    /// > originated from the connection from which it was received from.
    SentMsg {
        ty        : String,
        msg       : String,
    }
}

impl Comms {
    pub fn try_connect(serv_addr : &str, serv_name : &str,
                       nick : &str, hostname : &str, realname : &str)
                       -> io::Result<Comms> {
        let stream = try!(TcpStream::connect(serv_addr));
        try!(stream.set_read_timeout(Some(Duration::from_millis(10))));
        try!(stream.set_write_timeout(None));
        // try!(stream.set_nonblocking(true));

        let log_file = {
            if cfg!(debug_assertions) {
                fs::create_dir("logs");
                Some(File::create(format!("logs/{}.txt", serv_addr)).unwrap())
            } else {
                None
            }
        };

        let mut comms = Comms {
            stream:     stream,
            serv_name:  serv_name.to_owned(),
            msg_buf:    Vec::new(),
            log_file:   log_file,
        };

        try!(comms.introduce(nick, hostname, realname));
        Ok(comms)
    }

    /// Get the RawFd, to be used with select() or other I/O multiplexer.
    pub fn get_raw_fd(&self) -> RawFd {
        self.stream.as_raw_fd()
    }

    ////////////////////////////////////////////////////////////////////////////
    // Sending messages

    fn introduce(&mut self, nick : &str, hostname : &str, realname : &str) -> io::Result<()> {
        try!(Msg::user(hostname, realname, &mut self.stream));
        Msg::nick(nick, &mut self.stream)
    }

    ////////////////////////////////////////////////////////////////////////////
    // Receiving messages

    pub fn read_incoming_msg(&mut self) -> Vec<CommsRet> {
        let mut read_buf : [u8; 512] = [0; 512];

        // Handle disconnects
        match self.stream.read(&mut read_buf) {
            Err(_) => {
                // TODO: I don't understand why this happens. I'm ``randomly''
                // getting "temporarily unavailable" errors.
                // return vec![CommsRet::ShowErr(format!("error in read(): {:?}", err))];
                return vec![];
            },
            Ok(bytes_read) => {
                if bytes_read == 0 {
                    vec![CommsRet::Disconnected { fd: self.get_raw_fd() }]
                } else {
                    self.add_to_msg_buf(&read_buf[ 0 .. bytes_read ]);
                    self.handle_msgs()
                }
            }
        }
    }

    #[inline]
    fn add_to_msg_buf(&mut self, slice : &[u8]) {
        // Some invisible ASCII characters causing glitches on some terminals,
        // we filter those out here.
        self.msg_buf.extend(slice.iter().filter(|c| **c != 0x1 /* SOH */ ||
                                                    **c != 0x2 /* STX */ ||
                                                    **c != 0x0 /* NUL */ ||
                                                    **c != 0x4 /* EOT */));
    }

    fn handle_msgs(&mut self) -> Vec<CommsRet> {
        let mut ret = Vec::with_capacity(1);

        // Have we read any CRLFs? In that case just process the message and
        // update buffers. Otherwise just leave the partial message in the
        // buffer.
        loop {
            match find_byte(self.msg_buf.borrow(), b'\r') {
                None => { break; },
                Some(cr_idx) => {
                    // We have a CR, however, we don't have any guarantees that
                    // a single read() will read both CR and LF. So if we have a
                    // CR, but that's the last byte, we should just wait until
                    // we read NL too.
                    if cr_idx == self.msg_buf.len() - 1 {
                        break;
                    } else {
                        assert!(self.msg_buf[cr_idx + 1] == b'\n');

                        let msg = {
                            // Don't include CRLF
                            let msg_slice = &self.msg_buf[ 0 .. cr_idx ];
                            // Log the message in debug mode
                            if cfg!(debug_assertions) {
                                writeln!(self.log_file.as_ref().unwrap(), "{}",
                                         unsafe { str::from_utf8_unchecked(msg_slice) });
                            }
                            Msg::parse(msg_slice)
                        };

                        Comms::handle_msg(&mut self.stream, msg, &mut ret);
                        // Update the buffer (drop CRLF too)
                        self.msg_buf.drain(0 .. cr_idx + 2);
                    }
                }
            }
        }

        ret
    }

    fn handle_msg(stream : &mut TcpStream, msg : Result<Msg, String>,
                  ret : &mut Vec<CommsRet>) {
        match msg {
            Err(err_msg) => {
                ret.push(CommsRet::Err { err_msg: err_msg });
            },
            Ok(Msg { pfx, command, params }) => {
                match command {
                    Command::Str(str) =>
                        Comms::handle_str_command(stream, ret, pfx, str, params),
                    Command::Num(num) =>
                        Comms::handle_num_command(stream, ret, pfx, num, params),
                }
            }
        }
    }

    fn handle_str_command(stream : &mut TcpStream, ret : &mut Vec<CommsRet>,
                          pfx : Option<Pfx>, cmd : String, params : Vec<Vec<u8>>) {
        if cmd == "PING" {
            debug_assert!(params.len() == 1);
            Msg::pong(unsafe {
                        str::from_utf8_unchecked(params.into_iter().nth(0).unwrap().as_ref())
                      }, stream).unwrap();
        } else {
            match pfx {
                None => {
                    ret.push(CommsRet::SentMsg {
                        ty: cmd,
                        msg: params.into_iter().map(|s| unsafe {
                            String::from_utf8_unchecked(s)
                        }).collect::<Vec<_>>().join(" "), // FIXME: intermediate vector
                    });
                },
                Some(pfx) => {
                    ret.push(CommsRet::IncomingMsg {
                        pfx: pfx,
                        ty: cmd,
                        args: params.into_iter().map(|s| unsafe {
                            String::from_utf8_unchecked(s)
                        }).collect(),
                    });
                }
            }
        }
    }

    fn handle_num_command(stream : &mut TcpStream,
                          ret : &mut Vec<CommsRet>,
                          prefix : Option<Pfx>, num : u16, params : Vec<Vec<u8>>) {
        // TODO
        Comms::handle_str_command(stream, ret, prefix, "UNKNOWN".to_owned(), params)
    }
}

impl Write for Comms {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.stream.write(buf)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.stream.write_all(buf)
    }

    #[inline]
    fn write_fmt(&mut self, fmt: Arguments) -> io::Result<()> {
        self.stream.write_fmt(fmt)
    }

    #[inline]
    fn by_ref(&mut self) -> &mut Comms {
        self
    }
}
<|endoftext|>


<commit_before>extern crate crypto;
use self::crypto::digest::Digest;
use self::crypto::md5::Md5;

const DATA: &'static str = include_str!("../data/input_4.txt");

pub fn main() -> Vec<String> {
    let s1 = find_number(DATA.trim(), 5);
    let s2 = find_number(DATA.trim(), 6);
    vec![s1.to_string(), s2.to_string()]
}

fn find_number(key: &str, zeros: usize) -> u32 {
    let mut md5 = Md5::new();
    let prefix = format!("{:01$}", 0, zeros);

    (1..)
        .find(|n| {
            md5.reset();
            md5.input_str(&key);
            md5.input_str(&n.to_string());
            md5.result_str().starts_with(&prefix)
        })
        .unwrap()
}

#[test]
#[ignore]
fn test_find_number() {
    assert_eq!(find_number("abcdef", 5), 609043);
    assert_eq!(find_number("pqrstuv", 5), 1048970);
}
<commit_msg>day04: speed up by removing string allocations & utf8 checks<commit_after>extern crate crypto;
use self::crypto::digest::Digest;
use self::crypto::md5::Md5;

use std::io::Write;

const DATA: &'static str = include_str!("../data/input_4.txt");

pub fn main() -> Vec<String> {
    let s1 = find_number(DATA.trim(), 5);
    let s2 = find_number(DATA.trim(), 6);
    vec![s1.to_string(), s2.to_string()]
}

fn leading_zeros(buf: &[u8], count: usize) -> bool {
    let (bytes, nibble) = (count / 2, count % 2);
    buf[..bytes].iter().all(|b| *b == 0) &&
    if nibble != 0 {
        (buf[bytes] & 0xf0) == 0
    } else {
        true
    }
}

fn find_number(key: &str, zeros: usize) -> u32 {
    let mut md5 = Md5::new();
    let k = key.as_bytes();
    let mut buf = Vec::with_capacity(256);
    let mut output = [0u8; 16];

    (1..)
        .find(|n| {
            // it would be nice if I could just roll-back to this state instead of inputting the key each time
            md5.reset();
            md5.input(k);

            // use write! to a pre-allocated buffer
            buf.clear();
            write!(&mut buf, "{}", n).unwrap();
            md5.input(&buf);

            md5.result(&mut output);
            leading_zeros(&output, zeros)
        })
        .unwrap()
}

#[cfg(test)]
mod test {
    extern crate test;
    use self::test::Bencher;
    use super::{find_number, DATA};

    #[test]
    fn test_find_number() {
        assert_eq!(find_number("abcdef", 5), 609043);
        assert_eq!(find_number("pqrstuv", 5), 1048970);
    }

    #[bench]
    fn bench_day4(b: &mut Bencher) {
        b.iter(|| super::find_number(DATA.trim(), 5));
    }
}
<|endoftext|>


<commit_before>//! Error types

pub use hyper::Error as HyperError;
pub use json::Error as JsonError;

use message::Disconnect;
use std::error::Error as StdError;
use std::fmt::{self, Display, Formatter};
use std::io;
use std::str::Utf8Error;
use types::StatusCode;

/// An error occurred while trying to connect to a Stream.
#[derive(Debug)]
pub enum Error {
    /// An HTTP error from the Stream.
    Http(StatusCode),
    /// An error from the `hyper` crate.
    Hyper(HyperError),
    TimedOut,
}

/// An error occured while listening on a Stream.
#[derive(Debug)]
pub enum StreamError {
    /// The Stream has been disconnected by the server.
    Disconnect(Disconnect),
    Hyper(HyperError),
    /// An I/O error.
    Io(io::Error),
    /// Failed to parse a JSON message from a Stream.
    Json(JsonError),
    TimedOut,
    Utf8(Utf8Error),
}

pub enum JsonStreamError {
    Hyper(HyperError),
    TimedOut,
    Utf8(Utf8Error),
}

impl StdError for Error {
    fn description(&self) -> &str {
        use Error::*;

        match *self {
            Http(ref status) => status.canonical_reason().unwrap_or("<unknown status code>"),
            Hyper(ref e) => e.description(),
            TimedOut => "timed out",
        }
    }

    fn cause(&self) -> Option<&StdError> {
        use Error::*;

        match *self {
            Hyper(ref e) => Some(e),
            Http(_) | TimedOut => None,
        }
    }
}

impl StdError for StreamError {
    fn description(&self) -> &str {
        use StreamError::*;

        match *self {
            Disconnect(ref d) => &d.reason,
            Hyper(ref e) => e.description(),
            Io(ref e) => e.description(),
            Json(ref e) => e.description(),
            TimedOut => "timed out",
            Utf8(ref e) => e.description(),
        }
    }

    fn cause(&self) -> Option<&StdError> {
        use StreamError::*;

        match *self {
            Hyper(ref e) => Some(e),
            Io(ref e) => Some(e),
            Json(ref e) => Some(e),
            Utf8(ref e) => Some(e),
            Disconnect(_) | TimedOut => None,
        }
    }
}

impl Display for Error {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        use Error::*;

        match *self {
            Http(ref code) => Display::fmt(code, f),
            Hyper(ref e) => Display::fmt(e, f),
            TimedOut => Display::fmt(self.description(), f),
        }
    }
}

impl Display for StreamError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        use StreamError::*;

        match *self {
            Disconnect(ref d) => Display::fmt(d, f),
            Hyper(ref e) => Display::fmt(e, f),
            Io(ref e) => Display::fmt(e, f),
            Json(ref e) => Display::fmt(e, f),
            TimedOut => Display::fmt(self.description(), f),
            Utf8(ref e) => Display::fmt(e, f),
        }
    }
}

impl From<JsonStreamError> for StreamError {
    fn from(e: JsonStreamError) -> Self {
        match e {
            JsonStreamError::Hyper(e) => StreamError::Hyper(e),
            JsonStreamError::TimedOut => StreamError::TimedOut,
            JsonStreamError::Utf8(e) => StreamError::Utf8(e),
        }
    }
}

impl From<HyperError> for JsonStreamError {
    fn from(e: HyperError) -> Self {
        JsonStreamError::Hyper(e)
    }
}

impl From<Utf8Error> for JsonStreamError {
    fn from(e: Utf8Error) -> Self {
        JsonStreamError::Utf8(e)
    }
}

impl From<JsonError> for StreamError {
    fn from(e: JsonError) -> Self {
        StreamError::Json(e)
    }
}

impl From<StatusCode> for Error {
    fn from(e: StatusCode) -> Self {
        Error::Http(e)
    }
}

impl From<HyperError> for Error {
    fn from(e: HyperError) -> Self {
        Error::Hyper(e)
    }
}

impl From<io::Error> for StreamError {
    fn from(e: io::Error) -> Self {
        StreamError::Io(e)
    }
}
<commit_msg>`derive(Debug)` for `JsonStreamError`<commit_after>//! Error types

pub use hyper::Error as HyperError;
pub use json::Error as JsonError;

use message::Disconnect;
use std::error::Error as StdError;
use std::fmt::{self, Display, Formatter};
use std::io;
use std::str::Utf8Error;
use types::StatusCode;

/// An error occurred while trying to connect to a Stream.
#[derive(Debug)]
pub enum Error {
    /// An HTTP error from the Stream.
    Http(StatusCode),
    /// An error from the `hyper` crate.
    Hyper(HyperError),
    TimedOut,
}

/// An error occured while listening on a Stream.
#[derive(Debug)]
pub enum StreamError {
    /// The Stream has been disconnected by the server.
    Disconnect(Disconnect),
    Hyper(HyperError),
    /// An I/O error.
    Io(io::Error),
    /// Failed to parse a JSON message from a Stream.
    Json(JsonError),
    TimedOut,
    Utf8(Utf8Error),
}

#[derive(Debug)]
pub enum JsonStreamError {
    Hyper(HyperError),
    TimedOut,
    Utf8(Utf8Error),
}

impl StdError for Error {
    fn description(&self) -> &str {
        use Error::*;

        match *self {
            Http(ref status) => status.canonical_reason().unwrap_or("<unknown status code>"),
            Hyper(ref e) => e.description(),
            TimedOut => "timed out",
        }
    }

    fn cause(&self) -> Option<&StdError> {
        use Error::*;

        match *self {
            Hyper(ref e) => Some(e),
            Http(_) | TimedOut => None,
        }
    }
}

impl StdError for StreamError {
    fn description(&self) -> &str {
        use StreamError::*;

        match *self {
            Disconnect(ref d) => &d.reason,
            Hyper(ref e) => e.description(),
            Io(ref e) => e.description(),
            Json(ref e) => e.description(),
            TimedOut => "timed out",
            Utf8(ref e) => e.description(),
        }
    }

    fn cause(&self) -> Option<&StdError> {
        use StreamError::*;

        match *self {
            Hyper(ref e) => Some(e),
            Io(ref e) => Some(e),
            Json(ref e) => Some(e),
            Utf8(ref e) => Some(e),
            Disconnect(_) | TimedOut => None,
        }
    }
}

impl Display for Error {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        use Error::*;

        match *self {
            Http(ref code) => Display::fmt(code, f),
            Hyper(ref e) => Display::fmt(e, f),
            TimedOut => Display::fmt(self.description(), f),
        }
    }
}

impl Display for StreamError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        use StreamError::*;

        match *self {
            Disconnect(ref d) => Display::fmt(d, f),
            Hyper(ref e) => Display::fmt(e, f),
            Io(ref e) => Display::fmt(e, f),
            Json(ref e) => Display::fmt(e, f),
            TimedOut => Display::fmt(self.description(), f),
            Utf8(ref e) => Display::fmt(e, f),
        }
    }
}

impl From<JsonStreamError> for StreamError {
    fn from(e: JsonStreamError) -> Self {
        match e {
            JsonStreamError::Hyper(e) => StreamError::Hyper(e),
            JsonStreamError::TimedOut => StreamError::TimedOut,
            JsonStreamError::Utf8(e) => StreamError::Utf8(e),
        }
    }
}

impl From<HyperError> for JsonStreamError {
    fn from(e: HyperError) -> Self {
        JsonStreamError::Hyper(e)
    }
}

impl From<Utf8Error> for JsonStreamError {
    fn from(e: Utf8Error) -> Self {
        JsonStreamError::Utf8(e)
    }
}

impl From<JsonError> for StreamError {
    fn from(e: JsonError) -> Self {
        StreamError::Json(e)
    }
}

impl From<StatusCode> for Error {
    fn from(e: StatusCode) -> Self {
        Error::Http(e)
    }
}

impl From<HyperError> for Error {
    fn from(e: HyperError) -> Self {
        Error::Hyper(e)
    }
}

impl From<io::Error> for StreamError {
    fn from(e: io::Error) -> Self {
        StreamError::Io(e)
    }
}
<|endoftext|>


<commit_before>//! Mouse and key events.

use std::io::{Error, ErrorKind};
use std::str;

/// An event reported by the terminal.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Event {
    /// A key press.
    Key(Key),
    /// A mouse button press, release or wheel use at specific coordinates.
    Mouse(MouseEvent),
    /// An event that cannot currently be evaluated.
    Unsupported(Vec<u8>),
}

/// A mouse related event.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum MouseEvent {
    /// A mouse button was pressed.
    ///
    /// The coordinates are one-based.
    Press(MouseButton, u16, u16),
    /// A mouse button was released.
    ///
    /// The coordinates are one-based.
    Release(u16, u16),
    /// A mouse button is held over the given coordinates.
    ///
    /// The coordinates are one-based.
    Hold(u16, u16),
}

/// A mouse button.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum MouseButton {
    /// The left mouse button.
    Left,
    /// The right mouse button.
    Right,
    /// The middle mouse button.
    Middle,
    /// Mouse wheel is going up.
    ///
    /// This event is typically only used with Mouse::Press.
    WheelUp,
    /// Mouse wheel is going down.
    ///
    /// This event is typically only used with Mouse::Press.
    WheelDown,
}

/// A key.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum Key {
    /// Backspace.
    Backspace,
    /// Left arrow.
    Left,
    /// Right arrow.
    Right,
    /// Up arrow.
    Up,
    /// Down arrow.
    Down,
    /// Home key.
    Home,
    /// End key.
    End,
    /// Page Up key.
    PageUp,
    /// Page Down key.
    PageDown,
    /// Delete key.
    Delete,
    /// Insert key.
    Insert,
    /// Function keys.
    ///
    /// Only function keys 1 through 12 are supported.
    F(u8),
    /// Normal character.
    Char(char),
    /// Alt modified character.
    Alt(char),
    /// Ctrl modified character.
    ///
    /// Note that certain keys may not be modifiable with `ctrl`, due to limitations of terminals.
    Ctrl(char),
    /// Null byte.
    Null,
    /// Esc key.
    Esc,

    #[doc(hidden)]
    __IsNotComplete,
}

/// Parse an Event from `item` and possibly subsequent bytes through `iter`.
pub fn parse_event<I>(item: u8, iter: &mut I) -> Result<Event, Error>
    where I: Iterator<Item = Result<u8, Error>>
{
    let error = Error::new(ErrorKind::Other, "Could not parse an event");
    match item {
        b'\x1B' => {
            // This is an escape character, leading a control sequence.
            Ok(match iter.next() {
                   Some(Ok(b'O')) => {
                match iter.next() {
                    // F1-F4
                    Some(Ok(val @ b'P'...b'S')) => Event::Key(Key::F(1 + val - b'P')),
                    _ => return Err(error),
                }
            }
                   Some(Ok(b'[')) => {
                // This is a CSI sequence.
                parse_csi(iter).ok_or(error)?
            }
                   Some(Ok(c)) => {
                let ch = parse_utf8_char(c, iter);
                Event::Key(Key::Alt(try!(ch)))
            }
                   Some(Err(_)) | None => return Err(error),
               })
        }
        b'\n' | b'\r' => Ok(Event::Key(Key::Char('\n'))),
        b'\t' => Ok(Event::Key(Key::Char('\t'))),
        b'\x7F' => Ok(Event::Key(Key::Backspace)),
        c @ b'\x01'...b'\x1A' => Ok(Event::Key(Key::Ctrl((c as u8 - 0x1 + b'a') as char))),
        c @ b'\x1C'...b'\x1F' => Ok(Event::Key(Key::Ctrl((c as u8 - 0x1C + b'4') as char))),
        b'\0' => Ok(Event::Key(Key::Null)),
        c => {
            Ok({
                   let ch = parse_utf8_char(c, iter);
                   Event::Key(Key::Char(try!(ch)))
               })
        }
    }
}

/// Parses a CSI sequence, just after reading ^[
///
/// Returns None if an unrecognized sequence is found.
fn parse_csi<I>(iter: &mut I) -> Option<Event>
    where I: Iterator<Item = Result<u8, Error>>
{
    Some(match iter.next() {
             Some(Ok(b'[')) => match iter.next() {
                 Some(Ok(val @ b'A'...b'E')) => Event::Key(Key::F(1 + val - b'A')),
                 _ => return None,
             },
             Some(Ok(b'D')) => Event::Key(Key::Left),
             Some(Ok(b'C')) => Event::Key(Key::Right),
             Some(Ok(b'A')) => Event::Key(Key::Up),
             Some(Ok(b'B')) => Event::Key(Key::Down),
             Some(Ok(b'H')) => Event::Key(Key::Home),
             Some(Ok(b'F')) => Event::Key(Key::End),
             Some(Ok(b'M')) => {
        // X10 emulation mouse encoding: ESC [ CB Cx Cy (6 characters only).
        let mut next = || iter.next().unwrap().unwrap();

        let cb = next() as i8 - 32;
        // (1, 1) are the coords for upper left.
        let cx = next().saturating_sub(32) as u16;
        let cy = next().saturating_sub(32) as u16;
        Event::Mouse(match cb & 0b11 {
                         0 => {
                             if cb & 0x40 != 0 {
                                 MouseEvent::Press(MouseButton::WheelUp, cx, cy)
                             } else {
                                 MouseEvent::Press(MouseButton::Left, cx, cy)
                             }
                         }
                         1 => {
                             if cb & 0x40 != 0 {
                                 MouseEvent::Press(MouseButton::WheelDown, cx, cy)
                             } else {
                                 MouseEvent::Press(MouseButton::Middle, cx, cy)
                             }
                         }
                         2 => MouseEvent::Press(MouseButton::Right, cx, cy),
                         3 => MouseEvent::Release(cx, cy),
                         _ => return None,
                     })
    }
             Some(Ok(b'<')) => {
        // xterm mouse encoding:
        // ESC [ < Cb ; Cx ; Cy (;) (M or m)
        let mut buf = Vec::new();
        let mut c = iter.next().unwrap().unwrap();
        while match c {
                  b'm' | b'M' => false,
                  _ => true,
              } {
            buf.push(c);
            c = iter.next().unwrap().unwrap();
        }
        let str_buf = String::from_utf8(buf).unwrap();
        let nums = &mut str_buf.split(';');

        let cb = nums.next()
            .unwrap()
            .parse::<u16>()
            .unwrap();
        let cx = nums.next()
            .unwrap()
            .parse::<u16>()
            .unwrap();
        let cy = nums.next()
            .unwrap()
            .parse::<u16>()
            .unwrap();

        let event = match cb {
            0...2 | 64...65 => {
                let button = match cb {
                    0 => MouseButton::Left,
                    1 => MouseButton::Middle,
                    2 => MouseButton::Right,
                    64 => MouseButton::WheelUp,
                    65 => MouseButton::WheelDown,
                    _ => unreachable!(),
                };
                match c {
                    b'M' => MouseEvent::Press(button, cx, cy),
                    b'm' => MouseEvent::Release(cx, cy),
                    _ => return None,
                }
            }
            32 => MouseEvent::Hold(cx, cy),
            3 => MouseEvent::Release(cx, cy),
            _ => return None,
        };

        Event::Mouse(event)
    }
             Some(Ok(c @ b'0'...b'9')) => {
        // Numbered escape code.
        let mut buf = Vec::new();
        buf.push(c);
        let mut c = iter.next().unwrap().unwrap();
        // The final byte of a CSI sequence can be in the range 64-126, so
        // let's keep reading anything else.
        while c < 64 || c > 126 {
            buf.push(c);
            c = iter.next().unwrap().unwrap();
        }

        match c {
            // rxvt mouse encoding:
            // ESC [ Cb ; Cx ; Cy ; M
            b'M' => {
                let str_buf = String::from_utf8(buf).unwrap();

                let nums: Vec<u16> = str_buf.split(';').map(|n| n.parse().unwrap()).collect();

                let cb = nums[0];
                let cx = nums[1];
                let cy = nums[2];

                let event = match cb {
                    32 => MouseEvent::Press(MouseButton::Left, cx, cy),
                    33 => MouseEvent::Press(MouseButton::Middle, cx, cy),
                    34 => MouseEvent::Press(MouseButton::Right, cx, cy),
                    35 => MouseEvent::Release(cx, cy),
                    64 => MouseEvent::Hold(cx, cy),
                    96 | 97 => MouseEvent::Press(MouseButton::WheelUp, cx, cy),
                    _ => return None,
                };

                Event::Mouse(event)
            }
            // Special key code.
            b'~' => {
                let str_buf = String::from_utf8(buf).unwrap();

                // This CSI sequence can be a list of semicolon-separated
                // numbers.
                let nums: Vec<u8> = str_buf.split(';').map(|n| n.parse().unwrap()).collect();

                if nums.is_empty() {
                    return None;
                }

                // TODO: handle multiple values for key modififiers (ex: values
                // [3, 2] means Shift+Delete)
                if nums.len() > 1 {
                    return None;
                }

                match nums[0] {
                    1 | 7 => Event::Key(Key::Home),
                    2 => Event::Key(Key::Insert),
                    3 => Event::Key(Key::Delete),
                    4 | 8 => Event::Key(Key::End),
                    5 => Event::Key(Key::PageUp),
                    6 => Event::Key(Key::PageDown),
                    v @ 11...15 => Event::Key(Key::F(v - 10)),
                    v @ 17...21 => Event::Key(Key::F(v - 11)),
                    v @ 23...24 => Event::Key(Key::F(v - 12)),
                    _ => return None,
                }
            }
            _ => return None,
        }
    }
             _ => return None,
         })

}

/// Parse `c` as either a single byte ASCII char or a variable size UTF-8 char.
fn parse_utf8_char<I>(c: u8, iter: &mut I) -> Result<char, Error>
    where I: Iterator<Item = Result<u8, Error>>
{
    let error = Err(Error::new(ErrorKind::Other, "Input character is not valid UTF-8"));
    if c.is_ascii() {
        Ok(c as char)
    } else {
        let bytes = &mut Vec::new();
        bytes.push(c);

        loop {
            match iter.next() {
                Some(Ok(next)) => {
                    bytes.push(next);
                    if let Ok(st) = str::from_utf8(bytes) {
                        return Ok(st.chars().next().unwrap());
                    }
                    if bytes.len() >= 4 {
                        return error;
                    }
                }
                _ => return error,
            }
        }
    }
}

#[cfg(test)]
#[test]
fn test_parse_utf8() {
    let st = "abcéŷ¤£€ù%323";
    let ref mut bytes = st.bytes().map(|x| Ok(x));
    let chars = st.chars();
    for c in chars {
        let b = bytes.next().unwrap().unwrap();
        assert!(c == parse_utf8_char(b, bytes).unwrap());
    }
}
<commit_msg>backtab key<commit_after>//! Mouse and key events.

use std::io::{Error, ErrorKind};
use std::str;

/// An event reported by the terminal.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Event {
    /// A key press.
    Key(Key),
    /// A mouse button press, release or wheel use at specific coordinates.
    Mouse(MouseEvent),
    /// An event that cannot currently be evaluated.
    Unsupported(Vec<u8>),
}

/// A mouse related event.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum MouseEvent {
    /// A mouse button was pressed.
    ///
    /// The coordinates are one-based.
    Press(MouseButton, u16, u16),
    /// A mouse button was released.
    ///
    /// The coordinates are one-based.
    Release(u16, u16),
    /// A mouse button is held over the given coordinates.
    ///
    /// The coordinates are one-based.
    Hold(u16, u16),
}

/// A mouse button.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum MouseButton {
    /// The left mouse button.
    Left,
    /// The right mouse button.
    Right,
    /// The middle mouse button.
    Middle,
    /// Mouse wheel is going up.
    ///
    /// This event is typically only used with Mouse::Press.
    WheelUp,
    /// Mouse wheel is going down.
    ///
    /// This event is typically only used with Mouse::Press.
    WheelDown,
}

/// A key.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum Key {
    /// Backspace.
    Backspace,
    /// Left arrow.
    Left,
    /// Right arrow.
    Right,
    /// Up arrow.
    Up,
    /// Down arrow.
    Down,
    /// Home key.
    Home,
    /// End key.
    End,
    /// Page Up key.
    PageUp,
    /// Page Down key.
    PageDown,
    /// Backward Tab key.
    BackTab,
    /// Delete key.
    Delete,
    /// Insert key.
    Insert,
    /// Function keys.
    ///
    /// Only function keys 1 through 12 are supported.
    F(u8),
    /// Normal character.
    Char(char),
    /// Alt modified character.
    Alt(char),
    /// Ctrl modified character.
    ///
    /// Note that certain keys may not be modifiable with `ctrl`, due to limitations of terminals.
    Ctrl(char),
    /// Null byte.
    Null,
    /// Esc key.
    Esc,

    #[doc(hidden)]
    __IsNotComplete,
}

/// Parse an Event from `item` and possibly subsequent bytes through `iter`.
pub fn parse_event<I>(item: u8, iter: &mut I) -> Result<Event, Error>
    where I: Iterator<Item = Result<u8, Error>>
{
    let error = Error::new(ErrorKind::Other, "Could not parse an event");
    match item {
        b'\x1B' => {
            // This is an escape character, leading a control sequence.
            Ok(match iter.next() {
                   Some(Ok(b'O')) => {
                match iter.next() {
                    // F1-F4
                    Some(Ok(val @ b'P'...b'S')) => Event::Key(Key::F(1 + val - b'P')),
                    _ => return Err(error),
                }
            }
                   Some(Ok(b'[')) => {
                // This is a CSI sequence.
                parse_csi(iter).ok_or(error)?
            }
                   Some(Ok(c)) => {
                let ch = parse_utf8_char(c, iter);
                Event::Key(Key::Alt(try!(ch)))
            }
                   Some(Err(_)) | None => return Err(error),
               })
        }
        b'\n' | b'\r' => Ok(Event::Key(Key::Char('\n'))),
        b'\t' => Ok(Event::Key(Key::Char('\t'))),
        b'\x7F' => Ok(Event::Key(Key::Backspace)),
        c @ b'\x01'...b'\x1A' => Ok(Event::Key(Key::Ctrl((c as u8 - 0x1 + b'a') as char))),
        c @ b'\x1C'...b'\x1F' => Ok(Event::Key(Key::Ctrl((c as u8 - 0x1C + b'4') as char))),
        b'\0' => Ok(Event::Key(Key::Null)),
        c => {
            Ok({
                   let ch = parse_utf8_char(c, iter);
                   Event::Key(Key::Char(try!(ch)))
               })
        }
    }
}

/// Parses a CSI sequence, just after reading ^[
///
/// Returns None if an unrecognized sequence is found.
fn parse_csi<I>(iter: &mut I) -> Option<Event>
    where I: Iterator<Item = Result<u8, Error>>
{
    Some(match iter.next() {
             Some(Ok(b'[')) => match iter.next() {
                 Some(Ok(val @ b'A'...b'E')) => Event::Key(Key::F(1 + val - b'A')),
                 _ => return None,
             },
             Some(Ok(b'D')) => Event::Key(Key::Left),
             Some(Ok(b'C')) => Event::Key(Key::Right),
             Some(Ok(b'A')) => Event::Key(Key::Up),
             Some(Ok(b'B')) => Event::Key(Key::Down),
             Some(Ok(b'H')) => Event::Key(Key::Home),
             Some(Ok(b'F')) => Event::Key(Key::End),
             Some(Ok(b'Z')) => Event::Key(Key::BackTab),
             Some(Ok(b'M')) => {
        // X10 emulation mouse encoding: ESC [ CB Cx Cy (6 characters only).
        let mut next = || iter.next().unwrap().unwrap();

        let cb = next() as i8 - 32;
        // (1, 1) are the coords for upper left.
        let cx = next().saturating_sub(32) as u16;
        let cy = next().saturating_sub(32) as u16;
        Event::Mouse(match cb & 0b11 {
                         0 => {
                             if cb & 0x40 != 0 {
                                 MouseEvent::Press(MouseButton::WheelUp, cx, cy)
                             } else {
                                 MouseEvent::Press(MouseButton::Left, cx, cy)
                             }
                         }
                         1 => {
                             if cb & 0x40 != 0 {
                                 MouseEvent::Press(MouseButton::WheelDown, cx, cy)
                             } else {
                                 MouseEvent::Press(MouseButton::Middle, cx, cy)
                             }
                         }
                         2 => MouseEvent::Press(MouseButton::Right, cx, cy),
                         3 => MouseEvent::Release(cx, cy),
                         _ => return None,
                     })
    }
             Some(Ok(b'<')) => {
        // xterm mouse encoding:
        // ESC [ < Cb ; Cx ; Cy (;) (M or m)
        let mut buf = Vec::new();
        let mut c = iter.next().unwrap().unwrap();
        while match c {
                  b'm' | b'M' => false,
                  _ => true,
              } {
            buf.push(c);
            c = iter.next().unwrap().unwrap();
        }
        let str_buf = String::from_utf8(buf).unwrap();
        let nums = &mut str_buf.split(';');

        let cb = nums.next()
            .unwrap()
            .parse::<u16>()
            .unwrap();
        let cx = nums.next()
            .unwrap()
            .parse::<u16>()
            .unwrap();
        let cy = nums.next()
            .unwrap()
            .parse::<u16>()
            .unwrap();

        let event = match cb {
            0...2 | 64...65 => {
                let button = match cb {
                    0 => MouseButton::Left,
                    1 => MouseButton::Middle,
                    2 => MouseButton::Right,
                    64 => MouseButton::WheelUp,
                    65 => MouseButton::WheelDown,
                    _ => unreachable!(),
                };
                match c {
                    b'M' => MouseEvent::Press(button, cx, cy),
                    b'm' => MouseEvent::Release(cx, cy),
                    _ => return None,
                }
            }
            32 => MouseEvent::Hold(cx, cy),
            3 => MouseEvent::Release(cx, cy),
            _ => return None,
        };

        Event::Mouse(event)
    }
             Some(Ok(c @ b'0'...b'9')) => {
        // Numbered escape code.
        let mut buf = Vec::new();
        buf.push(c);
        let mut c = iter.next().unwrap().unwrap();
        // The final byte of a CSI sequence can be in the range 64-126, so
        // let's keep reading anything else.
        while c < 64 || c > 126 {
            buf.push(c);
            c = iter.next().unwrap().unwrap();
        }

        match c {
            // rxvt mouse encoding:
            // ESC [ Cb ; Cx ; Cy ; M
            b'M' => {
                let str_buf = String::from_utf8(buf).unwrap();

                let nums: Vec<u16> = str_buf.split(';').map(|n| n.parse().unwrap()).collect();

                let cb = nums[0];
                let cx = nums[1];
                let cy = nums[2];

                let event = match cb {
                    32 => MouseEvent::Press(MouseButton::Left, cx, cy),
                    33 => MouseEvent::Press(MouseButton::Middle, cx, cy),
                    34 => MouseEvent::Press(MouseButton::Right, cx, cy),
                    35 => MouseEvent::Release(cx, cy),
                    64 => MouseEvent::Hold(cx, cy),
                    96 | 97 => MouseEvent::Press(MouseButton::WheelUp, cx, cy),
                    _ => return None,
                };

                Event::Mouse(event)
            }
            // Special key code.
            b'~' => {
                let str_buf = String::from_utf8(buf).unwrap();

                // This CSI sequence can be a list of semicolon-separated
                // numbers.
                let nums: Vec<u8> = str_buf.split(';').map(|n| n.parse().unwrap()).collect();

                if nums.is_empty() {
                    return None;
                }

                // TODO: handle multiple values for key modififiers (ex: values
                // [3, 2] means Shift+Delete)
                if nums.len() > 1 {
                    return None;
                }

                match nums[0] {
                    1 | 7 => Event::Key(Key::Home),
                    2 => Event::Key(Key::Insert),
                    3 => Event::Key(Key::Delete),
                    4 | 8 => Event::Key(Key::End),
                    5 => Event::Key(Key::PageUp),
                    6 => Event::Key(Key::PageDown),
                    v @ 11...15 => Event::Key(Key::F(v - 10)),
                    v @ 17...21 => Event::Key(Key::F(v - 11)),
                    v @ 23...24 => Event::Key(Key::F(v - 12)),
                    _ => return None,
                }
            }
            _ => return None,
        }
    }
             _ => return None,
         })

}

/// Parse `c` as either a single byte ASCII char or a variable size UTF-8 char.
fn parse_utf8_char<I>(c: u8, iter: &mut I) -> Result<char, Error>
    where I: Iterator<Item = Result<u8, Error>>
{
    let error = Err(Error::new(ErrorKind::Other, "Input character is not valid UTF-8"));
    if c.is_ascii() {
        Ok(c as char)
    } else {
        let bytes = &mut Vec::new();
        bytes.push(c);

        loop {
            match iter.next() {
                Some(Ok(next)) => {
                    bytes.push(next);
                    if let Ok(st) = str::from_utf8(bytes) {
                        return Ok(st.chars().next().unwrap());
                    }
                    if bytes.len() >= 4 {
                        return error;
                    }
                }
                _ => return error,
            }
        }
    }
}

#[cfg(test)]
#[test]
fn test_parse_utf8() {
    let st = "abcéŷ¤£€ù%323";
    let ref mut bytes = st.bytes().map(|x| Ok(x));
    let chars = st.chars();
    for c in chars {
        let b = bytes.next().unwrap().unwrap();
        assert!(c == parse_utf8_char(b, bytes).unwrap());
    }
}
<|endoftext|>


<commit_before>use core::ops::{BitXorAssign, MulAssign};
use byteorder::{BigEndian, ByteOrder};

pub(crate) fn ghash(key: &[u8; 16], data: &[u8], ciphertext: &[u8]) -> [u8; 16] {
    let mut tag = [0; 16];
    let mut mac = GHash::new(key, data);
    mac.update(ciphertext);
    mac.write_tag(&mut tag);
    tag
}

const R0: u64 = 0xe1 << 56;

struct GHash {
    function: PolyFunction,
    data_len: u64,
    ciphertext_len: u64,
}

impl GHash {
    fn new(key: &[u8; 16], data: &[u8]) -> Self {
        let mut ghash = Self {
            function: PolyFunction::new(key),
            data_len: data.len() as u64,
            ciphertext_len: 0,
        };
        ghash.process(data);
        ghash
    }

    fn update(&mut self, input: &[u8]) {
        self.ciphertext_len += input.len() as u64;
        self.process(input);
    }

    fn write_tag(mut self, output: &mut [u8; 16]) {
        BigEndian::write_u64(&mut output[..8], 8 * self.data_len);
        BigEndian::write_u64(&mut output[8..], 8 * self.ciphertext_len);
        self.function.process(output);
        self.function.write_value(output);
    }

    fn process(&mut self, input: &[u8]) {
        for chunk in input.chunks(16) {
            if chunk.len() < 16 {
                let buffer = &mut [0; 16];
                buffer[..chunk.len()].copy_from_slice(chunk);
                self.function.process(buffer);
            } else {
                self.function.process(chunk);
            }
        }
    }
}

#[derive(Clone, Copy)]
struct GFBlock([u64; 2]);

struct PolyFunction {
    key_block: GFBlock,
    state: GFBlock,
}

impl PolyFunction {
    fn new(key: &[u8; 16]) -> Self {
        Self {
            key_block: GFBlock::new(key),
            state: GFBlock([0; 2]),
        }
    }

    fn process(&mut self, input: &[u8]) {
        self.state ^= GFBlock::new(input);
        self.state *= self.key_block;
    }

    fn write_value(self, output: &mut [u8; 16]) {
        BigEndian::write_u64_into(&self.state.0, output);
    }
}

impl GFBlock {
    fn new(bytes: &[u8]) -> Self {
        let mut block = [0; 2];
        BigEndian::read_u64_into(bytes, &mut block);
        GFBlock(block)
    }
}

impl BitXorAssign for GFBlock {
    fn bitxor_assign(&mut self, rhs: Self) {
        for (l, r) in self.0.iter_mut().zip(&rhs.0) {
            *l ^= r;
        }
    }
}

impl MulAssign for GFBlock {
    fn mul_assign(&mut self, rhs: Self) {
        let mut z = GFBlock([0; 2]);
        let mut v = rhs;
        for xp in &mut self.0 {
            for _ in 0..64 {
                let mut h = *xp & (1 << 63);
                let mut m = (h as i64 >> 63) as u64;
                for (l, r) in z.0.iter_mut().zip(&v.0) {
                    *l ^= r & m;
                }

                h = v.0[1] << 63;
                m = (h as i64 >> 7) as u64;
                v.0[1] >>= 1;
                v.0[1] |= v.0[0] << 63;
                v.0[0] >>= 1;
                v.0[0] ^= R0 & m;
                *xp <<= 1;
            }
        }
        self.0 = z.0;
    }
}

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

    fn check(expected: &str, h: &str, a: &str, c: &str) {
        let h_vec = &h2b(h);
        let h = &mut [0; 16];
        h.copy_from_slice(h_vec);
        let a = &h2b(a);
        let c = &h2b(c);
        let expected = h2b(expected);
        assert_eq!(expected, ghash(h, a, c));
    }

    #[test]
    fn test_case_1_2() {
        let h = "66e94bd4ef8a2c3b884cfa59ca342b2e";
        let a = "";
        let c = "";
        let expected = "00000000000000000000000000000000";
        check(expected, h, a, c);

        let c = "0388dace60b6a392f328c2b971b2fe78";
        let expected = "f38cbb1ad69223dcc3457ae5b6b0f885";
        check(expected, h, a, c);
    }

    #[test]
    fn test_case_15_16_17_18() {
        let h = "acbef20579b4b8ebce889bac8732dad7";
        let a = "";
        let c = "522dc1f099567d07f47f37a32a84427d643a8cdcbfe5c0c97598a2bd2555d1aa\
                 8cb08e48590dbb3da7b08b1056828838c5f61e6393ba7a0abcc9f662898015ad";
        let expected = "4db870d37cb75fcb46097c36230d1612";
        check(expected, h, a, c);

        let a = "feedfacedeadbeeffeedfacedeadbeefabaddad2";
        let c = &c[..120];
        let expected = "8bd0c4d8aacd391e67cca447e8c38f65";
        check(expected, h, a, c);

        let c = "c3762df1ca787d32ae47c13bf19844cbaf1ae14d0b976afac52ff7d79bba9de0\
                 feb582d33934a4f0954cc2363bc73f7862ac430e64abe499f47c9b1f";
        let expected = "75a34288b8c68f811c52b2e9a2f97f63";
        check(expected, h, a, c);

        let c = "5a8def2f0c9e53f1f75d7853659e2a20eeb2b22aafde6419a058ab4f6f746bf4\
                 0fc0c3b780f244452da3ebf1c5d82cdea2418997200ef82e44ae7e3f";
        let expected = "d5ffcf6fc5ac4d69722187421a7f170b";
        check(expected, h, a, c);
    }
}
<commit_msg>ghash: add tests<commit_after>use core::ops::{BitXorAssign, MulAssign};
use byteorder::{BigEndian, ByteOrder};

pub(crate) fn ghash(key: &[u8; 16], data: &[u8], ciphertext: &[u8]) -> [u8; 16] {
    let mut tag = [0; 16];
    let mut mac = GHash::new(key, data);
    mac.update(ciphertext);
    mac.write_tag(&mut tag);
    tag
}

const R0: u64 = 0xe1 << 56;

struct GHash {
    function: PolyFunction,
    data_len: u64,
    ciphertext_len: u64,
}

impl GHash {
    fn new(key: &[u8; 16], data: &[u8]) -> Self {
        let mut ghash = Self {
            function: PolyFunction::new(key),
            data_len: data.len() as u64,
            ciphertext_len: 0,
        };
        ghash.process(data);
        ghash
    }

    fn update(&mut self, input: &[u8]) {
        self.ciphertext_len += input.len() as u64;
        self.process(input);
    }

    fn write_tag(mut self, output: &mut [u8; 16]) {
        BigEndian::write_u64(&mut output[..8], 8 * self.data_len);
        BigEndian::write_u64(&mut output[8..], 8 * self.ciphertext_len);
        self.function.process(output);
        self.function.write_value(output);
    }

    fn process(&mut self, input: &[u8]) {
        for chunk in input.chunks(16) {
            if chunk.len() < 16 {
                let buffer = &mut [0; 16];
                buffer[..chunk.len()].copy_from_slice(chunk);
                self.function.process(buffer);
            } else {
                self.function.process(chunk);
            }
        }
    }
}

#[derive(Clone, Copy)]
struct GFBlock([u64; 2]);

struct PolyFunction {
    key_block: GFBlock,
    state: GFBlock,
}

impl PolyFunction {
    fn new(key: &[u8; 16]) -> Self {
        Self {
            key_block: GFBlock::new(key),
            state: GFBlock([0; 2]),
        }
    }

    fn process(&mut self, input: &[u8]) {
        self.state ^= GFBlock::new(input);
        self.state *= self.key_block;
    }

    fn write_value(self, output: &mut [u8; 16]) {
        BigEndian::write_u64_into(&self.state.0, output);
    }
}

impl GFBlock {
    fn new(bytes: &[u8]) -> Self {
        let mut block = [0; 2];
        BigEndian::read_u64_into(bytes, &mut block);
        GFBlock(block)
    }
}

impl BitXorAssign for GFBlock {
    fn bitxor_assign(&mut self, rhs: Self) {
        for (l, r) in self.0.iter_mut().zip(&rhs.0) {
            *l ^= r;
        }
    }
}

impl MulAssign for GFBlock {
    fn mul_assign(&mut self, rhs: Self) {
        let mut z = GFBlock([0; 2]);
        let mut v = rhs;
        for xp in &mut self.0 {
            for _ in 0..64 {
                let mut h = *xp & (1 << 63);
                let mut m = (h as i64 >> 63) as u64;
                for (l, r) in z.0.iter_mut().zip(&v.0) {
                    *l ^= r & m;
                }

                h = v.0[1] << 63;
                m = (h as i64 >> 7) as u64;
                v.0[1] >>= 1;
                v.0[1] |= v.0[0] << 63;
                v.0[0] >>= 1;
                v.0[0] ^= R0 & m;
                *xp <<= 1;
            }
        }
        self.0 = z.0;
    }
}

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

    fn check(expected: &str, h: &str, a: &str, c: &str) {
        let h_vec = &h2b(h);
        let h = &mut [0; 16];
        h.copy_from_slice(h_vec);
        let a = &h2b(a);
        let c = &h2b(c);
        let expected = h2b(expected);
        assert_eq!(expected, ghash(h, a, c));
    }

    #[test]
    fn test_case_1_2() {
        let h = "66e94bd4ef8a2c3b884cfa59ca342b2e";
        let a = "";
        let c = "";
        let expected = "00000000000000000000000000000000";
        check(expected, h, a, c);

        let c = "0388dace60b6a392f328c2b971b2fe78";
        let expected = "f38cbb1ad69223dcc3457ae5b6b0f885";
        check(expected, h, a, c);
    }

    #[test]
    fn test_case_3_4_5_6() {
        let h = "b83b533708bf535d0aa6e52980d53b78";
        let a = "";
        let c = "42831ec2217774244b7221b784d0d49ce3aa212f2c02a4e035c17e2329aca12e\
                 21d514b25466931c7d8f6a5aac84aa051ba30b396a0aac973d58e091473f5985";
        let expected = "7f1b32b81b820d02614f8895ac1d4eac";
        check(expected, h, a, c);

        let a = "feedfacedeadbeeffeedfacedeadbeefabaddad2";
        let c = &c[..120];
        let expected = "698e57f70e6ecc7fd9463b7260a9ae5f";
        check(expected, h, a, c);

        let c = "61353b4c2806934a777ff51fa22a4755699b2a714fcdc6f83766e5f97b6c7423\
                 73806900e49f24b22b097544d4896b424989b5e1ebac0f07c23f4598";
        let expected = "df586bb4c249b92cb6922877e444d37b";
        check(expected, h, a, c);

        let c = "8ce24998625615b603a033aca13fb894be9112a5c3a211a8ba262a3cca7e2ca7\
                 01e4a9a4fba43c90ccdcb281d48c7c6fd62875d2aca417034c34aee5";
        let expected = "1c5afe9760d3932f3c9a878aac3dc3de";
        check(expected, h, a, c);
    }

    #[test]
    fn test_case_7_8() {
        let h = "aae06992acbf52a3e8f4a96ec9300bd7";
        let a = "";
        let c = "";
        let expected = "00000000000000000000000000000000";
        check(expected, h, a, c);

        let c = "98e7247c07f0fe411c267e4384b0f600";
        let expected = "e2c63f0ac44ad0e02efa05ab6743d4ce";
        check(expected, h, a, c);
    }

    #[test]
    fn test_case_9_10_11_12() {
        let h = "466923ec9ae682214f2c082badb39249";
        let a = "";
        let c = "3980ca0b3c00e841eb06fac4872a2757859e1ceaa6efd984628593b40ca1e19c\
                 7d773d00c144c525ac619d18c84a3f4718e2448b2fe324d9ccda2710acade256";
        let expected = "51110d40f6c8fff0eb1ae33445a889f0";
        check(expected, h, a, c);

        let a = "feedfacedeadbeeffeedfacedeadbeefabaddad2";
        let c = &c[..120];
        let expected = "ed2ce3062e4a8ec06db8b4c490e8a268";
        check(expected, h, a, c);

        let c = "0f10f599ae14a154ed24b36e25324db8c566632ef2bbb34f8347280fc4507057\
                 fddc29df9a471f75c66541d4d4dad1c9e93a19a58e8b473fa0f062f7";
        let expected = "1e6a133806607858ee80eaf237064089";
        check(expected, h, a, c);

        let c = "d27e88681ce3243c4830165a8fdcf9ff1de9a1d8e6b447ef6ef7b79828666e45\
                 81e79012af34ddd9e2f037589b292db3e67c036745fa22e7e9b7373b";
        let expected = "82567fb0b4cc371801eadec005968e94";
        check(expected, h, a, c);
    }

    #[test]
    fn test_case_13_14() {
        let h = "dc95c078a2408989ad48a21492842087";
        let a = "";
        let c = "";
        let expected = "00000000000000000000000000000000";
        check(expected, h, a, c);

        let c = "cea7403d4d606b6e074ec5d3baf39d18";
        let expected = "83de425c5edc5d498f382c441041ca92";
        check(expected, h, a, c);
    }

    #[test]
    fn test_case_15_16_17_18() {
        let h = "acbef20579b4b8ebce889bac8732dad7";
        let a = "";
        let c = "522dc1f099567d07f47f37a32a84427d643a8cdcbfe5c0c97598a2bd2555d1aa\
                 8cb08e48590dbb3da7b08b1056828838c5f61e6393ba7a0abcc9f662898015ad";
        let expected = "4db870d37cb75fcb46097c36230d1612";
        check(expected, h, a, c);

        let a = "feedfacedeadbeeffeedfacedeadbeefabaddad2";
        let c = &c[..120];
        let expected = "8bd0c4d8aacd391e67cca447e8c38f65";
        check(expected, h, a, c);

        let c = "c3762df1ca787d32ae47c13bf19844cbaf1ae14d0b976afac52ff7d79bba9de0\
                 feb582d33934a4f0954cc2363bc73f7862ac430e64abe499f47c9b1f";
        let expected = "75a34288b8c68f811c52b2e9a2f97f63";
        check(expected, h, a, c);

        let c = "5a8def2f0c9e53f1f75d7853659e2a20eeb2b22aafde6419a058ab4f6f746bf4\
                 0fc0c3b780f244452da3ebf1c5d82cdea2418997200ef82e44ae7e3f";
        let expected = "d5ffcf6fc5ac4d69722187421a7f170b";
        check(expected, h, a, c);
    }
}
<|endoftext|>


<commit_before>use std::hash::Hash;
use std::collections::{HashMap, HashSet, VecDeque};

use crate::colour::Colour;
use crate::grid::Grid;

// Represent the state of the game as a graph.  One node for each cell.  There is an edge between
// two cells if and only if they are the same colour and have a distance of 1 according to the
// Manhattan metric.
//
// The game has only one real operation (i.e. not counting level generation): Switching the colour
// of the connected component containing the top left cell to whichever colour the clicked cell
// has.  This results in said component growing by every component it neighbours which has its new
// colour.  In essence, the aim of the gaim is to join connected components until only one is
// left.  Fewer clicks is better.
//
// So, we represent the game state as a collection of connected components, each of which has a
// colour, a list of coordinates of the constituent cells, and a list of it s neighbours.  As we
// do all our operations on connected components, we actually work not with the graph of
// neighbouring cells but with a graph of connected components.  The graph of connected components
// is made up of one node representing each connected component and an edge between any two
// connected components that have cells which share an edge, i.e. two cells which would have been
// connected by an edge in the original graph if they had had the same colour (which they did
// not).

#[derive(Clone, Eq, PartialEq, Hash)]
struct Position {
    row: usize,
    column: usize,
}

#[derive(Clone)]
struct ConnectedComponent {
    counter: usize,
    colour: Colour,
    cells: HashSet<Position>,
}

impl PartialEq for ConnectedComponent {
    fn eq(&self, other: &Self) -> bool {
        self.counter == other.counter
    }
}

impl Eq for ConnectedComponent {}

impl Hash for ConnectedComponent {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.counter.hash(state);
    }
}

pub struct Graph {
    neighbours: HashMap<ConnectedComponent, HashSet<ConnectedComponent>>,
}

fn find_connected_components(grid: &Grid) -> Vec<(ConnectedComponent, HashSet<usize>)> {
    let rows = grid.number_of_rows;
    let columns = grid.number_of_columns;

    let mut remaining_cells: HashSet<usize> = (0..grid.cells.len()).into_iter().collect();

    let mut components = vec![];
    let mut counter = 0;

    // Take any element to land in any component.
    while let Some(&start_cell_index) = remaining_cells.iter().next() {
        let mut visited = HashSet::new();
        let mut neighbours = HashSet::new();
        let mut queue = VecDeque::new();

        queue.push_back(start_cell_index);

        // Gather all cells in the connected component of start_cell_index.
        while let Some(i) = queue.pop_front() {
            visited.insert(i);
            remaining_cells.remove(&i);

            let mut neighbouring_cells = vec![];
            if i % columns > 0 {
                neighbouring_cells.push(i - 1);
            }
            if i % columns < columns - 1 {
                neighbouring_cells.push(i + 1);
            }
            if i >= columns {
                neighbouring_cells.push(i - columns)
            }
            if i < (rows - 1 ) * columns {
                neighbouring_cells.push(i + columns)
            }

            for neighbour in neighbouring_cells {
                if visited.contains(&neighbour) {
                    continue
                }

                if grid.cells[i] == grid.cells[neighbour] {
                    queue.push_back(neighbour);
                } else {
                    neighbours.insert(neighbour);
                }
            }
        }

        let cells = visited
            .iter()
            .map(|&i| Position {
                column: i % columns,
                row: i / columns,
            })
            .collect();

        let component = ConnectedComponent {
            counter,
            colour: grid.cells[*visited.iter().next().unwrap()],
            cells,
        };
        counter += 1;

        components.push((component, neighbours));
    }

    return components
}

impl Graph {
    pub fn create(grid: &Grid) -> Self {
        let components_and_neighbours = find_connected_components(&grid);

        let map_cell_to_component: HashMap<Position, ConnectedComponent> = components_and_neighbours
        .iter()
        .flat_map(|(component, _)| {
            component.cells
            .iter()
            .map(|cell| (cell.clone(), component.clone()))
            .collect::<Vec<(Position, ConnectedComponent)>>()
        })
        .collect();

        let mut neighbours:  HashMap<ConnectedComponent, HashSet<ConnectedComponent>> = HashMap::new();

        for (component, neighbour_cells) in components_and_neighbours {
            let mut neighbour_components: HashSet<ConnectedComponent> = HashSet::new();

            for cell in neighbour_cells {
                let position = Position{column: cell % grid.number_of_columns, row: cell / grid.number_of_columns};
                neighbour_components.insert(map_cell_to_component[&position].clone());
            }

            neighbours.insert(component, neighbour_components);
        }

        Self { neighbours }
    }

    fn to_grid(&self, rows: usize, columns: usize) -> Grid {
        let mut cells = vec![Colour::Red; rows * columns];

        for component in self.neighbours.keys() {
            for position in &component.cells {
                cells[position.column + position.row * columns] = component.colour;
            }
        }

        Grid { number_of_rows: rows, number_of_columns: columns, cells, }
    }
}

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

    #[test]
    fn should_return_same_grid_as_input() {
        let size = 4;
        let grid = Grid::generate(size);

        let graph = Graph::create(&grid);
        let reconstituted_grid = graph.to_grid(size, size);

        assert_eq!(reconstituted_grid, grid);
    }

    #[test]
    fn should_have_one_component() {
        let grid = Grid{
            number_of_columns: 2,
             number_of_rows: 2,
             cells: vec![Colour::Red; 4]
        };

        let graph = Graph::create(&grid);

        assert_eq!(graph.neighbours.len(), 1);
    }

    #[test]
    fn should_have_four_component() {
        let grid = Grid{
            number_of_columns: 2,
             number_of_rows: 2,
             cells: vec![Colour::Red, Colour::Yellow, Colour::Yellow, Colour::Red]
        };

        let graph = Graph::create(&grid);

        assert_eq!(graph.neighbours.len(), 4);
    }

    #[test]
    fn should_have_two_component() {
        let grid = Grid{
            number_of_columns: 2,
             number_of_rows: 2,
             cells: vec![Colour::Red, Colour::Red, Colour::Yellow, Colour::Yellow]
        };

        let graph = Graph::create(&grid);

        assert_eq!(graph.neighbours.len(), 2);
    }
}
<commit_msg>Make a bunch of stuff public for rendering<commit_after>use std::hash::Hash;
use std::collections::{HashMap, HashSet, VecDeque};

use crate::colour::Colour;
use crate::grid::Grid;

// Represent the state of the game as a graph.  One node for each cell.  There is an edge between
// two cells if and only if they are the same colour and have a distance of 1 according to the
// Manhattan metric.
//
// The game has only one real operation (i.e. not counting level generation): Switching the colour
// of the connected component containing the top left cell to whichever colour the clicked cell
// has.  This results in said component growing by every component it neighbours which has its new
// colour.  In essence, the aim of the gaim is to join connected components until only one is
// left.  Fewer clicks is better.
//
// So, we represent the game state as a collection of connected components, each of which has a
// colour, a list of coordinates of the constituent cells, and a list of it s neighbours.  As we
// do all our operations on connected components, we actually work not with the graph of
// neighbouring cells but with a graph of connected components.  The graph of connected components
// is made up of one node representing each connected component and an edge between any two
// connected components that have cells which share an edge, i.e. two cells which would have been
// connected by an edge in the original graph if they had had the same colour (which they did
// not).

#[derive(Clone, Eq, PartialEq, Hash)]
pub struct Position {
    pub row: usize,
    pub column: usize,
}

#[derive(Clone)]
pub struct ConnectedComponent {
    counter: usize,
    pub colour: Colour,
    pub cells: HashSet<Position>,
}

impl PartialEq for ConnectedComponent {
    fn eq(&self, other: &Self) -> bool {
        self.counter == other.counter
    }
}

impl Eq for ConnectedComponent {}

impl Hash for ConnectedComponent {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.counter.hash(state);
    }
}

pub struct Graph {
    pub neighbours: HashMap<ConnectedComponent, HashSet<ConnectedComponent>>,
}

fn find_connected_components(grid: &Grid) -> Vec<(ConnectedComponent, HashSet<usize>)> {
    let rows = grid.number_of_rows;
    let columns = grid.number_of_columns;

    let mut remaining_cells: HashSet<usize> = (0..grid.cells.len()).into_iter().collect();

    let mut components = vec![];
    let mut counter = 0;

    // Take any element to land in any component.
    while let Some(&start_cell_index) = remaining_cells.iter().next() {
        let mut visited = HashSet::new();
        let mut neighbours = HashSet::new();
        let mut queue = VecDeque::new();

        queue.push_back(start_cell_index);

        // Gather all cells in the connected component of start_cell_index.
        while let Some(i) = queue.pop_front() {
            visited.insert(i);
            remaining_cells.remove(&i);

            let mut neighbouring_cells = vec![];
            if i % columns > 0 {
                neighbouring_cells.push(i - 1);
            }
            if i % columns < columns - 1 {
                neighbouring_cells.push(i + 1);
            }
            if i >= columns {
                neighbouring_cells.push(i - columns)
            }
            if i < (rows - 1 ) * columns {
                neighbouring_cells.push(i + columns)
            }

            for neighbour in neighbouring_cells {
                if visited.contains(&neighbour) {
                    continue
                }

                if grid.cells[i] == grid.cells[neighbour] {
                    queue.push_back(neighbour);
                } else {
                    neighbours.insert(neighbour);
                }
            }
        }

        let cells = visited
            .iter()
            .map(|&i| Position {
                column: i % columns,
                row: i / columns,
            })
            .collect();

        let component = ConnectedComponent {
            counter,
            colour: grid.cells[*visited.iter().next().unwrap()],
            cells,
        };
        counter += 1;

        components.push((component, neighbours));
    }

    return components
}

impl Graph {
    pub fn create(grid: &Grid) -> Self {
        let components_and_neighbours = find_connected_components(&grid);

        let map_cell_to_component: HashMap<Position, ConnectedComponent> = components_and_neighbours
        .iter()
        .flat_map(|(component, _)| {
            component.cells
            .iter()
            .map(|cell| (cell.clone(), component.clone()))
            .collect::<Vec<(Position, ConnectedComponent)>>()
        })
        .collect();

        let mut neighbours:  HashMap<ConnectedComponent, HashSet<ConnectedComponent>> = HashMap::new();

        for (component, neighbour_cells) in components_and_neighbours {
            let mut neighbour_components: HashSet<ConnectedComponent> = HashSet::new();

            for cell in neighbour_cells {
                let position = Position{column: cell % grid.number_of_columns, row: cell / grid.number_of_columns};
                neighbour_components.insert(map_cell_to_component[&position].clone());
            }

            neighbours.insert(component, neighbour_components);
        }

        Self { neighbours }
    }

    fn to_grid(&self, rows: usize, columns: usize) -> Grid {
        let mut cells = vec![Colour::Red; rows * columns];

        for component in self.neighbours.keys() {
            for position in &component.cells {
                cells[position.column + position.row * columns] = component.colour;
            }
        }

        Grid { number_of_rows: rows, number_of_columns: columns, cells, }
    }
}

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

    #[test]
    fn should_return_same_grid_as_input() {
        let size = 4;
        let grid = Grid::generate(size);

        let graph = Graph::create(&grid);
        let reconstituted_grid = graph.to_grid(size, size);

        assert_eq!(reconstituted_grid, grid);
    }

    #[test]
    fn should_have_one_component() {
        let grid = Grid{
            number_of_columns: 2,
             number_of_rows: 2,
             cells: vec![Colour::Red; 4]
        };

        let graph = Graph::create(&grid);

        assert_eq!(graph.neighbours.len(), 1);
    }

    #[test]
    fn should_have_four_component() {
        let grid = Grid{
            number_of_columns: 2,
             number_of_rows: 2,
             cells: vec![Colour::Red, Colour::Yellow, Colour::Yellow, Colour::Red]
        };

        let graph = Graph::create(&grid);

        assert_eq!(graph.neighbours.len(), 4);
    }

    #[test]
    fn should_have_two_component() {
        let grid = Grid{
            number_of_columns: 2,
             number_of_rows: 2,
             cells: vec![Colour::Red, Colour::Red, Colour::Yellow, Colour::Yellow]
        };

        let graph = Graph::create(&grid);

        assert_eq!(graph.neighbours.len(), 2);
    }
}
<|endoftext|>