<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::error::{Error, Result};
use serde::de::{self, Deserialize, DeserializeSeed, IntoDeserializer, Visitor};

/// Deserializes a `&[u8]` into a type.
///
/// This function will attempt to interpret `bytes` as the LCS serialized form of `T` and
/// deserialize `T` from `bytes`.
///
/// # Examples
///
/// ```
/// use libra_canonical_serialization::from_bytes;
/// use serde::Deserialize;
///
/// #[derive(Deserialize)]
/// struct Ip([u8; 4]);
///
/// #[derive(Deserialize)]
/// struct Port(u16);
///
/// #[derive(Deserialize)]
/// struct SocketAddr {
///     ip: Ip,
///     port: Port,
/// }
///
/// let bytes = vec![0x7f, 0x00, 0x00, 0x01, 0x41, 0x1f];
/// let socket_addr: SocketAddr = from_bytes(&bytes).unwrap();
///
/// assert_eq!(socket_addr.ip.0, [127, 0, 0, 1]);
/// assert_eq!(socket_addr.port.0, 8001);
/// ```
pub fn from_bytes<'a, T>(bytes: &'a [u8]) -> Result<T>
where
    T: Deserialize<'a>,
{
    let mut deserializer = Deserializer::new(bytes);
    let t = T::deserialize(&mut deserializer)?;
    deserializer.end().map(move |_| t)
}

/// Perform a stateful deserialization from a `&[u8]` using the provided `seed`.
pub fn from_bytes_seed<'a, T>(seed: T, bytes: &'a [u8]) -> Result<T::Value>
where
    T: DeserializeSeed<'a>,
{
    let mut deserializer = Deserializer::new(bytes);
    let t = seed.deserialize(&mut deserializer)?;
    deserializer.end().map(move |_| t)
}

/// Deserialization implementation for LCS
struct Deserializer<'de> {
    input: &'de [u8],
}

impl<'de> Deserializer<'de> {
    /// Creates a new `Deserializer` which will be deserializing the provided
    /// input.
    fn new(input: &'de [u8]) -> Self {
        Deserializer { input }
    }

    /// The `Deserializer::end` method should be called after a type has been
    /// fully deserialized. This allows the `Deserializer` to validate that
    /// the there are no more bytes remaining in the input stream.
    fn end(&mut self) -> Result<()> {
        if self.input.is_empty() {
            Ok(())
        } else {
            Err(Error::RemainingInput)
        }
    }
}

impl<'de> Deserializer<'de> {
    fn peek(&mut self) -> Result<u8> {
        self.input.first().copied().ok_or(Error::Eof)
    }

    fn next(&mut self) -> Result<u8> {
        let byte = self.peek()?;
        self.input = &self.input[1..];
        Ok(byte)
    }

    fn parse_bool(&mut self) -> Result<bool> {
        let byte = self.next()?;

        match byte {
            0 => Ok(false),
            1 => Ok(true),
            _ => Err(Error::ExpectedBoolean),
        }
    }

    fn fill_slice(&mut self, slice: &mut [u8]) -> Result<()> {
        for byte in slice {
            *byte = self.next()?;
        }
        Ok(())
    }

    fn parse_u8(&mut self) -> Result<u8> {
        self.next()
    }

    fn parse_u16(&mut self) -> Result<u16> {
        let mut le_bytes = [0; 2];
        self.fill_slice(&mut le_bytes)?;
        Ok(u16::from_le_bytes(le_bytes))
    }

    fn parse_u32(&mut self) -> Result<u32> {
        let mut le_bytes = [0; 4];
        self.fill_slice(&mut le_bytes)?;
        Ok(u32::from_le_bytes(le_bytes))
    }

    fn parse_u64(&mut self) -> Result<u64> {
        let mut le_bytes = [0; 8];
        self.fill_slice(&mut le_bytes)?;
        Ok(u64::from_le_bytes(le_bytes))
    }

    fn parse_u128(&mut self) -> Result<u128> {
        let mut le_bytes = [0; 16];
        self.fill_slice(&mut le_bytes)?;
        Ok(u128::from_le_bytes(le_bytes))
    }

    fn parse_bytes(&mut self) -> Result<&'de [u8]> {
        let len = self.parse_u32()? as usize;
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }

        let slice = self.input.get(..len).ok_or(Error::Eof)?;
        self.input = &self.input[len..];
        Ok(slice)
    }

    fn parse_string(&mut self) -> Result<&'de str> {
        let slice = self.parse_bytes()?;
        std::str::from_utf8(slice).map_err(|_| Error::Utf8)
    }
}

impl<'de, 'a> de::Deserializer<'de> for &'a mut Deserializer<'de> {
    type Error = Error;

    // LCS is not a self-describing format so we can't implement `deserialize_any`
    fn deserialize_any<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_any"))
    }

    fn deserialize_bool<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_bool(self.parse_bool()?)
    }

    fn deserialize_i8<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i8(self.parse_u8()? as i8)
    }

    fn deserialize_i16<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i16(self.parse_u16()? as i16)
    }

    fn deserialize_i32<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i32(self.parse_u32()? as i32)
    }

    fn deserialize_i64<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i64(self.parse_u64()? as i64)
    }

    fn deserialize_i128<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i128(self.parse_u128()? as i128)
    }

    fn deserialize_u8<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u8(self.parse_u8()?)
    }

    fn deserialize_u16<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u16(self.parse_u16()?)
    }

    fn deserialize_u32<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u32(self.parse_u32()?)
    }

    fn deserialize_u64<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u64(self.parse_u64()?)
    }

    fn deserialize_u128<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u128(self.parse_u128()?)
    }

    fn deserialize_f32<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_f32"))
    }

    fn deserialize_f64<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_f64"))
    }

    fn deserialize_char<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_char"))
    }

    fn deserialize_str<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_borrowed_str(self.parse_string()?)
    }

    fn deserialize_string<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        self.deserialize_str(visitor)
    }

    fn deserialize_bytes<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_borrowed_bytes(self.parse_bytes()?)
    }

    fn deserialize_byte_buf<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        self.deserialize_bytes(visitor)
    }

    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        let byte = self.next()?;

        match byte {
            0 => visitor.visit_none(),
            1 => visitor.visit_some(self),
            _ => Err(Error::ExpectedOption),
        }
    }

    fn deserialize_unit<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_unit()
    }

    fn deserialize_unit_struct<V>(self, _name: &'static str, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        self.deserialize_unit(visitor)
    }

    fn deserialize_newtype_struct<V>(self, _name: &'static str, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_newtype_struct(self)
    }

    fn deserialize_seq<V>(mut self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        let len = self.parse_u32()? as usize;
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }

        visitor.visit_seq(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_tuple<V>(mut self, len: usize, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_tuple_struct<V>(
        mut self,
        _name: &'static str,
        len: usize,
        visitor: V,
    ) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_map<V>(mut self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        let len = self.parse_u32()? as usize;
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }

        visitor.visit_map(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_struct<V>(
        mut self,
        _name: &'static str,
        fields: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(&mut self, fields.len()))
    }

    fn deserialize_enum<V>(
        self,
        _name: &'static str,
        _variants: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_enum(self)
    }

    fn deserialize_identifier<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_identifier"))
    }

    // LCS is not a self-describing format so we can't implement `deserialize_ignored_any`
    fn deserialize_ignored_any<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_ignored_any"))
    }

    // LCS is not a human readable format
    fn is_human_readable(&self) -> bool {
        false
    }
}

struct SeqDeserializer<'a, 'de: 'a> {
    de: &'a mut Deserializer<'de>,
    remaining: usize,
}

impl<'a, 'de> SeqDeserializer<'a, 'de> {
    fn new(de: &'a mut Deserializer<'de>, remaining: usize) -> Self {
        Self { de, remaining }
    }
}

impl<'de, 'a> de::SeqAccess<'de> for SeqDeserializer<'a, 'de> {
    type Error = Error;

    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>>
    where
        T: DeserializeSeed<'de>,
    {
        if self.remaining == 0 {
            Ok(None)
        } else {
            self.remaining -= 1;
            seed.deserialize(&mut *self.de).map(Some)
        }
    }

    fn size_hint(&self) -> Option<usize> {
        Some(self.remaining)
    }
}

impl<'de, 'a> de::MapAccess<'de> for SeqDeserializer<'a, 'de> {
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>>
    where
        K: DeserializeSeed<'de>,
    {
        if self.remaining == 0 {
            Ok(None)
        } else {
            self.remaining -= 1;
            seed.deserialize(&mut *self.de).map(Some)
        }
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value>
    where
        V: DeserializeSeed<'de>,
    {
        seed.deserialize(&mut *self.de)
    }

    fn size_hint(&self) -> Option<usize> {
        Some(self.remaining)
    }
}

impl<'de, 'a> de::EnumAccess<'de> for &'a mut Deserializer<'de> {
    type Error = Error;
    type Variant = Self;

    fn variant_seed<V>(self, seed: V) -> Result<(V::Value, Self::Variant)>
    where
        V: DeserializeSeed<'de>,
    {
        let variant_index = self.parse_u32()?;
        let value = seed.deserialize(variant_index.into_deserializer())?;
        Ok((value, self))
    }
}

impl<'de, 'a> de::VariantAccess<'de> for &'a mut Deserializer<'de> {
    type Error = Error;

    fn unit_variant(self) -> Result<()> {
        Ok(())
    }

    fn newtype_variant_seed<T>(self, seed: T) -> Result<T::Value>
    where
        T: DeserializeSeed<'de>,
    {
        seed.deserialize(self)
    }

    fn tuple_variant<V>(self, len: usize, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        de::Deserializer::deserialize_tuple(self, len, visitor)
    }

    fn struct_variant<V>(self, fields: &'static [&'static str], visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        de::Deserializer::deserialize_tuple(self, fields.len(), visitor)
    }
}
<commit_msg>[lcs] enable support for when an enum has modifiers<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::error::{Error, Result};
use serde::de::{self, Deserialize, DeserializeSeed, IntoDeserializer, Visitor};

