<commit_before>/*!
A potentially non-blocking response decoder.

The decoder wraps a stream of chunks and produces a new stream of decompressed chunks.
The decompressed chunks aren't guaranteed to align to the compressed ones.

If the response is plaintext then no additional work is carried out.
Chunks are just passed along.

If the response is gzip, then the chunks are decompressed into a buffer.
Slices of that buffer are emitted as new chunks.

This module consists of a few main types:

- `ReadableChunks` is a `Read`-like wrapper around a stream
- `Decoder` is a layer over `ReadableChunks` that applies the right decompression

The following types directly support the gzip compression case:

- `Pending` is a non-blocking constructor for a `Decoder` in case the body needs to be checked for EOF
*/

use std::fmt;
use std::mem;
use std::cmp;
use std::io::{self, Read};

use bytes::{Buf, BufMut, BytesMut};
use flate2::read::GzDecoder;
use futures::{Async, Future, Poll, Stream};
use hyper::{HeaderMap};
use hyper::header::{CONTENT_ENCODING, CONTENT_LENGTH, TRANSFER_ENCODING};

use super::{Body, Chunk};
use error;

const INIT_BUFFER_SIZE: usize = 8192;

/// A response decompressor over a non-blocking stream of chunks.
///
/// The inner decoder may be constructed asynchronously.
pub struct Decoder {
    inner: Inner
}

enum Inner {
    /// A `PlainText` decoder just returns the response content as is.
    PlainText(Body),
    /// A `Gzip` decoder will uncompress the gzipped response content before returning it.
    Gzip(Gzip),
    /// A decoder that doesn't have a value yet.
    Pending(Pending)
}

/// A future attempt to poll the response body for EOF so we know whether to use gzip or not.
struct Pending {
    body: ReadableChunks<Body>,
}

/// A gzip decoder that reads from a `flate2::read::GzDecoder` into a `BytesMut` and emits the results
/// as a `Chunk`.
struct Gzip {
    inner: Box<GzDecoder<ReadableChunks<Body>>>,
    buf: BytesMut,
}

impl fmt::Debug for Decoder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Decoder")
            .finish()
    }
}

impl Decoder {
    /// An empty decoder.
    ///
    /// This decoder will produce a single 0 byte chunk.
    #[inline]
    pub fn empty() -> Decoder {
        Decoder {
            inner: Inner::PlainText(Body::empty())
        }
    }

    /// A plain text decoder.
    ///
    /// This decoder will emit the underlying chunks as-is.
    #[inline]
    fn plain_text(body: Body) -> Decoder {
        Decoder {
            inner: Inner::PlainText(body)
        }
    }

    /// A gzip decoder.
    ///
    /// This decoder will buffer and decompress chunks that are gzipped.
    #[inline]
    fn gzip(body: Body) -> Decoder {
        Decoder {
            inner: Inner::Pending(Pending { body: ReadableChunks::new(body) })
        }
    }

    /// Constructs a Decoder from a hyper request.
    ///
    /// A decoder is just a wrapper around the hyper request that knows
    /// how to decode the content body of the request.
    ///
    /// Uses the correct variant by inspecting the Content-Encoding header.
    pub(crate) fn detect(headers: &mut HeaderMap, body: Body, check_gzip: bool) -> Decoder {
        if !check_gzip {
            return Decoder::plain_text(body);
        }
        let content_encoding_gzip: bool;
        let mut is_gzip = {
            content_encoding_gzip = headers
                .get_all(CONTENT_ENCODING)
                .iter()
                .fold(false, |acc, enc| acc || enc == "gzip");
            content_encoding_gzip ||
            headers
                .get_all(TRANSFER_ENCODING)
                .iter()
                .fold(false, |acc, enc| acc || enc == "gzip")
        };
        if is_gzip {
            if let Some(content_length) = headers.get(CONTENT_LENGTH) {
                if content_length == "0" {
                    warn!("gzip response with content-length of 0");
                    is_gzip = false;
                }
            }
        }
        if content_encoding_gzip {
            headers.remove(CONTENT_ENCODING);
            headers.remove(CONTENT_LENGTH);
        }
        if is_gzip {
            Decoder::gzip(body)
        } else {
            Decoder::plain_text(body)
        }
    }

    pub(crate) fn content_length(&self) -> Option<u64> {
        match self.inner {
            Inner::PlainText(ref body) => body.content_length(),
            _ => None,
        }
    }
}

impl Stream for Decoder {
    type Item = Chunk;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // Do a read or poll for a pendidng decoder value.
        let new_value = match self.inner {
            Inner::Pending(ref mut future) => {
                match future.poll() {
                    Ok(Async::Ready(inner)) => inner,
                    Ok(Async::NotReady) => return Ok(Async::NotReady),
                    Err(e) => return Err(e)
                }
            },
            Inner::PlainText(ref mut body) => return body.poll(),
            Inner::Gzip(ref mut decoder) => return decoder.poll()
        };

        self.inner = new_value;
        self.poll()
    }
}

impl Future for Pending {
    type Item = Inner;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let body_state = match self.body.poll_stream() {
            Ok(Async::Ready(state)) => state,
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => return Err(e)
        };

        let body = mem::replace(&mut self.body, ReadableChunks::new(Body::empty()));
        match body_state {
            StreamState::Eof => Ok(Async::Ready(Inner::PlainText(Body::empty()))),
            StreamState::HasMore => Ok(Async::Ready(Inner::Gzip(Gzip::new(body))))
        }
    }
}

impl Gzip {
    fn new(stream: ReadableChunks<Body>) -> Self {
        Gzip {
            buf: BytesMut::with_capacity(INIT_BUFFER_SIZE),
            inner: Box::new(GzDecoder::new(stream)),
        }
    }
}

impl Stream for Gzip {
    type Item = Chunk;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.buf.remaining_mut() == 0 {
            self.buf.reserve(INIT_BUFFER_SIZE);
        }

        // The buffer contains uninitialised memory so getting a readable slice is unsafe.
        // We trust the `flate2` and `miniz` writer not to read from the memory given.
        //
        // To be safe, this memory could be zeroed before passing to `flate2`.
        // Otherwise we might need to deal with the case where `flate2` panics.
        let read = {
            let mut buf = unsafe { self.buf.bytes_mut() };
            self.inner.read(&mut buf)
        };

        match read {
            Ok(read) if read == 0 => {
                Ok(Async::Ready(None))
            },
            Ok(read) => {
                unsafe { self.buf.advance_mut(read) };
                let chunk = Chunk::from_chunk(self.buf.split_to(read).freeze());

                Ok(Async::Ready(Some(chunk)))
            },
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                Ok(Async::NotReady)
            },
            Err(e) => Err(error::from(e))
        }
    }
}

/// A `Read`able wrapper over a stream of chunks.
pub struct ReadableChunks<S> {
    state: ReadState,
    stream: S,
}

enum ReadState {
    /// A chunk is ready to be read from.
    Ready(Chunk),
    /// The next chunk isn't ready yet.
    NotReady,
    /// The stream has finished.
    Eof,
}

enum StreamState {
    /// More bytes can be read from the stream.
    HasMore,
    /// No more bytes can be read from the stream.
    Eof
}

impl<S> ReadableChunks<S> {
    #[inline]
    pub(crate) fn new(stream: S) -> Self {
        ReadableChunks {
            state: ReadState::NotReady,
            stream: stream,
        }
    }
}

impl<S> fmt::Debug for ReadableChunks<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("ReadableChunks")
            .finish()
    }
}

impl<S> Read for ReadableChunks<S>
where
    S: Stream<Item = Chunk, Error = error::Error>,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        loop {
            let ret;
            match self.state {
                ReadState::Ready(ref mut chunk) => {
                    let len = cmp::min(buf.len(), chunk.remaining());

                    buf[..len].copy_from_slice(&chunk[..len]);
                    chunk.advance(len);
                    if chunk.is_empty() {
                        ret = len;
                    } else {
                        return Ok(len);
                    }
                },
                ReadState::NotReady => {
                    match self.poll_stream() {
                        Ok(Async::Ready(StreamState::HasMore)) => continue,
                        Ok(Async::Ready(StreamState::Eof)) => {
                            return Ok(0)
                        },
                        Ok(Async::NotReady) => {
                            return Err(io::ErrorKind::WouldBlock.into())
                        },
                        Err(e) => {
                            return Err(error::into_io(e))
                        }
                    }
                },
                ReadState::Eof => return Ok(0),
            }
            self.state = ReadState::NotReady;
            return Ok(ret);
        }
    }
}

