<commit_before>//! A module to manage protobuf serialization

use std::io::Write;

use errors::Result;
use message::MessageWrite;

use byteorder::WriteBytesExt;
use byteorder::LittleEndian as LE;

/// A struct to write protobuf messages
///
/// # Examples
///
/// ```rust
/// // an automatically generated module which is in a separate file in general
/// mod foo_bar {
///     # use quick_protobuf::{MessageWrite, Writer, Result};
///     # use std::borrow::Cow;
///     # use std::io::Write;
///     pub struct Foo<'a> { pub name: Option<Cow<'a, str>>, }
///     pub struct Bar { pub id: Option<u32> }
///     pub struct FooBar<'a> { pub foos: Vec<Foo<'a>>, pub bars: Vec<Bar>, }
///     impl<'a> MessageWrite for FooBar<'a> {
///         // implements
///         // fn get_size(&self) -> usize { ... }
///         // fn write_message<W: Write>(&self, r: &mut Writer<W>) -> Result<()> { ... }
///         # fn get_size(&self) -> usize { 0 }
///         # fn write_message<W: Write>(&self, _: &mut Writer<W>) -> Result<()> { Ok(()) }
///     }
/// }
///
/// // FooBar is a message generated from a proto file
/// // in parcicular it contains a `write_message` function
/// use foo_bar::{FooBar, Foo, Bar};
/// use std::fs::File;
/// use std::borrow::Cow;
/// use quick_protobuf::{MessageWrite, Writer};
///
/// fn main() {
///     // let's say we want to write to a file
///     let mut file: File;
///     # let mut file = Vec::new();
///     let mut writer = Writer::new(&mut file);
///
///     // manually generates a FooBar for the example
///     let foobar = FooBar {
///         foos: vec![Foo { name: Some(Cow::Borrowed("test!")) }, Foo { name: None }],
///         bars: vec![Bar { id: Some(43) }, Bar { id: None }],
///     };
///
///     // now using the generated module
///     writer.write_message(&foobar).expect("Cannot write FooBar");
/// }
/// ```
pub struct Writer<W> {
    inner: W,
}

impl<W: Write> Writer<W> {

    /// Creates a new `ProtobufWriter`
    pub fn new(w: W) -> Writer<W> {
        Writer { inner: w }
    }

    /// Writes a `varint` (compacted `u64`)
    pub fn write_varint(&mut self, mut v: u64) -> Result<()> {
        while v > 0x7F {
            self.inner.write_u8(((v as u8) & 0x7F) | 0x80)?;
            v >>= 7;
        }
        self.inner.write_u8(v as u8).map_err(|e| e.into())
    }

    /// Writes a tag, which represents both the field number and the wire type
    pub fn write_tag(&mut self, tag: u32) -> Result<()> {
        self.write_varint(tag as u64)
    }

