<commit_before>use libc::c_void;

use parse_code;
use std::io::{self, Write};
use zstd_sys;

struct EncoderContext {
    s: *mut zstd_sys::ZSTD_CStream,
}

impl Default for EncoderContext {
    fn default() -> Self {
        EncoderContext { s: unsafe { zstd_sys::ZSTD_createCStream() } }
    }
}

impl Drop for EncoderContext {
    fn drop(&mut self) {
        let code = unsafe { zstd_sys::ZSTD_freeCStream(self.s) };
        parse_code(code).unwrap();
    }
}

/// An encoder that compress and forward data to another writer.
///
/// This allows to compress a stream of data
/// (good for files or heavy network stream).
///
/// Don't forget to call `finish()` before dropping it!
///
/// Note: The zstd library has its own internal input buffer (~128kb).
pub struct Encoder<W: Write> {
    // output writer (compressed data)
    writer: W,
    // output buffer
    buffer: Vec<u8>,
    // offset in the output buffer
    offset: usize,

    // compression context
    context: EncoderContext,
}

/// A wrapper around an `Encoder<W>` that finishes the stream on drop.
pub struct AutoFinishEncoder<W: Write> {
    // We wrap this in an option to take it during drop.
    encoder: Option<Encoder<W>>,
    // TODO: make this a FnOnce once it works in a Box
    on_finish: Option<Box<FnMut(io::Result<W>)>>,
}

impl<W: Write> AutoFinishEncoder<W> {
    fn new<F: 'static + FnMut(io::Result<W>)>(encoder: Encoder<W>,
                                              on_finish: F)
                                              -> Self {
        AutoFinishEncoder {
            encoder: Some(encoder),
            on_finish: Some(Box::new(on_finish)),
        }
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.encoder
            .as_ref()
            .unwrap()
            .get_ref()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.encoder
            .as_mut()
            .unwrap()
            .get_mut()
    }
}

impl<W: Write> Drop for AutoFinishEncoder<W> {
    fn drop(&mut self) {
        let result = self.encoder
            .take()
            .unwrap()
            .finish();
        if let Some(mut on_finish) = self.on_finish.take() {
            on_finish(result);
        }
    }
}

impl<W: Write> Write for AutoFinishEncoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.encoder
            .as_mut()
            .unwrap()
            .write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.encoder
            .as_mut()
            .unwrap()
            .flush()
    }
}

impl<W: Write> Encoder<W> {
    /// Creates a new encoder.
    ///
    /// `level`: compression level (1-21)
    pub fn new(writer: W, level: i32) -> io::Result<Self> {
        Self::with_dictionary(writer, level, &[])
    }

    /// Creates a new encoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_dictionary(writer: W, level: i32, dictionary: &[u8])
                           -> io::Result<Self> {
        let context = EncoderContext::default();

        // Initialize the stream with an existing dictionary
        parse_code(unsafe {
            zstd_sys::ZSTD_initCStream_usingDict(context.s,
                                           dictionary.as_ptr() as *const c_void,
                                           dictionary.len(),
                                           level)
        })?;

        Encoder::with_context(writer, context)
    }

    /// Returns a wrapper around `self` that will finish the stream on drop.
    ///
    /// # Panic
    ///
    /// Panics on drop if an error happens when finishing the stream.
    pub fn auto_finish(self) -> AutoFinishEncoder<W> {
        self.on_finish(|result| { result.unwrap(); })
    }

    /// Returns an encoder that will finish the stream on drop.
    ///
    /// Calls the given callback with the result from `finish()`.
    pub fn on_finish<F: 'static + FnMut(io::Result<W>)>
        (self, f: F)
         -> AutoFinishEncoder<W> {
        AutoFinishEncoder::new(self, f)
    }

    fn with_context(writer: W, context: EncoderContext) -> io::Result<Self> {
        // This is the output buffer size,
        // for compressed data we get from zstd.
        let buffer_size = unsafe { zstd_sys::ZSTD_CStreamOutSize() };

        Ok(Encoder {
               writer: writer,
               buffer: Vec::with_capacity(buffer_size),
               offset: 0,
               context: context,
           })
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        &self.writer
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.writer
    }

    /// Finishes the stream. You *need* to call this after writing your stuff.
    ///
    /// This returns the inner writer in case you need it.
    ///
    /// **Note**: If you don't want (or can't) call `finish()` manually after writing your data,
    /// consider using `auto_finish()` to get an `AutoFinishEncoder`.
    pub fn finish(mut self) -> io::Result<W> {

        // First, closes the stream.

        let mut buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };
        let remaining = parse_code(unsafe {
            zstd_sys::ZSTD_endStream(self.context.s,
                               &mut buffer as *mut zstd_sys::ZSTD_outBuffer)
        })?;
        unsafe {
            self.buffer.set_len(buffer.pos);
        }
        if remaining != 0 {
            // Need to flush?
            panic!("Need to flush, but I'm lazy.");
        }

        // Write the end out
        self.writer.write_all(&self.buffer)?;

        // Return the writer, because why not
        Ok(self.writer)
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        unsafe { zstd_sys::ZSTD_CStreamInSize() }
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {

        if self.offset < self.buffer.len() {
            // If we still had some things to write, do it first.
            self.offset += self.writer.write(&self.buffer[self.offset..])?;
            // Maybe next time!
            return Err(io::Error::new(io::ErrorKind::Interrupted,
                                      "Internal buffer full"));
        }

        // If we get to here, `self.buffer` can safely be discarded.

        let mut in_buffer = zstd_sys::ZSTD_inBuffer {
            src: buf.as_ptr() as *const c_void,
            size: buf.len(),
            pos: 0,
        };

        let mut out_buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };

        unsafe {
            // Time to fill our input buffer
            let code =
                zstd_sys::ZSTD_compressStream(self.context.s,
                                              &mut out_buffer as
                                              *mut zstd_sys::ZSTD_outBuffer,
                                              &mut in_buffer as
                                              *mut zstd_sys::ZSTD_inBuffer);
            // Note: this may very well be empty,
            // if it doesn't exceed zstd's own buffer
            self.buffer.set_len(out_buffer.pos);

            // Do we care about the hint?
            let _ = parse_code(code)?;
        }

        // This is the first time he sees this buffer.
        // Remember his delicate touch.
        self.offset = self.writer.write(&self.buffer)?;

        Ok(in_buffer.pos)
    }

    fn flush(&mut self) -> io::Result<()> {
        let mut buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };
        unsafe {
            let code =
                zstd_sys::ZSTD_flushStream(self.context.s,
                                           &mut buffer as
                                           *mut zstd_sys::ZSTD_outBuffer);
            self.buffer.set_len(buffer.pos);
            let _ = parse_code(code)?;
        }

        self.writer.write_all(&self.buffer)?;
        Ok(())
    }
}
<commit_msg>encoder: instead of using write_all, keep track of partial writes<commit_after>use libc::c_void;

use parse_code;
use std::io::{self, Write};
use zstd_sys;

struct EncoderContext {
    s: *mut zstd_sys::ZSTD_CStream,
}

impl Default for EncoderContext {
    fn default() -> Self {
        EncoderContext { s: unsafe { zstd_sys::ZSTD_createCStream() } }
    }
}

impl Drop for EncoderContext {
    fn drop(&mut self) {
        let code = unsafe { zstd_sys::ZSTD_freeCStream(self.s) };
        parse_code(code).unwrap();
    }
}

/// An encoder that compress and forward data to another writer.
///
/// This allows to compress a stream of data
/// (good for files or heavy network stream).
///
/// Don't forget to call `finish()` before dropping it!
///
/// Note: The zstd library has its own internal input buffer (~128kb).
pub struct Encoder<W: Write> {
    // output writer (compressed data)
    writer: W,
    // output buffer
    buffer: Vec<u8>,
    // offset in the output buffer
    offset: usize,

    // compression context
    context: EncoderContext,
}

/// A wrapper around an `Encoder<W>` that finishes the stream on drop.
pub struct AutoFinishEncoder<W: Write> {
    // We wrap this in an option to take it during drop.
    encoder: Option<Encoder<W>>,
    // TODO: make this a FnOnce once it works in a Box
    on_finish: Option<Box<FnMut(io::Result<W>)>>,
}

impl<W: Write> AutoFinishEncoder<W> {
    fn new<F: 'static + FnMut(io::Result<W>)>(encoder: Encoder<W>,
                                              on_finish: F)
                                              -> Self {
        AutoFinishEncoder {
            encoder: Some(encoder),
            on_finish: Some(Box::new(on_finish)),
        }
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.encoder
            .as_ref()
            .unwrap()
            .get_ref()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.encoder
            .as_mut()
            .unwrap()
            .get_mut()
    }
}

impl<W: Write> Drop for AutoFinishEncoder<W> {
    fn drop(&mut self) {
        let result = self.encoder
            .take()
            .unwrap()
            .finish();
        if let Some(mut on_finish) = self.on_finish.take() {
            on_finish(result);
        }
    }
}

impl<W: Write> Write for AutoFinishEncoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.encoder
            .as_mut()
            .unwrap()
            .write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.encoder
            .as_mut()
            .unwrap()
            .flush()
    }
}

