<commit_before>use crate::error;
use std::io;

pub trait LZBuffer {
    fn len(&self) -> usize;
    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8;
    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8>;
    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()>;
    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()>;
    // Flush the buffer to the output
    fn finish(self) -> io::Result<()>;
}

// An accumulating buffer for LZ sequences
pub struct LZAccumBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Buffer
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W) -> Self {
        Self {
            stream,
            buf: Vec::new(),
            len: 0,
        }
    }

    // Append bytes
    pub fn append_bytes(&mut self, buf: &[u8]) {
        self.buf.extend_from_slice(buf);
        self.len += buf.len();
    }

    // Reset the internal dictionary
    pub fn reset(&mut self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.buf.clear();
        // TODO: clear the length ?
        Ok(())
    }
}

impl<'a, W> LZBuffer for LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        let buf_len = self.buf.len();
        if buf_len == 0 {
            lit
        } else {
            self.buf[buf_len - 1]
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        Ok(self.buf[buf_len - dist])
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        self.buf.push(lit);
        self.len += 1;
        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        let mut offset = buf_len - dist;
        for _ in 0..len {
            let x = self.buf[offset];
            self.buf.push(x);
            offset += 1;
        }
        self.len += len;
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.stream.flush()?;
        Ok(())
    }
}

// A circular buffer for LZ sequences
pub struct LZCircularBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Circular buffer
    dict_size: usize,  // Length of the buffer
    cursor: usize,     // Current position
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W, dict_size: usize) -> Self {
        info!("Dict size in LZ buffer: {}", dict_size);
        Self {
            stream: stream,
            buf: Vec::new(),
            dict_size,
            cursor: 0,
            len: 0,
        }
    }

    fn get(&self, index: usize) -> u8 {
        match self.buf.get(index) {
            Some(x) => *x,
            None => 0,
        }
    }

    fn get_mut(&mut self, index: usize) -> &mut u8 {
        self.buf.resize(index + 1, 0);
        self.buf.get_mut(index).unwrap()
    }
}

impl<'a, W> LZBuffer for LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        if self.len == 0 {
            lit
        } else {
            self.get((self.dict_size + self.cursor - 1) % self.dict_size)
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        Ok(self.get(offset))
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        *self.get_mut(self.cursor) = lit;
        self.cursor += 1;
        self.len += 1;

        // Flush the circular buffer to the output
        if self.cursor == self.dict_size {
            self.stream.write_all(self.buf.as_slice())?;
            self.cursor = 0;
        }

        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let mut offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        for _ in 0..len {
            let x = self.get(offset);
            self.append_literal(x)?;
            offset += 1;
            if offset == self.dict_size {
                offset = 0
            }
        }
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        if self.cursor > 0 {
            self.stream.write_all(&self.buf[0..self.cursor])?;
            self.stream.flush()?;
        }
        Ok(())
    }
}
<commit_msg>Adapt to comments from review<commit_after>use crate::error;
use std::io;

pub trait LZBuffer {
    fn len(&self) -> usize;
    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8;
    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8>;
    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()>;
    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()>;
    // Flush the buffer to the output
    fn finish(self) -> io::Result<()>;
}

// An accumulating buffer for LZ sequences
pub struct LZAccumBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Buffer
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W) -> Self {
        Self {
            stream,
            buf: Vec::new(),
            len: 0,
        }
    }

    // Append bytes
    pub fn append_bytes(&mut self, buf: &[u8]) {
        self.buf.extend_from_slice(buf);
        self.len += buf.len();
    }

    // Reset the internal dictionary
    pub fn reset(&mut self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.buf.clear();
        // TODO: clear the length ?
        Ok(())
    }
}

impl<'a, W> LZBuffer for LZAccumBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        let buf_len = self.buf.len();
        if buf_len == 0 {
            lit
        } else {
            self.buf[buf_len - 1]
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        Ok(self.buf[buf_len - dist])
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        self.buf.push(lit);
        self.len += 1;
        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        let buf_len = self.buf.len();
        if dist > buf_len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, buf_len
            )));
        }

        let mut offset = buf_len - dist;
        for _ in 0..len {
            let x = self.buf[offset];
            self.buf.push(x);
            offset += 1;
        }
        self.len += len;
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        self.stream.write_all(self.buf.as_slice())?;
        self.stream.flush()?;
        Ok(())
    }
}

// A circular buffer for LZ sequences
pub struct LZCircularBuffer<'a, W>
where
    W: 'a + io::Write,
{
    stream: &'a mut W, // Output sink
    buf: Vec<u8>,      // Circular buffer
    dict_size: usize,  // Length of the buffer
    cursor: usize,     // Current position
    len: usize,        // Total number of bytes sent through the buffer
}

impl<'a, W> LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    pub fn from_stream(stream: &'a mut W, dict_size: usize) -> Self {
        info!("Dict size in LZ buffer: {}", dict_size);
        Self {
            stream,
            buf: Vec::new(),
            dict_size,
            cursor: 0,
            len: 0,
        }
    }

    fn get(&self, index: usize) -> u8 {
        *self.buf.get(index).unwrap_or(&0)
    }

    fn get_mut(&mut self, index: usize) -> &mut u8 {
        if self.buf.len() < index + 1 {
            self.buf.resize(index + 1, 0);
        }
        self.buf.get_mut(index).unwrap()
    }
}

