<commit_before>//! Visitor for a run-time value with a given layout: Traverse enums, structs and other compound
//! types until we arrive at the leaves, with custom handling for primitive types.

use rustc::ty::layout::{self, TyLayout};
use rustc::ty;
use rustc::mir::interpret::{
    EvalResult,
};

use super::{
    Machine, EvalContext, MPlaceTy, PlaceTy, OpTy,
};

// A thing that we can project into, and that has a layout.
// This wouldn't have to depend on `Machine` but with the current type inference,
// that's just more convenient to work with (avoids repeading all the `Machine` bounds).
pub trait Value<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>>: Copy
{
    // Get this value's layout.
    fn layout(&self) -> TyLayout<'tcx>;

    // Make this a `MPlaceTy`, or panic if that's not possible.
    fn to_mem_place(
        self,
        ectx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, MPlaceTy<'tcx, M::PointerTag>>;

    // Create this from an `MPlaceTy`.
    fn from_mem_place(MPlaceTy<'tcx, M::PointerTag>) -> Self;

    // Read the current enum discriminant, and downcast to that.  Also return the
    // variant index.
    fn project_downcast(
        self,
        ectx: &EvalContext<'a, 'mir, 'tcx, M>
    ) -> EvalResult<'tcx, (Self, usize)>;

    // Project to the n-th field.
    fn project_field(
        self,
        ectx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self>;
}

// Operands and places are both values
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for OpTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_mem_place(
        self,
        _ectx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, MPlaceTy<'tcx, M::PointerTag>> {
        Ok(self.to_mem_place())
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace.into()
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ectx: &EvalContext<'a, 'mir, 'tcx, M>
    ) -> EvalResult<'tcx, (Self, usize)> {
        let idx = ectx.read_discriminant(self)?.1;
        Ok((ectx.operand_downcast(self, idx)?, idx))
    }

    #[inline(always)]
    fn project_field(
        self,
        ectx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ectx.operand_field(self, field)
    }
}
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for MPlaceTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_mem_place(
        self,
        _ectx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, MPlaceTy<'tcx, M::PointerTag>> {
        Ok(self)
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ectx: &EvalContext<'a, 'mir, 'tcx, M>
    ) -> EvalResult<'tcx, (Self, usize)> {
        let idx = ectx.read_discriminant(self.into())?.1;
        Ok((ectx.mplace_downcast(self, idx)?, idx))
    }

    #[inline(always)]
    fn project_field(
        self,
        ectx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ectx.mplace_field(self, field)
    }
}
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for PlaceTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_mem_place(
        self,
        ectx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, MPlaceTy<'tcx, M::PointerTag>> {
        // If this refers to a local, assert that it already has an allocation.
        Ok(ectx.place_to_op(self)?.to_mem_place())
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace.into()
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ectx: &EvalContext<'a, 'mir, 'tcx, M>
    ) -> EvalResult<'tcx, (Self, usize)> {
        let idx = ectx.read_discriminant(ectx.place_to_op(self)?)?.1;
        Ok((ectx.place_downcast(self, idx)?, idx))
    }

    #[inline(always)]
    fn project_field(
        self,
        ectx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ectx.place_field(self, field)
    }
}

// How to traverse a value and what to do when we are at the leaves.
pub trait ValueVisitor<'a, 'mir, 'tcx: 'mir+'a, M: Machine<'a, 'mir, 'tcx>>: Sized {
    type V: Value<'a, 'mir, 'tcx, M>;

    // The visitor must have an `EvalContext` in it.
    fn ecx(&mut self) -> &mut EvalContext<'a, 'mir, 'tcx, M>;

    // Recursie actions, ready to be overloaded.
    /// Visit the given value, dispatching as appropriate to more speicalized visitors.
    #[inline(always)]
    fn visit_value(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        self.walk_value(v)
    }
    /// Visit the given value as a union.
    #[inline(always)]
    fn visit_union(&mut self, _v: Self::V) -> EvalResult<'tcx>
    {
        Ok(())
    }
    /// Visit the given value as an array.
    #[inline(always)]
    fn visit_array(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        self.walk_array(v)
    }
    /// Called each time we recurse down to a field, passing in old and new value.
    /// This gives the visitor the chance to track the stack of nested fields that
    /// we are descending through.
    #[inline(always)]
    fn visit_field(
        &mut self,
        _old_val: Self::V,
        _field: usize,
        new_val: Self::V,
    ) -> EvalResult<'tcx> {
        self.visit_value(new_val)
    }

    // Actions on the leaves, ready to be overloaded.
    /// Called whenever we reach a value with uninhabited layout.
    /// Recursing to fields will continue after this!
    #[inline(always)]
    fn visit_uninhabited(&mut self, _v: Self::V) -> EvalResult<'tcx>
    { Ok(()) }
    /// Called whenever we reach a value with scalar layout.
    /// Recursing to fields will continue after this!
    #[inline(always)]
    fn visit_scalar(&mut self, _v: Self::V, _layout: &layout::Scalar) -> EvalResult<'tcx>
    { Ok(()) }
    /// Called whenever we reach a value of primitive type.  There can be no recursion
    /// below such a value.
    #[inline(always)]
    fn visit_primitive(&mut self, _v: Self::V) -> EvalResult<'tcx>
    { Ok(()) }

    // Default recursors. Not meant to be overloaded.
    fn walk_array(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        // Let's get an mplace first.
        let mplace = if v.layout().is_zst() {
            // it's a ZST, the memory content cannot matter
            MPlaceTy::dangling(v.layout(), self.ecx())
        } else {
            // non-ZST array/slice/str cannot be immediate
            v.to_mem_place(self.ecx())?
        };
        // Now iterate over it.
        for (i, field) in self.ecx().mplace_array_fields(mplace)?.enumerate() {
            self.visit_field(v, i, Value::from_mem_place(field?))?;
        }
        Ok(())
    }
    fn walk_value(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        // If this is a multi-variant layout, we have find the right one and proceed with that.
        // (No benefit from making this recursion, but it is equivalent to that.)
        match v.layout().variants {
            layout::Variants::NicheFilling { .. } |
            layout::Variants::Tagged { .. } => {
                let (inner, idx) = v.project_downcast(self.ecx())?;
                trace!("variant layout: {:#?}", inner.layout());
                // recurse with the inner type
                return self.visit_field(v, idx, inner);
            }
            layout::Variants::Single { .. } => {}
        }

        // Even for single variants, we might be able to get a more refined type:
        // If it is a trait object, switch to the actual type that was used to create it.
        match v.layout().ty.sty {
            ty::Dynamic(..) => {
                // immediate trait objects are not a thing
                let dest = v.to_mem_place(self.ecx())?;
                let inner = self.ecx().unpack_dyn_trait(dest)?.1;
                trace!("dyn object layout: {:#?}", inner.layout);
                // recurse with the inner type
                return self.visit_field(v, 0, Value::from_mem_place(inner));
            },
            _ => {},
        };

        // If this is a scalar, visit it as such.
        // Things can be aggregates and have scalar layout at the same time, and that
        // is very relevant for `NonNull` and similar structs: We need to visit them
        // at their scalar layout *before* descending into their fields.
        // FIXME: We could avoid some redundant checks here. For newtypes wrapping
        // scalars, we do the same check on every "level" (e.g. first we check
        // MyNewtype and then the scalar in there).
        match v.layout().abi {
            layout::Abi::Uninhabited => {
                self.visit_uninhabited(v)?;
            }
            layout::Abi::Scalar(ref layout) => {
                self.visit_scalar(v, layout)?;
            }
            // FIXME: Should we do something for ScalarPair? Vector?
            _ => {}
        }

        // Check primitive types.  We do this after checking the scalar layout,
        // just to have that done as well.  Primitives can have varying layout,
        // so we check them separately and before aggregate handling.
        // It is CRITICAL that we get this check right, or we might be
        // validating the wrong thing!
        let primitive = match v.layout().fields {
            // Primitives appear as Union with 0 fields -- except for Boxes and fat pointers.
            layout::FieldPlacement::Union(0) => true,
            _ => v.layout().ty.builtin_deref(true).is_some(),
        };
        if primitive {
            return self.visit_primitive(v);
        }

        // Proceed into the fields.
        match v.layout().fields {
            layout::FieldPlacement::Union(fields) => {
                // Empty unions are not accepted by rustc. That's great, it means we can
                // use that as an unambiguous signal for detecting primitives.  Make sure
                // we did not miss any primitive.
                debug_assert!(fields > 0);
                self.visit_union(v)?;
            },
            layout::FieldPlacement::Arbitrary { ref offsets, .. } => {
                for i in 0..offsets.len() {
                    let val = v.project_field(self.ecx(), i as u64)?;
                    self.visit_field(v, i, val)?;
                }
            },
            layout::FieldPlacement::Array { .. } => {
                self.visit_array(v)?;
            }
        }
        Ok(())
    }
}
<commit_msg>all values can convert to operators<commit_after>//! Visitor for a run-time value with a given layout: Traverse enums, structs and other compound
//! types until we arrive at the leaves, with custom handling for primitive types.

use rustc::ty::layout::{self, TyLayout};
use rustc::ty;
use rustc::mir::interpret::{
    EvalResult,
};

use super::{
    Machine, EvalContext, MPlaceTy, PlaceTy, OpTy,
};

// A thing that we can project into, and that has a layout.
// This wouldn't have to depend on `Machine` but with the current type inference,
// that's just more convenient to work with (avoids repeading all the `Machine` bounds).
pub trait Value<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>>: Copy
{
    // Get this value's layout.
    fn layout(&self) -> TyLayout<'tcx>;

    // Make this into an `OpTy`.
    fn to_op(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>>;

    // Create this from an `MPlaceTy`.
    fn from_mem_place(MPlaceTy<'tcx, M::PointerTag>) -> Self;

    // Project to the given enum variant.
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: usize,
    ) -> EvalResult<'tcx, Self>;

    // Project to the n-th field.
    fn project_field(
        self,
        ecx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self>;
}

// Operands and places are both values
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for OpTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_op(
        self,
        _ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {
        Ok(self)
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace.into()
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: usize,
    ) -> EvalResult<'tcx, Self> {
        ecx.operand_downcast(self, variant)
    }

    #[inline(always)]
    fn project_field(
        self,
        ecx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ecx.operand_field(self, field)
    }
}
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for MPlaceTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_op(
        self,
        _ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {
        Ok(self.into())
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: usize,
    ) -> EvalResult<'tcx, Self> {
        ecx.mplace_downcast(self, variant)
    }

    #[inline(always)]
    fn project_field(
        self,
        ecx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ecx.mplace_field(self, field)
    }
}
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for PlaceTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_op(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {
        ecx.place_to_op(self)
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace.into()
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: usize,
    ) -> EvalResult<'tcx, Self> {
        ecx.place_downcast(self, variant)
    }

    #[inline(always)]
    fn project_field(
        self,
        ecx: &mut EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ecx.place_field(self, field)
    }
}

// How to traverse a value and what to do when we are at the leaves.
pub trait ValueVisitor<'a, 'mir, 'tcx: 'mir+'a, M: Machine<'a, 'mir, 'tcx>>: Sized {
    type V: Value<'a, 'mir, 'tcx, M>;