/// Deserializes a `&[u8]` into a type.
///
/// This function will attempt to interpret `bytes` as the LCS serialized form of `T` and
/// deserialize `T` from `bytes`.
///
/// # Examples
///
/// ```
/// use libra_canonical_serialization::from_bytes;
/// use serde::Deserialize;
///
/// #[derive(Deserialize)]
/// struct Ip([u8; 4]);
///
/// #[derive(Deserialize)]
/// struct Port(u16);
///
/// #[derive(Deserialize)]
/// struct SocketAddr {
///     ip: Ip,
///     port: Port,
/// }
///
/// let bytes = vec![0x7f, 0x00, 0x00, 0x01, 0x41, 0x1f];
/// let socket_addr: SocketAddr = from_bytes(&bytes).unwrap();
///
/// assert_eq!(socket_addr.ip.0, [127, 0, 0, 1]);
/// assert_eq!(socket_addr.port.0, 8001);
/// ```
pub fn from_bytes<'a, T>(bytes: &'a [u8]) -> Result<T>
where
    T: Deserialize<'a>,
{
    let mut deserializer = Deserializer::new(bytes);
    let t = T::deserialize(&mut deserializer)?;
    deserializer.end().map(move |_| t)
}

/// Perform a stateful deserialization from a `&[u8]` using the provided `seed`.
pub fn from_bytes_seed<'a, T>(seed: T, bytes: &'a [u8]) -> Result<T::Value>
where
    T: DeserializeSeed<'a>,
{
    let mut deserializer = Deserializer::new(bytes);
    let t = seed.deserialize(&mut deserializer)?;
    deserializer.end().map(move |_| t)
}

/// Deserialization implementation for LCS
struct Deserializer<'de> {
    input: &'de [u8],
}

impl<'de> Deserializer<'de> {
    /// Creates a new `Deserializer` which will be deserializing the provided
    /// input.
    fn new(input: &'de [u8]) -> Self {
        Deserializer { input }
    }

    /// The `Deserializer::end` method should be called after a type has been
    /// fully deserialized. This allows the `Deserializer` to validate that
    /// the there are no more bytes remaining in the input stream.
    fn end(&mut self) -> Result<()> {
        if self.input.is_empty() {
            Ok(())
        } else {
            Err(Error::RemainingInput)
        }
    }
}

impl<'de> Deserializer<'de> {
    fn peek(&mut self) -> Result<u8> {
        self.input.first().copied().ok_or(Error::Eof)
    }

    fn next(&mut self) -> Result<u8> {
        let byte = self.peek()?;
        self.input = &self.input[1..];
        Ok(byte)
    }

    fn parse_bool(&mut self) -> Result<bool> {
        let byte = self.next()?;

        match byte {
            0 => Ok(false),
            1 => Ok(true),
            _ => Err(Error::ExpectedBoolean),
        }
    }

    fn fill_slice(&mut self, slice: &mut [u8]) -> Result<()> {
        for byte in slice {
            *byte = self.next()?;
        }
        Ok(())
    }

    fn parse_u8(&mut self) -> Result<u8> {
        self.next()
    }

    fn parse_u16(&mut self) -> Result<u16> {
        let mut le_bytes = [0; 2];
        self.fill_slice(&mut le_bytes)?;
        Ok(u16::from_le_bytes(le_bytes))
    }

    fn parse_u32(&mut self) -> Result<u32> {
        let mut le_bytes = [0; 4];
        self.fill_slice(&mut le_bytes)?;
        Ok(u32::from_le_bytes(le_bytes))
    }

    fn parse_u64(&mut self) -> Result<u64> {
        let mut le_bytes = [0; 8];
        self.fill_slice(&mut le_bytes)?;
        Ok(u64::from_le_bytes(le_bytes))
    }

    fn parse_u128(&mut self) -> Result<u128> {
        let mut le_bytes = [0; 16];
        self.fill_slice(&mut le_bytes)?;
        Ok(u128::from_le_bytes(le_bytes))
    }

    fn parse_bytes(&mut self) -> Result<&'de [u8]> {
        let len = self.parse_u32()? as usize;
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }

        let slice = self.input.get(..len).ok_or(Error::Eof)?;
        self.input = &self.input[len..];
        Ok(slice)
    }

    fn parse_string(&mut self) -> Result<&'de str> {
        let slice = self.parse_bytes()?;
        std::str::from_utf8(slice).map_err(|_| Error::Utf8)
    }
}

impl<'de, 'a> de::Deserializer<'de> for &'a mut Deserializer<'de> {
    type Error = Error;

    // LCS is not a self-describing format so we can't implement `deserialize_any`
    fn deserialize_any<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_any"))
    }

    fn deserialize_bool<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_bool(self.parse_bool()?)
    }

    fn deserialize_i8<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i8(self.parse_u8()? as i8)
    }

    fn deserialize_i16<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i16(self.parse_u16()? as i16)
    }

    fn deserialize_i32<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i32(self.parse_u32()? as i32)
    }

    fn deserialize_i64<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i64(self.parse_u64()? as i64)
    }

    fn deserialize_i128<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_i128(self.parse_u128()? as i128)
    }

    fn deserialize_u8<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u8(self.parse_u8()?)
    }

    fn deserialize_u16<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u16(self.parse_u16()?)
    }

    fn deserialize_u32<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u32(self.parse_u32()?)
    }

    fn deserialize_u64<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u64(self.parse_u64()?)
    }

    fn deserialize_u128<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_u128(self.parse_u128()?)
    }

    fn deserialize_f32<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_f32"))
    }

    fn deserialize_f64<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_f64"))
    }

    fn deserialize_char<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_char"))
    }

    fn deserialize_str<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_borrowed_str(self.parse_string()?)
    }

    fn deserialize_string<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        self.deserialize_str(visitor)
    }

    fn deserialize_bytes<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_borrowed_bytes(self.parse_bytes()?)
    }

    fn deserialize_byte_buf<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        self.deserialize_bytes(visitor)
    }

    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        let byte = self.next()?;

        match byte {
            0 => visitor.visit_none(),
            1 => visitor.visit_some(self),
            _ => Err(Error::ExpectedOption),
        }
    }

    fn deserialize_unit<V>(self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_unit()
    }

    fn deserialize_unit_struct<V>(self, _name: &'static str, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        self.deserialize_unit(visitor)
    }

    fn deserialize_newtype_struct<V>(self, _name: &'static str, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_newtype_struct(self)
    }

    fn deserialize_seq<V>(mut self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        let len = self.parse_u32()? as usize;
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }

        visitor.visit_seq(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_tuple<V>(mut self, len: usize, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_tuple_struct<V>(
        mut self,
        _name: &'static str,
        len: usize,
        visitor: V,
    ) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_map<V>(mut self, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        let len = self.parse_u32()? as usize;
        if len > crate::MAX_SEQUENCE_LENGTH {
            return Err(Error::ExceededMaxLen(len));
        }

        visitor.visit_map(SeqDeserializer::new(&mut self, len))
    }

    fn deserialize_struct<V>(
        mut self,
        _name: &'static str,
        fields: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(&mut self, fields.len()))
    }

    fn deserialize_enum<V>(
        self,
        _name: &'static str,
        _variants: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        visitor.visit_enum(self)
    }

    // LCS does not utilize identifiers, so throw them away
    fn deserialize_identifier<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        self.deserialize_bytes(_visitor)
    }

    // LCS is not a self-describing format so we can't implement `deserialize_ignored_any`
    fn deserialize_ignored_any<V>(self, _visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        Err(Error::NotSupported("deserialize_ignored_any"))
    }

    // LCS is not a human readable format
    fn is_human_readable(&self) -> bool {
        false
    }
}

struct SeqDeserializer<'a, 'de: 'a> {
    de: &'a mut Deserializer<'de>,
    remaining: usize,
}

impl<'a, 'de> SeqDeserializer<'a, 'de> {
    fn new(de: &'a mut Deserializer<'de>, remaining: usize) -> Self {
        Self { de, remaining }
    }
}

impl<'de, 'a> de::SeqAccess<'de> for SeqDeserializer<'a, 'de> {
    type Error = Error;

    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>>
    where
        T: DeserializeSeed<'de>,
    {
        if self.remaining == 0 {
            Ok(None)
        } else {
            self.remaining -= 1;
            seed.deserialize(&mut *self.de).map(Some)
        }
    }

    fn size_hint(&self) -> Option<usize> {
        Some(self.remaining)
    }
}

impl<'de, 'a> de::MapAccess<'de> for SeqDeserializer<'a, 'de> {
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>>
    where
        K: DeserializeSeed<'de>,
    {
        if self.remaining == 0 {
            Ok(None)
        } else {
            self.remaining -= 1;
            seed.deserialize(&mut *self.de).map(Some)
        }
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value>
    where
        V: DeserializeSeed<'de>,
    {
        seed.deserialize(&mut *self.de)
    }

    fn size_hint(&self) -> Option<usize> {
        Some(self.remaining)
    }
}

impl<'de, 'a> de::EnumAccess<'de> for &'a mut Deserializer<'de> {
    type Error = Error;
    type Variant = Self;

    fn variant_seed<V>(self, seed: V) -> Result<(V::Value, Self::Variant)>
    where
        V: DeserializeSeed<'de>,
    {
        let variant_index = self.parse_u32()?;
        let value = seed.deserialize(variant_index.into_deserializer())?;
        Ok((value, self))
    }
}