impl<'a, W> LZBuffer for LZCircularBuffer<'a, W>
where
    W: io::Write,
{
    fn len(&self) -> usize {
        self.len
    }

    // Retrieve the last byte or return a default
    fn last_or(&self, lit: u8) -> u8 {
        if self.len == 0 {
            lit
        } else {
            self.get((self.dict_size + self.cursor - 1) % self.dict_size)
        }
    }

    // Retrieve the n-th last byte
    fn last_n(&self, dist: usize) -> error::Result<u8> {
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "Match distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        Ok(self.get(offset))
    }

    // Append a literal
    fn append_literal(&mut self, lit: u8) -> io::Result<()> {
        *self.get_mut(self.cursor) = lit;
        self.cursor += 1;
        self.len += 1;

        // Flush the circular buffer to the output
        if self.cursor == self.dict_size {
            self.stream.write_all(self.buf.as_slice())?;
            self.cursor = 0;
        }

        Ok(())
    }

    // Fetch an LZ sequence (length, distance) from inside the buffer
    fn append_lz(&mut self, len: usize, dist: usize) -> error::Result<()> {
        debug!("LZ {{ len: {}, dist: {} }}", len, dist);
        if dist > self.dict_size {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond dictionary size {}",
                dist, self.dict_size
            )));
        }
        if dist > self.len {
            return Err(error::Error::LZMAError(format!(
                "LZ distance {} is beyond output size {}",
                dist, self.len
            )));
        }

        let mut offset = (self.dict_size + self.cursor - dist) % self.dict_size;
        for _ in 0..len {
            let x = self.get(offset);
            self.append_literal(x)?;
            offset += 1;
            if offset == self.dict_size {
                offset = 0
            }
        }
        Ok(())
    }

    // Flush the buffer to the output
    fn finish(self) -> io::Result<()> {
        if self.cursor > 0 {
            self.stream.write_all(&self.buf[0..self.cursor])?;
            self.stream.flush()?;
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::mem;
use std::slice;

use super::ffi;
use super::libc::{c_char, c_int, c_uint};

use result::{Error, Kind, Result};

const ZLIB_VERSION: &'static str = "1.2.8\0";

trait Context {
    fn stream(&mut self) -> &mut ffi::z_stream;

    fn stream_apply<F>(&mut self, input: &[u8], output: &mut Vec<u8>, each: F) -> Result<()>
    where
        F: Fn(&mut ffi::z_stream) -> Option<Result<()>>,
    {
        debug_assert!(output.len() == 0, "Output vector is not empty.");

        let stream = self.stream();

        stream.next_in = input.as_ptr() as *mut _;
        stream.avail_in = input.len() as c_uint;

        let mut output_size;

        loop {
            output_size = output.len();

            if output_size == output.capacity() {
                output.reserve(input.len())
            }

            let out_slice = unsafe {
                slice::from_raw_parts_mut(
                    output.as_mut_ptr().offset(output_size as isize),
                    output.capacity() - output_size,
                )
            };

            stream.next_out = out_slice.as_mut_ptr();
            stream.avail_out = out_slice.len() as c_uint;

            let before = stream.total_out;
            let cont = each(stream);

            unsafe {
                output.set_len((stream.total_out - before) as usize + output_size);
            }

            if let Some(result) = cont {
                return result;
            }
        }
    }
}

pub struct Compressor {
    // Box the z_stream to ensure it isn't moved. Moving the z_stream
    // causes zlib to fail, because it maintains internal pointers.
    stream: Box<ffi::z_stream>,
}

impl Compressor {
    pub fn new(window_bits: i8) -> Compressor {
        debug_assert!(window_bits >= 9, "Received too small window size.");
        debug_assert!(window_bits <= 15, "Received too large window size.");

        unsafe {
            let mut stream: Box<ffi::z_stream> = Box::new(MaybeUninit::uninit().assume_init());
            let result = ffi::deflateInit2_(
                stream.as_mut(),
                9,
                ffi::Z_DEFLATED,
                -window_bits as c_int,
                9,
                ffi::Z_DEFAULT_STRATEGY,
                ZLIB_VERSION.as_ptr() as *const c_char,
                mem::size_of::<ffi::z_stream>() as c_int,
            );
            assert!(result == ffi::Z_OK, "Failed to initialize compresser.");
            Compressor { stream: stream }
        }
    }

    pub fn compress(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<()> {
        self.stream_apply(input, output, |stream| unsafe {
            match ffi::deflate(stream, ffi::Z_SYNC_FLUSH) {
                ffi::Z_OK | ffi::Z_BUF_ERROR => {
                    if stream.avail_in == 0 && stream.avail_out > 0 {
                        Some(Ok(()))
                    } else {
                        None
                    }
                }
                code => Some(Err(Error::new(
                    Kind::Protocol,
                    format!("Failed to perform compression: {}", code),
                ))),
            }
        })
    }

    pub fn reset(&mut self) -> Result<()> {
        match unsafe { ffi::deflateReset(self.stream.as_mut()) } {
            ffi::Z_OK => Ok(()),
            code => Err(Error::new(
                Kind::Protocol,
                format!("Failed to reset compression context: {}", code),
            )),
        }
    }
}

impl Context for Compressor {
    fn stream(&mut self) -> &mut ffi::z_stream {
        self.stream.as_mut()
    }
}

impl Drop for Compressor {
    fn drop(&mut self) {
        match unsafe { ffi::deflateEnd(self.stream.as_mut()) } {
            ffi::Z_STREAM_ERROR => error!("Compression stream encountered bad state."),
            // Ignore discarded data error because we are raw
            ffi::Z_OK | ffi::Z_DATA_ERROR => trace!("Deallocated compression context."),
            code => error!("Bad zlib status encountered: {}", code),
        }
    }
}

pub struct Decompressor {
    stream: Box<ffi::z_stream>,
}

impl Decompressor {
    pub fn new(window_bits: i8) -> Decompressor {
        debug_assert!(window_bits >= 8, "Received too small window size.");
        debug_assert!(window_bits <= 15, "Received too large window size.");

        unsafe {
            let mut stream: Box<ffi::z_stream> = Box::new(MaybeUninit::uninit().assume_init());
            let result = ffi::inflateInit2_(
                stream.as_mut(),
                -window_bits as c_int,
                ZLIB_VERSION.as_ptr() as *const c_char,
                mem::size_of::<ffi::z_stream>() as c_int,
            );
            assert!(result == ffi::Z_OK, "Failed to initialize decompresser.");
            Decompressor { stream: stream }
        }
    }

    pub fn decompress(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<()> {
        self.stream_apply(input, output, |stream| unsafe {
            match ffi::inflate(stream, ffi::Z_SYNC_FLUSH) {
                ffi::Z_OK | ffi::Z_BUF_ERROR => {
                    if stream.avail_in == 0 && stream.avail_out > 0 {
                        Some(Ok(()))
                    } else {
                        None
                    }
                }
                code => Some(Err(Error::new(
                    Kind::Protocol,
                    format!("Failed to perform decompression: {}", code),
                ))),
            }
        })
    }

    pub fn reset(&mut self) -> Result<()> {
        match unsafe { ffi::inflateReset(self.stream.as_mut()) } {
            ffi::Z_OK => Ok(()),
            code => Err(Error::new(
                Kind::Protocol,
                format!("Failed to reset compression context: {}", code),
            )),
        }
    }
}

impl Context for Decompressor {
    fn stream(&mut self) -> &mut ffi::z_stream {
        self.stream.as_mut()
    }
}

impl Drop for Decompressor {
    fn drop(&mut self) {
        match unsafe { ffi::inflateEnd(self.stream.as_mut()) } {
            ffi::Z_STREAM_ERROR => error!("Decompression stream encountered bad state."),
            ffi::Z_OK => trace!("Deallocated decompression context."),
            code => error!("Bad zlib status encountered: {}", code),
        }
    }
}

mod test {
    #![allow(unused_imports, unused_variables, dead_code)]
    use super::*;

    fn as_hex(s: &[u8]) {
        for byte in s {
            print!("0x{:x} ", byte);
        }
        print!("\n");
    }

    #[test]
    fn round_trip() {
        for i in 9..16 {
            let data = "HI THERE THIS IS some data. これはデータだよ。".as_bytes();
            let mut compressed = Vec::with_capacity(data.len());
            let mut decompressed = Vec::with_capacity(data.len());

            let com = Compressor::new(i);
            let mut moved_com = com;

            moved_com
                .compress(&data, &mut compressed)
                .expect("Failed to compress data.");

            let dec = Decompressor::new(i);
            let mut moved_dec = dec;

            moved_dec
                .decompress(&compressed, &mut decompressed)
                .expect("Failed to decompress data.");

            assert_eq!(data, &decompressed[..]);
        }
    }

    #[test]
    fn reset() {
        let data1 = "HI THERE 直子さん".as_bytes();
        let data2 = "HI THERE 人太郎".as_bytes();
        let mut compressed1 = Vec::with_capacity(data1.len());
        let mut compressed2 = Vec::with_capacity(data2.len());
        let mut compressed2_ind = Vec::with_capacity(data2.len());

        let mut decompressed1 = Vec::with_capacity(data1.len());
        let mut decompressed2 = Vec::with_capacity(data2.len());
        let mut decompressed2_ind = Vec::with_capacity(data2.len());

        let mut com = Compressor::new(9);

        com.compress(&data1, &mut compressed1).unwrap();
        com.compress(&data2, &mut compressed2).unwrap();
        com.reset().unwrap();
        com.compress(&data2, &mut compressed2_ind).unwrap();

        let mut dec = Decompressor::new(9);

        dec.decompress(&compressed1, &mut decompressed1).unwrap();
        dec.decompress(&compressed2, &mut decompressed2).unwrap();
        dec.reset().unwrap();
        dec.decompress(&compressed2_ind, &mut decompressed2_ind)
            .unwrap();

        assert_eq!(data1, &decompressed1[..]);
        assert_eq!(data2, &decompressed2[..]);
        assert_eq!(data2, &decompressed2_ind[..]);
        assert!(compressed2 != compressed2_ind);
        assert!(compressed2.len() < compressed2_ind.len());
    }
}
<commit_msg>Update context.rs<commit_after>use std::mem;
use std::slice;
use std::mem::MaybeUninit;

use super::ffi;
use super::libc::{c_char, c_int, c_uint};

use result::{Error, Kind, Result};

const ZLIB_VERSION: &'static str = "1.2.8\0";

trait Context {
    fn stream(&mut self) -> &mut ffi::z_stream;

    fn stream_apply<F>(&mut self, input: &[u8], output: &mut Vec<u8>, each: F) -> Result<()>
    where
        F: Fn(&mut ffi::z_stream) -> Option<Result<()>>,
    {
        debug_assert!(output.len() == 0, "Output vector is not empty.");

        let stream = self.stream();

        stream.next_in = input.as_ptr() as *mut _;
        stream.avail_in = input.len() as c_uint;

        let mut output_size;

        loop {
            output_size = output.len();

            if output_size == output.capacity() {
                output.reserve(input.len())
            }

            let out_slice = unsafe {
                slice::from_raw_parts_mut(
                    output.as_mut_ptr().offset(output_size as isize),
                    output.capacity() - output_size,
                )
            };

            stream.next_out = out_slice.as_mut_ptr();
            stream.avail_out = out_slice.len() as c_uint;

            let before = stream.total_out;
            let cont = each(stream);

            unsafe {
                output.set_len((stream.total_out - before) as usize + output_size);
            }

            if let Some(result) = cont {
                return result;
            }
        }
    }
}

pub struct Compressor {
    // Box the z_stream to ensure it isn't moved. Moving the z_stream
    // causes zlib to fail, because it maintains internal pointers.
    stream: Box<ffi::z_stream>,
}

impl Compressor {
    pub fn new(window_bits: i8) -> Compressor {
        debug_assert!(window_bits >= 9, "Received too small window size.");
        debug_assert!(window_bits <= 15, "Received too large window size.");

        unsafe {
            let mut stream: Box<ffi::z_stream> = Box::new(MaybeUninit::uninit().assume_init());
            let result = ffi::deflateInit2_(
                stream.as_mut(),
                9,
                ffi::Z_DEFLATED,
                -window_bits as c_int,
                9,
                ffi::Z_DEFAULT_STRATEGY,
                ZLIB_VERSION.as_ptr() as *const c_char,
                mem::size_of::<ffi::z_stream>() as c_int,
            );
            assert!(result == ffi::Z_OK, "Failed to initialize compresser.");
            Compressor { stream: stream }
        }
    }

    pub fn compress(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<()> {
        self.stream_apply(input, output, |stream| unsafe {
            match ffi::deflate(stream, ffi::Z_SYNC_FLUSH) {
                ffi::Z_OK | ffi::Z_BUF_ERROR => {
                    if stream.avail_in == 0 && stream.avail_out > 0 {
                        Some(Ok(()))
                    } else {
                        None
                    }
                }
                code => Some(Err(Error::new(
                    Kind::Protocol,
                    format!("Failed to perform compression: {}", code),
                ))),
            }
        })
    }

    pub fn reset(&mut self) -> Result<()> {
        match unsafe { ffi::deflateReset(self.stream.as_mut()) } {
            ffi::Z_OK => Ok(()),
            code => Err(Error::new(
                Kind::Protocol,
                format!("Failed to reset compression context: {}", code),
            )),
        }
    }
}

impl Context for Compressor {
    fn stream(&mut self) -> &mut ffi::z_stream {
        self.stream.as_mut()
    }
}

impl Drop for Compressor {
    fn drop(&mut self) {
        match unsafe { ffi::deflateEnd(self.stream.as_mut()) } {
            ffi::Z_STREAM_ERROR => error!("Compression stream encountered bad state."),
            // Ignore discarded data error because we are raw
            ffi::Z_OK | ffi::Z_DATA_ERROR => trace!("Deallocated compression context."),
            code => error!("Bad zlib status encountered: {}", code),
        }
    }
}

pub struct Decompressor {
    stream: Box<ffi::z_stream>,
}

impl Decompressor {
    pub fn new(window_bits: i8) -> Decompressor {
        debug_assert!(window_bits >= 8, "Received too small window size.");
        debug_assert!(window_bits <= 15, "Received too large window size.");

        unsafe {
            let mut stream: Box<ffi::z_stream> = Box::new(MaybeUninit::uninit().assume_init());
            let result = ffi::inflateInit2_(
                stream.as_mut(),
                -window_bits as c_int,
                ZLIB_VERSION.as_ptr() as *const c_char,
                mem::size_of::<ffi::z_stream>() as c_int,
            );
            assert!(result == ffi::Z_OK, "Failed to initialize decompresser.");
            Decompressor { stream: stream }
        }
    }

    pub fn decompress(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<()> {
        self.stream_apply(input, output, |stream| unsafe {
            match ffi::inflate(stream, ffi::Z_SYNC_FLUSH) {
                ffi::Z_OK | ffi::Z_BUF_ERROR => {
                    if stream.avail_in == 0 && stream.avail_out > 0 {
                        Some(Ok(()))
                    } else {
                        None
                    }
                }
                code => Some(Err(Error::new(
                    Kind::Protocol,
                    format!("Failed to perform decompression: {}", code),
                ))),
            }
        })
    }

    pub fn reset(&mut self) -> Result<()> {
        match unsafe { ffi::inflateReset(self.stream.as_mut()) } {
            ffi::Z_OK => Ok(()),
            code => Err(Error::new(
                Kind::Protocol,
                format!("Failed to reset compression context: {}", code),
            )),
        }
    }
}

impl Context for Decompressor {
    fn stream(&mut self) -> &mut ffi::z_stream {
        self.stream.as_mut()
    }
}

impl Drop for Decompressor {
    fn drop(&mut self) {
        match unsafe { ffi::inflateEnd(self.stream.as_mut()) } {
            ffi::Z_STREAM_ERROR => error!("Decompression stream encountered bad state."),
            ffi::Z_OK => trace!("Deallocated decompression context."),
            code => error!("Bad zlib status encountered: {}", code),
        }
    }
}

mod test {
    #![allow(unused_imports, unused_variables, dead_code)]
    use super::*;

    fn as_hex(s: &[u8]) {
        for byte in s {
            print!("0x{:x} ", byte);
        }
        print!("\n");
    }

    #[test]
    fn round_trip() {
        for i in 9..16 {
            let data = "HI THERE THIS IS some data. これはデータだよ。".as_bytes();
            let mut compressed = Vec::with_capacity(data.len());
            let mut decompressed = Vec::with_capacity(data.len());

            let com = Compressor::new(i);
            let mut moved_com = com;

            moved_com
                .compress(&data, &mut compressed)
                .expect("Failed to compress data.");

            let dec = Decompressor::new(i);
            let mut moved_dec = dec;

            moved_dec
                .decompress(&compressed, &mut decompressed)
                .expect("Failed to decompress data.");

            assert_eq!(data, &decompressed[..]);
        }
    }

    #[test]
    fn reset() {
        let data1 = "HI THERE 直子さん".as_bytes();
        let data2 = "HI THERE 人太郎".as_bytes();
        let mut compressed1 = Vec::with_capacity(data1.len());
        let mut compressed2 = Vec::with_capacity(data2.len());
        let mut compressed2_ind = Vec::with_capacity(data2.len());

        let mut decompressed1 = Vec::with_capacity(data1.len());
        let mut decompressed2 = Vec::with_capacity(data2.len());
        let mut decompressed2_ind = Vec::with_capacity(data2.len());

        let mut com = Compressor::new(9);

        com.compress(&data1, &mut compressed1).unwrap();
        com.compress(&data2, &mut compressed2).unwrap();
        com.reset().unwrap();
        com.compress(&data2, &mut compressed2_ind).unwrap();

        let mut dec = Decompressor::new(9);

        dec.decompress(&compressed1, &mut decompressed1).unwrap();
        dec.decompress(&compressed2, &mut decompressed2).unwrap();
        dec.reset().unwrap();
        dec.decompress(&compressed2_ind, &mut decompressed2_ind)
            .unwrap();

        assert_eq!(data1, &decompressed1[..]);
        assert_eq!(data2, &decompressed2[..]);
        assert_eq!(data2, &decompressed2_ind[..]);
        assert!(compressed2 != compressed2_ind);
        assert!(compressed2.len() < compressed2_ind.len());
    }
}
<|endoftext|>


<commit_before>use super::*;
use log;

#[test]
fn merge_maps_both_empty() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 0);
}

#[test]
fn merge_maps_first_empty() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    map2.insert("test".to_string(), "value".to_string());

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    assert_eq!(output.get("test").unwrap(), &"value".to_string());
}

#[test]
fn merge_maps_second_empty() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    map1.insert("test".to_string(), "value".to_string());

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    assert_eq!(output.get("test").unwrap(), &"value".to_string());
}

#[test]
fn merge_maps_both_with_values() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    map1.insert("test1".to_string(), "value1".to_string());
    map2.insert("test21".to_string(), "value21".to_string());
    map2.insert("test22".to_string(), "value22".to_string());

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 3);
    assert_eq!(output.get("test1").unwrap(), &"value1".to_string());
    assert_eq!(output.get("test21").unwrap(), &"value21".to_string());
    assert_eq!(output.get("test22").unwrap(), &"value22".to_string());
}