    /// Writes a `int32` which is internally coded as a `varint`
    pub fn write_int32(&mut self, v: i32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `int64` which is internally coded as a `varint`
    pub fn write_int64(&mut self, v: i64) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint32` which is internally coded as a `varint`
    pub fn write_uint32(&mut self, v: u32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint64` which is internally coded as a `varint`
    pub fn write_uint64(&mut self, v: u64) -> Result<()> {
        self.write_varint(v)
    }

    /// Writes a `sint32` which is internally coded as a `varint`
    pub fn write_sint32(&mut self, v: i32) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 31)) as u64)
    }

    /// Writes a `sint64` which is internally coded as a `varint`
    pub fn write_sint64(&mut self, v: i64) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 63)) as u64)
    }

    /// Writes a `fixed64` which is little endian coded `u64`
    pub fn write_fixed64(&mut self, v: u64) -> Result<()> {
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `fixed32` which is little endian coded `u32`
    pub fn write_fixed32(&mut self, v: u32) -> Result<()> {
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed64` which is little endian coded `i64`
    pub fn write_sfixed64(&mut self, v: i64) -> Result<()> {
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed32` which is little endian coded `i32`
    pub fn write_sfixed32(&mut self, v: i32) -> Result<()> {
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `float`
    pub fn write_float(&mut self, v: f32) -> Result<()> {
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `double`
    pub fn write_double(&mut self, v: f64) -> Result<()> {
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `bool` 1 = true, 0 = false
    pub fn write_bool(&mut self, v: bool) -> Result<()> {
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes an `enum` converting it to a `i32` first
    pub fn write_enum(&mut self, v: i32) -> Result<()> {
        self.write_int32(v)
    }

    /// Writes `bytes`: length first then the chunk of data
    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes `string`: length first then the chunk of data
    pub fn write_string(&mut self, s: &str) -> Result<()> {
        self.write_bytes(s.as_bytes())
    }

    /// Writes packed repeated field: length first then the chunk of data
    pub fn write_packed_repeated_field<M, F, S>(&mut self, v: &[M], mut write: F, size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes packed repeated field when we know the size of items
    ///
    /// `item_size` is internally used to compute the total length
    /// As the length is fixed (and the same as rust internal representation, we can directly dump
    /// all data at once
    pub fn write_packed_fixed_size<M>(&mut self, v: &[M], item_size: usize) -> Result<()> {
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes a message which implements `MessageWrite`
    pub fn write_message<M: MessageWrite>(&mut self, m: &M) -> Result<()> {
        let len = m.get_size();
        self.write_varint(len as u64)?;
        m.write_message(self)
    }

    /// Writes tag then `int32`
    pub fn write_int32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `int64`
    pub fn write_int64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint32`
    pub fn write_uint32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint64`
    pub fn write_uint64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v)
    }

    /// Writes tag then `sint32`
    pub fn write_sint32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint32(v)
    }

    /// Writes tag then `sint64`
    pub fn write_sint64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint64(v)
    }

    /// Writes tag then `fixed64`
    pub fn write_fixed64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `fixed32`
    pub fn write_fixed32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed64`
    pub fn write_sfixed64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed32`
    pub fn write_sfixed32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `float`
    pub fn write_float_with_tag(&mut self, tag: u32, v: f32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `double`
    pub fn write_double_with_tag(&mut self, tag: u32, v: f64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `bool`
    pub fn write_bool_with_tag(&mut self, tag: u32, v: bool) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes tag then `bytes`
    pub fn write_bytes_with_tag(&mut self, tag: u32, bytes: &[u8]) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes tag then `string`
    pub fn write_string_with_tag(&mut self, tag: u32, s: &str) -> Result<()> {
        self.write_tag(tag)?;
        self.write_bytes(s.as_bytes())
    }

    /// Writes tag then repeated field
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_repeated_field_with_tag<M, F, S>(&mut self,
                                                         tag: u32,
                                                         v: &[M],
                                                         mut write: F,
                                                         size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }

        self.write_tag(tag)?;
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes tag then repeated field with fixed length item size
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_fixed_size_with_tag<M>(&mut self,
                                               tag: u32,
                                               v: &[M],
                                               item_size: usize) -> Result<()> {
        if v.is_empty() {
            return Ok(());
        }
        self.write_tag(tag)?;
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes tag then message
    pub fn write_message_with_tag<M: MessageWrite>(&mut self, tag: u32, m: &M) -> Result<()> {
        self.write_tag(tag)?;
        self.write_message(m)
    }

    /// Writes tag then enum
    pub fn write_enum_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_int32(v)
    }
}
<commit_msg>fix doc on beta<commit_after>//! A module to manage protobuf serialization

use std::io::Write;

use errors::Result;
use message::MessageWrite;

use byteorder::WriteBytesExt;
use byteorder::LittleEndian as LE;

/// A struct to write protobuf messages
///
/// # Examples
///
/// ```rust
///
/// // an automatically generated module which is in a separate file in general
/// mod foo_bar {
///     # use quick_protobuf::{MessageWrite, Writer, Result};
///     # use std::borrow::Cow;
///     # use std::io::Write;
///     pub struct Foo<'a> { pub name: Option<Cow<'a, str>>, }
///     pub struct Bar { pub id: Option<u32> }
///     pub struct FooBar<'a> { pub foos: Vec<Foo<'a>>, pub bars: Vec<Bar>, }
///     impl<'a> MessageWrite for FooBar<'a> {
///         // implements
///         // fn get_size(&self) -> usize { ... }
///         // fn write_message<W: Write>(&self, r: &mut Writer<W>) -> Result<()> { ... }
///         # fn get_size(&self) -> usize { 0 }
///         # fn write_message<W: Write>(&self, _: &mut Writer<W>) -> Result<()> { Ok(()) }
///     }
/// }
///
/// // FooBar is a message generated from a proto file
/// // in parcicular it contains a `write_message` function
/// use foo_bar::{FooBar, Foo, Bar};
/// use std::borrow::Cow;
/// use quick_protobuf::Writer;
///
/// fn main() {
///     // let mut r = File::create("...").expect("Cannot create file");
///     // for the sake of example, we'll use a simpler struct which impl `Write`
///     let mut r = Vec::new();
///     let mut writer = Writer::new(&mut r);
///
///     // manually generates a FooBar for the example
///     let foobar = FooBar {
///         foos: vec![Foo { name: Some(Cow::Borrowed("test!")) }, Foo { name: None }],
///         bars: vec![Bar { id: Some(43) }, Bar { id: None }],
///     };
///
///     // now using the generated module
///     writer.write_message(&foobar).expect("Cannot write FooBar");
/// }
/// ```
pub struct Writer<W> {
    inner: W,
}

impl<W: Write> Writer<W> {

    /// Creates a new `ProtobufWriter`
    pub fn new(w: W) -> Writer<W> {
        Writer { inner: w }
    }

    /// Writes a `varint` (compacted `u64`)
    pub fn write_varint(&mut self, mut v: u64) -> Result<()> {
        while v > 0x7F {
            self.inner.write_u8(((v as u8) & 0x7F) | 0x80)?;
            v >>= 7;
        }
        self.inner.write_u8(v as u8).map_err(|e| e.into())
    }

    /// Writes a tag, which represents both the field number and the wire type
    pub fn write_tag(&mut self, tag: u32) -> Result<()> {
        self.write_varint(tag as u64)
    }

    /// Writes a `int32` which is internally coded as a `varint`
    pub fn write_int32(&mut self, v: i32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `int64` which is internally coded as a `varint`
    pub fn write_int64(&mut self, v: i64) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint32` which is internally coded as a `varint`
    pub fn write_uint32(&mut self, v: u32) -> Result<()> {
        self.write_varint(v as u64)
    }

    /// Writes a `uint64` which is internally coded as a `varint`
    pub fn write_uint64(&mut self, v: u64) -> Result<()> {
        self.write_varint(v)
    }

    /// Writes a `sint32` which is internally coded as a `varint`
    pub fn write_sint32(&mut self, v: i32) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 31)) as u64)
    }

    /// Writes a `sint64` which is internally coded as a `varint`
    pub fn write_sint64(&mut self, v: i64) -> Result<()> {
        self.write_varint(((v << 1) ^ (v >> 63)) as u64)
    }

    /// Writes a `fixed64` which is little endian coded `u64`
    pub fn write_fixed64(&mut self, v: u64) -> Result<()> {
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `fixed32` which is little endian coded `u32`
    pub fn write_fixed32(&mut self, v: u32) -> Result<()> {
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed64` which is little endian coded `i64`
    pub fn write_sfixed64(&mut self, v: i64) -> Result<()> {
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `sfixed32` which is little endian coded `i32`
    pub fn write_sfixed32(&mut self, v: i32) -> Result<()> {
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `float`
    pub fn write_float(&mut self, v: f32) -> Result<()> {
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `double`
    pub fn write_double(&mut self, v: f64) -> Result<()> {
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes a `bool` 1 = true, 0 = false
    pub fn write_bool(&mut self, v: bool) -> Result<()> {
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes an `enum` converting it to a `i32` first
    pub fn write_enum(&mut self, v: i32) -> Result<()> {
        self.write_int32(v)
    }

    /// Writes `bytes`: length first then the chunk of data
    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes `string`: length first then the chunk of data
    pub fn write_string(&mut self, s: &str) -> Result<()> {
        self.write_bytes(s.as_bytes())
    }

    /// Writes packed repeated field: length first then the chunk of data
    pub fn write_packed_repeated_field<M, F, S>(&mut self, v: &[M], mut write: F, size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes packed repeated field when we know the size of items
    ///
    /// `item_size` is internally used to compute the total length
    /// As the length is fixed (and the same as rust internal representation, we can directly dump
    /// all data at once
    pub fn write_packed_fixed_size<M>(&mut self, v: &[M], item_size: usize) -> Result<()> {
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes a message which implements `MessageWrite`
    pub fn write_message<M: MessageWrite>(&mut self, m: &M) -> Result<()> {
        let len = m.get_size();
        self.write_varint(len as u64)?;
        m.write_message(self)
    }

    /// Writes tag then `int32`
    pub fn write_int32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `int64`
    pub fn write_int64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint32`
    pub fn write_uint32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v as u64)
    }

    /// Writes tag then `uint64`
    pub fn write_uint64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(v)
    }

    /// Writes tag then `sint32`
    pub fn write_sint32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint32(v)
    }

    /// Writes tag then `sint64`
    pub fn write_sint64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.write_sint64(v)
    }

    /// Writes tag then `fixed64`
    pub fn write_fixed64_with_tag(&mut self, tag: u32, v: u64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `fixed32`
    pub fn write_fixed32_with_tag(&mut self, tag: u32, v: u32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_u32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed64`
    pub fn write_sfixed64_with_tag(&mut self, tag: u32, v: i64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `sfixed32`
    pub fn write_sfixed32_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_i32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `float`
    pub fn write_float_with_tag(&mut self, tag: u32, v: f32) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f32::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `double`
    pub fn write_double_with_tag(&mut self, tag: u32, v: f64) -> Result<()> {
        self.write_tag(tag)?;
        self.inner.write_f64::<LE>(v).map_err(|e| e.into())
    }

    /// Writes tag then `bool`
    pub fn write_bool_with_tag(&mut self, tag: u32, v: bool) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(if v { 1 } else { 0 })
    }

    /// Writes tag then `bytes`
    pub fn write_bytes_with_tag(&mut self, tag: u32, bytes: &[u8]) -> Result<()> {
        self.write_tag(tag)?;
        self.write_varint(bytes.len() as u64)?;
        self.inner.write_all(bytes).map_err(|e| e.into())
    }

    /// Writes tag then `string`
    pub fn write_string_with_tag(&mut self, tag: u32, s: &str) -> Result<()> {
        self.write_tag(tag)?;
        self.write_bytes(s.as_bytes())
    }

    /// Writes tag then repeated field
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_repeated_field_with_tag<M, F, S>(&mut self,
                                                         tag: u32,
                                                         v: &[M],
                                                         mut write: F,
                                                         size: &S) -> Result<()>
        where F: FnMut(&mut Self, &M) -> Result<()>,
              S: Fn(&M) -> usize,
    {
        if v.is_empty() {
            return Ok(());
        }

        self.write_tag(tag)?;
        let len: usize = v.iter().map(|m| size(m)).sum();
        self.write_varint(len as u64)?;
        for m in v {
            write(self, m)?;
        }
        Ok(())
    }

    /// Writes tag then repeated field with fixed length item size
    ///
    /// If array is empty, then do nothing (do not even write the tag)
    pub fn write_packed_fixed_size_with_tag<M>(&mut self,
                                               tag: u32,
                                               v: &[M],
                                               item_size: usize) -> Result<()> {
        if v.is_empty() {
            return Ok(());
        }
        self.write_tag(tag)?;
        let len = v.len() * item_size;
        let bytes = unsafe { ::std::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
        self.write_bytes(bytes)
    }

    /// Writes tag then message
    pub fn write_message_with_tag<M: MessageWrite>(&mut self, tag: u32, m: &M) -> Result<()> {
        self.write_tag(tag)?;
        self.write_message(m)
    }

    /// Writes tag then enum
    pub fn write_enum_with_tag(&mut self, tag: u32, v: i32) -> Result<()> {
        self.write_tag(tag)?;
        self.write_int32(v)
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::fmt::{Show,Formatter};
use std::str;
use std::io::{UserDir,AllPermissions};
use std::io::fs::{mkdir_recursive,rmdir_recursive,chmod};
use serialize::{Encodable,Encoder};

use core::source::{Location, Local, Remote};
use util::{CargoResult, ChainError, ProcessBuilder, process, human};

#[deriving(PartialEq,Clone,Encodable)]
pub enum GitReference {
    Master,
    Other(String)
}

impl GitReference {
    pub fn for_str<S: Str>(string: S) -> GitReference {
        if string.as_slice() == "master" {
            Master
        } else {
            Other(string.as_slice().to_str())
        }
    }
}

impl Str for GitReference {
    fn as_slice<'a>(&'a self) -> &'a str {
        match *self {
            Master => "master",
            Other(ref string) => string.as_slice()
        }
    }
}

impl Show for GitReference {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

macro_rules! git(
    ($config:expr, $str:expr, $($rest:expr),*) => (
        try!(git_inherit(&$config, format!($str, $($rest),*)))
    );

    ($config:expr, $str:expr) => (
        try!(git_inherit(&$config, format!($str)))
    );
)

macro_rules! git_output(
    ($config:expr, $str:expr, $($rest:expr),*) => ({
        try!(git_output(&$config, format!($str, $($rest),*)))
    });

    ($config:expr, $str:expr) => ({
        try!(git_output(&$config, format!($str)))
    });
)

macro_rules! errln(
    ($($arg:tt)*) => (let _ = writeln!(::std::io::stdio::stderr(), $($arg)*))
)

/// GitRemote represents a remote repository. It gets cloned into a local
/// GitDatabase.
#[deriving(PartialEq,Clone,Show)]
pub struct GitRemote {
    location: Location,
}

#[deriving(PartialEq,Clone,Encodable)]
struct EncodableGitRemote {
    location: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitRemote {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitRemote {
            location: self.location.to_str()
        }.encode(s)
    }
}

/// GitDatabase is a local clone of a remote repository's database. Multiple
/// GitCheckouts can be cloned from this GitDatabase.
#[deriving(PartialEq,Clone)]
pub struct GitDatabase {
    remote: GitRemote,
    path: Path,
}

#[deriving(Encodable)]
pub struct EncodableGitDatabase {
    remote: GitRemote,
    path: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitDatabase {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitDatabase {
            remote: self.remote.clone(),
            path: self.path.display().to_str()
        }.encode(s)
    }
}

/// GitCheckout is a local checkout of a particular revision. Calling
/// `clone_into` with a reference will resolve the reference into a revision,
/// and return a CargoError if no revision for that reference was found.
pub struct GitCheckout {
    database: GitDatabase,
    location: Path,
    reference: GitReference,
    revision: String,
}

#[deriving(Encodable)]
pub struct EncodableGitCheckout {
    database: GitDatabase,
    location: String,
    reference: String,
    revision: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitCheckout {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitCheckout {
            database: self.database.clone(),
            location: self.location.display().to_str(),
            reference: self.reference.to_str(),
            revision: self.revision.to_str()
        }.encode(s)
    }
}

// Implementations

impl GitRemote {
    pub fn new(location: &Location) -> GitRemote {
        GitRemote { location: location.clone() }
    }

    pub fn get_location<'a>(&'a self) -> &'a Location {
        &self.location
    }

    pub fn has_ref<S: Str>(&self, path: &Path, reference: S) -> CargoResult<()> {
        git_output!(*path, "rev-parse {}", reference.as_slice());
        Ok(())
    }

    pub fn checkout(&self, into: &Path) -> CargoResult<GitDatabase> {
        if into.exists() {
            try!(self.fetch_into(into));
        } else {
            try!(self.clone_into(into));
        }

        Ok(GitDatabase { remote: self.clone(), path: into.clone() })
    }

    pub fn db_at(&self, db_path: &Path) -> GitDatabase {
        GitDatabase { remote: self.clone(), path: db_path.clone() }
    }

    fn fetch_into(&self, path: &Path) -> CargoResult<()> {
        Ok(git!(*path, "fetch --force --quiet --tags {} \
                        refs/heads/*:refs/heads/*", self.fetch_location()))
    }

    fn clone_into(&self, path: &Path) -> CargoResult<()> {
        let dirname = Path::new(path.dirname());

        try!(mkdir_recursive(path, UserDir));

        Ok(git!(dirname, "clone {} {} --bare --no-hardlinks --quiet",
                self.fetch_location(), path.display()))
    }

    fn fetch_location(&self) -> String {
        match self.location {
            Local(ref p) => p.display().to_str(),
            Remote(ref u) => u.to_str(),
        }
    }
}

impl GitDatabase {
    fn get_path<'a>(&'a self) -> &'a Path {
        &self.path
    }

    pub fn copy_to<S: Str>(&self, reference: S,
                           dest: &Path) -> CargoResult<GitCheckout> {
        let checkout = try!(GitCheckout::clone_into(dest, self.clone(),
                                  GitReference::for_str(reference.as_slice())));

        try!(checkout.fetch());
        try!(checkout.update_submodules());

        Ok(checkout)
    }

    pub fn rev_for<S: Str>(&self, reference: S) -> CargoResult<String> {
        Ok(git_output!(self.path, "rev-parse {}", reference.as_slice()))
    }

}

impl GitCheckout {
    fn clone_into(into: &Path, database: GitDatabase,
                  reference: GitReference) -> CargoResult<GitCheckout> {
        let revision = try!(database.rev_for(reference.as_slice()));
        let checkout = GitCheckout {
            location: into.clone(),
            database: database,
            reference: reference,
            revision: revision,
        };

        // If the git checkout already exists, we don't need to clone it again
        if !checkout.location.join(".git").exists() {
            try!(checkout.clone_repo());
        }

        Ok(checkout)
    }

    fn get_source<'a>(&'a self) -> &'a Path {
        self.database.get_path()
    }

    fn clone_repo(&self) -> CargoResult<()> {
        let dirname = Path::new(self.location.dirname());

        try!(mkdir_recursive(&dirname, UserDir).chain_error(|| {
            human(format!("Couldn't mkdir {}",
                          Path::new(self.location.dirname()).display()))
        }));

        if self.location.exists() {
            try!(rmdir_recursive(&self.location).chain_error(|| {
                human(format!("Couldn't rmdir {}",
                              Path::new(&self.location).display()))
            }));
        }

        git!(dirname, "clone --no-checkout --quiet {} {}",
             self.get_source().display(), self.location.display());
        try!(chmod(&self.location, AllPermissions));

        Ok(())
    }

    fn fetch(&self) -> CargoResult<()> {
        // In git 1.8, apparently --tags explicitly *only* fetches tags, it does
        // not fetch anything else. In git 1.9, however, git apparently fetches
        // everything when --tags is passed.
        //
        // This means that if we want to fetch everything we need to execute
        // both with and without --tags on 1.8 (apparently), and only with
        // --tags on 1.9. For simplicity, we execute with and without --tags for
        // all gits.
        //
        // FIXME: This is suspicious. I have been informated that, for example,
        //        bundler does not do this, yet bundler appears to work!
        git!(self.location, "fetch --force --quiet {}",
             self.get_source().display());
        git!(self.location, "fetch --force --quiet --tags {}",
             self.get_source().display());
        try!(self.reset(self.revision.as_slice()));
        Ok(())
    }

    fn reset<T: Show>(&self, revision: T) -> CargoResult<()> {
        Ok(git!(self.location, "reset -q --hard {}", revision))
    }

    fn update_submodules(&self) -> CargoResult<()> {
        Ok(git!(self.location, "submodule update --init --recursive --quiet"))
    }
}

fn git(path: &Path, str: &str) -> ProcessBuilder {
    debug!("Executing git {} @ {}", str, path.display());

    process("git").args(str.split(' ').collect::<Vec<&str>>().as_slice())
                  .cwd(path.clone())
}

fn git_inherit(path: &Path, str: String) -> CargoResult<()> {
    git(path, str.as_slice()).exec().chain_error(|| {
        human(format!("Executing `git {}` failed", str))
    })
}

fn git_output(path: &Path, str: String) -> CargoResult<String> {
    let output = try!(git(path, str.as_slice()).exec_with_output()
                                                     .chain_error(||
        human(format!("Executing `git {}` failed", str))));

    Ok(to_str(output.output.as_slice()).as_slice().trim_right().to_str())
}

fn to_str(vec: &[u8]) -> String {
    str::from_utf8_lossy(vec).to_str()
}

<commit_msg>auto merge of #101 : alexcrichton/cargo/fix-git-1.7, r=wycats<commit_after>use std::fmt;
use std::fmt::{Show,Formatter};
use std::str;
use std::io::{UserDir,AllPermissions};
use std::io::fs::{mkdir_recursive,rmdir_recursive,chmod};
use serialize::{Encodable,Encoder};

use core::source::{Location, Local, Remote};
use util::{CargoResult, ChainError, ProcessBuilder, process, human};

#[deriving(PartialEq,Clone,Encodable)]
pub enum GitReference {
    Master,
    Other(String)
}

impl GitReference {
    pub fn for_str<S: Str>(string: S) -> GitReference {
        if string.as_slice() == "master" {
            Master
        } else {
            Other(string.as_slice().to_str())
        }
    }
}

impl Str for GitReference {
    fn as_slice<'a>(&'a self) -> &'a str {
        match *self {
            Master => "master",
            Other(ref string) => string.as_slice()
        }
    }
}

impl Show for GitReference {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

macro_rules! git(
    ($config:expr, $str:expr, $($rest:expr),*) => (
        try!(git_inherit(&$config, format!($str, $($rest),*)))
    );

    ($config:expr, $str:expr) => (
        try!(git_inherit(&$config, format!($str)))
    );
)

macro_rules! git_output(
    ($config:expr, $str:expr, $($rest:expr),*) => ({
        try!(git_output(&$config, format!($str, $($rest),*)))
    });

    ($config:expr, $str:expr) => ({
        try!(git_output(&$config, format!($str)))
    });
)

macro_rules! errln(
    ($($arg:tt)*) => (let _ = writeln!(::std::io::stdio::stderr(), $($arg)*))
)

/// GitRemote represents a remote repository. It gets cloned into a local
/// GitDatabase.
#[deriving(PartialEq,Clone,Show)]
pub struct GitRemote {
    location: Location,
}

#[deriving(PartialEq,Clone,Encodable)]
struct EncodableGitRemote {
    location: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitRemote {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitRemote {
            location: self.location.to_str()
        }.encode(s)
    }
}

/// GitDatabase is a local clone of a remote repository's database. Multiple
/// GitCheckouts can be cloned from this GitDatabase.
#[deriving(PartialEq,Clone)]
pub struct GitDatabase {
    remote: GitRemote,
    path: Path,
}

#[deriving(Encodable)]
pub struct EncodableGitDatabase {
    remote: GitRemote,
    path: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitDatabase {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitDatabase {
            remote: self.remote.clone(),
            path: self.path.display().to_str()
        }.encode(s)
    }
}

/// GitCheckout is a local checkout of a particular revision. Calling
/// `clone_into` with a reference will resolve the reference into a revision,
/// and return a CargoError if no revision for that reference was found.
pub struct GitCheckout {
    database: GitDatabase,
    location: Path,
    reference: GitReference,
    revision: String,
}

#[deriving(Encodable)]
pub struct EncodableGitCheckout {
    database: GitDatabase,
    location: String,
    reference: String,
    revision: String,
}

impl<E, S: Encoder<E>> Encodable<S, E> for GitCheckout {
    fn encode(&self, s: &mut S) -> Result<(), E> {
        EncodableGitCheckout {
            database: self.database.clone(),
            location: self.location.display().to_str(),
            reference: self.reference.to_str(),
            revision: self.revision.to_str()
        }.encode(s)
    }
}

// Implementations

impl GitRemote {
    pub fn new(location: &Location) -> GitRemote {
        GitRemote { location: location.clone() }
    }

    pub fn get_location<'a>(&'a self) -> &'a Location {
        &self.location
    }

    pub fn has_ref<S: Str>(&self, path: &Path, reference: S) -> CargoResult<()> {
        git_output!(*path, "rev-parse {}", reference.as_slice());
        Ok(())
    }

    pub fn checkout(&self, into: &Path) -> CargoResult<GitDatabase> {
        if into.exists() {
            try!(self.fetch_into(into));
        } else {
            try!(self.clone_into(into));
        }

        Ok(GitDatabase { remote: self.clone(), path: into.clone() })
    }

    pub fn db_at(&self, db_path: &Path) -> GitDatabase {
        GitDatabase { remote: self.clone(), path: db_path.clone() }
    }

    fn fetch_into(&self, path: &Path) -> CargoResult<()> {
        Ok(git!(*path, "fetch --force --quiet --tags {} \
                        refs/heads/*:refs/heads/*", self.fetch_location()))
    }

    fn clone_into(&self, path: &Path) -> CargoResult<()> {
        let dirname = Path::new(path.dirname());

        try!(mkdir_recursive(path, UserDir));

        Ok(git!(dirname, "clone {} {} --bare --no-hardlinks --quiet",
                self.fetch_location(), path.display()))
    }

    fn fetch_location(&self) -> String {
        match self.location {
            Local(ref p) => p.display().to_str(),
            Remote(ref u) => u.to_str(),
        }
    }
}

impl GitDatabase {
    fn get_path<'a>(&'a self) -> &'a Path {
        &self.path
    }

    pub fn copy_to<S: Str>(&self, reference: S,
                           dest: &Path) -> CargoResult<GitCheckout> {
        let checkout = try!(GitCheckout::clone_into(dest, self.clone(),
                                  GitReference::for_str(reference.as_slice())));

        try!(checkout.fetch());
        try!(checkout.update_submodules());

        Ok(checkout)
    }

    pub fn rev_for<S: Str>(&self, reference: S) -> CargoResult<String> {
        Ok(git_output!(self.path, "rev-parse {}", reference.as_slice()))
    }

}

impl GitCheckout {
    fn clone_into(into: &Path, database: GitDatabase,
                  reference: GitReference) -> CargoResult<GitCheckout> {
        let revision = try!(database.rev_for(reference.as_slice()));
        let checkout = GitCheckout {
            location: into.clone(),
            database: database,
            reference: reference,
            revision: revision,
        };

        // If the git checkout already exists, we don't need to clone it again
        if !checkout.location.join(".git").exists() {
            try!(checkout.clone_repo());
        }

        Ok(checkout)
    }

    fn get_source<'a>(&'a self) -> &'a Path {
        self.database.get_path()
    }

    fn clone_repo(&self) -> CargoResult<()> {
        let dirname = Path::new(self.location.dirname());

        try!(mkdir_recursive(&dirname, UserDir).chain_error(|| {
            human(format!("Couldn't mkdir {}",
                          Path::new(self.location.dirname()).display()))
        }));

        if self.location.exists() {
            try!(rmdir_recursive(&self.location).chain_error(|| {
                human(format!("Couldn't rmdir {}",
                              Path::new(&self.location).display()))
            }));
        }

        git!(dirname, "clone --no-checkout --quiet {} {}",
             self.get_source().display(), self.location.display());
        try!(chmod(&self.location, AllPermissions));

        Ok(())
    }

    fn fetch(&self) -> CargoResult<()> {
        // In git 1.8, apparently --tags explicitly *only* fetches tags, it does
        // not fetch anything else. In git 1.9, however, git apparently fetches
        // everything when --tags is passed.
        //
        // This means that if we want to fetch everything we need to execute
        // both with and without --tags on 1.8 (apparently), and only with
        // --tags on 1.9. For simplicity, we execute with and without --tags for
        // all gits.
        //
        // FIXME: This is suspicious. I have been informed that, for example,
        //        bundler does not do this, yet bundler appears to work!
        //
        // And to continue the fun, git before 1.7.3 had the fun bug that if a
        // branch was tracking a remote, then `git fetch $url` doesn't work!
        //
        // For details, see
        // https://www.kernel.org/pub/software/scm/git/docs/RelNotes-1.7.3.txt
        //
        // In this case we just use `origin` here instead of the database path.
        git!(self.location, "fetch --force --quiet origin");
        git!(self.location, "fetch --force --quiet --tags origin");
        try!(self.reset(self.revision.as_slice()));
        Ok(())
    }

    fn reset<T: Show>(&self, revision: T) -> CargoResult<()> {
        Ok(git!(self.location, "reset -q --hard {}", revision))
    }

    fn update_submodules(&self) -> CargoResult<()> {
        Ok(git!(self.location, "submodule update --init --recursive --quiet"))
    }
}

fn git(path: &Path, str: &str) -> ProcessBuilder {
    debug!("Executing git {} @ {}", str, path.display());

    process("git").args(str.split(' ').collect::<Vec<&str>>().as_slice())
                  .cwd(path.clone())
}

fn git_inherit(path: &Path, str: String) -> CargoResult<()> {
    git(path, str.as_slice()).exec().chain_error(|| {
        human(format!("Executing `git {}` failed", str))
    })
}

fn git_output(path: &Path, str: String) -> CargoResult<String> {
    let output = try!(git(path, str.as_slice()).exec_with_output()
                                                     .chain_error(||
        human(format!("Executing `git {}` failed", str))));

    Ok(to_str(output.output.as_slice()).as_slice().trim_right().to_str())
}

fn to_str(vec: &[u8]) -> String {
    str::from_utf8_lossy(vec).to_str()
}

<|endoftext|>


<commit_before>use crate::codec::*;
use crate::error::*;
use crate::*;
use crate::{SocketType, ZmqResult};
use async_trait::async_trait;
use crossbeam::atomic::AtomicCell;
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use futures::lock::Mutex;
use futures::stream::FuturesUnordered;
use futures_util::sink::SinkExt;
use std::sync::Arc;
use std::time::Duration;
use tokio::stream::StreamExt;

struct ReqSocketBackend {
    pub(crate) peers: Arc<DashMap<PeerIdentity, Peer>>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    pub(crate) current_request_peer_id: AtomicCell<usize>,
}

pub struct ReqSocket {
    backend: Arc<ReqSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl BlockingSend for ReqSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        if self.current_request.is_some() {
            return Err(ZmqError::Socket(
                "Unable to send message. Request already in progress",
            ));
        }
        // In normal scenario this will always be only 1 iteration
        // There can be special case when peer has disconnected and his id is still in RR queue
        // This happens because SegQueue don't have an api to delete items from queue.
        // So in such case we'll just pop item and skip it if we don't have a matching peer in peers map
        loop {
            let next_peer_id = match self.backend.round_robin.pop() {
                Ok(peer) => peer,
                Err(_) => {
                    return Err(ZmqError::Other(
                        "Not connected to peers. Unable to send messages",
                    ))
                }
            };
            match self.backend.peers.get_mut(&next_peer_id) {
                Some(mut peer) => {
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    self.current_request = Some(next_peer_id.clone());
                    self.backend
                        .current_request_peer_id
                        .store(self.backend.peers.hash_usize(&next_peer_id));
                    self.backend.round_robin.push(next_peer_id);
                    return Ok(());
                }
                None => continue,
            }
        }
    }
}

#[async_trait]
impl BlockingRecv for ReqSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(recv_queue) = self
                    .backend
                    .peers
                    .get(&peer_id)
                    .map(|p| p.recv_queue.clone())
                {
                    let message = recv_queue.lock().await.next().await;
                    match message {
                        Some(Message::MultipartMessage(mut message)) => {
                            assert!(message.len() == 2);
                            assert!(message[0].data.is_empty()); // Ensure that we have delimeter as first part
                            Ok(message.pop().unwrap())
                        }
                        Some(_) => Err(ZmqError::Other("Wrong message type received")),
                        None => Err(ZmqError::NoMessage),
                    }
                } else {
                    Err(ZmqError::Other("Server disconnected"))
                }
            }
            None => Err(ZmqError::Other("Unable to recv. No request in progress")),
        }
    }
}

