<commit_before>use std::error;
use std::fmt::{self, Display};
use std::io::Write;

use serde;

use rmp;
use rmp::Marker;
use rmp::encode::{write_nil, write_bool, write_uint, write_sint, write_f32, write_f64, write_str,
                  write_array_len, write_map_len, write_bin_len, ValueWriteError};

#[derive(Debug)]
pub enum Error {
    InvalidValueWrite(ValueWriteError),

    /// Failed to serialize struct, sequence or map, because its length is unknown.
    UnknownLength,

    /// Depth limit exceeded
    DepthLimitExceeded,
    Custom(String)
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::InvalidValueWrite(..) => "invalid value write",
            Error::UnknownLength => "attempt to serialize struct, sequence or map with unknown length",
            Error::DepthLimitExceeded => "depth limit exceeded",
            Error::Custom(..) => "custom message",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::InvalidValueWrite(ref err) => Some(err),
            Error::UnknownLength => None,
            Error::DepthLimitExceeded => None,
            Error::Custom(..) => None,
        }
    }
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        error::Error::description(self).fmt(f)
    }
}

impl From<ValueWriteError> for Error {
    fn from(err: ValueWriteError) -> Error {
        Error::InvalidValueWrite(err)
    }
}

impl serde::ser::Error for Error {
    /// Raised when there is general error when deserializing a type.
    fn custom<T: Into<String>>(msg: T) -> Error {
        Error::Custom(msg.into())
    }
}

pub trait VariantWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError>
        where W: Write;
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError>
        where W: Write;
}

/// Writes struct as MessagePack array with no field names
pub struct StructArrayWriter;

impl VariantWriter for StructArrayWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError>
        where W: Write
    {
        write_array_len(wr, len)
    }

    /// This implementation does not write field names
    #[allow(unused_variables)]
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError>
        where W: Write
    {
        Ok(())
    }
}

/// Represents MessagePack serialization implementation.
///
/// # Note
///
/// MessagePack has no specification about how to encode variant types. Thus we are free to do
/// whatever we want, so the given chose may be not ideal for you.
///
/// Every Rust variant value can be represented as a tuple of index and a value.
///
/// All instances of `ErrorKind::Interrupted` are handled by this function and the underlying
/// operation is retried.
// TODO: Docs. Examples.
pub struct Serializer<'a, W: VariantWriter> {
    wr: &'a mut Write,
    vw: W,
    depth: usize,
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Changes the maximum nesting depth that is allowed
    pub fn set_max_depth(&mut self, depth: usize) {
        self.depth = depth;
    }
}

impl<'a> Serializer<'a, StructArrayWriter> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn new(wr: &'a mut Write) -> Serializer<'a, StructArrayWriter> {
        Serializer {
            wr: wr,
            vw: StructArrayWriter,
            depth: 1024,
        }
    }
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn with(wr: &'a mut Write, vw: W) -> Serializer<'a, W> {
        Serializer {
            wr: wr,
            vw: vw,
            depth: 1024,
        }
    }
}

impl<'a, W: VariantWriter> serde::Serializer for Serializer<'a, W> {
    type Error = Error;
    type SeqState = ();
    type TupleState = ();
    type TupleStructState = ();
    type TupleVariantState = ();
    type MapState = ();
    type StructState = ();
    type StructVariantState = ();

    fn serialize_unit(&mut self) -> Result<(), Error> {
        write_nil(&mut self.wr)
            .map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidMarkerWrite(err)))
    }

    fn serialize_bool(&mut self, val: bool) -> Result<(), Error> {
        write_bool(&mut self.wr, val)
            .map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidMarkerWrite(err)))
    }

    fn serialize_u8(&mut self, val: u8) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u16(&mut self, val: u16) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u32(&mut self, val: u32) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u64(&mut self, val: u64) -> Result<(), Error> {
        println!("serialize_u64");
        try!(write_uint(&mut self.wr, val));
        Ok(())
    }

    fn serialize_usize(&mut self, val: usize) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_i8(&mut self, val: i8) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i16(&mut self, val: i16) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i32(&mut self, val: i32) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i64(&mut self, val: i64) -> Result<(), Error> {
        try!(write_sint(&mut self.wr, val));
        Ok(())
    }

    fn serialize_isize(&mut self, val: isize) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_f32(&mut self, val: f32) -> Result<(), Error> {
        write_f32(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_f64(&mut self, val: f64) -> Result<(), Error> {
        write_f64(&mut self.wr, val).map_err(From::from)
    }

    // TODO: The implementation involves heap allocation and is unstable.
    fn serialize_char(&mut self, val: char) -> Result<(), Error> {
        let mut buf = String::new();
        buf.push(val);
        self.serialize_str(&buf)
    }

    fn serialize_str(&mut self, val: &str) -> Result<(), Error> {
        write_str(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_unit_variant(&mut self,
                          _name: &str,
                          variant_index: usize,
                          _variant: &str) -> Result<(), Error>
    {
        // Mark that we want to encode a variant type.
        try!(write_array_len(&mut self.wr, 2));

        // Encode a value position...
        try!(self.serialize_usize(variant_index));

        // ... and its arguments length.
        try!(write_array_len(&mut self.wr, 0));

        Ok(())
    }

    /// Encodes and attempts to write the enum value into the Write.
    ///
    /// Currently we encode variant types as a tuple of id with array of args, like: [id, [args...]]
    fn serialize_tuple_variant(&mut self, name: &'static str, id: usize, _variant : &'static str, len: usize) ->
        Result<Self::TupleVariantState, Error>
    {
        try!(rmp::encode::write_array_len(&mut self.wr, 2));
        try!(self.serialize_usize(id));
        self.serialize_tuple_struct(name, len)
    }

    fn serialize_tuple_variant_elt<T>(&mut self, state: &mut Self::TupleVariantState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_struct_elt(state, value)
    }

    fn serialize_tuple_variant_end(&mut self, state: Self::TupleVariantState) -> Result<(), Self::Error> {
        self.serialize_tuple_struct_end(state)
    }

    fn serialize_none(&mut self) -> Result<(), Error> {
        self.serialize_unit()
    }

    fn serialize_some<V>(&mut self, v: V) -> Result<(), Error>
        where V: serde::Serialize,
    {
        v.serialize(self)
    }

    // TODO: Check len, overflow is possible.
    fn serialize_seq(&mut self, length: Option<usize>) -> Result<Self::SeqState, Error>
    {
        let len = match length {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_array_len(&mut self.wr, len as u32));

        Ok(())
    }

    fn serialize_seq_elt<V>(&mut self, _state: &mut Self::SeqState, value: V) -> Result<Self::SeqState, Error>
        where V: serde::Serialize
    {
        let _ = value.serialize(self);
        Ok(())
    }

    fn serialize_seq_end(&mut self, _state: Self::SeqState) -> Result<(), Self::Error>
    {
        Ok(())
    }

    fn serialize_seq_fixed_size(&mut self, size: usize) -> Result<Self::SeqState, Self::Error>
    {
        self.serialize_seq(Some(size))
    }

    fn serialize_map(&mut self, length: Option<usize>) -> Result<Self::MapState, Error>
    {
        let len = match length {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_map_len(&mut self.wr, len as u32));

        Ok(())
    }

    fn serialize_map_key<T>(&mut self, _state: &mut Self::MapState, key: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        let _ = key.serialize(self);

        Ok(())
    }

    fn serialize_map_value<T>(&mut self, _state: &mut Self::MapState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        let _ = value.serialize(self);

        Ok(())
    }

    fn serialize_map_end(&mut self, _state: Self::MapState) -> Result<(), Self::Error>
    {
        Ok(())
    }

    fn serialize_unit_struct(&mut self, _name: &'static str) -> Result<(), Error>
    {
        try!(self.vw.write_struct_len(&mut self.wr, 0));

        Ok(())
    }

    fn serialize_tuple(&mut self, len: usize) -> Result<Self::TupleState, Self::Error>
    {
        self.serialize_seq(Some(len))
    }

    fn serialize_tuple_elt<T>(&mut self, state: &mut Self::TupleState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_seq_elt(state, value)
    }

    fn serialize_tuple_end(&mut self, state: Self::TupleState) -> Result<(), Self::Error>
    {
        self.serialize_seq_end(state)
    }

    fn serialize_tuple_struct(&mut self, _name: &'static str, len: usize) -> Result<Self::TupleStructState, Self::Error>
    {
        self.serialize_tuple(len)
    }

    fn serialize_tuple_struct_elt<T>(&mut self, state: &mut Self::TupleStructState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_elt(state, value)
    }

    fn serialize_tuple_struct_end(&mut self, state: Self::TupleStructState) -> Result<(), Self::Error>
    {
        self.serialize_tuple_end(state)
    }

    fn serialize_newtype_struct<T>(&mut self, name: &'static str, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_struct(name, 1)
        .and_then(|mut state: Self::TupleState| self.serialize_tuple_struct_elt(&mut state, value))
        .and_then(|state: Self::TupleState| self.serialize_tuple_struct_end(state))
    }

    fn serialize_newtype_variant<T>(&mut self, name: &'static str, variant_index: usize, variant: &'static str, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_variant(name, variant_index, variant, 1)
        .and_then(|mut state: Self::TupleState| self.serialize_tuple_variant_elt(&mut state, value))
        .and_then(|state: Self::TupleState| self.serialize_tuple_variant_end(state))
    }

    fn serialize_bytes(&mut self, value: &[u8]) -> Result<(), Error> {
        try!(write_bin_len(&mut self.wr, value.len() as u32));
        self.wr.write_all(value)
            .map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidDataWrite(err)))
    }

    /// Begins to serialize a struct. This call must be followed by zero or more
    /// calls to `serialize_struct_elt`, then a call to `serialize_struct_end`.
    fn serialize_struct(&mut self, _name: &'static str, len: usize) -> Result<Self::StructState, Self::Error>
    {
        try!(self.vw.write_struct_len(&mut self.wr, len as u32));
        Ok(())
    }

    /// Serializes a struct field. Must have previously called
    /// `serialize_struct`.
    fn serialize_struct_elt<V>(&mut self, _state: &mut Self::StructState, key: &'static str, value: V) ->  Result<(), Self::Error>
        where V: serde::Serialize
    {
        self.vw.write_field_name(&mut self.wr, key)
            .map_err(|e| e.into())
            .and_then(|_| value.serialize(self))
            .map(|_| ())
    }

    /// Finishes serializing a struct.
    fn serialize_struct_end(&mut self, _state: Self::StructState) -> Result<(), Error> {
        Ok(())
    }

    /// Begins to serialize a struct variant.
    ///
    /// This call must be followed by zero or more calls to `serialize_struct_variant_elt()`, then
    /// a call to `serialize_struct_variant_end()`.
    fn serialize_struct_variant(&mut self, name: &'static str, id: usize, _variant: &'static str, len: usize) ->
        Result<Self::StructVariantState, Error>
    {
        try!(write_array_len(&mut self.wr, 2));
        try!(serde::Serializer::serialize_usize(self, id));
        self.serialize_struct(name, len)
    }

    /// Serialize a struct variant element.
    ///
    /// Must have previously called `serialize_struct_variant()`.
    fn serialize_struct_variant_elt<V>(&mut self, state: &mut Self::StructVariantState, key: &'static str, value: V) -> Result<(), Self::Error>
        where V: serde::Serialize
    {
        self.serialize_struct_elt(state, key, value)
    }

    /// Finishes serializing a struct variant.
    fn serialize_struct_variant_end(&mut self,state: Self::StructVariantState) -> Result<(), Self::Error>
    {
        self.serialize_struct_end(state)
    }
}
<commit_msg>refactor: and again<commit_after>use std::error;
use std::fmt::{self, Display};
use std::io::Write;

use serde;

use rmp;
use rmp::Marker;
use rmp::encode::{write_nil, write_bool, write_uint, write_sint, write_f32, write_f64, write_str,
                  write_array_len, write_map_len, write_bin_len, ValueWriteError};

#[derive(Debug)]
pub enum Error {
    InvalidValueWrite(ValueWriteError),

    /// Failed to serialize struct, sequence or map, because its length is unknown.
    UnknownLength,

    /// Depth limit exceeded
    DepthLimitExceeded,
    Custom(String)
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::InvalidValueWrite(..) => "invalid value write",
            Error::UnknownLength => "attempt to serialize struct, sequence or map with unknown length",
            Error::DepthLimitExceeded => "depth limit exceeded",
            Error::Custom(..) => "custom message",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::InvalidValueWrite(ref err) => Some(err),
            Error::UnknownLength => None,
            Error::DepthLimitExceeded => None,
            Error::Custom(..) => None,
        }
    }
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        error::Error::description(self).fmt(f)
    }
}

impl From<ValueWriteError> for Error {
    fn from(err: ValueWriteError) -> Error {
        Error::InvalidValueWrite(err)
    }
}

impl serde::ser::Error for Error {
    /// Raised when there is general error when deserializing a type.
    fn custom<T: Into<String>>(msg: T) -> Error {
        Error::Custom(msg.into())
    }
}

pub trait VariantWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError>
        where W: Write;
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError>
        where W: Write;
}

/// Writes struct as MessagePack array with no field names
pub struct StructArrayWriter;

impl VariantWriter for StructArrayWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError>
        where W: Write
    {
        write_array_len(wr, len)
    }

    /// This implementation does not write field names
    #[allow(unused_variables)]
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError>
        where W: Write
    {
        Ok(())
    }
}

/// Represents MessagePack serialization implementation.
///
/// # Note
///
/// MessagePack has no specification about how to encode variant types. Thus we are free to do
/// whatever we want, so the given chose may be not ideal for you.
///
/// Every Rust variant value can be represented as a tuple of index and a value.
///
/// All instances of `ErrorKind::Interrupted` are handled by this function and the underlying
/// operation is retried.
// TODO: Docs. Examples.
pub struct Serializer<'a, W: VariantWriter> {
    wr: &'a mut Write,
    vw: W,
    depth: usize,
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Changes the maximum nesting depth that is allowed
    pub fn set_max_depth(&mut self, depth: usize) {
        self.depth = depth;
    }
}

