<commit_before>use std::sync::Arc;
use std::io::{self, Read, Write};

use futures::{Future, Wake, Tokens};
use futures::stream::{Stream, StreamResult, Fuse};
use futuremio::*;

pub trait Parse: Sized + Send + 'static {
    type Parser: Default + Send + 'static;
    type Error: Send + 'static + From<io::Error>;

    fn parse(parser: &mut Self::Parser, buf: &[u8]) -> Option<Result<(Self, usize), Self::Error>>;
}

/// A stream for parsing from an underlying reader, using an unbounded internal
/// buffer.
pub struct ParseStream<R, P: Parse> {
    source: R,
    source_ready: ReadinessStream,
    parser: P::Parser,
    buf: Vec<u8>,

    // how far into the buffer have we parsed? note: we drain lazily
    pos: usize,

    // is there new data that we need to try to parse?
    need_parse: bool,

    // has `source` yielded an EOF?
    eof: bool,
}

impl<R, P> ParseStream<R, P>
    where R: Read + Send + 'static,
          P: Parse
{
    pub fn new(source: R, source_ready: ReadinessStream) -> ParseStream<R, P> {
        ParseStream {
            source: source,
            source_ready: source_ready,
            parser: Default::default(),
            buf: Vec::with_capacity(2048),
            pos: 0,
            need_parse: false,
            eof: false,
        }
    }
}

// TODO: move this into method
fn read<R: Read>(socket: &mut R, input: &mut Vec<u8>) -> io::Result<(usize, bool)> {
    loop {
        match socket.read(unsafe { slice_to_end(input) }) {
            Ok(0) => {
                trace!("socket EOF");
                return Ok((0, true))
            }
            Ok(n) => {
                trace!("socket read {} bytes", n);
                let len = input.len();
                unsafe {
                    input.set_len(len + n);
                }
                return Ok((n, false));
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok((0, false)),
            Err(e) => return Err(e),
        }
    }

    unsafe fn slice_to_end(v: &mut Vec<u8>) -> &mut [u8] {
        use std::slice;
        if v.capacity() == 0 {
            v.reserve(16);
        }
        if v.capacity() == v.len() {
            v.reserve(1);
        }
        slice::from_raw_parts_mut(v.as_mut_ptr().offset(v.len() as isize),
                                  v.capacity() - v.len())
    }
}

impl<R, P: Parse> ParseStream<R, P> {
    fn drain(&mut self) {
        self.buf.drain(..self.pos);
        self.pos = 0;
    }
}

impl<R, P> Stream for ParseStream<R, P>
    where R: Read + Send + 'static,
          P: Parse
{
    type Item = P;
    type Error = P::Error;

    fn poll(&mut self, tokens: &Tokens) -> Option<StreamResult<P, P::Error>> {
        if !self.eof {
            match self.source_ready.poll(tokens) {
                // TODO: consider refactoring `poll` API to make this more
                //       readable...
                None => {}
                Some(Err(e)) => return Some(Err(e.into())),
                Some(Ok(Some(()))) => {
                    // drain any stale contents, to make as much space as we can
                    // before reading
                    self.drain();

                    match read(&mut self.source, &mut self.buf) {
                        Ok((n, eof)) => {
                            self.eof = eof;
                            if n > 0 {
                                self.need_parse = true
                            }
                        }
                        Err(e) => return Some(Err(e.into())),
                    }
                }
                _ => unreachable!(),
            }
        }

        if self.need_parse {
            debug!("attempting to parse");
            if let Some(res) = P::parse(&mut self.parser, &self.buf) {
                return Some(res.map(|(i, n)| {
                    self.pos += n;
                    Some(i)
                }))
            }

            // don't try to parse again until we've seen more bytes
            self.need_parse = false;
        }

        if self.eof {
            Some(Ok(None))
        } else {
            None
        }
    }

    fn schedule(&mut self, wake: Arc<Wake>) {
        // TODO: think through this carefully...
        if self.need_parse {
            // Empty tokens because in a `need_parse` situation, we'll attempt
            // to parse regardless of tokens
            wake.wake(&Tokens::empty())
        } else {
            self.source_ready.schedule(wake)
        }
    }
}

// TODO: make this a method
fn write<W: Write>(sink: &mut W, buf: &mut Vec<u8>) -> io::Result<()> {
    loop {
        match sink.write(&buf) {
            Ok(0) => {
                // TODO: copied from mio example, clean up
                return Err(io::Error::new(io::ErrorKind::Other, "early eof2"));
            }
            Ok(n) => {
                // TODO: consider draining more lazily, i.e. only just before
                //       returning
                buf.drain(..n);
                if buf.len() == 0 {
                    return Ok(());
                }
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(()),
            Err(e) => return Err(e),
        }
    }
}

pub trait Serialize: Send + 'static {
    fn serialize(&self, buf: &mut Vec<u8>);
}

/// Serialize a stream of items into a writer, using an unbounded internal
/// buffer.
///
/// Represented as a future which yields () on successfully writing the entire
/// stream (which requires the stream to terminate), or an error if there is any
/// error along the way.
pub struct StreamWriter<W, S> {
    sink: W,
    sink_ready: ReadinessStream,
    items: Fuse<S>,
    buf: Vec<u8>,
}

impl<W, S> StreamWriter<W, S>
    where W: Write + Send + 'static,
          S: Stream,
          S::Item: Serialize,
          S::Error: From<io::Error>
{
    pub fn new(sink: W, sink_ready: ReadinessStream, items: S) -> StreamWriter<W, S> {
        StreamWriter {
            sink: sink,
            sink_ready: sink_ready,
            items: items.fuse(),
            buf: Vec::with_capacity(2048),
        }
    }
}

impl<W, S> Future for StreamWriter<W, S>
    where W: Write + Send + 'static,
          S: Stream,
          S::Item: Serialize,
          S::Error: From<io::Error>
{
    type Item = ();
    type Error = S::Error;

    fn poll(&mut self, tokens: &Tokens) -> Option<Result<(), S::Error>> {
        // make sure to pass down `tokens` only on the *first* poll for items
        let mut tokens_for_items = Some(tokens);
        loop {
            match self.items.poll(tokens_for_items.take().unwrap_or(&Tokens::all())) {
                Some(Err(e)) => return Some(Err(e)),
                Some(Ok(Some(item))) => {
                    debug!("got an item to serialize!");
                    item.serialize(&mut self.buf)
                }
                Some(Ok(None)) |
                None => break,
            }
        }

        // TODO: optimization for case where we just transitioned from no bytes
        // to write to having bytes to write; in that case, we should try to
        // write regardless of sink_ready.poll, because we haven't asked for a
        // readiness notifcation. Saves a trip around the event loop.

        if self.buf.len() > 0 {
            match self.sink_ready.poll(tokens) {
                Some(Err(e)) => Some(Err(e.into())),
                Some(Ok(Some(()))) => {
                    debug!("trying to write some data");
                    if let Err(e) = write(&mut self.sink, &mut self.buf) {
                        Some(Err(e.into()))
                    } else {
                        None
                    }
                }
                Some(Ok(None)) | // TODO: this should translate to an error
                None => None,
            }
        } else if self.items.is_done() {
            // Nothing more to write to sink, and no more incoming items; we're done!
            Some(Ok(()))
        } else {
            None
        }
    }

    fn schedule(&mut self, wake: Arc<Wake>) {
        // wake up on writability only if we have something to write
        if self.buf.len() > 0 {
            self.sink_ready.schedule(wake.clone());
        }

        // for now, we are always happy to write more items into our unbounded buffer
        self.items.schedule(wake);
    }
}
<commit_msg>Be sure to advance the parser<commit_after>use std::sync::Arc;
use std::io::{self, Read, Write};

use futures::{Future, Wake, Tokens};
use futures::stream::{Stream, StreamResult, Fuse};
use futuremio::*;

pub trait Parse: Sized + Send + 'static {
    type Parser: Default + Send + 'static;
    type Error: Send + 'static + From<io::Error>;

    fn parse(parser: &mut Self::Parser, buf: &[u8]) -> Option<Result<(Self, usize), Self::Error>>;
}

/// A stream for parsing from an underlying reader, using an unbounded internal
/// buffer.
pub struct ParseStream<R, P: Parse> {
    source: R,
    source_ready: ReadinessStream,
    parser: P::Parser,
    buf: Vec<u8>,

    // how far into the buffer have we parsed? note: we drain lazily
    pos: usize,

    // is there new data that we need to try to parse?
    need_parse: bool,

    // has `source` yielded an EOF?
    eof: bool,
}

impl<R, P> ParseStream<R, P>
    where R: Read + Send + 'static,
          P: Parse
{
    pub fn new(source: R, source_ready: ReadinessStream) -> ParseStream<R, P> {
        ParseStream {
            source: source,
            source_ready: source_ready,
            parser: Default::default(),
            buf: Vec::with_capacity(2048),
            pos: 0,
            need_parse: false,
            eof: false,
        }
    }
}