impl<'de, 'a> de::VariantAccess<'de> for &'a mut Deserializer<'de> {
    type Error = Error;

    fn unit_variant(self) -> Result<()> {
        Ok(())
    }

    fn newtype_variant_seed<T>(self, seed: T) -> Result<T::Value>
    where
        T: DeserializeSeed<'de>,
    {
        seed.deserialize(self)
    }

    fn tuple_variant<V>(self, len: usize, visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        de::Deserializer::deserialize_tuple(self, len, visitor)
    }

    fn struct_variant<V>(self, fields: &'static [&'static str], visitor: V) -> Result<V::Value>
    where
        V: Visitor<'de>,
    {
        de::Deserializer::deserialize_tuple(self, fields.len(), visitor)
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::ffi::OsStr;
use std::path::PathBuf;
use std::process::Command;

use crate::support::git;
use crate::support::paths;
use crate::support::project;
use crate::support::registry::Package;
use git2;

use url::Url;

fn find_index() -> PathBuf {
    let dir = paths::home().join(".cargo/registry/index");
    dir.read_dir().unwrap().next().unwrap().unwrap().path()
}

fn run_test(path_env: Option<&OsStr>) {
    const N: usize = 50;

    let foo = project()
        .file(
            "Cargo.toml",
            r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []

            [dependencies]
            bar = "*"
        "#,
        )
        .file("src/lib.rs", "")
        .build();
    Package::new("bar", "0.1.0").publish();

    foo.cargo("build").run();

    let index = find_index();
    let path = paths::home().join("tmp");
    let url = Url::from_file_path(&path).unwrap().to_string();
    let repo = git2::Repository::init(&path).unwrap();
    let index = git2::Repository::open(&index).unwrap();
    let mut cfg = repo.config().unwrap();
    cfg.set_str("user.email", "foo@bar.com").unwrap();
    cfg.set_str("user.name", "Foo Bar").unwrap();
    let mut cfg = index.config().unwrap();
    cfg.set_str("user.email", "foo@bar.com").unwrap();
    cfg.set_str("user.name", "Foo Bar").unwrap();

    for _ in 0..N {
        git::commit(&repo);
        index
            .remote_anonymous(&url)
            .unwrap()
            .fetch(&["refs/heads/master:refs/remotes/foo/master"], None, None)
            .unwrap();
    }
    drop((repo, index));
    Package::new("bar", "0.1.1").publish();

    let before = find_index()
        .join(".git/objects/pack")
        .read_dir()
        .unwrap()
        .count();
    assert!(before > N);

    let mut cmd = foo.cargo("update");
    cmd.env("__CARGO_PACKFILE_LIMIT", "10");
    if let Some(path) = path_env {
        cmd.env("PATH", path);
    }
    cmd.env("CARGO_LOG", "trace");
    cmd.run();
    let after = find_index()
        .join(".git/objects/pack")
        .read_dir()
        .unwrap()
        .count();
    assert!(
        after < before,
        "packfiles before: {}\n\
         packfiles after:  {}",
        before,
        after
    );
}

#[cargo_test]
fn use_git_gc() {
    if Command::new("git").arg("--version").output().is_err() {
        return;
    }
    run_test(None);
}

#[cargo_test]
// it looks like this test passes on some windows machines but not others,
// notably not on AppVeyor's machines. Sounds like another but for another day.
#[cfg_attr(windows, ignore)]
fn avoid_using_git() {
    let path = env::var_os("PATH").unwrap_or_default();
    let mut paths = env::split_paths(&path).collect::<Vec<_>>();
    let idx = paths
        .iter()
        .position(|p| p.join("git").exists() || p.join("git.exe").exists());
    match idx {
        Some(i) => {
            paths.remove(i);
        }
        None => return,
    }
    run_test(Some(&env::join_paths(&paths).unwrap()));
}
<commit_msg>Remove appveyor-aimed #[ignore]<commit_after>use std::env;
use std::ffi::OsStr;
use std::path::PathBuf;
use std::process::Command;

use crate::support::git;
use crate::support::paths;
use crate::support::project;
use crate::support::registry::Package;
use git2;

use url::Url;

fn find_index() -> PathBuf {
    let dir = paths::home().join(".cargo/registry/index");
    dir.read_dir().unwrap().next().unwrap().unwrap().path()
}

fn run_test(path_env: Option<&OsStr>) {
    const N: usize = 50;

    let foo = project()
        .file(
            "Cargo.toml",
            r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []

            [dependencies]
            bar = "*"
        "#,
        )
        .file("src/lib.rs", "")
        .build();
    Package::new("bar", "0.1.0").publish();

    foo.cargo("build").run();

    let index = find_index();
    let path = paths::home().join("tmp");
    let url = Url::from_file_path(&path).unwrap().to_string();
    let repo = git2::Repository::init(&path).unwrap();
    let index = git2::Repository::open(&index).unwrap();
    let mut cfg = repo.config().unwrap();
    cfg.set_str("user.email", "foo@bar.com").unwrap();
    cfg.set_str("user.name", "Foo Bar").unwrap();
    let mut cfg = index.config().unwrap();
    cfg.set_str("user.email", "foo@bar.com").unwrap();
    cfg.set_str("user.name", "Foo Bar").unwrap();

    for _ in 0..N {
        git::commit(&repo);
        index
            .remote_anonymous(&url)
            .unwrap()
            .fetch(&["refs/heads/master:refs/remotes/foo/master"], None, None)
            .unwrap();
    }
    drop((repo, index));
    Package::new("bar", "0.1.1").publish();

    let before = find_index()
        .join(".git/objects/pack")
        .read_dir()
        .unwrap()
        .count();
    assert!(before > N);

    let mut cmd = foo.cargo("update");
    cmd.env("__CARGO_PACKFILE_LIMIT", "10");
    if let Some(path) = path_env {
        cmd.env("PATH", path);
    }
    cmd.env("CARGO_LOG", "trace");
    cmd.run();
    let after = find_index()
        .join(".git/objects/pack")
        .read_dir()
        .unwrap()
        .count();
    assert!(
        after < before,
        "packfiles before: {}\n\
         packfiles after:  {}",
        before,
        after
    );
}

#[cargo_test]
fn use_git_gc() {
    if Command::new("git").arg("--version").output().is_err() {
        return;
    }
    run_test(None);
}

#[cargo_test]
fn avoid_using_git() {
    let path = env::var_os("PATH").unwrap_or_default();
    let mut paths = env::split_paths(&path).collect::<Vec<_>>();
    let idx = paths
        .iter()
        .position(|p| p.join("git").exists() || p.join("git.exe").exists());
    match idx {
        Some(i) => {
            paths.remove(i);
        }
        None => return,
    }
    run_test(Some(&env::join_paths(&paths).unwrap()));
}
<|endoftext|>


<commit_before>use std::any::Any;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Sender, TryRecvError};
use std::thread;

use {Actor, ActorRef, CanReceive, Props};
use actor_cell::ActorCell;

/// Wrapper around the threads handle and termination sender.
type ConsumerThread = (thread::JoinHandle<()>, Sender<()>);

/// ActorSystem, the struct that manages the creation of everything and that everything does what
/// it is supposed to do.
///
/// NOTE: It currently holds the consumer threads and do not create the user, system and root
/// actors.
pub struct ActorSystem {
    name: Arc<String>,
    // For now we will have the worker pool in the system.
    // TODO(find a way to have a clean way to separate system and user threads).
    consumer_threads: Arc<Mutex<Vec<ConsumerThread>>>,
    // TODO(gamazeps): Have a CanHandle Trait for that.
    actors_queue: Arc<Mutex<VecDeque<Arc<CanReceive >>>>,
}

impl ActorSystem {
    /// Creates a new ActorSystem.
    ///
    /// Note that no threads are started.
    pub fn new(name: String) -> ActorSystem {
        ActorSystem {
            name: Arc::new(name),
            consumer_threads: Arc::new(Mutex::new(Vec::new())),
            actors_queue: Arc::new(Mutex::new(VecDeque::new())),
        }
    }

    /// Spawns an Actor of type A, created using the Props given.
    pub fn actor_of<Args: Copy + Send + Sync + 'static, M: Copy + Send + Sync + 'static + Any, A: Actor<M> + 'static>(&self, props: Props<Args, M, A>) -> ActorRef<Args, M, A> {
        let actor = props.create();
        let actor_cell = ActorCell::new(actor, props, self.clone());
        ActorRef::with_cell(actor_cell)
    }

    /// Enqueues the given Actor on the queue of Actors with something to handle.
    pub fn enqueue_actor<Args: Copy + Send + Sync + 'static, M: Copy + Send + Sync + 'static + Any, A: Actor<M> + 'static>(&self, actor_ref: ActorRef<Args, M, A>) {
        self.actors_queue.lock().unwrap().push_back(Arc::new(actor_ref));
    }

    /// Spawns a thread that will consume messages from the `ActorRef` in `actors_queue`.
    /// This thread can be terminated by calling `terminate_thread`.
    pub fn spawn_thread(&self) {
        let (tx, rx) = channel();
        let thread_system = self.clone();
        let handle = thread::spawn(move || {
            loop {
                // If we received a () we kill the thread.
                match rx.try_recv() {
                    Ok(_) | Err(TryRecvError::Disconnected) => {
                        println!("Terminating a consumer thread.");
                        break;
                    },
                    Err(TryRecvError::Empty) => {}
                };
                // Else we try to prcess a message.
                let actor_ref = {thread_system.actors_queue.lock().unwrap().pop_front()};
                for actor in actor_ref.iter() {
                    actor.handle();
                }
            }
        });
        self.add_thread(handle, tx);
    }

    fn add_thread(&self, handle: thread::JoinHandle<()>, sender: Sender<()>) {
        self.consumer_threads.lock().unwrap().push((handle, sender));
    }

    /// Kills a consumer thread of the `ActorSystem`.
    pub fn terminate_thread(&self) {
        self.terminate_threads(1);
    }

    /// Spawns n threads that will consume messages from the `ActorRef` in `actors_queue`.
    pub fn spawn_threads(&self, n: u32) {
        for _ in 0..n {
            self.spawn_thread();
        }
    }

    /// Kills n consumer threads.
    pub fn terminate_threads(&self, n: u32) {
        let mut handles = Vec::new();
        for _ in 0..n {
            let (handle, tx) = {self.consumer_threads.lock().unwrap().pop().unwrap()};
            let _res = tx.send(());
            handles.push(handle);
        }
        for handle in handles {
            let _res = handle.join();
        }
    }
}

impl Clone for ActorSystem {
    fn clone(&self) -> ActorSystem {
        ActorSystem {
            name: self.name.clone(),
            consumer_threads: self.consumer_threads.clone(),
            actors_queue: self.actors_queue.clone(),
        }
    }
}
<commit_msg>Use channel for enqueuing actors instead of a VecDeque.<commit_after>use std::any::Any;
use std::process::exit;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use {Actor, ActorRef, CanReceive, Props};
use actor_cell::ActorCell;

/// Wrapper around the threads handle and termination sender.
type ConsumerThread = (thread::JoinHandle<()>, Sender<()>);

/// ActorSystem, the struct that manages the creation of everything and that everything does what
/// it is supposed to do.
///
/// NOTE: It currently holds the consumer threads and do not create the user, system and root
/// actors.
pub struct ActorSystem {
    name: Arc<String>,
    // For now we will have the worker pool in the system.
    // TODO(find a way to have a clean way to separate system and user threads).
    consumer_threads: Arc<Mutex<Vec<ConsumerThread>>>,
    // Sends Canreceive to be handled on that channel.
    actors_queue_sender: Arc<Mutex<Sender<Arc<CanReceive>>>>,
    // Receiving end to give to the thread pool.
    actors_queue_receiver: Arc<Mutex<Receiver<Arc<CanReceive>>>>,
}