#[async_trait]
impl SocketFrontend for ReqSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(ReqSocketBackend {
                peers: Arc::new(DashMap::new()),
                round_robin: SegQueue::new(),
                current_request_peer_id: AtomicCell::new(0),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        if self._accept_close_handle.is_some() {
            return Err(ZmqError::Other(
                "Socket server already started. Currently only one server is supported",
            ));
        }
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        let addr = endpoint.parse::<SocketAddr>()?;
        let raw_socket = tokio::net::TcpStream::connect(addr).await?;
        util::peer_connected(raw_socket, self.backend.clone()).await;
        Ok(())
    }
}

impl MultiPeer for ReqSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 100;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );
        self.round_robin.push(peer_id.clone());

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for ReqSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        let hash = self.peers.hash_usize(&peer_id);
        let current = self.current_request_peer_id.compare_and_swap(hash, 0);
        // This is needed to ensure that we only store messages that we are expecting to get
        // Other messages are silently discarded according to spec
        if current == hash {
            // We've got reply that we were waiting for
            self.peers
                .get_mut(peer_id)
                .expect("Not found peer by id")
                .recv_queue_in
                .send(message)
                .await
                .expect("Failed to send");
        }
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REQ
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

struct RepSocketBackend {
    pub(crate) peers: Arc<DashMap<PeerIdentity, Peer>>,
}