#[test]
fn merge_tasks_both_empty() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 0);
}

#[test]
fn merge_tasks_first_empty() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task = Task::new();
    task.install_crate = Some("my crate".to_string());
    task.command = Some("test".to_string());

    map2.insert("test".to_string(), task);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    let task = output.get("test").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());
}

#[test]
fn merge_tasks_second_empty() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task = Task::new();
    task.install_crate = Some("my crate".to_string());
    task.command = Some("test".to_string());

    map1.insert("test".to_string(), task);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    let task = output.get("test").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());
}

#[test]
fn merge_tasks_both_with_values() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task1 = Task::new();
    task1.install_crate = Some("my crate".to_string());
    task1.command = Some("test".to_string());

    map1.insert("test".to_string(), task1);

    let mut task2 = Task::new();
    task2.command = Some("test".to_string());

    map2.insert("test2".to_string(), task2);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 2);

    let mut task = output.get("test").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());

    task = output.get("test2").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_none());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());
}

#[test]
fn merge_tasks_extend_task() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task1 = Task::new();
    task1.disabled = Some(false);
    task1.install_crate = Some("my crate".to_string());
    task1.command = Some("test1".to_string());

    map1.insert("test".to_string(), task1);

    let mut task2 = Task::new();
    task2.disabled = Some(true);
    task2.command = Some("test2".to_string());

    map2.insert("test".to_string(), task2);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);

    let task = output.get("test").unwrap();
    assert!(task.disabled.is_some());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());

    let task_clone = task.clone();
    assert!(task_clone.disabled.unwrap());
    assert_eq!(task_clone.install_crate.unwrap(), "my crate");
    assert_eq!(task_clone.command.unwrap(), "test2");
}

#[test]
fn load_external_descriptor_no_file() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "bad_file.toml2", &logger);

    assert!(config.env.is_none());
    assert!(config.tasks.is_none());
}

#[test]
fn load_external_descriptor_simple_file() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "./examples/alias.toml", &logger);

    assert!(config.env.is_none());
    assert!(config.tasks.is_some());

    let tasks = config.tasks.unwrap();
    let test_task = tasks.get("D2").unwrap();
    let alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D");
}

#[test]
fn load_external_descriptor_extending_file() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "examples/extending.toml", &logger);

    assert!(config.env.is_some());
    assert!(config.tasks.is_some());

    assert_eq!(config.env.unwrap().len(), 0);

    let tasks = config.tasks.unwrap();
    let mut test_task = tasks.get("D2").unwrap();
    let mut alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D");

    test_task = tasks.get("extended").unwrap();
    alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D2");
}

#[test]
fn load_external_descriptor_extending_file_sub_folder() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "examples/files/extending.toml", &logger);

    assert!(config.env.is_some());
    assert!(config.tasks.is_some());

    assert_eq!(config.env.unwrap().len(), 0);

    let tasks = config.tasks.unwrap();
    let mut test_task = tasks.get("D2").unwrap();
    let mut alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D");

    test_task = tasks.get("extended").unwrap();
    alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D2");

    test_task = tasks.get("extended2").unwrap();
    alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "extended");
}
<commit_msg>update tests<commit_after>use super::*;
use log;

#[test]
fn merge_maps_both_empty() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 0);
}

#[test]
fn merge_maps_first_empty() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    map2.insert("test".to_string(), "value".to_string());

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    assert_eq!(output.get("test").unwrap(), &"value".to_string());
}

#[test]
fn merge_maps_second_empty() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    map1.insert("test".to_string(), "value".to_string());

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    assert_eq!(output.get("test").unwrap(), &"value".to_string());
}

#[test]
fn merge_maps_both_with_values() {
    let mut map1 = HashMap::<String, String>::new();
    let mut map2 = HashMap::<String, String>::new();

    map1.insert("test1".to_string(), "value1".to_string());
    map2.insert("test21".to_string(), "value21".to_string());
    map2.insert("test22".to_string(), "value22".to_string());

    let output = merge_maps(&mut map1, &mut map2);
    assert_eq!(output.len(), 3);
    assert_eq!(output.get("test1").unwrap(), &"value1".to_string());
    assert_eq!(output.get("test21").unwrap(), &"value21".to_string());
    assert_eq!(output.get("test22").unwrap(), &"value22".to_string());
}

#[test]
fn merge_tasks_both_empty() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 0);
}

#[test]
fn merge_tasks_first_empty() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task = Task::new();
    task.install_crate = Some("my crate".to_string());
    task.command = Some("test".to_string());

    map2.insert("test".to_string(), task);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    let task = output.get("test").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());
}

#[test]
fn merge_tasks_second_empty() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task = Task::new();
    task.install_crate = Some("my crate".to_string());
    task.command = Some("test".to_string());

    map1.insert("test".to_string(), task);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);
    let task = output.get("test").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());
}

#[test]
fn merge_tasks_both_with_values() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task1 = Task::new();
    task1.install_crate = Some("my crate".to_string());
    task1.command = Some("test".to_string());

    map1.insert("test".to_string(), task1);

    let mut task2 = Task::new();
    task2.command = Some("test".to_string());

    map2.insert("test2".to_string(), task2);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 2);

    let mut task = output.get("test").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());

    task = output.get("test2").unwrap();
    assert!(task.disabled.is_none());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_none());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());
}

#[test]
fn merge_tasks_extend_task() {
    let mut map1 = HashMap::<String, Task>::new();
    let mut map2 = HashMap::<String, Task>::new();

    let mut task1 = Task::new();
    task1.disabled = Some(false);
    task1.install_crate = Some("my crate".to_string());
    task1.command = Some("test1".to_string());

    map1.insert("test".to_string(), task1);

    let mut task2 = Task::new();
    task2.disabled = Some(true);
    task2.command = Some("test2".to_string());

    map2.insert("test".to_string(), task2);

    let output = merge_tasks(&mut map1, &mut map2);
    assert_eq!(output.len(), 1);

    let task = output.get("test").unwrap();
    assert!(task.disabled.is_some());
    assert!(task.alias.is_none());
    assert!(task.linux_alias.is_none());
    assert!(task.windows_alias.is_none());
    assert!(task.mac_alias.is_none());
    assert!(task.install_crate.is_some());
    assert!(task.install_script.is_none());
    assert!(task.command.is_some());
    assert!(task.args.is_none());
    assert!(task.script.is_none());
    assert!(task.dependencies.is_none());
    assert!(task.linux.is_none());
    assert!(task.windows.is_none());
    assert!(task.mac.is_none());

    let task_clone = task.clone();
    assert!(task_clone.disabled.unwrap());
    assert_eq!(task_clone.install_crate.unwrap(), "my crate");
    assert_eq!(task_clone.command.unwrap(), "test2");
}