impl<S> ReadableChunks<S>
    where S: Stream<Item = Chunk, Error = error::Error>
{
    /// Poll the readiness of the inner reader.
    ///
    /// This function will update the internal state and return a simplified
    /// version of the `ReadState`.
    fn poll_stream(&mut self) -> Poll<StreamState, error::Error> {
        match self.stream.poll() {
            Ok(Async::Ready(Some(chunk))) => {
                self.state = ReadState::Ready(chunk);

                Ok(Async::Ready(StreamState::HasMore))
            },
            Ok(Async::Ready(None)) => {
                self.state = ReadState::Eof;

                Ok(Async::Ready(StreamState::Eof))
            },
            Ok(Async::NotReady) => {
                Ok(Async::NotReady)
            },
            Err(e) => Err(e)
        }
    }
}
<commit_msg>fix gzip + chunked encoding reuse of connection (#509)<commit_after>/*!
A potentially non-blocking response decoder.

The decoder wraps a stream of chunks and produces a new stream of decompressed chunks.
The decompressed chunks aren't guaranteed to align to the compressed ones.

If the response is plaintext then no additional work is carried out.
Chunks are just passed along.

If the response is gzip, then the chunks are decompressed into a buffer.
Slices of that buffer are emitted as new chunks.

This module consists of a few main types:

- `ReadableChunks` is a `Read`-like wrapper around a stream
- `Decoder` is a layer over `ReadableChunks` that applies the right decompression

The following types directly support the gzip compression case:

- `Pending` is a non-blocking constructor for a `Decoder` in case the body needs to be checked for EOF
*/

use std::fmt;
use std::mem;
use std::cmp;
use std::io::{self, Read};

use bytes::{Buf, BufMut, BytesMut};
use flate2::read::GzDecoder;
use futures::{Async, Future, Poll, Stream};
use hyper::{HeaderMap};
use hyper::header::{CONTENT_ENCODING, CONTENT_LENGTH, TRANSFER_ENCODING};

use super::{Body, Chunk};
use error;

const INIT_BUFFER_SIZE: usize = 8192;

/// A response decompressor over a non-blocking stream of chunks.
///
/// The inner decoder may be constructed asynchronously.
pub struct Decoder {
    inner: Inner
}

enum Inner {
    /// A `PlainText` decoder just returns the response content as is.
    PlainText(Body),
    /// A `Gzip` decoder will uncompress the gzipped response content before returning it.
    Gzip(Gzip),
    /// A decoder that doesn't have a value yet.
    Pending(Pending)
}

/// A future attempt to poll the response body for EOF so we know whether to use gzip or not.
struct Pending {
    body: ReadableChunks<Body>,
}

/// A gzip decoder that reads from a `flate2::read::GzDecoder` into a `BytesMut` and emits the results
/// as a `Chunk`.
struct Gzip {
    inner: Box<GzDecoder<ReadableChunks<Body>>>,
    buf: BytesMut,
}

impl fmt::Debug for Decoder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Decoder")
            .finish()
    }
}

impl Decoder {
    /// An empty decoder.
    ///
    /// This decoder will produce a single 0 byte chunk.
    #[inline]
    pub fn empty() -> Decoder {
        Decoder {
            inner: Inner::PlainText(Body::empty())
        }
    }

    /// A plain text decoder.
    ///
    /// This decoder will emit the underlying chunks as-is.
    #[inline]
    fn plain_text(body: Body) -> Decoder {
        Decoder {
            inner: Inner::PlainText(body)
        }
    }

    /// A gzip decoder.
    ///
    /// This decoder will buffer and decompress chunks that are gzipped.
    #[inline]
    fn gzip(body: Body) -> Decoder {
        Decoder {
            inner: Inner::Pending(Pending { body: ReadableChunks::new(body) })
        }
    }

    /// Constructs a Decoder from a hyper request.
    ///
    /// A decoder is just a wrapper around the hyper request that knows
    /// how to decode the content body of the request.
    ///
    /// Uses the correct variant by inspecting the Content-Encoding header.
    pub(crate) fn detect(headers: &mut HeaderMap, body: Body, check_gzip: bool) -> Decoder {
        if !check_gzip {
            return Decoder::plain_text(body);
        }
        let content_encoding_gzip: bool;
        let mut is_gzip = {
            content_encoding_gzip = headers
                .get_all(CONTENT_ENCODING)
                .iter()
                .fold(false, |acc, enc| acc || enc == "gzip");
            content_encoding_gzip ||
            headers
                .get_all(TRANSFER_ENCODING)
                .iter()
                .fold(false, |acc, enc| acc || enc == "gzip")
        };
        if is_gzip {
            if let Some(content_length) = headers.get(CONTENT_LENGTH) {
                if content_length == "0" {
                    warn!("gzip response with content-length of 0");
                    is_gzip = false;
                }
            }
        }
        if content_encoding_gzip {
            headers.remove(CONTENT_ENCODING);
            headers.remove(CONTENT_LENGTH);
        }
        if is_gzip {
            Decoder::gzip(body)
        } else {
            Decoder::plain_text(body)
        }
    }

    pub(crate) fn content_length(&self) -> Option<u64> {
        match self.inner {
            Inner::PlainText(ref body) => body.content_length(),
            _ => None,
        }
    }
}

impl Stream for Decoder {
    type Item = Chunk;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // Do a read or poll for a pendidng decoder value.
        let new_value = match self.inner {
            Inner::Pending(ref mut future) => {
                match future.poll() {
                    Ok(Async::Ready(inner)) => inner,
                    Ok(Async::NotReady) => return Ok(Async::NotReady),
                    Err(e) => return Err(e)
                }
            },
            Inner::PlainText(ref mut body) => return body.poll(),
            Inner::Gzip(ref mut decoder) => return decoder.poll()
        };

        self.inner = new_value;
        self.poll()
    }
}

impl Future for Pending {
    type Item = Inner;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let body_state = match self.body.poll_stream() {
            Ok(Async::Ready(state)) => state,
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => return Err(e)
        };

        let body = mem::replace(&mut self.body, ReadableChunks::new(Body::empty()));
        match body_state {
            StreamState::Eof => Ok(Async::Ready(Inner::PlainText(Body::empty()))),
            StreamState::HasMore => Ok(Async::Ready(Inner::Gzip(Gzip::new(body))))
        }
    }
}

impl Gzip {
    fn new(stream: ReadableChunks<Body>) -> Self {
        Gzip {
            buf: BytesMut::with_capacity(INIT_BUFFER_SIZE),
            inner: Box::new(GzDecoder::new(stream)),
        }
    }
}

impl Stream for Gzip {
    type Item = Chunk;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.buf.remaining_mut() == 0 {
            self.buf.reserve(INIT_BUFFER_SIZE);
        }

        // The buffer contains uninitialised memory so getting a readable slice is unsafe.
        // We trust the `flate2` and `miniz` writer not to read from the memory given.
        //
        // To be safe, this memory could be zeroed before passing to `flate2`.
        // Otherwise we might need to deal with the case where `flate2` panics.
        let read = {
            let mut buf = unsafe { self.buf.bytes_mut() };
            self.inner.read(&mut buf)
        };

        match read {
            Ok(read) if read == 0 => match self.inner.get_mut().read(&mut [0]) {
                Ok(0) => Ok(Async::Ready(None)),
                Ok(_) => Err(error::from(io::Error::new(io::ErrorKind::InvalidData, "Unexpected Data"))),
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(Async::NotReady),
                Err(e) => Err(error::from(e))
            },
            Ok(read) => {
                unsafe { self.buf.advance_mut(read) };
                let chunk = Chunk::from_chunk(self.buf.split_to(read).freeze());

                Ok(Async::Ready(Some(chunk)))
            },
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                Ok(Async::NotReady)
            },
            Err(e) => Err(error::from(e))
        }
    }
}

/// A `Read`able wrapper over a stream of chunks.
pub struct ReadableChunks<S> {
    state: ReadState,
    stream: S,
}

enum ReadState {
    /// A chunk is ready to be read from.
    Ready(Chunk),
    /// The next chunk isn't ready yet.
    NotReady,
    /// The stream has finished.
    Eof,
}

enum StreamState {
    /// More bytes can be read from the stream.
    HasMore,
    /// No more bytes can be read from the stream.
    Eof
}

impl<S> ReadableChunks<S> {
    #[inline]
    pub(crate) fn new(stream: S) -> Self {
        ReadableChunks {
            state: ReadState::NotReady,
            stream: stream,
        }
    }
}

impl<S> fmt::Debug for ReadableChunks<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("ReadableChunks")
            .finish()
    }
}

impl<S> Read for ReadableChunks<S>
where
    S: Stream<Item = Chunk, Error = error::Error>,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        loop {
            let ret;
            match self.state {
                ReadState::Ready(ref mut chunk) => {
                    let len = cmp::min(buf.len(), chunk.remaining());

                    buf[..len].copy_from_slice(&chunk[..len]);
                    chunk.advance(len);
                    if chunk.is_empty() {
                        ret = len;
                    } else {
                        return Ok(len);
                    }
                },
                ReadState::NotReady => {
                    match self.poll_stream() {
                        Ok(Async::Ready(StreamState::HasMore)) => continue,
                        Ok(Async::Ready(StreamState::Eof)) => {
                            return Ok(0)
                        },
                        Ok(Async::NotReady) => {
                            return Err(io::ErrorKind::WouldBlock.into())
                        },
                        Err(e) => {
                            return Err(error::into_io(e))
                        }
                    }
                },
                ReadState::Eof => return Ok(0),
            }
            self.state = ReadState::NotReady;
            return Ok(ret);
        }
    }
}

impl<S> ReadableChunks<S>
    where S: Stream<Item = Chunk, Error = error::Error>
{
    /// Poll the readiness of the inner reader.
    ///
    /// This function will update the internal state and return a simplified
    /// version of the `ReadState`.
    fn poll_stream(&mut self) -> Poll<StreamState, error::Error> {
        match self.stream.poll() {
            Ok(Async::Ready(Some(chunk))) => {
                self.state = ReadState::Ready(chunk);

                Ok(Async::Ready(StreamState::HasMore))
            },
            Ok(Async::Ready(None)) => {
                self.state = ReadState::Eof;

                Ok(Async::Ready(StreamState::Eof))
            },
            Ok(Async::NotReady) => {
                Ok(Async::NotReady)
            },
            Err(e) => Err(e)
        }
    }
}
<|endoftext|>


<commit_before>use super::{Open, Sink};
use std::io;
use alsa::{PCM, Stream, Mode, Format, Access};

pub struct AlsaSink(Option<PCM>, String);

impl Open for AlsaSink {
   fn open(device: Option<String>) -> AlsaSink {
        info!("Using alsa sink");

        let name = device.unwrap_or("default".to_string());

        AlsaSink(None, name)
    }
}