// TODO: move this into method
fn read<R: Read>(socket: &mut R, input: &mut Vec<u8>) -> io::Result<(usize, bool)> {
    loop {
        match socket.read(unsafe { slice_to_end(input) }) {
            Ok(0) => {
                trace!("socket EOF");
                return Ok((0, true))
            }
            Ok(n) => {
                trace!("socket read {} bytes", n);
                let len = input.len();
                unsafe {
                    input.set_len(len + n);
                }
                return Ok((n, false));
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok((0, false)),
            Err(e) => return Err(e),
        }
    }

    unsafe fn slice_to_end(v: &mut Vec<u8>) -> &mut [u8] {
        use std::slice;
        if v.capacity() == 0 {
            v.reserve(16);
        }
        if v.capacity() == v.len() {
            v.reserve(1);
        }
        slice::from_raw_parts_mut(v.as_mut_ptr().offset(v.len() as isize),
                                  v.capacity() - v.len())
    }
}

impl<R, P: Parse> ParseStream<R, P> {
    fn drain(&mut self) {
        self.buf.drain(..self.pos);
        self.pos = 0;
    }
}

impl<R, P> Stream for ParseStream<R, P>
    where R: Read + Send + 'static,
          P: Parse
{
    type Item = P;
    type Error = P::Error;

    fn poll(&mut self, tokens: &Tokens) -> Option<StreamResult<P, P::Error>> {
        if !self.eof {
            match self.source_ready.poll(tokens) {
                // TODO: consider refactoring `poll` API to make this more
                //       readable...
                None => {}
                Some(Err(e)) => return Some(Err(e.into())),
                Some(Ok(Some(()))) => {
                    // drain any stale contents, to make as much space as we can
                    // before reading
                    self.drain();

                    match read(&mut self.source, &mut self.buf) {
                        Ok((n, eof)) => {
                            self.eof = eof;
                            if n > 0 {
                                self.need_parse = true
                            }
                        }
                        Err(e) => return Some(Err(e.into())),
                    }
                }
                _ => unreachable!(),
            }
        }

        if self.need_parse {
            debug!("attempting to parse");
            if let Some(res) = P::parse(&mut self.parser, &self.buf[self.pos..]) {
                return Some(res.map(|(i, n)| {
                    self.pos += n;
                    Some(i)
                }))
            }

            // don't try to parse again until we've seen more bytes
            self.need_parse = false;
        }

        if self.eof {
            Some(Ok(None))
        } else {
            None
        }
    }

    fn schedule(&mut self, wake: Arc<Wake>) {
        // TODO: think through this carefully...
        if self.need_parse {
            // Empty tokens because in a `need_parse` situation, we'll attempt
            // to parse regardless of tokens
            wake.wake(&Tokens::empty())
        } else {
            self.source_ready.schedule(wake)
        }
    }
}

// TODO: make this a method
fn write<W: Write>(sink: &mut W, buf: &mut Vec<u8>) -> io::Result<()> {
    loop {
        match sink.write(&buf) {
            Ok(0) => {
                // TODO: copied from mio example, clean up
                return Err(io::Error::new(io::ErrorKind::Other, "early eof2"));
            }
            Ok(n) => {
                // TODO: consider draining more lazily, i.e. only just before
                //       returning
                buf.drain(..n);
                if buf.len() == 0 {
                    return Ok(());
                }
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(()),
            Err(e) => return Err(e),
        }
    }
}

pub trait Serialize: Send + 'static {
    fn serialize(&self, buf: &mut Vec<u8>);
}

/// Serialize a stream of items into a writer, using an unbounded internal
/// buffer.
///
/// Represented as a future which yields () on successfully writing the entire
/// stream (which requires the stream to terminate), or an error if there is any
/// error along the way.
pub struct StreamWriter<W, S> {
    sink: W,
    sink_ready: ReadinessStream,
    items: Fuse<S>,
    buf: Vec<u8>,
}

impl<W, S> StreamWriter<W, S>
    where W: Write + Send + 'static,
          S: Stream,
          S::Item: Serialize,
          S::Error: From<io::Error>
{
    pub fn new(sink: W, sink_ready: ReadinessStream, items: S) -> StreamWriter<W, S> {
        StreamWriter {
            sink: sink,
            sink_ready: sink_ready,
            items: items.fuse(),
            buf: Vec::with_capacity(2048),
        }
    }
}

impl<W, S> Future for StreamWriter<W, S>
    where W: Write + Send + 'static,
          S: Stream,
          S::Item: Serialize,
          S::Error: From<io::Error>
{
    type Item = ();
    type Error = S::Error;

    fn poll(&mut self, tokens: &Tokens) -> Option<Result<(), S::Error>> {
        // make sure to pass down `tokens` only on the *first* poll for items
        let mut tokens_for_items = Some(tokens);
        loop {
            match self.items.poll(tokens_for_items.take().unwrap_or(&Tokens::all())) {
                Some(Err(e)) => return Some(Err(e)),
                Some(Ok(Some(item))) => {
                    debug!("got an item to serialize!");
                    item.serialize(&mut self.buf)
                }
                Some(Ok(None)) |
                None => break,
            }
        }

        // TODO: optimization for case where we just transitioned from no bytes
        // to write to having bytes to write; in that case, we should try to
        // write regardless of sink_ready.poll, because we haven't asked for a
        // readiness notifcation. Saves a trip around the event loop.

        if self.buf.len() > 0 {
            match self.sink_ready.poll(tokens) {
                Some(Err(e)) => Some(Err(e.into())),
                Some(Ok(Some(()))) => {
                    debug!("trying to write some data");
                    if let Err(e) = write(&mut self.sink, &mut self.buf) {
                        Some(Err(e.into()))
                    } else {
                        None
                    }
                }
                Some(Ok(None)) | // TODO: this should translate to an error
                None => None,
            }
        } else if self.items.is_done() {
            // Nothing more to write to sink, and no more incoming items; we're done!
            Some(Ok(()))
        } else {
            None
        }
    }

    fn schedule(&mut self, wake: Arc<Wake>) {
        // wake up on writability only if we have something to write
        if self.buf.len() > 0 {
            self.sink_ready.schedule(wake.clone());
        }

        // for now, we are always happy to write more items into our unbounded buffer
        self.items.schedule(wake);
    }
}
<|endoftext|>


<commit_before>use std::ops::{Not, BitAnd, BitOr, BitXor, Shl, Shr};
use std::num::Wrapping;

use {Num, NumCast};
use bounds::Bounded;
use ops::checked::*;
use ops::saturating::Saturating;

pub trait PrimInt
    : Sized
    + Copy
    + Num + NumCast
    + Bounded
    + PartialOrd + Ord + Eq
    + Not<Output=Self>
    + BitAnd<Output=Self>
    + BitOr<Output=Self>
    + BitXor<Output=Self>
    + Shl<usize, Output=Self>
    + Shr<usize, Output=Self>
    + CheckedAdd<Output=Self>
    + CheckedSub<Output=Self>
    + CheckedMul<Output=Self>
    + CheckedDiv<Output=Self>
    + Saturating
{
    /// Returns the number of ones in the binary representation of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b01001100u8;
    ///
    /// assert_eq!(n.count_ones(), 3);
    /// ```
    fn count_ones(self) -> u32;

    /// Returns the number of zeros in the binary representation of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b01001100u8;
    ///
    /// assert_eq!(n.count_zeros(), 5);
    /// ```
    fn count_zeros(self) -> u32;

    /// Returns the number of leading zeros in the binary representation
    /// of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b0101000u16;
    ///
    /// assert_eq!(n.leading_zeros(), 10);
    /// ```
    fn leading_zeros(self) -> u32;

    /// Returns the number of trailing zeros in the binary representation
    /// of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b0101000u16;
    ///
    /// assert_eq!(n.trailing_zeros(), 3);
    /// ```
    fn trailing_zeros(self) -> u32;

    /// Shifts the bits to the left by a specified amount amount, `n`, wrapping
    /// the truncated bits to the end of the resulting integer.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0x3456789ABCDEF012u64;
    ///
    /// assert_eq!(n.rotate_left(12), m);
    /// ```
    fn rotate_left(self, n: u32) -> Self;

    /// Shifts the bits to the right by a specified amount amount, `n`, wrapping
    /// the truncated bits to the beginning of the resulting integer.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0xDEF0123456789ABCu64;
    ///
    /// assert_eq!(n.rotate_right(12), m);
    /// ```
    fn rotate_right(self, n: u32) -> Self;

    /// Shifts the bits to the left by a specified amount amount, `n`, filling
    /// zeros in the least significant bits.
    ///
    /// This is bitwise equivalent to signed `Shl`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0x3456789ABCDEF000u64;
    ///
    /// assert_eq!(n.signed_shl(12), m);
    /// ```
    fn signed_shl(self, n: u32) -> Self;

    /// Shifts the bits to the right by a specified amount amount, `n`, copying
    /// the "sign bit" in the most significant bits even for unsigned types.
    ///
    /// This is bitwise equivalent to signed `Shr`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0xFEDCBA9876543210u64;
    /// let m = 0xFFFFEDCBA9876543u64;
    ///
    /// assert_eq!(n.signed_shr(12), m);
    /// ```
    fn signed_shr(self, n: u32) -> Self;

    /// Shifts the bits to the left by a specified amount amount, `n`, filling
    /// zeros in the least significant bits.
    ///
    /// This is bitwise equivalent to unsigned `Shl`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFi64;
    /// let m = 0x3456789ABCDEF000i64;
    ///
    /// assert_eq!(n.unsigned_shl(12), m);
    /// ```
    fn unsigned_shl(self, n: u32) -> Self;

    /// Shifts the bits to the right by a specified amount amount, `n`, filling
    /// zeros in the most significant bits.
    ///
    /// This is bitwise equivalent to unsigned `Shr`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0xFEDCBA9876543210i64;
    /// let m = 0x000FEDCBA9876543i64;
    ///
    /// assert_eq!(n.unsigned_shr(12), m);
    /// ```
    fn unsigned_shr(self, n: u32) -> Self;

    /// Reverses the byte order of the integer.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0xEFCDAB8967452301u64;
    ///
    /// assert_eq!(n.swap_bytes(), m);
    /// ```
    fn swap_bytes(self) -> Self;

    /// Convert an integer from big endian to the target's endianness.
    ///
    /// On big endian this is a no-op. On little endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "big") {
    ///     assert_eq!(u64::from_be(n), n)
    /// } else {
    ///     assert_eq!(u64::from_be(n), n.swap_bytes())
    /// }
    /// ```
    fn from_be(x: Self) -> Self;

    /// Convert an integer from little endian to the target's endianness.
    ///
    /// On little endian this is a no-op. On big endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "little") {
    ///     assert_eq!(u64::from_le(n), n)
    /// } else {
    ///     assert_eq!(u64::from_le(n), n.swap_bytes())
    /// }
    /// ```
    fn from_le(x: Self) -> Self;

    /// Convert `self` to big endian from the target's endianness.
    ///
    /// On big endian this is a no-op. On little endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "big") {
    ///     assert_eq!(n.to_be(), n)
    /// } else {
    ///     assert_eq!(n.to_be(), n.swap_bytes())
    /// }
    /// ```
    fn to_be(self) -> Self;

    /// Convert `self` to little endian from the target's endianness.
    ///
    /// On little endian this is a no-op. On big endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "little") {
    ///     assert_eq!(n.to_le(), n)
    /// } else {
    ///     assert_eq!(n.to_le(), n.swap_bytes())
    /// }
    /// ```
    fn to_le(self) -> Self;

    /// Raises self to the power of `exp`, using exponentiation by squaring.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// assert_eq!(2i32.pow(4), 16);
    /// ```
    fn pow(self, exp: u32) -> Self;
}

macro_rules! prim_int_impl {
    ($T:ty, $S:ty, $U:ty) => (
        impl PrimInt for $T {
            #[inline]
            fn count_ones(self) -> u32 {
                <$T>::count_ones(self)
            }

            #[inline]
            fn count_zeros(self) -> u32 {
                <$T>::count_zeros(self)
            }

            #[inline]
            fn leading_zeros(self) -> u32 {
                <$T>::leading_zeros(self)
            }

            #[inline]
            fn trailing_zeros(self) -> u32 {
                <$T>::trailing_zeros(self)
            }

            #[inline]
            fn rotate_left(self, n: u32) -> Self {
                <$T>::rotate_left(self, n)
            }

            #[inline]
            fn rotate_right(self, n: u32) -> Self {
                <$T>::rotate_right(self, n)
            }

            #[inline]
            fn signed_shl(self, n: u32) -> Self {
                ((self as $S) << n) as $T
            }

            #[inline]
            fn signed_shr(self, n: u32) -> Self {
                ((self as $S) >> n) as $T
            }

            #[inline]
            fn unsigned_shl(self, n: u32) -> Self {
                ((self as $U) << n) as $T
            }

            #[inline]
            fn unsigned_shr(self, n: u32) -> Self {
                ((self as $U) >> n) as $T
            }

            #[inline]
            fn swap_bytes(self) -> Self {
                <$T>::swap_bytes(self)
            }

            #[inline]
            fn from_be(x: Self) -> Self {
                <$T>::from_be(x)
            }

            #[inline]
            fn from_le(x: Self) -> Self {
                <$T>::from_le(x)
            }

            #[inline]
            fn to_be(self) -> Self {
                <$T>::to_be(self)
            }

            #[inline]
            fn to_le(self) -> Self {
                <$T>::to_le(self)
            }

            #[inline]
            fn pow(self, exp: u32) -> Self {
                <$T>::pow(self, exp)
            }
        }
    )
}

// prim_int_impl!(type, signed, unsigned);
prim_int_impl!(u8,    i8,    u8);
prim_int_impl!(u16,   i16,   u16);
prim_int_impl!(u32,   i32,   u32);
prim_int_impl!(u64,   i64,   u64);
prim_int_impl!(usize, isize, usize);
prim_int_impl!(i8,    i8,    u8);
prim_int_impl!(i16,   i16,   u16);
prim_int_impl!(i32,   i32,   u32);
prim_int_impl!(i64,   i64,   u64);
prim_int_impl!(isize, isize, usize);

// While this might violate the possible assumption that
// PrimInt represents genuine primitive integer types,
// Wrapping<T> is _practically_ one too, so this
// shouldn't be a concern.
impl<T: PrimInt> PrimInt for Wrapping<T>
    where Wrapping<T>:
          Num + NumCast
        + Bounded
        + Not<Output=Wrapping<T>>
        + BitAnd<Output=Wrapping<T>>
        + BitOr<Output=Wrapping<T>>
        + BitXor<Output=Wrapping<T>>
        + Shl<usize, Output=Wrapping<T>>
        + Shr<usize, Output=Wrapping<T>>
        + CheckedAdd<Output=Wrapping<T>>
        + CheckedSub<Output=Wrapping<T>>
        + CheckedMul<Output=Wrapping<T>>
        + CheckedDiv<Output=Wrapping<T>>
        + Saturating
{
    #[inline]
    fn count_ones(self) -> u32 {
        self.0.count_ones()
    }

    #[inline]
    fn count_zeros(self) -> u32 {
        self.0.count_zeros()
    }

    #[inline]
    fn leading_zeros(self) -> u32 {
        self.0.leading_zeros()
    }

    #[inline]
    fn trailing_zeros(self) -> u32 {
        self.0.trailing_zeros()
    }

    #[inline]
    fn rotate_left(self, n: u32) -> Self {
        Wrapping(self.0.rotate_left(n))
    }

    #[inline]
    fn rotate_right(self, n: u32) -> Self {
        Wrapping(self.0.rotate_right(n))
    }

    #[inline]
    fn signed_shl(self, n: u32) -> Self {
        Wrapping(self.0.signed_shl(n))
    }

    #[inline]
    fn signed_shr(self, n: u32) -> Self {
        Wrapping(self.0.signed_shr(n))
    }

    #[inline]
    fn unsigned_shl(self, n: u32) -> Self {
        Wrapping(self.0.unsigned_shl(n))
    }

    #[inline]
    fn unsigned_shr(self, n: u32) -> Self {
        Wrapping(self.0.unsigned_shr(n))
    }

    #[inline]
    fn swap_bytes(self) -> Self {
        Wrapping(self.0.swap_bytes())
    }

    #[inline]
    fn from_be(x: Self) -> Self {
        Wrapping(T::from_be(x.0))
    }

    #[inline]
    fn from_le(x: Self) -> Self {
        Wrapping(T::from_le(x.0))
    }

    #[inline]
    fn to_be(self) -> Self {
        Wrapping(self.0.to_be())
    }

    #[inline]
    fn to_le(self) -> Self {
        Wrapping(self.0.to_le())
    }

    // This is (or should be ?) fine because Wrapping<T> only guarantees that
    // _standard_ operations (Add, Mul, etc) have wrapping semantics.
    #[inline]
    fn pow(self, exp: u32) -> Self {
        Wrapping(self.0.pow(exp))
    }
}
<commit_msg>Revert "impl PrimInt for Wrapping<T>"<commit_after>use std::ops::{Not, BitAnd, BitOr, BitXor, Shl, Shr};

use {Num, NumCast};
use bounds::Bounded;
use ops::checked::*;
use ops::saturating::Saturating;

pub trait PrimInt
    : Sized
    + Copy
    + Num + NumCast
    + Bounded
    + PartialOrd + Ord + Eq
    + Not<Output=Self>
    + BitAnd<Output=Self>
    + BitOr<Output=Self>
    + BitXor<Output=Self>
    + Shl<usize, Output=Self>
    + Shr<usize, Output=Self>
    + CheckedAdd<Output=Self>
    + CheckedSub<Output=Self>
    + CheckedMul<Output=Self>
    + CheckedDiv<Output=Self>
    + Saturating
{
    /// Returns the number of ones in the binary representation of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b01001100u8;
    ///
    /// assert_eq!(n.count_ones(), 3);
    /// ```
    fn count_ones(self) -> u32;

    /// Returns the number of zeros in the binary representation of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b01001100u8;
    ///
    /// assert_eq!(n.count_zeros(), 5);
    /// ```
    fn count_zeros(self) -> u32;

    /// Returns the number of leading zeros in the binary representation
    /// of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b0101000u16;
    ///
    /// assert_eq!(n.leading_zeros(), 10);
    /// ```
    fn leading_zeros(self) -> u32;

    /// Returns the number of trailing zeros in the binary representation
    /// of `self`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0b0101000u16;
    ///
    /// assert_eq!(n.trailing_zeros(), 3);
    /// ```
    fn trailing_zeros(self) -> u32;

    /// Shifts the bits to the left by a specified amount amount, `n`, wrapping
    /// the truncated bits to the end of the resulting integer.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0x3456789ABCDEF012u64;
    ///
    /// assert_eq!(n.rotate_left(12), m);
    /// ```
    fn rotate_left(self, n: u32) -> Self;

    /// Shifts the bits to the right by a specified amount amount, `n`, wrapping
    /// the truncated bits to the beginning of the resulting integer.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0xDEF0123456789ABCu64;
    ///
    /// assert_eq!(n.rotate_right(12), m);
    /// ```
    fn rotate_right(self, n: u32) -> Self;

    /// Shifts the bits to the left by a specified amount amount, `n`, filling
    /// zeros in the least significant bits.
    ///
    /// This is bitwise equivalent to signed `Shl`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0x3456789ABCDEF000u64;
    ///
    /// assert_eq!(n.signed_shl(12), m);
    /// ```
    fn signed_shl(self, n: u32) -> Self;

    /// Shifts the bits to the right by a specified amount amount, `n`, copying
    /// the "sign bit" in the most significant bits even for unsigned types.
    ///
    /// This is bitwise equivalent to signed `Shr`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0xFEDCBA9876543210u64;
    /// let m = 0xFFFFEDCBA9876543u64;
    ///
    /// assert_eq!(n.signed_shr(12), m);
    /// ```
    fn signed_shr(self, n: u32) -> Self;

    /// Shifts the bits to the left by a specified amount amount, `n`, filling
    /// zeros in the least significant bits.
    ///
    /// This is bitwise equivalent to unsigned `Shl`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFi64;
    /// let m = 0x3456789ABCDEF000i64;
    ///
    /// assert_eq!(n.unsigned_shl(12), m);
    /// ```
    fn unsigned_shl(self, n: u32) -> Self;

    /// Shifts the bits to the right by a specified amount amount, `n`, filling
    /// zeros in the most significant bits.
    ///
    /// This is bitwise equivalent to unsigned `Shr`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0xFEDCBA9876543210i64;
    /// let m = 0x000FEDCBA9876543i64;
    ///
    /// assert_eq!(n.unsigned_shr(12), m);
    /// ```
    fn unsigned_shr(self, n: u32) -> Self;

    /// Reverses the byte order of the integer.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    /// let m = 0xEFCDAB8967452301u64;
    ///
    /// assert_eq!(n.swap_bytes(), m);
    /// ```
    fn swap_bytes(self) -> Self;

    /// Convert an integer from big endian to the target's endianness.
    ///
    /// On big endian this is a no-op. On little endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "big") {
    ///     assert_eq!(u64::from_be(n), n)
    /// } else {
    ///     assert_eq!(u64::from_be(n), n.swap_bytes())
    /// }
    /// ```
    fn from_be(x: Self) -> Self;

    /// Convert an integer from little endian to the target's endianness.
    ///
    /// On little endian this is a no-op. On big endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "little") {
    ///     assert_eq!(u64::from_le(n), n)
    /// } else {
    ///     assert_eq!(u64::from_le(n), n.swap_bytes())
    /// }
    /// ```
    fn from_le(x: Self) -> Self;

    /// Convert `self` to big endian from the target's endianness.
    ///
    /// On big endian this is a no-op. On little endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "big") {
    ///     assert_eq!(n.to_be(), n)
    /// } else {
    ///     assert_eq!(n.to_be(), n.swap_bytes())
    /// }
    /// ```
    fn to_be(self) -> Self;

    /// Convert `self` to little endian from the target's endianness.
    ///
    /// On little endian this is a no-op. On big endian the bytes are swapped.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// let n = 0x0123456789ABCDEFu64;
    ///
    /// if cfg!(target_endian = "little") {
    ///     assert_eq!(n.to_le(), n)
    /// } else {
    ///     assert_eq!(n.to_le(), n.swap_bytes())
    /// }
    /// ```
    fn to_le(self) -> Self;

    /// Raises self to the power of `exp`, using exponentiation by squaring.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::PrimInt;
    ///
    /// assert_eq!(2i32.pow(4), 16);
    /// ```
    fn pow(self, exp: u32) -> Self;
}

macro_rules! prim_int_impl {
    ($T:ty, $S:ty, $U:ty) => (
        impl PrimInt for $T {
            #[inline]
            fn count_ones(self) -> u32 {
                <$T>::count_ones(self)
            }

            #[inline]
            fn count_zeros(self) -> u32 {
                <$T>::count_zeros(self)
            }

            #[inline]
            fn leading_zeros(self) -> u32 {
                <$T>::leading_zeros(self)
            }

            #[inline]
            fn trailing_zeros(self) -> u32 {
                <$T>::trailing_zeros(self)
            }

            #[inline]
            fn rotate_left(self, n: u32) -> Self {
                <$T>::rotate_left(self, n)
            }

            #[inline]
            fn rotate_right(self, n: u32) -> Self {
                <$T>::rotate_right(self, n)
            }

            #[inline]
            fn signed_shl(self, n: u32) -> Self {
                ((self as $S) << n) as $T
            }

            #[inline]
            fn signed_shr(self, n: u32) -> Self {
                ((self as $S) >> n) as $T
            }

            #[inline]
            fn unsigned_shl(self, n: u32) -> Self {
                ((self as $U) << n) as $T
            }

            #[inline]
            fn unsigned_shr(self, n: u32) -> Self {
                ((self as $U) >> n) as $T
            }

            #[inline]
            fn swap_bytes(self) -> Self {
                <$T>::swap_bytes(self)
            }

            #[inline]
            fn from_be(x: Self) -> Self {
                <$T>::from_be(x)
            }

            #[inline]
            fn from_le(x: Self) -> Self {
                <$T>::from_le(x)
            }

            #[inline]
            fn to_be(self) -> Self {
                <$T>::to_be(self)
            }

            #[inline]
            fn to_le(self) -> Self {
                <$T>::to_le(self)
            }

            #[inline]
            fn pow(self, exp: u32) -> Self {
                <$T>::pow(self, exp)
            }
        }
    )
}

// prim_int_impl!(type, signed, unsigned);
prim_int_impl!(u8,    i8,    u8);
prim_int_impl!(u16,   i16,   u16);
prim_int_impl!(u32,   i32,   u32);
prim_int_impl!(u64,   i64,   u64);
prim_int_impl!(usize, isize, usize);
prim_int_impl!(i8,    i8,    u8);
prim_int_impl!(i16,   i16,   u16);
prim_int_impl!(i32,   i32,   u32);
prim_int_impl!(i64,   i64,   u64);
prim_int_impl!(isize, isize, usize);
<|endoftext|>


<commit_before>use crypt32;
use winapi;

use Inner;

pub struct CtlContext(winapi::PCCTL_CONTEXT);

impl Drop for CtlContext {
	fn drop(&mut self) {
		unsafe {
			crypt32::CertFreeCTLContext(self.0);
		}
	}
}

impl Inner<winapi::PCCTL_CONTEXT> for CtlContext {
	unsafe fn from_inner(t: winapi::PCCTL_CONTEXT) -> CtlContext {
		CtlContext(t)
	}

	fn as_inner(&self) -> winapi::PCCTL_CONTEXT {
		self.0
	}
}<commit_msg>CTL creation<commit_after>use crypt32;
use std::io;
use std::mem;
use std::ptr;
use winapi;

use cert_context::CertContext;
use Inner;

pub struct CtlContext(winapi::PCCTL_CONTEXT);

impl Drop for CtlContext {
	fn drop(&mut self) {
		unsafe {
			crypt32::CertFreeCTLContext(self.0);
		}
	}
}

impl Inner<winapi::PCCTL_CONTEXT> for CtlContext {
	unsafe fn from_inner(t: winapi::PCCTL_CONTEXT) -> CtlContext {
		CtlContext(t)
	}

	fn as_inner(&self) -> winapi::PCCTL_CONTEXT {
		self.0
	}
}

impl CtlContext {
	pub fn builder() -> Builder {
		Builder {
			entries: vec![],
			usages: vec![],
		}
	}
}

pub struct Builder {
	entries: Vec<CtlEntry>,
	usages: Vec<Vec<u8>>,
}

impl Builder {
	pub fn entry(&mut self, entry: CtlEntry) -> &mut Builder {
		self.entries.push(entry);
		self
	}

	pub fn usage(&mut self, usage: &str) -> &mut Builder {
		let mut usage = usage.as_bytes().to_owned();
		usage.push(0);
		self.usages.push(usage);
		self
	}

	pub fn build(&self) -> io::Result<Vec<u8>> {
		unsafe {
			let encoding = winapi::X509_ASN_ENCODING | winapi::PKCS_7_ASN_ENCODING;

			let mut usages = self.usages.iter().map(|u| u.as_ptr()).collect::<Vec<_>>();
			let mut entries = self.entries.iter()
							      .map(|e| *(e.0.as_ptr() as *const winapi::CTL_ENTRY))
							      .collect::<Vec<_>>();

			let mut ctl_info: winapi::CTL_INFO = mem::zeroed();
			ctl_info.dwVersion = winapi::CTL_V1;
			ctl_info.SubjectUsage.cUsageIdentifier = usages.len() as winapi::DWORD;
			ctl_info.SubjectUsage.rgpszUsageIdentifier = usages.as_mut_ptr() as *mut winapi::LPSTR;
			let mut algorithm = winapi::szOID_OIWSEC_sha1.as_bytes().to_owned();
			algorithm.push(0);
			ctl_info.SubjectAlgorithm.pszObjId = algorithm.as_ptr() as winapi::LPSTR;
			ctl_info.cCTLEntry = entries.len() as winapi::DWORD;
			ctl_info.rgCTLEntry = entries.as_mut_ptr();

			let mut sign_info: winapi::CMSG_SIGNED_ENCODE_INFO = mem::zeroed();
			sign_info.cbSize = mem::size_of_val(&sign_info) as winapi::DWORD;

			let flags = winapi::CMSG_ENCODE_SORTED_CTL_FLAG |
				winapi::CMSG_ENCODE_HASHED_SUBJECT_IDENTIFIER_FLAG;

			let mut size = 0;

			let res = crypt32::CryptMsgEncodeAndSignCTL(
				encoding,
				&mut ctl_info,
				&mut sign_info,
				flags,
				ptr::null_mut(),
				&mut size);
			if res == winapi::FALSE {
				return Err(io::Error::last_os_error());
			}

			let mut encoded = vec![0; size as usize];

			let res = crypt32::CryptMsgEncodeAndSignCTL(
				encoding,
				&mut ctl_info,
				&mut sign_info,
				flags,
				encoded.as_mut_ptr() as *mut winapi::BYTE,
				&mut size);
			if res == winapi::FALSE {
				return Err(io::Error::last_os_error());
			}

			Ok(encoded)
		}
	}
}

pub struct CtlEntry(Vec<u8>);

impl CtlEntry {
	pub fn from_certificate(cert: &mut CertContext) -> io::Result<CtlEntry> {
		unsafe {
			let mut size = 0;

			let res = crypt32::CertCreateCTLEntryFromCertificateContextProperties(
				cert.as_inner(),
				0,
				ptr::null_mut(),
				winapi::CTL_ENTRY_FROM_PROP_CHAIN_FLAG,
				ptr::null_mut(),
				ptr::null_mut(),
				&mut size);
			if res == winapi::FALSE {
				return Err(io::Error::last_os_error());
			}

			let mut entry = vec![0; size as usize];
			let res = crypt32::CertCreateCTLEntryFromCertificateContextProperties(
				cert.as_inner(),
				0,
				ptr::null_mut(),
				winapi::CTL_ENTRY_FROM_PROP_CHAIN_FLAG,
				ptr::null_mut(),
				entry.as_mut_ptr() as winapi::PCTL_ENTRY,
				&mut size);
			if res == winapi::FALSE {
				Err(io::Error::last_os_error())
			} else {
				Ok(CtlEntry(entry))
			}
		}
	}
}

#[cfg(test)]
mod test {
	use std::io::Read;
	use std::fs::File;
	use cert_store::CertStore;
	use ctl_context::{CtlContext, CtlEntry};

	#[test]
	fn create_ctl() {
        let mut file = File::open(concat!(env!("CARGO_MANIFEST_DIR"),
                                          "/test/self-signed.badssl.com.cer"))
                           .unwrap();
        let mut cert = vec![];
        file.read_to_end(&mut cert).unwrap();

        let mut store = CertStore::memory().unwrap();
        let mut cert = store.add_der_certificate(&cert).unwrap();

        let entry = CtlEntry::from_certificate(&mut cert).unwrap();
        CtlContext::builder()
            .entry(entry)
            .usage("1.3.6.1.4.1.311.2.2.2")
            .build()
            .unwrap();
	}
}<|endoftext|>


<commit_before>use evaluation::Level;
use types::PokeType;

mod generated_species;
pub use self::generated_species::*;

pub const NUM_SPECIES: u16 = 251;

#[derive(Debug, Eq, PartialEq, Hash)]
/// Information about a specific Pokemon species.
pub struct PokeSpecies {
    id: u16,
    attack: u16,
    defense: u16,
    stamina: u16,
    primary_type: PokeType,
    secondary_type: Option<PokeType>,
    english: &'static str,
    japanese: &'static str,
    japanese_transliterated: &'static str,
    korean: &'static str,
    chinese: &'static str,
    french: &'static str,
    german: &'static str,
    spanish: &'static str,
    italian: &'static str,
}

impl PokeSpecies {
    /// Given the id of a species, return its corresponding data. Return `None` if no data is found
    /// for the given id.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::species::PokeSpecies;
    ///
    /// let dragonite = PokeSpecies::from_id(149).unwrap();
    /// assert_eq!((263, 201, 182), dragonite.base_stats());
    ///
    /// let no_such_pokemon = PokeSpecies::from_id(500);
    /// assert!(no_such_pokemon.is_none());
    /// ```
    pub fn from_id(id: u16) -> Option<&'static PokeSpecies> {
        if id == 0 || id > NUM_SPECIES - 1 {
            return None
        }

        Some(generated_species::ALL_SPECIES[(id - 1) as usize])
    }

    /// Return a list of all `Species`, sorted by id.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::species::PokeSpecies;
    ///
    /// let all_species = PokeSpecies::all_species();
    ///
    /// assert_eq!(251, all_species.len());
    /// assert_eq!(1, all_species[0].id());
    /// ```
    pub fn all_species() -> &'static [&'static PokeSpecies] {
        generated_species::ALL_SPECIES
    }

    /// This species' id.
    pub fn id(&self) -> u16 {
        self.id
    }

    /// This species' base attack value.
    pub fn attack(&self) -> u16 {
        self.attack
    }

    /// This species's base defense value.
    pub fn defense(&self) -> u16 {
        self.defense
    }

    /// This species' base stamina value.
    pub fn stamina(&self) -> u16 {
        self.stamina
    }

    /// This species' base stats, as an `(attack, defense, stamina)` triple.
    pub fn base_stats(&self) -> (u16, u16, u16) {
        (self.attack, self.defense, self.stamina)
    }

    /// This species' primary `PokeType`.
    pub fn primary_type(&self) -> PokeType {
        self.primary_type
    }

    /// If this species is dual-typed, returns the secondary `PokeType`.
    pub fn secondary_type(&self) -> Option<PokeType> {
        self.secondary_type
    }

    /// Returns `true` if this species has two `PokeType`s.
    pub fn is_dual_type(&self) -> bool {
        self.secondary_type.is_some()
    }

    pub fn english(&self) -> &'static str {
        self.english
    }

    pub fn japanese(&self) -> &'static str {
        self.japanese
    }

    pub fn japanese_transliterated(&self) -> &'static str {
        self.japanese_transliterated
    }

    pub fn korean(&self) -> &'static str {
        self.korean
    }

    pub fn chinese(&self) -> &'static str {
        self.chinese
    }

    pub fn french(&self) -> &'static str {
        self.french
    }

    pub fn german(&self) -> &'static str {
        self.german
    }

    pub fn spanish(&self) -> &'static str {
        self.spanish
    }

    pub fn italian(&self) -> &'static str {
        self.italian
    }

    /// Returns the highest possible stats for this pokemon (e.g., perfect IVs).
    pub fn perfect_stats(&self) -> (u16, u16, u16) {
        (self.attack() + 15, self.defense() + 15, self.stamina() + 15)
    }

    /// Returns the species' max combat power at the given `Level`, assuming perfect IVs.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::{Level, species};
    ///
    /// let level_30 = Level::new(30).unwrap();
    /// assert_eq!(188, species::Magikarp.max_cp_at_level(level_30));
    /// ```
   pub fn max_cp_at_level(&self, level: Level) -> u16 {
        let (attack, defense, stamina) = self.perfect_stats();

        PokeSpecies::calculate_cp(attack, defense, stamina, level)
    }

    /// Returns this species' max Combat Power at level 39, assuming perfect IVs.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::species;
    ///
    /// assert_eq!(3_617, species::Tyranitar.max_cp());
    /// ```
    pub fn max_cp(&self) -> u16 {
        let (attack, defense, stamina) = self.perfect_stats();

        PokeSpecies::calculate_cp(attack, defense, stamina, Level::max())
    }

    fn calculate_cp(attack: u16, defense: u16, stamina: u16, level: Level) -> u16 {
        let attack_calc = attack as f64;
        let defense_calc = (defense as f64).sqrt();
        let stamina_calc = (stamina as f64).sqrt();

        let cp_multiplier = level.cp_multiplier().powi(2);

        let inner = attack_calc * defense_calc * stamina_calc * cp_multiplier / 10.;

        ::std::cmp::max(10, inner.floor() as u16)
    }
}