impl ActorSystem {
    /// Creates a new ActorSystem.
    ///
    /// Note that no threads are started.
    pub fn new(name: String) -> ActorSystem {
        let (tx, rx) = channel();
        ActorSystem {
            name: Arc::new(name),
            consumer_threads: Arc::new(Mutex::new(Vec::new())),
            actors_queue_sender: Arc::new(Mutex::new(tx)),
            actors_queue_receiver: Arc::new(Mutex::new(rx)),
        }
    }

    /// Spawns an Actor of type A, created using the Props given.
    pub fn actor_of<Args: Copy + Send + Sync + 'static, M: Copy + Send + Sync + 'static + Any, A: Actor<M> + 'static>(&self, props: Props<Args, M, A>) -> ActorRef<Args, M, A> {
        let actor = props.create();
        let actor_cell = ActorCell::new(actor, props, self.clone());
        ActorRef::with_cell(actor_cell)
    }

    /// Enqueues the given Actor on the queue of Actors with something to handle.
    pub fn enqueue_actor<Args: Copy + Send + Sync + 'static, M: Copy + Send + Sync + 'static + Any, A: Actor<M> + 'static>(&self, actor_ref: ActorRef<Args, M, A>) {
        match self.actors_queue_sender.lock().unwrap().send(Arc::new(actor_ref)) {
            Ok(_) => return,
            Err(_) => {
                println!("The communication channel for messages is disconnected, this is bad!");
                exit(1);
            },
        }
    }

    /// Spawns a thread that will consume messages from the `ActorRef` in `actors_queue`.
    /// This thread can be terminated by calling `terminate_thread`.
    pub fn spawn_thread(&self) {
        let (tx, rx) = channel();
        let actors_queue = self.actors_queue_receiver.clone();
        let handle = thread::spawn(move || {
            loop {
                // If we received a () we kill the thread.
                match rx.try_recv() {
                    Ok(_) | Err(TryRecvError::Disconnected) => {
                        println!("Terminating a consumer thread.");
                        break;
                    },
                    Err(TryRecvError::Empty) => {}
                };

                // Else we try to process a message.
                let actor_ref = {
                    let lock = actors_queue.lock().unwrap();
                    lock.try_recv()
                };

                match actor_ref {
                    Ok(actor_ref) => actor_ref.handle(),
                    Err(TryRecvError::Empty) => continue,
                    Err(TryRecvError::Disconnected) => {
                        println!("The actors queue failed, something is very wrong");
                        exit(1);
                    }
                }
            }
        });
        self.add_thread(handle, tx);
    }

    fn add_thread(&self, handle: thread::JoinHandle<()>, sender: Sender<()>) {
        self.consumer_threads.lock().unwrap().push((handle, sender));
    }

    /// Kills a consumer thread of the `ActorSystem`.
    pub fn terminate_thread(&self) {
        self.terminate_threads(1);
    }

    /// Spawns n threads that will consume messages from the `ActorRef` in `actors_queue`.
    pub fn spawn_threads(&self, n: u32) {
        for _ in 0..n {
            self.spawn_thread();
        }
    }

    /// Kills n consumer threads.
    pub fn terminate_threads(&self, n: u32) {
        let mut handles = Vec::new();
        for _ in 0..n {
            let (handle, tx) = {self.consumer_threads.lock().unwrap().pop().unwrap()};
            let _res = tx.send(());
            handles.push(handle);
        }
        for handle in handles {
            let _res = handle.join();
        }
    }
}

impl Clone for ActorSystem {
    fn clone(&self) -> ActorSystem {
        ActorSystem {
            name: self.name.clone(),
            consumer_threads: self.consumer_threads.clone(),
            actors_queue_sender: self.actors_queue_sender.clone(),
            actors_queue_receiver: self.actors_queue_receiver.clone(),
        }
    }
}
<|endoftext|>


<commit_before>#![feature(core)]
#![feature(std_misc)]

extern crate algorithmia;
extern crate getopts;

use algorithmia::{Service};
use getopts::Options;
use std::ascii::AsciiExt;
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Semaphore};

static DEFAULT_UPLOAD_CONCURRENCY: u32 = 8;

fn print_usage(opts: &Options) {
    let brief = vec![
        "Usage: algodata USER/COLLECTION [CMD [CMD_ARGS...]]",
        "Supported CMDs",
        "  CREATE",
        "  UPLOAD FILE..."
    ];
    println!("{}", opts.usage(&*brief.connect("\n")));
    env::set_exit_status(1);
}

struct AlgoData {
    service: Service,
}

impl AlgoData {
    fn new(api_key: &str) -> AlgoData {
        AlgoData { service: Service::new(api_key) }
    }

    fn show_collection(self, username: &str, collection_name: &str) {
        let mut my_bucket = self.service.collection(username, collection_name);
        match my_bucket.show() {
            Ok(output) => println!("{:?}", output),
            Err(why) => println!("ERROR: {:?}", why),
        };
    }

    fn create_collection(self, username: &str, collection_name: &str) {
        let mut my_bucket = self.service.collection(username, collection_name);
        match my_bucket.create() {
            Ok(output) => println!("{:?}", output),
            Err(why) => println!("ERROR: {:?}", why),
        };
    }

    fn upload_files(self, username: &str, collection_name: &str, file_paths: &[String], concurrency: u32) {
        println!("Uploading {} file(s)...", file_paths.len());
        let arc_sem = Arc::new(Semaphore::new(concurrency as isize));

        let _: Vec<_> = file_paths.iter().map(|file_path| {
            // Acquire semaphore before we start the thread
            let child_sem = arc_sem.clone();
            child_sem.acquire();
            // println!("Uploading {}", file_path);

            let service = self.service.clone();
            thread::scoped( move || {
                let mut my_bucket = service.collection(username, collection_name);
                match File::open(file_path) {
                    Ok(mut file) => {
                        let ref mut bucket = my_bucket;
                        match bucket.upload_file(&mut file) {
                            Ok(file_added) => println!("Uploaded {}", file_added.result),
                            Err(e) => println!("ERROR uploading {}: {:?}", file_path, e),
                        };
                    },
                    Err(e) => {
                        println!("Failed to open {}: {}", file_path, e);
                    }
                };

                // Release the semaphore
                child_sem.release();
            })
        }).collect();
        println!("Finished uploading {} file(s)", file_paths.len())
    }
}

fn main() {
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help");
    opts.optopt("c", "concurrency", &*format!("max concurrent threads to use for uploading files (default = {})", DEFAULT_UPLOAD_CONCURRENCY), "THREADS");

    let argopts = match opts.parse(env::args()) {
        Ok(m) => m,
        Err(f) => {
            println!("Failed to parse args: {}", f);
            print_usage(&opts);
            return;
        }
    };

    let mut args_iter = argopts.free.clone().into_iter().skip(1);
    if argopts.opt_present("help") || args_iter.len() == 0 {
        print_usage(&opts);
        return;
    }

    let api_key = match env::var("ALGORITHMIA_API_KEY") {
        Ok(val) => val,
        Err(_) => {
            println!("Must set ALGORITHMIA_API_KEY");
            print_usage(&opts);
            return;
        }
    };

    // Get the --concurrency
    let concurrency: u32 = match argopts.opt_str("concurrency") {
        Some(nstr) => {
            match nstr.parse::<u32>() {
                Ok(n) => n,
                Err(_) => {
                    println!("Invalid concurrency option: {}", nstr);
                    print_usage(&opts);
                    return;
                }
            }
        },
        None => DEFAULT_UPLOAD_CONCURRENCY,
    };

    let data = AlgoData::new(&*api_key);

    // Get the USERNAME/COLLECTION arg
    let first_arg = args_iter.next();
    let user_collection: Vec<&str> = match first_arg {
        Some(ref arg) => arg.split('/').collect(),
        None => {
            println!("Did not specify USERNAME/COLLECTION");
            print_usage(&opts);
            return;
        }
    };

    // Get the CMD arg
    let cmd = match args_iter.next() {
        Some(ref arg) => arg.to_ascii_lowercase(),
        None => "show".to_string(),
    };

    match user_collection.as_slice() {
        [user, collection] => {
            match cmd.as_slice() {
                "show" => data.show_collection(user, collection),
                "create" => data.create_collection(user, collection),
                "upload" => {
                    let files: Vec<String> = args_iter.collect();
                    data.upload_files(user, collection, files.as_slice(), concurrency);
                },
                invalid => {
                    println!("Not a valid command: {}", invalid);
                    print_usage(&opts);
                    return;
                }
            }
        },
        invalid => {
            println!("Invalid data repository: {:?}", invalid );
            print_usage(&opts);
            return;
        }
    }
}<commit_msg>--help mentions SHOW cmd<commit_after>#![feature(core)]
#![feature(std_misc)]

extern crate algorithmia;
extern crate getopts;

use algorithmia::{Service};
use getopts::Options;
use std::ascii::AsciiExt;
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Semaphore};

static DEFAULT_UPLOAD_CONCURRENCY: u32 = 8;

fn print_usage(opts: &Options) {
    let brief = vec![
        "Usage: algodata USER/COLLECTION [CMD [CMD_ARGS...]]",
        "Supported CMDs",
        "  SHOW",
        "  CREATE",
        "  UPLOAD FILE..."
    ];
    println!("{}", opts.usage(&*brief.connect("\n")));
    env::set_exit_status(1);
}

struct AlgoData {
    service: Service,
}

impl AlgoData {
    fn new(api_key: &str) -> AlgoData {
        AlgoData { service: Service::new(api_key) }
    }

    fn show_collection(self, username: &str, collection_name: &str) {
        let mut my_bucket = self.service.collection(username, collection_name);
        match my_bucket.show() {
            Ok(output) => println!("{:?}", output),
            Err(why) => println!("ERROR: {:?}", why),
        };
    }

    fn create_collection(self, username: &str, collection_name: &str) {
        let mut my_bucket = self.service.collection(username, collection_name);
        match my_bucket.create() {
            Ok(output) => println!("{:?}", output),
            Err(why) => println!("ERROR: {:?}", why),
        };
    }

    fn upload_files(self, username: &str, collection_name: &str, file_paths: &[String], concurrency: u32) {
        println!("Uploading {} file(s)...", file_paths.len());
        let arc_sem = Arc::new(Semaphore::new(concurrency as isize));

        let _: Vec<_> = file_paths.iter().map(|file_path| {
            // Acquire semaphore before we start the thread
            let child_sem = arc_sem.clone();
            child_sem.acquire();
            // println!("Uploading {}", file_path);

            let service = self.service.clone();
            thread::scoped( move || {
                let mut my_bucket = service.collection(username, collection_name);
                match File::open(file_path) {
                    Ok(mut file) => {
                        let ref mut bucket = my_bucket;
                        match bucket.upload_file(&mut file) {
                            Ok(file_added) => println!("Uploaded {}", file_added.result),
                            Err(e) => println!("ERROR uploading {}: {:?}", file_path, e),
                        };
                    },
                    Err(e) => {
                        println!("Failed to open {}: {}", file_path, e);
                    }
                };

                // Release the semaphore
                child_sem.release();
            })
        }).collect();
        println!("Finished uploading {} file(s)", file_paths.len())
    }
}