pub struct RepSocket {
    backend: Arc<RepSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl SocketFrontend for RepSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(RepSocketBackend {
                peers: Arc::new(DashMap::new()),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        unimplemented!()
    }
}

impl MultiPeer for RepSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 100;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for RepSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        self.peers
            .get_mut(peer_id)
            .expect("Not found peer by id")
            .recv_queue_in
            .send(message)
            .await
            .expect("Failed to send");
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REP
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

#[async_trait]
impl BlockingSend for RepSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(mut peer) = self.backend.peers.get_mut(&peer_id) {
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    Ok(())
                } else {
                    Err(ZmqError::Other("Client disconnected"))
                }
            }
            None => Err(ZmqError::Other(
                "Unable to send reply. No request in progress",
            )),
        }
    }
}

#[async_trait]
impl BlockingRecv for RepSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        let mut messages = FuturesUnordered::new();
        loop {
            for peer in self.backend.peers.iter() {
                let peer_id = peer.identity.clone();
                let recv_queue = peer.recv_queue.clone();
                messages.push(async move { (peer_id, recv_queue.lock().await.next().await) });
            }

            if messages.is_empty() {
                // TODO this is super stupid way of waiting for new connections
                // Needs to be refactored
                tokio::time::delay_for(Duration::from_millis(100)).await;
                continue;
            } else {
                break;
            }
        }
        loop {
            match messages.next().await {
                Some((peer_id, Some(Message::MultipartMessage(mut messages)))) => {
                    assert!(messages.len() == 2);
                    assert!(messages[0].data.is_empty()); // Ensure that we have delimeter as first part
                    self.current_request = Some(peer_id);
                    return Ok(messages.pop().unwrap());
                }
                Some((peer_id, None)) => {
                    println!("Peer disconnected {:?}", peer_id);
                    self.backend.peers.remove(&peer_id);
                }
                Some((_peer_id, _)) => todo!(),
                None => continue,
            };
        }
    }
}
<commit_msg>Removed useless Arc<commit_after>use crate::codec::*;
use crate::error::*;
use crate::*;
use crate::{SocketType, ZmqResult};
use async_trait::async_trait;
use crossbeam::atomic::AtomicCell;
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use futures::lock::Mutex;
use futures::stream::FuturesUnordered;
use futures_util::sink::SinkExt;
use std::sync::Arc;
use std::time::Duration;
use tokio::stream::StreamExt;