    // The visitor must have an `EvalContext` in it.
    fn ecx(&mut self) -> &mut EvalContext<'a, 'mir, 'tcx, M>;

    // Recursie actions, ready to be overloaded.
    /// Visit the given value, dispatching as appropriate to more speicalized visitors.
    #[inline(always)]
    fn visit_value(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        self.walk_value(v)
    }
    /// Visit the given value as a union.
    #[inline(always)]
    fn visit_union(&mut self, _v: Self::V) -> EvalResult<'tcx>
    {
        Ok(())
    }
    /// Visit the given value as an array.
    #[inline(always)]
    fn visit_array(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        self.walk_array(v)
    }
    /// Called each time we recurse down to a field, passing in old and new value.
    /// This gives the visitor the chance to track the stack of nested fields that
    /// we are descending through.
    #[inline(always)]
    fn visit_field(
        &mut self,
        _old_val: Self::V,
        _field: usize,
        new_val: Self::V,
    ) -> EvalResult<'tcx> {
        self.visit_value(new_val)
    }

    // Actions on the leaves, ready to be overloaded.
    /// Called whenever we reach a value with uninhabited layout.
    /// Recursing to fields will continue after this!
    #[inline(always)]
    fn visit_uninhabited(&mut self, _v: Self::V) -> EvalResult<'tcx>
    { Ok(()) }
    /// Called whenever we reach a value with scalar layout.
    /// Recursing to fields will continue after this!
    #[inline(always)]
    fn visit_scalar(&mut self, _v: Self::V, _layout: &layout::Scalar) -> EvalResult<'tcx>
    { Ok(()) }
    /// Called whenever we reach a value of primitive type.  There can be no recursion
    /// below such a value.
    #[inline(always)]
    fn visit_primitive(&mut self, _v: Self::V) -> EvalResult<'tcx>
    { Ok(()) }

    // Default recursors. Not meant to be overloaded.
    fn walk_array(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        // Let's get an mplace first.
        let mplace = if v.layout().is_zst() {
            // it's a ZST, the memory content cannot matter
            MPlaceTy::dangling(v.layout(), self.ecx())
        } else {
            // non-ZST array/slice/str cannot be immediate
            v.to_op(self.ecx())?.to_mem_place()
        };
        // Now iterate over it.
        for (i, field) in self.ecx().mplace_array_fields(mplace)?.enumerate() {
            self.visit_field(v, i, Value::from_mem_place(field?))?;
        }
        Ok(())
    }
    fn walk_value(&mut self, v: Self::V) -> EvalResult<'tcx>
    {
        // If this is a multi-variant layout, we have find the right one and proceed with that.
        // (No benefit from making this recursion, but it is equivalent to that.)
        match v.layout().variants {
            layout::Variants::NicheFilling { .. } |
            layout::Variants::Tagged { .. } => {
                let op = v.to_op(self.ecx())?;
                let idx = self.ecx().read_discriminant(op)?.1;
                let inner = v.project_downcast(self.ecx(), idx)?;
                trace!("walk_value: variant layout: {:#?}", inner.layout());
                // recurse with the inner type
                return self.visit_field(v, idx, inner);
            }
            layout::Variants::Single { .. } => {}
        }

        // Even for single variants, we might be able to get a more refined type:
        // If it is a trait object, switch to the actual type that was used to create it.
        match v.layout().ty.sty {
            ty::Dynamic(..) => {
                // immediate trait objects are not a thing
                let dest = v.to_op(self.ecx())?.to_mem_place();
                let inner = self.ecx().unpack_dyn_trait(dest)?.1;
                trace!("walk_value: dyn object layout: {:#?}", inner.layout);
                // recurse with the inner type
                return self.visit_field(v, 0, Value::from_mem_place(inner));
            },
            _ => {},
        };

        // If this is a scalar, visit it as such.
        // Things can be aggregates and have scalar layout at the same time, and that
        // is very relevant for `NonNull` and similar structs: We need to visit them
        // at their scalar layout *before* descending into their fields.
        // FIXME: We could avoid some redundant checks here. For newtypes wrapping
        // scalars, we do the same check on every "level" (e.g. first we check
        // MyNewtype and then the scalar in there).
        match v.layout().abi {
            layout::Abi::Uninhabited => {
                self.visit_uninhabited(v)?;
            }
            layout::Abi::Scalar(ref layout) => {
                self.visit_scalar(v, layout)?;
            }
            // FIXME: Should we do something for ScalarPair? Vector?
            _ => {}
        }

        // Check primitive types.  We do this after checking the scalar layout,
        // just to have that done as well.  Primitives can have varying layout,
        // so we check them separately and before aggregate handling.
        // It is CRITICAL that we get this check right, or we might be
        // validating the wrong thing!
        let primitive = match v.layout().fields {
            // Primitives appear as Union with 0 fields -- except for Boxes and fat pointers.
            layout::FieldPlacement::Union(0) => true,
            _ => v.layout().ty.builtin_deref(true).is_some(),
        };
        if primitive {
            return self.visit_primitive(v);
        }

        // Proceed into the fields.
        match v.layout().fields {
            layout::FieldPlacement::Union(fields) => {
                // Empty unions are not accepted by rustc. That's great, it means we can
                // use that as an unambiguous signal for detecting primitives.  Make sure
                // we did not miss any primitive.
                debug_assert!(fields > 0);
                self.visit_union(v)?;
            },
            layout::FieldPlacement::Arbitrary { ref offsets, .. } => {
                for i in 0..offsets.len() {
                    let val = v.project_field(self.ecx(), i as u64)?;
                    self.visit_field(v, i, val)?;
                }
            },
            layout::FieldPlacement::Array { .. } => {
                self.visit_array(v)?;
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Mappings of objects received from the API, with optional helper methods for
//! ease of use.
//!
//! Models can optionally have additional helper methods compiled, by enabling
//! the `model` feature.

#[macro_use]
mod utils;

pub mod event;
pub mod permissions;

mod channel;
mod error;
mod gateway;
mod guild;
mod invite;
mod misc;
mod user;
mod voice;
mod webhook;

pub use self::channel::*;
pub use self::error::Error as ModelError;
pub use self::gateway::*;
pub use self::guild::*;
pub use self::invite::*;
pub use self::misc::*;
pub use self::permissions::Permissions;
pub use self::user::*;
pub use self::voice::*;
pub use self::webhook::*;

use chrono::NaiveDateTime;
use internal::prelude::*;
use parking_lot::RwLock;
use self::utils::*;
use serde::de::Visitor;
use std::collections::HashMap;
use std::fmt::{Display, Formatter, Result as FmtResult};
use std::sync::Arc;

#[cfg(feature = "utils")]
use utils::Colour;

fn default_true() -> bool { true }

macro_rules! id_u64 {
    ($(#[$attr:meta] $name:ident;)*) => {
        $(
            #[$attr]
            #[derive(Copy, Clone, Default, Debug, Eq, Hash, PartialOrd, Ord, Serialize)]
            #[allow(derive_hash_xor_eq)]
            pub struct $name(pub u64);

            impl $name {
                /// Retrieves the time that the Id was created at.
                pub fn created_at(&self) -> NaiveDateTime {
                    let offset = (self.0 >> 22) / 1000;

                    NaiveDateTime::from_timestamp(1_420_070_400 + offset as i64, 0)
                }
            }

            // This is a hack so functions can accept iterators that either:
            // 1. return the id itself (e.g: `MessageId`)
            // 2. return a reference to it (`&MessageId`).
            impl AsRef<$name> for $name {
                fn as_ref(&self) -> &Self {
                    self
                }
            }

            impl From<u64> for $name {
                fn from(id_as_u64: u64) -> $name {
                    $name(id_as_u64)
                }
            }

            impl PartialEq for $name {
                fn eq(&self, other: &Self) -> bool {
                    self.0 == other.0
                }
            }

            impl PartialEq<u64> for $name {
                fn eq(&self, u: &u64) -> bool {
                    self.0 == *u
                }
            }

            impl Display for $name {
                fn fmt(&self, f: &mut Formatter) -> FmtResult {
                    Display::fmt(&self.0, f)
                }
            }

            impl<'de> Deserialize<'de> for $name {
                fn deserialize<D: Deserializer<'de>>(deserializer: D) -> StdResult<Self, D::Error> {
                    deserializer.deserialize_u64(U64Visitor).map($name)
                }
            }
        )*
    }
}

id_u64! {
    /// An identifier for a Channel
    ChannelId;
    /// An identifier for an Emoji
    EmojiId;
    /// An identifier for a Guild
    GuildId;
    /// An identifier for an Integration
    IntegrationId;
    /// An identifier for a Message
    MessageId;
    /// An identifier for a Role
    RoleId;
    /// An identifier for a User
    UserId;
    /// An identifier for a [`Webhook`](struct.Webhook.html).
    WebhookId;
    /// An identifier for an audit log entry.
    AuditLogEntryId;
}

/// A container for guilds.
///
/// This is used to differentiate whether a guild itself can be used or whether
/// a guild needs to be retrieved from the cache.
#[allow(large_enum_variant)]
#[derive(Clone, Debug)]
pub enum GuildContainer {
    /// A guild which can have its contents directly searched.
    Guild(PartialGuild),
    /// A guild's id, which can be used to search the cache for a guild.
    Id(GuildId),
}

/// Information about a user's application. An application does not necessarily
/// have an associated bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct ApplicationInfo {
    /// The bot user associated with the application. See [`BotApplication`] for
    /// more information.
    ///
    /// [`BotApplication`]: struct.BotApplication.html
    pub bot: Option<BotApplication>,
    /// Indicator of whether the bot is public.
    ///
    /// If a bot is public, anyone may invite it to their [`Guild`]. While a bot
    /// is private, only the owner may add it to a guild.
    ///
    /// [`Guild`]: struct.Guild.html
    #[serde(default = "default_true")]
    pub bot_public: bool,
    /// Indicator of whether the bot requires an OAuth2 code grant.
    pub bot_require_code_grant: bool,
    /// A description of the application, assigned by the application owner.
    pub description: String,
    /// A set of bitflags assigned to the application, which represent gated
    /// feature flags that have been enabled for the application.
    pub flags: Option<u64>,
    /// A hash pointing to the application's icon.
    ///
    /// This is not necessarily equivalent to the bot user's avatar.
    pub icon: Option<String>,
    /// The unique numeric Id of the application.
    pub id: UserId,
    /// The name assigned to the application by the application owner.
    pub name: String,
    /// A list of redirect URIs assigned to the application.
    pub redirect_uris: Vec<String>,
    /// A list of RPC Origins assigned to the application.
    pub rpc_origins: Vec<String>,
    /// The given secret to the application.
    ///
    /// This is not equivalent to the application's bot user's token.
    pub secret: String,
}

/// Information about an application with an application's bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct BotApplication {
    /// The unique Id of the bot user.
    pub id: UserId,
    /// A hash of the avatar, if one is assigned.
    ///
    /// Can be used to generate a full URL to the avatar.
    pub avatar: Option<String>,
    /// Indicator of whether it is a bot.
    #[serde(default)]
    pub bot: bool,
    /// The discriminator assigned to the bot user.
    ///
    /// While discriminators are not unique, the `username#discriminator` pair
    /// is.
    pub discriminator: u16,
    /// The bot user's username.
    pub name: String,
    /// The token used to authenticate as the bot user.
    ///
    /// **Note**: Keep this information private, as untrusted sources can use it
    /// to perform any action with a bot user.
    pub token: String,
}

/// Information about the current application and its owner.
#[derive(Clone, Debug, Deserialize)]
pub struct CurrentApplicationInfo {
    pub description: String,
    pub icon: Option<String>,
    pub id: UserId,
    pub name: String,
    pub owner: User,
    #[serde(default)] pub rpc_origins: Vec<String>,
}

/// The name of a region that a voice server can be located in.
#[derive(Copy, Clone, Debug, Deserialize, Eq, Hash, PartialEq, PartialOrd, Ord, Serialize)]
pub enum Region {
    #[serde(rename = "amsterdam")] Amsterdam,
    #[serde(rename = "brazil")] Brazil,
    #[serde(rename = "eu-central")] EuCentral,
    #[serde(rename = "eu-west")] EuWest,
    #[serde(rename = "frankfurt")] Frankfurt,
    #[serde(rename = "london")] London,
    #[serde(rename = "sydney")] Sydney,
    #[serde(rename = "us-central")] UsCentral,
    #[serde(rename = "us-east")] UsEast,
    #[serde(rename = "us-south")] UsSouth,
    #[serde(rename = "us-west")] UsWest,
    #[serde(rename = "vip-amsterdam")] VipAmsterdam,
    #[serde(rename = "vip-us-east")] VipUsEast,
    #[serde(rename = "vip-us-west")] VipUsWest,
}

impl Region {
    pub fn name(&self) -> &str {
        match *self {
            Region::Amsterdam => "amsterdam",
            Region::Brazil => "brazil",
            Region::EuCentral => "eu-central",
            Region::EuWest => "eu-west",
            Region::Frankfurt => "frankfurt",
            Region::London => "london",
            Region::Sydney => "sydney",
            Region::UsCentral => "us-central",
            Region::UsEast => "us-east",
            Region::UsSouth => "us-south",
            Region::UsWest => "us-west",
            Region::VipAmsterdam => "vip-amsterdam",
            Region::VipUsEast => "vip-us-east",
            Region::VipUsWest => "vip-us-west",
        }
    }
}

use serde::{Deserialize, Deserializer};
use std::result::Result as StdResult;

fn deserialize_sync_user<'de, D: Deserializer<'de>>(deserializer: D)
                                                    -> StdResult<Arc<RwLock<User>>, D::Error> {
    Ok(Arc::new(RwLock::new(User::deserialize(deserializer)?)))
}
<commit_msg>Fix snowflake deserializer<commit_after>//! Mappings of objects received from the API, with optional helper methods for
//! ease of use.
//!
//! Models can optionally have additional helper methods compiled, by enabling
//! the `model` feature.

#[macro_use]
mod utils;

pub mod event;
pub mod permissions;

mod channel;
mod error;
mod gateway;
mod guild;
mod invite;
mod misc;
mod user;
mod voice;
mod webhook;

pub use self::channel::*;
pub use self::error::Error as ModelError;
pub use self::gateway::*;
pub use self::guild::*;
pub use self::invite::*;
pub use self::misc::*;
pub use self::permissions::Permissions;
pub use self::user::*;
pub use self::voice::*;
pub use self::webhook::*;

use chrono::NaiveDateTime;
use internal::prelude::*;
use parking_lot::RwLock;
use self::utils::*;
use serde::de::Visitor;
use std::collections::HashMap;
use std::fmt::{Display, Formatter, Result as FmtResult};
use std::sync::Arc;

#[cfg(feature = "utils")]
use utils::Colour;

fn default_true() -> bool { true }

macro_rules! id_u64 {
    ($(#[$attr:meta] $name:ident;)*) => {
        $(
            #[$attr]
            #[derive(Copy, Clone, Default, Debug, Eq, Hash, PartialOrd, Ord, Serialize)]
            #[allow(derive_hash_xor_eq)]
            pub struct $name(pub u64);

            impl $name {
                /// Retrieves the time that the Id was created at.
                pub fn created_at(&self) -> NaiveDateTime {
                    let offset = (self.0 >> 22) / 1000;

                    NaiveDateTime::from_timestamp(1_420_070_400 + offset as i64, 0)
                }
            }

            // This is a hack so functions can accept iterators that either:
            // 1. return the id itself (e.g: `MessageId`)
            // 2. return a reference to it (`&MessageId`).
            impl AsRef<$name> for $name {
                fn as_ref(&self) -> &Self {
                    self
                }
            }

            impl From<u64> for $name {
                fn from(id_as_u64: u64) -> $name {
                    $name(id_as_u64)
                }
            }

            impl PartialEq for $name {
                fn eq(&self, other: &Self) -> bool {
                    self.0 == other.0
                }
            }

            impl PartialEq<u64> for $name {
                fn eq(&self, u: &u64) -> bool {
                    self.0 == *u
                }
            }

            impl Display for $name {
                fn fmt(&self, f: &mut Formatter) -> FmtResult {
                    Display::fmt(&self.0, f)
                }
            }

            impl<'de> Deserialize<'de> for $name {
                fn deserialize<D: Deserializer<'de>>(deserializer: D) -> StdResult<Self, D::Error> {
                    deserializer.deserialize_any(U64Visitor).map($name)
                }
            }
        )*
    }
}

id_u64! {
    /// An identifier for a Channel
    ChannelId;
    /// An identifier for an Emoji
    EmojiId;
    /// An identifier for a Guild
    GuildId;
    /// An identifier for an Integration
    IntegrationId;
    /// An identifier for a Message
    MessageId;
    /// An identifier for a Role
    RoleId;
    /// An identifier for a User
    UserId;
    /// An identifier for a [`Webhook`](struct.Webhook.html).
    WebhookId;
    /// An identifier for an audit log entry.
    AuditLogEntryId;
}

/// A container for guilds.
///
/// This is used to differentiate whether a guild itself can be used or whether
/// a guild needs to be retrieved from the cache.
#[allow(large_enum_variant)]
#[derive(Clone, Debug)]
pub enum GuildContainer {
    /// A guild which can have its contents directly searched.
    Guild(PartialGuild),
    /// A guild's id, which can be used to search the cache for a guild.
    Id(GuildId),
}

/// Information about a user's application. An application does not necessarily
/// have an associated bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct ApplicationInfo {
    /// The bot user associated with the application. See [`BotApplication`] for
    /// more information.
    ///
    /// [`BotApplication`]: struct.BotApplication.html
    pub bot: Option<BotApplication>,
    /// Indicator of whether the bot is public.
    ///
    /// If a bot is public, anyone may invite it to their [`Guild`]. While a bot
    /// is private, only the owner may add it to a guild.
    ///
    /// [`Guild`]: struct.Guild.html
    #[serde(default = "default_true")]
    pub bot_public: bool,
    /// Indicator of whether the bot requires an OAuth2 code grant.
    pub bot_require_code_grant: bool,
    /// A description of the application, assigned by the application owner.
    pub description: String,
    /// A set of bitflags assigned to the application, which represent gated
    /// feature flags that have been enabled for the application.
    pub flags: Option<u64>,
    /// A hash pointing to the application's icon.
    ///
    /// This is not necessarily equivalent to the bot user's avatar.
    pub icon: Option<String>,
    /// The unique numeric Id of the application.
    pub id: UserId,
    /// The name assigned to the application by the application owner.
    pub name: String,
    /// A list of redirect URIs assigned to the application.
    pub redirect_uris: Vec<String>,
    /// A list of RPC Origins assigned to the application.
    pub rpc_origins: Vec<String>,
    /// The given secret to the application.
    ///
    /// This is not equivalent to the application's bot user's token.
    pub secret: String,
}