impl Sink for AlsaSink {
    fn start(&mut self) -> io::Result<()> {
        if self.0.is_some() {
        } else {
            match PCM::open(&*self.1,
                                    Stream::Playback, Mode::Blocking,
                                    Format::Signed16, Access::Interleaved,
                                    2, 44100) {
                Ok(f) => self.0 = Some(f),
                Err(e) => {
                    self.0 = None;
                    error!("Alsa error PCM open {}", e);
                    return Err(io::Error::new(io::ErrorKind::Other, "Alsa error: PCM open failed"));
                }
            }
        }
        Ok(())
    }

    fn stop(&mut self) -> io::Result<()> {
        self.0 = None;
        Ok(())
    }

    fn write(&mut self, data: &[i16]) -> io::Result<()> {
        self.0.as_mut().unwrap().write_interleaved(&data).unwrap();
        Ok(())
    }
}
<commit_msg>minor code cleanup<commit_after>use super::{Open, Sink};
use std::io;
use alsa::{PCM, Stream, Mode, Format, Access};

pub struct AlsaSink(Option<PCM>, String);

impl Open for AlsaSink {
   fn open(device: Option<String>) -> AlsaSink {
        info!("Using alsa sink");

        let name = device.unwrap_or("default".to_string());

        AlsaSink(None, name)
    }
}

impl Sink for AlsaSink {
    fn start(&mut self) -> io::Result<()> {
        if self.0.is_none() {
            match PCM::open(&*self.1,
                                    Stream::Playback, Mode::Blocking,
                                    Format::Signed16, Access::Interleaved,
                                    2, 44100) {
                Ok(f) => self.0 = Some(f),
                Err(e) => {
                    error!("Alsa error PCM open {}", e);
                    return Err(io::Error::new(io::ErrorKind::Other, "Alsa error: PCM open failed"));
                }
            }
        }
        Ok(())
    }

    fn stop(&mut self) -> io::Result<()> {
        self.0 = None;
        Ok(())
    }

    fn write(&mut self, data: &[i16]) -> io::Result<()> {
        self.0.as_mut().unwrap().write_interleaved(&data).unwrap();
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Implements all functionality related to encoding header blocks using
//! HPACK.
//!
//! Clients should use the `Encoder` struct as the API for performing HPACK
//! encoding.
//!
//! # Examples
//!
//! Encodes a header using a literal encoding.
//!
//! ```rust
//! use solicit::hpack::Encoder;
//!
//! let mut encoder = Encoder::new();
//!
//! let headers = vec![
//!     (b"custom-key".to_vec(), b"custom-value".to_vec()),
//! ];
//! // First encoding...
//! let result = encoder.encode(&headers);
//! // The result is a literal encoding of the header name and value, with an
//! // initial byte representing the type of the encoding
//! // (incremental indexing).
//! assert_eq!(
//!     vec![0x40,
//!          10, b'c', b'u', b's', b't', b'o', b'm', b'-', b'k', b'e', b'y',
//!          12, b'c', b'u', b's', b't', b'o', b'm', b'-', b'v', b'a', b'l',
//!          b'u', b'e'],
//!     result);
//! ```
use super::STATIC_TABLE;
use super::HeaderTable;

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
pub fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// Represents an HPACK encoder. Allows clients to encode arbitrary header sets
/// and tracks the encoding context. That is, encoding subsequent header sets
/// will use the context built by previous encode calls.
///
/// This is the main API for performing HPACK encoding of headers.
///
/// # Examples
///
/// Encoding a header two times in a row produces two different
/// representations, due to the utilization of HPACK compression.
///
/// ```rust
/// use solicit::hpack::Encoder;
///
/// let mut encoder = Encoder::new();
///
/// let headers = vec![
///     (b"custom-key".to_vec(), b"custom-value".to_vec()),
/// ];
/// // First encoding...
/// let result = encoder.encode(&headers);
/// // The result is a literal encoding of the header name and value, with an
/// // initial byte representing the type of the encoding
/// // (incremental indexing).
/// assert_eq!(
///     vec![0x40,
///          10, b'c', b'u', b's', b't', b'o', b'm', b'-', b'k', b'e', b'y',
///          12, b'c', b'u', b's', b't', b'o', b'm', b'-', b'v', b'a', b'l',
///          b'u', b'e'],
///     result);
///
/// // Encode the same headers again!
/// let result = encoder.encode(&headers);
/// // The result is simply the index of the header in the header table (62),
/// // with a flag representing that the decoder should use the index.
/// assert_eq!(vec![0x80 | 62], result);
/// ```
pub struct Encoder<'a> {
    /// The header table represents the encoder's context
    header_table: HeaderTable<'a>,
}

impl<'a> Encoder<'a> {
    /// Creates a new `Encoder` with a default static table, as defined by the
    /// HPACK spec (Appendix A).
    pub fn new() -> Encoder<'a> {
        Encoder {
            header_table: HeaderTable::with_static_table(STATIC_TABLE),
        }
    }

    /// Encodes the given headers using the HPACK rules and returns a newly
    /// allocated `Vec` containing the bytes representing the encoded header
    /// set.
    ///
    /// The encoder so far supports only a single, extremely simple encoding
    /// strategy, whereby each header is represented as an indexed header if
    /// already found in the header table and a literal otherwise. When a
    /// header isn't found in the table, it is added if the header name wasn't
    /// found either (i.e. there are never two header names with different
    /// values in the produced header table). Strings are always encoded as
    /// literals (Huffman encoding is not used).
    pub fn encode(&mut self, headers: &Vec<(Vec<u8>, Vec<u8>)>) -> Vec<u8> {
        let mut encoded: Vec<u8> = Vec::new();

        for header in headers.iter() {
            match self.header_table.find_header((&header.0, &header.1)) {
                None => {
                    // The name of the header is in no tables: need to encode
                    // it with both a literal name and value.
                    self.encode_literal(header, true, &mut encoded);
                    self.header_table.add_header(header.0.clone(), header.1.clone());
                },
                Some((index, false)) => {
                    // The name of the header is at the given index, but the
                    // value does not match the current one: need to encode
                    // only the value as a literal.
                    self.encode_literal(header, false, &mut encoded);
                },
                Some((index, true)) => {
                    // The full header was found in one of the tables, so we
                    // just encode the index.
                    self.encode_indexed(index, &mut encoded);
                }
            };
        }

        encoded
    }

    /// Encodes a header as a literal (i.e. both the name and the value are
    /// encoded as a string literal) and places the result in the given buffer
    /// `buf`.
    ///
    /// # Parameters
    ///
    /// - `header` - the header to be encoded
    /// - `should_index` - indicates whether the given header should be indexed, i.e.
    ///                    inserted into the dynamic table
    /// - `buf` - The buffer into which the result is placed
    ///
    fn encode_literal(&mut self,
                      header: &(Vec<u8>, Vec<u8>),
                      should_index: bool,
                      buf: &mut Vec<u8>) {
        let mask = if should_index {
            0x40
        } else {
            0x0
        };

        buf.push(mask);
        self.encode_string_literal(&header.0, buf);
        self.encode_string_literal(&header.1, buf);
    }

    /// Encodes a string literal and places the result in the given buffer
    /// `buf`.
    ///
    /// The function does not consider Huffman encoding for now, but always
    /// produces a string literal representations, according to the HPACK spec
    /// section 5.2.
    fn encode_string_literal(&mut self, octet_str: &[u8], buf: &mut Vec<u8>) {
        buf.extend(encode_integer(octet_str.len(), 7).into_iter());
        buf.extend(octet_str.to_vec().into_iter());
    }

    /// Encodes a header whose name is indexed and places the result in the
    /// given buffer `buf`.
    fn encode_indexed_name(&mut self, header: (usize, &Vec<u8>), should_index: bool, buf: &mut Vec<u8>) {
        let (mask, prefix) = if should_index {
            (0x40, 6)
        } else {
            (0x0, 4)
        };

        let mut encoded_index = encode_integer(header.0, prefix);
        encoded_index[0] |= mask;
        buf.extend(encoded_index.into_iter());
        // So far, we rely on just one strategy for encoding string literals.
        self.encode_string_literal(&header.1, buf);
    }

    /// Encodes an indexed header (a header that is fully in the header table)
    /// and places the result in the given buffer `buf`.
    ///
    /// The encoding is according to the rules of the HPACK spec, section 6.1.
    fn encode_indexed(&self, index: usize, buf: &mut Vec<u8>) {
        let mut encoded = encode_integer(index, 7);
        // We need to set the most significant bit, since the bit-pattern is
        // `1xxxxxxx` for indexed headers.
        encoded[0] |= 0x80;

        buf.extend(encoded.into_iter());
    }
}

#[cfg(test)]
mod tests {
    use super::encode_integer;
    use super::Encoder;