struct ReqSocketBackend {
    pub(crate) peers: DashMap<PeerIdentity, Peer>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    pub(crate) current_request_peer_id: AtomicCell<usize>,
}

pub struct ReqSocket {
    backend: Arc<ReqSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl BlockingSend for ReqSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        if self.current_request.is_some() {
            return Err(ZmqError::Socket(
                "Unable to send message. Request already in progress",
            ));
        }
        // In normal scenario this will always be only 1 iteration
        // There can be special case when peer has disconnected and his id is still in RR queue
        // This happens because SegQueue don't have an api to delete items from queue.
        // So in such case we'll just pop item and skip it if we don't have a matching peer in peers map
        loop {
            let next_peer_id = match self.backend.round_robin.pop() {
                Ok(peer) => peer,
                Err(_) => {
                    return Err(ZmqError::Other(
                        "Not connected to peers. Unable to send messages",
                    ))
                }
            };
            match self.backend.peers.get_mut(&next_peer_id) {
                Some(mut peer) => {
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    self.current_request = Some(next_peer_id.clone());
                    self.backend
                        .current_request_peer_id
                        .store(self.backend.peers.hash_usize(&next_peer_id));
                    self.backend.round_robin.push(next_peer_id);
                    return Ok(());
                }
                None => continue,
            }
        }
    }
}