#[test]
fn load_external_descriptor_no_file() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "bad_file.toml2", &logger);

    assert!(config.env.is_none());
    assert!(config.tasks.is_none());
}

#[test]
fn load_external_descriptor_simple_file() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "./examples/alias.toml", &logger);

    assert!(config.env.is_none());
    assert!(config.tasks.is_some());

    let tasks = config.tasks.unwrap();
    let test_task = tasks.get("D2").unwrap();
    let alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D");
}

#[test]
fn load_external_descriptor_extending_file() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "examples/extending.toml", &logger);

    assert!(config.env.is_some());
    assert!(config.tasks.is_some());

    assert_eq!(config.env.unwrap().len(), 0);

    let tasks = config.tasks.unwrap();
    let mut test_task = tasks.get("D2").unwrap();
    let mut alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D");

    test_task = tasks.get("extended").unwrap();
    alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D2");
}

#[test]
#[cfg(target_os = "linux")]
fn load_external_descriptor_extending_file_sub_folder() {
    let logger = log::create("error");
    let config = load_external_descriptor(".", "examples/files/extending.toml", &logger);

    assert!(config.env.is_some());
    assert!(config.tasks.is_some());

    assert_eq!(config.env.unwrap().len(), 0);

    let tasks = config.tasks.unwrap();
    let mut test_task = tasks.get("D2").unwrap();
    let mut alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D");

    test_task = tasks.get("extended").unwrap();
    alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "D2");

    test_task = tasks.get("extended2").unwrap();
    alias = test_task.alias.clone();
    assert_eq!(alias.unwrap(), "extended");
}
<|endoftext|>


<commit_before>use beacon_chain::{BeaconChain, BeaconChainTypes};
use environment::RuntimeContext;
use exit_future::Signal;
use futures::{Future, Stream};
use network::Service as NetworkService;
use parking_lot::Mutex;
use slog::{debug, error, info, warn};
use slot_clock::SlotClock;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::timer::Interval;
use types::{EthSpec, Slot};

/// Create a warning log whenever the peer count is at or below this value.
pub const WARN_PEER_COUNT: usize = 1;

const SECS_PER_MINUTE: u64 = 60;
const SECS_PER_HOUR: u64 = 3600;
const SECS_PER_DAY: u64 = 86400; // non-leap
const SECS_PER_WEEK: u64 = 604800; // non-leap
const DAYS_PER_WEEK: u64 = 7;
const HOURS_PER_DAY: u64 = 24;
const MINUTES_PER_HOUR: u64 = 60;

/// How long to wait for the lock on `network.libp2p_service()` before we give up.
const LIBP2P_LOCK_TIMEOUT: Duration = Duration::from_millis(50);

/// Spawns a notifier service which periodically logs information about the node.
pub fn spawn_notifier<T: BeaconChainTypes>(
    context: RuntimeContext<T::EthSpec>,
    beacon_chain: Arc<BeaconChain<T>>,
    network: Arc<NetworkService<T>>,
    milliseconds_per_slot: u64,
) -> Result<Signal, String> {
    let log_1 = context.log.clone();
    let log_2 = context.log.clone();

    let slot_duration = Duration::from_millis(milliseconds_per_slot);
    let duration_to_next_slot = beacon_chain
        .slot_clock
        .duration_to_next_slot()
        .ok_or_else(|| "slot_notifier unable to determine time to next slot")?;

    // Run this half way through each slot.
    let start_instant = Instant::now() + duration_to_next_slot + (slot_duration / 2);

    // Run this each slot.
    let interval_duration = slot_duration;

    let previous_head_slot = Mutex::new(Slot::new(0));

    let interval_future = Interval::new(start_instant, interval_duration)
        .map_err(
            move |e| error!(log_1, "Slot notifier timer failed"; "error" => format!("{:?}", e)),
        )
        .for_each(move |_| {
            let log = log_2.clone();

            let connected_peer_count = if let Some(libp2p) = network
                .libp2p_service()
                .try_lock_until(Instant::now() + LIBP2P_LOCK_TIMEOUT)
            {
                libp2p.swarm.connected_peers()
            } else {
                // Use max_value here and we'll print something pretty later.
                usize::max_value()
            };

            let head = beacon_chain.head();

            let head_slot = head.beacon_block.slot;
            let head_epoch = head_slot.epoch(T::EthSpec::slots_per_epoch());
            let current_slot = beacon_chain.slot().map_err(|e| {
                error!(
                    log,
                    "Unable to read current slot";
                    "error" => format!("{:?}", e)
                )
            })?;
            let current_epoch = current_slot.epoch(T::EthSpec::slots_per_epoch());
            let finalized_epoch = head.beacon_state.finalized_checkpoint.epoch;
            let finalized_root = head.beacon_state.finalized_checkpoint.root;
            let head_root = head.beacon_block_root;

            let mut previous_head_slot = previous_head_slot.lock();

            // The next two lines take advantage of saturating subtraction on `Slot`.
            let head_distance = current_slot - head_slot;
            let slots_since_last_update = head_slot - *previous_head_slot;

            *previous_head_slot = head_slot;

            if connected_peer_count <= WARN_PEER_COUNT {
                warn!(log, "Low peer count"; "peer_count" => peer_count_pretty(connected_peer_count));
            }

            debug!(
                log,
                "Slot timer";
                "peers" => peer_count_pretty(connected_peer_count),
                "finalized_root" => format!("{}", finalized_root),
                "finalized_epoch" => finalized_epoch,
                "head_block" => format!("{}", head_root),
                "head_slot" => head_slot,
                "current_slot" => current_slot,
            );

            if head_epoch + 1 < current_epoch {
                let distance = format!(
                    "{} slots ({})",
                    head_distance.as_u64(),
                    slot_distance_pretty(head_distance, slot_duration)
                );

                info!(
                    log,
                    "Syncing";
                    "peers" => peer_count_pretty(connected_peer_count),
                    "speed" => sync_rate_pretty(slots_since_last_update, interval_duration.as_secs()),
                    "distance" => distance
                );

                return Ok(());
            };

            macro_rules! not_quite_synced_log {
                ($message: expr) => {
                    info!(
                        log_2,
                        $message;
                        "peers" => peer_count_pretty(connected_peer_count),
                        "finalized_root" => format!("{}", finalized_root),
                        "finalized_epoch" => finalized_epoch,
                        "head_slot" => head_slot,
                        "current_slot" => current_slot,
                    );
                }
            }

            if head_epoch + 1 == current_epoch {
                not_quite_synced_log!("Synced to previous epoch")
            } else if head_slot != current_slot {
                not_quite_synced_log!("Synced to current epoch")
            } else {
                info!(
                    log_2,
                    "Synced";
                    "peers" => peer_count_pretty(connected_peer_count),
                    "finalized_root" => format!("{}", finalized_root),
                    "finalized_epoch" => finalized_epoch,
                    "epoch" => current_epoch,
                    "slot" => current_slot,
                );
            };

            Ok(())
        });

    let (exit_signal, exit) = exit_future::signal();
    context
        .executor
        .spawn(exit.until(interval_future).map(|_| ()));

    Ok(exit_signal)
}

/// Returns the peer count, returning something helpful if it's `usize::max_value` (effectively a
/// `None` value).
fn peer_count_pretty(peer_count: usize) -> String {
    if peer_count == usize::max_value() {
        String::from("--")
    } else {
        format!("{}", peer_count)
    }
}

/// Returns a nicely formated string describing the rate of slot imports per second.
fn sync_rate_pretty(slots_since_last_update: Slot, update_interval_secs: u64) -> String {
    if update_interval_secs == 0 {
        return "Error".into();
    }

    if slots_since_last_update == 0 {
        "No progress".into()
    } else {
        let distance = f64::from(slots_since_last_update.as_u64() as u32);
        let time = f64::from(update_interval_secs as u32);

        format!("{:.2} slots/sec", distance / time)
    }
}