/// Information about an application with an application's bot user.
#[derive(Clone, Debug, Deserialize)]
pub struct BotApplication {
    /// The unique Id of the bot user.
    pub id: UserId,
    /// A hash of the avatar, if one is assigned.
    ///
    /// Can be used to generate a full URL to the avatar.
    pub avatar: Option<String>,
    /// Indicator of whether it is a bot.
    #[serde(default)]
    pub bot: bool,
    /// The discriminator assigned to the bot user.
    ///
    /// While discriminators are not unique, the `username#discriminator` pair
    /// is.
    pub discriminator: u16,
    /// The bot user's username.
    pub name: String,
    /// The token used to authenticate as the bot user.
    ///
    /// **Note**: Keep this information private, as untrusted sources can use it
    /// to perform any action with a bot user.
    pub token: String,
}

/// Information about the current application and its owner.
#[derive(Clone, Debug, Deserialize)]
pub struct CurrentApplicationInfo {
    pub description: String,
    pub icon: Option<String>,
    pub id: UserId,
    pub name: String,
    pub owner: User,
    #[serde(default)] pub rpc_origins: Vec<String>,
}

/// The name of a region that a voice server can be located in.
#[derive(Copy, Clone, Debug, Deserialize, Eq, Hash, PartialEq, PartialOrd, Ord, Serialize)]
pub enum Region {
    #[serde(rename = "amsterdam")] Amsterdam,
    #[serde(rename = "brazil")] Brazil,
    #[serde(rename = "eu-central")] EuCentral,
    #[serde(rename = "eu-west")] EuWest,
    #[serde(rename = "frankfurt")] Frankfurt,
    #[serde(rename = "london")] London,
    #[serde(rename = "sydney")] Sydney,
    #[serde(rename = "us-central")] UsCentral,
    #[serde(rename = "us-east")] UsEast,
    #[serde(rename = "us-south")] UsSouth,
    #[serde(rename = "us-west")] UsWest,
    #[serde(rename = "vip-amsterdam")] VipAmsterdam,
    #[serde(rename = "vip-us-east")] VipUsEast,
    #[serde(rename = "vip-us-west")] VipUsWest,
}

impl Region {
    pub fn name(&self) -> &str {
        match *self {
            Region::Amsterdam => "amsterdam",
            Region::Brazil => "brazil",
            Region::EuCentral => "eu-central",
            Region::EuWest => "eu-west",
            Region::Frankfurt => "frankfurt",
            Region::London => "london",
            Region::Sydney => "sydney",
            Region::UsCentral => "us-central",
            Region::UsEast => "us-east",
            Region::UsSouth => "us-south",
            Region::UsWest => "us-west",
            Region::VipAmsterdam => "vip-amsterdam",
            Region::VipUsEast => "vip-us-east",
            Region::VipUsWest => "vip-us-west",
        }
    }
}

use serde::{Deserialize, Deserializer};
use std::result::Result as StdResult;

fn deserialize_sync_user<'de, D: Deserializer<'de>>(deserializer: D)
                                                    -> StdResult<Arc<RwLock<User>>, D::Error> {
    Ok(Arc::new(RwLock::new(User::deserialize(deserializer)?)))
}
<|endoftext|>