#[async_trait]
impl BlockingRecv for ReqSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(recv_queue) = self
                    .backend
                    .peers
                    .get(&peer_id)
                    .map(|p| p.recv_queue.clone())
                {
                    let message = recv_queue.lock().await.next().await;
                    match message {
                        Some(Message::MultipartMessage(mut message)) => {
                            assert!(message.len() == 2);
                            assert!(message[0].data.is_empty()); // Ensure that we have delimeter as first part
                            Ok(message.pop().unwrap())
                        }
                        Some(_) => Err(ZmqError::Other("Wrong message type received")),
                        None => Err(ZmqError::NoMessage),
                    }
                } else {
                    Err(ZmqError::Other("Server disconnected"))
                }
            }
            None => Err(ZmqError::Other("Unable to recv. No request in progress")),
        }
    }
}

#[async_trait]
impl SocketFrontend for ReqSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(ReqSocketBackend {
                peers: DashMap::new(),
                round_robin: SegQueue::new(),
                current_request_peer_id: AtomicCell::new(0),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        if self._accept_close_handle.is_some() {
            return Err(ZmqError::Other(
                "Socket server already started. Currently only one server is supported",
            ));
        }
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        let addr = endpoint.parse::<SocketAddr>()?;
        let raw_socket = tokio::net::TcpStream::connect(addr).await?;
        util::peer_connected(raw_socket, self.backend.clone()).await;
        Ok(())
    }
}