impl<'a> Serializer<'a, StructArrayWriter> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn new(wr: &'a mut Write) -> Serializer<'a, StructArrayWriter> {
        Serializer {
            wr: wr,
            vw: StructArrayWriter,
            depth: 1024,
        }
    }
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn with(wr: &'a mut Write, vw: W) -> Serializer<'a, W> {
        Serializer {
            wr: wr,
            vw: vw,
            depth: 1024,
        }
    }
}

impl<'a, W: VariantWriter> serde::Serializer for Serializer<'a, W> {
    type Error = Error;
    type SeqState = ();
    type TupleState = ();
    type TupleStructState = ();
    type TupleVariantState = ();
    type MapState = ();
    type StructState = ();
    type StructVariantState = ();

    fn serialize_unit(&mut self) -> Result<(), Error> {
        write_nil(&mut self.wr)
            .map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidMarkerWrite(err)))
    }

    fn serialize_bool(&mut self, val: bool) -> Result<(), Error> {
        write_bool(&mut self.wr, val)
            .map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidMarkerWrite(err)))
    }

    fn serialize_u8(&mut self, val: u8) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u16(&mut self, val: u16) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u32(&mut self, val: u32) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u64(&mut self, val: u64) -> Result<(), Error> {
        try!(write_uint(&mut self.wr, val));
        Ok(())
    }

    fn serialize_usize(&mut self, val: usize) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_i8(&mut self, val: i8) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i16(&mut self, val: i16) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i32(&mut self, val: i32) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i64(&mut self, val: i64) -> Result<(), Error> {
        try!(write_sint(&mut self.wr, val));
        Ok(())
    }

    fn serialize_isize(&mut self, val: isize) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_f32(&mut self, val: f32) -> Result<(), Error> {
        write_f32(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_f64(&mut self, val: f64) -> Result<(), Error> {
        write_f64(&mut self.wr, val).map_err(From::from)
    }

    // TODO: The implementation involves heap allocation and is unstable.
    fn serialize_char(&mut self, val: char) -> Result<(), Error> {
        let mut buf = String::new();
        buf.push(val);
        self.serialize_str(&buf)
    }

    fn serialize_str(&mut self, val: &str) -> Result<(), Error> {
        write_str(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_unit_variant(&mut self,
                          _name: &str,
                          variant_index: usize,
                          _variant: &str) -> Result<(), Error>
    {
        // Mark that we want to encode a variant type.
        try!(write_array_len(&mut self.wr, 2));

        // Encode a value position...
        try!(self.serialize_usize(variant_index));

        // ... and its arguments length.
        try!(write_array_len(&mut self.wr, 0));

        Ok(())
    }

    /// Encodes and attempts to write the enum value into the Write.
    ///
    /// Currently we encode variant types as a tuple of id with array of args, like: [id, [args...]]
    fn serialize_tuple_variant(&mut self, name: &'static str, id: usize, _variant : &'static str, len: usize) ->
        Result<Self::TupleVariantState, Error>
    {
        try!(rmp::encode::write_array_len(&mut self.wr, 2));
        try!(self.serialize_usize(id));
        self.serialize_tuple_struct(name, len)
    }

    fn serialize_tuple_variant_elt<T>(&mut self, state: &mut Self::TupleVariantState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_struct_elt(state, value)
    }

    fn serialize_tuple_variant_end(&mut self, state: Self::TupleVariantState) -> Result<(), Self::Error> {
        self.serialize_tuple_struct_end(state)
    }

    fn serialize_none(&mut self) -> Result<(), Error> {
        self.serialize_unit()
    }

    fn serialize_some<V>(&mut self, v: V) -> Result<(), Error>
        where V: serde::Serialize,
    {
        v.serialize(self)
    }

    // TODO: Check len, overflow is possible.
    fn serialize_seq(&mut self, length: Option<usize>) -> Result<Self::SeqState, Error>
    {
        let len = match length {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_array_len(&mut self.wr, len as u32));

        Ok(())
    }

    fn serialize_seq_elt<V>(&mut self, _state: &mut Self::SeqState, value: V) -> Result<Self::SeqState, Error>
        where V: serde::Serialize
    {
        let _ = value.serialize(self);
        Ok(())
    }

    fn serialize_seq_end(&mut self, _state: Self::SeqState) -> Result<(), Self::Error>
    {
        Ok(())
    }

    fn serialize_seq_fixed_size(&mut self, size: usize) -> Result<Self::SeqState, Self::Error>
    {
        self.serialize_seq(Some(size))
    }

    fn serialize_map(&mut self, length: Option<usize>) -> Result<Self::MapState, Error>
    {
        let len = match length {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_map_len(&mut self.wr, len as u32));

        Ok(())
    }

    fn serialize_map_key<T>(&mut self, _state: &mut Self::MapState, key: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        let _ = key.serialize(self);

        Ok(())
    }

    fn serialize_map_value<T>(&mut self, _state: &mut Self::MapState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        let _ = value.serialize(self);

        Ok(())
    }

    fn serialize_map_end(&mut self, _state: Self::MapState) -> Result<(), Self::Error>
    {
        Ok(())
    }

    fn serialize_unit_struct(&mut self, _name: &'static str) -> Result<(), Error>
    {
        try!(self.vw.write_struct_len(&mut self.wr, 0));

        Ok(())
    }

    fn serialize_tuple(&mut self, len: usize) -> Result<Self::TupleState, Self::Error>
    {
        self.serialize_seq(Some(len))
    }

    fn serialize_tuple_elt<T>(&mut self, state: &mut Self::TupleState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_seq_elt(state, value)
    }

    fn serialize_tuple_end(&mut self, state: Self::TupleState) -> Result<(), Self::Error>
    {
        self.serialize_seq_end(state)
    }

    fn serialize_tuple_struct(&mut self, _name: &'static str, len: usize) -> Result<Self::TupleStructState, Self::Error>
    {
        self.serialize_tuple(len)
    }

    fn serialize_tuple_struct_elt<T>(&mut self, state: &mut Self::TupleStructState, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_elt(state, value)
    }

    fn serialize_tuple_struct_end(&mut self, state: Self::TupleStructState) -> Result<(), Self::Error>
    {
        self.serialize_tuple_end(state)
    }

    fn serialize_newtype_struct<T>(&mut self, name: &'static str, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_struct(name, 1)
        .and_then(|mut state: Self::TupleState| self.serialize_tuple_struct_elt(&mut state, value))
        .and_then(|state: Self::TupleState| self.serialize_tuple_struct_end(state))
    }

    fn serialize_newtype_variant<T>(&mut self, name: &'static str, variant_index: usize, variant: &'static str, value: T) -> Result<(), Self::Error>
        where T: serde::Serialize
    {
        self.serialize_tuple_variant(name, variant_index, variant, 1)
        .and_then(|mut state: Self::TupleState| self.serialize_tuple_variant_elt(&mut state, value))
        .and_then(|state: Self::TupleState| self.serialize_tuple_variant_end(state))
    }

    fn serialize_bytes(&mut self, value: &[u8]) -> Result<(), Error> {
        try!(write_bin_len(&mut self.wr, value.len() as u32));
        self.wr.write_all(value)
            .map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidDataWrite(err)))
    }

    /// Begins to serialize a struct. This call must be followed by zero or more
    /// calls to `serialize_struct_elt`, then a call to `serialize_struct_end`.
    fn serialize_struct(&mut self, _name: &'static str, len: usize) -> Result<Self::StructState, Self::Error>
    {
        try!(self.vw.write_struct_len(&mut self.wr, len as u32));
        Ok(())
    }

    /// Serializes a struct field. Must have previously called
    /// `serialize_struct`.
    fn serialize_struct_elt<V>(&mut self, _state: &mut Self::StructState, key: &'static str, value: V) ->  Result<(), Self::Error>
        where V: serde::Serialize
    {
        self.vw.write_field_name(&mut self.wr, key)
            .map_err(|e| e.into())
            .and_then(|_| value.serialize(self))
            .map(|_| ())
    }

    /// Finishes serializing a struct.
    fn serialize_struct_end(&mut self, _state: Self::StructState) -> Result<(), Error> {
        Ok(())
    }

    /// Begins to serialize a struct variant.
    ///
    /// This call must be followed by zero or more calls to `serialize_struct_variant_elt()`, then
    /// a call to `serialize_struct_variant_end()`.
    fn serialize_struct_variant(&mut self, name: &'static str, id: usize, _variant: &'static str, len: usize) ->
        Result<Self::StructVariantState, Error>
    {
        try!(write_array_len(&mut self.wr, 2));
        try!(serde::Serializer::serialize_usize(self, id));
        self.serialize_struct(name, len)
    }

    /// Serialize a struct variant element.
    ///
    /// Must have previously called `serialize_struct_variant()`.
    fn serialize_struct_variant_elt<V>(&mut self, state: &mut Self::StructVariantState, key: &'static str, value: V) -> Result<(), Self::Error>
        where V: serde::Serialize
    {
        self.serialize_struct_elt(state, key, value)
    }

    /// Finishes serializing a struct variant.
    fn serialize_struct_variant_end(&mut self,state: Self::StructVariantState) -> Result<(), Self::Error>
    {
        self.serialize_struct_end(state)
    }
}
<|endoftext|>


<commit_before>use rusqlite::{Connection, SQLITE_OPEN_READ_WRITE, SQLITE_OPEN_CREATE,
               SQLITE_OPEN_FULL_MUTEX, SQLITE_OPEN_URI};

/// Using a slightly modified set of rusqlite open flags, open and return a connection
pub fn db_connect(path: &String) -> Connection {
    let flags = { SQLITE_OPEN_READ_WRITE | SQLITE_OPEN_CREATE |
                  SQLITE_OPEN_FULL_MUTEX | SQLITE_OPEN_URI };

    let conn = match path.as_ref() {
        "" => {
            Connection::open_in_memory_with_flags(flags)
        },
        x => {
            Connection::open_with_flags(x, flags)
        }
    };

    debug!("Connection to {:?} has been made", path);
    conn.unwrap()
}

// Initialize the database
pub fn db_init(conn: &Connection) {
    conn.execute("
        CREATE TABLE IF NOT EXISTS torrent (
            info_hash   TEXT,
            ip          INTEGER,
            port        INTEGER,
            peer_id     TEXT,
            remaining   INTEGER,
            last_active INTEGER,
            PRIMARY KEY (info_hash, ip, port, peer_id)
        );",
        &[]
    ).unwrap();
}

pub fn db_prune(conn: &Connection) {
    conn.execute(
        "DELETE FROM torrent
        WHERE (strftime('%s','now') - last_active) > 1860;",
        &[]
    ).unwrap();
}
<commit_msg>db: connection code cleanup<commit_after>use std::path::Path;

use rusqlite::{Connection, SQLITE_OPEN_READ_WRITE, SQLITE_OPEN_CREATE,
               SQLITE_OPEN_FULL_MUTEX, SQLITE_OPEN_URI};

/// Using a slightly modified set of rusqlite open flags, open and return a connection
pub fn db_connect(path: &String) -> Connection {
    let flags = { SQLITE_OPEN_READ_WRITE | SQLITE_OPEN_CREATE |
                  SQLITE_OPEN_FULL_MUTEX | SQLITE_OPEN_URI };

    let conn = if Path::new(path).exists() {
            debug!("Connection to {:?} has been established", path);
            Connection::open_with_flags(path, flags)
    } else {
            debug!("Connection to in memory db has been established");
            Connection::open_in_memory_with_flags(flags)
    };

    conn.unwrap()
}

// Initialize the database
pub fn db_init(conn: &Connection) {
    conn.execute("
        CREATE TABLE IF NOT EXISTS torrent (
            info_hash   TEXT,
            ip          INTEGER,
            port        INTEGER,
            peer_id     TEXT,
            remaining   INTEGER,
            last_active INTEGER,
            PRIMARY KEY (info_hash, ip, port, peer_id)
        );",
        &[]
    ).unwrap();
}

pub fn db_prune(conn: &Connection) {
    conn.execute(
        "DELETE FROM torrent
        WHERE (strftime('%s','now') - last_active) > 1860;",
        &[]
    ).unwrap();
}
<|endoftext|>


<commit_before>use std::process;
use std::collections::HashSet;
use std::io::Write;
use std::io;

use gb_proc::cpu::{Cpu, print_cpu_status};
use gb_proc::opcodes::OpCode;

pub struct Debugger {
    stepping: bool,
    address_breakpoints: HashSet<u16>,
    op_breakpoints: HashSet<u8>,
}

fn print_help() {
    println!("Help: ");
    println!("clear            -- clear breakpoints");
    println!("[l]ist           -- list breakpoints");
    println!("bo [u8]          -- breakpoint for opcode [u8]");
    println!("ba [u16]         -- breakpoint for address [u16]");
    println!("p cpu            -- display cpu information and registers");
    println!("[p]rint [u16]    -- print memory at [u16]");
    println!("[s]et [u16] [u8] -- put value [u8] at memory address [u16]");
    println!("[c]ontinue       -- continue execution");
    println!("d                -- continue execution but print cpu information");
    println!("[q]uit           -- quit application");
    println!("");
}

fn to_value(arg: &str) -> Result<u8, ()> {
    u8::from_str_radix(arg, 16).map_err(|_| ())
}

fn to_address(arg: &str) -> Result<u16, ()> {
    u16::from_str_radix(arg, 16).map_err(|_| ())
}

impl Debugger {
    pub fn new() -> Debugger {
        Debugger {
            stepping: false,
            address_breakpoints: HashSet::new(),
            op_breakpoints: HashSet::new(),
        }
    }

    pub fn handle_unary(&mut self, command: &str, cpu: &mut Cpu) -> Result<bool, ()> {
        match command {
            "l" | "list" => {
                for op in self.op_breakpoints.clone() {
                    println!("Breakpoint for opcode {}",
                             OpCode::from_byte(op, false).to_string());
                }
                for address in self.address_breakpoints.clone() {
                    println!("Breakpoint for address {:04X}", address);
                }
            },
            "clear" => {
                self.op_breakpoints.clear();
                self.address_breakpoints.clear();
            },
            "c" | "continue" => {
                self.stepping = false;
                cpu.set_debug(false);
                println!("Continuing.");

                return Ok(true);
            },
            "d" => {
                self.stepping = false;

                return Ok(true);
            },
            "q" | "quit" => {
                println!("Quitting.");
                process::exit(0);
            },
            "h" | "help" => {
                print_help();
            },
            _ => {
                return Err(());
            },
        }

        Ok(false)
    }

    pub fn handle_binary(&mut self, command: &str, arg: &str, cpu: &mut Cpu) -> Result<bool, ()> {
        match command {
            "bo" => {
                let op = try!(to_value(arg));
                println!("Breakpoint for opcode {}",
                         OpCode::from_byte(op, false).to_string());
                self.op_breakpoints.insert(op);
            },
            "ba" => {
                let address = try!(to_address(arg));
                println!("Breakpoint for address {:04X}", address);
                self.address_breakpoints.insert(address);
            },
            "p" | "print" => {
                if arg == "cpu" {
                    print_cpu_status(&cpu);
                } else {
                    let address = try!(to_address(arg));
                    println!("${:04X}={:02X}", address, cpu.deref(address));
                    self.address_breakpoints.insert(address);
                }
            },
            _ => {
                return Err(());
            },
        }

        Ok(false)
    }

    pub fn handle_trinary(&mut self, command: &str, arg1: &str, arg2: &str,
                          cpu: &mut Cpu) -> Result<bool, ()> {
        match command {
            "s" | "set" => {
                let address = try!(to_address(arg1));
                let v = try!(to_value(arg2));

                println!("Setting ${:04X}={:02X}h", address, v);
                cpu.set_deref(address, v);
            },
            _ => {
                return Err(());
            },
        }

        Ok(false)
    }

    fn handle(&mut self, args: Vec<&str>, cpu: &mut Cpu) -> Result<bool, ()> {
        match args.len() {
            0 => Ok(true),
            1 => self.handle_unary(args[0], cpu),
            2 => self.handle_binary(args[0], args[1], cpu),
            3 => self.handle_trinary(args[0], args[1], args[2], cpu),
            _ => Err(()),
        }
    }

    pub fn check_breakpoints(&mut self, cpu: &mut Cpu) {
        let address = cpu.get_PC();
        let op = cpu.deref_PC();

        if !self.address_breakpoints.contains(&address) &&
                !self.op_breakpoints.contains(&op) {
            return;
        }

        cpu.set_debug(true);
        self.stepping = true;
        println!("Brakepoint hit! at {:04X}", cpu.get_PC());
    }

    pub fn next_instruction(&mut self, cpu: &mut Cpu) {
        if !self.stepping {
            return;
        }

        print_cpu_status(cpu);
        loop {
            let mut input = String::new();
            print!(">");
            io::stdout().flush().unwrap();
            io::stdin().read_line(&mut input).unwrap();
            input = input.replace("\n", "")
                .replace("\r", "");

            let commands: Vec<&str> = input.split(' ').collect();

            match self.handle(commands, cpu) {
                Err(()) => {
                    println!("Command not understood {:?}", input);
                    print_help();
                },
                Ok(true) => {
                    break;
                },
                Ok(false) => {
                    println!("");
                },
            }
        }
    }

    pub fn breakpoint(&mut self, cpu: &mut Cpu) {
        self.stepping = true;
        cpu.set_debug(true);
        print_cpu_status(cpu);
    }
}
<commit_msg>Add a step command to the debugger, remember last command.<commit_after>use std::process;
use std::collections::HashSet;
use std::io::Write;
use std::io;

use gb_proc::cpu::{Cpu, print_cpu_status};
use gb_proc::opcodes::OpCode;

pub struct Debugger {
    stepping: bool,
    address_breakpoints: HashSet<u16>,
    op_breakpoints: HashSet<u8>,
    last_command: String,
}

fn print_help() {
    println!("Help: ");
    println!("clear            -- clear breakpoints");
    println!("[l]ist           -- list breakpoints");
    println!("bo [u8]          -- breakpoint for opcode [u8]");
    println!("ba [u16]         -- breakpoint for address [u16]");
    println!("p cpu            -- display cpu information and registers");
    println!("[p]rint [u16]    -- print memory at [u16]");
    println!("[s]et [u16] [u8] -- put value [u8] at memory address [u16]");
    println!("[c]ontinue       -- continue execution");
    println!("[s]tep           -- go to next instruction");
    println!("d                -- continue execution but print cpu information");
    println!("[q]uit           -- quit application");
    println!("");
}

fn to_value(arg: &str) -> Result<u8, ()> {
    u8::from_str_radix(arg, 16).map_err(|_| ())
}

fn to_address(arg: &str) -> Result<u16, ()> {
    u16::from_str_radix(arg, 16).map_err(|_| ())
}

impl Debugger {
    pub fn new() -> Debugger {
        Debugger {
            stepping: false,
            address_breakpoints: HashSet::new(),
            op_breakpoints: HashSet::new(),
            last_command: "s".to_string(),
        }
    }

    pub fn handle_unary(&mut self, command: &str, cpu: &mut Cpu) -> Result<bool, ()> {
        match command {
            "l" | "list" => {
                for op in self.op_breakpoints.clone() {
                    println!("Breakpoint for opcode {}",
                             OpCode::from_byte(op, false).to_string());
                }
                for address in self.address_breakpoints.clone() {
                    println!("Breakpoint for address {:04X}", address);
                }
            },
            "clear" => {
                self.op_breakpoints.clear();
                self.address_breakpoints.clear();
            },
            "c" | "continue" => {
                self.stepping = false;
                cpu.set_debug(false);
                println!("Continuing.");

                return Ok(true);
            },
            "d" => {
                self.stepping = false;

                return Ok(true);
            },
            "s" | "step" => {
                return Ok(true);
            },
            "q" | "quit" => {
                println!("Quitting.");
                process::exit(0);
            },
            "h" | "help" => {
                print_help();
            },
            _ => {
                return Err(());
            },
        }

        Ok(false)
    }

    pub fn handle_binary(&mut self, command: &str, arg: &str, cpu: &mut Cpu) -> Result<bool, ()> {
        match command {
            "bo" => {
                let op = try!(to_value(arg));
                println!("Breakpoint for opcode {}",
                         OpCode::from_byte(op, false).to_string());
                self.op_breakpoints.insert(op);
            },
            "ba" => {
                let address = try!(to_address(arg));
                println!("Breakpoint for address {:04X}", address);
                self.address_breakpoints.insert(address);
            },
            "p" | "print" => {
                if arg == "cpu" {
                    print_cpu_status(&cpu);
                } else {
                    let address = try!(to_address(arg));
                    println!("${:04X}={:02X}", address, cpu.deref(address));
                    self.address_breakpoints.insert(address);
                }
            },
            _ => {
                return Err(());
            },
        }

        Ok(false)
    }

    pub fn handle_trinary(&mut self, command: &str, arg1: &str, arg2: &str,
                          cpu: &mut Cpu) -> Result<bool, ()> {
        match command {
            "s" | "set" => {
                let address = try!(to_address(arg1));
                let v = try!(to_value(arg2));

                println!("Setting ${:04X}={:02X}h", address, v);
                cpu.set_deref(address, v);
            },
            _ => {
                return Err(());
            },
        }

        Ok(false)
    }

    fn handle(&mut self, args: Vec<&str>, cpu: &mut Cpu) -> Result<bool, ()> {
        match args.len() {
            0 => Ok(true),
            1 => self.handle_unary(args[0], cpu),
            2 => self.handle_binary(args[0], args[1], cpu),
            3 => self.handle_trinary(args[0], args[1], args[2], cpu),
            _ => Err(()),
        }
    }

    pub fn check_breakpoints(&mut self, cpu: &mut Cpu) {
        let address = cpu.get_PC();
        let op = cpu.deref_PC();

        if !self.address_breakpoints.contains(&address) &&
                !self.op_breakpoints.contains(&op) {
            return;
        }

        cpu.set_debug(true);
        self.stepping = true;
        println!("Brakepoint hit! at {:04X}", cpu.get_PC());
    }

    pub fn next_instruction(&mut self, cpu: &mut Cpu) {
        if !self.stepping {
            return;
        }

        print_cpu_status(cpu);
        loop {
            let mut input = String::new();
            print!(">");
            io::stdout().flush().unwrap();
            io::stdin().read_line(&mut input).unwrap();
            input = input.replace("\n", "")
                .replace("\r", "");

            if input != "" {
                self.last_command = input.clone();
            } else {
                input = self.last_command.clone();
                println!("{}", input);
            }

            let commands: Vec<&str> = input.split(' ').collect();

            match self.handle(commands, cpu) {
                Err(()) => {
                    println!("Command not understood {:?}", input);
                    print_help();
                },
                Ok(true) => {
                    break;
                },
                Ok(false) => {
                    println!("");
                },
            }
        }
    }

    pub fn breakpoint(&mut self, cpu: &mut Cpu) {
        self.stepping = true;
        cpu.set_debug(true);
        print_cpu_status(cpu);
    }
}
<|endoftext|>


<commit_before>use wiggle_runtime::{GuestError, GuestErrorType, GuestPtr, GuestPtrMut, GuestString};
use wiggle_test::WasiCtx;

lucet_wiggle::from_witx!({
    witx: ["../wasi/phases/snapshot/witx/wasi_snapshot_preview1.witx"],
    ctx: WasiCtx,
});

type Result<T> = std::result::Result<T, types::Errno>;

impl GuestErrorType for types::Errno {
    type Context = WasiCtx;
    fn success() -> types::Errno {
        types::Errno::Success
    }
    fn from_error(e: GuestError, ctx: &mut WasiCtx) -> types::Errno {
        eprintln!("GUEST ERROR: {:?}", e);
        ctx.guest_errors.push(e);
        types::Errno::Io
    }
}

impl crate::wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiCtx {
    fn args_get(
        &mut self,
        _argv: GuestPtrMut<GuestPtrMut<u8>>,
        _argv_buf: GuestPtrMut<u8>,
    ) -> Result<()> {
        unimplemented!("args_get")
    }

    fn args_sizes_get(&mut self) -> Result<(types::Size, types::Size)> {
        unimplemented!("args_sizes_get")
    }

    fn environ_get(
        &mut self,
        _environ: GuestPtrMut<GuestPtrMut<u8>>,
        _environ_buf: GuestPtrMut<u8>,
    ) -> Result<()> {
        unimplemented!("environ_get")
    }

    fn environ_sizes_get(&mut self) -> Result<(types::Size, types::Size)> {
        unimplemented!("environ_sizes_get")
    }

    fn clock_res_get(&mut self, _id: types::Clockid) -> Result<types::Timestamp> {
        unimplemented!("clock_res_get")
    }

    fn clock_time_get(
        &mut self,
        _id: types::Clockid,
        _precision: types::Timestamp,
    ) -> Result<types::Timestamp> {
        unimplemented!("clock_time_get")
    }

    fn fd_advise(
        &mut self,
        _fd: types::Fd,
        _offset: types::Filesize,
        _len: types::Filesize,
        _advice: types::Advice,
    ) -> Result<()> {
        unimplemented!("fd_advise")
    }

    fn fd_allocate(
        &mut self,
        _fd: types::Fd,
        _offset: types::Filesize,
        _len: types::Filesize,
    ) -> Result<()> {
        unimplemented!("fd_allocate")
    }

    fn fd_close(&mut self, _fd: types::Fd) -> Result<()> {
        unimplemented!("fd_close")
    }

    fn fd_datasync(&mut self, _fd: types::Fd) -> Result<()> {
        unimplemented!("fd_datasync")
    }

    fn fd_fdstat_get(&mut self, _fd: types::Fd) -> Result<types::Fdstat> {
        unimplemented!("fd_fdstat_get")
    }

    fn fd_fdstat_set_flags(&mut self, _fd: types::Fd, _flags: types::Fdflags) -> Result<()> {
        unimplemented!("fd_fdstat_set_flags")
    }

    fn fd_fdstat_set_rights(
        &mut self,
        _fd: types::Fd,
        _fs_rights_base: types::Rights,
        _fs_rights_inherting: types::Rights,
    ) -> Result<()> {
        unimplemented!("fd_fdstat_set_rights")
    }

    fn fd_filestat_get(&mut self, _fd: types::Fd) -> Result<types::Filestat> {
        unimplemented!("fd_filestat_get")
    }

    fn fd_filestat_set_size(&mut self, _fd: types::Fd, _size: types::Filesize) -> Result<()> {
        unimplemented!("fd_filestat_set_size")
    }

    fn fd_filestat_set_times(
        &mut self,
        _fd: types::Fd,
        _atim: types::Timestamp,
        _mtim: types::Timestamp,
        _fst_flags: types::Fstflags,
    ) -> Result<()> {
        unimplemented!("fd_filestat_set_times")
    }

    fn fd_pread(
        &mut self,
        _fd: types::Fd,
        _iovs: &types::IovecArray<'_>,
        _offset: types::Filesize,
    ) -> Result<types::Size> {
        unimplemented!("fd_pread")
    }

    fn fd_prestat_get(&mut self, _fd: types::Fd) -> Result<types::Prestat> {
        unimplemented!("fd_prestat_get")
    }

    fn fd_prestat_dir_name(
        &mut self,
        _fd: types::Fd,
        _path: GuestPtrMut<u8>,
        _path_len: types::Size,
    ) -> Result<()> {
        unimplemented!("fd_prestat_dir_name")
    }

    fn fd_pwrite(
        &mut self,
        _fd: types::Fd,
        _ciovs: &types::CiovecArray<'_>,
        _offset: types::Filesize,
    ) -> Result<types::Size> {
        unimplemented!("fd_pwrite")
    }

    fn fd_read(&mut self, _fd: types::Fd, _iovs: &types::IovecArray<'_>) -> Result<types::Size> {
        unimplemented!("fd_read")
    }

    fn fd_readdir(
        &mut self,
        _fd: types::Fd,
        _buf: GuestPtrMut<u8>,
        _buf_len: types::Size,
        _cookie: types::Dircookie,
    ) -> Result<types::Size> {
        unimplemented!("fd_readdir")
    }

    fn fd_renumber(&mut self, _fd: types::Fd, _to: types::Fd) -> Result<()> {
        unimplemented!("fd_renumber")
    }

    fn fd_seek(
        &mut self,
        _fd: types::Fd,
        _offset: types::Filedelta,
        _whence: types::Whence,
    ) -> Result<types::Filesize> {
        unimplemented!("fd_seek")
    }

    fn fd_sync(&mut self, _fd: types::Fd) -> Result<()> {
        unimplemented!("fd_sync")
    }

    fn fd_tell(&mut self, _fd: types::Fd) -> Result<types::Filesize> {
        unimplemented!("fd_tell")
    }

    fn fd_write(&mut self, _fd: types::Fd, _ciovs: &types::CiovecArray<'_>) -> Result<types::Size> {
        unimplemented!("fd_write")
    }

    fn path_create_directory(&mut self, _fd: types::Fd, _path: &GuestString<'_>) -> Result<()> {
        unimplemented!("path_create_directory")
    }

    fn path_filestat_get(
        &mut self,
        _fd: types::Fd,
        _flags: types::Lookupflags,
        _path: &GuestString<'_>,
    ) -> Result<types::Filestat> {
        unimplemented!("path_filestat_get")
    }

    fn path_filestat_set_times(
        &mut self,
        _fd: types::Fd,
        _flags: types::Lookupflags,
        _path: &GuestString<'_>,
        _atim: types::Timestamp,
        _mtim: types::Timestamp,
        _fst_flags: types::Fstflags,
    ) -> Result<()> {
        unimplemented!("path_filestat_set_times")
    }

    fn path_link(
        &mut self,
        _old_fd: types::Fd,
        _old_flags: types::Lookupflags,
        _old_path: &GuestString<'_>,
        _new_fd: types::Fd,
        _new_path: &GuestString<'_>,
    ) -> Result<()> {
        unimplemented!("path_link")
    }

    fn path_open(
        &mut self,
        _fd: types::Fd,
        _dirflags: types::Lookupflags,
        _path: &GuestString<'_>,
        _oflags: types::Oflags,
        _fs_rights_base: types::Rights,
        _fs_rights_inherting: types::Rights,
        _fdflags: types::Fdflags,
    ) -> Result<types::Fd> {
        unimplemented!("path_open")
    }

    fn path_readlink(
        &mut self,
        _fd: types::Fd,
        _path: &GuestString<'_>,
        _buf: GuestPtrMut<u8>,
        _buf_len: types::Size,
    ) -> Result<types::Size> {
        unimplemented!("path_readlink")
    }

    fn path_remove_directory(&mut self, _fd: types::Fd, _path: &GuestString<'_>) -> Result<()> {
        unimplemented!("path_remove_directory")
    }

    fn path_rename(
        &mut self,
        _fd: types::Fd,
        _old_path: &GuestString<'_>,
        _new_fd: types::Fd,
        _new_path: &GuestString<'_>,
    ) -> Result<()> {
        unimplemented!("path_rename")
    }

    fn path_symlink(
        &mut self,
        _old_path: &GuestString<'_>,
        _fd: types::Fd,
        _new_path: &GuestString<'_>,
    ) -> Result<()> {
        unimplemented!("path_symlink")
    }

    fn path_unlink_file(&mut self, _fd: types::Fd, _path: &GuestString<'_>) -> Result<()> {
        unimplemented!("path_unlink_file")
    }

    fn poll_oneoff(
        &mut self,
        _in_: GuestPtr<types::Subscription>,
        _out: GuestPtrMut<types::Event>,
        _nsubscriptions: types::Size,
    ) -> Result<types::Size> {
        unimplemented!("poll_oneoff")
    }

    fn proc_exit(&mut self, _rval: types::Exitcode) -> std::result::Result<(), ()> {
        unimplemented!("proc_exit")
    }

    fn proc_raise(&mut self, _sig: types::Signal) -> Result<()> {
        unimplemented!("proc_raise")
    }

    fn sched_yield(&mut self) -> Result<()> {
        unimplemented!("sched_yield")
    }

    fn random_get(&mut self, _buf: GuestPtrMut<u8>, _buf_len: types::Size) -> Result<()> {
        unimplemented!("random_get")
    }

    fn sock_recv(
        &mut self,
        _fd: types::Fd,
        _ri_data: &types::IovecArray<'_>,
        _ri_flags: types::Riflags,
    ) -> Result<(types::Size, types::Roflags)> {
        unimplemented!("sock_recv")
    }

    fn sock_send(
        &mut self,
        _fd: types::Fd,
        _si_data: &types::CiovecArray<'_>,
        _si_flags: types::Siflags,
    ) -> Result<types::Size> {
        unimplemented!("sock_send")
    }

    fn sock_shutdown(&mut self, _fd: types::Fd, _how: types::Sdflags) -> Result<()> {
        unimplemented!("sock_shutdown")
    }
}
<commit_msg>wasi wiggle test: latest wiggle api<commit_after>use wiggle_runtime::{GuestError, GuestErrorType, GuestPtr};
use wiggle_test::WasiCtx;

lucet_wiggle::from_witx!({
    witx: ["../wasi/phases/snapshot/witx/wasi_snapshot_preview1.witx"],
    ctx: WasiCtx,
});

type Result<T> = std::result::Result<T, types::Errno>;

impl GuestErrorType for types::Errno {
    type Context = WasiCtx;
    fn success() -> types::Errno {
        types::Errno::Success
    }
    fn from_error(e: GuestError, ctx: &WasiCtx) -> types::Errno {
        eprintln!("GUEST ERROR: {:?}", e);
        ctx.guest_errors.borrow().push(e);
        types::Errno::Io
    }
}

impl crate::wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiCtx {
    fn args_get(&self, _argv: GuestPtr<GuestPtr<u8>>, _argv_buf: GuestPtr<u8>) -> Result<()> {
        unimplemented!("args_get")
    }

    fn args_sizes_get(&self) -> Result<(types::Size, types::Size)> {
        unimplemented!("args_sizes_get")
    }

    fn environ_get(
        &self,
        _environ: GuestPtr<GuestPtr<u8>>,
        _environ_buf: GuestPtr<u8>,
    ) -> Result<()> {
        unimplemented!("environ_get")
    }

    fn environ_sizes_get(&self) -> Result<(types::Size, types::Size)> {
        unimplemented!("environ_sizes_get")
    }

    fn clock_res_get(&self, _id: types::Clockid) -> Result<types::Timestamp> {
        unimplemented!("clock_res_get")
    }

    fn clock_time_get(
        &self,
        _id: types::Clockid,
        _precision: types::Timestamp,
    ) -> Result<types::Timestamp> {
        unimplemented!("clock_time_get")
    }

    fn fd_advise(
        &self,
        _fd: types::Fd,
        _offset: types::Filesize,
        _len: types::Filesize,
        _advice: types::Advice,
    ) -> Result<()> {
        unimplemented!("fd_advise")
    }

    fn fd_allocate(
        &self,
        _fd: types::Fd,
        _offset: types::Filesize,
        _len: types::Filesize,
    ) -> Result<()> {
        unimplemented!("fd_allocate")
    }

    fn fd_close(&self, _fd: types::Fd) -> Result<()> {
        unimplemented!("fd_close")
    }

    fn fd_datasync(&self, _fd: types::Fd) -> Result<()> {
        unimplemented!("fd_datasync")
    }

    fn fd_fdstat_get(&self, _fd: types::Fd) -> Result<types::Fdstat> {
        unimplemented!("fd_fdstat_get")
    }

    fn fd_fdstat_set_flags(&self, _fd: types::Fd, _flags: types::Fdflags) -> Result<()> {
        unimplemented!("fd_fdstat_set_flags")
    }

    fn fd_fdstat_set_rights(
        &self,
        _fd: types::Fd,
        _fs_rights_base: types::Rights,
        _fs_rights_inherting: types::Rights,
    ) -> Result<()> {
        unimplemented!("fd_fdstat_set_rights")
    }

    fn fd_filestat_get(&self, _fd: types::Fd) -> Result<types::Filestat> {
        unimplemented!("fd_filestat_get")
    }

    fn fd_filestat_set_size(&self, _fd: types::Fd, _size: types::Filesize) -> Result<()> {
        unimplemented!("fd_filestat_set_size")
    }

    fn fd_filestat_set_times(
        &self,
        _fd: types::Fd,
        _atim: types::Timestamp,
        _mtim: types::Timestamp,
        _fst_flags: types::Fstflags,
    ) -> Result<()> {
        unimplemented!("fd_filestat_set_times")
    }

    fn fd_pread(
        &self,
        _fd: types::Fd,
        _iovs: &types::IovecArray<'_>,
        _offset: types::Filesize,
    ) -> Result<types::Size> {
        unimplemented!("fd_pread")
    }

    fn fd_prestat_get(&self, _fd: types::Fd) -> Result<types::Prestat> {
        unimplemented!("fd_prestat_get")
    }

    fn fd_prestat_dir_name(
        &self,
        _fd: types::Fd,
        _path: GuestPtr<u8>,
        _path_len: types::Size,
    ) -> Result<()> {
        unimplemented!("fd_prestat_dir_name")
    }

    fn fd_pwrite(
        &self,
        _fd: types::Fd,
        _ciovs: &types::CiovecArray<'_>,
        _offset: types::Filesize,
    ) -> Result<types::Size> {
        unimplemented!("fd_pwrite")
    }

    fn fd_read(&self, _fd: types::Fd, _iovs: &types::IovecArray<'_>) -> Result<types::Size> {
        unimplemented!("fd_read")
    }

    fn fd_readdir(
        &self,
        _fd: types::Fd,
        _buf: GuestPtr<u8>,
        _buf_len: types::Size,
        _cookie: types::Dircookie,
    ) -> Result<types::Size> {
        unimplemented!("fd_readdir")
    }

    fn fd_renumber(&self, _fd: types::Fd, _to: types::Fd) -> Result<()> {
        unimplemented!("fd_renumber")
    }

    fn fd_seek(
        &self,
        _fd: types::Fd,
        _offset: types::Filedelta,
        _whence: types::Whence,
    ) -> Result<types::Filesize> {
        unimplemented!("fd_seek")
    }

    fn fd_sync(&self, _fd: types::Fd) -> Result<()> {
        unimplemented!("fd_sync")
    }

    fn fd_tell(&self, _fd: types::Fd) -> Result<types::Filesize> {
        unimplemented!("fd_tell")
    }

    fn fd_write(&self, _fd: types::Fd, _ciovs: &types::CiovecArray<'_>) -> Result<types::Size> {
        unimplemented!("fd_write")
    }

    fn path_create_directory(&self, _fd: types::Fd, _path: &GuestPtr<'_, str>) -> Result<()> {
        unimplemented!("path_create_directory")
    }

    fn path_filestat_get(
        &self,
        _fd: types::Fd,
        _flags: types::Lookupflags,
        _path: &GuestPtr<'_, str>,
    ) -> Result<types::Filestat> {
        unimplemented!("path_filestat_get")
    }

    fn path_filestat_set_times(
        &self,
        _fd: types::Fd,
        _flags: types::Lookupflags,
        _path: &GuestPtr<'_, str>,
        _atim: types::Timestamp,
        _mtim: types::Timestamp,
        _fst_flags: types::Fstflags,
    ) -> Result<()> {
        unimplemented!("path_filestat_set_times")
    }

    fn path_link(
        &self,
        _old_fd: types::Fd,
        _old_flags: types::Lookupflags,
        _old_path: &GuestPtr<'_, str>,
        _new_fd: types::Fd,
        _new_path: &GuestPtr<'_, str>,
    ) -> Result<()> {
        unimplemented!("path_link")
    }

    fn path_open(
        &self,
        _fd: types::Fd,
        _dirflags: types::Lookupflags,
        _path: &GuestPtr<'_, str>,
        _oflags: types::Oflags,
        _fs_rights_base: types::Rights,
        _fs_rights_inherting: types::Rights,
        _fdflags: types::Fdflags,
    ) -> Result<types::Fd> {
        unimplemented!("path_open")
    }

    fn path_readlink(
        &self,
        _fd: types::Fd,
        _path: &GuestPtr<'_, str>,
        _buf: GuestPtr<u8>,
        _buf_len: types::Size,
    ) -> Result<types::Size> {
        unimplemented!("path_readlink")
    }

    fn path_remove_directory(&self, _fd: types::Fd, _path: &GuestPtr<'_, str>) -> Result<()> {
        unimplemented!("path_remove_directory")
    }

    fn path_rename(
        &self,
        _fd: types::Fd,
        _old_path: &GuestPtr<'_, str>,
        _new_fd: types::Fd,
        _new_path: &GuestPtr<'_, str>,
    ) -> Result<()> {
        unimplemented!("path_rename")
    }

    fn path_symlink(
        &self,
        _old_path: &GuestPtr<'_, str>,
        _fd: types::Fd,
        _new_path: &GuestPtr<'_, str>,
    ) -> Result<()> {
        unimplemented!("path_symlink")
    }

    fn path_unlink_file(&self, _fd: types::Fd, _path: &GuestPtr<'_, str>) -> Result<()> {
        unimplemented!("path_unlink_file")
    }

    fn poll_oneoff(
        &self,
        _in_: GuestPtr<types::Subscription>,
        _out: GuestPtr<types::Event>,
        _nsubscriptions: types::Size,
    ) -> Result<types::Size> {
        unimplemented!("poll_oneoff")
    }

    fn proc_exit(&self, _rval: types::Exitcode) -> std::result::Result<(), ()> {
        unimplemented!("proc_exit")
    }

    fn proc_raise(&self, _sig: types::Signal) -> Result<()> {
        unimplemented!("proc_raise")
    }

    fn sched_yield(&self) -> Result<()> {
        unimplemented!("sched_yield")
    }

    fn random_get(&self, _buf: GuestPtr<u8>, _buf_len: types::Size) -> Result<()> {
        unimplemented!("random_get")
    }

    fn sock_recv(
        &self,
        _fd: types::Fd,
        _ri_data: &types::IovecArray<'_>,
        _ri_flags: types::Riflags,
    ) -> Result<(types::Size, types::Roflags)> {
        unimplemented!("sock_recv")
    }

    fn sock_send(
        &self,
        _fd: types::Fd,
        _si_data: &types::CiovecArray<'_>,
        _si_flags: types::Siflags,
    ) -> Result<types::Size> {
        unimplemented!("sock_send")
    }

    fn sock_shutdown(&self, _fd: types::Fd, _how: types::Sdflags) -> Result<()> {
        unimplemented!("sock_shutdown")
    }
}
<|endoftext|>


<commit_before>use std::{convert, io};
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info.info_hash) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

struct PeerConnection {
    stream: TcpStream,
}

impl PeerConnection {
    fn connect(peer: &Peer, info_hash: &[u8]) -> Result<PeerConnection, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let mut conn = PeerConnection { stream: stream };
        try!(conn.handshake(info_hash));
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Parsing Incoming Messages<commit_after>use std::{convert, io};
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info.info_hash) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

struct PeerConnection {
    stream: TcpStream,
}

impl PeerConnection {
    fn connect(peer: &Peer, info_hash: &[u8]) -> Result<PeerConnection, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let mut conn = PeerConnection { stream: stream };
        try!(conn.handshake(info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }
}

#[derive(Debug)]
enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece,   // TODO add params
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => Message::Piece,
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }
}

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * 16777216 +
    bytes[1] as usize * 65536 +
    bytes[2] as usize * 256 +
    bytes[3] as usize
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>


<commit_before>use std::cmp::min;

use Byte;

/// Represents a range between two Byte types
#[derive(Debug, PartialEq)]
struct RangeBytes(Byte, Byte);

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Byte,
                    global_chunk_length: Byte)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Byte = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Byte = min(content_length - 1,
                            ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

#[cfg(test)]
mod test_chunk_length {
    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}<commit_msg>Add functions to download easily some chunks, for a given remote content<commit_after>use {Byte, Chunk, SChunks};
use client::GetResponse;
use hyper::client::Client;
use hyper::error::Error;
use hyper::header::{ByteRangeSpec, Headers, Range};
use std::cmp::min;
use std::io::Read;
use std::thread;

/// Represents a range between two Byte types
#[derive(Debug, PartialEq)]
struct RangeBytes(Byte, Byte);

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Byte,
                    global_chunk_length: Byte)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Byte = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Byte = min(content_length - 1,
                            ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

/// Function to get the HTTP header to send to the file server, for a chunk (specified by its index)
fn get_header_from_index(chunk_index: u64,
                         content_length: Byte,
                         global_chunk_length: Byte)
                         -> Option<(Headers, Byte)> {

    match get_chunk_length(chunk_index, content_length, global_chunk_length) {
        Some(range) => {
            let mut header = Headers::new();
            header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(range.0, range.1)]));
            Some((header, range.1 - range.0))
        }
        None => None,
    }

}

/// Function to get from the server the content of a chunk.
/// This function returns a Result type - Byte if the content of the header is accessible, an Error type otherwise.
fn download_a_chunk(http_client: &Client,
                    http_header: Headers,
                    chunk_vector: &mut Chunk,
                    url: &str)
                    -> Result<Byte, Error> {

    match http_client.get_http_response_using_headers(url, http_header) {
        Ok(mut body) => {
            match body.read_to_end(chunk_vector) {
                Ok(nb_bytes) => Ok(nb_bytes as u64),
                Err(_) => Ok(0u64),
            }
        }
        Err(http_error) => Err(http_error),
    }

}

/// Function to download each chunk of a remote content (given by its URL).
/// This function takes as parameters:
/// * the remote content length,
/// * a mutable reference to share between threads, which contains each chunk,
/// * the number of chunks that contains the remote content,
/// * the URL of the remote content server.
pub fn download_chunks(content_length: u64,
                       downloaded_chunks: &SChunks,
                       nb_chunks: u64,
                       url: &str) {

    // let mut downloaded_chunks: Arc<Mutex<Chunks>> =
    //     Arc::new(Mutex::new(Chunks::with_capacity(nb_chunks as usize)));
    let global_chunk_length: u64 = (content_length / nb_chunks) + 1;
    let mut jobs = vec![];

    for chunk_index in 0..nb_chunks {

        let (http_header, chunk_length) =
            get_header_from_index(chunk_index, content_length, global_chunk_length).unwrap();
        let mut chunk_content = Chunk::with_capacity(chunk_length as usize);
        let hyper_client = Client::new();
        let url_clone = String::from(url);
        let clone_chunks = downloaded_chunks.clone();

        jobs.push(thread::spawn(move || {
            match download_a_chunk(&hyper_client, http_header, &mut chunk_content, &url_clone) {
                Ok(bytes_written) => {
                    if bytes_written > 0 {
                        let mut shared_clone_chunks = clone_chunks.lock().unwrap();
                        shared_clone_chunks.insert(chunk_index as usize, chunk_content);
                    } else {
                        panic!("The downloaded chunk {} is empty", chunk_index);
                    }
                }
                Err(error) => {
                    panic!("Canno't download the chunk {}, due to error {}",
                           chunk_index,
                           error);
                }
            }
        }));
    }

    for child in jobs {
        let _ = child.join();
    }

}

#[cfg(test)]
mod test_chunk_length {
    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}<|endoftext|>


<commit_before>extern crate ktensor as k;
use k::{Arc, Vec2, Tensor, Context, Graph, State, Variable};
use std::path::{Path};
use std::fs::{File};
use std::io::{BufReader, Read};
use std::error::{Error};
use std::cmp;

fn read_u8(reader: &mut Read) -> u8 {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u8: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_u32(reader: &mut Read) -> u32 {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u32: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_mnist(labels_path: &Path, labels_checknum: u32, data_path: &Path, data_checknum: u32, batch_size: Option<usize>, sample_size: Option<usize>) -> Vec<(Tensor<f32>, Tensor<f32>)> {
    let labels_file = match File::open(&labels_path) {
        Err(reason) => panic!("failed to open {}: {}", labels_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let data_file = match File::open(&data_path) {
        Err(reason) => panic!("failed to open {}: {}", data_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let ref mut labels_reader = BufReader::new(labels_file);
    let ref mut data_reader = BufReader::new(data_file);

    let labels_magic = u32::from_be(read_u32(labels_reader));
    let data_magic = u32::from_be(read_u32(data_reader));

    assert_eq!(labels_magic, labels_checknum);
    assert_eq!(data_magic, data_checknum);

    let labels_count = u32::from_be(read_u32(labels_reader)) as usize;
    let data_count = u32::from_be(read_u32(data_reader)) as usize;
    let sample_count = cmp::min(labels_count, data_count);
    let sample_count = cmp::min(sample_count, sample_size.unwrap_or(sample_count));

    let batch_count = cmp::min(sample_count, batch_size.unwrap_or(sample_count));

    let rows = u32::from_be(read_u32(data_reader)) as usize;
    let columns = u32::from_be(read_u32(data_reader)) as usize;

    let mut sample_vec: Vec<(Tensor<f32>, Tensor<f32>)> = Vec::with_capacity(sample_count/batch_count);
    for _ in 0..sample_count/batch_count {
        let mut sample_data = vec![0u8; batch_count*rows*columns];
        match data_reader.read_exact(sample_data.as_mut()) {
            Err(reason) => panic!("failed to read data byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        let mut sample_labels_byte = vec![0u8; batch_count];
        match labels_reader.read_exact(sample_labels_byte.as_mut()) {
            Err(reason) => panic!("failed to read labels byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }

        let sample_data = sample_data.iter().map(|&x| x as f32 / 128.0).collect();
        let mut sample_labels: Vec<f32> = Vec::with_capacity(batch_count * 10);
        for i in sample_labels_byte {
            match i {
                0 => sample_labels.extend([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                1 => sample_labels.extend([0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                2 => sample_labels.extend([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                3 => sample_labels.extend([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                4 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                5 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                6 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0].iter().cloned()),
                7 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0].iter().cloned()),
                8 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0].iter().cloned()),
                9 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0].iter().cloned()),
                _ => (),
            }
        }

        sample_vec.push((Tensor::from_vec(Vec2(batch_count, rows*columns), sample_data), Tensor::from_vec(Vec2(batch_count, 10), sample_labels)));
    }

    sample_vec
}

#[test]
#[ignore]
fn mnist(){
    ///////////////
    // Variables //
    ///////////////

    let input_x = Arc::new(Variable::new("input_x".to_string(), Vec2(0, 28 * 28)));
    let target_y = Arc::new(Variable::new("target_y".to_string(), Vec2(0, 10)));
    let relu_threshold = Arc::new(Variable::new("relu_threshold".to_string(), Vec2(1, 1)));

    // Initialize
    let mut variable_context = Context::<f32>::with_capacity(2);
    Variable::init_f32(vec![input_x.clone(), target_y.clone(), relu_threshold.clone()], &mut variable_context);

    ///////////
    // Graph //
    ///////////

    let layers: usize = 2;
    let mut states = Vec::<Arc<State>>::with_capacity(2 * layers);
    let mut graph_head: Arc<Graph<f32>> = input_x.clone();

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 1), Vec2(28 * 28, 10)));
        let b = Arc::new(State::new(format!("weight_b_{}", 1), Vec2(1, 10)));

        let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 1), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 1), dot, b.clone()));

        let relu = Arc::new(k::op::relu_f32(format!("layer_{}_relu", 1), add, relu_threshold.clone()));

        graph_head = relu;

        states.push(w);
        states.push(b);
    }

    // {
    //     let w = Arc::new(State::new(format!("weight_w_{}", 2), Vec2(16, 10)));
    //     let b = Arc::new(State::new(format!("weight_b_{}", 2), Vec2(1, 10)));
    //
    //     let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 2), graph_head.clone(), w.clone()));
    //     let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 2), dot, b.clone()));
    //
    //     let relu = Arc::new(k::op::relu_f32(format!("layer_{}_relu", 2), add, relu_threshold.clone()));
    //
    //     graph_head = relu;
    //
    //     states.push(w);
    //     states.push(b);
    // }

    let softmax = Arc::new(k::op::softmax_f32(format!("layer_{}_softmax", 3), graph_head.clone()));
    let xentropy = Arc::new(k::cost::softmax_cross_entropy_f32(format!("layer_{}_xentropy", 3), softmax.clone(), target_y.clone()));

    // initialize states
    let mut state_context = Context::<f32>::with_capacity(2 * layers);
    State::init_f32(states, &mut state_context);

    //////////////
    // Training //
    //////////////
    {
        // Parameters
        let learning_rate = -0.0001;
        let threshold: f32 = 2.0;
        let iterations = 8192;
        let batch_size = Some(16);
        let sample_size = Some(16 * 8192);
        let print_rate = 64;

        let train_labels_path = Path::new("data/train-labels-idx1-ubyte");
        let train_data_path = Path::new("data/train-images-idx3-ubyte");
        let training_vec = read_mnist(&train_labels_path, 2049, &train_data_path, 2051, batch_size, sample_size);

        let mut history = Context::<f32>::with_capacity(5 * layers + 4);

        variable_context.set(relu_threshold.get_id(), Tensor::from_vec(Vec2(1, 1), vec![threshold]));

        let training_vec_length = training_vec.len();
        for i in 0..iterations {
            let (ref a, ref b) = training_vec[i % training_vec_length];
            variable_context.set(input_x.get_id(), a.clone());
            variable_context.set(target_y.get_id(), b.clone());

            k::train(xentropy.clone(), &mut state_context, &variable_context, &mut history, learning_rate);

            // test print
            if i % print_rate == 0 {
                println!("iteration: {} | cross entropy cost: {}", i,
                k::execute(xentropy.clone(), &state_context, &variable_context).get(Vec2(0, 0)));
            }
        }
    }

    ////////////////
    // Final test //
    ////////////////
    {
        println!("\nFinal Test:\n");
        // Parameters
        let batch_size = None;
        let sample_size = None;

        let test_labels_path = Path::new("data/t10k-labels-idx1-ubyte");
        let test_data_path = Path::new("data/t10k-images-idx3-ubyte");
        let test_vec = read_mnist(&test_labels_path, 2049, &test_data_path, 2051, batch_size, sample_size);

        let (ref a, ref b) = test_vec[0];
        variable_context.set(input_x.get_id(), a.clone());
        variable_context.set(target_y.get_id(), b.clone());

        let result_tensor = k::op::softmax_round_f32(k::execute(softmax.clone(), &state_context, &variable_context)) + (k::op::softmax_round_f32(b.clone()) * -1);
        let total = result_tensor.buffer().len() as f64;
        let mut score = 0.0;
        for &i in result_tensor.buffer().iter() {
            if i == 0 {
                score += 1.0;
            }
        }

        println!("{}", score/total);
    }
}
<commit_msg>using sigmoid<commit_after>extern crate ktensor as k;
use k::{Arc, Vec2, Tensor, Context, Graph, State, Variable};
use std::path::{Path};
use std::fs::{File};
use std::io::{BufReader, Read};
use std::error::{Error};
use std::cmp;

fn read_u8(reader: &mut Read) -> u8 {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u8: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_u32(reader: &mut Read) -> u32 {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u32: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_mnist(labels_path: &Path, labels_checknum: u32, data_path: &Path, data_checknum: u32, batch_size: Option<usize>, sample_size: Option<usize>) -> Vec<(Tensor<f32>, Tensor<f32>)> {
    let labels_file = match File::open(&labels_path) {
        Err(reason) => panic!("failed to open {}: {}", labels_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let data_file = match File::open(&data_path) {
        Err(reason) => panic!("failed to open {}: {}", data_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let ref mut labels_reader = BufReader::new(labels_file);
    let ref mut data_reader = BufReader::new(data_file);

    let labels_magic = u32::from_be(read_u32(labels_reader));
    let data_magic = u32::from_be(read_u32(data_reader));

    assert_eq!(labels_magic, labels_checknum);
    assert_eq!(data_magic, data_checknum);

    let labels_count = u32::from_be(read_u32(labels_reader)) as usize;
    let data_count = u32::from_be(read_u32(data_reader)) as usize;
    let sample_count = cmp::min(labels_count, data_count);
    let sample_count = cmp::min(sample_count, sample_size.unwrap_or(sample_count));

    let batch_count = cmp::min(sample_count, batch_size.unwrap_or(sample_count));

    let rows = u32::from_be(read_u32(data_reader)) as usize;
    let columns = u32::from_be(read_u32(data_reader)) as usize;

    let mut sample_vec: Vec<(Tensor<f32>, Tensor<f32>)> = Vec::with_capacity(sample_count/batch_count);
    for _ in 0..sample_count/batch_count {
        let mut sample_data = vec![0u8; batch_count*rows*columns];
        match data_reader.read_exact(sample_data.as_mut()) {
            Err(reason) => panic!("failed to read data byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        let mut sample_labels_byte = vec![0u8; batch_count];
        match labels_reader.read_exact(sample_labels_byte.as_mut()) {
            Err(reason) => panic!("failed to read labels byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }

        let sample_data = sample_data.iter().map(|&x| x as f32 / 128.0).collect();
        let mut sample_labels: Vec<f32> = Vec::with_capacity(batch_count * 10);
        for i in sample_labels_byte {
            match i {
                0 => sample_labels.extend([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                1 => sample_labels.extend([0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                2 => sample_labels.extend([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                3 => sample_labels.extend([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                4 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                5 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                6 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0].iter().cloned()),
                7 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0].iter().cloned()),
                8 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0].iter().cloned()),
                9 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0].iter().cloned()),
                _ => (),
            }
        }

        sample_vec.push((Tensor::from_vec(Vec2(batch_count, rows*columns), sample_data), Tensor::from_vec(Vec2(batch_count, 10), sample_labels)));
    }

    sample_vec
}

#[test]
#[ignore]
fn mnist(){
    ///////////////
    // Variables //
    ///////////////

    let input_x = Arc::new(Variable::new("input_x".to_string(), Vec2(0, 28 * 28)));
    let target_y = Arc::new(Variable::new("target_y".to_string(), Vec2(0, 10)));

    // Initialize
    let mut variable_context = Context::<f32>::with_capacity(2);
    Variable::init_f32(vec![input_x.clone(), target_y.clone()], &mut variable_context);

    ///////////
    // Graph //
    ///////////

    let layers: usize = 2;
    let mut states = Vec::<Arc<State>>::with_capacity(2 * layers);
    let mut graph_head: Arc<Graph<f32>> = input_x.clone();

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 1), Vec2(28 * 28, 10)));
        let b = Arc::new(State::new(format!("weight_b_{}", 1), Vec2(1, 10)));

        let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 1), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 1), dot, b.clone()));

        let sigmoid = Arc::new(k::op::sigmoid_f32(format!("layer_{}_sigmoid", 1), add));

        graph_head = sigmoid;

        states.push(w);
        states.push(b);
    }

    // {
    //     let w = Arc::new(State::new(format!("weight_w_{}", 2), Vec2(16, 10)));
    //     let b = Arc::new(State::new(format!("weight_b_{}", 2), Vec2(1, 10)));
    //
    //     let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 2), graph_head.clone(), w.clone()));
    //     let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 2), dot, b.clone()));
    //
    //     let sigmoid = Arc::new(k::op::sigmoid_f32(format!("layer_{}_sigmoid", 2), add));
    //
    //     graph_head = relu;
    //
    //     states.push(w);
    //     states.push(b);
    // }

    let softmax = Arc::new(k::op::softmax_f32(format!("layer_{}_softmax", 3), graph_head.clone()));
    let xentropy = Arc::new(k::cost::softmax_cross_entropy_f32(format!("layer_{}_xentropy", 3), softmax.clone(), target_y.clone()));

    // initialize states
    let mut state_context = Context::<f32>::with_capacity(2 * layers);
    State::init_f32(states, &mut state_context);

    //////////////
    // Training //
    //////////////
    {
        // Parameters
        let learning_rate = -0.0001;
        let iterations = 8192;
        let batch_size = Some(16);
        let sample_size = Some(16 * 8192);
        let print_rate = 64;

        let train_labels_path = Path::new("data/train-labels-idx1-ubyte");
        let train_data_path = Path::new("data/train-images-idx3-ubyte");
        let training_vec = read_mnist(&train_labels_path, 2049, &train_data_path, 2051, batch_size, sample_size);

        let mut history = Context::<f32>::with_capacity(5 * layers + 4);

        let training_vec_length = training_vec.len();
        for i in 0..iterations {
            let (ref a, ref b) = training_vec[i % training_vec_length];
            variable_context.set(input_x.get_id(), a.clone());
            variable_context.set(target_y.get_id(), b.clone());

            k::train(xentropy.clone(), &mut state_context, &variable_context, &mut history, learning_rate);

            // test print
            if i % print_rate == 0 {
                println!("iteration: {} | cross entropy cost: {}", i,
                k::execute(xentropy.clone(), &state_context, &variable_context).get(Vec2(0, 0)));
            }
        }
    }

    ////////////////
    // Final test //
    ////////////////
    {
        println!("\nFinal Test:\n");
        // Parameters
        let batch_size = None;
        let sample_size = None;

        let test_labels_path = Path::new("data/t10k-labels-idx1-ubyte");
        let test_data_path = Path::new("data/t10k-images-idx3-ubyte");
        let test_vec = read_mnist(&test_labels_path, 2049, &test_data_path, 2051, batch_size, sample_size);

        let (ref a, ref b) = test_vec[0];
        variable_context.set(input_x.get_id(), a.clone());
        variable_context.set(target_y.get_id(), b.clone());

        let result_tensor = k::op::softmax_round_f32(k::execute(softmax.clone(), &state_context, &variable_context)) + (k::op::softmax_round_f32(b.clone()) * -1);
        let total = result_tensor.buffer().len() as f64;
        let mut score = 0.0;
        for &i in result_tensor.buffer().iter() {
            if i == 0 {
                score += 1.0;
            }
        }

        println!("{}", score/total);
    }
}
<|endoftext|>