#[cfg(test)]
mod tests {
    use evaluation::Level;
    use types::PokeType;
    use species::*;

    #[test]
    fn species_by_id() {
        let pokemon = PokeSpecies::from_id(1).unwrap();
        assert_eq!(Bulbasaur, pokemon);

        let pokemon = PokeSpecies::from_id(161).unwrap();
        assert_eq!(Sentret, pokemon);

        let pokemon = PokeSpecies::from_id(249).unwrap();
        assert_eq!(Lugia, pokemon);

        assert!(PokeSpecies::from_id(252).is_none());
    }

    #[test]
    fn poketype() {
        assert_eq!(PokeType::Fire, Charizard.primary_type());
        assert_eq!(Some(PokeType::Flying), Charizard.secondary_type());

        assert_eq!(PokeType::Water, Squirtle.primary_type());
        assert!(Squirtle.secondary_type().is_none());
    }

    #[test]
    fn all_species() {
        let all_pokemon = PokeSpecies::all_species();
        let num_pokemon = 251;

        assert_eq!(num_pokemon, all_pokemon.len());

        for i in 0..num_pokemon {
            assert!(
                all_pokemon[i].id() == (i + 1) as u16,
                format!("Couldn't find pokemon #{}", i)
            );
        }
    }

    #[test]
    fn max_cp_at_level() {
        let level = Level::new(20).unwrap();

        assert_eq!(224, Kakuna.max_cp_at_level(level));
    }