<commit_before>use crate::hasher::{Digest, Size};
use crate::Error;
use core::convert::TryFrom;
#[cfg(feature = "std")]
use core::convert::TryInto;
use core::fmt::Debug;
use generic_array::{ArrayLength, GenericArray};

/// Trait that implements hashing.
///
/// It is usually implemented by a custom code table enum that derives the [`Multihash` derive].
///
/// [`Multihash` derive]: crate::derive
pub trait MultihashCode:
    TryFrom<u64> + Into<u64> + Send + Sync + Unpin + Copy + Eq + Debug + 'static
{
    /// The maximum size a hash will allocate.
    type AllocSize: Size;

    /// Calculate the hash of some input data.
    ///
    /// # Example
    ///
    /// ```
    /// // `Code` implements `MultihashCode`
    /// use multihash::{Code, MultihashCode};
    ///
    /// let hash = Code::Sha3_256.digest(b"Hello world!");
    /// println!("{:02x?}", hash);
    /// ```
    fn digest(&self, input: &[u8]) -> Multihash<Self::AllocSize>;

    /// Create a multihash from an existing [`Digest`].
    ///
    /// # Example
    ///
    /// ```
    /// use multihash::{Code, MultihashCode, Sha3_256, StatefulHasher};
    ///
    /// let mut hasher = Sha3_256::default();
    /// hasher.update(b"Hello world!");
    /// let hash = Code::multihash_from_digest(&hasher.finalize());
    /// println!("{:02x?}", hash);
    /// ```
    fn multihash_from_digest<'a, S, D>(digest: &'a D) -> Multihash<Self::AllocSize>
    where
        S: Size,
        D: Digest<S>,
        Self: From<&'a D>;
}

/// A Multihash instance that only supports the basic functionality and no hashing.
///
/// With this Multihash implementation you can operate on Multihashes in a generic way, but
/// no hasher implementation is associated with the code.
///
/// # Example
///
/// ```
/// use multihash::{Multihash, U64};
///
/// const Sha3_256: u64 = 0x16;
/// let digest_bytes = [
///     0x16, 0x20, 0x64, 0x4b, 0xcc, 0x7e, 0x56, 0x43, 0x73, 0x04, 0x09, 0x99, 0xaa, 0xc8, 0x9e,
///     0x76, 0x22, 0xf3, 0xca, 0x71, 0xfb, 0xa1, 0xd9, 0x72, 0xfd, 0x94, 0xa3, 0x1c, 0x3b, 0xfb,
///     0xf2, 0x4e, 0x39, 0x38,
/// ];
/// let mh = Multihash::<U64>::from_bytes(&digest_bytes).unwrap();
/// assert_eq!(mh.code(), Sha3_256);
/// assert_eq!(mh.size(), 32);
/// assert_eq!(mh.digest(), &digest_bytes[2..]);
/// ```
#[cfg_attr(feature = "serde-codec", derive(serde::Deserialize))]
#[cfg_attr(feature = "serde-codec", derive(serde::Serialize))]
#[cfg_attr(feature = "serde-codec", serde(bound = "S: Size"))]
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Multihash<S: Size> {
    /// The code of the Multihash.
    code: u64,
    /// The actual size of the digest in bytes (not the allocated size).
    size: u8,
    /// The digest.
    digest: GenericArray<u8, S>,
}

impl<S: Size> Copy for Multihash<S> where <S as ArrayLength<u8>>::ArrayType: Copy {}

impl<S: Size> Multihash<S> {
    /// Wraps the digest in a multihash.
    pub fn wrap(code: u64, input_digest: &[u8]) -> Result<Self, Error> {
        if input_digest.len() > S::to_usize() {
            return Err(Error::InvalidSize(input_digest.len() as _));
        }
        let size = input_digest.len();
        let mut digest = GenericArray::default();
        digest[..size].copy_from_slice(input_digest);
        Ok(Self {
            code,
            size: size as u8,
            digest,
        })
    }

    /// Returns the code of the multihash.
    pub fn code(&self) -> u64 {
        self.code
    }

    /// Returns the size of the digest.
    pub fn size(&self) -> u8 {
        self.size
    }

    /// Returns the digest.
    pub fn digest(&self) -> &[u8] {
        &self.digest[..self.size as usize]
    }

    /// Reads a multihash from a byte stream.
    #[cfg(feature = "std")]
    pub fn read<R: std::io::Read>(r: R) -> Result<Self, Error>
    where
        Self: Sized,
    {
        let (code, size, digest) = read_multihash(r)?;
        Ok(Self { code, size, digest })
    }

    /// Parses a multihash from a bytes.
    ///
    /// You need to make sure the passed in bytes have the correct length. The digest length
    /// needs to match the `size` value of the multihash.
    #[cfg(feature = "std")]
    pub fn from_bytes(mut bytes: &[u8]) -> Result<Self, Error>
    where
        Self: Sized,
    {
        let result = Self::read(&mut bytes)?;
        // There were more bytes supplied than read
        if !bytes.is_empty() {
            return Err(Error::InvalidSize(bytes.len().try_into().expect(
                "Currently the maximum size is 255, therefore always fits into usize",
            )));
        }

        Ok(result)
    }

    /// Writes a multihash to a byte stream.
    #[cfg(feature = "std")]
    pub fn write<W: std::io::Write>(&self, w: W) -> Result<(), Error> {
        write_multihash(w, self.code(), self.size(), self.digest())
    }