/// Returns a nicely formatted string describing the `slot_span` in terms of weeks, days, hours
/// and/or minutes.
fn slot_distance_pretty(slot_span: Slot, slot_duration: Duration) -> String {
    if slot_duration == Duration::from_secs(0) {
        return String::from("Unknown");
    }

    let secs = (slot_duration * slot_span.as_u64() as u32).as_secs();

    let weeks = secs / SECS_PER_WEEK;
    let days = secs / SECS_PER_DAY;
    let hours = secs / SECS_PER_HOUR;
    let minutes = secs / SECS_PER_MINUTE;

    if weeks > 0 {
        format!("{} weeks {} days", weeks, days % DAYS_PER_WEEK)
    } else if days > 0 {
        format!("{} days {} hrs", days, hours % HOURS_PER_DAY)
    } else if hours > 0 {
        format!("{} hrs {} mins", hours, minutes % MINUTES_PER_HOUR)
    } else {
        format!("{} mins", minutes)
    }
}
<commit_msg>Improve sync speed indicator (#683)<commit_after>use beacon_chain::{BeaconChain, BeaconChainTypes};
use environment::RuntimeContext;
use exit_future::Signal;
use futures::{Future, Stream};
use network::Service as NetworkService;
use parking_lot::Mutex;
use slog::{debug, error, info, warn};
use slot_clock::SlotClock;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::timer::Interval;
use types::{EthSpec, Slot};

/// Create a warning log whenever the peer count is at or below this value.
pub const WARN_PEER_COUNT: usize = 1;

const SECS_PER_MINUTE: f64 = 60.0;
const SECS_PER_HOUR: f64 = 3600.0;
const SECS_PER_DAY: f64 = 86400.0; // non-leap
const SECS_PER_WEEK: f64 = 604800.0; // non-leap
const DAYS_PER_WEEK: f64 = 7.0;
const HOURS_PER_DAY: f64 = 24.0;
const MINUTES_PER_HOUR: f64 = 60.0;

/// How long to wait for the lock on `network.libp2p_service()` before we give up.
const LIBP2P_LOCK_TIMEOUT: Duration = Duration::from_millis(50);

/// The number of historical observations that should be used to determine the average sync time.
const SPEEDO_OBSERVATIONS: usize = 4;

/// Spawns a notifier service which periodically logs information about the node.
pub fn spawn_notifier<T: BeaconChainTypes>(
    context: RuntimeContext<T::EthSpec>,
    beacon_chain: Arc<BeaconChain<T>>,
    network: Arc<NetworkService<T>>,
    milliseconds_per_slot: u64,
) -> Result<Signal, String> {
    let log_1 = context.log.clone();
    let log_2 = context.log.clone();

    let slot_duration = Duration::from_millis(milliseconds_per_slot);
    let duration_to_next_slot = beacon_chain
        .slot_clock
        .duration_to_next_slot()
        .ok_or_else(|| "slot_notifier unable to determine time to next slot")?;

    // Run this half way through each slot.
    let start_instant = Instant::now() + duration_to_next_slot + (slot_duration / 2);

    // Run this each slot.
    let interval_duration = slot_duration;

    let speedo = Mutex::new(Speedo::default());

    let interval_future = Interval::new(start_instant, interval_duration)
        .map_err(
            move |e| error!(log_1, "Slot notifier timer failed"; "error" => format!("{:?}", e)),
        )
        .for_each(move |_| {
            let log = log_2.clone();

            let connected_peer_count = if let Some(libp2p) = network
                .libp2p_service()
                .try_lock_until(Instant::now() + LIBP2P_LOCK_TIMEOUT)
            {
                libp2p.swarm.connected_peers()
            } else {
                // Use max_value here and we'll print something pretty later.
                usize::max_value()
            };

            let head = beacon_chain.head();

            let head_slot = head.beacon_block.slot;
            let head_epoch = head_slot.epoch(T::EthSpec::slots_per_epoch());
            let current_slot = beacon_chain.slot().map_err(|e| {
                error!(
                    log,
                    "Unable to read current slot";
                    "error" => format!("{:?}", e)
                )
            })?;
            let current_epoch = current_slot.epoch(T::EthSpec::slots_per_epoch());
            let finalized_epoch = head.beacon_state.finalized_checkpoint.epoch;
            let finalized_root = head.beacon_state.finalized_checkpoint.root;
            let head_root = head.beacon_block_root;

            let mut speedo = speedo.lock();
            speedo.observe(head_slot, Instant::now());

            // The next two lines take advantage of saturating subtraction on `Slot`.
            let head_distance = current_slot - head_slot;

            if connected_peer_count <= WARN_PEER_COUNT {
                warn!(log, "Low peer count"; "peer_count" => peer_count_pretty(connected_peer_count));
            }

            debug!(
                log,
                "Slot timer";
                "peers" => peer_count_pretty(connected_peer_count),
                "finalized_root" => format!("{}", finalized_root),
                "finalized_epoch" => finalized_epoch,
                "head_block" => format!("{}", head_root),
                "head_slot" => head_slot,
                "current_slot" => current_slot,
            );

            if head_epoch + 1 < current_epoch {
                let distance = format!(
                    "{} slots ({})",
                    head_distance.as_u64(),
                    slot_distance_pretty(head_distance, slot_duration)
                );

                info!(
                    log,
                    "Syncing";
                    "peers" => peer_count_pretty(connected_peer_count),
                    "est_time" => estimated_time_pretty(speedo.estimated_time_till_slot(current_slot)),
                    "speed" => sync_speed_pretty(speedo.slots_per_second()),
                    "distance" => distance
                );

                return Ok(());
            };

            macro_rules! not_quite_synced_log {
                ($message: expr) => {
                    info!(
                        log_2,
                        $message;
                        "peers" => peer_count_pretty(connected_peer_count),
                        "finalized_root" => format!("{}", finalized_root),
                        "finalized_epoch" => finalized_epoch,
                        "head_slot" => head_slot,
                        "current_slot" => current_slot,
                    );
                }
            }

            if head_epoch + 1 == current_epoch {
                not_quite_synced_log!("Synced to previous epoch")
            } else if head_slot != current_slot {
                not_quite_synced_log!("Synced to current epoch")
            } else {
                info!(
                    log_2,
                    "Synced";
                    "peers" => peer_count_pretty(connected_peer_count),
                    "finalized_root" => format!("{}", finalized_root),
                    "finalized_epoch" => finalized_epoch,
                    "epoch" => current_epoch,
                    "slot" => current_slot,
                );
            };

            Ok(())
        });

    let (exit_signal, exit) = exit_future::signal();
    context
        .executor
        .spawn(exit.until(interval_future).map(|_| ()));

    Ok(exit_signal)
}

/// Returns the peer count, returning something helpful if it's `usize::max_value` (effectively a
/// `None` value).
fn peer_count_pretty(peer_count: usize) -> String {
    if peer_count == usize::max_value() {
        String::from("--")
    } else {
        format!("{}", peer_count)
    }
}

/// Returns a nicely formatted string describing the rate of slot imports per second.
fn sync_speed_pretty(slots_per_second: Option<f64>) -> String {
    if let Some(slots_per_second) = slots_per_second {
        format!("{:.2} slots/sec", slots_per_second)
    } else {
        "--".into()
    }
}

/// Returns a nicely formatted string how long will we reach the target slot.
fn estimated_time_pretty(seconds_till_slot: Option<f64>) -> String {
    if let Some(seconds_till_slot) = seconds_till_slot {
        seconds_pretty(seconds_till_slot)
    } else {
        "--".into()
    }
}

/// Returns a nicely formatted string describing the `slot_span` in terms of weeks, days, hours
/// and/or minutes.
fn slot_distance_pretty(slot_span: Slot, slot_duration: Duration) -> String {
    if slot_duration == Duration::from_secs(0) {
        return String::from("Unknown");
    }

    let secs = (slot_duration * slot_span.as_u64() as u32).as_secs();
    seconds_pretty(secs as f64)
}

/// Returns a nicely formatted string describing the `slot_span` in terms of weeks, days, hours
/// and/or minutes.
fn seconds_pretty(secs: f64) -> String {
    if secs <= 0.0 {
        return "--".into();
    }

    let weeks = secs / SECS_PER_WEEK;
    let days = secs / SECS_PER_DAY;
    let hours = secs / SECS_PER_HOUR;
    let minutes = secs / SECS_PER_MINUTE;

    if weeks.floor() > 0.0 {
        format!(
            "{:.0} weeks {:.0} days",
            weeks,
            (days % DAYS_PER_WEEK).round()
        )
    } else if days.floor() > 0.0 {
        format!(
            "{:.0} days {:.0} hrs",
            days,
            (hours % HOURS_PER_DAY).round()
        )
    } else if hours.floor() > 0.0 {
        format!(
            "{:.0} hrs {:.0} mins",
            hours,
            (minutes % MINUTES_PER_HOUR).round()
        )
    } else {
        format!("{:.0} mins", minutes.round())
    }
}

/// "Speedo" is Australian for speedometer. This struct observes syncing times.
#[derive(Default)]
pub struct Speedo(Vec<(Slot, Instant)>);

impl Speedo {
    /// Observe that we were at some `slot` at the given `instant`.
    pub fn observe(&mut self, slot: Slot, instant: Instant) {
        if self.0.len() > SPEEDO_OBSERVATIONS {
            self.0.remove(0);
        }

        self.0.push((slot, instant));
    }

    /// Returns the average of the speeds between each observation.
    ///
    /// Does not gracefully handle slots that are above `u32::max_value()`.
    pub fn slots_per_second(&self) -> Option<f64> {
        let speeds = self
            .0
            .windows(2)
            .filter_map(|windows| {
                let (slot_a, instant_a) = windows[0];
                let (slot_b, instant_b) = windows[1];

                // Taking advantage of saturating subtraction on `Slot`.
                let distance = f64::from((slot_b - slot_a).as_u64() as u32);

                let seconds = f64::from((instant_b - instant_a).as_millis() as u32) / 1_000.0;

                if seconds > 0.0 {
                    Some(distance / seconds)
                } else {
                    None
                }
            })
            .collect::<Vec<f64>>();

        let count = speeds.len();
        let sum: f64 = speeds.iter().sum();

        if count > 0 {
            Some(sum / f64::from(count as u32))
        } else {
            None
        }
    }

    /// Returns the time we should reach the given `slot`, judging by the latest observation and
    /// historical average syncing time.
    ///
    /// Returns `None` if the slot is prior to our latest observed slot or we have not made any
    /// observations.
    pub fn estimated_time_till_slot(&self, target_slot: Slot) -> Option<f64> {
        let (prev_slot, _) = self.0.last()?;
        let slots_per_second = self.slots_per_second()?;

        if target_slot > *prev_slot {
            let distance = (target_slot - *prev_slot).as_u64() as f64;
            Some(distance / slots_per_second)
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use] extern crate lazy_static;
#[macro_use] extern crate log;
extern crate regex;

pub mod ismatch;
pub mod key_value_split;
pub mod trace;
pub mod variants;
<commit_msg>Enable lints<commit_after>#![deny(
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_must_use,
    unused_mut,
    unused_qualifications,
    while_true,
)]

#[macro_use] extern crate lazy_static;
#[macro_use] extern crate log;
extern crate regex;

pub mod ismatch;
pub mod key_value_split;
pub mod trace;
pub mod variants;
<|endoftext|>


<commit_before>extern crate octobuild;
extern crate petgraph;
extern crate tempdir;
extern crate num_cpus;

use octobuild::common::BuildTask;
use octobuild::cache::Cache;
use octobuild::wincmd;
use octobuild::xg;
use octobuild::version;
use octobuild::vs::compiler::VsCompiler;
use octobuild::compiler::*;

use petgraph::{Graph, EdgeDirection};
use petgraph::graph::NodeIndex;
use tempdir::TempDir;

use std::fs::File;
use std::env;
use std::io::{BufReader, Error, ErrorKind, Write};
use std::io;
use std::iter::FromIterator;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::process;
use std::thread;

#[derive(Debug)]
struct TaskMessage {
	index: NodeIndex,
	task: BuildTask
}

#[derive(Debug)]
struct ResultMessage {
	index: NodeIndex,
	task: BuildTask,
	worker: usize,
	result: Result<OutputInfo, Error>
}

fn main() {
	println!("XGConsole ({}):", version::full_version());
	let args = Vec::from_iter(env::args());
	for arg in args.iter() {
		println!("  {}", arg);
	}
	process::exit(match execute(&args[1..]) {
		Ok(result) => {
			match result {
				Some(r) => r,
				None => 501
			}
		}
		Err(e) => {
			println!("FATAL ERROR: {:?}", e);
			500
		}
	})
}

fn execute(args: &[String]) -> Result<Option<i32>, Error> {
	let cache = Cache::new();
	let temp_dir = try! (TempDir::new("octobuild"));
	for arg in args.iter() {
		if arg.starts_with("/") {continue}

		let (tx_result, rx_result): (Sender<ResultMessage>, Receiver<ResultMessage>) = channel();
		let (tx_task, rx_task): (Sender<TaskMessage>, Receiver<TaskMessage>) = channel();

		let mutex_rx_task = create_threads(rx_task, tx_result, num_cpus::get(), |worker_id:usize| {
			let temp_path = temp_dir.path().to_path_buf();
			let temp_cache = cache.clone();
			move |task:TaskMessage| -> ResultMessage {
				execute_task(&temp_cache, &temp_path, worker_id, task)
			}
		});

		let path = Path::new(arg);
		let file = try! (File::open(&path));
		let graph = try! (xg::parser::parse(BufReader::new(file)));
		let validated_graph = try! (validate_graph(graph));
		match try! (execute_graph(&validated_graph, tx_task, mutex_rx_task, rx_result)) {
			Some(v) if v == 0 => {}
			v => {return Ok(v)}
		}
	}
	cache.cleanup(16 * 1024 * 1024 * 1024);
	Ok(Some(0))
}

fn create_threads<R: 'static + Send, T: 'static + Send, Worker:'static + Fn(T) -> R + Send, Factory:Fn(usize) -> Worker>(rx_task: Receiver<T>, tx_result: Sender<R>, num_cpus: usize, factory: Factory) ->  Arc<Mutex<Receiver<T>>> {
	let mutex_rx_task = Arc::new(Mutex::new(rx_task));
	for cpu_id in 0..num_cpus {
		let local_rx_task = mutex_rx_task.clone();
		let local_tx_result = tx_result.clone();
		let worker = factory(cpu_id);
		thread::spawn(move || {
			loop {
				let task: T;
				match local_rx_task.lock().unwrap().recv() {
					Ok(v) => {task = v;}
					Err(_) => {break;}
				}
				match local_tx_result.send(worker(task)) {
					Ok(_) => {}
					Err(_) => {break;}
				}
			}
		});
	}
	mutex_rx_task
}