fn main() {
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help");
    opts.optopt("c", "concurrency", &*format!("max concurrent threads to use for uploading files (default = {})", DEFAULT_UPLOAD_CONCURRENCY), "THREADS");

    let argopts = match opts.parse(env::args()) {
        Ok(m) => m,
        Err(f) => {
            println!("Failed to parse args: {}", f);
            print_usage(&opts);
            return;
        }
    };

    let mut args_iter = argopts.free.clone().into_iter().skip(1);
    if argopts.opt_present("help") || args_iter.len() == 0 {
        print_usage(&opts);
        return;
    }

    let api_key = match env::var("ALGORITHMIA_API_KEY") {
        Ok(val) => val,
        Err(_) => {
            println!("Must set ALGORITHMIA_API_KEY");
            print_usage(&opts);
            return;
        }
    };

    // Get the --concurrency
    let concurrency: u32 = match argopts.opt_str("concurrency") {
        Some(nstr) => {
            match nstr.parse::<u32>() {
                Ok(n) => n,
                Err(_) => {
                    println!("Invalid concurrency option: {}", nstr);
                    print_usage(&opts);
                    return;
                }
            }
        },
        None => DEFAULT_UPLOAD_CONCURRENCY,
    };

    let data = AlgoData::new(&*api_key);

    // Get the USERNAME/COLLECTION arg
    let first_arg = args_iter.next();
    let user_collection: Vec<&str> = match first_arg {
        Some(ref arg) => arg.split('/').collect(),
        None => {
            println!("Did not specify USERNAME/COLLECTION");
            print_usage(&opts);
            return;
        }
    };

    // Get the CMD arg
    let cmd = match args_iter.next() {
        Some(ref arg) => arg.to_ascii_lowercase(),
        None => "show".to_string(),
    };

    match user_collection.as_slice() {
        [user, collection] => {
            match cmd.as_slice() {
                "show" => data.show_collection(user, collection),
                "create" => data.create_collection(user, collection),
                "upload" => {
                    let files: Vec<String> = args_iter.collect();
                    data.upload_files(user, collection, files.as_slice(), concurrency);
                },
                invalid => {
                    println!("Not a valid command: {}", invalid);
                    print_usage(&opts);
                    return;
                }
            }
        },
        invalid => {
            println!("Invalid data repository: {:?}", invalid );
            print_usage(&opts);
            return;
        }
    }
}<|endoftext|>


<commit_before>use std::fs::{self, OpenOptions};
use std::path::PathBuf;

use clap::{App, SubCommand, Arg, ArgMatches};
use memmap::{Mmap, Protection};
use futures::prelude::*;

use attaca::context::Context;
use attaca::marshal::{ObjectHash, Object, DataObject, SubtreeObject, SubtreeEntry};
use attaca::repository::Repository;
use attaca::store::ObjectStore;
use attaca::trace::Trace;

use errors::*;

pub fn command() -> App<'static, 'static> {
    SubCommand::with_name("checkout")
        .help(
            "Checkout a file or working directory from a previous revision.",
        )
        .arg(Arg::with_name("COMMIT").index(1).required(true).help(
            "The commit hash to checkout.",
        ))
}

fn write_data_object<T: Trace, S: ObjectStore>(
    ctx: &Context<T, S>,
    path: PathBuf,
    object_hash: ObjectHash,
    size: u64,
) -> Result<()> {
    let mut parent = path.clone();
    assert!(parent.pop());
    fs::create_dir_all(parent)?;

    let file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(true)
        .open(path)?;
    file.set_len(size)?;

    // PermissionError indicates mmap opened with wrong permissions.
    let mut mmap = Mmap::open(&file, Protection::ReadWrite)?;

    let slice = unsafe { mmap.as_mut_slice() };
    let mut stack = vec![(0u64, object_hash)];
    while let Some((mut offset, object_hash)) = stack.pop() {
        match ctx.read_object(object_hash).wait()? {
            Object::Data(DataObject::Small(small_object)) => {
                slice[offset as usize..].copy_from_slice(&small_object.chunk);
            }
            Object::Data(DataObject::Large(large_object)) => {
                for (child_size, child_hash) in large_object.children {
                    stack.push((child_size, child_hash));
                    offset += child_size;
                }
            }
            _ => bail!("meep"),
        }
    }

    Ok(())
}

// TODO: Tree diff in order to remove files.
pub fn go(repository: &mut Repository, matches: &ArgMatches) -> Result<()> {
    let commit_hash = matches.value_of("COMMIT").unwrap().parse()?;

    {
        let ctx = repository.local(())?;
        let commit = ctx.read_object(commit_hash).wait()?;
        let commit_object = match commit {
            Object::Commit(commit_object) => commit_object,
            _ => bail!(ErrorKind::NotACommit(commit_hash)),
        };

        let mut stack = vec![(PathBuf::new(), commit_object.subtree)];
        while let Some((path, object)) = stack.pop() {
            match ctx.read_object(object).wait()? {
                Object::Subtree(SubtreeObject { entries }) => {
                    for (component, entry) in entries {
                        let joined = path.join(component);

                        match entry {
                            SubtreeEntry::File(object_hash, size) => {
                                write_data_object(&ctx, joined, object_hash, size)
                                    .chain_err(|| "While trying to write file")?;
                            }
                            SubtreeEntry::Subtree(object_hash) => {
                                stack.push((joined, object_hash));
                            }
                        }
                    }
                }
                _ => bail!("Invalid subtree!"),
            }
        }

        ctx.close().wait()?;

        Ok(())
    }
}
<commit_msg>Fix checkout (bounds error, large object bug)<commit_after>use std::fs::{self, OpenOptions};
use std::path::PathBuf;

use clap::{App, Arg, ArgMatches, SubCommand};
use memmap::{Mmap, Protection};
use futures::prelude::*;

use attaca::context::Context;
use attaca::marshal::{DataObject, Object, ObjectHash, SubtreeEntry, SubtreeObject};
use attaca::repository::Repository;
use attaca::store::ObjectStore;
use attaca::trace::Trace;

use errors::*;

pub fn command() -> App<'static, 'static> {
    SubCommand::with_name("checkout")
        .help("Checkout a file or working directory from a previous revision.")
        .arg(
            Arg::with_name("COMMIT")
                .index(1)
                .required(true)
                .help("The commit hash to checkout."),
        )
}

fn write_data_object<T: Trace, S: ObjectStore>(
    ctx: &Context<T, S>,
    path: PathBuf,
    object_hash: ObjectHash,
    size: u64,
) -> Result<()> {
    let mut parent = path.clone();
    assert!(parent.pop());
    fs::create_dir_all(parent)?;

    let file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(true)
        .open(&path)?;
    file.set_len(size)?;

    // PermissionError indicates mmap opened with wrong permissions.
    let mut mmap = Mmap::open(&file, Protection::ReadWrite)?;

    let slice = unsafe { mmap.as_mut_slice() };
    let mut stack = vec![(0u64, object_hash)];
    while let Some((mut offset, object_hash)) = stack.pop() {
        match ctx.read_object(object_hash).wait()? {
            Object::Data(DataObject::Small(small_object)) => {
                slice[offset as usize..(offset as usize + small_object.chunk.len())]
                    .copy_from_slice(&small_object.chunk);
            }
            Object::Data(DataObject::Large(large_object)) => {
                for (child_size, child_hash) in large_object.children {
                    stack.push((offset, child_hash));
                    offset += child_size;
                }
            }
            _ => bail!("meep"),
        }
    }

    Ok(())
}