impl MultiPeer for ReqSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 100;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );
        self.round_robin.push(peer_id.clone());

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for ReqSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        let hash = self.peers.hash_usize(&peer_id);
        let current = self.current_request_peer_id.compare_and_swap(hash, 0);
        // This is needed to ensure that we only store messages that we are expecting to get
        // Other messages are silently discarded according to spec
        if current == hash {
            // We've got reply that we were waiting for
            self.peers
                .get_mut(peer_id)
                .expect("Not found peer by id")
                .recv_queue_in
                .send(message)
                .await
                .expect("Failed to send");
        }
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REQ
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

struct RepSocketBackend {
    pub(crate) peers: Arc<DashMap<PeerIdentity, Peer>>,
}

pub struct RepSocket {
    backend: Arc<RepSocketBackend>,
    _accept_close_handle: Option<oneshot::Sender<bool>>,
    current_request: Option<PeerIdentity>,
}

#[async_trait]
impl SocketFrontend for RepSocket {
    fn new() -> Self {
        Self {
            backend: Arc::new(RepSocketBackend {
                peers: Arc::new(DashMap::new()),
            }),
            _accept_close_handle: None,
            current_request: None,
        }
    }

    async fn bind(&mut self, endpoint: &str) -> ZmqResult<()> {
        let stop_handle = util::start_accepting_connections(endpoint, self.backend.clone()).await?;
        self._accept_close_handle = Some(stop_handle);
        Ok(())
    }

    async fn connect(&mut self, endpoint: &str) -> ZmqResult<()> {
        unimplemented!()
    }
}

impl MultiPeer for RepSocketBackend {
    fn peer_connected(
        &self,
        peer_id: &PeerIdentity,
    ) -> (mpsc::Receiver<Message>, oneshot::Receiver<bool>) {
        let default_queue_size = 100;
        let (out_queue, out_queue_receiver) = mpsc::channel(default_queue_size);
        let (in_queue, in_queue_receiver) = mpsc::channel(default_queue_size);
        let (stop_handle, stop_callback) = oneshot::channel::<bool>();

        self.peers.insert(
            peer_id.clone(),
            Peer {
                identity: peer_id.clone(),
                send_queue: out_queue,
                recv_queue: Arc::new(Mutex::new(in_queue_receiver)),
                recv_queue_in: in_queue,
                _io_close_handle: stop_handle,
            },
        );

        (out_queue_receiver, stop_callback)
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

#[async_trait]
impl SocketBackend for RepSocketBackend {
    async fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        self.peers
            .get_mut(peer_id)
            .expect("Not found peer by id")
            .recv_queue_in
            .send(message)
            .await
            .expect("Failed to send");
    }

    fn socket_type(&self) -> SocketType {
        SocketType::REP
    }

    fn shutdown(&self) {
        self.peers.clear();
    }
}

#[async_trait]
impl BlockingSend for RepSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        match self.current_request.take() {
            Some(peer_id) => {
                if let Some(mut peer) = self.backend.peers.get_mut(&peer_id) {
                    let frames = vec![
                        "".into(), // delimiter frame
                        message,
                    ];
                    peer.send_queue
                        .send(Message::MultipartMessage(frames))
                        .await?;
                    Ok(())
                } else {
                    Err(ZmqError::Other("Client disconnected"))
                }
            }
            None => Err(ZmqError::Other(
                "Unable to send reply. No request in progress",
            )),
        }
    }
}

#[async_trait]
impl BlockingRecv for RepSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        let mut messages = FuturesUnordered::new();
        loop {
            for peer in self.backend.peers.iter() {
                let peer_id = peer.identity.clone();
                let recv_queue = peer.recv_queue.clone();
                messages.push(async move { (peer_id, recv_queue.lock().await.next().await) });
            }

            if messages.is_empty() {
                // TODO this is super stupid way of waiting for new connections
                // Needs to be refactored
                tokio::time::delay_for(Duration::from_millis(100)).await;
                continue;
            } else {
                break;
            }
        }
        loop {
            match messages.next().await {
                Some((peer_id, Some(Message::MultipartMessage(mut messages)))) => {
                    assert!(messages.len() == 2);
                    assert!(messages[0].data.is_empty()); // Ensure that we have delimeter as first part
                    self.current_request = Some(peer_id);
                    return Ok(messages.pop().unwrap());
                }
                Some((peer_id, None)) => {
                    println!("Peer disconnected {:?}", peer_id);
                    self.backend.peers.remove(&peer_id);
                }
                Some((_peer_id, _)) => todo!(),
                None => continue,
            };
        }
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
use ::{card, deck, hand};