    #[test]
    fn max_cp() {
        assert_eq!(955, Eevee.max_cp());
        assert_eq!(3_530, Dragonite.max_cp());
    }
}
<commit_msg>Fix reference to old Species<commit_after>use evaluation::Level;
use types::PokeType;

mod generated_species;
pub use self::generated_species::*;

pub const NUM_SPECIES: u16 = 251;

#[derive(Debug, Eq, PartialEq, Hash)]
/// Information about a specific Pokemon species.
pub struct PokeSpecies {
    id: u16,
    attack: u16,
    defense: u16,
    stamina: u16,
    primary_type: PokeType,
    secondary_type: Option<PokeType>,
    english: &'static str,
    japanese: &'static str,
    japanese_transliterated: &'static str,
    korean: &'static str,
    chinese: &'static str,
    french: &'static str,
    german: &'static str,
    spanish: &'static str,
    italian: &'static str,
}

impl PokeSpecies {
    /// Given the id of a species, return its corresponding data. Return `None` if no data is found
    /// for the given id.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::species::PokeSpecies;
    ///
    /// let dragonite = PokeSpecies::from_id(149).unwrap();
    /// assert_eq!((263, 201, 182), dragonite.base_stats());
    ///
    /// let no_such_pokemon = PokeSpecies::from_id(500);
    /// assert!(no_such_pokemon.is_none());
    /// ```
    pub fn from_id(id: u16) -> Option<&'static PokeSpecies> {
        if id == 0 || id > NUM_SPECIES - 1 {
            return None
        }