    use super::super::Decoder;

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }

    /// A helper function that checks whether the given buffer can be decoded
    /// into a set of headers that corresponds to the given `headers` list.
    /// Relies on using the `solicit::hpack::decoder::Decoder`` struct for
    /// performing the decoding.
    ///
    /// # Returns
    ///
    /// A `bool` indicating whether such a decoding can be performed.
    fn is_decodable(buf: &Vec<u8>, headers: &Vec<(Vec<u8>, Vec<u8>)>) -> bool {
        let mut decoder = Decoder::new();
        match decoder.decode(buf).ok() {
            Some(h) => h == *headers,
            None => false,
        }
    }

    /// Tests that encoding only the `:method` header works.
    #[test]
    fn test_encode_only_method() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b":method".to_vec(), b"GET".to_vec()),
        ];

        let result = encoder.encode(&headers);

        debug!("{:?}", result);
        assert!(is_decodable(&result, &headers));
    }

    /// Tests that when a single custom header is sent it gets indexed by the
    /// coder.
    #[test]
    fn test_custom_header_gets_indexed() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
        ];

        let result = encoder.encode(&headers);
        assert!(is_decodable(&result, &headers));
        // The header is in the encoder's dynamic table.
        assert_eq!(encoder.header_table.dynamic_table.to_vec(), headers);
        // ...but also indicated as such in the output.
        assert!(0x40 == (0x40 & result[0]));
        debug!("{:?}", result);
    }

    /// Tests that when a header gets added to the dynamic table, the encoder
    /// will use the index, instead of the literal representation on the next
    /// encoding of the same header.
    #[test]
    fn test_uses_index_on_second_iteration() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
        ];
        // First encoding...
        let _ = encoder.encode(&headers);

        // Encode the same headers again!
        let result = encoder.encode(&headers);

        // The header is in the encoder's dynamic table.
        assert_eq!(encoder.header_table.dynamic_table.to_vec(), headers);
        // The output is a single index byte?
        assert_eq!(result.len(), 1);
        // The index is correctly encoded:
        // - The most significant bit is set
        assert_eq!(0x80 & result[0], 0x80);
        // - The other 7 bits decode to an integer giving the index in the full
        //   header address space.
        assert_eq!(result[0] ^ 0x80, 62);
        // assert_eq!(encoder.header_table.get_from_table(62).unwrap(), headers[0]);
    }

    /// Tests that multiple headers are correctly encoded (i.e. can be decoded
    /// back to their original representation).
    #[test]
    fn test_multiple_headers_encoded() {
        let mut encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
            (b":method".to_vec(), b"GET".to_vec()),
            (b":path".to_vec(), b"/some/path".to_vec()),
        ];

        let result = encoder.encode(&headers);

        assert!(is_decodable(&result, &headers));
    }
}
<commit_msg>hpack: Improve a unit test<commit_after>//! Implements all functionality related to encoding header blocks using
//! HPACK.
//!
//! Clients should use the `Encoder` struct as the API for performing HPACK
//! encoding.
//!
//! # Examples
//!
//! Encodes a header using a literal encoding.
//!
//! ```rust
//! use solicit::hpack::Encoder;
//!
//! let mut encoder = Encoder::new();
//!
//! let headers = vec![
//!     (b"custom-key".to_vec(), b"custom-value".to_vec()),
//! ];
//! // First encoding...
//! let result = encoder.encode(&headers);
//! // The result is a literal encoding of the header name and value, with an
//! // initial byte representing the type of the encoding
//! // (incremental indexing).
//! assert_eq!(
//!     vec![0x40,
//!          10, b'c', b'u', b's', b't', b'o', b'm', b'-', b'k', b'e', b'y',
//!          12, b'c', b'u', b's', b't', b'o', b'm', b'-', b'v', b'a', b'l',
//!          b'u', b'e'],
//!     result);
//! ```
use super::STATIC_TABLE;
use super::HeaderTable;

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
pub fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// Represents an HPACK encoder. Allows clients to encode arbitrary header sets
/// and tracks the encoding context. That is, encoding subsequent header sets
/// will use the context built by previous encode calls.
///
/// This is the main API for performing HPACK encoding of headers.
///
/// # Examples
///
/// Encoding a header two times in a row produces two different
/// representations, due to the utilization of HPACK compression.
///
/// ```rust
/// use solicit::hpack::Encoder;
///
/// let mut encoder = Encoder::new();
///
/// let headers = vec![
///     (b"custom-key".to_vec(), b"custom-value".to_vec()),
/// ];
/// // First encoding...
/// let result = encoder.encode(&headers);
/// // The result is a literal encoding of the header name and value, with an
/// // initial byte representing the type of the encoding
/// // (incremental indexing).
/// assert_eq!(
///     vec![0x40,
///          10, b'c', b'u', b's', b't', b'o', b'm', b'-', b'k', b'e', b'y',
///          12, b'c', b'u', b's', b't', b'o', b'm', b'-', b'v', b'a', b'l',
///          b'u', b'e'],
///     result);
///
/// // Encode the same headers again!
/// let result = encoder.encode(&headers);
/// // The result is simply the index of the header in the header table (62),
/// // with a flag representing that the decoder should use the index.
/// assert_eq!(vec![0x80 | 62], result);
/// ```
pub struct Encoder<'a> {
    /// The header table represents the encoder's context
    header_table: HeaderTable<'a>,
}

impl<'a> Encoder<'a> {
    /// Creates a new `Encoder` with a default static table, as defined by the
    /// HPACK spec (Appendix A).
    pub fn new() -> Encoder<'a> {
        Encoder {
            header_table: HeaderTable::with_static_table(STATIC_TABLE),
        }
    }

    /// Encodes the given headers using the HPACK rules and returns a newly
    /// allocated `Vec` containing the bytes representing the encoded header
    /// set.
    ///
    /// The encoder so far supports only a single, extremely simple encoding
    /// strategy, whereby each header is represented as an indexed header if
    /// already found in the header table and a literal otherwise. When a
    /// header isn't found in the table, it is added if the header name wasn't
    /// found either (i.e. there are never two header names with different
    /// values in the produced header table). Strings are always encoded as
    /// literals (Huffman encoding is not used).
    pub fn encode(&mut self, headers: &Vec<(Vec<u8>, Vec<u8>)>) -> Vec<u8> {
        let mut encoded: Vec<u8> = Vec::new();

        for header in headers.iter() {
            match self.header_table.find_header((&header.0, &header.1)) {
                None => {
                    // The name of the header is in no tables: need to encode
                    // it with both a literal name and value.
                    self.encode_literal(header, true, &mut encoded);
                    self.header_table.add_header(header.0.clone(), header.1.clone());
                },
                Some((index, false)) => {
                    // The name of the header is at the given index, but the
                    // value does not match the current one: need to encode
                    // only the value as a literal.
                    self.encode_literal(header, false, &mut encoded);
                },
                Some((index, true)) => {
                    // The full header was found in one of the tables, so we
                    // just encode the index.
                    self.encode_indexed(index, &mut encoded);
                }
            };
        }

        encoded
    }

    /// Encodes a header as a literal (i.e. both the name and the value are
    /// encoded as a string literal) and places the result in the given buffer
    /// `buf`.
    ///
    /// # Parameters
    ///
    /// - `header` - the header to be encoded
    /// - `should_index` - indicates whether the given header should be indexed, i.e.
    ///                    inserted into the dynamic table
    /// - `buf` - The buffer into which the result is placed
    ///
    fn encode_literal(&mut self,
                      header: &(Vec<u8>, Vec<u8>),
                      should_index: bool,
                      buf: &mut Vec<u8>) {
        let mask = if should_index {
            0x40
        } else {
            0x0
        };

        buf.push(mask);
        self.encode_string_literal(&header.0, buf);
        self.encode_string_literal(&header.1, buf);
    }

    /// Encodes a string literal and places the result in the given buffer
    /// `buf`.
    ///
    /// The function does not consider Huffman encoding for now, but always
    /// produces a string literal representations, according to the HPACK spec
    /// section 5.2.
    fn encode_string_literal(&mut self, octet_str: &[u8], buf: &mut Vec<u8>) {
        buf.extend(encode_integer(octet_str.len(), 7).into_iter());
        buf.extend(octet_str.to_vec().into_iter());
    }

    /// Encodes a header whose name is indexed and places the result in the
    /// given buffer `buf`.
    fn encode_indexed_name(&mut self, header: (usize, &Vec<u8>), should_index: bool, buf: &mut Vec<u8>) {
        let (mask, prefix) = if should_index {
            (0x40, 6)
        } else {
            (0x0, 4)
        };

        let mut encoded_index = encode_integer(header.0, prefix);
        encoded_index[0] |= mask;
        buf.extend(encoded_index.into_iter());
        // So far, we rely on just one strategy for encoding string literals.
        self.encode_string_literal(&header.1, buf);
    }

    /// Encodes an indexed header (a header that is fully in the header table)
    /// and places the result in the given buffer `buf`.
    ///
    /// The encoding is according to the rules of the HPACK spec, section 6.1.
    fn encode_indexed(&self, index: usize, buf: &mut Vec<u8>) {
        let mut encoded = encode_integer(index, 7);
        // We need to set the most significant bit, since the bit-pattern is
        // `1xxxxxxx` for indexed headers.
        encoded[0] |= 0x80;

        buf.extend(encoded.into_iter());
    }
}

#[cfg(test)]
mod tests {
    use super::encode_integer;
    use super::Encoder;