impl<W: Write> Encoder<W> {
    /// Creates a new encoder.
    ///
    /// `level`: compression level (1-21)
    pub fn new(writer: W, level: i32) -> io::Result<Self> {
        Self::with_dictionary(writer, level, &[])
    }

    /// Creates a new encoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_dictionary(writer: W, level: i32, dictionary: &[u8])
                           -> io::Result<Self> {
        let context = EncoderContext::default();

        // Initialize the stream with an existing dictionary
        parse_code(unsafe {
            zstd_sys::ZSTD_initCStream_usingDict(context.s,
                                           dictionary.as_ptr() as *const c_void,
                                           dictionary.len(),
                                           level)
        })?;

        Encoder::with_context(writer, context)
    }

    /// Returns a wrapper around `self` that will finish the stream on drop.
    ///
    /// # Panic
    ///
    /// Panics on drop if an error happens when finishing the stream.
    pub fn auto_finish(self) -> AutoFinishEncoder<W> {
        self.on_finish(|result| { result.unwrap(); })
    }

    /// Returns an encoder that will finish the stream on drop.
    ///
    /// Calls the given callback with the result from `finish()`.
    pub fn on_finish<F: 'static + FnMut(io::Result<W>)>
        (self, f: F)
         -> AutoFinishEncoder<W> {
        AutoFinishEncoder::new(self, f)
    }

    fn with_context(writer: W, context: EncoderContext) -> io::Result<Self> {
        // This is the output buffer size,
        // for compressed data we get from zstd.
        let buffer_size = unsafe { zstd_sys::ZSTD_CStreamOutSize() };

        Ok(Encoder {
               writer: writer,
               buffer: Vec::with_capacity(buffer_size),
               offset: 0,
               context: context,
           })
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        &self.writer
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.writer
    }

    /// Finishes the stream. You *need* to call this after writing your stuff.
    ///
    /// This returns the inner writer in case you need it.
    ///
    /// **Note**: If you don't want (or can't) call `finish()` manually after writing your data,
    /// consider using `auto_finish()` to get an `AutoFinishEncoder`.
    pub fn finish(mut self) -> io::Result<W> {

        // First, closes the stream.

        let mut buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };
        let remaining = parse_code(unsafe {
            zstd_sys::ZSTD_endStream(self.context.s,
                               &mut buffer as *mut zstd_sys::ZSTD_outBuffer)
        })?;
        unsafe {
            self.buffer.set_len(buffer.pos);
        }
        if remaining != 0 {
            // Need to flush?
            panic!("Need to flush, but I'm lazy.");
        }
        self.offset = 0;

        self.write_from_offset()?;

        // Return the writer, because why not
        Ok(self.writer)
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        unsafe { zstd_sys::ZSTD_CStreamInSize() }
    }

    /// write_all, except keep track of partial writes for non-blocking IO.
    fn write_from_offset(&mut self) -> io::Result<()> {
        while self.offset < self.buffer.len() {
            match self.writer.write(&self.buffer[self.offset..]) {
                Ok(n) => self.offset += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {},
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {

        if self.offset < self.buffer.len() {
            // If we still had some things to write, do it first.
            self.offset += self.writer.write(&self.buffer[self.offset..])?;
            // Maybe next time!
            return Err(io::Error::new(io::ErrorKind::Interrupted,
                                      "Internal buffer full"));
        }

        // If we get to here, `self.buffer` can safely be discarded.

        let mut in_buffer = zstd_sys::ZSTD_inBuffer {
            src: buf.as_ptr() as *const c_void,
            size: buf.len(),
            pos: 0,
        };

        let mut out_buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };

        unsafe {
            // Time to fill our input buffer
            let code =
                zstd_sys::ZSTD_compressStream(self.context.s,
                                              &mut out_buffer as
                                              *mut zstd_sys::ZSTD_outBuffer,
                                              &mut in_buffer as
                                              *mut zstd_sys::ZSTD_inBuffer);
            // Note: this may very well be empty,
            // if it doesn't exceed zstd's own buffer
            self.buffer.set_len(out_buffer.pos);

            // Do we care about the hint?
            let _ = parse_code(code)?;
        }

        // This is the first time he sees this buffer.
        // Remember his delicate touch.
        self.offset = self.writer.write(&self.buffer)?;

        Ok(in_buffer.pos)
    }

    fn flush(&mut self) -> io::Result<()> {
        let mut buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };
        unsafe {
            let code =
                zstd_sys::ZSTD_flushStream(self.context.s,
                                           &mut buffer as
                                           *mut zstd_sys::ZSTD_outBuffer);
            self.buffer.set_len(buffer.pos);
            let _ = parse_code(code)?;
        }
        self.offset = 0;

        self.write_from_offset()?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, ReadBytesExt};
use rustc_serialize;
use std;
use super::error::Error;

pub struct Decoder {
    input: std::vec::IntoIter<u8>,
    bytes_read: usize,
}

impl Decoder {
    pub fn new(data: Vec<u8>) -> Decoder {
        Decoder {
            input: data.into_iter(),
            bytes_read: 0,
        }
    }

    fn pop_length(&mut self) -> Result<u32, Error> {
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_u32::<LittleEndian>()))
    }

    fn read_bytes(&mut self, count: u32) -> Result<Vec<u8>, Error> {
        let mut buffer = vec![];
        for _ in 0..count {
            buffer.push(try!(self.input.next().ok_or(Error::Truncated)));
            self.bytes_read += 1;
        }
        Ok(buffer)
    }
}

macro_rules! match_length {
    ($s:expr, $x:expr) => (if $x != try!($s.pop_length()) {return Err(Error::Mismatch)});
}

impl rustc_serialize::Decoder for Decoder {
    type Error = Error;

    fn read_nil(&mut self) -> Result<(), Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_usize(&mut self) -> Result<usize, Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_u64(&mut self) -> Result<u64, Self::Error> {
        match_length!(self, 8);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(8)));
        Ok(try!(reader.read_u64::<LittleEndian>()))
    }

    fn read_u32(&mut self) -> Result<u32, Self::Error> {
        match_length!(self, 4);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_u32::<LittleEndian>()))
    }

    fn read_u16(&mut self) -> Result<u16, Self::Error> {
        match_length!(self, 2);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(2)));
        Ok(try!(reader.read_u16::<LittleEndian>()))
    }

    fn read_u8(&mut self) -> Result<u8, Self::Error> {
        match_length!(self, 1);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(1)));
        Ok(try!(reader.read_u8()))
    }

    fn read_isize(&mut self) -> Result<isize, Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_i64(&mut self) -> Result<i64, Self::Error> {
        match_length!(self, 8);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(8)));
        Ok(try!(reader.read_i64::<LittleEndian>()))
    }

    fn read_i32(&mut self) -> Result<i32, Self::Error> {
        match_length!(self, 4);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_i32::<LittleEndian>()))
    }

    fn read_i16(&mut self) -> Result<i16, Self::Error> {
        match_length!(self, 2);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(2)));
        Ok(try!(reader.read_i16::<LittleEndian>()))
    }

    fn read_i8(&mut self) -> Result<i8, Self::Error> {
        match_length!(self, 1);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(1)));
        Ok(try!(reader.read_i8()))
    }

    fn read_bool(&mut self) -> Result<bool, Self::Error> {
        match_length!(self, 1);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(1)));
        Ok(try!(reader.read_u8()) != 0)
    }

    fn read_f64(&mut self) -> Result<f64, Self::Error> {
        match_length!(self, 8);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(8)));
        Ok(try!(reader.read_f64::<LittleEndian>()))
    }

    fn read_f32(&mut self) -> Result<f32, Self::Error> {
        match_length!(self, 4);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_f32::<LittleEndian>()))
    }

    fn read_char(&mut self) -> Result<char, Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_str(&mut self) -> Result<String, Self::Error> {
        let length = try!(self.pop_length());
        String::from_utf8(try!(self.read_bytes(length))).or(Err(Error::UnsupportedData))
    }

    fn read_enum<T, F>(&mut self, _: &str, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_variant<T, F>(&mut self, _: &[&str], _: F) -> Result<T, Self::Error>
        where F: FnMut(&mut Self, usize) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_variant_arg<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_struct_variant<T, F>(&mut self, _: &[&str], _: F) -> Result<T, Self::Error>
        where F: FnMut(&mut Self, usize) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_struct_variant_field<T, F>(&mut self,
                                            _: &str,
                                            _: usize,
                                            _: F)
                                            -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_struct<T, F>(&mut self, _: &str, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        let length = try!(self.pop_length());
        let data_start = self.bytes_read;
        let retval = f(self);
        if self.bytes_read == data_start + length as usize {
            retval
        } else {
            Err(Error::Truncated)
        }
    }

    fn read_struct_field<T, F>(&mut self, _: &str, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        f(self)
    }

    fn read_tuple<T, F>(&mut self, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        let length = try!(self.pop_length());
        let data_start = self.bytes_read;
        let retval = f(self);
        if self.bytes_read == data_start + length as usize {
            retval
        } else {
            Err(Error::Truncated)
        }
    }

    fn read_tuple_arg<T, F>(&mut self, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        f(self)
    }

    fn read_tuple_struct<T, F>(&mut self, _: &str, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_tuple_struct_arg<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_option<T, F>(&mut self, _: F) -> Result<T, Self::Error>
        where F: FnMut(&mut Self, bool) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_seq<T, F>(&mut self, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self, usize) -> Result<T, Self::Error>
    {
        let length = try!(self.pop_length());
        let data_start = self.bytes_read;
        let count = try!(self.pop_length());
        let retval = f(self, count as usize);
        if self.bytes_read == data_start + length as usize {
            retval
        } else {
            Err(Error::Truncated)
        }
    }

    fn read_seq_elt<T, F>(&mut self, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        f(self)
    }

    fn read_map<T, F>(&mut self, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self, usize) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_map_elt_key<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_map_elt_val<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn error(&mut self, err: &str) -> Self::Error {
        Error::Other(err.to_owned())
    }
}
<commit_msg>Allow public access to pop_length in TCPROS<commit_after>use byteorder::{LittleEndian, ReadBytesExt};
use rustc_serialize;
use std;
use super::error::Error;

pub struct Decoder {
    input: std::vec::IntoIter<u8>,
    bytes_read: usize,
}

impl Decoder {
    pub fn new(data: Vec<u8>) -> Decoder {
        Decoder {
            input: data.into_iter(),
            bytes_read: 0,
        }
    }

    pub fn pop_length(&mut self) -> Result<u32, Error> {
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_u32::<LittleEndian>()))
    }

    fn read_bytes(&mut self, count: u32) -> Result<Vec<u8>, Error> {
        let mut buffer = vec![];
        for _ in 0..count {
            buffer.push(try!(self.input.next().ok_or(Error::Truncated)));
            self.bytes_read += 1;
        }
        Ok(buffer)
    }
}