    /// Returns the bytes of a multihash.
    #[cfg(feature = "std")]
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = vec![];
        self.write(&mut bytes)
            .expect("writing to a vec should never fail");
        bytes
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Encode for Multihash<crate::U32> {
    fn encode_to<EncOut: parity_scale_codec::Output>(&self, dest: &mut EncOut) {
        let mut digest = [0; 32];
        digest.copy_from_slice(&self.digest);
        dest.push(&self.code);
        dest.push(&self.size);
        dest.push(&digest);
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::EncodeLike for Multihash<crate::U32> {}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Decode for Multihash<crate::U32> {
    fn decode<DecIn: parity_scale_codec::Input>(
        input: &mut DecIn,
    ) -> Result<Self, parity_scale_codec::Error> {
        Ok(Multihash {
            code: parity_scale_codec::Decode::decode(input)?,
            size: parity_scale_codec::Decode::decode(input)?,
            digest: {
                let digest = <[u8; 32]>::decode(input)?;
                GenericArray::clone_from_slice(&digest)
            },
        })
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Encode for Multihash<crate::U64> {
    fn encode_to<EncOut: parity_scale_codec::Output>(&self, dest: &mut EncOut) {
        let mut digest = [0; 64];
        digest.copy_from_slice(&self.digest);
        dest.push(&self.code);
        dest.push(&self.size);
        dest.push(&digest);
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::EncodeLike for Multihash<crate::U64> {}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Decode for Multihash<crate::U64> {
    fn decode<DecIn: parity_scale_codec::Input>(
        input: &mut DecIn,
    ) -> Result<Self, parity_scale_codec::Error> {
        Ok(Multihash {
            code: parity_scale_codec::Decode::decode(input)?,
            size: parity_scale_codec::Decode::decode(input)?,
            digest: {
                let digest = <[u8; 64]>::decode(input)?;
                GenericArray::clone_from_slice(&digest)
            },
        })
    }
}

/// Writes the multihash to a byte stream.
#[cfg(feature = "std")]
pub fn write_multihash<W>(mut w: W, code: u64, size: u8, digest: &[u8]) -> Result<(), Error>
where
    W: std::io::Write,
{
    use unsigned_varint::encode as varint_encode;

    let mut code_buf = varint_encode::u64_buffer();
    let code = varint_encode::u64(code, &mut code_buf);

    let mut size_buf = varint_encode::u8_buffer();
    let size = varint_encode::u8(size, &mut size_buf);

    w.write_all(code)?;
    w.write_all(size)?;
    w.write_all(digest)?;
    Ok(())
}

/// Reads a multihash from a byte stream that contains a full multihash (code, size and the digest)
///
/// Returns the code, size and the digest. The size is the actual size and not the
/// maximum/allocated size of the digest.
///
/// Currently the maximum size for a digest is 255 bytes.
#[cfg(feature = "std")]
pub fn read_multihash<R, S>(mut r: R) -> Result<(u64, u8, GenericArray<u8, S>), Error>
where
    R: std::io::Read,
    S: Size,
{
    use unsigned_varint::io::read_u64;

    let code = read_u64(&mut r)?;
    let size = read_u64(&mut r)?;

    if size > S::to_u64() || size > u8::MAX as u64 {
        return Err(Error::InvalidSize(size));
    }

    let mut digest = GenericArray::default();
    r.read_exact(&mut digest[..size as usize])?;
    Ok((code, size as u8, digest))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::multihash_impl::Code;

    #[test]
    fn roundtrip() {
        let hash = Code::Sha2_256.digest(b"hello world");
        let mut buf = [0u8; 35];
        hash.write(&mut buf[..]).unwrap();
        let hash2 = Multihash::read(&buf[..]).unwrap();
        assert_eq!(hash, hash2);
    }

    #[test]
    #[cfg(feature = "scale-codec")]
    fn test_scale() {
        use parity_scale_codec::{Decode, Encode};

        let mh = Multihash::<crate::U32>::default();
        let bytes = mh.encode();
        let mh2: Multihash<crate::U32> = Decode::decode(&mut &bytes[..]).unwrap();
        assert_eq!(mh, mh2);
    }

    #[test]
    #[cfg(feature = "serde-codec")]
    fn test_serde() {
        let mh = Multihash::<crate::U32>::default();
        let bytes = serde_json::to_string(&mh).unwrap();
        let mh2 = serde_json::from_str(&bytes).unwrap();
        assert_eq!(mh, mh2);
    }
}
<commit_msg>fix: allocate approriate size<commit_after>use crate::hasher::{Digest, Size};
use crate::Error;
use core::convert::TryFrom;
#[cfg(feature = "std")]
use core::convert::TryInto;
use core::fmt::Debug;
use generic_array::{ArrayLength, GenericArray};

/// Trait that implements hashing.
///
/// It is usually implemented by a custom code table enum that derives the [`Multihash` derive].
///
/// [`Multihash` derive]: crate::derive
pub trait MultihashCode:
    TryFrom<u64> + Into<u64> + Send + Sync + Unpin + Copy + Eq + Debug + 'static
{
    /// The maximum size a hash will allocate.
    type AllocSize: Size;

    /// Calculate the hash of some input data.
    ///
    /// # Example
    ///
    /// ```
    /// // `Code` implements `MultihashCode`
    /// use multihash::{Code, MultihashCode};
    ///
    /// let hash = Code::Sha3_256.digest(b"Hello world!");
    /// println!("{:02x?}", hash);
    /// ```
    fn digest(&self, input: &[u8]) -> Multihash<Self::AllocSize>;

    /// Create a multihash from an existing [`Digest`].
    ///
    /// # Example
    ///
    /// ```
    /// use multihash::{Code, MultihashCode, Sha3_256, StatefulHasher};
    ///
    /// let mut hasher = Sha3_256::default();
    /// hasher.update(b"Hello world!");
    /// let hash = Code::multihash_from_digest(&hasher.finalize());
    /// println!("{:02x?}", hash);
    /// ```
    fn multihash_from_digest<'a, S, D>(digest: &'a D) -> Multihash<Self::AllocSize>
    where
        S: Size,
        D: Digest<S>,
        Self: From<&'a D>;
}

/// A Multihash instance that only supports the basic functionality and no hashing.
///
/// With this Multihash implementation you can operate on Multihashes in a generic way, but
/// no hasher implementation is associated with the code.
///
/// # Example
///
/// ```
/// use multihash::{Multihash, U64};
///
/// const Sha3_256: u64 = 0x16;
/// let digest_bytes = [
///     0x16, 0x20, 0x64, 0x4b, 0xcc, 0x7e, 0x56, 0x43, 0x73, 0x04, 0x09, 0x99, 0xaa, 0xc8, 0x9e,
///     0x76, 0x22, 0xf3, 0xca, 0x71, 0xfb, 0xa1, 0xd9, 0x72, 0xfd, 0x94, 0xa3, 0x1c, 0x3b, 0xfb,
///     0xf2, 0x4e, 0x39, 0x38,
/// ];
/// let mh = Multihash::<U64>::from_bytes(&digest_bytes).unwrap();
/// assert_eq!(mh.code(), Sha3_256);
/// assert_eq!(mh.size(), 32);
/// assert_eq!(mh.digest(), &digest_bytes[2..]);
/// ```
#[cfg_attr(feature = "serde-codec", derive(serde::Deserialize))]
#[cfg_attr(feature = "serde-codec", derive(serde::Serialize))]
#[cfg_attr(feature = "serde-codec", serde(bound = "S: Size"))]
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Multihash<S: Size> {
    /// The code of the Multihash.
    code: u64,
    /// The actual size of the digest in bytes (not the allocated size).
    size: u8,
    /// The digest.
    digest: GenericArray<u8, S>,
}

impl<S: Size> Copy for Multihash<S> where <S as ArrayLength<u8>>::ArrayType: Copy {}

impl<S: Size> Multihash<S> {
    /// Wraps the digest in a multihash.
    pub fn wrap(code: u64, input_digest: &[u8]) -> Result<Self, Error> {
        if input_digest.len() > S::to_usize() {
            return Err(Error::InvalidSize(input_digest.len() as _));
        }
        let size = input_digest.len();
        let mut digest = GenericArray::default();
        digest[..size].copy_from_slice(input_digest);
        Ok(Self {
            code,
            size: size as u8,
            digest,
        })
    }

    /// Returns the code of the multihash.
    pub fn code(&self) -> u64 {
        self.code
    }

    /// Returns the size of the digest.
    pub fn size(&self) -> u8 {
        self.size
    }

    /// Returns the digest.
    pub fn digest(&self) -> &[u8] {
        &self.digest[..self.size as usize]
    }

    /// Reads a multihash from a byte stream.
    #[cfg(feature = "std")]
    pub fn read<R: std::io::Read>(r: R) -> Result<Self, Error>
    where
        Self: Sized,
    {
        let (code, size, digest) = read_multihash(r)?;
        Ok(Self { code, size, digest })
    }

    /// Parses a multihash from a bytes.
    ///
    /// You need to make sure the passed in bytes have the correct length. The digest length
    /// needs to match the `size` value of the multihash.
    #[cfg(feature = "std")]
    pub fn from_bytes(mut bytes: &[u8]) -> Result<Self, Error>
    where
        Self: Sized,
    {
        let result = Self::read(&mut bytes)?;
        // There were more bytes supplied than read
        if !bytes.is_empty() {
            return Err(Error::InvalidSize(bytes.len().try_into().expect(
                "Currently the maximum size is 255, therefore always fits into usize",
            )));
        }

        Ok(result)
    }

    /// Writes a multihash to a byte stream.
    #[cfg(feature = "std")]
    pub fn write<W: std::io::Write>(&self, w: W) -> Result<(), Error> {
        write_multihash(w, self.code(), self.size(), self.digest())
    }

    /// Returns the bytes of a multihash.
    #[cfg(feature = "std")]
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(self.size().into());
        self.write(&mut bytes)
            .expect("writing to a vec should never fail");
        bytes
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Encode for Multihash<crate::U32> {
    fn encode_to<EncOut: parity_scale_codec::Output>(&self, dest: &mut EncOut) {
        let mut digest = [0; 32];
        digest.copy_from_slice(&self.digest);
        dest.push(&self.code);
        dest.push(&self.size);
        dest.push(&digest);
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::EncodeLike for Multihash<crate::U32> {}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Decode for Multihash<crate::U32> {
    fn decode<DecIn: parity_scale_codec::Input>(
        input: &mut DecIn,
    ) -> Result<Self, parity_scale_codec::Error> {
        Ok(Multihash {
            code: parity_scale_codec::Decode::decode(input)?,
            size: parity_scale_codec::Decode::decode(input)?,
            digest: {
                let digest = <[u8; 32]>::decode(input)?;
                GenericArray::clone_from_slice(&digest)
            },
        })
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Encode for Multihash<crate::U64> {
    fn encode_to<EncOut: parity_scale_codec::Output>(&self, dest: &mut EncOut) {
        let mut digest = [0; 64];
        digest.copy_from_slice(&self.digest);
        dest.push(&self.code);
        dest.push(&self.size);
        dest.push(&digest);
    }
}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::EncodeLike for Multihash<crate::U64> {}

#[cfg(feature = "scale-codec")]
impl parity_scale_codec::Decode for Multihash<crate::U64> {
    fn decode<DecIn: parity_scale_codec::Input>(
        input: &mut DecIn,
    ) -> Result<Self, parity_scale_codec::Error> {
        Ok(Multihash {
            code: parity_scale_codec::Decode::decode(input)?,
            size: parity_scale_codec::Decode::decode(input)?,
            digest: {
                let digest = <[u8; 64]>::decode(input)?;
                GenericArray::clone_from_slice(&digest)
            },
        })
    }
}

/// Writes the multihash to a byte stream.
#[cfg(feature = "std")]
pub fn write_multihash<W>(mut w: W, code: u64, size: u8, digest: &[u8]) -> Result<(), Error>
where
    W: std::io::Write,
{
    use unsigned_varint::encode as varint_encode;

    let mut code_buf = varint_encode::u64_buffer();
    let code = varint_encode::u64(code, &mut code_buf);

    let mut size_buf = varint_encode::u8_buffer();
    let size = varint_encode::u8(size, &mut size_buf);

    w.write_all(code)?;
    w.write_all(size)?;
    w.write_all(digest)?;
    Ok(())
}

/// Reads a multihash from a byte stream that contains a full multihash (code, size and the digest)
///
/// Returns the code, size and the digest. The size is the actual size and not the
/// maximum/allocated size of the digest.
///
/// Currently the maximum size for a digest is 255 bytes.
#[cfg(feature = "std")]
pub fn read_multihash<R, S>(mut r: R) -> Result<(u64, u8, GenericArray<u8, S>), Error>
where
    R: std::io::Read,
    S: Size,
{
    use unsigned_varint::io::read_u64;

    let code = read_u64(&mut r)?;
    let size = read_u64(&mut r)?;

    if size > S::to_u64() || size > u8::MAX as u64 {
        return Err(Error::InvalidSize(size));
    }

    let mut digest = GenericArray::default();
    r.read_exact(&mut digest[..size as usize])?;
    Ok((code, size as u8, digest))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::multihash_impl::Code;

    #[test]
    fn roundtrip() {
        let hash = Code::Sha2_256.digest(b"hello world");
        let mut buf = [0u8; 35];
        hash.write(&mut buf[..]).unwrap();
        let hash2 = Multihash::read(&buf[..]).unwrap();
        assert_eq!(hash, hash2);
    }

    #[test]
    #[cfg(feature = "scale-codec")]
    fn test_scale() {
        use parity_scale_codec::{Decode, Encode};

        let mh = Multihash::<crate::U32>::default();
        let bytes = mh.encode();
        let mh2: Multihash<crate::U32> = Decode::decode(&mut &bytes[..]).unwrap();
        assert_eq!(mh, mh2);
    }

    #[test]
    #[cfg(feature = "serde-codec")]
    fn test_serde() {
        let mh = Multihash::<crate::U32>::default();
        let bytes = serde_json::to_string(&mh).unwrap();
        let mh2 = serde_json::from_str(&bytes).unwrap();
        assert_eq!(mh, mh2);
    }
}
<|endoftext|>


<commit_before>use std::io::{TempDir, MemWriter};
use std::str::from_utf8;
use schema;
use buffer;

#[deriving(Show, Eq, Hash)]
struct Register {
	record: schema::Record,
	datatype: schema::SqlType,
}

impl Register {
	fn new(rec: schema::Record, typ: schema::SqlType) -> Register {
		Register {
			record: rec,
			datatype: typ,
		}
	}

	fn set_int(&mut self, val: int) {
		// TODO
	}

	fn get_int(&self) -> int {
		self.record.to_int()
	}

	fn set_str(&mut self) {
		// TODO
	}

	fn get_str<'a>(&'a self) -> &'a str {
		let d = self.record.get_data();
		from_utf8(d).unwrap()
	}
}

impl Ord for Register {
	fn lt(&self, other: &Register) -> bool {
		// not the most robust implementation
		self.get_int() < other.get_int()
	}
}

trait Operatorish<T>: Iterator<T> {
	//fn open();
	//fn close(&self);
}

struct TableScan<'a, 'b> {
	relation: schema::Relation,
	segment: &'a mut schema::SPSegment<'b>,
	current: uint,
}

impl<'a, 'b> TableScan<'a, 'b> {
	fn new(rel: schema::Relation, seg: &'a mut schema::SPSegment<'b>) -> TableScan<'a, 'b> {
		TableScan {
			relation: rel,
			segment: seg,
			current: 0,
		}
	}
}

impl<'a, 'b> Operatorish<Vec<Register>> for TableScan<'a, 'b> {
}

impl<'a, 'b> Iterator<Vec<Register>> for TableScan<'a, 'b> {
	fn next(&mut self) -> Option<Vec<Register>> {
		if (self.current as u64) < self.relation.inserted {
			let tup = self.relation.get(self.segment, self.current);
			let mut res = tup.move_iter().map(|(v, t)| Register::new(v, t)).
				collect::<Vec<Register>>();
			debug!("TS: {}", res);
			self.current += 1;
			Some(res)
		} else {
			None
		}
	}
}

struct Print<'a, T, V> {
	input: T,
	output: &'a mut V,
}

impl<'a, T: Operatorish<Vec<Register>>, V: Writer> Print<'a, T, V> {
	fn new(input: T, output: &'a mut V) -> Print<'a, T, V> {
		Print {
			input: input,
			output: output
		}
	}
}

impl<'a, T: Operatorish<Vec<Register>>, V: Writer> Iterator<Vec<Register>> for Print<'a, T, V> {
	fn next(&mut self) -> Option<Vec<Register>> {
		let cur = self.input.next();
		match cur {
			Some(reg) => {
				for item in reg.iter() {
					match item.datatype {
						schema::Varchar(_) => self.output.write(item.get_str().as_bytes()),
						schema::Integer => self.output.write(
							format!("{}", item.get_int()).as_bytes())
					};
					self.output.write(", ".as_bytes());
				}
				self.output.write("\n".as_bytes());
				Some(reg)
			},
			None => None,
		}
	}
}

struct Project<T> {
	input: T,
	registerids: Vec<uint>,
}

impl<T: Operatorish<Vec<Register>>> Project<T> {
	fn new(input: T, regs: Vec<uint>) -> Project<T> {
		Project {
			input: input,
			registerids: regs,
		}
	}
}

impl<T: Operatorish<Vec<Register>>> Iterator<Vec<Register>> for Project<T> {
	fn next(&mut self) -> Option<Vec<Register>> {
		let cur = self.input.next();
		match cur {
			None => None,
			Some(mut reg) => {
				let mut res = Vec::with_capacity(self.registerids.len());
				let mut regids = self.registerids.clone();
				regids.reverse();
				for index in regids.move_iter() {
					let value = reg.remove(index);
					let v = match value {
						None => fail!("Projection doesn't have vield {}", index),
						Some(v) => v,
					};
					res.push(v);
				}
				Some(res)
			}
		}
	}
}

impl<T: Operatorish<Vec<Register>>> Operatorish<Vec<Register>> for Project<T> {
}

#[deriving(Show)]
enum Selectable {
	Varchar(~str),
	Integer(int),
}

struct Select<T> {
	input: T,
	index: uint,
	value: Selectable,
}

impl<T: Operatorish<Vec<Register>>> Select<T> {
	fn new(input: T, index: uint, value: Selectable) -> Select<T> {
		Select {
			input: input,
			index: index,
			value: value,
		}
	}
}

impl<T: Operatorish<Vec<Register>>> Iterator<Vec<Register>> for Select<T> {
	fn next(&mut self) -> Option<Vec<Register>> {
		// initial value
		let mut cur = self.input.next();
		while cur.is_some() {
			// we made sure cur is Some(…), so we can unwrap safely
			let reg = cur.unwrap();

			// check what we got to select for and retrieve values
			// accordingly
			match self.value {
				// comparing to a string
				Varchar(ref v) => {
					println!("Comparing {} with {}",
						reg.get(self.index).get_str(),
						v);
					// getting the register value as string
					if reg.get(self.index).get_str() == *v {
						return Some(reg)
					} else {
						// not equal, skip to next entry
						cur = self.input.next();
						continue
					}
				},
				// comparing to a number
				Integer(v) => {
					// compare against reg as integer
					if reg.get(self.index).get_int() == v {
						return Some(reg)
					} else {
						// not equal, skip to next entry
						cur = self.input.next();
						continue
					}
				}
			}
		}
		// no more entries, signal iterator exhaustion
		None
	}
}

impl<T: Operatorish<Vec<Register>>> Operatorish<Vec<Register>> for Select<T> {
}

#[test]
fn simple_tablescan() {
	let dir = match TempDir::new("tablescan") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};

	//let p = dir.path();
	let p = Path::new(".");

	let mut manager = buffer::BufferManager::new(1024, p.clone());
	let mut seg = schema::SPSegment {id: 1, manager: &mut manager};

	let name = schema::Column::new(~"name", schema::Varchar(128), vec!(schema::NotNull));
	let age = schema::Column::new(~"age", schema::Integer, vec!(schema::NotNull));
	let mut relation = schema::Relation::new(~"Person");
	relation.add_column(name);
	relation.add_column(age);
	relation.insert(&mut seg, vec!(schema::Record::from_str(~"Alice"), schema::Record::from_int(20)));
	relation.insert(&mut seg, vec!(schema::Record::from_str(~"Bob"), schema::Record::from_int(40)));

	{
		let mut ts = TableScan::new(relation.clone(), &mut seg);
		for tuple in ts {
			println!("Got entry {}", tuple);
		}
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut pr = Print::new(ts, &mut mw);
			// force write by iterating, strange API
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut pr = Project::new(ts, vec!(0));
			let mut pr = Print::new(pr, &mut mw);
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut se = Select::new(ts, 1, Integer(20));
			let mut pr = Print::new(se, &mut mw);
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut se = Select::new(ts, 0, Varchar(~"Bob"));
			let mut pr = Print::new(se, &mut mw);
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	assert!(false);
}
<commit_msg>HashJoin constructor<commit_after>use std::io::{TempDir, MemWriter};
use std::str::from_utf8;
use schema;
use buffer;

#[deriving(Show, Eq, Hash)]
struct Register {
	record: schema::Record,
	datatype: schema::SqlType,
}

impl Register {
	fn new(rec: schema::Record, typ: schema::SqlType) -> Register {
		Register {
			record: rec,
			datatype: typ,
		}
	}

	fn set_int(&mut self, val: int) {
		// TODO
	}

	fn get_int(&self) -> int {
		self.record.to_int()
	}

	fn set_str(&mut self) {
		// TODO
	}

	fn get_str<'a>(&'a self) -> &'a str {
		let d = self.record.get_data();
		from_utf8(d).unwrap()
	}
}

impl Ord for Register {
	fn lt(&self, other: &Register) -> bool {
		// not the most robust implementation
		self.get_int() < other.get_int()
	}
}

trait Operatorish<T>: Iterator<T> {
	//fn open();
	//fn close(&self);
}

struct TableScan<'a, 'b> {
	relation: schema::Relation,
	segment: &'a mut schema::SPSegment<'b>,
	current: uint,
}

impl<'a, 'b> TableScan<'a, 'b> {
	fn new(rel: schema::Relation, seg: &'a mut schema::SPSegment<'b>) -> TableScan<'a, 'b> {
		TableScan {
			relation: rel,
			segment: seg,
			current: 0,
		}
	}
}

impl<'a, 'b> Operatorish<Vec<Register>> for TableScan<'a, 'b> {
}

impl<'a, 'b> Iterator<Vec<Register>> for TableScan<'a, 'b> {
	fn next(&mut self) -> Option<Vec<Register>> {
		if (self.current as u64) < self.relation.inserted {
			let tup = self.relation.get(self.segment, self.current);
			let mut res = tup.move_iter().map(|(v, t)| Register::new(v, t)).
				collect::<Vec<Register>>();
			debug!("TS: {}", res);
			self.current += 1;
			Some(res)
		} else {
			None
		}
	}
}

struct Print<'a, T, V> {
	input: T,
	output: &'a mut V,
}

impl<'a, T: Operatorish<Vec<Register>>, V: Writer> Print<'a, T, V> {
	fn new(input: T, output: &'a mut V) -> Print<'a, T, V> {
		Print {
			input: input,
			output: output
		}
	}
}

impl<'a, T: Operatorish<Vec<Register>>, V: Writer> Iterator<Vec<Register>> for Print<'a, T, V> {
	fn next(&mut self) -> Option<Vec<Register>> {
		let cur = self.input.next();
		match cur {
			Some(reg) => {
				for item in reg.iter() {
					match item.datatype {
						schema::Varchar(_) => self.output.write(item.get_str().as_bytes()),
						schema::Integer => self.output.write(
							format!("{}", item.get_int()).as_bytes())
					};
					self.output.write(", ".as_bytes());
				}
				self.output.write("\n".as_bytes());
				Some(reg)
			},
			None => None,
		}
	}
}

struct Project<T> {
	input: T,
	registerids: Vec<uint>,
}

impl<T: Operatorish<Vec<Register>>> Project<T> {
	fn new(input: T, regs: Vec<uint>) -> Project<T> {
		Project {
			input: input,
			registerids: regs,
		}
	}
}

impl<T: Operatorish<Vec<Register>>> Iterator<Vec<Register>> for Project<T> {
	fn next(&mut self) -> Option<Vec<Register>> {
		let cur = self.input.next();
		match cur {
			None => None,
			Some(mut reg) => {
				let mut res = Vec::with_capacity(self.registerids.len());
				let mut regids = self.registerids.clone();
				regids.reverse();
				for index in regids.move_iter() {
					let value = reg.remove(index);
					let v = match value {
						None => fail!("Projection doesn't have vield {}", index),
						Some(v) => v,
					};
					res.push(v);
				}
				Some(res)
			}
		}
	}
}

impl<T: Operatorish<Vec<Register>>> Operatorish<Vec<Register>> for Project<T> {
}

#[deriving(Show)]
enum Selectable {
	Varchar(~str),
	Integer(int),
}

struct Select<T> {
	input: T,
	index: uint,
	value: Selectable,
}

impl<T: Operatorish<Vec<Register>>> Select<T> {
	fn new(input: T, index: uint, value: Selectable) -> Select<T> {
		Select {
			input: input,
			index: index,
			value: value,
		}
	}
}

impl<T: Operatorish<Vec<Register>>> Iterator<Vec<Register>> for Select<T> {
	fn next(&mut self) -> Option<Vec<Register>> {
		// initial value
		let mut cur = self.input.next();
		while cur.is_some() {
			// we made sure cur is Some(…), so we can unwrap safely
			let reg = cur.unwrap();

			// check what we got to select for and retrieve values
			// accordingly
			match self.value {
				// comparing to a string
				Varchar(ref v) => {
					println!("Comparing {} with {}",
						reg.get(self.index).get_str(),
						v);
					// getting the register value as string
					if reg.get(self.index).get_str() == *v {
						return Some(reg)
					} else {
						// not equal, skip to next entry
						cur = self.input.next();
						continue
					}
				},
				// comparing to a number
				Integer(v) => {
					// compare against reg as integer
					if reg.get(self.index).get_int() == v {
						return Some(reg)
					} else {
						// not equal, skip to next entry
						cur = self.input.next();
						continue
					}
				}
			}
		}
		// no more entries, signal iterator exhaustion
		None
	}
}

impl<T: Operatorish<Vec<Register>>> Operatorish<Vec<Register>> for Select<T> {
}

struct HashJoin<T> {
	left: T,
	right: T,
	on: (uint, uint),
}

impl<T: Operatorish<Vec<Register>>> HashJoin<T> {
	fn new(left: T, right: T, on: (uint, uint)) -> HashJoin<T> {
		HashJoin {
			left: left,
			right: right,
			on: on,
		}
	}
}

#[test]
fn simple_tablescan() {
	let dir = match TempDir::new("tablescan") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};

	//let p = dir.path();
	let p = Path::new(".");

	let mut manager = buffer::BufferManager::new(1024, p.clone());
	let mut seg = schema::SPSegment {id: 1, manager: &mut manager};

	let name = schema::Column::new(~"name", schema::Varchar(128), vec!(schema::NotNull));
	let age = schema::Column::new(~"age", schema::Integer, vec!(schema::NotNull));
	let mut relation = schema::Relation::new(~"Person");
	relation.add_column(name);
	relation.add_column(age);
	relation.insert(&mut seg, vec!(schema::Record::from_str(~"Alice"), schema::Record::from_int(20)));
	relation.insert(&mut seg, vec!(schema::Record::from_str(~"Bob"), schema::Record::from_int(40)));

	{
		let mut ts = TableScan::new(relation.clone(), &mut seg);
		for tuple in ts {
			println!("Got entry {}", tuple);
		}
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut pr = Print::new(ts, &mut mw);
			// force write by iterating, strange API
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut pr = Project::new(ts, vec!(0));
			let mut pr = Print::new(pr, &mut mw);
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut se = Select::new(ts, 1, Integer(20));
			let mut pr = Print::new(se, &mut mw);
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	{
		let mut mw = MemWriter::new();
		{
			let mut ts = TableScan::new(relation.clone(), &mut seg);
			let mut se = Select::new(ts, 0, Varchar(~"Bob"));
			let mut pr = Print::new(se, &mut mw);
			for _ in pr {}
		}
		println!("Saved: {}", from_utf8(mw.unwrap()).unwrap());
	}

	assert!(false);
}
<|endoftext|>


<commit_before>use material::textures::ImageTexture;
use vec3::Vec3;

#[allow(dead_code)]
pub struct CubeMap {
    pub faces: Vec<ImageTexture>
}

impl CubeMap {
    #[allow(dead_code)]
    pub fn load(x: &str, x_neg: &str, y: &str, y_neg: &str, z: &str, z_neg: &str) -> CubeMap {
        CubeMap {
            faces: vec![
                ImageTexture::load(x),
                ImageTexture::load(x_neg),
                ImageTexture::load(y),
                ImageTexture::load(y_neg),
                ImageTexture::load(z),
                ImageTexture::load(z_neg)
            ]
        }
    }

    #[allow(dead_code)]
    pub fn color(&self, dir: Vec3) -> Vec3 {
        // Select cube face
        let x_mag = dir.x.abs();
        let y_mag = dir.y.abs();
        let z_mag = dir.z.abs();

        let mut face = -1;
        let mut s = 0.0;
        let mut t = 0.0;

        if x_mag >= y_mag && x_mag > z_mag {
            // +x -x direction
            face = if dir.x < 0.0 { 0 } else { 1 };
            s = dir.y / dir.x;
            t = dir.z / dir.x;
        } else if y_mag >= x_mag && y_mag > z_mag {
            // +y -y direction
            face = if dir.y < 0.0 { 2 } else { 3 };
            s = dir.x / dir.y;
            t = dir.z / dir.y;
        } else if z_mag >= y_mag && z_mag > x_mag {
            // +z -z direction
            face = if dir.z < 0.0 { 4 } else { 5 };
            s = dir.x / dir.z;
            t = dir.y / dir.z;
        }

        // [-1..1] -> [0..1]
        s = s * 0.5 + 0.5;
        t = t * 0.5 + 0.5;

        self.faces[face].sample(s, t)
    }
}
<commit_msg>Avoid seams in CubeMap skyboxes.<commit_after>use material::textures::ImageTexture;
use vec3::Vec3;

#[allow(dead_code)]
pub struct CubeMap {
    pub faces: Vec<ImageTexture>
}

impl CubeMap {
    #[allow(dead_code)]
    pub fn load(x: &str, x_neg: &str, y: &str, y_neg: &str, z: &str, z_neg: &str) -> CubeMap {
        CubeMap {
            faces: vec![
                ImageTexture::load(x),
                ImageTexture::load(x_neg),
                ImageTexture::load(y),
                ImageTexture::load(y_neg),
                ImageTexture::load(z),
                ImageTexture::load(z_neg)
            ]
        }
    }

    #[allow(dead_code)]
    pub fn color(&self, dir: Vec3) -> Vec3 {
        // Select cube face
        let x_mag = dir.x.abs();
        let y_mag = dir.y.abs();
        let z_mag = dir.z.abs();

        let mut face = -1;
        let mut s = 0.0;
        let mut t = 0.0;

        if x_mag >= y_mag && x_mag > z_mag {
            // +x -x direction
            face = if dir.x < 0.0 { 0 } else { 1 };
            s = dir.y / dir.x;
            t = dir.z / dir.x;
        } else if y_mag >= x_mag && y_mag > z_mag {
            // +y -y direction
            face = if dir.y < 0.0 { 2 } else { 3 };
            s = dir.x / dir.y;
            t = dir.z / dir.y;
        } else if z_mag >= y_mag && z_mag > x_mag {
            // +z -z direction
            face = if dir.z < 0.0 { 4 } else { 5 };
            s = dir.x / dir.z;
            t = dir.y / dir.z;
        }

        // [-1..1] -> [0..1]
        let seam_delta = 0.0001;
        s = (s * 0.5 + 0.5).max(seam_delta).min(1.0 - seam_delta);
        t = (t * 0.5 + 0.5).max(seam_delta).min(1.0 - seam_delta);

        self.faces[face].sample(s, t)
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

//! An implementation of the RESP protocol

use std::io;
use std::str;

use bytes::{BufMut, BytesMut};

use tokio_io::codec::{Decoder, Encoder};

use super::error::Error;

/// A single RESP value, this owns the data that is read/to-be written to Redis.
///
/// It is cloneable to allow multiple copies to be delivered in certain circumstances, e.g. multiple
/// subscribers to the same topic. TODO - revisit this assumption to make sure its sound, perhaps a
/// single-subscriber enforcement would make more sense, or sharing via `Arc`
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum RespValue {
    /// Zero, one or more other `RespValue`s.
    Array(Vec<RespValue>),

    /// A bulk string.  In Redis terminology a string is a byte-array, so this is stored as a
    /// vector of `u8`s to allow clients to interpret the bytes as appropriate.
    BulkString(Vec<u8>),

    /// An error from the Redis server
    Error(String),

    Integer(usize),

    SimpleString(String),
}

pub trait FromResp: Sized {
    fn from_resp(resp: RespValue) -> Result<Self, Error>;
}

impl FromResp for RespValue {
    fn from_resp(resp: RespValue) -> Result<RespValue, Error> {
        match resp {
            RespValue::Error(string) => Err(Error::Remote(string)),
            x => Ok(x),
        }
    }
}

impl FromResp for String {
    fn from_resp(resp: RespValue) -> Result<String, Error> {
        match resp {
            RespValue::Array(_) => Err(Error::RESP("Cannot be converted into a string".into())),
            RespValue::BulkString(ref bytes) => Ok(String::from_utf8_lossy(bytes).into_owned()),
            RespValue::Error(string) => Err(Error::Remote(string)),
            RespValue::Integer(i) => Ok(i.to_string()),
            RespValue::SimpleString(string) => Ok(string),
        }
    }
}

/// A trait to be implemented on types that can be automatically converted into `RespValue`s.
///
/// A `From<T>` where `T: ToResp` has been implemented so that everything that implements `ToResp`
/// can be converted with conversion traits.
pub trait ToResp {
    fn to_resp(&self) -> RespValue;
}

impl<'a> ToResp for &'a str {
    fn to_resp(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<'a, T> ToResp for &'a [T]
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.as_ref().iter().map(|x| x.to_resp()).collect())
    }
}

impl<'a> ToResp for String {
    fn to_resp(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<'a, T> ToResp for Vec<T>
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.into_iter().map(|v| v.to_resp()).collect())
    }
}

impl<T: ToResp> From<T> for RespValue {
    fn from(from: T) -> RespValue {
        from.to_resp()
    }
}

/// Codec to read frames
pub struct RespCodec;

fn write_rn(buf: &mut BytesMut) {
    buf.put_u8(b'\r');
    buf.put_u8(b'\n');
}

fn check_and_reserve(buf: &mut BytesMut, amt: usize) {
    let remaining_bytes = buf.remaining_mut();
    if remaining_bytes < amt {
        buf.reserve(amt);
    }
}

fn write_header(symb: u8, len: usize, buf: &mut BytesMut) {
    let len_as_string = len.to_string();
    let len_as_bytes = len_as_string.as_bytes();
    let header_bytes = 1 + len_as_bytes.len() + 2;
    check_and_reserve(buf, header_bytes);
    buf.put_u8(symb);
    buf.extend(len_as_bytes);
    write_rn(buf);
}

fn write_simple_string(symb: u8, string: &str, buf: &mut BytesMut) {
    let bytes = string.as_bytes();
    let size = 1 + bytes.len() + 2;
    check_and_reserve(buf, size);
    buf.put_u8(symb);
    buf.extend(bytes);
    write_rn(buf);
}

impl Encoder for RespCodec {
    type Item = RespValue;
    type Error = io::Error;

    fn encode(&mut self, msg: RespValue, buf: &mut BytesMut) -> Result<(), Self::Error> {
        match msg {
            RespValue::Array(ary) => {
                write_header(b'*', ary.len(), buf);
                for v in ary {
                    self.encode(v, buf)?;
                }
            }
            RespValue::BulkString(bstr) => {
                let len = bstr.len();
                write_header(b'$', len, buf);
                check_and_reserve(buf, len + 2);
                buf.extend(bstr);
                write_rn(buf);
            }
            RespValue::Error(ref string) => {
                write_simple_string(b'-', string, buf);
            }
            RespValue::Integer(val) => {
                // Simple integer are just the header
                write_header(b':', val, buf);
            }
            RespValue::SimpleString(ref string) => {
                write_simple_string(b'+', string, buf);
            }
        }
        Ok(())
    }
}

#[inline]
fn parse_error(message: String) -> Error {
    Error::RESP(message)
}

/// Many RESP types have their length (which is either bytes or "number of elements", depending on context)
/// encoded as a string, terminated by "\r\n", this looks for them.
///
/// Only return the string if the whole sequence is complete, including the terminator bytes (but those final
/// two bytes will not be returned)
///
/// TODO - rename this function potentially, it's used for simple integers too
fn scan_integer<'a>(buf: &'a mut BytesMut, idx: usize) -> Result<Option<(usize, &'a [u8])>, Error> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return Ok(None);
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => return Ok(Some((pos + 1, &buf[idx..pos - 1]))),
            (false, b'\r') => at_end = true,
            (false, b'0'...b'9') => (),
            (_, val) => return Err(parse_error(format!("Unexpected byte in size_string: {}", val))),
        }
        pos += 1;
    }
}

fn scan_string(buf: &mut BytesMut, idx: usize) -> Option<(usize, String)> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return None;
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => {
                let value = String::from_utf8_lossy(&buf[idx..pos - 1]).into_owned();
                return Some((pos + 1, value));
            }
            (true, _) => at_end = false,
            (false, b'\r') => at_end = true,
            (false, _) => (),
        }
        pos += 1;
    }
}