        Some(generated_species::ALL_SPECIES[(id - 1) as usize])
    }

    /// Return a list of all `PokeSpecies`, sorted by id.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::species::PokeSpecies;
    ///
    /// let all_species = PokeSpecies::all_species();
    ///
    /// assert_eq!(251, all_species.len());
    /// assert_eq!(1, all_species[0].id());
    /// ```
    pub fn all_species() -> &'static [&'static PokeSpecies] {
        generated_species::ALL_SPECIES
    }

    /// This species' id.
    pub fn id(&self) -> u16 {
        self.id
    }

    /// This species' base attack value.
    pub fn attack(&self) -> u16 {
        self.attack
    }

    /// This species's base defense value.
    pub fn defense(&self) -> u16 {
        self.defense
    }

    /// This species' base stamina value.
    pub fn stamina(&self) -> u16 {
        self.stamina
    }

    /// This species' base stats, as an `(attack, defense, stamina)` triple.
    pub fn base_stats(&self) -> (u16, u16, u16) {
        (self.attack, self.defense, self.stamina)
    }

    /// This species' primary `PokeType`.
    pub fn primary_type(&self) -> PokeType {
        self.primary_type
    }

    /// If this species is dual-typed, returns the secondary `PokeType`.
    pub fn secondary_type(&self) -> Option<PokeType> {
        self.secondary_type
    }

    /// Returns `true` if this species has two `PokeType`s.
    pub fn is_dual_type(&self) -> bool {
        self.secondary_type.is_some()
    }

    pub fn english(&self) -> &'static str {
        self.english
    }

    pub fn japanese(&self) -> &'static str {
        self.japanese
    }

    pub fn japanese_transliterated(&self) -> &'static str {
        self.japanese_transliterated
    }

    pub fn korean(&self) -> &'static str {
        self.korean
    }

    pub fn chinese(&self) -> &'static str {
        self.chinese
    }

    pub fn french(&self) -> &'static str {
        self.french
    }

    pub fn german(&self) -> &'static str {
        self.german
    }

    pub fn spanish(&self) -> &'static str {
        self.spanish
    }

    pub fn italian(&self) -> &'static str {
        self.italian
    }

    /// Returns the highest possible stats for this pokemon (e.g., perfect IVs).
    pub fn perfect_stats(&self) -> (u16, u16, u16) {
        (self.attack() + 15, self.defense() + 15, self.stamina() + 15)
    }

    /// Returns the species' max combat power at the given `Level`, assuming perfect IVs.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::{Level, species};
    ///
    /// let level_30 = Level::new(30).unwrap();
    /// assert_eq!(188, species::Magikarp.max_cp_at_level(level_30));
    /// ```
   pub fn max_cp_at_level(&self, level: Level) -> u16 {
        let (attack, defense, stamina) = self.perfect_stats();

        PokeSpecies::calculate_cp(attack, defense, stamina, level)
    }

    /// Returns this species' max Combat Power at level 39, assuming perfect IVs.
    ///
    /// ## Example:
    ///
    /// ```
    /// use pokemon_go_data::species;
    ///
    /// assert_eq!(3_617, species::Tyranitar.max_cp());
    /// ```
    pub fn max_cp(&self) -> u16 {
        let (attack, defense, stamina) = self.perfect_stats();

        PokeSpecies::calculate_cp(attack, defense, stamina, Level::max())
    }

    fn calculate_cp(attack: u16, defense: u16, stamina: u16, level: Level) -> u16 {
        let attack_calc = attack as f64;
        let defense_calc = (defense as f64).sqrt();
        let stamina_calc = (stamina as f64).sqrt();

        let cp_multiplier = level.cp_multiplier().powi(2);

        let inner = attack_calc * defense_calc * stamina_calc * cp_multiplier / 10.;

        ::std::cmp::max(10, inner.floor() as u16)
    }
}

#[cfg(test)]
mod tests {
    use evaluation::Level;
    use types::PokeType;
    use species::*;

    #[test]
    fn species_by_id() {
        let pokemon = PokeSpecies::from_id(1).unwrap();
        assert_eq!(Bulbasaur, pokemon);

        let pokemon = PokeSpecies::from_id(161).unwrap();
        assert_eq!(Sentret, pokemon);

        let pokemon = PokeSpecies::from_id(249).unwrap();
        assert_eq!(Lugia, pokemon);

        assert!(PokeSpecies::from_id(252).is_none());
    }

    #[test]
    fn poketype() {
        assert_eq!(PokeType::Fire, Charizard.primary_type());
        assert_eq!(Some(PokeType::Flying), Charizard.secondary_type());

        assert_eq!(PokeType::Water, Squirtle.primary_type());
        assert!(Squirtle.secondary_type().is_none());
    }

    #[test]
    fn all_species() {
        let all_pokemon = PokeSpecies::all_species();
        let num_pokemon = 251;

        assert_eq!(num_pokemon, all_pokemon.len());

        for i in 0..num_pokemon {
            assert!(
                all_pokemon[i].id() == (i + 1) as u16,
                format!("Couldn't find pokemon #{}", i)
            );
        }
    }