fn validate_graph(graph: Graph<BuildTask, ()>) -> Result<Graph<BuildTask, ()>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	let mut queue:Vec<NodeIndex> = Vec::new();
	for index in 0 .. graph.node_count() {
		completed.push(false);
		queue.push(NodeIndex::new(index));
	}
	let mut count:usize = 0;
	let mut i:usize = 0;
	while i < queue.len() {
		let index = queue[i];
		if (!completed[index.index()]) && (is_ready(&graph, &completed, &index)) {
			completed[index.index()] = true;
			for neighbor in graph.neighbors_directed(index, EdgeDirection::Incoming) {
				queue.push(neighbor);
			}
			count += 1;
			if count ==completed.len() {
				return Ok(graph);
			}
		}
		i = i + 1;
	}
	Err(Error::new(ErrorKind::InvalidInput, "Found cycles in build dependencies"))
}

fn execute_task(cache: &Cache, temp_dir: &Path, worker: usize, message: TaskMessage) -> ResultMessage {
	let args = wincmd::expand_args(&message.task.args, &|name:&str|->Option<String>{env::var(name).ok()});
	let output = execute_compiler(cache, temp_dir, &message.task, &args);
	ResultMessage {
		index: message.index,
		task: message.task,
		worker: worker,
		result: output,
	}
}

fn execute_compiler(cache: &Cache, temp_dir: &Path, task: &BuildTask, args: &[String]) -> Result<OutputInfo, Error> {
	let command = CommandInfo {
		program: Path::new(&task.exec).to_path_buf(),
		current_dir: Some(Path::new(&task.working_dir).to_path_buf()),
		env: task.env.clone(),
	};
	if Path::new(&task.exec).ends_with("cl.exe") {
		let compiler = VsCompiler::new(cache, temp_dir);
		compiler.compile(command, args)
	} else {
		command.to_command()
			.args(&args)
			.output()
			.map(|o| OutputInfo::new(o))
	}
}

fn execute_graph(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, mutex_rx_task: Arc<Mutex<Receiver<TaskMessage>>>, rx_result: Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	// Run all tasks.
	let result = execute_until_failed(graph, tx_task, &rx_result);
	// Cleanup task queue.
	for _ in mutex_rx_task.lock().unwrap().iter() {
	}
	// Wait for in progress task completion.
	for _ in rx_result.iter() {
	}
	result
}

fn execute_until_failed(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, rx_result: &Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	for _ in 0 .. graph.node_count() {
		completed.push(false);
	}
	for index in graph.without_edges(EdgeDirection::Outgoing) {
		try! (tx_task.send(TaskMessage{
			index: index,
			task: graph.node_weight(index).unwrap().clone(),
		}).map_err(|e| Error::new(ErrorKind::Other, e)));
	}

	let mut count:usize = 0;
	for message in rx_result.iter() {
		assert!(!completed[message.index.index()]);
		count += 1;
		println!("#{} {}/{}: {}", message.worker, count, completed.len(), message.task.title);
		let result = try! (message.result);
		try! (io::stdout().write_all(&result.stdout));
		try! (io::stderr().write_all(&result.stderr));
		if !result.success() {
			return Ok(result.status);
		}
		completed[message.index.index()] = true;

		for source in graph.neighbors_directed(message.index, EdgeDirection::Incoming) {
			if is_ready(graph, &completed, &source) {
				try! (tx_task.send(TaskMessage{
					index: source,
					task: graph.node_weight(source).unwrap().clone(),
				}).map_err(|e| Error::new(ErrorKind::Other, e)));
			}
		}

		if count == completed.len() {
			return Ok(Some(0));
		}
	}
	panic! ("Unexpected end of result pipe");
}

fn is_ready(graph: &Graph<BuildTask, ()>, completed: &Vec<bool>, source: &NodeIndex) -> bool {
	for neighbor in graph.neighbors_directed(*source, EdgeDirection::Outgoing) {
		if !completed[neighbor.index()]{
			return false
		}
	}
	true
}
<commit_msg>Fix hello xgConsole message<commit_after>extern crate octobuild;
extern crate petgraph;
extern crate tempdir;
extern crate num_cpus;

use octobuild::common::BuildTask;
use octobuild::cache::Cache;
use octobuild::wincmd;
use octobuild::xg;
use octobuild::version;
use octobuild::vs::compiler::VsCompiler;
use octobuild::compiler::*;

use petgraph::{Graph, EdgeDirection};
use petgraph::graph::NodeIndex;
use tempdir::TempDir;

use std::fs::File;
use std::env;
use std::io::{BufReader, Error, ErrorKind, Write};
use std::io;
use std::iter::FromIterator;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::process;
use std::thread;

#[derive(Debug)]
struct TaskMessage {
	index: NodeIndex,
	task: BuildTask
}

#[derive(Debug)]
struct ResultMessage {
	index: NodeIndex,
	task: BuildTask,
	worker: usize,
	result: Result<OutputInfo, Error>
}

fn main() {
	println!("xgConsole ({}):", version::full_version());
	let args = Vec::from_iter(env::args());
	for arg in args.iter() {
		println!("  {}", arg);
	}
	process::exit(match execute(&args[1..]) {
		Ok(result) => {
			match result {
				Some(r) => r,
				None => 501
			}
		}
		Err(e) => {
			println!("FATAL ERROR: {:?}", e);
			500
		}
	})
}

fn execute(args: &[String]) -> Result<Option<i32>, Error> {
	let cache = Cache::new();
	let temp_dir = try! (TempDir::new("octobuild"));
	for arg in args.iter() {
		if arg.starts_with("/") {continue}

		let (tx_result, rx_result): (Sender<ResultMessage>, Receiver<ResultMessage>) = channel();
		let (tx_task, rx_task): (Sender<TaskMessage>, Receiver<TaskMessage>) = channel();

		let mutex_rx_task = create_threads(rx_task, tx_result, num_cpus::get(), |worker_id:usize| {
			let temp_path = temp_dir.path().to_path_buf();
			let temp_cache = cache.clone();
			move |task:TaskMessage| -> ResultMessage {
				execute_task(&temp_cache, &temp_path, worker_id, task)
			}
		});

		let path = Path::new(arg);
		let file = try! (File::open(&path));
		let graph = try! (xg::parser::parse(BufReader::new(file)));
		let validated_graph = try! (validate_graph(graph));
		match try! (execute_graph(&validated_graph, tx_task, mutex_rx_task, rx_result)) {
			Some(v) if v == 0 => {}
			v => {return Ok(v)}
		}
	}
	cache.cleanup(16 * 1024 * 1024 * 1024);
	Ok(Some(0))
}

fn create_threads<R: 'static + Send, T: 'static + Send, Worker:'static + Fn(T) -> R + Send, Factory:Fn(usize) -> Worker>(rx_task: Receiver<T>, tx_result: Sender<R>, num_cpus: usize, factory: Factory) ->  Arc<Mutex<Receiver<T>>> {
	let mutex_rx_task = Arc::new(Mutex::new(rx_task));
	for cpu_id in 0..num_cpus {
		let local_rx_task = mutex_rx_task.clone();
		let local_tx_result = tx_result.clone();
		let worker = factory(cpu_id);
		thread::spawn(move || {
			loop {
				let task: T;
				match local_rx_task.lock().unwrap().recv() {
					Ok(v) => {task = v;}
					Err(_) => {break;}
				}
				match local_tx_result.send(worker(task)) {
					Ok(_) => {}
					Err(_) => {break;}
				}
			}
		});
	}
	mutex_rx_task
}