macro_rules! match_length {
    ($s:expr, $x:expr) => (if $x != try!($s.pop_length()) {return Err(Error::Mismatch)});
}

impl rustc_serialize::Decoder for Decoder {
    type Error = Error;

    fn read_nil(&mut self) -> Result<(), Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_usize(&mut self) -> Result<usize, Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_u64(&mut self) -> Result<u64, Self::Error> {
        match_length!(self, 8);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(8)));
        Ok(try!(reader.read_u64::<LittleEndian>()))
    }

    fn read_u32(&mut self) -> Result<u32, Self::Error> {
        match_length!(self, 4);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_u32::<LittleEndian>()))
    }

    fn read_u16(&mut self) -> Result<u16, Self::Error> {
        match_length!(self, 2);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(2)));
        Ok(try!(reader.read_u16::<LittleEndian>()))
    }

    fn read_u8(&mut self) -> Result<u8, Self::Error> {
        match_length!(self, 1);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(1)));
        Ok(try!(reader.read_u8()))
    }

    fn read_isize(&mut self) -> Result<isize, Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_i64(&mut self) -> Result<i64, Self::Error> {
        match_length!(self, 8);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(8)));
        Ok(try!(reader.read_i64::<LittleEndian>()))
    }

    fn read_i32(&mut self) -> Result<i32, Self::Error> {
        match_length!(self, 4);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_i32::<LittleEndian>()))
    }

    fn read_i16(&mut self) -> Result<i16, Self::Error> {
        match_length!(self, 2);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(2)));
        Ok(try!(reader.read_i16::<LittleEndian>()))
    }

    fn read_i8(&mut self) -> Result<i8, Self::Error> {
        match_length!(self, 1);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(1)));
        Ok(try!(reader.read_i8()))
    }

    fn read_bool(&mut self) -> Result<bool, Self::Error> {
        match_length!(self, 1);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(1)));
        Ok(try!(reader.read_u8()) != 0)
    }

    fn read_f64(&mut self) -> Result<f64, Self::Error> {
        match_length!(self, 8);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(8)));
        Ok(try!(reader.read_f64::<LittleEndian>()))
    }

    fn read_f32(&mut self) -> Result<f32, Self::Error> {
        match_length!(self, 4);
        let mut reader = std::io::Cursor::new(try!(self.read_bytes(4)));
        Ok(try!(reader.read_f32::<LittleEndian>()))
    }

    fn read_char(&mut self) -> Result<char, Self::Error> {
        Err(Error::UnsupportedData)
    }

    fn read_str(&mut self) -> Result<String, Self::Error> {
        let length = try!(self.pop_length());
        String::from_utf8(try!(self.read_bytes(length))).or(Err(Error::UnsupportedData))
    }

    fn read_enum<T, F>(&mut self, _: &str, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_variant<T, F>(&mut self, _: &[&str], _: F) -> Result<T, Self::Error>
        where F: FnMut(&mut Self, usize) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_variant_arg<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_struct_variant<T, F>(&mut self, _: &[&str], _: F) -> Result<T, Self::Error>
        where F: FnMut(&mut Self, usize) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_enum_struct_variant_field<T, F>(&mut self,
                                            _: &str,
                                            _: usize,
                                            _: F)
                                            -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_struct<T, F>(&mut self, _: &str, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        let length = try!(self.pop_length());
        let data_start = self.bytes_read;
        let retval = f(self);
        if self.bytes_read == data_start + length as usize {
            retval
        } else {
            Err(Error::Truncated)
        }
    }

    fn read_struct_field<T, F>(&mut self, _: &str, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        f(self)
    }

    fn read_tuple<T, F>(&mut self, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        let length = try!(self.pop_length());
        let data_start = self.bytes_read;
        let retval = f(self);
        if self.bytes_read == data_start + length as usize {
            retval
        } else {
            Err(Error::Truncated)
        }
    }

    fn read_tuple_arg<T, F>(&mut self, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        f(self)
    }

    fn read_tuple_struct<T, F>(&mut self, _: &str, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_tuple_struct_arg<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_option<T, F>(&mut self, _: F) -> Result<T, Self::Error>
        where F: FnMut(&mut Self, bool) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_seq<T, F>(&mut self, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self, usize) -> Result<T, Self::Error>
    {
        let length = try!(self.pop_length());
        let data_start = self.bytes_read;
        let count = try!(self.pop_length());
        let retval = f(self, count as usize);
        if self.bytes_read == data_start + length as usize {
            retval
        } else {
            Err(Error::Truncated)
        }
    }

    fn read_seq_elt<T, F>(&mut self, _: usize, f: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        f(self)
    }

    fn read_map<T, F>(&mut self, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self, usize) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_map_elt_key<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn read_map_elt_val<T, F>(&mut self, _: usize, _: F) -> Result<T, Self::Error>
        where F: FnOnce(&mut Self) -> Result<T, Self::Error>
    {
        Err(Error::UnsupportedData)
    }

    fn error(&mut self, err: &str) -> Self::Error {
        Error::Other(err.to_owned())
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]
extern crate hyper;
extern crate futures;
extern crate tokio_core;
extern crate pretty_env_logger;

use std::io::{self, Read, Write};
use std::net::TcpListener;
use std::thread;
use std::time::Duration;

use hyper::client::{Client, Request, HttpConnector};
use hyper::{Method, StatusCode};

use futures::Future;
use futures::sync::oneshot;

use tokio_core::reactor::{Core, Handle};

fn client(handle: &Handle) -> Client<HttpConnector> {
    Client::new(handle)
}

fn s(buf: &[u8]) -> &str {
    ::std::str::from_utf8(buf).unwrap()
}

macro_rules! test {
    (
        name: $name:ident,
        server:
            expected: $server_expected:expr,
            reply: $server_reply:expr,
        client:
            request:
                method: $client_method:ident,
                url: $client_url:expr,
                headers: [ $($request_headers:expr,)* ],
                body: $request_body:expr,
                proxy: $request_proxy:expr,

            response:
                status: $client_status:ident,
                headers: [ $($response_headers:expr,)* ],
                body: $response_body:expr,
    ) => (
        #[test]
        fn $name() {
            #![allow(unused)]
            use hyper::header::*;
            let _ = pretty_env_logger::init();
            let server = TcpListener::bind("127.0.0.1:0").unwrap();
            let addr = server.local_addr().unwrap();
            let mut core = Core::new().unwrap();
            let client = client(&core.handle());
            let mut req = Request::new(Method::$client_method, format!($client_url, addr=addr).parse().unwrap());
            $(
                req.headers_mut().set($request_headers);
            )*

            if let Some(body) = $request_body {
                let body: &'static str = body;
                req.set_body(body);
            }
            req.set_proxy($request_proxy);

            let res = client.request(req);

            let (tx, rx) = oneshot::channel();

            thread::spawn(move || {
                let mut inc = server.accept().unwrap().0;
                inc.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
                inc.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
                let expected = format!($server_expected, addr=addr);
                let mut buf = [0; 4096];
                let mut n = 0;
                while n < buf.len() && n < expected.len() {
                    n += inc.read(&mut buf[n..]).unwrap();
                }
                assert_eq!(s(&buf[..n]), expected, "expected is invalid");

                inc.write_all($server_reply.as_ref()).unwrap();
                tx.complete(());
            });

            let rx = rx.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));

            let work = res.join(rx).map(|r| r.0);

            let res = core.run(work).unwrap();
            assert_eq!(res.status(), &StatusCode::$client_status, "status is invalid");
            $(
                assert_eq!(res.headers().get(), Some(&$response_headers), "headers are invalid");
            )*
        }
    );
}