    use super::super::Decoder;

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }

    /// A helper function that checks whether the given buffer can be decoded
    /// into a set of headers that corresponds to the given `headers` list.
    /// Relies on using the `solicit::hpack::decoder::Decoder`` struct for
    /// performing the decoding.
    ///
    /// # Returns
    ///
    /// A `bool` indicating whether such a decoding can be performed.
    fn is_decodable(buf: &Vec<u8>, headers: &Vec<(Vec<u8>, Vec<u8>)>) -> bool {
        let mut decoder = Decoder::new();
        match decoder.decode(buf).ok() {
            Some(h) => h == *headers,
            None => false,
        }
    }

    /// Tests that encoding only the `:method` header works.
    #[test]
    fn test_encode_only_method() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b":method".to_vec(), b"GET".to_vec()),
        ];

        let result = encoder.encode(&headers);

        debug!("{:?}", result);
        assert!(is_decodable(&result, &headers));
    }

    /// Tests that when a single custom header is sent it gets indexed by the
    /// coder.
    #[test]
    fn test_custom_header_gets_indexed() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
        ];

        let result = encoder.encode(&headers);
        assert!(is_decodable(&result, &headers));
        // The header is in the encoder's dynamic table.
        assert_eq!(encoder.header_table.dynamic_table.to_vec(), headers);
        // ...but also indicated as such in the output.
        assert!(0x40 == (0x40 & result[0]));
        debug!("{:?}", result);
    }

    /// Tests that when a header gets added to the dynamic table, the encoder
    /// will use the index, instead of the literal representation on the next
    /// encoding of the same header.
    #[test]
    fn test_uses_index_on_second_iteration() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
        ];
        // First encoding...
        let _ = encoder.encode(&headers);

        // Encode the same headers again!
        let result = encoder.encode(&headers);

        // The header is in the encoder's dynamic table.
        assert_eq!(encoder.header_table.dynamic_table.to_vec(), headers);
        // The output is a single index byte?
        assert_eq!(result.len(), 1);
        // The index is correctly encoded:
        // - The most significant bit is set
        assert_eq!(0x80 & result[0], 0x80);
        // - The other 7 bits decode to an integer giving the index in the full
        //   header address space.
        assert_eq!(result[0] ^ 0x80, 62);
        // The header table actually contains the header at that index?
        assert_eq!(
            encoder.header_table.get_from_table(62).unwrap(),
            (&headers[0].0[..], &headers[0].1[..]));
    }

    /// Tests that multiple headers are correctly encoded (i.e. can be decoded
    /// back to their original representation).
    #[test]
    fn test_multiple_headers_encoded() {
        let mut encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
            (b":method".to_vec(), b"GET".to_vec()),
            (b":path".to_vec(), b"/some/path".to_vec()),
        ];

        let result = encoder.encode(&headers);

        assert!(is_decodable(&result, &headers));
    }
}
<|endoftext|>


<commit_before>/// A module exposing utilities for encoding and decoding Huffman-coded octet
/// strings, under the Huffman code defined by HPACK.
/// (HPACK-draft-10, Appendix B)

use std::collections::HashMap;

/// A helper struct that represents an iterator over individual bits of all
/// bytes found in a wrapped Iterator over bytes.
/// Bits are represented as `bool`s, where `true` corresponds to a set bit and
/// `false` to a 0 bit.
///
/// Bits are yielded in order of significance, starting from the
/// most-significant bit.
struct BitIterator<'a, I: Iterator> {
    buffer_iterator: I,
    current_byte: Option<&'a u8>,
    /// The bit-position within the current byte
    pos: u8,
}

impl<'a, I: Iterator> BitIterator<'a, I>
        where I: Iterator<Item=&'a u8> {
    pub fn new(iterator: I) -> BitIterator<'a, I> {
        BitIterator::<'a, I> {
            buffer_iterator: iterator,
            current_byte: None,
            pos: 7,
        }
    }
}

impl<'a, I> Iterator for BitIterator<'a, I>
        where I: Iterator<Item=&'a u8> {
    type Item = bool;

    fn next(&mut self) -> Option<bool> {
        if self.current_byte.is_none() {
            self.current_byte = self.buffer_iterator.next();
            self.pos = 7;
        }

        // If we still have `None`, it means the buffer has been exhausted
        if self.current_byte.is_none() {
            return None;
        }

        let b = *self.current_byte.unwrap();

        let is_set = (b & (1 << self.pos)) == (1 << self.pos);
        if self.pos == 0 {
            // We have exhausted all bits from the current byte -- try to get
            // a new one on the next pass.
            self.current_byte = None;
        } else {
            // Still more bits left here...
            self.pos -= 1;
        }

        Some(is_set)
    }
}

mod tests {
    use super::BitIterator;

    /// A helper function that converts the given slice containing values `1`
    /// and `0` to a `Vec` of `bool`s, according to the number.
    fn to_expected_bit_result(numbers: &[u8]) -> Vec<bool> {
        numbers.iter().map(|b| -> bool {
            *b == 1
        }).collect()
    }

    #[test]
    fn test_bit_iterator_single_byte() {
        let expected_result = to_expected_bit_result(
            &[0, 0, 0, 0, 1, 0, 1, 0]);

        let mut res: Vec<bool> = Vec::new();
        for b in BitIterator::new(vec![10u8].iter()) {
            res.push(b);
        }

        assert_eq!(res, expected_result);
    }

    #[test]
    fn test_bit_iterator_multiple_bytes() {
        let expected_result = to_expected_bit_result(
            &[0, 0, 0, 0, 1, 0, 1, 0,
              1, 1, 1, 1, 1, 1, 1, 1,
              1, 0, 0, 0, 0, 0, 0, 0,
              0, 0, 0, 0, 0, 0, 0, 1,
              0, 0, 0, 0, 0, 0, 0, 0,
              1, 0, 1, 0, 1, 0, 1, 0]);

        let mut res: Vec<bool> = Vec::new();
        for b in BitIterator::new(vec![10u8, 255, 128, 1, 0, 170].iter()) {
            res.push(b);
        }

        assert_eq!(res, expected_result);
    }
}
<commit_msg>Add a Huffman code symbol representation<commit_after>/// A module exposing utilities for encoding and decoding Huffman-coded octet
/// strings, under the Huffman code defined by HPACK.
/// (HPACK-draft-10, Appendix B)

use std::collections::HashMap;

/// Represents a symbol that can be inserted into a Huffman-encoded octet
/// string.
enum HuffmanCodeSymbol {
    /// Any octet is a valid symbol
    Symbol(u8),
    /// A special symbol represents the end of the string
    EndOfString,
}

impl HuffmanCodeSymbol {
    pub fn new(symbol: usize) -> HuffmanCodeSymbol {
        if symbol == 256 {
            HuffmanCodeSymbol::EndOfString
        } else {
            // It is safe to downcast since now we know that the value
            // is in the half-open interval [0, 256)
            HuffmanCodeSymbol::Symbol(symbol as u8)
        }
    }
}

/// A helper struct that represents an iterator over individual bits of all
/// bytes found in a wrapped Iterator over bytes.
/// Bits are represented as `bool`s, where `true` corresponds to a set bit and
/// `false` to a 0 bit.
///
/// Bits are yielded in order of significance, starting from the
/// most-significant bit.
struct BitIterator<'a, I: Iterator> {
    buffer_iterator: I,
    current_byte: Option<&'a u8>,
    /// The bit-position within the current byte
    pos: u8,
}

impl<'a, I: Iterator> BitIterator<'a, I>
        where I: Iterator<Item=&'a u8> {
    pub fn new(iterator: I) -> BitIterator<'a, I> {
        BitIterator::<'a, I> {
            buffer_iterator: iterator,
            current_byte: None,
            pos: 7,
        }
    }
}

impl<'a, I> Iterator for BitIterator<'a, I>
        where I: Iterator<Item=&'a u8> {
    type Item = bool;

    fn next(&mut self) -> Option<bool> {
        if self.current_byte.is_none() {
            self.current_byte = self.buffer_iterator.next();
            self.pos = 7;
        }

        // If we still have `None`, it means the buffer has been exhausted
        if self.current_byte.is_none() {
            return None;
        }

        let b = *self.current_byte.unwrap();

        let is_set = (b & (1 << self.pos)) == (1 << self.pos);
        if self.pos == 0 {
            // We have exhausted all bits from the current byte -- try to get
            // a new one on the next pass.
            self.current_byte = None;
        } else {
            // Still more bits left here...
            self.pos -= 1;
        }

        Some(is_set)
    }
}

mod tests {
    use super::BitIterator;

    /// A helper function that converts the given slice containing values `1`
    /// and `0` to a `Vec` of `bool`s, according to the number.
    fn to_expected_bit_result(numbers: &[u8]) -> Vec<bool> {
        numbers.iter().map(|b| -> bool {
            *b == 1
        }).collect()
    }

    #[test]
    fn test_bit_iterator_single_byte() {
        let expected_result = to_expected_bit_result(
            &[0, 0, 0, 0, 1, 0, 1, 0]);

        let mut res: Vec<bool> = Vec::new();
        for b in BitIterator::new(vec![10u8].iter()) {
            res.push(b);
        }

        assert_eq!(res, expected_result);
    }