fn decode_raw_integer(buf: &mut BytesMut, idx: usize) -> Result<Option<(usize, usize)>, Error> {
    match scan_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int_str))) => {
            let int: usize = str::from_utf8(int_str)
                .expect("Not a string")
                .parse()
                .expect("Not an integer");
            Ok(Some((pos, int)))
        }
        Err(e) => Err(e),
    }
}

type DecodeResult = Result<Option<(usize, RespValue)>, Error>;

fn decode_bulk_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let remaining = buf.len() - pos;
            let required_bytes = size + 2;

            if remaining < required_bytes {
                return Ok(None);
            }

            let bulk_string = RespValue::BulkString(buf[pos..(pos + size)].to_vec());
            Ok(Some((pos + required_bytes, bulk_string)))
        }
        Err(e) => Err(e),
    }
}

fn decode_array(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let mut pos = pos;
            let mut values = Vec::with_capacity(size);
            for _ in 0..size {
                match decode(buf, pos) {
                    Ok(None) => return Ok(None),
                    Ok(Some((new_pos, value))) => {
                        values.push(value);
                        pos = new_pos;
                    }
                    Err(e) => return Err(e),
                }
            }
            Ok(Some((pos, RespValue::Array(values))))
        }
        Err(e) => Err(e),
    }
}