fn validate_graph(graph: Graph<BuildTask, ()>) -> Result<Graph<BuildTask, ()>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	let mut queue:Vec<NodeIndex> = Vec::new();
	for index in 0 .. graph.node_count() {
		completed.push(false);
		queue.push(NodeIndex::new(index));
	}
	let mut count:usize = 0;
	let mut i:usize = 0;
	while i < queue.len() {
		let index = queue[i];
		if (!completed[index.index()]) && (is_ready(&graph, &completed, &index)) {
			completed[index.index()] = true;
			for neighbor in graph.neighbors_directed(index, EdgeDirection::Incoming) {
				queue.push(neighbor);
			}
			count += 1;
			if count ==completed.len() {
				return Ok(graph);
			}
		}
		i = i + 1;
	}
	Err(Error::new(ErrorKind::InvalidInput, "Found cycles in build dependencies"))
}

fn execute_task(cache: &Cache, temp_dir: &Path, worker: usize, message: TaskMessage) -> ResultMessage {
	let args = wincmd::expand_args(&message.task.args, &|name:&str|->Option<String>{env::var(name).ok()});
	let output = execute_compiler(cache, temp_dir, &message.task, &args);
	ResultMessage {
		index: message.index,
		task: message.task,
		worker: worker,
		result: output,
	}
}

fn execute_compiler(cache: &Cache, temp_dir: &Path, task: &BuildTask, args: &[String]) -> Result<OutputInfo, Error> {
	let command = CommandInfo {
		program: Path::new(&task.exec).to_path_buf(),
		current_dir: Some(Path::new(&task.working_dir).to_path_buf()),
		env: task.env.clone(),
	};
	if Path::new(&task.exec).ends_with("cl.exe") {
		let compiler = VsCompiler::new(cache, temp_dir);
		compiler.compile(command, args)
	} else {
		command.to_command()
			.args(&args)
			.output()
			.map(|o| OutputInfo::new(o))
	}
}

fn execute_graph(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, mutex_rx_task: Arc<Mutex<Receiver<TaskMessage>>>, rx_result: Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	// Run all tasks.
	let result = execute_until_failed(graph, tx_task, &rx_result);
	// Cleanup task queue.
	for _ in mutex_rx_task.lock().unwrap().iter() {
	}
	// Wait for in progress task completion.
	for _ in rx_result.iter() {
	}
	result
}

fn execute_until_failed(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, rx_result: &Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	for _ in 0 .. graph.node_count() {
		completed.push(false);
	}
	for index in graph.without_edges(EdgeDirection::Outgoing) {
		try! (tx_task.send(TaskMessage{
			index: index,
			task: graph.node_weight(index).unwrap().clone(),
		}).map_err(|e| Error::new(ErrorKind::Other, e)));
	}

	let mut count:usize = 0;
	for message in rx_result.iter() {
		assert!(!completed[message.index.index()]);
		count += 1;
		println!("#{} {}/{}: {}", message.worker, count, completed.len(), message.task.title);
		let result = try! (message.result);
		try! (io::stdout().write_all(&result.stdout));
		try! (io::stderr().write_all(&result.stderr));
		if !result.success() {
			return Ok(result.status);
		}
		completed[message.index.index()] = true;

		for source in graph.neighbors_directed(message.index, EdgeDirection::Incoming) {
			if is_ready(graph, &completed, &source) {
				try! (tx_task.send(TaskMessage{
					index: source,
					task: graph.node_weight(source).unwrap().clone(),
				}).map_err(|e| Error::new(ErrorKind::Other, e)));
			}
		}

		if count == completed.len() {
			return Ok(Some(0));
		}
	}
	panic! ("Unexpected end of result pipe");
}

fn is_ready(graph: &Graph<BuildTask, ()>, completed: &Vec<bool>, source: &NodeIndex) -> bool {
	for neighbor in graph.neighbors_directed(*source, EdgeDirection::Outgoing) {
		if !completed[neighbor.index()]{
			return false
		}
	}
	true
}
<|endoftext|>


<commit_before>#![feature(exit_status)]
#![feature(io)]
#![feature(os)]
extern crate octobuild;
extern crate tempdir;

use octobuild::common::BuildTask;
use octobuild::cache::Cache;
use octobuild::wincmd;
use octobuild::xg;
use octobuild::graph::{Graph, NodeIndex, Node, EdgeIndex, Edge};
use octobuild::vs::compiler::VsCompiler;
use octobuild::compiler::*;

use tempdir::TempDir;

use std::fs::File;
use std::env;
use std::io::{BufReader, Error, ErrorKind, Write};
use std::io;
use std::iter::FromIterator;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::thread;

#[derive(Debug)]
struct TaskMessage {
	index: NodeIndex,
	task: BuildTask
}

#[derive(Debug)]
struct ResultMessage {
	index: NodeIndex,
	task: BuildTask,
	worker: usize,
	result: Result<OutputInfo, Error>
}

fn main() {
	println!("XGConsole:");
	let args = Vec::from_iter(env::args());
	for arg in args.iter() {
		println!("  {}", arg);
	}
	match execute(&args[1..]) {
		Ok(result) => {
			env::set_exit_status(match result {
				Some(r) => r,
				None => 500
			});
		}
		Err(e) => {
			println!("FATAL ERROR: {:?}", e);
			env::set_exit_status(500);
		}
	}
}

fn execute(args: &[String]) -> Result<Option<i32>, Error> {
	let cache = Cache::new();
	let temp_dir = try! (TempDir::new("octobuild"));
	for arg in args.iter() {
		if arg.starts_with("/") {continue}

		let (tx_result, rx_result): (Sender<ResultMessage>, Receiver<ResultMessage>) = channel();
		let (tx_task, rx_task): (Sender<TaskMessage>, Receiver<TaskMessage>) = channel();

		create_threads(rx_task, tx_result, std::os::num_cpus(), |worker_id:usize| {
			let temp_path = temp_dir.path().to_path_buf();
			let temp_cache = cache.clone();
			move |task:TaskMessage| -> ResultMessage {
				execute_task(&temp_cache, &temp_path, worker_id, task)
			}
		});

		let path = Path::new(arg);
		let file = try! (File::open(&path));
		let graph = try! (xg::parser::parse(BufReader::new(file)));
		let validated_graph = try! (validate_graph(graph));
		match try! (execute_graph(&validated_graph, tx_task, rx_result)) {
			Some(v) if v == 0 => {}
			v => {return Ok(v)}
		}
	}
	cache.cleanup(16 * 1024 * 1024 * 1024);
	Ok(Some(0))
}

fn create_threads<R: 'static + Send, T: 'static + Send, Worker:'static + Fn(T) -> R + Send, Factory:Fn(usize) -> Worker>(rx_task: Receiver<T>, tx_result: Sender<R>, num_cpus: usize, factory: Factory) ->  Arc<Mutex<Receiver<T>>> {
	let mutex_rx_task = Arc::new(Mutex::new(rx_task));
	for cpu_id in 0..num_cpus {
		let local_rx_task = mutex_rx_task.clone();
		let local_tx_result = tx_result.clone();
		let worker = factory(cpu_id);
		thread::spawn(move || {
			loop {
				let task: T;
				match local_rx_task.lock().unwrap().recv() {
					Ok(v) => {task = v;}
					Err(_) => {break;}
				}
				match local_tx_result.send(worker(task)) {
					Ok(_) => {}
					Err(_) => {break;}
				}
			}
		});
	}
	mutex_rx_task
}

fn validate_graph(graph: Graph<BuildTask, ()>) -> Result<Graph<BuildTask, ()>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	let mut queue:Vec<NodeIndex> = Vec::new();
	graph.each_node(|index: NodeIndex, _:&Node<BuildTask>|->bool {
		completed.push(false);
		queue.push(index);
		true
	});
	let mut count:usize = 0;
	let mut i:usize = 0;
	while i < queue.len() {
		let index = queue[i];
		if (!completed[index.node_id()]) && (is_ready(&graph, &completed, &index)) {
			completed[index.node_id()] = true;
			graph.each_incoming_edge(index, |_:EdgeIndex, edge:&Edge<()>| -> bool {
				queue.push(edge.source());
				true
			});
			count += 1;
			if count ==completed.len() {
				return Ok(graph);
			}
		}
		i = i + 1;
	}
	Err(Error::new(ErrorKind::InvalidInput, "Found cycles in build dependencies", None))
}

fn execute_task(cache: &Cache, temp_dir: &Path, worker: usize, message: TaskMessage) -> ResultMessage {
	let args = wincmd::expand_args(&message.task.args, &|name:&str|->Option<String>{env::var(name).ok()});
	let output = execute_compiler(cache, temp_dir, &message.task.exec, &Path::new(&message.task.working_dir), &args);
	ResultMessage {
		index: message.index,
		task: message.task,
		worker: worker,
		result: output,
	}
}

fn execute_compiler(cache: &Cache, temp_dir: &Path, program: &str, cwd: &Path, args: &[String]) -> Result<OutputInfo, Error> {
	let command = CommandInfo {
		program: Path::new(program).to_path_buf(),
		current_dir: Some(cwd.to_path_buf()),
	};
	if Path::new(program).ends_with("cl.exe") {
		let compiler = VsCompiler::new(cache, temp_dir);
		compiler.compile(command, args)
	} else {
		command.to_command()
			.args(&args)
			.output()
			.map(|o| OutputInfo::new(o))
	}
}

fn execute_graph(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, rx_result: Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	// Run all tasks.
	let result = execute_until_failed(graph, tx_task, &rx_result);
	// Wait for in progress task completion.
	for _ in rx_result.iter() {
	}
	result
}