#[test]
fn create_hand() {
    let mut deck = deck::Deck::new();
    let mut cards: Vec<card::Card> = vec![];
    for _i in 0..5 {
        cards.push(deck.draw().unwrap());
    }
    let hand = hand::Hand::new(cards);

    assert_eq!(hand.cards.len(), 5);
}

#[test]
fn has_flush() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Six, card::Suit::Hearts);
    let c = card::Card::new(card::Value::Ten, card::Suit::Hearts);
    let d = card::Card::new(card::Value::Queen, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Two, card::Suit::Hearts);
    let cards: Vec<card::Card> = vec![a, b, c, d, e];
    let hand: hand::Hand = hand::Hand::new(cards);

    assert_eq!(hand.has_flush(), true);
}

#[test]
fn no_flush() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Six, card::Suit::Spades);
    let c = card::Card::new(card::Value::Ten, card::Suit::Hearts);
    let d = card::Card::new(card::Value::Queen, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Two, card::Suit::Hearts);
    let cards: Vec<card::Card> = vec![a, b, c, d, e];
    let hand: hand::Hand = hand::Hand::new(cards);

    assert_eq!(hand.has_flush(), false);
}

#[test]
fn high_card_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let poker_hand = hand::PokerHand::HighCard(a);

    assert_eq!(poker_hand.raw_value(), 0);
}

#[test]
fn one_pair_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Three, card::Suit::Diamonds);
    let poker_hand = hand::PokerHand::OnePair((a, b));

    assert_eq!(poker_hand.raw_value(), 1);
}

#[test]
fn two_pair_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Three, card::Suit::Diamonds);
    let c = card::Card::new(card::Value::King, card::Suit::Spades);
    let d = card::Card::new(card::Value::King, card::Suit::Hearts);
    let poker_hand = hand::PokerHand::TwoPair((a, b), (c, d));

    assert_eq!(poker_hand.raw_value(), 2);
}

#[test]
fn trips_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Three, card::Suit::Diamonds);
    let c = card::Card::new(card::Value::Three, card::Suit::Spades);
    let poker_hand = hand::PokerHand::Trips((a, b, c));

    assert_eq!(poker_hand.raw_value(), 3);
}

#[test]
fn straight_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Four, card::Suit::Diamonds);
    let c = card::Card::new(card::Value::Five, card::Suit::Spades);
    let d = card::Card::new(card::Value::Six, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Seven, card::Suit::Hearts);
    let poker_hand = hand::PokerHand::Straight((a, b, c, d, e));

    assert_eq!(poker_hand.raw_value(), 4);
}

#[test]
fn flush_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::King, card::Suit::Hearts);
    let c = card::Card::new(card::Value::Five, card::Suit::Hearts);
    let d = card::Card::new(card::Value::Jack, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Ten, card::Suit::Hearts);
    let poker_hand = hand::PokerHand::Flush((a, b, c, d, e));

    assert_eq!(poker_hand.raw_value(), 5);
}

#[test]
fn full_house_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Three, card::Suit::Spades);
    let c = card::Card::new(card::Value::Three, card::Suit::Clubs);
    let d = card::Card::new(card::Value::Ten, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Ten, card::Suit::Diamonds);
    let poker_hand = hand::PokerHand::FullHouse((a, b, c), (d, e));

    assert_eq!(poker_hand.raw_value(), 6);
}

#[test]
fn quads_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Diamonds);
    let b = card::Card::new(card::Value::Three, card::Suit::Spades);
    let c = card::Card::new(card::Value::Three, card::Suit::Clubs);
    let d = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let poker_hand = hand::PokerHand::Quads((a, b, c, d));

    assert_eq!(poker_hand.raw_value(), 7);
}

#[test]
fn straight_flush_raw_value() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Four, card::Suit::Hearts);
    let c = card::Card::new(card::Value::Five, card::Suit::Hearts);
    let d = card::Card::new(card::Value::Six, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Seven, card::Suit::Hearts);
    let poker_hand = hand::PokerHand::StraightFlush((a, b, c, d, e));

    assert_eq!(poker_hand.raw_value(), 8);
}
<commit_msg>Removed raw_value() tests and replaced them with compare tests<commit_after>#[cfg(test)]
use ::{card, deck, hand};

#[test]
fn create_hand() {
    let mut deck = deck::Deck::new();
    let mut cards: Vec<card::Card> = vec![];
    for _i in 0..5 {
        cards.push(deck.draw().unwrap());
    }
    let hand = hand::Hand::new(cards);

    assert_eq!(hand.cards.len(), 5);
}

#[test]
fn has_flush() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Six, card::Suit::Hearts);
    let c = card::Card::new(card::Value::Ten, card::Suit::Hearts);
    let d = card::Card::new(card::Value::Queen, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Two, card::Suit::Hearts);
    let cards: Vec<card::Card> = vec![a, b, c, d, e];
    let hand: hand::Hand = hand::Hand::new(cards);

    assert_eq!(hand.has_flush(), true);
}

#[test]
fn no_flush() {
    let a = card::Card::new(card::Value::Three, card::Suit::Hearts);
    let b = card::Card::new(card::Value::Six, card::Suit::Spades);
    let c = card::Card::new(card::Value::Ten, card::Suit::Hearts);
    let d = card::Card::new(card::Value::Queen, card::Suit::Hearts);
    let e = card::Card::new(card::Value::Two, card::Suit::Hearts);
    let cards: Vec<card::Card> = vec![a, b, c, d, e];
    let hand: hand::Hand = hand::Hand::new(cards);

    assert_eq!(hand.has_flush(), false);
}

#[test]
fn compare_unequal_poker_hands() {
    assert_eq!(hand::PokerRanking::HighCard < hand::PokerRanking::OnePair, true);
    assert_eq!(hand::PokerRanking::HighCard == hand::PokerRanking::OnePair, false);
    assert_eq!(hand::PokerRanking::HighCard > hand::PokerRanking::OnePair, false);
}

#[test]
fn compare_equal_poker_hands() {
    assert_eq!(hand::PokerRanking::Trips < hand::PokerRanking::Trips, false);
    assert_eq!(hand::PokerRanking::Trips == hand::PokerRanking::Trips, true);
    assert_eq!(hand::PokerRanking::Trips > hand::PokerRanking::Trips, false);
}
<|endoftext|>