    #[test]
    fn test_bit_iterator_multiple_bytes() {
        let expected_result = to_expected_bit_result(
            &[0, 0, 0, 0, 1, 0, 1, 0,
              1, 1, 1, 1, 1, 1, 1, 1,
              1, 0, 0, 0, 0, 0, 0, 0,
              0, 0, 0, 0, 0, 0, 0, 1,
              0, 0, 0, 0, 0, 0, 0, 0,
              1, 0, 1, 0, 1, 0, 1, 0]);

        let mut res: Vec<bool> = Vec::new();
        for b in BitIterator::new(vec![10u8, 255, 128, 1, 0, 170].iter()) {
            res.push(b);
        }

        assert_eq!(res, expected_result);
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Pierre Talbot (IRCAM)

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub use compile_error::*;
pub use syntax::parse::token::InternedString;
pub use syntax::ast::*;

use syntax::codemap::DUMMY_SP;
use std::rc::Rc;

pub struct AttributeValue<T>
{
  value: Option<T>,
  duplicate: DuplicateAttribute,
  span: Span
}

impl<T: Clone> AttributeValue<T>
{
  pub fn new(duplicate: DuplicateAttribute) -> AttributeValue<T>
  {
    AttributeValue {
      value: None,
      duplicate: duplicate,
      span: DUMMY_SP
    }
  }

  pub fn simple() -> AttributeValue<T>
  {
    AttributeValue::new(DuplicateAttribute::simple(Warn))
  }

  pub fn update(self, cx: &ExtCtxt, value: T, span: Span) -> AttributeValue<T>
  {
    let mut this = self;
    if this.value.is_some() {
      this.duplicate.issue(cx, this.span, span);
    } else {
      this.value = Some(value);
      this.span = span;
    }
    this
  }

  pub fn value_or(&self, default: T) -> T
  {
    match self.value {
      None => default,
      Some(ref value) => value.clone()
    }
  }

  pub fn span(&self) -> Span
  {
    self.span.clone()
  }
}

pub struct AttributeInfo
{
  pub name: InternedString,
  pub desc: &'static str,
  pub model: AttributeModel
}

impl AttributeInfo
{
  pub fn new(name: &'static str, desc: &'static str,
    model: AttributeModel) -> AttributeInfo
  {
    AttributeInfo {
      name: InternedString::new(name),
      desc: desc,
      model: model
    }
  }

  pub fn simple(name: &'static str, desc: &'static str) -> AttributeInfo
  {
    AttributeInfo::new(name, desc, NoProperty(AttributeValue::simple()))
  }

  pub fn update(self, model: AttributeModel) -> AttributeInfo
  {
    let mut this = self;
    this.model = model;
    this
  }

  pub fn plain_value<'a>(&'a self) -> &'a AttributeValue<bool>
  {
    match self.model {
      NoProperty(ref v) => v,
      _ => fail!("No plain value for the current attribute.")
    }
  }

  pub fn sub_model<'a>(&'a self) -> &'a AttributeDict
  {
    match self.model {
      SubAttribute(ref dict) => dict,
      _ => fail!("No sub value for the attribute.")
    }
  }
}

pub enum AttributeModel
{
  NoProperty(AttributeValue<bool>),
  KeyValue(AttributeLitModel),
  SubAttribute(AttributeDict)
}

pub enum AttributeLitModel
{
  MLitStr(AttributeValue<(InternedString, StrStyle)>),
  MLitBinary(AttributeValue<Rc<Vec<u8>>>),
  MLitByte(AttributeValue<u8>),
  MLitChar(AttributeValue<char>),
  MLitInt(AttributeValue<(u64, LitIntType)>),
  MLitFloat(AttributeValue<(InternedString, FloatTy)>),
  MLitFloatUnsuffixed(AttributeValue<InternedString>),
  MLitNil(AttributeValue<()>),
  MLitBool(AttributeValue<bool>)
}

pub struct AttributeDict
{
  dict: Vec<AttributeInfo>
}

impl AttributeDict
{
  pub fn new(dict: Vec<AttributeInfo>) -> AttributeDict
  {
    AttributeDict {
      dict: dict
    }
  }

  pub fn move_map(self, f: |AttributeInfo| -> AttributeInfo) -> AttributeDict
  {
    let mut this = self;
    this.dict = this.dict.move_iter().map(f).collect();
    this
  }

  pub fn push(&mut self, attr: AttributeInfo)
  {
    self.dict.push(attr);
  }

  pub fn push_all(&mut self, attrs: Vec<AttributeInfo>)
  {
    self.dict.push_all_move(attrs);
  }

  pub fn get<'a>(&'a self, name: &'static str) -> &'a AttributeInfo
  {
    let interned = InternedString::new(name);
    for info in self.dict.iter() {
      if info.name == interned {
        return info;
      }
    }
    fail!("Try to get an attribute that doesn't exist.")
  }

  pub fn plain_value<'a>(&'a self, name: &'static str) -> &'a AttributeValue<bool>
  {
    self.get(name).plain_value()
  }

  pub fn sub_model<'a>(&'a self, name: &'static str) -> &'a AttributeDict
  {
    self.get(name).sub_model()
  }

  pub fn plain_value_or(&self, name: &'static str, def: bool) -> bool
  {
    self.plain_value(name).value_or(def)
  }

  // fn attribute_doc(&self)
  // {
  //   let mut doc = format!("Attribute `#[{}(<attribute list>)]`: {}\n",
  //       self.root_name.get(), self.root_desc);
  //   for info in self.infos.iter() {
  //     doc.add(format!("  * `#[{}({})]`: {}\n",
  //       self.root_name.get(), info.name, info.desc));
  //   }
  //   self.cx.parse_sess.span_diagnostic.handler.note(doc.as_slice());
  // }
}

impl AttributeLitModel
{
  pub fn to_lit_printer(&self) -> LitTypePrinter
  {
    match *self {
      MLitStr(_) => PLitStr,
      MLitBinary(_) => PLitBinary,
      MLitByte(_) => PLitByte,
      MLitChar(_) => PLitChar,
      MLitInt(_) => PLitInt,
      MLitFloat(_) => PLitFloat,
      MLitFloatUnsuffixed(_) => PLitFloatUnsuffixed,
      MLitNil(_) => PLitNil,
      MLitBool(_) => PLitBool
    }
  }
}

pub fn lit_to_lit_printer(lit: &Lit_) -> LitTypePrinter
{
  match *lit {
    LitStr(_, _) => PLitStr,
    LitBinary(_) => PLitBinary,
    LitByte(_) => PLitByte,
    LitChar(_) => PLitChar,
    LitInt(_, _) => PLitInt,
    LitFloat(_, _) => PLitFloat,
    LitFloatUnsuffixed(_) => PLitFloatUnsuffixed,
    LitNil => PLitNil,
    LitBool(_) => PLitBool
  }
}

pub enum LitTypePrinter
{
  PLitStr,
  PLitBinary,
  PLitByte,
  PLitChar,
  PLitInt,
  PLitFloat,
  PLitFloatUnsuffixed,
  PLitNil,
  PLitBool
}

impl LitTypePrinter
{
  pub fn type_to_str(&self) -> &'static str
  {
    match *self {
      PLitStr => "string",
      PLitBinary => "binary",
      PLitByte => "byte",
      PLitChar => "char",
      PLitInt => "int",
      PLitFloat => "float",
      PLitFloatUnsuffixed => "unsuffixed float",
      PLitNil => "nil",
      PLitBool => "boolean"
    }
  }

  pub fn type_example_to_str(&self) -> &'static str
  {
    match *self {
      PLitStr => "\"hello world\"",
      PLitBinary => "0b01010101",
      PLitByte => "b'9'",
      PLitChar => "'a'",
      PLitInt => "38",
      PLitFloat => "0.01f32",
      PLitFloatUnsuffixed => "0.1",
      PLitNil => "()",
      PLitBool => "true"
    }
  }
}
<commit_msg>[AttributeValue] The value and span are now public members.<commit_after>// Copyright 2014 Pierre Talbot (IRCAM)

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub use compile_error::*;
pub use syntax::parse::token::InternedString;
pub use syntax::ast::*;

use syntax::codemap::DUMMY_SP;
use std::rc::Rc;

pub struct AttributeValue<T>
{
  pub value: Option<T>,
  pub span: Span,
  duplicate: DuplicateAttribute
}

impl<T: Clone> AttributeValue<T>
{
  pub fn new(duplicate: DuplicateAttribute) -> AttributeValue<T>
  {
    AttributeValue {
      value: None,
      duplicate: duplicate,
      span: DUMMY_SP
    }
  }

  pub fn simple() -> AttributeValue<T>
  {
    AttributeValue::new(DuplicateAttribute::simple(Warn))
  }

  pub fn update(self, cx: &ExtCtxt, value: T, span: Span) -> AttributeValue<T>
  {
    let mut this = self;
    if this.value.is_some() {
      this.duplicate.issue(cx, this.span, span);
    } else {
      this.value = Some(value);
      this.span = span;
    }
    this
  }

  pub fn value_or(&self, default: T) -> T
  {
    match self.value {
      None => default,
      Some(ref value) => value.clone()
    }
  }

  pub fn span(&self) -> Span
  {
    self.span.clone()
  }
}

pub struct AttributeInfo
{
  pub name: InternedString,
  pub desc: &'static str,
  pub model: AttributeModel
}

impl AttributeInfo
{
  pub fn new(name: &'static str, desc: &'static str,
    model: AttributeModel) -> AttributeInfo
  {
    AttributeInfo {
      name: InternedString::new(name),
      desc: desc,
      model: model
    }
  }

  pub fn simple(name: &'static str, desc: &'static str) -> AttributeInfo
  {
    AttributeInfo::new(name, desc, NoProperty(AttributeValue::simple()))
  }

  pub fn update(self, model: AttributeModel) -> AttributeInfo
  {
    let mut this = self;
    this.model = model;
    this
  }

  pub fn plain_value<'a>(&'a self) -> &'a AttributeValue<bool>
  {
    match self.model {
      NoProperty(ref v) => v,
      _ => fail!("No plain value for the current attribute.")
    }
  }

  pub fn sub_model<'a>(&'a self) -> &'a AttributeDict
  {
    match self.model {
      SubAttribute(ref dict) => dict,
      _ => fail!("No sub value for the attribute.")
    }
  }
}

pub enum AttributeModel
{
  NoProperty(AttributeValue<bool>),
  KeyValue(AttributeLitModel),
  SubAttribute(AttributeDict)
}

pub enum AttributeLitModel
{
  MLitStr(AttributeValue<(InternedString, StrStyle)>),
  MLitBinary(AttributeValue<Rc<Vec<u8>>>),
  MLitByte(AttributeValue<u8>),
  MLitChar(AttributeValue<char>),
  MLitInt(AttributeValue<(u64, LitIntType)>),
  MLitFloat(AttributeValue<(InternedString, FloatTy)>),
  MLitFloatUnsuffixed(AttributeValue<InternedString>),
  MLitNil(AttributeValue<()>),
  MLitBool(AttributeValue<bool>)
}