fn decode_integer(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int))) => Ok(Some((pos, RespValue::Integer(int)))),
        Err(e) => Err(e),
    }
}

/// A simple string is any series of bytes that ends with `\r\n`
fn decode_simple_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::SimpleString(string)))),
    }
}

fn decode_error(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::Error(string)))),
    }
}

fn decode(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    let length = buf.len();
    if length <= idx {
        return Ok(None);
    }

    let first_byte = buf[idx];
    match first_byte {
        b'$' => decode_bulk_string(buf, idx + 1),
        b'*' => decode_array(buf, idx + 1),
        b':' => decode_integer(buf, idx + 1),
        b'+' => decode_simple_string(buf, idx + 1),
        b'-' => decode_error(buf, idx + 1),
        _ => Err(parse_error(format!("Unexpected byte: {}", first_byte))),
    }
}

impl Decoder for RespCodec {
    type Item = RespValue;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        match decode(buf, 0) {
            Ok(None) => Ok(None),
            Ok(Some((pos, item))) => {
                buf.split_to(pos);
                Ok(Some(item))
            }
            Err(e) => Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use bytes::BytesMut;

    use tokio_io::codec::{Decoder, Encoder};

    use super::{RespCodec, RespValue};

    #[test]
    fn test_bulk_string() {
        let resp_object = RespValue::BulkString("THISISATEST".as_bytes().to_vec());
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"$11\r\nTHISISATEST\r\n".to_vec(), bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }

    #[test]
    fn test_array() {
        let resp_object = RespValue::Array(vec!["TEST1".into(), "TEST2".into()]);
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"*2\r\n$5\r\nTEST1\r\n$5\r\nTEST2\r\n".to_vec(),
                   bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }
}<commit_msg>Documentation<commit_after>/*
 * Copyright 2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

//! An implementation of the RESP protocol

use std::io;
use std::str;

use bytes::{BufMut, BytesMut};

use tokio_io::codec::{Decoder, Encoder};

use super::error::Error;

/// A single RESP value, this owns the data that is read/to-be written to Redis.
///
/// It is cloneable to allow multiple copies to be delivered in certain circumstances, e.g. multiple
/// subscribers to the same topic. TODO - revisit this assumption to make sure its sound, perhaps a
/// single-subscriber enforcement would make more sense, or sharing via `Arc`
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum RespValue {
    /// Zero, one or more other `RespValue`s.
    Array(Vec<RespValue>),

    /// A bulk string.  In Redis terminology a string is a byte-array, so this is stored as a
    /// vector of `u8`s to allow clients to interpret the bytes as appropriate.
    BulkString(Vec<u8>),

    /// An error from the Redis server
    Error(String),

    Integer(usize),

    SimpleString(String),
}

/// A trait to be implemented for every time which can be read from a RESP value.
///
/// Implementing this trait on a type means that type becomes a valid return type for calls such as `send` on
/// `client::PairedConnection`
pub trait FromResp: Sized {
    /// Return a `Result` containing either `Self` or `Error`.  Errors can occur due to either: a) the particular
    /// `RespValue` being incompatible with the required type, or b) a remote Redis error occuring.
    fn from_resp(resp: RespValue) -> Result<Self, Error>;
}

impl FromResp for RespValue {
    fn from_resp(resp: RespValue) -> Result<RespValue, Error> {
        match resp {
            RespValue::Error(string) => Err(Error::Remote(string)),
            x => Ok(x),
        }
    }
}

impl FromResp for String {
    fn from_resp(resp: RespValue) -> Result<String, Error> {
        match resp {
            RespValue::Array(_) => Err(Error::RESP("Cannot be converted into a string".into())),
            RespValue::BulkString(ref bytes) => Ok(String::from_utf8_lossy(bytes).into_owned()),
            RespValue::Error(string) => Err(Error::Remote(string)),
            RespValue::Integer(i) => Ok(i.to_string()),
            RespValue::SimpleString(string) => Ok(string),
        }
    }
}

/// A trait to be implemented on types that can be automatically converted into `RespValue`s.
///
/// A `From<T>` where `T: ToResp` has been implemented so that everything that implements `ToResp`
/// can be converted with conversion traits.
pub trait ToResp {
    fn to_resp(&self) -> RespValue;
}

impl<'a> ToResp for &'a str {
    fn to_resp(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<'a, T> ToResp for &'a [T]
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.as_ref().iter().map(|x| x.to_resp()).collect())
    }
}

impl<'a> ToResp for String {
    fn to_resp(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<'a, T> ToResp for Vec<T>
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.into_iter().map(|v| v.to_resp()).collect())
    }
}

impl<T: ToResp> From<T> for RespValue {
    fn from(from: T) -> RespValue {
        from.to_resp()
    }
}

/// Codec to read frames
pub struct RespCodec;

fn write_rn(buf: &mut BytesMut) {
    buf.put_u8(b'\r');
    buf.put_u8(b'\n');
}

fn check_and_reserve(buf: &mut BytesMut, amt: usize) {
    let remaining_bytes = buf.remaining_mut();
    if remaining_bytes < amt {
        buf.reserve(amt);
    }
}

fn write_header(symb: u8, len: usize, buf: &mut BytesMut) {
    let len_as_string = len.to_string();
    let len_as_bytes = len_as_string.as_bytes();
    let header_bytes = 1 + len_as_bytes.len() + 2;
    check_and_reserve(buf, header_bytes);
    buf.put_u8(symb);
    buf.extend(len_as_bytes);
    write_rn(buf);
}

fn write_simple_string(symb: u8, string: &str, buf: &mut BytesMut) {
    let bytes = string.as_bytes();
    let size = 1 + bytes.len() + 2;
    check_and_reserve(buf, size);
    buf.put_u8(symb);
    buf.extend(bytes);
    write_rn(buf);
}

impl Encoder for RespCodec {
    type Item = RespValue;
    type Error = io::Error;

    fn encode(&mut self, msg: RespValue, buf: &mut BytesMut) -> Result<(), Self::Error> {
        match msg {
            RespValue::Array(ary) => {
                write_header(b'*', ary.len(), buf);
                for v in ary {
                    self.encode(v, buf)?;
                }
            }
            RespValue::BulkString(bstr) => {
                let len = bstr.len();
                write_header(b'$', len, buf);
                check_and_reserve(buf, len + 2);
                buf.extend(bstr);
                write_rn(buf);
            }
            RespValue::Error(ref string) => {
                write_simple_string(b'-', string, buf);
            }
            RespValue::Integer(val) => {
                // Simple integer are just the header
                write_header(b':', val, buf);
            }
            RespValue::SimpleString(ref string) => {
                write_simple_string(b'+', string, buf);
            }
        }
        Ok(())
    }
}

#[inline]
fn parse_error(message: String) -> Error {
    Error::RESP(message)
}

/// Many RESP types have their length (which is either bytes or "number of elements", depending on context)
/// encoded as a string, terminated by "\r\n", this looks for them.
///
/// Only return the string if the whole sequence is complete, including the terminator bytes (but those final
/// two bytes will not be returned)
///
/// TODO - rename this function potentially, it's used for simple integers too
fn scan_integer<'a>(buf: &'a mut BytesMut, idx: usize) -> Result<Option<(usize, &'a [u8])>, Error> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return Ok(None);
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => return Ok(Some((pos + 1, &buf[idx..pos - 1]))),
            (false, b'\r') => at_end = true,
            (false, b'0'...b'9') => (),
            (_, val) => return Err(parse_error(format!("Unexpected byte in size_string: {}", val))),
        }
        pos += 1;
    }
}

fn scan_string(buf: &mut BytesMut, idx: usize) -> Option<(usize, String)> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return None;
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => {
                let value = String::from_utf8_lossy(&buf[idx..pos - 1]).into_owned();
                return Some((pos + 1, value));
            }
            (true, _) => at_end = false,
            (false, b'\r') => at_end = true,
            (false, _) => (),
        }
        pos += 1;
    }
}

fn decode_raw_integer(buf: &mut BytesMut, idx: usize) -> Result<Option<(usize, usize)>, Error> {
    match scan_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int_str))) => {
            let int: usize = str::from_utf8(int_str)
                .expect("Not a string")
                .parse()
                .expect("Not an integer");
            Ok(Some((pos, int)))
        }
        Err(e) => Err(e),
    }
}

type DecodeResult = Result<Option<(usize, RespValue)>, Error>;

fn decode_bulk_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let remaining = buf.len() - pos;
            let required_bytes = size + 2;

            if remaining < required_bytes {
                return Ok(None);
            }

            let bulk_string = RespValue::BulkString(buf[pos..(pos + size)].to_vec());
            Ok(Some((pos + required_bytes, bulk_string)))
        }
        Err(e) => Err(e),
    }
}

fn decode_array(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let mut pos = pos;
            let mut values = Vec::with_capacity(size);
            for _ in 0..size {
                match decode(buf, pos) {
                    Ok(None) => return Ok(None),
                    Ok(Some((new_pos, value))) => {
                        values.push(value);
                        pos = new_pos;
                    }
                    Err(e) => return Err(e),
                }
            }
            Ok(Some((pos, RespValue::Array(values))))
        }
        Err(e) => Err(e),
    }
}

fn decode_integer(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int))) => Ok(Some((pos, RespValue::Integer(int)))),
        Err(e) => Err(e),
    }
}

/// A simple string is any series of bytes that ends with `\r\n`
fn decode_simple_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::SimpleString(string)))),
    }
}

fn decode_error(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::Error(string)))),
    }
}

fn decode(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    let length = buf.len();
    if length <= idx {
        return Ok(None);
    }

    let first_byte = buf[idx];
    match first_byte {
        b'$' => decode_bulk_string(buf, idx + 1),
        b'*' => decode_array(buf, idx + 1),
        b':' => decode_integer(buf, idx + 1),
        b'+' => decode_simple_string(buf, idx + 1),
        b'-' => decode_error(buf, idx + 1),
        _ => Err(parse_error(format!("Unexpected byte: {}", first_byte))),
    }
}

impl Decoder for RespCodec {
    type Item = RespValue;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        match decode(buf, 0) {
            Ok(None) => Ok(None),
            Ok(Some((pos, item))) => {
                buf.split_to(pos);
                Ok(Some(item))
            }
            Err(e) => Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use bytes::BytesMut;

    use tokio_io::codec::{Decoder, Encoder};

    use super::{RespCodec, RespValue};

    #[test]
    fn test_bulk_string() {
        let resp_object = RespValue::BulkString("THISISATEST".as_bytes().to_vec());
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"$11\r\nTHISISATEST\r\n".to_vec(), bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }

    #[test]
    fn test_array() {
        let resp_object = RespValue::Array(vec!["TEST1".into(), "TEST2".into()]);
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"*2\r\n$5\r\nTEST1\r\n$5\r\nTEST2\r\n".to_vec(),
                   bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }
}<|endoftext|>