    #[test]
    fn max_cp_at_level() {
        let level = Level::new(20).unwrap();

        assert_eq!(224, Kakuna.max_cp_at_level(level));
    }

    #[test]
    fn max_cp() {
        assert_eq!(955, Eevee.max_cp());
        assert_eq!(3_530, Dragonite.max_cp());
    }
}
<|endoftext|>


<commit_before>//! Header generation for data packets.
//!
//! A header has several user-filled fields followed by a 16-bit checksum over those
//! fields.

use data::crc;
use data::values;

/// Packet header block for confirmed data packet.
pub type ConfirmedHeader = Header<ConfirmedFields>;

/// Packet header block for unconfirmed data packet.
pub type UnconfirmedHeader = Header<UnconfirmedFields>;

/// Write some bytes into a buffer.
pub trait BufWrite {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T);
}

/// Field is only a single byte.
pub trait ByteField {
    fn byte(&self) -> u8;
}

/// Just write the single byte.
impl<B: ByteField> BufWrite for B {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        *buf.next().unwrap() = self.byte();
    }
}

/// Preamble header field.
pub struct HeaderPreamble {
    /// Whether the packet requires confirmation.
    pub confirmed: bool,
    /// Whether the packet is an outbound message.
    pub outbound: bool,
    /// Packet type.
    pub format: values::DataPacket,
}

impl ByteField for HeaderPreamble {
    fn byte(&self) -> u8 {
        bool_to_bit(self.confirmed) << 6 | bool_to_bit(self.outbound) << 5 |
            self.format.to_bits()
    }
}

/// Preamble for confirmed data packet.
pub struct ConfirmedPreamble(HeaderPreamble);

impl ConfirmedPreamble {
    fn new(outbound: bool) -> ConfirmedPreamble {
        ConfirmedPreamble(HeaderPreamble {
            confirmed: true,
            outbound: outbound,
            format: values::DataPacket::ConfirmedPacket,
        })
    }

    pub fn outbound() -> Self { Self::new(true) }
    pub fn inbound() -> Self { Self::new(false) }
}

impl ByteField for ConfirmedPreamble {
    fn byte(&self) -> u8 { self.0.byte() }
}

/// Preamble for unconfirmed data packet.
pub struct UnconfirmedPreamble(HeaderPreamble);

impl UnconfirmedPreamble {
    fn new(outbound: bool) -> UnconfirmedPreamble {
        UnconfirmedPreamble(HeaderPreamble {
            confirmed: false,
            outbound: outbound,
            format: values::DataPacket::UnconfirmedPacket
        })
    }

    pub fn outbound() -> Self { Self::new(true) }
    pub fn inbound() -> Self { Self::new(false) }
}

impl ByteField for UnconfirmedPreamble {
    fn byte(&self) -> u8 { self.0.byte() }
}

/// Service access point (SAP) field.
pub struct ServiceAccessPoint(pub values::ServiceAccessPoint);

impl ByteField for ServiceAccessPoint {
    fn byte(&self) -> u8 {
        0b11000000 | self.0.to_bits()
    }
}

/// Manufacturer's ID field.
pub struct Manufacturer(pub u8);

impl ByteField for Manufacturer {
    fn byte(&self) -> u8 { self.0 }
}

/// Logical link ID field for addressing source or destination subscriber.
pub struct LogicalLink(pub u32);

impl BufWrite for LogicalLink {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        assert!(self.0 >> 24 == 0);

        *buf.next().unwrap() = (self.0 >> 16) as u8;
        *buf.next().unwrap() = (self.0 >> 8) as u8;
        *buf.next().unwrap() = self.0 as u8;
    }
}

/// FMF and blocks-to-follow fields.
pub struct BlockCount {
    /// Whether the packet is "complete", not being partially retransmitted.
    pub full_pkt: bool,
    /// Number of data blocks in the packet.
    pub count: u8,
}

impl ByteField for BlockCount {
    fn byte(&self) -> u8 {
        assert!(self.count >> 7 == 0);
        bool_to_bit(self.full_pkt) << 7 | self.count
    }
}

/// Number of pad bytes at the end of the data.
pub struct PadCount(pub u8);

impl ByteField for PadCount {
    fn byte(&self) -> u8 {
        assert!(self.0 >> 5 == 0);
        self.0
    }
}

/// Syn, N(S), and FSNF fields.
pub struct Sequencing {
    /// Whether the receiver should resynchronize its sequence numbers using `pkt_seq` and
    /// `frag_seq`.
    pub resync: bool,
    /// Packet sequence number, used for ordering and duplicate removal.
    pub pkt_seq: u8,
    /// Fragment sequence number, used when data is split across multiple fragments.
    pub frag_seq: u8,
}

impl ByteField for Sequencing {
    fn byte(&self) -> u8 {
        assert!(self.pkt_seq >> 3 == 0);
        assert!(self.frag_seq >> 4 == 0);

        bool_to_bit(self.resync) << 7 | self.pkt_seq << 4 | self.frag_seq
    }
}

/// Byte offset into data payload where data header stops and data information begins.
pub struct DataOffset(pub u8);

impl ByteField for DataOffset {
    fn byte(&self) -> u8 {
        assert!(self.0 >> 6 == 0);
        self.0
    }
}

/// Header fields for confirmed packet.
pub struct ConfirmedFields {
    pub preamble: ConfirmedPreamble,
    pub sap: ServiceAccessPoint,
    pub mfg: Manufacturer,
    pub addr: LogicalLink,
    pub blocks: BlockCount,
    pub pads: PadCount,
    pub seq: Sequencing,
    pub data_offset: DataOffset,
}

impl BufWrite for ConfirmedFields {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        self.preamble.write(buf);
        self.sap.write(buf);
        self.mfg.write(buf);
        self.addr.write(buf);
        self.blocks.write(buf);
        self.pads.write(buf);
        self.seq.write(buf);
        self.data_offset.write(buf);
    }
}

/// Header fields for unconfirmed packet.
pub struct UnconfirmedFields {
    pub preamble: UnconfirmedPreamble,
    pub sap: ServiceAccessPoint,
    pub mfg: Manufacturer,
    pub addr: LogicalLink,
    pub blocks: BlockCount,
    pub pads: PadCount,
    pub data_offset: DataOffset,
}

impl BufWrite for UnconfirmedFields {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        self.preamble.write(buf);
        self.sap.write(buf);
        self.mfg.write(buf);
        self.addr.write(buf);
        self.blocks.write(buf);
        self.pads.write(buf);
        *buf.next().unwrap() = 0;
        self.data_offset.write(buf);
    }
}

/// Builds a checksummed header based on the given fields.
pub struct Header<F: BufWrite>(F);

impl<F: BufWrite> Header<F> {
    /// Construct a new `Header` with the given fields.
    pub fn new(fields: F) -> Header<F> {
        Header(fields)
    }

    /// Get the fields and checksum that make up the header, in that order.
    pub fn build(self) -> ([u8; 10], [u8; 2]) {
        let fields = self.fields();
        let checksum = self.checksum(&fields);

        (fields, checksum)
    }

    /// Build a byte buffer from the header fields.
    fn fields(&self) -> [u8; 10] {
        let mut buf = [0; 10];
        self.0.write(&mut buf.iter_mut());
        buf
    }

    /// Calculate the checksum of the header fields.
    fn checksum(&self, fields: &[u8]) -> [u8; 2] {
        assert!(fields.len() == 10);

        let checksum = crc::CRC16::new()
            .feed_bytes(fields.iter().cloned())
            .finish();

        [(checksum >> 8) as u8, checksum as u8]
    }
}

/// Convert the given Boolean to a single bit.
fn bool_to_bit(b: bool) -> u8 {
    if b { 1 } else { 0 }
}

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

    #[test]
    fn test_preamble() {
        let p = ConfirmedPreamble::outbound();
        assert_eq!(p.byte(), 0b01110110);
        let p = ConfirmedPreamble::inbound();
        assert_eq!(p.byte(), 0b01010110);
    }

    #[test]
    fn test_sap() {
        let s = ServiceAccessPoint(values::ServiceAccessPoint::ExtendedAddressing);
        assert_eq!(s.byte(), 0b11011111);
    }

    #[test]
    fn test_mfg() {
        let m = Manufacturer(0b11011011);
        assert_eq!(m.byte(), 0b11011011);
    }

    #[test]
    fn test_ll() {
        let l = LogicalLink(0xABCDEF);
        let mut buf = [0; 3];
        l.write(&mut buf.iter_mut());
        assert_eq!(&buf, &[0xAB, 0xCD, 0xEF]);
    }

    #[test]
    fn test_bc() {
        let b = BlockCount {
            full_pkt: true,
            count: 127,
        };

        assert_eq!(b.byte(), 0b11111111);
    }

    #[test]
    fn test_pc() {
        let p = PadCount(12);
        assert_eq!(p.byte(), 0b00001100);
    }

    #[test]
    fn test_seq() {
        let s = Sequencing {
            resync: false,
            pkt_seq: 6,
            frag_seq: 10,
        };
        assert_eq!(s.byte(), 0b01101010);
    }

    #[test]
    fn test_do() {
        let d = DataOffset(11);
        assert_eq!(d.byte(), 0b00001011);
    }

    #[test]
    fn test_confirmed_fields() {
        let fields = ConfirmedFields {
            preamble: ConfirmedPreamble::outbound(),
            sap: ServiceAccessPoint(values::ServiceAccessPoint::Paging),
            mfg: Manufacturer(0x12),
            addr: LogicalLink(0x342134),
            blocks: BlockCount {
                full_pkt: true,
                count: 127,
            },
            pads: PadCount(3),
            seq: Sequencing {
                resync: false,
                pkt_seq: 5,
                frag_seq: 2,
            },
            data_offset: DataOffset(0),
        };

        let mut buf = [0; 10];
        fields.write(&mut buf.iter_mut());

        assert_eq!(&buf, &[
            0b01110110,
            0b11100110,
            0b00010010,
            0b00110100,
            0b00100001,
            0b00110100,
            0b11111111,
            0b00000011,
            0b01010010,
            0b00000000,
        ]);
    }

    #[test]
    fn test_confirmed_header() {
        let (fields, checksum) = ConfirmedHeader::new(ConfirmedFields {
            preamble: ConfirmedPreamble::outbound(),
            sap: ServiceAccessPoint(values::ServiceAccessPoint::PacketData),
            mfg: Manufacturer(0x12),
            addr: LogicalLink(0x342134),
            blocks: BlockCount {
                full_pkt: true,
                count: 127,
            },
            pads: PadCount(3),
            seq: Sequencing {
                resync: false,
                pkt_seq: 5,
                frag_seq: 2,
            },
            data_offset: DataOffset(0),
        }).build();

        assert_eq!(fields, [
            0b01110110,
            0b11000100,
            0b00010010,
            0b00110100,
            0b00100001,
            0b00110100,
            0b11111111,
            0b00000011,
            0b01010010,
            0b00000000,
        ]);

        assert_eq!(checksum, [
            0b10001010,
            0b01110010,
        ]);
    }

    #[test]
    #[should_panic]
    fn test_ll_validate() {
        let mut buf = [0; 4];
        let mut iter = buf.iter_mut();
        LogicalLink(0x20FFFFFF).write(&mut iter);
    }

    #[test]
    #[should_panic]
    fn test_bc_validate() {
        BlockCount {
            full_pkt: true,
            count: 0b11111111,
        }.byte();
    }

    #[test]
    #[should_panic]
    fn test_pc_validate() {
        PadCount(0b11111111).byte();
    }

    #[test]
    #[should_panic]
    fn test_seq_validate() {
        Sequencing {
            resync: true,
            pkt_seq: 0b11111111,
            frag_seq: 0b11111111,
        }.byte();
    }

    #[test]
    #[should_panic]
    fn test_do_validate() {
        DataOffset(0b11111111).byte();
    }
}
<commit_msg>Fix nightly compiler warning<commit_after>//! Header generation for data packets.
//!
//! A header has several user-filled fields followed by a 16-bit checksum over those
//! fields.