static REPLY_OK: &'static str = "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n";

test! {
    name: client_get,

    server:
        expected: "GET / HTTP/1.1\r\nHost: {addr}\r\n\r\n",
        reply: REPLY_OK,

    client:
        request:
            method: Get,
            url: "http://{addr}/",
            headers: [],
            body: None,
            proxy: false,
        response:
            status: Ok,
            headers: [
                ContentLength(0),
            ],
            body: None,
}

test! {
    name: client_get_query,

    server:
        expected: "GET /foo?key=val HTTP/1.1\r\nHost: {addr}\r\n\r\n",
        reply: REPLY_OK,

    client:
        request:
            method: Get,
            url: "http://{addr}/foo?key=val#dont_send_me",
            headers: [],
            body: None,
            proxy: false,
        response:
            status: Ok,
            headers: [
                ContentLength(0),
            ],
            body: None,
}

test! {
    name: client_post_sized,

    server:
        expected: "\
            POST /length HTTP/1.1\r\n\
            Host: {addr}\r\n\
            Content-Length: 7\r\n\
            \r\n\
            foo bar\
            ",
        reply: REPLY_OK,

    client:
        request:
            method: Post,
            url: "http://{addr}/length",
            headers: [
                ContentLength(7),
            ],
            body: Some("foo bar"),
            proxy: false,
        response:
            status: Ok,
            headers: [],
            body: None,
}

test! {
    name: client_post_chunked,

    server:
        expected: "\
            POST /chunks HTTP/1.1\r\n\
            Host: {addr}\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            B\r\n\
            foo bar baz\r\n\
            0\r\n\r\n\
            ",
        reply: REPLY_OK,

    client:
        request:
            method: Post,
            url: "http://{addr}/chunks",
            headers: [
                TransferEncoding::chunked(),
            ],
            body: Some("foo bar baz"),
            proxy: false,
        response:
            status: Ok,
            headers: [],
            body: None,
}

test! {
    name: client_http_proxy,

    server:
        expected: "\
            GET http://{addr}/proxy HTTP/1.1\r\n\
            Host: {addr}\r\n\
            \r\n\
            ",
        reply: REPLY_OK,

    client:
        request:
            method: Get,
            url: "http://{addr}/proxy",
            headers: [],
            body: None,
            proxy: true,
        response:
            status: Ok,
            headers: [],
            body: None,
}

#[test]
fn client_keep_alive() {
    let server = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = server.local_addr().unwrap();
    let mut core = Core::new().unwrap();
    let client = client(&core.handle());

    let (tx1, rx1) = oneshot::channel();
    let (tx2, rx2) = oneshot::channel();
    thread::spawn(move || {
        let mut sock = server.accept().unwrap().0;
        sock.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
        sock.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
        let mut buf = [0; 4096];
        sock.read(&mut buf).expect("read 1");
        sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n").expect("write 1");
        tx1.complete(());

        sock.read(&mut buf).expect("read 2");
        let second_get = b"GET /b HTTP/1.1\r\n";
        assert_eq!(&buf[..second_get.len()], second_get);
        sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n").expect("write 2");
        tx2.complete(());
    });

    let rx = rx1.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/a", addr).parse().unwrap());
    core.run(res.join(rx).map(|r| r.0)).unwrap();

    let rx = rx2.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/b", addr).parse().unwrap());
    core.run(res.join(rx).map(|r| r.0)).unwrap();
}

/* TODO: re-enable once retry works, its currently a flaky test
#[test]
fn client_pooled_socket_disconnected() {
    let _ = pretty_env_logger::init();
    let server = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = server.local_addr().unwrap();
    let mut core = Core::new().unwrap();
    let client = client(&core.handle());

    let (tx1, rx1) = oneshot::channel();
    let (tx2, rx2) = oneshot::channel();
    thread::spawn(move || {
        let mut sock = server.accept().unwrap().0;
        sock.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
        sock.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
        let mut buf = [0; 4096];
        sock.read(&mut buf).expect("read 1");
        let remote_addr = sock.peer_addr().unwrap().to_string();
        let out = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}", remote_addr.len(), remote_addr);
        sock.write_all(out.as_bytes()).expect("write 1");
        drop(sock);
        tx1.complete(());

        let mut sock = server.accept().unwrap().0;
        sock.read(&mut buf).expect("read 2");
        let second_get = b"GET /b HTTP/1.1\r\n";
        assert_eq!(&buf[..second_get.len()], second_get);
        let remote_addr = sock.peer_addr().unwrap().to_string();
        let out = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}", remote_addr.len(), remote_addr);
        sock.write_all(out.as_bytes()).expect("write 2");
        tx2.complete(());
    });

    // spin shortly so we receive the hangup on the client socket
    let sleep = Timeout::new(Duration::from_millis(500), &core.handle()).unwrap();
    core.run(sleep).unwrap();

    let rx = rx1.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/a", addr).parse().unwrap())
        .and_then(|res| {
            res.body()
                .map(|chunk| chunk.to_vec())
                .collect()
                .map(|vec| vec.concat())
                .map(|vec| String::from_utf8(vec).unwrap())
        });
    let addr1 = core.run(res.join(rx).map(|r| r.0)).unwrap();

    let rx = rx2.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/b", addr).parse().unwrap())
        .and_then(|res| {
            res.body()
                .map(|chunk| chunk.to_vec())
                .collect()
                .map(|vec| vec.concat())
                .map(|vec| String::from_utf8(vec).unwrap())
        });
    let addr2 = core.run(res.join(rx).map(|r| r.0)).unwrap();

    assert_ne!(addr1, addr2);
}
*/
<commit_msg>test(client): replace deprecated Sender.completer with send<commit_after>#![deny(warnings)]
extern crate hyper;
extern crate futures;
extern crate tokio_core;
extern crate pretty_env_logger;

use std::io::{self, Read, Write};
use std::net::TcpListener;
use std::thread;
use std::time::Duration;

use hyper::client::{Client, Request, HttpConnector};
use hyper::{Method, StatusCode};

use futures::Future;
use futures::sync::oneshot;

use tokio_core::reactor::{Core, Handle};

fn client(handle: &Handle) -> Client<HttpConnector> {
    Client::new(handle)
}

fn s(buf: &[u8]) -> &str {
    ::std::str::from_utf8(buf).unwrap()
}

macro_rules! test {
    (
        name: $name:ident,
        server:
            expected: $server_expected:expr,
            reply: $server_reply:expr,
        client:
            request:
                method: $client_method:ident,
                url: $client_url:expr,
                headers: [ $($request_headers:expr,)* ],
                body: $request_body:expr,
                proxy: $request_proxy:expr,

            response:
                status: $client_status:ident,
                headers: [ $($response_headers:expr,)* ],
                body: $response_body:expr,
    ) => (
        #[test]
        fn $name() {
            #![allow(unused)]
            use hyper::header::*;
            let _ = pretty_env_logger::init();
            let server = TcpListener::bind("127.0.0.1:0").unwrap();
            let addr = server.local_addr().unwrap();
            let mut core = Core::new().unwrap();
            let client = client(&core.handle());
            let mut req = Request::new(Method::$client_method, format!($client_url, addr=addr).parse().unwrap());
            $(
                req.headers_mut().set($request_headers);
            )*

            if let Some(body) = $request_body {
                let body: &'static str = body;
                req.set_body(body);
            }
            req.set_proxy($request_proxy);

            let res = client.request(req);

            let (tx, rx) = oneshot::channel();

            thread::spawn(move || {
                let mut inc = server.accept().unwrap().0;
                inc.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
                inc.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
                let expected = format!($server_expected, addr=addr);
                let mut buf = [0; 4096];
                let mut n = 0;
                while n < buf.len() && n < expected.len() {
                    n += inc.read(&mut buf[n..]).unwrap();
                }
                assert_eq!(s(&buf[..n]), expected, "expected is invalid");

                inc.write_all($server_reply.as_ref()).unwrap();
                let _ = tx.send(());
            });

            let rx = rx.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));

            let work = res.join(rx).map(|r| r.0);

            let res = core.run(work).unwrap();
            assert_eq!(res.status(), &StatusCode::$client_status, "status is invalid");
            $(
                assert_eq!(res.headers().get(), Some(&$response_headers), "headers are invalid");
            )*
        }
    );
}

static REPLY_OK: &'static str = "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n";

test! {
    name: client_get,

    server:
        expected: "GET / HTTP/1.1\r\nHost: {addr}\r\n\r\n",
        reply: REPLY_OK,

    client:
        request:
            method: Get,
            url: "http://{addr}/",
            headers: [],
            body: None,
            proxy: false,
        response:
            status: Ok,
            headers: [
                ContentLength(0),
            ],
            body: None,
}

test! {
    name: client_get_query,

    server:
        expected: "GET /foo?key=val HTTP/1.1\r\nHost: {addr}\r\n\r\n",
        reply: REPLY_OK,

    client:
        request:
            method: Get,
            url: "http://{addr}/foo?key=val#dont_send_me",
            headers: [],
            body: None,
            proxy: false,
        response:
            status: Ok,
            headers: [
                ContentLength(0),
            ],
            body: None,
}