pub struct AttributeDict
{
  dict: Vec<AttributeInfo>
}

impl AttributeDict
{
  pub fn new(dict: Vec<AttributeInfo>) -> AttributeDict
  {
    AttributeDict {
      dict: dict
    }
  }

  pub fn move_map(self, f: |AttributeInfo| -> AttributeInfo) -> AttributeDict
  {
    let mut this = self;
    this.dict = this.dict.move_iter().map(f).collect();
    this
  }

  pub fn push(&mut self, attr: AttributeInfo)
  {
    self.dict.push(attr);
  }

  pub fn push_all(&mut self, attrs: Vec<AttributeInfo>)
  {
    self.dict.push_all_move(attrs);
  }

  pub fn get<'a>(&'a self, name: &'static str) -> &'a AttributeInfo
  {
    let interned = InternedString::new(name);
    for info in self.dict.iter() {
      if info.name == interned {
        return info;
      }
    }
    fail!("Try to get an attribute that doesn't exist.")
  }

  pub fn plain_value<'a>(&'a self, name: &'static str) -> &'a AttributeValue<bool>
  {
    self.get(name).plain_value()
  }

  pub fn sub_model<'a>(&'a self, name: &'static str) -> &'a AttributeDict
  {
    self.get(name).sub_model()
  }

  pub fn plain_value_or(&self, name: &'static str, def: bool) -> bool
  {
    self.plain_value(name).value_or(def)
  }

  // fn attribute_doc(&self)
  // {
  //   let mut doc = format!("Attribute `#[{}(<attribute list>)]`: {}\n",
  //       self.root_name.get(), self.root_desc);
  //   for info in self.infos.iter() {
  //     doc.add(format!("  * `#[{}({})]`: {}\n",
  //       self.root_name.get(), info.name, info.desc));
  //   }
  //   self.cx.parse_sess.span_diagnostic.handler.note(doc.as_slice());
  // }
}

impl AttributeLitModel
{
  pub fn to_lit_printer(&self) -> LitTypePrinter
  {
    match *self {
      MLitStr(_) => PLitStr,
      MLitBinary(_) => PLitBinary,
      MLitByte(_) => PLitByte,
      MLitChar(_) => PLitChar,
      MLitInt(_) => PLitInt,
      MLitFloat(_) => PLitFloat,
      MLitFloatUnsuffixed(_) => PLitFloatUnsuffixed,
      MLitNil(_) => PLitNil,
      MLitBool(_) => PLitBool
    }
  }
}

pub fn lit_to_lit_printer(lit: &Lit_) -> LitTypePrinter
{
  match *lit {
    LitStr(_, _) => PLitStr,
    LitBinary(_) => PLitBinary,
    LitByte(_) => PLitByte,
    LitChar(_) => PLitChar,
    LitInt(_, _) => PLitInt,
    LitFloat(_, _) => PLitFloat,
    LitFloatUnsuffixed(_) => PLitFloatUnsuffixed,
    LitNil => PLitNil,
    LitBool(_) => PLitBool
  }
}

pub enum LitTypePrinter
{
  PLitStr,
  PLitBinary,
  PLitByte,
  PLitChar,
  PLitInt,
  PLitFloat,
  PLitFloatUnsuffixed,
  PLitNil,
  PLitBool
}

impl LitTypePrinter
{
  pub fn type_to_str(&self) -> &'static str
  {
    match *self {
      PLitStr => "string",
      PLitBinary => "binary",
      PLitByte => "byte",
      PLitChar => "char",
      PLitInt => "int",
      PLitFloat => "float",
      PLitFloatUnsuffixed => "unsuffixed float",
      PLitNil => "nil",
      PLitBool => "boolean"
    }
  }