use data::crc;
use data::values;

/// Packet header block for confirmed data packet.
pub type ConfirmedHeader = Header<ConfirmedFields>;

/// Packet header block for unconfirmed data packet.
pub type UnconfirmedHeader = Header<UnconfirmedFields>;

/// Write some bytes into a buffer.
pub trait BufWrite {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, buf: &'b mut T);
}

/// Field is only a single byte.
pub trait ByteField {
    fn byte(&self) -> u8;
}

/// Just write the single byte.
impl<B: ByteField> BufWrite for B {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        *buf.next().unwrap() = self.byte();
    }
}

/// Preamble header field.
pub struct HeaderPreamble {
    /// Whether the packet requires confirmation.
    pub confirmed: bool,
    /// Whether the packet is an outbound message.
    pub outbound: bool,
    /// Packet type.
    pub format: values::DataPacket,
}

impl ByteField for HeaderPreamble {
    fn byte(&self) -> u8 {
        bool_to_bit(self.confirmed) << 6 | bool_to_bit(self.outbound) << 5 |
            self.format.to_bits()
    }
}

/// Preamble for confirmed data packet.
pub struct ConfirmedPreamble(HeaderPreamble);

impl ConfirmedPreamble {
    fn new(outbound: bool) -> ConfirmedPreamble {
        ConfirmedPreamble(HeaderPreamble {
            confirmed: true,
            outbound: outbound,
            format: values::DataPacket::ConfirmedPacket,
        })
    }

    pub fn outbound() -> Self { Self::new(true) }
    pub fn inbound() -> Self { Self::new(false) }
}

impl ByteField for ConfirmedPreamble {
    fn byte(&self) -> u8 { self.0.byte() }
}

/// Preamble for unconfirmed data packet.
pub struct UnconfirmedPreamble(HeaderPreamble);

impl UnconfirmedPreamble {
    fn new(outbound: bool) -> UnconfirmedPreamble {
        UnconfirmedPreamble(HeaderPreamble {
            confirmed: false,
            outbound: outbound,
            format: values::DataPacket::UnconfirmedPacket
        })
    }

    pub fn outbound() -> Self { Self::new(true) }
    pub fn inbound() -> Self { Self::new(false) }
}

impl ByteField for UnconfirmedPreamble {
    fn byte(&self) -> u8 { self.0.byte() }
}

/// Service access point (SAP) field.
pub struct ServiceAccessPoint(pub values::ServiceAccessPoint);

impl ByteField for ServiceAccessPoint {
    fn byte(&self) -> u8 {
        0b11000000 | self.0.to_bits()
    }
}

/// Manufacturer's ID field.
pub struct Manufacturer(pub u8);

impl ByteField for Manufacturer {
    fn byte(&self) -> u8 { self.0 }
}

/// Logical link ID field for addressing source or destination subscriber.
pub struct LogicalLink(pub u32);

impl BufWrite for LogicalLink {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        assert!(self.0 >> 24 == 0);

        *buf.next().unwrap() = (self.0 >> 16) as u8;
        *buf.next().unwrap() = (self.0 >> 8) as u8;
        *buf.next().unwrap() = self.0 as u8;
    }
}

/// FMF and blocks-to-follow fields.
pub struct BlockCount {
    /// Whether the packet is "complete", not being partially retransmitted.
    pub full_pkt: bool,
    /// Number of data blocks in the packet.
    pub count: u8,
}

impl ByteField for BlockCount {
    fn byte(&self) -> u8 {
        assert!(self.count >> 7 == 0);
        bool_to_bit(self.full_pkt) << 7 | self.count
    }
}

/// Number of pad bytes at the end of the data.
pub struct PadCount(pub u8);

impl ByteField for PadCount {
    fn byte(&self) -> u8 {
        assert!(self.0 >> 5 == 0);
        self.0
    }
}

/// Syn, N(S), and FSNF fields.
pub struct Sequencing {
    /// Whether the receiver should resynchronize its sequence numbers using `pkt_seq` and
    /// `frag_seq`.
    pub resync: bool,
    /// Packet sequence number, used for ordering and duplicate removal.
    pub pkt_seq: u8,
    /// Fragment sequence number, used when data is split across multiple fragments.
    pub frag_seq: u8,
}

impl ByteField for Sequencing {
    fn byte(&self) -> u8 {
        assert!(self.pkt_seq >> 3 == 0);
        assert!(self.frag_seq >> 4 == 0);

        bool_to_bit(self.resync) << 7 | self.pkt_seq << 4 | self.frag_seq
    }
}

/// Byte offset into data payload where data header stops and data information begins.
pub struct DataOffset(pub u8);

impl ByteField for DataOffset {
    fn byte(&self) -> u8 {
        assert!(self.0 >> 6 == 0);
        self.0
    }
}

/// Header fields for confirmed packet.
pub struct ConfirmedFields {
    pub preamble: ConfirmedPreamble,
    pub sap: ServiceAccessPoint,
    pub mfg: Manufacturer,
    pub addr: LogicalLink,
    pub blocks: BlockCount,
    pub pads: PadCount,
    pub seq: Sequencing,
    pub data_offset: DataOffset,
}

impl BufWrite for ConfirmedFields {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        self.preamble.write(buf);
        self.sap.write(buf);
        self.mfg.write(buf);
        self.addr.write(buf);
        self.blocks.write(buf);
        self.pads.write(buf);
        self.seq.write(buf);
        self.data_offset.write(buf);
    }
}

/// Header fields for unconfirmed packet.
pub struct UnconfirmedFields {
    pub preamble: UnconfirmedPreamble,
    pub sap: ServiceAccessPoint,
    pub mfg: Manufacturer,
    pub addr: LogicalLink,
    pub blocks: BlockCount,
    pub pads: PadCount,
    pub data_offset: DataOffset,
}

impl BufWrite for UnconfirmedFields {
    fn write<'a, 'b, T: Iterator<Item = &'a mut u8>>(&self, mut buf: &'b mut T) {
        self.preamble.write(buf);
        self.sap.write(buf);
        self.mfg.write(buf);
        self.addr.write(buf);
        self.blocks.write(buf);
        self.pads.write(buf);
        *buf.next().unwrap() = 0;
        self.data_offset.write(buf);
    }
}

/// Builds a checksummed header based on the given fields.
pub struct Header<F: BufWrite>(F);

impl<F: BufWrite> Header<F> {
    /// Construct a new `Header` with the given fields.
    pub fn new(fields: F) -> Header<F> {
        Header(fields)
    }

    /// Get the fields and checksum that make up the header, in that order.
    pub fn build(self) -> ([u8; 10], [u8; 2]) {
        let fields = self.fields();
        let checksum = self.checksum(&fields);

        (fields, checksum)
    }

    /// Build a byte buffer from the header fields.
    fn fields(&self) -> [u8; 10] {
        let mut buf = [0; 10];
        self.0.write(&mut buf.iter_mut());
        buf
    }

    /// Calculate the checksum of the header fields.
    fn checksum(&self, fields: &[u8]) -> [u8; 2] {
        assert!(fields.len() == 10);

        let checksum = crc::CRC16::new()
            .feed_bytes(fields.iter().cloned())
            .finish();

        [(checksum >> 8) as u8, checksum as u8]
    }
}