test! {
    name: client_post_sized,

    server:
        expected: "\
            POST /length HTTP/1.1\r\n\
            Host: {addr}\r\n\
            Content-Length: 7\r\n\
            \r\n\
            foo bar\
            ",
        reply: REPLY_OK,

    client:
        request:
            method: Post,
            url: "http://{addr}/length",
            headers: [
                ContentLength(7),
            ],
            body: Some("foo bar"),
            proxy: false,
        response:
            status: Ok,
            headers: [],
            body: None,
}

test! {
    name: client_post_chunked,

    server:
        expected: "\
            POST /chunks HTTP/1.1\r\n\
            Host: {addr}\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            B\r\n\
            foo bar baz\r\n\
            0\r\n\r\n\
            ",
        reply: REPLY_OK,

    client:
        request:
            method: Post,
            url: "http://{addr}/chunks",
            headers: [
                TransferEncoding::chunked(),
            ],
            body: Some("foo bar baz"),
            proxy: false,
        response:
            status: Ok,
            headers: [],
            body: None,
}

test! {
    name: client_http_proxy,

    server:
        expected: "\
            GET http://{addr}/proxy HTTP/1.1\r\n\
            Host: {addr}\r\n\
            \r\n\
            ",
        reply: REPLY_OK,

    client:
        request:
            method: Get,
            url: "http://{addr}/proxy",
            headers: [],
            body: None,
            proxy: true,
        response:
            status: Ok,
            headers: [],
            body: None,
}

#[test]
fn client_keep_alive() {
    let server = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = server.local_addr().unwrap();
    let mut core = Core::new().unwrap();
    let client = client(&core.handle());

    let (tx1, rx1) = oneshot::channel();
    let (tx2, rx2) = oneshot::channel();
    thread::spawn(move || {
        let mut sock = server.accept().unwrap().0;
        sock.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
        sock.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
        let mut buf = [0; 4096];
        sock.read(&mut buf).expect("read 1");
        sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n").expect("write 1");
        let _ = tx1.send(());

        sock.read(&mut buf).expect("read 2");
        let second_get = b"GET /b HTTP/1.1\r\n";
        assert_eq!(&buf[..second_get.len()], second_get);
        sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n").expect("write 2");
        let _ = tx2.send(());
    });

    let rx = rx1.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/a", addr).parse().unwrap());
    core.run(res.join(rx).map(|r| r.0)).unwrap();

    let rx = rx2.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/b", addr).parse().unwrap());
    core.run(res.join(rx).map(|r| r.0)).unwrap();
}

/* TODO: re-enable once retry works, its currently a flaky test
#[test]
fn client_pooled_socket_disconnected() {
    let _ = pretty_env_logger::init();
    let server = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = server.local_addr().unwrap();
    let mut core = Core::new().unwrap();
    let client = client(&core.handle());

    let (tx1, rx1) = oneshot::channel();
    let (tx2, rx2) = oneshot::channel();
    thread::spawn(move || {
        let mut sock = server.accept().unwrap().0;
        sock.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
        sock.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
        let mut buf = [0; 4096];
        sock.read(&mut buf).expect("read 1");
        let remote_addr = sock.peer_addr().unwrap().to_string();
        let out = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}", remote_addr.len(), remote_addr);
        sock.write_all(out.as_bytes()).expect("write 1");
        drop(sock);
        tx1.send(());

        let mut sock = server.accept().unwrap().0;
        sock.read(&mut buf).expect("read 2");
        let second_get = b"GET /b HTTP/1.1\r\n";
        assert_eq!(&buf[..second_get.len()], second_get);
        let remote_addr = sock.peer_addr().unwrap().to_string();
        let out = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}", remote_addr.len(), remote_addr);
        sock.write_all(out.as_bytes()).expect("write 2");
        tx2.send(());
    });

    // spin shortly so we receive the hangup on the client socket
    let sleep = Timeout::new(Duration::from_millis(500), &core.handle()).unwrap();
    core.run(sleep).unwrap();

    let rx = rx1.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/a", addr).parse().unwrap())
        .and_then(|res| {
            res.body()
                .map(|chunk| chunk.to_vec())
                .collect()
                .map(|vec| vec.concat())
                .map(|vec| String::from_utf8(vec).unwrap())
        });
    let addr1 = core.run(res.join(rx).map(|r| r.0)).unwrap();

    let rx = rx2.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked")));
    let res = client.get(format!("http://{}/b", addr).parse().unwrap())
        .and_then(|res| {
            res.body()
                .map(|chunk| chunk.to_vec())
                .collect()
                .map(|vec| vec.concat())
                .map(|vec| String::from_utf8(vec).unwrap())
        });
    let addr2 = core.run(res.join(rx).map(|r| r.0)).unwrap();

    assert_ne!(addr1, addr2);
}
*/
<|endoftext|>


<commit_before>include!(concat!(env!("OUT_DIR"), "/tests.rs"));

fn run_file (name: &str) {

}
<commit_msg>load files<commit_after>use std::fs::File;
use std::io::Read;
use std::path::Path;

include!(concat!(env!("OUT_DIR"), "/tests.rs"));

fn run_file (name: &str) {
    let mut file_path = Path::new("tests/fodder").join(name);
    file_path.set_extension("txt");

    let mut file = match File::open(&file_path) {
        Err(e) => {
            println!("Couldn't open {}: {}", file_path.display(), e);
            panic!(e);
        },
        Ok(f) => f
    };

    let mut s = String::new();
    file.read_to_string(&mut s).unwrap();
    println!("{}", s);
}
<|endoftext|>


<commit_before>extern crate lichen;

use std::io::Cursor;

use lichen::parse::{StreamParser,Block, Env};
use lichen::source::Src;
use lichen::var::Var;
use lichen::eval::{Eval,Evaluator};

struct Data;
impl Eval for Data {
    #[allow(unused_variables)]
    fn get (&self, path: Option<Vec<&str>>, lookup: &str) -> Option<Var> {
        None
    }

    #[allow(unused_variables)]
    fn set (&mut self, path: Option<Vec<&str>>, lookup: &str, var: Var) {}

    #[allow(unused_variables)]
    fn call (&mut self, var: Var, fun: &str, vars: &Vec<Var>) -> Option<Var> {
        None
    }
}

#[test]
fn stream_parser() {
    let src = "root\n
    emit \"hi\"\n
    next:now some_block\n
;\n
some_block\n
\n
emi";  //unfinished source

    let bytes = src.as_bytes();
    let c = Cursor::new(&bytes[..]);

    let mut s = StreamParser::new(c,None);
    let idx = s.parse();

    assert!(idx.is_some());
    if let Some(idx) = idx {
        let b = s.blocks.get(idx).expect("ERROR: No block");
        match b {
            &Block::Src(ref b) => {
                assert_eq!(b.name, "root".to_owned());
                assert_eq!(b.src.get(0), Some(&Src::Emit(vec![Var::String("hi".to_owned())])));
                assert!(b.src.len() < 3);
            },
            _ => { panic!("ERROR: Invalid block type") }
        }

    }

    let mut env = Env::empty();
    assert!(s.sink(&mut env).is_err());
    assert_eq!(s.blocks.len(), 1);

    let src = "t \"hi again\"\n
;"; //finish source to parse

    let bytes = src.as_bytes();
    let c = Cursor::new(&bytes[..]);

    s.stream = c; //swap in new 'stream'
    let idx = s.parse();

    assert!(idx.is_some());
    if let Some(idx) = idx {
        let b = s.blocks.get(idx).expect("ERROR: No block");
        match b {
            &Block::Src(ref b) => {
                assert_eq!(b.name, "some_block".to_owned());
                assert_eq!(b.src.get(0), Some(&Src::Emit(vec![Var::String("hi again".to_owned())])))
            },
            _ => { panic!("ERROR: Invalid block type") }
        }
    }

    // assert we can sink parsed blocks into an environment
    assert_eq!(s.blocks.len(), 2);
    assert!(s.sink(&mut env).is_ok());
    assert_eq!(s.blocks.len(), 0);

    let mut data = Data;
    let mut ev = Evaluator::new(&mut env, &mut data);
    let (vars,_) = ev.next().unwrap();
    assert_eq!(vars.get(0), Some(&Var::String("hi".to_owned())));
    let (vars,_) = ev.nth(1).expect("ERROR: Block failed to transition");
    assert_eq!(vars.get(0), Some(&Var::String("hi again".to_owned())));
}
<commit_msg>use empty<commit_after>extern crate lichen;

use std::io::Cursor;

use lichen::parse::{StreamParser,Block, Env};
use lichen::source::Src;
use lichen::var::Var;
use lichen::eval::{Evaluator,Empty};