  pub fn type_example_to_str(&self) -> &'static str
  {
    match *self {
      PLitStr => "\"hello world\"",
      PLitBinary => "0b01010101",
      PLitByte => "b'9'",
      PLitChar => "'a'",
      PLitInt => "38",
      PLitFloat => "0.01f32",
      PLitFloatUnsuffixed => "0.1",
      PLitNil => "()",
      PLitBool => "true"
    }
  }
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate evo;
extern crate petgraph;
extern crate graph_annealing;
extern crate clap;
extern crate crossbeam;
extern crate simple_parallel;
extern crate num_cpus;

use std::f32;
use std::fmt::Debug;
use std::str::FromStr;
use clap::{App, Arg};
use rand::{Rng, SeedableRng};
use rand::isaac::Isaac64Rng;
use rand::distributions::{IndependentSample, Range};

use evo::Probability;
use evo::nsga2::{self, FitnessEval, Mate, MultiObjective3};
use graph_annealing::repr::edge_ops_genome::EdgeOpsGenome;
use graph_annealing::helper::{draw_graph, line_graph};
use graph_annealing::goal::Goal;
use simple_parallel::Pool;

struct Mating {
    prob_mutate_elem: Probability,
}

impl Mate<EdgeOpsGenome> for Mating {
    fn mate<R: Rng>(&mut self,
                    rng: &mut R,
                    p1: &EdgeOpsGenome,
                    p2: &EdgeOpsGenome)
                    -> EdgeOpsGenome {
        let mut child = EdgeOpsGenome::mate(rng, p1, p2);
        child.mutate(rng, self.prob_mutate_elem);
        child
    }
}

fn fitness(goal: &Goal, ind: &EdgeOpsGenome) -> MultiObjective3<f32> {
    let g = ind.to_graph();
    let cc_dist = goal.strongly_connected_components_distance(&g);
    let nmc = goal.neighbor_matching_score(&g);
    let triadic_dist = goal.triadic_distance(g);
    MultiObjective3::from((cc_dist as f32,
                           triadic_dist as f32,
                           nmc /* ind.num_edges() as f32 */))
}

struct MyEval {
    goal: Goal,
    pool: Pool,
}

impl FitnessEval<EdgeOpsGenome, MultiObjective3<f32>> for MyEval {
    fn fitness(&mut self, pop: &[EdgeOpsGenome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;

        crossbeam::scope(|scope| pool.map(scope, pop, |ind| fitness(goal, ind)).collect())
    }
}

#[derive(Debug)]
struct Stat<T:Debug> {
   min: T, max: T, avg: T
}

fn main() {
    let ncpus = num_cpus::get();
    println!("Using {} CPUs", ncpus);
    let matches = App::new("nsga2_edge")
                      .arg(Arg::with_name("NGEN")
                               .long("ngen")
                               .help("Number of generations to run")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("N")
                               .long("n")
                               .help("Problem size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("MU")
                               .long("mu")
                               .help("Population size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("LAMBDA")
                               .long("lambda")
                               .help("Size of offspring population")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("SEED")
                               .long("seed")
                               .help("Seed value for Rng")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("PMUT")
                               .long("pmut")
                               .help("Probability for element mutation")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("ILEN")
                               .long("ilen")
                               .help("Initial genome length (random range)")
                               .takes_value(true)
                               .required(true))
                      .get_matches();

    // number of generations
    let NGEN: usize = FromStr::from_str(matches.value_of("NGEN").unwrap()).unwrap();
    println!("NGEN: {}", NGEN);

    // problem size to solve (number of nodes of line-graph)
    let N: usize = FromStr::from_str(matches.value_of("N").unwrap()).unwrap();
    println!("N: {}", N);

    // size of population
    let MU: usize = FromStr::from_str(matches.value_of("MU").unwrap()).unwrap();
    println!("MU: {}", MU);

    // size of offspring population
    let LAMBDA: usize = FromStr::from_str(matches.value_of("LAMBDA").unwrap()).unwrap();
    println!("LAMBDA: {}", LAMBDA);

    let seed_str = matches.value_of("SEED").unwrap();
    let seed: Vec<u64> = seed_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    println!("SEED: {:?}", seed);

    let PMUT: f32 = FromStr::from_str(matches.value_of("PMUT").unwrap()).unwrap();
    println!("PMUT: {:?}", PMUT);

    let ilen_str = matches.value_of("ILEN").unwrap();
    let ilen: Vec<usize> = ilen_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    assert!(ilen.len() == 1 || ilen.len() == 2);

    let ilen_from = ilen[0];
    let ilen_to = if ilen.len() == 1 {
        ilen[0]
    } else {
        ilen[1]
    };
    assert!(ilen_from <= ilen_to);

    let mut evaluator = MyEval {
        goal: Goal::new(line_graph(N as u32)),
        pool: Pool::new(ncpus),
    };

    let mut rng: Isaac64Rng = SeedableRng::from_seed(&seed[..]);

    // create initial random population
    let initial_population: Vec<EdgeOpsGenome> = (0..MU)
                                                     .map(|_| {
                                                         let len = if ilen_from == ilen_to {
                                                             ilen_from
                                                         } else {
                                                             Range::new(ilen_from, ilen_to)
                                                                 .ind_sample(&mut rng)
                                                         };
                                                         EdgeOpsGenome::random(&mut rng, len)
                                                     })
                                                     .collect();

    // evaluate fitness
    let fitness: Vec<_> = evaluator.fitness(&initial_population[..]);
    assert!(fitness.len() == initial_population.len());

    let mut pop = initial_population;
    let mut fit = fitness;

    let mut mating = Mating { prob_mutate_elem: Probability::new(PMUT) };

    for i in 0..NGEN {
        println!("Iteration: {}", i);
        let (new_pop, new_fit) = nsga2::iterate(&mut rng,
                                                pop,
                                                fit,
                                                &mut evaluator,
                                                MU,
                                                LAMBDA,
                                                2,
                                                &mut mating);
        pop = new_pop;
        fit = new_fit;

        let stats: Vec<Stat<f32>> = [0,1,2].iter().map(|&i| {
            let min = fit.iter().fold(f32::INFINITY, |acc, f| {
                let x = f.objectives[i];
                if acc < x { acc } else { x }
                });
            let max = fit.iter().fold(0.0, |acc, f| {
                let x = f.objectives[i];
                if acc > x { acc } else { x }
                });
            let sum = fit.iter().fold(0.0, |acc, f| acc + f.objectives[i]);
            Stat{min: min, max: max, avg: sum / fit.len() as f32}
        }).collect();
        println!("stats: {:?}", stats);
        let mut found_optimum = false;
        for f in fit.iter() {
            if f.objectives[0] < 1.0 && f.objectives[1] < 3.0 && f.objectives[2] < 0.01 {
                found_optimum = true;
                break;
            }
        }
        if found_optimum {
            println!("Found premature optimum in Iteration {}", i);
            break;
        }
    }
    println!("===========================================================");

    // finally evaluate rank and crowding distance (using select()).
    let rank_dist = nsga2::select(&fit[..], MU);
    assert!(rank_dist.len() == MU);

    let mut j = 0;

    for rd in rank_dist.iter() {
        if rd.rank == 0 {
            println!("-------------------------------------------");
            println!("rd: {:?}", rd);
            println!("fitness: {:?}", fit[rd.idx]);
            println!("genome: {:?}", pop[rd.idx]);

            if fit[rd.idx].objectives[0] < 1.0 {
                draw_graph(&pop[rd.idx].to_graph(),
                           &format!("line_{}_nsga2edge_ops_n{}_f{}_i{}.svg",
                                    N,
                                    NGEN,
                                    fit[rd.idx].objectives[1] as usize,
                                    j));
                j += 1;
            }
        }
    }
}
<commit_msg>Tournament size K as commandline argument --k<commit_after>extern crate rand;
extern crate evo;
extern crate petgraph;
extern crate graph_annealing;
extern crate clap;
extern crate crossbeam;
extern crate simple_parallel;
extern crate num_cpus;

use std::f32;
use std::fmt::Debug;
use std::str::FromStr;
use clap::{App, Arg};
use rand::{Rng, SeedableRng};
use rand::isaac::Isaac64Rng;
use rand::distributions::{IndependentSample, Range};

use evo::Probability;
use evo::nsga2::{self, FitnessEval, Mate, MultiObjective3};
use graph_annealing::repr::edge_ops_genome::EdgeOpsGenome;
use graph_annealing::helper::{draw_graph, line_graph};
use graph_annealing::goal::Goal;
use simple_parallel::Pool;

struct Mating {
    prob_mutate_elem: Probability,
}

impl Mate<EdgeOpsGenome> for Mating {
    fn mate<R: Rng>(&mut self,
                    rng: &mut R,
                    p1: &EdgeOpsGenome,
                    p2: &EdgeOpsGenome)
                    -> EdgeOpsGenome {
        let mut child = EdgeOpsGenome::mate(rng, p1, p2);
        child.mutate(rng, self.prob_mutate_elem);
        child
    }
}

fn fitness(goal: &Goal, ind: &EdgeOpsGenome) -> MultiObjective3<f32> {
    let g = ind.to_graph();
    let cc_dist = goal.strongly_connected_components_distance(&g);
    let nmc = goal.neighbor_matching_score(&g);
    let triadic_dist = goal.triadic_distance(g);
    MultiObjective3::from((cc_dist as f32,
                           triadic_dist as f32,
                           nmc /* ind.num_edges() as f32 */))
}

struct MyEval {
    goal: Goal,
    pool: Pool,
}

impl FitnessEval<EdgeOpsGenome, MultiObjective3<f32>> for MyEval {
    fn fitness(&mut self, pop: &[EdgeOpsGenome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;

        crossbeam::scope(|scope| pool.map(scope, pop, |ind| fitness(goal, ind)).collect())
    }
}

#[derive(Debug)]
struct Stat<T:Debug> {
   min: T, max: T, avg: T
}

fn main() {
    let ncpus = num_cpus::get();
    println!("Using {} CPUs", ncpus);
    let matches = App::new("nsga2_edge")
                      .arg(Arg::with_name("NGEN")
                               .long("ngen")
                               .help("Number of generations to run")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("N")
                               .long("n")
                               .help("Problem size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("MU")
                               .long("mu")
                               .help("Population size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("LAMBDA")
                               .long("lambda")
                               .help("Size of offspring population")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("K")
                               .long("k")
                               .help("Tournament selection (default: 2)")
                               .takes_value(true))
                      .arg(Arg::with_name("SEED")
                               .long("seed")
                               .help("Seed value for Rng")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("PMUT")
                               .long("pmut")
                               .help("Probability for element mutation")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("ILEN")
                               .long("ilen")
                               .help("Initial genome length (random range)")
                               .takes_value(true)
                               .required(true))
                      .get_matches();

    // number of generations
    let NGEN: usize = FromStr::from_str(matches.value_of("NGEN").unwrap()).unwrap();
    println!("NGEN: {}", NGEN);

    // problem size to solve (number of nodes of line-graph)
    let N: usize = FromStr::from_str(matches.value_of("N").unwrap()).unwrap();
    println!("N: {}", N);

    // size of population
    let MU: usize = FromStr::from_str(matches.value_of("MU").unwrap()).unwrap();
    println!("MU: {}", MU);

    // size of offspring population
    let LAMBDA: usize = FromStr::from_str(matches.value_of("LAMBDA").unwrap()).unwrap();
    println!("LAMBDA: {}", LAMBDA);

    // tournament selection
    let K: usize = FromStr::from_str(matches.value_of("K").unwrap_or("2")).unwrap();
    assert!(K > 0);
    println!("K: {}", K);

    let seed_str = matches.value_of("SEED").unwrap();
    let seed: Vec<u64> = seed_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    println!("SEED: {:?}", seed);

    let PMUT: f32 = FromStr::from_str(matches.value_of("PMUT").unwrap()).unwrap();
    println!("PMUT: {:?}", PMUT);

    let ilen_str = matches.value_of("ILEN").unwrap();
    let ilen: Vec<usize> = ilen_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    assert!(ilen.len() == 1 || ilen.len() == 2);

    let ilen_from = ilen[0];
    let ilen_to = if ilen.len() == 1 {
        ilen[0]
    } else {
        ilen[1]
    };
    assert!(ilen_from <= ilen_to);

    let mut evaluator = MyEval {
        goal: Goal::new(line_graph(N as u32)),
        pool: Pool::new(ncpus),
    };

    let mut rng: Isaac64Rng = SeedableRng::from_seed(&seed[..]);

    // create initial random population
    let initial_population: Vec<EdgeOpsGenome> = (0..MU)
                                                     .map(|_| {
                                                         let len = if ilen_from == ilen_to {
                                                             ilen_from
                                                         } else {
                                                             Range::new(ilen_from, ilen_to)
                                                                 .ind_sample(&mut rng)
                                                         };
                                                         EdgeOpsGenome::random(&mut rng, len)
                                                     })
                                                     .collect();

    // evaluate fitness
    let fitness: Vec<_> = evaluator.fitness(&initial_population[..]);
    assert!(fitness.len() == initial_population.len());

    let mut pop = initial_population;
    let mut fit = fitness;

    let mut mating = Mating { prob_mutate_elem: Probability::new(PMUT) };

    for i in 0..NGEN {
        println!("Iteration: {}", i);
        let (new_pop, new_fit) = nsga2::iterate(&mut rng,
                                                pop,
                                                fit,
                                                &mut evaluator,
                                                MU,
                                                LAMBDA,
                                                K,
                                                &mut mating);
        pop = new_pop;
        fit = new_fit;

        let stats: Vec<Stat<f32>> = [0,1,2].iter().map(|&i| {
            let min = fit.iter().fold(f32::INFINITY, |acc, f| {
                let x = f.objectives[i];
                if acc < x { acc } else { x }
                });
            let max = fit.iter().fold(0.0, |acc, f| {
                let x = f.objectives[i];
                if acc > x { acc } else { x }
                });
            let sum = fit.iter().fold(0.0, |acc, f| acc + f.objectives[i]);
            Stat{min: min, max: max, avg: sum / fit.len() as f32}
        }).collect();
        println!("stats: {:?}", stats);
        let mut found_optimum = false;
        for f in fit.iter() {
            if f.objectives[0] < 1.0 && f.objectives[1] < 3.0 && f.objectives[2] < 0.01 {
                found_optimum = true;
                break;
            }
        }
        if found_optimum {
            println!("Found premature optimum in Iteration {}", i);
            break;
        }
    }
    println!("===========================================================");

    // finally evaluate rank and crowding distance (using select()).
    let rank_dist = nsga2::select(&fit[..], MU);
    assert!(rank_dist.len() == MU);

    let mut j = 0;

    for rd in rank_dist.iter() {
        if rd.rank == 0 {
            println!("-------------------------------------------");
            println!("rd: {:?}", rd);
            println!("fitness: {:?}", fit[rd.idx]);
            println!("genome: {:?}", pop[rd.idx]);

            if fit[rd.idx].objectives[0] < 1.0 {
                draw_graph(&pop[rd.idx].to_graph(),
                           &format!("line_{}_nsga2edge_ops_n{}_f{}_i{}.svg",
                                    N,
                                    NGEN,
                                    fit[rd.idx].objectives[1] as usize,
                                    j));
                j += 1;
            }
        }
    }
}
<|endoftext|>