fn execute_until_failed(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, rx_result: &Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	if !graph.each_node(|index: NodeIndex, node:&Node<BuildTask>|->bool {
		let mut has_edges = false;
		graph.each_outgoing_edge(index, |_:EdgeIndex, _:&Edge<()>| -> bool {
			has_edges = true;
			false
		});
		completed.push(false);
		if !has_edges {
			match tx_task.send(TaskMessage{
				index: index,
				task: node.data.clone(),
			}) {
				Ok(_) => true,
				Err(_) => false,
			}
		} else {
			true
		}
	}) {
		return Err(Error::new(ErrorKind::BrokenPipe, "Can't schedule root tasks", None));
	}

	let mut count:usize = 0;
	for message in rx_result.iter() {
		assert!(!completed[message.index.node_id()]);
		count += 1;
		println!("#{} {}/{}: {}", message.worker, count, completed.len(), message.task.title);
		let result = try! (message.result);
		try! (io::stdout().write_all(&result.stdout));
		try! (io::stderr().write_all(&result.stderr));
		if !result.success() {
			return Ok(result.status);
		}
		completed[message.index.node_id()] = true;
		if !graph.each_incoming_edge(message.index, |_:EdgeIndex, edge:&Edge<()>| -> bool {
			let source = edge.source();
			if !completed[source.node_id()] {
				if is_ready(graph, &completed, &source) {
					match tx_task.send(TaskMessage{
						index: source,
						task: graph.node(source).data.clone(),
					}) {
						Ok(_) => {},
						Err(_) => {return false;}
					}
				}
			}
			true
		}) {
			return Err(Error::new(ErrorKind::BrokenPipe, "Can't schedule child task", None));
		};
		if count == completed.len() {
			return Ok(Some(0));
		}
	}
	panic! ("Unexpected end of result pipe");
}

fn is_ready(graph: &Graph<BuildTask, ()>, completed: &Vec<bool>, source: &NodeIndex) -> bool {
	let mut ready = true;
		graph.each_outgoing_edge(*source, |_:EdgeIndex, deps:&Edge<()>| -> bool {
		if !completed[deps.target().node_id()]{
			ready = false;
			false
		} else {
			true
		}
	});
	ready
}
<commit_msg>Fix xgconsole interruption on compilation failure<commit_after>#![feature(exit_status)]
#![feature(io)]
#![feature(os)]
extern crate octobuild;
extern crate tempdir;

use octobuild::common::BuildTask;
use octobuild::cache::Cache;
use octobuild::wincmd;
use octobuild::xg;
use octobuild::graph::{Graph, NodeIndex, Node, EdgeIndex, Edge};
use octobuild::vs::compiler::VsCompiler;
use octobuild::compiler::*;

use tempdir::TempDir;

use std::fs::File;
use std::env;
use std::io::{BufReader, Error, ErrorKind, Write};
use std::io;
use std::iter::FromIterator;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::thread;

#[derive(Debug)]
struct TaskMessage {
	index: NodeIndex,
	task: BuildTask
}

#[derive(Debug)]
struct ResultMessage {
	index: NodeIndex,
	task: BuildTask,
	worker: usize,
	result: Result<OutputInfo, Error>
}

fn main() {
	println!("XGConsole:");
	let args = Vec::from_iter(env::args());
	for arg in args.iter() {
		println!("  {}", arg);
	}
	match execute(&args[1..]) {
		Ok(result) => {
			env::set_exit_status(match result {
				Some(r) => r,
				None => 500
			});
		}
		Err(e) => {
			println!("FATAL ERROR: {:?}", e);
			env::set_exit_status(500);
		}
	}
}

fn execute(args: &[String]) -> Result<Option<i32>, Error> {
	let cache = Cache::new();
	let temp_dir = try! (TempDir::new("octobuild"));
	for arg in args.iter() {
		if arg.starts_with("/") {continue}

		let (tx_result, rx_result): (Sender<ResultMessage>, Receiver<ResultMessage>) = channel();
		let (tx_task, rx_task): (Sender<TaskMessage>, Receiver<TaskMessage>) = channel();

		let mutex_rx_task = create_threads(rx_task, tx_result, std::os::num_cpus(), |worker_id:usize| {
			let temp_path = temp_dir.path().to_path_buf();
			let temp_cache = cache.clone();
			move |task:TaskMessage| -> ResultMessage {
				execute_task(&temp_cache, &temp_path, worker_id, task)
			}
		});

		let path = Path::new(arg);
		let file = try! (File::open(&path));
		let graph = try! (xg::parser::parse(BufReader::new(file)));
		let validated_graph = try! (validate_graph(graph));
		match try! (execute_graph(&validated_graph, tx_task, mutex_rx_task, rx_result)) {
			Some(v) if v == 0 => {}
			v => {return Ok(v)}
		}
	}
	cache.cleanup(16 * 1024 * 1024 * 1024);
	Ok(Some(0))
}

fn create_threads<R: 'static + Send, T: 'static + Send, Worker:'static + Fn(T) -> R + Send, Factory:Fn(usize) -> Worker>(rx_task: Receiver<T>, tx_result: Sender<R>, num_cpus: usize, factory: Factory) ->  Arc<Mutex<Receiver<T>>> {
	let mutex_rx_task = Arc::new(Mutex::new(rx_task));
	for cpu_id in 0..num_cpus {
		let local_rx_task = mutex_rx_task.clone();
		let local_tx_result = tx_result.clone();
		let worker = factory(cpu_id);
		thread::spawn(move || {
			loop {
				let task: T;
				match local_rx_task.lock().unwrap().recv() {
					Ok(v) => {task = v;}
					Err(_) => {break;}
				}
				match local_tx_result.send(worker(task)) {
					Ok(_) => {}
					Err(_) => {break;}
				}
			}
		});
	}
	mutex_rx_task
}

fn validate_graph(graph: Graph<BuildTask, ()>) -> Result<Graph<BuildTask, ()>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	let mut queue:Vec<NodeIndex> = Vec::new();
	graph.each_node(|index: NodeIndex, _:&Node<BuildTask>|->bool {
		completed.push(false);
		queue.push(index);
		true
	});
	let mut count:usize = 0;
	let mut i:usize = 0;
	while i < queue.len() {
		let index = queue[i];
		if (!completed[index.node_id()]) && (is_ready(&graph, &completed, &index)) {
			completed[index.node_id()] = true;
			graph.each_incoming_edge(index, |_:EdgeIndex, edge:&Edge<()>| -> bool {
				queue.push(edge.source());
				true
			});
			count += 1;
			if count ==completed.len() {
				return Ok(graph);
			}
		}
		i = i + 1;
	}
	Err(Error::new(ErrorKind::InvalidInput, "Found cycles in build dependencies", None))
}

fn execute_task(cache: &Cache, temp_dir: &Path, worker: usize, message: TaskMessage) -> ResultMessage {
	let args = wincmd::expand_args(&message.task.args, &|name:&str|->Option<String>{env::var(name).ok()});
	let output = execute_compiler(cache, temp_dir, &message.task.exec, &Path::new(&message.task.working_dir), &args);
	ResultMessage {
		index: message.index,
		task: message.task,
		worker: worker,
		result: output,
	}
}

fn execute_compiler(cache: &Cache, temp_dir: &Path, program: &str, cwd: &Path, args: &[String]) -> Result<OutputInfo, Error> {
	let command = CommandInfo {
		program: Path::new(program).to_path_buf(),
		current_dir: Some(cwd.to_path_buf()),
	};
	if Path::new(program).ends_with("cl.exe") {
		let compiler = VsCompiler::new(cache, temp_dir);
		compiler.compile(command, args)
	} else {
		command.to_command()
			.args(&args)
			.output()
			.map(|o| OutputInfo::new(o))
	}
}

fn execute_graph(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, mutex_rx_task: Arc<Mutex<Receiver<TaskMessage>>>, rx_result: Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	// Run all tasks.
	let result = execute_until_failed(graph, tx_task, &rx_result);
	// Cleanup task queue.
	for _ in mutex_rx_task.lock().unwrap().iter() {
	}
	// Wait for in progress task completion.
	for _ in rx_result.iter() {
	}
	result
}

fn execute_until_failed(graph: &Graph<BuildTask, ()>, tx_task: Sender<TaskMessage>, rx_result: &Receiver<ResultMessage>) -> Result<Option<i32>, Error> {
	let mut completed:Vec<bool> = Vec::new();
	if !graph.each_node(|index: NodeIndex, node:&Node<BuildTask>|->bool {
		let mut has_edges = false;
		graph.each_outgoing_edge(index, |_:EdgeIndex, _:&Edge<()>| -> bool {
			has_edges = true;
			false
		});
		completed.push(false);
		if !has_edges {
			match tx_task.send(TaskMessage{
				index: index,
				task: node.data.clone(),
			}) {
				Ok(_) => true,
				Err(_) => false,
			}
		} else {
			true
		}
	}) {
		return Err(Error::new(ErrorKind::BrokenPipe, "Can't schedule root tasks", None));
	}

	let mut count:usize = 0;
	for message in rx_result.iter() {
		assert!(!completed[message.index.node_id()]);
		count += 1;
		println!("#{} {}/{}: {}", message.worker, count, completed.len(), message.task.title);
		let result = try! (message.result);
		try! (io::stdout().write_all(&result.stdout));
		try! (io::stderr().write_all(&result.stderr));
		if !result.success() {
			return Ok(result.status);
		}
		completed[message.index.node_id()] = true;
		if !graph.each_incoming_edge(message.index, |_:EdgeIndex, edge:&Edge<()>| -> bool {
			let source = edge.source();
			if !completed[source.node_id()] {
				if is_ready(graph, &completed, &source) {
					match tx_task.send(TaskMessage{
						index: source,
						task: graph.node(source).data.clone(),
					}) {
						Ok(_) => {},
						Err(_) => {return false;}
					}
				}
			}
			true
		}) {
			return Err(Error::new(ErrorKind::BrokenPipe, "Can't schedule child task", None));
		};
		if count == completed.len() {
			return Ok(Some(0));
		}
	}
	panic! ("Unexpected end of result pipe");
}

fn is_ready(graph: &Graph<BuildTask, ()>, completed: &Vec<bool>, source: &NodeIndex) -> bool {
	let mut ready = true;
		graph.each_outgoing_edge(*source, |_:EdgeIndex, deps:&Edge<()>| -> bool {
		if !completed[deps.target().node_id()]{
			ready = false;
			false
		} else {
			true
		}
	});
	ready
}
<|endoftext|>