#[test]
fn stream_parser() {
    let src = "root\n
    emit \"hi\"\n
    next:now some_block\n
;\n
some_block\n
\n
emi";  //unfinished source

    let bytes = src.as_bytes();
    let c = Cursor::new(&bytes[..]);

    let mut s = StreamParser::new(c,None);
    let idx = s.parse();

    assert!(idx.is_some());
    if let Some(idx) = idx {
        let b = s.blocks.get(idx).expect("ERROR: No block");
        match b {
            &Block::Src(ref b) => {
                assert_eq!(b.name, "root".to_owned());
                assert_eq!(b.src.get(0), Some(&Src::Emit(vec![Var::String("hi".to_owned())])));
                assert!(b.src.len() < 3);
            },
            _ => { panic!("ERROR: Invalid block type") }
        }

    }

    let mut env = Env::empty();
    assert!(s.sink(&mut env).is_err());
    assert_eq!(s.blocks.len(), 1);

    let src = "t \"hi again\"\n
;"; //finish source to parse

    let bytes = src.as_bytes();
    let c = Cursor::new(&bytes[..]);

    s.stream = c; //swap in new 'stream'
    let idx = s.parse();

    assert!(idx.is_some());
    if let Some(idx) = idx {
        let b = s.blocks.get(idx).expect("ERROR: No block");
        match b {
            &Block::Src(ref b) => {
                assert_eq!(b.name, "some_block".to_owned());
                assert_eq!(b.src.get(0), Some(&Src::Emit(vec![Var::String("hi again".to_owned())])))
            },
            _ => { panic!("ERROR: Invalid block type") }
        }
    }

    // assert we can sink parsed blocks into an environment
    assert_eq!(s.blocks.len(), 2);
    assert!(s.sink(&mut env).is_ok());
    assert_eq!(s.blocks.len(), 0);

    let mut data = Empty;
    let mut ev = Evaluator::new(&mut env, &mut data);
    let (vars,_) = ev.next().unwrap();
    assert_eq!(vars.get(0), Some(&Var::String("hi".to_owned())));
    let (vars,_) = ev.nth(1).expect("ERROR: Block failed to transition");
    assert_eq!(vars.get(0), Some(&Var::String("hi again".to_owned())));
}
<|endoftext|>


<commit_before>extern crate xfdocs;
use xfdocs::xflow::xfstruct::*;

use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;

fn read_json_file(filename:&str) -> String {
    // Create a path to the desired file
    let path = Path::new(filename);
    let display = path.display();

    let mut file = match File::open(&path) {
        Err(why) => panic!(
            "couldn't open {}: {}",
            display,
            Error::description(&why)
            ),
        Ok(file) => file,
    };

    // Read the file contents into a string, returns `io::Result<usize>`
    let mut s = String::new();
    match file.read_to_string(&mut s) {
        Err(why) => panic!("couldn't read {}: {}", display,
                                                   Error::description(&why)),
        Ok(_)    => print!("{} contains:\n{}", display, s),
    };

    s
}

//
// Test code
//

fn x_get_rand() -> i32 {
    5 // rand::thread_rng().gen_range(1, 11)
}

fn create_node(id:i32) -> XFlowNode {
    XFlowNode {
        id: id,
        nodetype: "flow".to_string(),
        action: "some action".to_string(),
        label: "some name".to_string()
    }
}

fn create_nodes(amount:i32) -> Vec<XFlowNode> {
    let mut nodes = Vec::<XFlowNode>::new();

    for i in 0..amount {
        nodes.push(create_node(i))
    }

    return nodes
}

fn create_edges(amount:i32) -> Vec<XFlowEdge> {
    let left   = x_get_rand();
    let right  = x_get_rand();

    let mut edges = Vec::<XFlowEdge>::new();

    for i in 0..amount {
        edges.push([left, right + i]);
    }

    edges
}

fn create_branches(amount:i32) -> Vec<XFlowBranch> {

    let mut branches = Vec::<XFlowBranch>::new();

    for i in 0..amount {
        let left   = x_get_rand();
        let right  = x_get_rand();
        branches.push(
            XFlowBranch {
                name: "Some branch".to_string(),
                edge : [left, right + i]
            })
    }

    branches
}

fn create_xflow_struct() -> XFlowStruct {
    XFlowStruct {
        id:       "id1".to_string(),
        version:  1,
        name:     "Some name".to_string(),
        requirements: Vec::<XFlowRequirement>::new(),
        variables:    XFlowVariables {
            input:  Vec::<XFlowVariable>::new(),
            local:  Vec::<XFlowVariable>::new(),
            output: Vec::<XFlowVariableDefinition>::new(),
        },
        nodes:        create_nodes(5),
        edges:        create_edges(5),
        branches:     create_branches(5)
    }
}

#[test]
fn test_xfs() {
    let xfs = create_xflow_struct();

    assert_eq!(xfs.nodes.len(), 5);
    assert_eq!(xfs.edges.len(), 5);
    assert_eq!(xfs.branches.len(), 5);
}

#[test]
fn test_xfs_fields() {
    let xfs = create_xflow_struct();

    assert_eq!(xfs.version, 1);
    assert_eq!(xfs.id, "id1".to_string());
    assert_eq!(xfs.name, "Some name");
}

#[test]
// #[should_panic]
fn test_xfs_entry() {
    let json_string = read_json_file("data/flows/10_steps.json");
    let xfs = XFlowStruct::from_json(&json_string);

    assert_eq!(xfs.get_nodes_by("flow", "start").len(), 1);
}

#[test]
fn test_xfs_get_nodes_of_type() {
    let json_string = read_json_file("data/flows/10_steps.json");
    let xfs = XFlowStruct::from_json(&json_string);

    assert_eq!(xfs.get_nodes_of_type("flow").len(), 2);
}

#[test]
fn test_xfs_from_json() {
    let json_string = read_json_file("data/flows/10_steps.json");
    let xfs = XFlowStruct::from_json(&json_string);

    assert_eq!(xfs.name, "steps".to_string());
    assert_eq!(xfs.nodes.len(), 10);
    assert_eq!(xfs.edges.len(), 9);
    assert_eq!(xfs.branches.len(), 0);

    assert_eq!(xfs.requirements.len(), 2);

    assert_eq!(xfs.variables.input.len(), 1);
    assert_eq!(xfs.variables.local.len(), 0);
    assert_eq!(xfs.variables.output.len(), 1);

    assert!(xfs.get_entry_node().is_ok());
    assert!(xfs.get_terminal_nodes().is_ok());

    match xfs.get_entry_node() {
        Ok(res)  => assert_eq!(res.id, 1),
        Err(err) => println!("Error: {:?}", err),
    }

    match xfs.get_terminal_nodes() {
        Ok(res)  => assert_eq!(res.len(), 1),
        Err(err) => println!("Error: {:?}", err),
    }

}

#[test]
fn test_xfs_from_json_string() {
    let empty_flow = read_json_file("data/flows/empty.json");
    let xfs = XFlowStruct::from_json(&empty_flow);

    assert_eq!(xfs.name, "empty".to_string());
    assert_eq!(xfs.nodes.len(), 0);
    assert_eq!(xfs.edges.len(), 0);
    assert_eq!(xfs.branches.len(), 0);
}

<commit_msg>Add basic test coverage for get_edges<commit_after>extern crate xfdocs;
use xfdocs::xflow::xfstruct::*;

use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;

fn read_json_file(filename:&str) -> String {
    // Create a path to the desired file
    let path = Path::new(filename);
    let display = path.display();

    let mut file = match File::open(&path) {
        Err(why) => panic!(
            "couldn't open {}: {}",
            display,
            Error::description(&why)
            ),
        Ok(file) => file,
    };

    // Read the file contents into a string, returns `io::Result<usize>`
    let mut s = String::new();
    match file.read_to_string(&mut s) {
        Err(why) => panic!("couldn't read {}: {}", display,
                                                   Error::description(&why)),
        Ok(_)    => print!("{} contains:\n{}", display, s),
    };

    s
}

//
// Test code
//

fn x_get_rand() -> i32 {
    5 // rand::thread_rng().gen_range(1, 11)
}

fn create_node(id:i32) -> XFlowNode {
    XFlowNode {
        id: id,
        nodetype: "flow".to_string(),
        action: "some action".to_string(),
        label: "some name".to_string()
    }
}

fn create_nodes(amount:i32) -> Vec<XFlowNode> {
    let mut nodes = Vec::<XFlowNode>::new();

    for i in 0..amount {
        nodes.push(create_node(i))
    }

    return nodes
}

fn create_edges(amount:i32) -> Vec<XFlowEdge> {
    let left   = x_get_rand();
    let right  = x_get_rand();

    let mut edges = Vec::<XFlowEdge>::new();

    for i in 0..amount {
        edges.push([left, right + i]);
    }

    edges
}

fn create_branches(amount:i32) -> Vec<XFlowBranch> {

    let mut branches = Vec::<XFlowBranch>::new();

    for i in 0..amount {
        let left   = x_get_rand();
        let right  = x_get_rand();
        branches.push(
            XFlowBranch {
                name: "Some branch".to_string(),
                edge : [left, right + i]
            })
    }

    branches
}

fn create_xflow_struct() -> XFlowStruct {
    XFlowStruct {
        id:       "id1".to_string(),
        version:  1,
        name:     "Some name".to_string(),
        requirements: Vec::<XFlowRequirement>::new(),
        variables:    XFlowVariables {
            input:  Vec::<XFlowVariable>::new(),
            local:  Vec::<XFlowVariable>::new(),
            output: Vec::<XFlowVariableDefinition>::new(),
        },
        nodes:        create_nodes(5),
        edges:        create_edges(5),
        branches:     create_branches(5)
    }
}