/// Convert the given Boolean to a single bit.
fn bool_to_bit(b: bool) -> u8 {
    if b { 1 } else { 0 }
}

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

    #[test]
    fn test_preamble() {
        let p = ConfirmedPreamble::outbound();
        assert_eq!(p.byte(), 0b01110110);
        let p = ConfirmedPreamble::inbound();
        assert_eq!(p.byte(), 0b01010110);
    }

    #[test]
    fn test_sap() {
        let s = ServiceAccessPoint(values::ServiceAccessPoint::ExtendedAddressing);
        assert_eq!(s.byte(), 0b11011111);
    }

    #[test]
    fn test_mfg() {
        let m = Manufacturer(0b11011011);
        assert_eq!(m.byte(), 0b11011011);
    }

    #[test]
    fn test_ll() {
        let l = LogicalLink(0xABCDEF);
        let mut buf = [0; 3];
        l.write(&mut buf.iter_mut());
        assert_eq!(&buf, &[0xAB, 0xCD, 0xEF]);
    }

    #[test]
    fn test_bc() {
        let b = BlockCount {
            full_pkt: true,
            count: 127,
        };

        assert_eq!(b.byte(), 0b11111111);
    }

    #[test]
    fn test_pc() {
        let p = PadCount(12);
        assert_eq!(p.byte(), 0b00001100);
    }

    #[test]
    fn test_seq() {
        let s = Sequencing {
            resync: false,
            pkt_seq: 6,
            frag_seq: 10,
        };
        assert_eq!(s.byte(), 0b01101010);
    }

    #[test]
    fn test_do() {
        let d = DataOffset(11);
        assert_eq!(d.byte(), 0b00001011);
    }

    #[test]
    fn test_confirmed_fields() {
        let fields = ConfirmedFields {
            preamble: ConfirmedPreamble::outbound(),
            sap: ServiceAccessPoint(values::ServiceAccessPoint::Paging),
            mfg: Manufacturer(0x12),
            addr: LogicalLink(0x342134),
            blocks: BlockCount {
                full_pkt: true,
                count: 127,
            },
            pads: PadCount(3),
            seq: Sequencing {
                resync: false,
                pkt_seq: 5,
                frag_seq: 2,
            },
            data_offset: DataOffset(0),
        };

        let mut buf = [0; 10];
        fields.write(&mut buf.iter_mut());

        assert_eq!(&buf, &[
            0b01110110,
            0b11100110,
            0b00010010,
            0b00110100,
            0b00100001,
            0b00110100,
            0b11111111,
            0b00000011,
            0b01010010,
            0b00000000,
        ]);
    }

    #[test]
    fn test_confirmed_header() {
        let (fields, checksum) = ConfirmedHeader::new(ConfirmedFields {
            preamble: ConfirmedPreamble::outbound(),
            sap: ServiceAccessPoint(values::ServiceAccessPoint::PacketData),
            mfg: Manufacturer(0x12),
            addr: LogicalLink(0x342134),
            blocks: BlockCount {
                full_pkt: true,
                count: 127,
            },
            pads: PadCount(3),
            seq: Sequencing {
                resync: false,
                pkt_seq: 5,
                frag_seq: 2,
            },
            data_offset: DataOffset(0),
        }).build();

        assert_eq!(fields, [
            0b01110110,
            0b11000100,
            0b00010010,
            0b00110100,
            0b00100001,
            0b00110100,
            0b11111111,
            0b00000011,
            0b01010010,
            0b00000000,
        ]);

        assert_eq!(checksum, [
            0b10001010,
            0b01110010,
        ]);
    }

    #[test]
    #[should_panic]
    fn test_ll_validate() {
        let mut buf = [0; 4];
        let mut iter = buf.iter_mut();
        LogicalLink(0x20FFFFFF).write(&mut iter);
    }

    #[test]
    #[should_panic]
    fn test_bc_validate() {
        BlockCount {
            full_pkt: true,
            count: 0b11111111,
        }.byte();
    }

    #[test]
    #[should_panic]
    fn test_pc_validate() {
        PadCount(0b11111111).byte();
    }

    #[test]
    #[should_panic]
    fn test_seq_validate() {
        Sequencing {
            resync: true,
            pkt_seq: 0b11111111,
            frag_seq: 0b11111111,
        }.byte();
    }

    #[test]
    #[should_panic]
    fn test_do_validate() {
        DataOffset(0b11111111).byte();
    }
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fs::File as FSFile;
use std::ops::Deref;
use std::io::Write;

pub mod path;
pub mod file;
pub mod parser;
pub mod json;

use module::Module;
use runtime::Runtime;
use storage::file::File;
use storage::file::id::FileID;
use storage::file::id_type::FileIDType;
use storage::file::hash::FileHash;
use storage::parser::{FileHeaderParser, Parser, ParserError};
use storage::file::header::data::FileHeaderData;

type Cache = HashMap<FileID, Rc<RefCell<File>>>;

pub struct Store {
    storepath: String,
    cache : RefCell<Cache>,
}

impl Store {

    pub fn new(storepath: String) -> Store {
        Store {
            storepath: storepath,
            cache: RefCell::new(HashMap::new()),
        }
    }

    fn put_in_cache(&self, f: File) -> FileID {
        let res = f.id().clone();
        self.cache.borrow_mut().insert(f.id().clone(), Rc::new(RefCell::new(f)));
        res
    }

    pub fn new_file(&self, module: &Module)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: FileHeaderData::Null,
            data: String::from(""),
            id: self.get_new_file_id(),
        };

        debug!("Create new File object: {:?}", &f);
        self.put_in_cache(f)
    }

    pub fn new_file_from_parser_result(&self,
                                       module: &Module,
                                       id: FileID,
                                       header: FileHeaderData,
                                       data: String)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: header,
            data: data,
            id: id,
        };
        debug!("Create new File object from parser result: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn new_file_with_header(&self,
                                module: &Module,
                                h: FileHeaderData)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: h,
            data: String::from(""),
            id: self.get_new_file_id(),
        };
        debug!("Create new File object with header: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn new_file_with_data(&self, module: &Module, d: String)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: FileHeaderData::Null,
            data: d,
            id: self.get_new_file_id(),
        };
        debug!("Create new File object with data: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn new_file_with_content(&self,
                                 module: &Module,
                                 h: FileHeaderData,
                                 d: String)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: h,
            data: d,
            id: self.get_new_file_id(),
        };
        debug!("Create new File object with content: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn persist<HP>(&self,
                       p: &Parser<HP>,
                       f: Rc<RefCell<File>>) -> bool
        where HP: FileHeaderParser
    {
        let file = f.deref().borrow();
        let text = p.write(file.contents());
        if text.is_err() {
            error!("Error: {}", text.err().unwrap());
            return false;
        }

        let path = {
            let ids : String = file.id().clone().into();
            format!("{}/{}-{}.imag", self.storepath, file.owning_module_name, ids)
        };

        FSFile::create(&path).map(|mut fsfile| {
            fsfile.write_all(&text.unwrap().clone().into_bytes()[..])
        }).map_err(|writeerr|  {
            debug!("Could not create file at '{}'", path);
        }).and(Ok(true)).unwrap()
    }

    pub fn load(&self, id: &FileID) -> Option<Rc<RefCell<File>>> {
        debug!("Loading '{:?}'", id);
        self.cache.borrow().get(id).cloned()
    }

    fn get_new_file_id(&self) -> FileID {
        use uuid::Uuid;
        let hash = FileHash::from(Uuid::new_v4().to_hyphenated_string());
        FileID::new(FileIDType::UUID, hash)
    }

}
<commit_msg>Add Store::remove(FileID)<commit_after>use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fs::File as FSFile;
use std::ops::Deref;
use std::io::Write;

pub mod path;
pub mod file;
pub mod parser;
pub mod json;

use module::Module;
use runtime::Runtime;
use storage::file::File;
use storage::file::id::FileID;
use storage::file::id_type::FileIDType;
use storage::file::hash::FileHash;
use storage::parser::{FileHeaderParser, Parser, ParserError};
use storage::file::header::data::FileHeaderData;

type Cache = HashMap<FileID, Rc<RefCell<File>>>;

pub struct Store {
    storepath: String,
    cache : RefCell<Cache>,
}

impl Store {

    pub fn new(storepath: String) -> Store {
        Store {
            storepath: storepath,
            cache: RefCell::new(HashMap::new()),
        }
    }

    fn put_in_cache(&self, f: File) -> FileID {
        let res = f.id().clone();
        self.cache.borrow_mut().insert(f.id().clone(), Rc::new(RefCell::new(f)));
        res
    }

    pub fn new_file(&self, module: &Module)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: FileHeaderData::Null,
            data: String::from(""),
            id: self.get_new_file_id(),
        };

        debug!("Create new File object: {:?}", &f);
        self.put_in_cache(f)
    }

    pub fn new_file_from_parser_result(&self,
                                       module: &Module,
                                       id: FileID,
                                       header: FileHeaderData,
                                       data: String)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: header,
            data: data,
            id: id,
        };
        debug!("Create new File object from parser result: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn new_file_with_header(&self,
                                module: &Module,
                                h: FileHeaderData)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: h,
            data: String::from(""),
            id: self.get_new_file_id(),
        };
        debug!("Create new File object with header: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn new_file_with_data(&self, module: &Module, d: String)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: FileHeaderData::Null,
            data: d,
            id: self.get_new_file_id(),
        };
        debug!("Create new File object with data: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn new_file_with_content(&self,
                                 module: &Module,
                                 h: FileHeaderData,
                                 d: String)
        -> FileID
    {
        let f = File {
            owning_module_name: module.name(),
            header: h,
            data: d,
            id: self.get_new_file_id(),
        };
        debug!("Create new File object with content: {:?}", f);
        self.put_in_cache(f)
    }

    pub fn persist<HP>(&self,
                       p: &Parser<HP>,
                       f: Rc<RefCell<File>>) -> bool
        where HP: FileHeaderParser
    {
        let file = f.deref().borrow();
        let text = p.write(file.contents());
        if text.is_err() {
            error!("Error: {}", text.err().unwrap());
            return false;
        }

        let path = {
            let ids : String = file.id().clone().into();
            format!("{}/{}-{}.imag", self.storepath, file.owning_module_name, ids)
        };

        FSFile::create(&path).map(|mut fsfile| {
            fsfile.write_all(&text.unwrap().clone().into_bytes()[..])
        }).map_err(|writeerr|  {
            debug!("Could not create file at '{}'", path);
        }).and(Ok(true)).unwrap()
    }

    pub fn load(&self, id: &FileID) -> Option<Rc<RefCell<File>>> {
        debug!("Loading '{:?}'", id);
        self.cache.borrow().get(id).cloned()
    }

    pub fn remove(&self, id: FileID) -> bool {
        use std::fs::remove_file;

        self.cache
            .borrow_mut()
            .remove(&id)
            .map(|file| {
                let idstr : String = id.into();
                let path = format!("{}/{}-{}.imag",
                                   self.storepath,
                                   file.deref().borrow().owner_name(),
                                   idstr);
                remove_file(path).is_ok()
            })
            .unwrap_or(false)
    }

    fn get_new_file_id(&self) -> FileID {
        use uuid::Uuid;
        let hash = FileHash::from(Uuid::new_v4().to_hyphenated_string());
        FileID::new(FileIDType::UUID, hash)
    }

}
<|endoftext|>


<commit_before>use std::fs::File;
use std::os::unix::io::AsRawFd;
use std::io;

use libc::c_int;

#[link(name = "fallocate")]
extern "C" {
    fn native_fallocate(fd: c_int, len: u64) -> c_int;
}

// TODO: Consider doing the stupid fallback of zeroing the file out.
pub fn fallocate(f: &File, len: u64) -> io::Result<()> {
    // We ignore the len here, if you actually have a u64 max, then you're kinda fucked either way.
    match unsafe { native_fallocate(f.as_raw_fd(), len) } {
        0 => Ok(()),
        e => Err(io::Error::from_raw_os_error(e)),
    }
}
<commit_msg>Fall back to truncation on fallocate failure<commit_after>use std::fs::File;
use std::os::unix::io::AsRawFd;
use std::io;

use libc::c_int;

#[link(name = "fallocate")]
extern "C" {
    fn native_fallocate(fd: c_int, len: u64) -> c_int;
}

pub fn fallocate(f: &File, len: u64) -> io::Result<()> {
    // We ignore the len here, if you actually have a u64 max, then you're kinda fucked either way.
    match unsafe { native_fallocate(f.as_raw_fd(), len) } {
        0 => Ok(()),
        e => {
            error!("fallocate failed: {}", e);
            f.set_len(len)?;
            Ok(())
        }
    }
}
<|endoftext|>