// TODO: Tree diff in order to remove files.
pub fn go(repository: &mut Repository, matches: &ArgMatches) -> Result<()> {
    let commit_hash = matches.value_of("COMMIT").unwrap().parse()?;

    {
        let ctx = repository.local(())?;
        let commit = ctx.read_object(commit_hash).wait()?;
        let commit_object = match commit {
            Object::Commit(commit_object) => commit_object,
            _ => bail!(ErrorKind::NotACommit(commit_hash)),
        };

        let mut stack = vec![(PathBuf::new(), commit_object.subtree)];
        while let Some((path, object)) = stack.pop() {
            match ctx.read_object(object).wait()? {
                Object::Subtree(SubtreeObject { entries }) => for (component, entry) in entries {
                    let joined = path.join(component);

                    match entry {
                        SubtreeEntry::File(object_hash, size) => {
                            write_data_object(&ctx, joined, object_hash, size)
                                .chain_err(|| "While trying to write file")?;
                        }
                        SubtreeEntry::Subtree(object_hash) => {
                            stack.push((joined, object_hash));
                        }
                    }
                },
                _ => bail!("Invalid subtree!"),
            }
        }

        ctx.close().wait()?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>use std;
use std::collections::HashMap;
use std::sync::Arc;
use std::thread;
use std::sync::atomic::{AtomicU64, Ordering};
use parking_lot::{RwLock, RwLockWriteGuard};
use serde;

use bifrost_hasher::{hash_str, hash_bytes};
use membership::client::{ObserverClient as MembershipClient, Member};
use conshash::weights::DEFAULT_SERVICE_ID;
use conshash::weights::client::{SMClient as WeightSMClient};
use raft::client::{RaftClient, SubscriptionError, SubscriptionReceipt};
use raft::state_machine::master::ExecError;
use utils::bincode::{serialize};
use rand;

use futures::prelude::*;

pub mod weights;

pub static DEFAULT_NODE_LIST_SIZE: f64 = 2048f64;

#[derive(Debug)]
pub enum Action {
    Joined,
    Left,
}

#[derive(Debug)]
pub enum InitTableError {
    GroupNotExisted,
    NoWeightService(ExecError),
    NoWeightGroup,
    NoWeightInfo,
    Unknown,
}

#[derive(Debug)]
pub enum CHError {
    WatchError(Result<Result<SubscriptionReceipt, SubscriptionError>, ExecError>),
    InitTableError(InitTableError),
}

#[derive(Clone)]
struct Node {
    start: u64,
    server: u64,
}

struct LookupTables {
    nodes: Vec<Node>,
    addrs: HashMap<u64, String>
}

pub struct ConsistentHashing {
    tables: RwLock<LookupTables>,
    membership: Arc<MembershipClient>,
    weight_sm_client: WeightSMClient,
    group_name: String,
    watchers: RwLock<Vec<Box<Fn(&Member, &Action, &LookupTables, &Vec<Node>) + Send + Sync>>>,
    version: AtomicU64
}

impl ConsistentHashing {
    pub fn new_with_id<'a>(id: u64, group: &'a str, raft_client: &Arc<RaftClient>) -> Result<Arc<ConsistentHashing>, CHError> {
        let membership = Arc::new(MembershipClient::new(raft_client));
        let ch = Arc::new(ConsistentHashing {
            tables: RwLock::new(LookupTables {
                nodes: Vec::new(),
                addrs: HashMap::new()
            }),
            membership: membership.clone(),
            weight_sm_client: WeightSMClient::new(id, &raft_client),
            group_name: group.to_string(),
            watchers: RwLock::new(Vec::new()),
            version: AtomicU64::new(0)
        });
        {
            let ch = ch.clone();
            let res = membership.on_group_member_joined(move |r| {
                if let Ok((member, version)) = r { server_joined(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        {
            let ch = ch.clone();
            let res = membership.on_group_member_online(move |r| {
                if let Ok((member, version)) = r { server_joined(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        {
            let ch = ch.clone();
            let res = membership.on_group_member_left(move |r| {
                if let Ok((member, version)) = r { server_left(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        {
            let ch = ch.clone();
            let res = membership.on_group_member_offline(move |r| {
                if let Ok((member, version)) = r { server_left(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        Ok(ch)
    }
    pub fn new<'a>(group: &'a str, raft_client: &Arc<RaftClient>) -> Result<Arc<ConsistentHashing>, CHError> {
        Self::new_with_id(DEFAULT_SERVICE_ID, group, raft_client)
    }
    pub fn new_client<'a>(group: &'a str, raft_client: &Arc<RaftClient>) -> Result<Arc<ConsistentHashing>, CHError> {
        match ConsistentHashing::new(group, raft_client) {
            Err(e) => Err(e),
            Ok(ch) => {
                match ch.init_table() {
                    Err(e) => Err(CHError::InitTableError(e)),
                    Ok(_) =>  Ok(ch.clone())
                }
            },
        }
    }
    pub fn init_table(&self) -> Result<(), InitTableError> {
        let mut table = &mut self.tables.write();
        self.init_table_(&mut table)
    }
    pub fn to_server_name(&self, server_id: Option<u64>) -> Option<String> {
        let lookup_table = self.tables.read();
        match server_id {
            Some(id) => Some(lookup_table.addrs.get(&id).unwrap().clone()),
            None => None
        }
    }
    pub fn get_server_id(&self, hash: u64) -> Option<u64> {
        let lookup_table = self.tables.read();
        let nodes = &lookup_table.nodes;
        let len = nodes.len();
        if len == 0 {return None;}
        let first_node = nodes.first().unwrap();
        let last_node = nodes.last().unwrap();
        if hash < first_node.start {return Some(first_node.server)}
        if hash >= last_node.start {return Some(last_node.server)}
        let mut low = 0;
        let mut high = len - 1;
        while low <= high {
            let mid = low + (high - low) / 2;
            let curr_node = &nodes[mid];
            let next_node = &nodes[mid + 1];
            if hash >= curr_node.start && hash < next_node.start {
                return Some(curr_node.server);
            } else if hash < curr_node.start {
                high = mid - 1;
            } else {
                low = mid + 1;
            }
        }
        None
    }
    pub fn get_server(&self, hash: u64) -> Option<String> {
        self.to_server_name(self.get_server_id(hash))
    }
    pub fn get_server_by_string(&self, string: &String) -> Option<String> {
        self.get_server(hash_str(string))
    }
    pub fn get_server_by<T>(&self, obj: &T) -> Option<String> where T: serde::Serialize {
        self.get_server(hash_bytes(serialize(obj).as_slice()))
    }
    pub fn get_server_id_by_string(&self, string: &String) -> Option<u64> {
        self.get_server_id(hash_str(string))
    }
    pub fn get_server_id_by<T>(&self, obj: &T) -> Option<u64> where T: serde::Serialize {
        self.get_server_id(hash_bytes(serialize(obj).as_slice()))
    }
    pub fn rand_server(&self) -> Option<String> {
        let rand = rand::random::<u64>();
        self.get_server(rand)
    }
    pub fn nodes_count(&self) -> usize {
        let lookup_table = self.tables.read();
        return lookup_table.nodes.len();
    }
    pub fn set_weight(&self, server_name: &String, weight: u64) -> impl Future<Item = bool, Error = ExecError> {
        let group_id = hash_str(&self.group_name);
        let server_id = hash_str(server_name);
        self.weight_sm_client.set_weight(&group_id, &server_id, &weight)
            .map(|res| if let Ok(_) = res { true } else { false })
    }
    pub fn watch_all_actions<F>(&self, f: F)
        where F: Fn(&Member, &Action, &LookupTables, &Vec<Node>) + 'static + Send + Sync {
        let mut watchers = self.watchers.write();
        watchers.push(Box::new(f));
    }
    pub fn watch_server_nodes_range_changed<F>(&self, server: &String, f: F)
        // return ranges [...,...)
        where F: Fn(Vec<(u64, u64)>) + 'static + Send + Sync {
        let server_id = hash_str(server);
        let wrapper = move |
            _: &Member,_: &Action,
            lookup_table: &LookupTables, _: &Vec<Node>| {
            let nodes = &lookup_table.nodes;
            let node_len = nodes.len();
            let mut ranges: Vec<(u64, u64)> = Vec::new();
            for ni in 0..node_len {
                let node = &nodes[ni];
                if node.server == server_id {
                    ranges.push((node.start, if ni == node_len - 1 {
                        std::u64::MAX
                    } else {
                        nodes[ni + 1].start
                    }));
                }
            }
            if !ranges.is_empty() {f(ranges);}
        };
        self.watch_all_actions(wrapper);
    }

    fn init_table_(&self, lookup_table: &mut RwLockWriteGuard<LookupTables>) -> Result<(), InitTableError> {
        if let Ok(Ok((members, version))) = self.membership.group_members(&self.group_name, true).wait() {
            let group_id = hash_str(&self.group_name);
            match self.weight_sm_client.get_weights(&group_id).wait() {
                Ok(Ok(Some(weights))) =>  {
                    if let Some(min_weight) = weights.values().min() {
                        lookup_table.nodes.clear(); // refresh nodes
                        let mut factors: HashMap<u64, f64> = HashMap::new();
                        let min_weight = *min_weight as f64;
                        for member in members.iter() {
                            let k = member.id;
                            let w = match weights.get(&k) {
                                Some(w) => *w as f64,
                                None => min_weight,
                            };
                            factors.insert(k, w / min_weight);
                        }
                        let factor_sum: f64 = factors.values().sum();
                        let nodes_size = if factor_sum > DEFAULT_NODE_LIST_SIZE {factor_sum} else {DEFAULT_NODE_LIST_SIZE};
                        let factor_scale = nodes_size / factor_sum;
                        for member in members.iter() {lookup_table.addrs.insert(member.id, member.address.clone());}
                        for (k, f) in factors.iter() {
                            let weight = (*f * factor_scale) as u64;
                            for i in 0..weight {
                                let node_key = format!("{}_{}", k, i);
                                let hash = hash_str(&node_key);
                                lookup_table.nodes.push(Node{
                                    start: hash, server: *k
                                });
                            }
                        };
                        lookup_table.nodes.sort_by(|n1, n2| n1.start.cmp(&n2.start));
                        self.version.store(version, Ordering::Relaxed);
                        Ok(())
                    } else {
                        Err(InitTableError::NoWeightInfo)
                    }
                },
                Err(e) => Err(InitTableError::NoWeightService(e)),
                Ok(Ok(None)) => Err(InitTableError::NoWeightGroup),
                _ => Err(InitTableError::Unknown)

            }
        } else {
            Err(InitTableError::GroupNotExisted)
        }
    }
}

fn server_joined(ch: &Arc<ConsistentHashing>, member: Member, version: u64) {
    server_changed(ch, member, Action::Joined, version);
}
fn server_left(ch: &Arc<ConsistentHashing>, member: Member, version: u64) {
    server_changed(ch, member, Action::Left, version);
}
fn server_changed(ch: &Arc<ConsistentHashing>, member: Member, action: Action, version: u64) {
    let ch_version = ch.version.load(Ordering::Relaxed);
    if ch_version < version {
        let ch = ch.clone();
        thread::spawn(move || {
            let mut lookup_table = ch.tables.write();
            let watchers = ch.watchers.read();
            let ch_version = ch.version.load(Ordering::Relaxed);
            if ch_version >= version {return;}
            let old_nodes = lookup_table.nodes.clone();
            ch.init_table_(&mut lookup_table);
            for watch in watchers.iter() {
                watch(&member, &action, &*lookup_table, &old_nodes);
            }
        });
    }
}<commit_msg>new client with id<commit_after>use std;
use std::collections::HashMap;
use std::sync::Arc;
use std::thread;
use std::sync::atomic::{AtomicU64, Ordering};
use parking_lot::{RwLock, RwLockWriteGuard};
use serde;

use bifrost_hasher::{hash_str, hash_bytes};
use membership::client::{ObserverClient as MembershipClient, Member};
use conshash::weights::DEFAULT_SERVICE_ID;
use conshash::weights::client::{SMClient as WeightSMClient};
use raft::client::{RaftClient, SubscriptionError, SubscriptionReceipt};
use raft::state_machine::master::ExecError;
use utils::bincode::{serialize};
use rand;

use futures::prelude::*;

pub mod weights;

pub static DEFAULT_NODE_LIST_SIZE: f64 = 2048f64;

#[derive(Debug)]
pub enum Action {
    Joined,
    Left,
}

#[derive(Debug)]
pub enum InitTableError {
    GroupNotExisted,
    NoWeightService(ExecError),
    NoWeightGroup,
    NoWeightInfo,
    Unknown,
}

#[derive(Debug)]
pub enum CHError {
    WatchError(Result<Result<SubscriptionReceipt, SubscriptionError>, ExecError>),
    InitTableError(InitTableError),
}

#[derive(Clone)]
struct Node {
    start: u64,
    server: u64,
}

struct LookupTables {
    nodes: Vec<Node>,
    addrs: HashMap<u64, String>
}

pub struct ConsistentHashing {
    tables: RwLock<LookupTables>,
    membership: Arc<MembershipClient>,
    weight_sm_client: WeightSMClient,
    group_name: String,
    watchers: RwLock<Vec<Box<Fn(&Member, &Action, &LookupTables, &Vec<Node>) + Send + Sync>>>,
    version: AtomicU64
}

impl ConsistentHashing {
    pub fn new_with_id<'a>(id: u64, group: &'a str, raft_client: &Arc<RaftClient>) -> Result<Arc<ConsistentHashing>, CHError> {
        let membership = Arc::new(MembershipClient::new(raft_client));
        let ch = Arc::new(ConsistentHashing {
            tables: RwLock::new(LookupTables {
                nodes: Vec::new(),
                addrs: HashMap::new()
            }),
            membership: membership.clone(),
            weight_sm_client: WeightSMClient::new(id, &raft_client),
            group_name: group.to_string(),
            watchers: RwLock::new(Vec::new()),
            version: AtomicU64::new(0)
        });
        {
            let ch = ch.clone();
            let res = membership.on_group_member_joined(move |r| {
                if let Ok((member, version)) = r { server_joined(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        {
            let ch = ch.clone();
            let res = membership.on_group_member_online(move |r| {
                if let Ok((member, version)) = r { server_joined(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        {
            let ch = ch.clone();
            let res = membership.on_group_member_left(move |r| {
                if let Ok((member, version)) = r { server_left(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        {
            let ch = ch.clone();
            let res = membership.on_group_member_offline(move |r| {
                if let Ok((member, version)) = r { server_left(&ch, member, version); }
            }, group).wait();
            if let Ok(Ok(_)) = res {} else {return Err(CHError::WatchError(res));}
        }
        Ok(ch)
    }
    pub fn new<'a>(group: &'a str, raft_client: &Arc<RaftClient>) -> Result<Arc<ConsistentHashing>, CHError> {
        Self::new_with_id(DEFAULT_SERVICE_ID, group, raft_client)
    }
    pub fn new_client<'a>(group: &'a str, raft_client: &Arc<RaftClient>) -> Result<Arc<ConsistentHashing>, CHError> {
        Self::new_client_with_id(DEFAULT_SERVICE_ID, group, raft_client)
    }
    pub fn new_client_with_id<'a>(id: u64, group: &'a str, raft_client: &Arc<RaftClient>) -> Result<Arc<ConsistentHashing>, CHError> {
        match ConsistentHashing::new_with_id(id, group, raft_client) {
            Err(e) => Err(e),
            Ok(ch) => {
                match ch.init_table() {
                    Err(e) => Err(CHError::InitTableError(e)),
                    Ok(_) =>  Ok(ch.clone())
                }
            },
        }
    }
    pub fn init_table(&self) -> Result<(), InitTableError> {
        let mut table = &mut self.tables.write();
        self.init_table_(&mut table)
    }
    pub fn to_server_name(&self, server_id: Option<u64>) -> Option<String> {
        let lookup_table = self.tables.read();
        match server_id {
            Some(id) => Some(lookup_table.addrs.get(&id).unwrap().clone()),
            None => None
        }
    }
    pub fn get_server_id(&self, hash: u64) -> Option<u64> {
        let lookup_table = self.tables.read();
        let nodes = &lookup_table.nodes;
        let len = nodes.len();
        if len == 0 {return None;}
        let first_node = nodes.first().unwrap();
        let last_node = nodes.last().unwrap();
        if hash < first_node.start {return Some(first_node.server)}
        if hash >= last_node.start {return Some(last_node.server)}
        let mut low = 0;
        let mut high = len - 1;
        while low <= high {
            let mid = low + (high - low) / 2;
            let curr_node = &nodes[mid];
            let next_node = &nodes[mid + 1];
            if hash >= curr_node.start && hash < next_node.start {
                return Some(curr_node.server);
            } else if hash < curr_node.start {
                high = mid - 1;
            } else {
                low = mid + 1;
            }
        }
        None
    }
    pub fn get_server(&self, hash: u64) -> Option<String> {
        self.to_server_name(self.get_server_id(hash))
    }
    pub fn get_server_by_string(&self, string: &String) -> Option<String> {
        self.get_server(hash_str(string))
    }
    pub fn get_server_by<T>(&self, obj: &T) -> Option<String> where T: serde::Serialize {
        self.get_server(hash_bytes(serialize(obj).as_slice()))
    }
    pub fn get_server_id_by_string(&self, string: &String) -> Option<u64> {
        self.get_server_id(hash_str(string))
    }
    pub fn get_server_id_by<T>(&self, obj: &T) -> Option<u64> where T: serde::Serialize {
        self.get_server_id(hash_bytes(serialize(obj).as_slice()))
    }
    pub fn rand_server(&self) -> Option<String> {
        let rand = rand::random::<u64>();
        self.get_server(rand)
    }
    pub fn nodes_count(&self) -> usize {
        let lookup_table = self.tables.read();
        return lookup_table.nodes.len();
    }
    pub fn set_weight(&self, server_name: &String, weight: u64) -> impl Future<Item = bool, Error = ExecError> {
        let group_id = hash_str(&self.group_name);
        let server_id = hash_str(server_name);
        self.weight_sm_client.set_weight(&group_id, &server_id, &weight)
            .map(|res| if let Ok(_) = res { true } else { false })
    }
    pub fn watch_all_actions<F>(&self, f: F)
        where F: Fn(&Member, &Action, &LookupTables, &Vec<Node>) + 'static + Send + Sync {
        let mut watchers = self.watchers.write();
        watchers.push(Box::new(f));
    }
    pub fn watch_server_nodes_range_changed<F>(&self, server: &String, f: F)
        // return ranges [...,...)
        where F: Fn(Vec<(u64, u64)>) + 'static + Send + Sync {
        let server_id = hash_str(server);
        let wrapper = move |
            _: &Member,_: &Action,
            lookup_table: &LookupTables, _: &Vec<Node>| {
            let nodes = &lookup_table.nodes;
            let node_len = nodes.len();
            let mut ranges: Vec<(u64, u64)> = Vec::new();
            for ni in 0..node_len {
                let node = &nodes[ni];
                if node.server == server_id {
                    ranges.push((node.start, if ni == node_len - 1 {
                        std::u64::MAX
                    } else {
                        nodes[ni + 1].start
                    }));
                }
            }
            if !ranges.is_empty() {f(ranges);}
        };
        self.watch_all_actions(wrapper);
    }

    fn init_table_(&self, lookup_table: &mut RwLockWriteGuard<LookupTables>) -> Result<(), InitTableError> {
        if let Ok(Ok((members, version))) = self.membership.group_members(&self.group_name, true).wait() {
            let group_id = hash_str(&self.group_name);
            match self.weight_sm_client.get_weights(&group_id).wait() {
                Ok(Ok(Some(weights))) =>  {
                    if let Some(min_weight) = weights.values().min() {
                        lookup_table.nodes.clear(); // refresh nodes
                        let mut factors: HashMap<u64, f64> = HashMap::new();
                        let min_weight = *min_weight as f64;
                        for member in members.iter() {
                            let k = member.id;
                            let w = match weights.get(&k) {
                                Some(w) => *w as f64,
                                None => min_weight,
                            };
                            factors.insert(k, w / min_weight);
                        }
                        let factor_sum: f64 = factors.values().sum();
                        let nodes_size = if factor_sum > DEFAULT_NODE_LIST_SIZE {factor_sum} else {DEFAULT_NODE_LIST_SIZE};
                        let factor_scale = nodes_size / factor_sum;
                        for member in members.iter() {lookup_table.addrs.insert(member.id, member.address.clone());}
                        for (k, f) in factors.iter() {
                            let weight = (*f * factor_scale) as u64;
                            for i in 0..weight {
                                let node_key = format!("{}_{}", k, i);
                                let hash = hash_str(&node_key);
                                lookup_table.nodes.push(Node{
                                    start: hash, server: *k
                                });
                            }
                        };
                        lookup_table.nodes.sort_by(|n1, n2| n1.start.cmp(&n2.start));
                        self.version.store(version, Ordering::Relaxed);
                        Ok(())
                    } else {
                        Err(InitTableError::NoWeightInfo)
                    }
                },
                Err(e) => Err(InitTableError::NoWeightService(e)),
                Ok(Ok(None)) => Err(InitTableError::NoWeightGroup),
                _ => Err(InitTableError::Unknown)

            }
        } else {
            Err(InitTableError::GroupNotExisted)
        }
    }
}

fn server_joined(ch: &Arc<ConsistentHashing>, member: Member, version: u64) {
    server_changed(ch, member, Action::Joined, version);
}
fn server_left(ch: &Arc<ConsistentHashing>, member: Member, version: u64) {
    server_changed(ch, member, Action::Left, version);
}
fn server_changed(ch: &Arc<ConsistentHashing>, member: Member, action: Action, version: u64) {
    let ch_version = ch.version.load(Ordering::Relaxed);
    if ch_version < version {
        let ch = ch.clone();
        thread::spawn(move || {
            let mut lookup_table = ch.tables.write();
            let watchers = ch.watchers.read();
            let ch_version = ch.version.load(Ordering::Relaxed);
            if ch_version >= version {return;}
            let old_nodes = lookup_table.nodes.clone();
            ch.init_table_(&mut lookup_table);
            for watch in watchers.iter() {
                watch(&member, &action, &*lookup_table, &old_nodes);
            }
        });
    }
}<|endoftext|>


<commit_before>//! Associative Range Query Tree based on [Al.Cash's compact representation]
//! (http://codeforces.com/blog/entry/18051).

/// Colloquially known as a "segtree" in the sport programming literature, it
/// represents a sequence of elements a_i (0 <= i < size) from a monoid (M, +)
/// on which we want to support fast range operations:
///
/// - modify(l, r, f) replaces a_i (l <= i <= r) by f(a_i) for an endomorphism f
/// - query(l, r) returns the aggregate a_l + a_{l+1} + ... + a_r
pub struct ArqTree<T: ArqSpec> {
    d: Vec<Option<T::F>>,
    t: Vec<T::M>,
}

impl<T: ArqSpec> ArqTree<T>
where
    T::F: Clone,
{
    /// Initializes a static balanced tree on top of the given sequence.
    pub fn new(mut init: Vec<T::M>) -> Self {
        let size = init.len();
        let mut t = (0..size).map(|_| T::identity()).collect::<Vec<_>>();
        t.append(&mut init);
        for i in (0..size).rev() {
            t[i] = T::op(&t[i << 1], &t[i << 1 | 1]);
        }
        Self {
            d: vec![None; size],
            t: t,
        }
    }

    fn apply(&mut self, p: usize, f: &T::F) {
        self.t[p] = T::apply(f, &self.t[p]);
        if p < self.d.len() {
            let h = match self.d[p] {
                Some(ref g) => T::compose(f, g),
                None => f.clone(),
            };
            self.d[p] = Some(h);
        }
    }

    fn push(&mut self, p: usize) {
        for s in (1..32).rev() {
            let i = p >> s;
            if let Some(ref f) = self.d[i].take() {
                self.apply(i << 1, f);
                self.apply(i << 1 | 1, f);
            }
        }
    }

    fn pull(&mut self, mut p: usize) {
        while p > 1 {
            p >>= 1;
            if self.d[p].is_none() {
                self.t[p] = T::op(&self.t[p << 1], &self.t[p << 1 | 1]);
            }
        }
    }

    /// Applies the endomorphism f to all entries from l to r, inclusive.
    pub fn modify(&mut self, mut l: usize, mut r: usize, f: &T::F) {
        l += self.d.len();
        r += self.d.len();
        let (l0, r0) = (l, r);
        self.push(l0);
        self.push(r0);
        while l <= r {
            if l & 1 == 1 {
                self.apply(l, f);
                l += 1;
            }
            if r & 1 == 0 {
                self.apply(r, f);
                r -= 1;
            }
            l >>= 1;
            r >>= 1;
        }
        self.pull(l0);
        self.pull(r0);
    }

    /// Returns the aggregate range query on all entries from l to r, inclusive.
    pub fn query(&mut self, mut l: usize, mut r: usize) -> T::M {
        l += self.d.len();
        r += self.d.len();
        self.push(l);
        self.push(r);
        let (mut l_agg, mut r_agg) = (T::identity(), T::identity());
        while l <= r {
            if l & 1 == 1 {
                l_agg = T::op(&l_agg, &self.t[l]);
                l += 1;
            }
            if r & 1 == 0 {
                r_agg = T::op(&self.t[r], &r_agg);
                r -= 1;
            }
            l >>= 1;
            r >>= 1;
        }
        T::op(&l_agg, &r_agg)
    }
}

pub trait ArqSpec {
    type F;
    type M;
    /// Require for all f,g,a: apply(compose(f, g), a) = apply(f, apply(g, a))
    fn compose(f: &Self::F, g: &Self::F) -> Self::F;
    /// Require for all f,a,b: apply(f, op(a, b)) = op(apply(f, a), apply(f, b))
    fn apply(f: &Self::F, a: &Self::M) -> Self::M;
    /// Require for all a,b,c: op(a, op(b, c)) = op(op(a, b), c)
    fn op(a: &Self::M, b: &Self::M) -> Self::M;
    /// Require for all a: op(a, identity()) = op(identity(), a) = a
    fn identity() -> Self::M;
}

/// In this example, we want to support range sum queries and range constant
/// assignments. Note that constant assignment f_c(a) = c is not a endomorphism
/// on integers because f_c(a+b) = c != 2*c = f_c(a) + f_c(b). In intuitive
/// terms, the problem is that the internal nodes of the tree should really be
/// set to a multiple of c, corresponding to the subtree size. So let's augment
/// the monoid type with size information, using the 2D vector (a_i,1) instead
/// of a_i. Now check that f_c((a, s)) = (c*s, s) is indeed an endomorphism:
/// f_c((a,s)+(b,t)) = f_c((a+b,s+t)) = (c*(s+t),s+t) = (c*s,s)+(c*t,t) =
/// f_c((a,s)) + f_c((b,t)).
///
/// # Panics
///
/// Associated functions will panic on overflow.
pub struct AssignAdd;

impl ArqSpec for AssignAdd {
    type F = i64;
    type M = (i64, i64);
    fn compose(f: &Self::F, _: &Self::F) -> Self::F {
        *f
    }
    fn apply(f: &Self::F, a: &Self::M) -> Self::M {
        (f * a.1, a.1)
    }
    fn op(a: &Self::M, b: &Self::M) -> Self::M {
        (a.0 + b.0, a.1 + b.1)
    }
    fn identity() -> Self::M {
        (0, 0)
    }
}

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

    #[test]
    fn test_arq_tree() {
        let mut arq = ArqTree::<AssignAdd>::new(vec![(0, 1); 10]);

        assert_eq!(arq.query(0, 9), (0, 10));

        arq.modify(1, 3, &10);
        arq.modify(3, 5, &1);

        assert_eq!(arq.query(0, 9), (23, 10));
    }
}
<commit_msg>ArqTree panic documentation<commit_after>//! Associative Range Query Tree based on [Al.Cash's compact representation]
//! (http://codeforces.com/blog/entry/18051).

/// Colloquially known as a "segtree" in the sport programming literature, it
/// represents a sequence of elements a_i (0 <= i < size) from a monoid (M, +)
/// on which we want to support fast range operations:
///
/// - modify(l, r, f) replaces a_i (l <= i <= r) by f(a_i) for an endomorphism f
/// - query(l, r) returns the aggregate a_l + a_{l+1} + ... + a_r
pub struct ArqTree<T: ArqSpec> {
    d: Vec<Option<T::F>>,
    t: Vec<T::M>,
}

impl<T: ArqSpec> ArqTree<T>
where
    T::F: Clone,
{
    /// Initializes a static balanced tree on top of the given sequence.
    pub fn new(mut init: Vec<T::M>) -> Self {
        let size = init.len();
        let mut t = (0..size).map(|_| T::identity()).collect::<Vec<_>>();
        t.append(&mut init);
        for i in (0..size).rev() {
            t[i] = T::op(&t[i << 1], &t[i << 1 | 1]);
        }
        Self {
            d: vec![None; size],
            t: t,
        }
    }

    fn apply(&mut self, p: usize, f: &T::F) {
        self.t[p] = T::apply(f, &self.t[p]);
        if p < self.d.len() {
            let h = match self.d[p] {
                Some(ref g) => T::compose(f, g),
                None => f.clone(),
            };
            self.d[p] = Some(h);
        }
    }

    fn push(&mut self, p: usize) {
        for s in (1..32).rev() {
            let i = p >> s;
            if let Some(ref f) = self.d[i].take() {
                self.apply(i << 1, f);
                self.apply(i << 1 | 1, f);
            }
        }
    }

    fn pull(&mut self, mut p: usize) {
        while p > 1 {
            p >>= 1;
            if self.d[p].is_none() {
                self.t[p] = T::op(&self.t[p << 1], &self.t[p << 1 | 1]);
            }
        }
    }

    /// Applies the endomorphism f to all entries from l to r, inclusive.
    ///
    /// # Panics
    ///
    /// Panics if l or r is out of range.
    pub fn modify(&mut self, mut l: usize, mut r: usize, f: &T::F) {
        l += self.d.len();
        r += self.d.len();
        let (l0, r0) = (l, r);
        self.push(l0);
        self.push(r0);
        while l <= r {
            if l & 1 == 1 {
                self.apply(l, f);
                l += 1;
            }
            if r & 1 == 0 {
                self.apply(r, f);
                r -= 1;
            }
            l >>= 1;
            r >>= 1;
        }
        self.pull(l0);
        self.pull(r0);
    }

    /// Returns the aggregate range query on all entries from l to r, inclusive.
    ///
    /// # Panics
    ///
    /// Panics if l or r is out of range.
    pub fn query(&mut self, mut l: usize, mut r: usize) -> T::M {
        l += self.d.len();
        r += self.d.len();
        self.push(l);
        self.push(r);
        let (mut l_agg, mut r_agg) = (T::identity(), T::identity());
        while l <= r {
            if l & 1 == 1 {
                l_agg = T::op(&l_agg, &self.t[l]);
                l += 1;
            }
            if r & 1 == 0 {
                r_agg = T::op(&self.t[r], &r_agg);
                r -= 1;
            }
            l >>= 1;
            r >>= 1;
        }
        T::op(&l_agg, &r_agg)
    }
}

pub trait ArqSpec {
    type F;
    type M;
    /// Require for all f,g,a: apply(compose(f, g), a) = apply(f, apply(g, a))
    fn compose(f: &Self::F, g: &Self::F) -> Self::F;
    /// Require for all f,a,b: apply(f, op(a, b)) = op(apply(f, a), apply(f, b))
    fn apply(f: &Self::F, a: &Self::M) -> Self::M;
    /// Require for all a,b,c: op(a, op(b, c)) = op(op(a, b), c)
    fn op(a: &Self::M, b: &Self::M) -> Self::M;
    /// Require for all a: op(a, identity()) = op(identity(), a) = a
    fn identity() -> Self::M;
}

/// In this example, we want to support range sum queries and range constant
/// assignments. Note that constant assignment f_c(a) = c is not a endomorphism
/// on integers because f_c(a+b) = c != 2*c = f_c(a) + f_c(b). In intuitive
/// terms, the problem is that the internal nodes of the tree should really be
/// set to a multiple of c, corresponding to the subtree size. So let's augment
/// the monoid type with size information, using the 2D vector (a_i,1) instead
/// of a_i. Now check that f_c((a, s)) = (c*s, s) is indeed an endomorphism:
/// f_c((a,s)+(b,t)) = f_c((a+b,s+t)) = (c*(s+t),s+t) = (c*s,s)+(c*t,t) =
/// f_c((a,s)) + f_c((b,t)).
///
/// # Panics
///
/// Associated functions will panic on overflow.
pub struct AssignAdd;

impl ArqSpec for AssignAdd {
    type F = i64;
    type M = (i64, i64);
    fn compose(f: &Self::F, _: &Self::F) -> Self::F {
        *f
    }
    fn apply(f: &Self::F, a: &Self::M) -> Self::M {
        (f * a.1, a.1)
    }
    fn op(a: &Self::M, b: &Self::M) -> Self::M {
        (a.0 + b.0, a.1 + b.1)
    }
    fn identity() -> Self::M {
        (0, 0)
    }
}

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

    #[test]
    fn test_arq_tree() {
        let mut arq = ArqTree::<AssignAdd>::new(vec![(0, 1); 10]);

        assert_eq!(arq.query(0, 9), (0, 10));

        arq.modify(1, 3, &10);
        arq.modify(3, 5, &1);

        assert_eq!(arq.query(0, 9), (23, 10));
    }
}
<|endoftext|>