#[test]
fn test_xfs() {
    let xfs = create_xflow_struct();

    assert_eq!(xfs.nodes.len(), 5);
    assert_eq!(xfs.edges.len(), 5);
    assert_eq!(xfs.branches.len(), 5);
}

#[test]
fn test_xfs_fields() {
    let xfs = create_xflow_struct();

    assert_eq!(xfs.version, 1);
    assert_eq!(xfs.id, "id1".to_string());
    assert_eq!(xfs.name, "Some name");
}

#[test]
// #[should_panic]
fn test_xfs_entry() {
    let json_string = read_json_file("data/flows/10_steps.json");
    let xfs = XFlowStruct::from_json(&json_string);

    assert_eq!(xfs.get_nodes_by("flow", "start").len(), 1);
}

#[test]
fn test_xfs_get_nodes_of_type() {
    let json_string = read_json_file("data/flows/10_steps.json");
    let xfs = XFlowStruct::from_json(&json_string);

    assert_eq!(xfs.get_nodes_of_type("flow").len(), 2);
}

#[test]
fn test_xfs_from_json() {
    let json_string = read_json_file("data/flows/10_steps.json");
    let xfs = XFlowStruct::from_json(&json_string);

    assert_eq!(xfs.name, "steps".to_string());
    assert_eq!(xfs.nodes.len(), 10);
    assert_eq!(xfs.edges.len(), 9);
    assert_eq!(xfs.branches.len(), 0);

    assert_eq!(xfs.requirements.len(), 2);

    assert_eq!(xfs.variables.input.len(), 1);
    assert_eq!(xfs.variables.local.len(), 0);
    assert_eq!(xfs.variables.output.len(), 1);

    assert!(xfs.get_entry_node().is_ok());
    assert!(xfs.get_terminal_nodes().is_ok());

    match xfs.get_entry_node() {
        Ok(res)  => assert_eq!(res.id, 1),
        Err(err) => println!("Error: {:?}", err),
    }

    match xfs.get_terminal_nodes() {
        Ok(res)  => assert_eq!(res.len(), 1),
        Err(err) => println!("Error: {:?}", err),
    }

    match xfs.get_entry_node() {
        Ok(res)  => {

            let in_edges  = xfs.get_in_edges(res);
            let out_edges = xfs.get_out_edges(res);

            assert_eq!(in_edges.len(), 0);
            assert_eq!(out_edges.len(), 1);

            assert_eq!(xfs.get_branches_for(out_edges[0]).len(), 0);

        },
        Err(err) => println!("Error: {:?}", err),
    }
}

#[test]
fn test_xfs_from_json_string() {
    let empty_flow = read_json_file("data/flows/empty.json");
    let xfs = XFlowStruct::from_json(&empty_flow);

    assert_eq!(xfs.name, "empty".to_string());
    assert_eq!(xfs.nodes.len(), 0);
    assert_eq!(xfs.edges.len(), 0);
    assert_eq!(xfs.branches.len(), 0);
}

<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate rustfmt;
extern crate diff;
extern crate regex;
extern crate term;

use std::collections::HashMap;
use std::fs;
use std::io::{self, Read, BufRead, BufReader};

use rustfmt::*;
use rustfmt::config::{Config, ReportTactic};
use rustfmt::rustfmt_diff::*;

static DIFF_CONTEXT_SIZE: usize = 3;

fn get_path_string(dir_entry: io::Result<fs::DirEntry>) -> String {
    let path = dir_entry.ok().expect("Couldn't get DirEntry.").path();

    path.to_str().expect("Couldn't stringify path.").to_owned()
}

// Integration tests. The files in the tests/source are formatted and compared
// to their equivalent in tests/target. The target file and config can be
// overriden by annotations in the source file. The input and output must match
// exactly.
// FIXME(#28) would be good to check for error messages and fail on them, or at
// least report.
#[test]
fn system_tests() {
    // Get all files in the tests/source directory.
    let files = fs::read_dir("tests/source").ok().expect("Couldn't read source dir.");
    // Turn a DirEntry into a String that represents the relative path to the
    // file.
    let files = files.map(get_path_string);
    let (_reports, count, fails) = check_files(files);

    // Display results.
    println!("Ran {} system tests.", count);
    assert!(fails == 0, "{} system tests failed", fails);
}

// Idempotence tests. Files in tests/target are checked to be unaltered by
// rustfmt.
#[test]
fn idempotence_tests() {
    // Get all files in the tests/target directory.
    let files = fs::read_dir("tests/target")
                    .ok()
                    .expect("Couldn't read target dir.")
                    .map(get_path_string);
    let (_reports, count, fails) = check_files(files);

    // Display results.
    println!("Ran {} idempotent tests.", count);
    assert!(fails == 0, "{} idempotent tests failed", fails);
}

// Run rustfmt on itself. This operation must be idempotent. We also check that
// no warnings are emitted.
#[test]
fn self_tests() {
    let files = fs::read_dir("src/bin")
                    .ok()
                    .expect("Couldn't read src dir.")
                    .chain(fs::read_dir("tests").ok().expect("Couldn't read tests dir."))
                    .map(get_path_string);
    // Hack because there's no `IntoIterator` impl for `[T; N]`.
    let files = files.chain(Some("src/lib.rs".to_owned()).into_iter());

    let (reports, count, fails) = check_files(files);
    let mut warnings = 0;

    // Display results.
    println!("Ran {} self tests.", count);
    assert!(fails == 0, "{} self tests failed", fails);

    for format_report in reports {
        println!("{}", format_report);
        warnings += format_report.warning_count();
    }

    assert!(warnings == 0,
            "Rustfmt's code generated {} warnings",
            warnings);
}

// For each file, run rustfmt and collect the output.
// Returns the number of files checked and the number of failures.
fn check_files<I>(files: I) -> (Vec<FormatReport>, u32, u32)
    where I: Iterator<Item = String>
{
    let mut count = 0;
    let mut fails = 0;
    let mut reports = vec![];

    for file_name in files.filter(|f| f.ends_with(".rs")) {
        println!("Testing '{}'...", file_name);

        match idempotent_check(file_name) {
            Ok(report) => reports.push(report),
            Err(msg) => {
                print_mismatches(msg);
                fails += 1;
            }
        }

        count += 1;
    }

    (reports, count, fails)
}

fn print_mismatches(result: HashMap<String, Vec<Mismatch>>) {
    let mut t = term::stdout().unwrap();

    for (file_name, diff) in result {
        print_diff(diff,
                   |line_num| format!("\nMismatch at {}:{}:", file_name, line_num));
    }

    assert!(t.reset().unwrap());
}

pub fn idempotent_check(filename: String) -> Result<FormatReport, HashMap<String, Vec<Mismatch>>> {
    let sig_comments = read_significant_comments(&filename);
    let mut config = get_config(sig_comments.get("config").map(|x| &(*x)[..]));
    let args = vec!["rustfmt".to_owned(), filename];

    for (key, val) in &sig_comments {
        if key != "target" && key != "config" {
            config.override_value(key, val);
        }
    }

    // Don't generate warnings for to-do items.
    config.report_todo = ReportTactic::Never;

    let mut file_map = format(args, &config);
    let format_report = fmt_lines(&mut file_map, &config);

    // Won't panic, as we're not doing any IO.
    let write_result = filemap::write_all_files(&file_map, WriteMode::Return, &config).unwrap();
    let target = sig_comments.get("target").map(|x| &(*x)[..]);

    handle_result(write_result, target).map(|_| format_report)
}

// Reads test config file from comments and reads its contents.
fn get_config(config_file: Option<&str>) -> Config {
    let config_file_name = match config_file {
        None => return Default::default(),
        Some(file_name) => {
            let mut full_path = "tests/config/".to_owned();
            full_path.push_str(&file_name);
            full_path
        }
    };

    let mut def_config_file = fs::File::open(config_file_name)
                                  .ok()
                                  .expect("Couldn't open config.");
    let mut def_config = String::new();
    def_config_file.read_to_string(&mut def_config).ok().expect("Couldn't read config.");

    Config::from_toml(&def_config)
}

// Reads significant comments of the form: // rustfmt-key: value
// into a hash map.
fn read_significant_comments(file_name: &str) -> HashMap<String, String> {
    let file = fs::File::open(file_name)
                   .ok()
                   .expect(&format!("Couldn't read file {}.", file_name));
    let reader = BufReader::new(file);
    let pattern = r"^\s*//\s*rustfmt-([^:]+):\s*(\S+)";
    let regex = regex::Regex::new(&pattern).ok().expect("Failed creating pattern 1.");

    // Matches lines containing significant comments or whitespace.
    let line_regex = regex::Regex::new(r"(^\s*$)|(^\s*//\s*rustfmt-[^:]+:\s*\S+)")
                         .ok()
                         .expect("Failed creating pattern 2.");

    reader.lines()
          .map(|line| line.ok().expect("Failed getting line."))
          .take_while(|line| line_regex.is_match(&line))
          .filter_map(|line| {
              regex.captures_iter(&line).next().map(|capture| {
                  (capture.at(1).expect("Couldn't unwrap capture.").to_owned(),
                   capture.at(2).expect("Couldn't unwrap capture.").to_owned())
              })
          })
          .collect()
}

// Compare output to input.
// TODO: needs a better name, more explanation.
fn handle_result(result: HashMap<String, String>,
                 target: Option<&str>)
                 -> Result<(), HashMap<String, Vec<Mismatch>>> {
    let mut failures = HashMap::new();

    for (file_name, fmt_text) in result {
        // If file is in tests/source, compare to file with same name in tests/target.
        let target = get_target(&file_name, target);
        let mut f = fs::File::open(&target).ok().expect("Couldn't open target.");

        let mut text = String::new();
        f.read_to_string(&mut text).ok().expect("Failed reading target.");

        if fmt_text != text {
            let diff = make_diff(&text, &fmt_text, DIFF_CONTEXT_SIZE);
            failures.insert(file_name, diff);
        }
    }

    if failures.is_empty() {
        Ok(())
    } else {
        Err(failures)
    }
}

// Map source file paths to their target paths.
fn get_target(file_name: &str, target: Option<&str>) -> String {
    if file_name.starts_with("tests/source/") {
        let base = target.unwrap_or(file_name.trim_left_matches("tests/source/"));

        format!("tests/target/{}", base)
    } else {
        file_name.to_owned()
    }
}
<commit_msg>Fix path issue with tests on Windows<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate rustfmt;
extern crate diff;
extern crate regex;
extern crate term;

use std::collections::HashMap;
use std::fs;
use std::io::{self, Read, BufRead, BufReader};
use std::path::Path;

use rustfmt::*;
use rustfmt::config::{Config, ReportTactic};
use rustfmt::rustfmt_diff::*;

static DIFF_CONTEXT_SIZE: usize = 3;

fn get_path_string(dir_entry: io::Result<fs::DirEntry>) -> String {
    let path = dir_entry.ok().expect("Couldn't get DirEntry.").path();

    path.to_str().expect("Couldn't stringify path.").to_owned()
}

// Integration tests. The files in the tests/source are formatted and compared
// to their equivalent in tests/target. The target file and config can be
// overriden by annotations in the source file. The input and output must match
// exactly.
// FIXME(#28) would be good to check for error messages and fail on them, or at
// least report.
#[test]
fn system_tests() {
    // Get all files in the tests/source directory.
    let files = fs::read_dir("tests/source").ok().expect("Couldn't read source dir.");
    // Turn a DirEntry into a String that represents the relative path to the
    // file.
    let files = files.map(get_path_string);
    let (_reports, count, fails) = check_files(files);

    // Display results.
    println!("Ran {} system tests.", count);
    assert!(fails == 0, "{} system tests failed", fails);
}

// Idempotence tests. Files in tests/target are checked to be unaltered by
// rustfmt.
#[test]
fn idempotence_tests() {
    // Get all files in the tests/target directory.
    let files = fs::read_dir("tests/target")
                    .ok()
                    .expect("Couldn't read target dir.")
                    .map(get_path_string);
    let (_reports, count, fails) = check_files(files);

    // Display results.
    println!("Ran {} idempotent tests.", count);
    assert!(fails == 0, "{} idempotent tests failed", fails);
}

// Run rustfmt on itself. This operation must be idempotent. We also check that
// no warnings are emitted.
#[test]
fn self_tests() {
    let files = fs::read_dir("src/bin")
                    .ok()
                    .expect("Couldn't read src dir.")
                    .chain(fs::read_dir("tests").ok().expect("Couldn't read tests dir."))
                    .map(get_path_string);
    // Hack because there's no `IntoIterator` impl for `[T; N]`.
    let files = files.chain(Some("src/lib.rs".to_owned()).into_iter());

    let (reports, count, fails) = check_files(files);
    let mut warnings = 0;

    // Display results.
    println!("Ran {} self tests.", count);
    assert!(fails == 0, "{} self tests failed", fails);

    for format_report in reports {
        println!("{}", format_report);
        warnings += format_report.warning_count();
    }

    assert!(warnings == 0,
            "Rustfmt's code generated {} warnings",
            warnings);
}

// For each file, run rustfmt and collect the output.
// Returns the number of files checked and the number of failures.
fn check_files<I>(files: I) -> (Vec<FormatReport>, u32, u32)
    where I: Iterator<Item = String>
{
    let mut count = 0;
    let mut fails = 0;
    let mut reports = vec![];

    for file_name in files.filter(|f| f.ends_with(".rs")) {
        println!("Testing '{}'...", file_name);

        match idempotent_check(file_name) {
            Ok(report) => reports.push(report),
            Err(msg) => {
                print_mismatches(msg);
                fails += 1;
            }
        }

        count += 1;
    }

    (reports, count, fails)
}

fn print_mismatches(result: HashMap<String, Vec<Mismatch>>) {
    let mut t = term::stdout().unwrap();

    for (file_name, diff) in result {
        print_diff(diff,
                   |line_num| format!("\nMismatch at {}:{}:", file_name, line_num));
    }

    assert!(t.reset().unwrap());
}

pub fn idempotent_check(filename: String) -> Result<FormatReport, HashMap<String, Vec<Mismatch>>> {
    let sig_comments = read_significant_comments(&filename);
    let mut config = get_config(sig_comments.get("config").map(|x| &(*x)[..]));
    let args = vec!["rustfmt".to_owned(), filename];

    for (key, val) in &sig_comments {
        if key != "target" && key != "config" {
            config.override_value(key, val);
        }
    }

    // Don't generate warnings for to-do items.
    config.report_todo = ReportTactic::Never;

    let mut file_map = format(args, &config);
    let format_report = fmt_lines(&mut file_map, &config);

    // Won't panic, as we're not doing any IO.
    let write_result = filemap::write_all_files(&file_map, WriteMode::Return, &config).unwrap();
    let target = sig_comments.get("target").map(|x| &(*x)[..]);

    handle_result(write_result, target).map(|_| format_report)
}

// Reads test config file from comments and reads its contents.
fn get_config(config_file: Option<&str>) -> Config {
    let config_file_name = match config_file {
        None => return Default::default(),
        Some(file_name) => {
            let mut full_path = "tests/config/".to_owned();
            full_path.push_str(&file_name);
            full_path
        }
    };

    let mut def_config_file = fs::File::open(config_file_name)
                                  .ok()
                                  .expect("Couldn't open config.");
    let mut def_config = String::new();
    def_config_file.read_to_string(&mut def_config).ok().expect("Couldn't read config.");

    Config::from_toml(&def_config)
}

// Reads significant comments of the form: // rustfmt-key: value
// into a hash map.
fn read_significant_comments(file_name: &str) -> HashMap<String, String> {
    let file = fs::File::open(file_name)
                   .ok()
                   .expect(&format!("Couldn't read file {}.", file_name));
    let reader = BufReader::new(file);
    let pattern = r"^\s*//\s*rustfmt-([^:]+):\s*(\S+)";
    let regex = regex::Regex::new(&pattern).ok().expect("Failed creating pattern 1.");

    // Matches lines containing significant comments or whitespace.
    let line_regex = regex::Regex::new(r"(^\s*$)|(^\s*//\s*rustfmt-[^:]+:\s*\S+)")
                         .ok()
                         .expect("Failed creating pattern 2.");

    reader.lines()
          .map(|line| line.ok().expect("Failed getting line."))
          .take_while(|line| line_regex.is_match(&line))
          .filter_map(|line| {
              regex.captures_iter(&line).next().map(|capture| {
                  (capture.at(1).expect("Couldn't unwrap capture.").to_owned(),
                   capture.at(2).expect("Couldn't unwrap capture.").to_owned())
              })
          })
          .collect()
}

// Compare output to input.
// TODO: needs a better name, more explanation.
fn handle_result(result: HashMap<String, String>,
                 target: Option<&str>)
                 -> Result<(), HashMap<String, Vec<Mismatch>>> {
    let mut failures = HashMap::new();

    for (file_name, fmt_text) in result {
        // If file is in tests/source, compare to file with same name in tests/target.
        let target = get_target(&file_name, target);
        let mut f = fs::File::open(&target).ok().expect("Couldn't open target.");

        let mut text = String::new();
        f.read_to_string(&mut text).ok().expect("Failed reading target.");

        if fmt_text != text {
            let diff = make_diff(&text, &fmt_text, DIFF_CONTEXT_SIZE);
            failures.insert(file_name, diff);
        }
    }

    if failures.is_empty() {
        Ok(())
    } else {
        Err(failures)
    }
}

// Map source file paths to their target paths.
fn get_target(file_name: &str, target: Option<&str>) -> String {
    let file_path = Path::new(file_name);
    let source_path_prefix = Path::new("tests/source/");
    if file_path.starts_with(source_path_prefix) {
        let mut components = file_path.components();
        // Can't skip(2) as the resulting iterator can't as_path()
        components.next();
        components.next();

        let new_target = match components.as_path().to_str() {
            Some(string) => string,
            None => file_name,
        };
        let base = target.unwrap_or(new_target);

        format!("tests/target/{}", base)
    } else {
        file_name.to_owned()
    }
}
<|endoftext|>