<commit_before>#![crate_name = "raft"]
#![crate_type="lib"]
#![doc(html_logo_url = "https://raw.githubusercontent.com/Hoverbear/raft/master/raft.png")]
#![doc(html_root_url = "https://hoverbear.github.io/raft/raft/")]

#![feature(test)]

//! This is the Raft Distributed Consensus Protocol implemented for Rust.
//! [Raft](http://raftconsensus.github.io/) is described as:
//!
//! > Raft is a consensus algorithm that is designed to be easy to understand. It's equivalent to
//! > Paxos in fault-tolerance and performance. The difference is that it's decomposed into
//! > relatively independent subproblems, and it cleanly addresses all major pieces needed for
//! > practical systems.
//!
//! This implementation utilizes [Cap'n Proto](https://kentonv.github.io/capnproto/) for its RPC,
//! [`mio`](https://github.com/carllerche/mio) for it's async event loop.
//!
//! If this package fails to build for you it is possibly because you do not have the
//! [`capnp`](https://capnproto.org/capnp-tool.html) utility installed. You should be able to find
//! appropriate packages for most popular distributions.
//!
//! # Consuming this library
//!
//! Consuming this library works in a few parts:
//!
//! 1. Implement or consume a Persistent Log and a State Machine such that they will hook into
//!    your application desirably.
//! 2. Create a `Server` with those implementations. It will independently fire up and join the
//!    cluster.
//! 3. Interact with the cluster by issuing `.propose()` and `.query()` calls via the `Client`
//! 4. React to calls to `.propose()` and `.query()` from the implemented `StateMachine`
//!
//! ## Peristent Log
//!
//! A `Log` represents the **replicated, persistent log** of your application. It has a
//! strong ordering such that `A → B → C` and should **only** act to store information. Entries
//! placed into the log should not be acted on in any way by the consuming application as they
//! have not been committed to the cluster.
//!
//! Some ideas for a Persistent Log implementation:
//!
//!   * A PostgreSQL / SQLite instance.
//!   * A plain old file.
//!   * A vector in memory *(Note: Log compaction is still pending, so be aware of running out!)*
//!
//! > It is our belief that in many cases the implementation of `Log` will be generic to
//! > application purposes. You are encouraged to submit your own implementations to us!
//!
//! ## State Machine
//!
//! The `StateMachine` represents the **stateful representation** of your application. Events
//! are applied to the `StateMachine` in the correct ordering at the time of commit. This is where
//! your application **should** act on information.
//!
//! In the `StateMachine` there are both mutable (`.apply()`) and immutable (`.query()`) calls.
//! There is a considerable performance difference, as `.query()` calls do not pass through the
//! durablable `Log` while `.apply()` events do.
//!
//! Some ideas for a State Machine implementation:
//!
//!   * A Hashmap or key-value store (Example provided)
//!   * A single register (Example provided)
//!   * Basically anything from `std::collections`
//!
//! ## Client Requests
//!
//! Client requests are the **only** way to interact with the Raft cluster. Calls to `.propose()`
//! and `.query()` are automatically routed to the relevant `Leader` node and behave as blocking
//! calls.
//!
//! This means `.propose()` won't return until the entry is durably replicated into the log of at
//! least the majority of the cluster and has been commited.
//!

extern crate bufstream;
extern crate capnp;
extern crate mio;
extern crate rand;
extern crate uuid;
extern crate test;
#[macro_use] extern crate log;
#[macro_use] extern crate scoped_log;
#[macro_use] extern crate wrapped_enum;
#[cfg(test)] extern crate env_logger;
#[cfg(feature = "serde")] extern crate serde;

/// Prepares the environment testing. Should be called as the first line of every test with the
/// name of the test as the only argument.
///
/// TODO: Make this an annotation like #[rust_test] instead of a macro.
#[cfg(test)]
macro_rules! setup_test {
    ($test_name:expr) => (
        let _ = env_logger::init();
        push_log_scope!($test_name);
    );
}

pub mod state_machine;
pub mod persistent_log;
pub mod messages_capnp {
    #![allow(dead_code)]
    include!(concat!(env!("OUT_DIR"), "/messages_capnp.rs"));
}

mod backoff;
mod client;
mod connection;
mod messages;
mod consensus;
mod server;
mod state;

pub use server::Server;
pub use state_machine::StateMachine;
pub use persistent_log::Log;
pub use client::Client;

use std::{io, net, ops, fmt};

use uuid::Uuid;

/// A simple convienence type.
pub type Result<T> = std::result::Result<T, Error>;

wrapped_enum!{
    #[doc = "The generic `raft::Error` is composed of one of the errors that can originate from the"]
    #[doc = "various libraries consumed by the library."]
    #[doc = "With the exception of the `Raft` variant these are generated from `try!()` macros invoking"]
    #[doc = "on `io::Error` or `capnp::Error` by using"]
    #[doc = "[`FromError`](https://doc.rust-lang.org/std/error/#the-fromerror-trait)."]
    #[derive(Debug)]
    pub enum Error {
        /// An error originating from the [Cap'n Proto](https://github.com/dwrensha/capnproto-rust) library.
        CapnProto(capnp::Error),
        /// A specific error produced when a bad Cap'n proto message is discovered.
        SchemaError(capnp::NotInSchema),
        /// Errors originating from `std::io`.
        Io(io::Error),
        /// Raft specific errors.
        Raft(RaftError),
        /// Errors related to parsing addresses.
        AddrParse(net::AddrParseError),
    }
}

/// A Raft Error represents a Raft specific error that consuming code is expected to handle
/// gracefully.
#[derive(Debug)]
pub enum RaftError {
    /// The server ran out of slots in the slab for new connections
    ConnectionLimitReached,
    /// A client reported an invalid client id
    InvalidClientId,
    /// A consensus module reported back a leader not in the cluster.
    ClusterViolation,
    /// A remote connection attempted to use an unknown connection type in the connection preamble
    UnknownConnectionType,
    /// An invalid peer in in the peer set. Returned Server::new().
    InvalidPeerSet,
    /// Registering a connection failed
    ConnectionRegisterFailed,
    /// Failed to find a leader in the cluster. Try again later.
    LeaderSearchExhausted,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::CapnProto(ref error) => fmt::Display::fmt(error, f),
            Error::SchemaError(ref error) => fmt::Display::fmt(error, f),
            Error::Io(ref error) => fmt::Display::fmt(error, f),
            Error::Raft(ref error) => fmt::Debug::fmt(error, f),
            Error::AddrParse(ref error) => fmt::Debug::fmt(error, f),
        }
    }
}

/// The term of a log entry.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Term(u64);
impl From<u64> for Term {
    fn from(val: u64) -> Term {
        Term(val)
    }
}
impl Into<u64> for Term {
    fn into(self) -> u64 {
        self.0
    }
}
impl ops::Add<u64> for Term {
    type Output = Term;
    fn add(self, rhs: u64) -> Term {
        Term(self.0.checked_add(rhs).expect("overflow while incrementing Term"))
    }
}
impl ops::Sub<u64> for Term {
    type Output = Term;
    fn sub(self, rhs: u64) -> Term {
        Term(self.0.checked_sub(rhs).expect("underflow while decrementing Term"))
    }
}
impl fmt::Display for Term {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

/// The index of a log entry.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct LogIndex(u64);
impl From<u64> for LogIndex {
    fn from(val: u64) -> LogIndex {
        LogIndex(val)
    }
}
impl Into<u64> for LogIndex {
    fn into(self) -> u64 {
        self.0
    }
}
impl ops::Add<u64> for LogIndex {
    type Output = LogIndex;
    fn add(self, rhs: u64) -> LogIndex {
        LogIndex(self.0.checked_add(rhs).expect("overflow while incrementing LogIndex"))
    }
}
impl ops::Sub<u64> for LogIndex {
    type Output = LogIndex;
    fn sub(self, rhs: u64) -> LogIndex {
        LogIndex(self.0.checked_sub(rhs).expect("underflow while decrementing LogIndex"))
    }
}
impl fmt::Display for LogIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

/// The id of a Raft server. Must be unique among the participants in a
/// consensus group.
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
pub struct ServerId(u64);
impl From<u64> for ServerId {
    fn from(val: u64) -> ServerId {
        ServerId(val)
    }
}
impl Into<u64> for ServerId {
    fn into(self) -> u64 {
        self.0
    }
}
impl fmt::Debug for ServerId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ServerId({})", self.0)
    }
}
impl fmt::Display for ServerId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

/// The ID of a Raft client.
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
pub struct ClientId(Uuid);
impl ClientId {
    fn new() -> ClientId {
        ClientId(Uuid::new_v4())
    }
    fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }
    fn from_bytes(bytes: &[u8]) -> Result<ClientId> {
        match Uuid::from_bytes(bytes) {
            Some(uuid) => Ok(ClientId(uuid)),
            None => Err(Error::Raft(RaftError::InvalidClientId)),
        }
    }
}
impl fmt::Debug for ClientId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ClientId({})", self.0)
    }
}
impl fmt::Display for ClientId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}
<commit_msg>ClientId from/into<commit_after>#![crate_name = "raft"]
#![crate_type="lib"]
#![doc(html_logo_url = "https://raw.githubusercontent.com/Hoverbear/raft/master/raft.png")]
#![doc(html_root_url = "https://hoverbear.github.io/raft/raft/")]

#![feature(test)]

//! This is the Raft Distributed Consensus Protocol implemented for Rust.
//! [Raft](http://raftconsensus.github.io/) is described as:
//!
//! > Raft is a consensus algorithm that is designed to be easy to understand. It's equivalent to
//! > Paxos in fault-tolerance and performance. The difference is that it's decomposed into
//! > relatively independent subproblems, and it cleanly addresses all major pieces needed for
//! > practical systems.
//!
//! This implementation utilizes [Cap'n Proto](https://kentonv.github.io/capnproto/) for its RPC,
//! [`mio`](https://github.com/carllerche/mio) for it's async event loop.
//!
//! If this package fails to build for you it is possibly because you do not have the
//! [`capnp`](https://capnproto.org/capnp-tool.html) utility installed. You should be able to find
//! appropriate packages for most popular distributions.
//!
//! # Consuming this library
//!
//! Consuming this library works in a few parts:
//!
//! 1. Implement or consume a Persistent Log and a State Machine such that they will hook into
//!    your application desirably.
//! 2. Create a `Server` with those implementations. It will independently fire up and join the
//!    cluster.
//! 3. Interact with the cluster by issuing `.propose()` and `.query()` calls via the `Client`
//! 4. React to calls to `.propose()` and `.query()` from the implemented `StateMachine`
//!
//! ## Peristent Log
//!
//! A `Log` represents the **replicated, persistent log** of your application. It has a
//! strong ordering such that `A → B → C` and should **only** act to store information. Entries
//! placed into the log should not be acted on in any way by the consuming application as they
//! have not been committed to the cluster.
//!
//! Some ideas for a Persistent Log implementation:
//!
//!   * A PostgreSQL / SQLite instance.
//!   * A plain old file.
//!   * A vector in memory *(Note: Log compaction is still pending, so be aware of running out!)*
//!
//! > It is our belief that in many cases the implementation of `Log` will be generic to
//! > application purposes. You are encouraged to submit your own implementations to us!
//!
//! ## State Machine
//!
//! The `StateMachine` represents the **stateful representation** of your application. Events
//! are applied to the `StateMachine` in the correct ordering at the time of commit. This is where
//! your application **should** act on information.
//!
//! In the `StateMachine` there are both mutable (`.apply()`) and immutable (`.query()`) calls.
//! There is a considerable performance difference, as `.query()` calls do not pass through the
//! durablable `Log` while `.apply()` events do.
//!
//! Some ideas for a State Machine implementation:
//!
//!   * A Hashmap or key-value store (Example provided)
//!   * A single register (Example provided)
//!   * Basically anything from `std::collections`
//!
//! ## Client Requests
//!
//! Client requests are the **only** way to interact with the Raft cluster. Calls to `.propose()`
//! and `.query()` are automatically routed to the relevant `Leader` node and behave as blocking
//! calls.
//!
//! This means `.propose()` won't return until the entry is durably replicated into the log of at
//! least the majority of the cluster and has been commited. `.query()` will perform better if
//! you wish to only read data and not have it pass through the persisted log.
//!

extern crate bufstream;
extern crate capnp;
extern crate mio;
extern crate rand;
extern crate uuid;
extern crate test;
#[macro_use] extern crate log;
#[macro_use] extern crate scoped_log;
#[macro_use] extern crate wrapped_enum;
#[cfg(test)] extern crate env_logger;
#[cfg(feature = "serde")] extern crate serde;

/// Prepares the environment testing. Should be called as the first line of every test with the
/// name of the test as the only argument.
///
/// TODO: Make this an annotation like #[rust_test] instead of a macro.
#[cfg(test)]
macro_rules! setup_test {
    ($test_name:expr) => (
        let _ = env_logger::init();
        push_log_scope!($test_name);
    );
}

pub mod state_machine;
pub mod persistent_log;
pub mod messages_capnp {
    #![allow(dead_code)]
    include!(concat!(env!("OUT_DIR"), "/messages_capnp.rs"));
}

mod backoff;
mod client;
mod connection;
mod messages;
mod consensus;
mod server;
mod state;

pub use server::Server;
pub use state_machine::StateMachine;
pub use persistent_log::Log;
pub use client::Client;

use std::{io, net, ops, fmt};

use uuid::Uuid;

/// A simple convienence type.
pub type Result<T> = std::result::Result<T, Error>;

wrapped_enum!{
    #[doc = "The generic `raft::Error` is composed of one of the errors that can originate from the"]
    #[doc = "various libraries consumed by the library."]
    #[doc = "With the exception of the `Raft` variant these are generated from `try!()` macros invoking"]
    #[doc = "on `io::Error` or `capnp::Error` by using"]
    #[doc = "[`FromError`](https://doc.rust-lang.org/std/error/#the-fromerror-trait)."]
    #[derive(Debug)]
    pub enum Error {
        /// An error originating from the [Cap'n Proto](https://github.com/dwrensha/capnproto-rust) library.
        CapnProto(capnp::Error),
        /// A specific error produced when a bad Cap'n proto message is discovered.
        SchemaError(capnp::NotInSchema),
        /// Errors originating from `std::io`.
        Io(io::Error),
        /// Raft specific errors.
        Raft(RaftError),
        /// Errors related to parsing addresses.
        AddrParse(net::AddrParseError),
    }
}

/// A Raft Error represents a Raft specific error that consuming code is expected to handle
/// gracefully.
#[derive(Debug)]
pub enum RaftError {
    /// The server ran out of slots in the slab for new connections
    ConnectionLimitReached,
    /// A client reported an invalid client id
    InvalidClientId,
    /// A consensus module reported back a leader not in the cluster.
    ClusterViolation,
    /// A remote connection attempted to use an unknown connection type in the connection preamble
    UnknownConnectionType,
    /// An invalid peer in in the peer set. Returned Server::new().
    InvalidPeerSet,
    /// Registering a connection failed
    ConnectionRegisterFailed,
    /// Failed to find a leader in the cluster. Try again later.
    LeaderSearchExhausted,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::CapnProto(ref error) => fmt::Display::fmt(error, f),
            Error::SchemaError(ref error) => fmt::Display::fmt(error, f),
            Error::Io(ref error) => fmt::Display::fmt(error, f),
            Error::Raft(ref error) => fmt::Debug::fmt(error, f),
            Error::AddrParse(ref error) => fmt::Debug::fmt(error, f),
        }
    }
}

/// The term of a log entry.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Term(u64);
impl From<u64> for Term {
    fn from(val: u64) -> Term {
        Term(val)
    }
}
impl Into<u64> for Term {
    fn into(self) -> u64 {
        self.0
    }
}
impl ops::Add<u64> for Term {
    type Output = Term;
    fn add(self, rhs: u64) -> Term {
        Term(self.0.checked_add(rhs).expect("overflow while incrementing Term"))
    }
}
impl ops::Sub<u64> for Term {
    type Output = Term;
    fn sub(self, rhs: u64) -> Term {
        Term(self.0.checked_sub(rhs).expect("underflow while decrementing Term"))
    }
}
impl fmt::Display for Term {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

/// The index of a log entry.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct LogIndex(u64);
impl From<u64> for LogIndex {
    fn from(val: u64) -> LogIndex {
        LogIndex(val)
    }
}
impl Into<u64> for LogIndex {
    fn into(self) -> u64 {
        self.0
    }
}
impl ops::Add<u64> for LogIndex {
    type Output = LogIndex;
    fn add(self, rhs: u64) -> LogIndex {
        LogIndex(self.0.checked_add(rhs).expect("overflow while incrementing LogIndex"))
    }
}
impl ops::Sub<u64> for LogIndex {
    type Output = LogIndex;
    fn sub(self, rhs: u64) -> LogIndex {
        LogIndex(self.0.checked_sub(rhs).expect("underflow while decrementing LogIndex"))
    }
}
impl fmt::Display for LogIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

/// The id of a Raft server. Must be unique among the participants in a
/// consensus group. Represented as a `u64`.
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
pub struct ServerId(u64);
impl From<u64> for ServerId {
    fn from(val: u64) -> ServerId {
        ServerId(val)
    }
}
impl Into<u64> for ServerId {
    fn into(self) -> u64 {
        self.0
    }
}
impl fmt::Debug for ServerId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ServerId({})", self.0)
    }
}
impl fmt::Display for ServerId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

/// The ID of a Raft client. Represented as a `Uuid::v4`.
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
pub struct ClientId(Uuid);
impl ClientId {
    fn new() -> ClientId {
        ClientId(Uuid::new_v4())
    }
    fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }
    fn from_bytes(bytes: &[u8]) -> Result<ClientId> {
        match Uuid::from_bytes(bytes) {
            Some(uuid) => Ok(ClientId(uuid)),
            None => Err(Error::Raft(RaftError::InvalidClientId)),
        }
    }
}
impl From<Uuid> for ClientId {
    fn from(val: Uuid) -> ClientId {
        ClientId(val)
    }
}
impl Into<Uuid> for ClientId {
    fn into(self) -> Uuid {
        self.0
    }
}
impl fmt::Debug for ClientId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ClientId({})", self.0)
    }
}
impl fmt::Display for ClientId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}
<|endoftext|>


<commit_before>//! A Little implementation of the roulette-wheel principle, `RouletteWheel<T>`.
//! https://wikipedia.org/wiki/Fitness_proportionate_selection
//!
//! ![Fitness proportionate selection](https://upload.wikimedia.org/wikipedia/commons/2/2a/Fitness_proportionate_selection_example.png)
//!
//! # Usage
//!
//! # Examples
//!
//! You can explicitly create a `RouletteWheel<T>` with `new()`:
//!
//! ```
//! use roulette_wheel::RouletteWheel;
//!
//! let rw: RouletteWheel<u8> = RouletteWheel::new();
//! ```
//!
//! You can `push` values in the roulette-wheel (which will grow the wheel as needed):
//! ```
//! use roulette_wheel::RouletteWheel;
//!
//! let mut rw = RouletteWheel::new();
//!
//! rw.push(5., 'a');
//! rw.push(1., 'b');
//!
//! // .iter() will not consume the roulette wheel
//! for (fit, &ind) in rw.iter() {
//!    // do things with individuals here
//! }
//!
//! // .into_iter() consume the roulette wheel
//! for (fit, ind) in rw.into_iter() {
//!     // do things with individuals here
//! }
//!
//! ```

extern crate rand;

use std::iter::{FromIterator, Iterator, IntoIterator};
use rand::{Rng, ThreadRng, thread_rng};
use rand::distributions::{Range, IndependentSample};

/// A roulette-wheel container
pub struct RouletteWheel<T> {
    total_fitness: f32,
    fitnesses: Vec<f32>,
    population: Vec<T>
}

impl<T: Clone> Clone for RouletteWheel<T> {
    fn clone(&self) -> RouletteWheel<T> {
        RouletteWheel {
            total_fitness: self.total_fitness,
            fitnesses: self.fitnesses.clone(),
            population: self.population.clone()
        }
    }
}

impl<T> FromIterator<(f32, T)> for RouletteWheel<T> {
    fn from_iter<A>(iter: A) -> Self where A: IntoIterator<Item=(f32, T)> {
        let iter = iter.into_iter();
        let (lower, _) = iter.size_hint();

        let mut total_fitness = 0.0;
        let mut fitnesses = Vec::with_capacity(lower);
        let mut population = Vec::with_capacity(lower);

        for (fitness, individual) in iter {
            total_fitness += fitness;
            fitnesses.push(fitness);
            population.push(individual);
        }

        RouletteWheel {
            total_fitness: total_fitness,
            fitnesses: fitnesses,
            population: population
        }

        // #![feature(field_init_shorthand)]
        // RouletteWheel { total_fitness, fitnesses, population }
    }
}

impl<T> RouletteWheel<T> {
    /// create a new empty random-wheel.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let rw = RouletteWheel::<u8>::new();
    /// ```
    pub fn new() -> RouletteWheel<T> {
        RouletteWheel {
            total_fitness: 0.0,
            fitnesses: Vec::new(),
            population: Vec::new()
        }
    }

    /// Creates an empty RouletteWheel with space for at least n elements.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let rw = RouletteWheel::<u8>::with_capacity(15);
    ///
    /// assert_eq!(rw.len(), 0);
    /// ```
    pub fn with_capacity(cap: usize) -> RouletteWheel<T> {
        RouletteWheel {
            total_fitness: 0.0,
            fitnesses: Vec::with_capacity(cap),
            population: Vec::with_capacity(cap)
        }
    }

    /// Reserves capacity for at least `additional` more elements to be inserted.
    /// The collection may reserve more space to avoid frequent reallocations.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::<u8>::new();
    /// rw.reserve(20);
    ///
    /// assert_eq!(rw.len(), 0);
    /// ```
    pub fn reserve(&mut self, additional: usize) {
        self.fitnesses.reserve(additional);
        self.population.reserve(additional);
    }

    /// Returns the number of elements in the wheel.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let rw: RouletteWheel<_> = [(0.1, 10), (0.2, 15), (0.5, 20)].iter().cloned().collect();
    ///
    /// assert_eq!(rw.len(), 3);
    /// ```
    pub fn len(&self) -> usize {
        self.population.len()
    }

    /// Returns `true` if empty else return `false`.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let empty_rw = RouletteWheel::<u8>::new();
    ///
    /// assert_eq!(empty_rw.is_empty(), true);
    ///
    /// let non_empty_rw: RouletteWheel<_> = [(0.1, 10), (0.2, 15), (0.5, 20)].iter().cloned().collect();
    ///
    /// assert_eq!(non_empty_rw.is_empty(), false);
    /// ```
    pub fn is_empty(&self) -> bool {
        self.population.is_empty()
    }

    /// Remove all elements in this wheel.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw: RouletteWheel<_> = [(0.1, 10), (0.2, 15), (0.5, 20)].iter().cloned().collect();
    ///
    /// assert_eq!(rw.len(), 3);
    ///
    /// rw.clear();
    ///
    /// assert_eq!(rw.len(), 0);
    /// ```
    pub fn clear(&mut self) {
        self.fitnesses.clear();
        self.population.clear();
    }

    /// Add an element associated with a probability.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::new();
    ///
    /// rw.push(1.0, 'r');
    /// rw.push(1.0, 'c');
    /// rw.push(1.0, 'a');
    ///
    /// assert_eq!(rw.len(), 3);
    /// ```
    pub fn push(&mut self, fitness: f32, individual: T) {
        assert!(fitness >= 0.0, "Can't push the less than zero fitness: {:?}", fitness);
        assert!((self.total_fitness + fitness).is_finite(), "Fitnesses sum reached a non-finite value!");
        unsafe { self.unchecked_push(fitness, individual) }
    }

    /// Add an element associated with a probability.
    /// This unsafe function doesn't check for total fitness overflow
    /// nether fitness positivity.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::new();
    ///
    /// unsafe { rw.unchecked_push(1.0, 'r') };
    /// unsafe { rw.unchecked_push(1.0, 'c') };
    /// unsafe { rw.unchecked_push(1.0, 'a') };
    ///
    /// assert_eq!(rw.len(), 3);
    /// ```
    pub unsafe fn unchecked_push(&mut self, fitness: f32, individual: T) {
        self.total_fitness += fitness;
        self.fitnesses.push(fitness);
        self.population.push(individual);
    }

    /// Returns the sum of all individual fitnesses.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::new();
    ///
    /// rw.push(3.0, 'r');
    /// rw.push(2.0, 'c');
    /// rw.push(1.5, 'a');
    ///
    /// assert_eq!(rw.total_fitness(), 6.5);
    /// ```
    pub fn total_fitness(&self) -> f32 {
        self.total_fitness
    }

    pub fn select_iter(&self) -> SelectIter<ThreadRng, T> {
        SelectIter::<ThreadRng, _>::new(&self)
    }
}

pub struct SelectIter<'a, R: Rng, T: 'a> {
    distribution_range: Range<f32>,
    rng: R,
    total_fitness: f32,
    fitnesses_ids: Vec<(usize, f32)>,
    roulette_wheel: &'a RouletteWheel<T>
}

impl<'a, R: Rng, T> SelectIter<'a, R, T> {
    pub fn new(roulette_wheel: &'a RouletteWheel<T>) -> SelectIter<'a, ThreadRng, T> {
        SelectIter::from_rng(roulette_wheel, thread_rng())
    }

    pub fn from_rng(roulette_wheel: &'a RouletteWheel<T>, rng: R) -> SelectIter<'a, R, T> {
        SelectIter {
            distribution_range: Range::new(0.0, 1.0),
            rng: rng,
            total_fitness: roulette_wheel.total_fitness,
            fitnesses_ids: roulette_wheel.fitnesses.iter().cloned().enumerate().collect(),
            roulette_wheel: roulette_wheel
        }
    }
}

impl<'a, R: Rng, T: 'a> Iterator for SelectIter<'a, R, T> {
    type Item = (f32, &'a T);

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.fitnesses_ids.len(), Some(self.fitnesses_ids.len()))
    }

    fn next(&mut self) -> Option<Self::Item> {
        if !self.fitnesses_ids.is_empty() {
            let sample = self.distribution_range.ind_sample(&mut self.rng);
            let mut selection = sample * self.total_fitness;
            let index = self.fitnesses_ids.iter().position(|&(_, fit)| {
                            selection -= fit;
                            selection <= 0.0
                        }).unwrap();
            let (index, fitness) = self.fitnesses_ids.swap_remove(index);
            self.total_fitness -= fitness;
            Some((fitness, &self.roulette_wheel.population[index]))
        }
        else { None }
    }
}

impl<T> IntoIterator for RouletteWheel<T> {
    type Item = (f32, T);
    type IntoIter = IntoSelectIter<ThreadRng, T>;

    fn into_iter(self) -> IntoSelectIter<ThreadRng, T> {
        IntoSelectIter::<ThreadRng, _>::new(self)
    }
}

pub struct IntoSelectIter<R: Rng, T> {
    distribution_range: Range<f32>,
    rng: R,
    total_fitness: f32,
    fitnesses: Vec<f32>,
    population: Vec<T>
}

impl<R: Rng, T> IntoSelectIter<R, T> {
    pub fn new(roulette_wheel: RouletteWheel<T>) -> IntoSelectIter<ThreadRng, T> {
        IntoSelectIter::from_rng(roulette_wheel, thread_rng())
    }

    pub fn from_rng(roulette_wheel: RouletteWheel<T>, rng: R) -> IntoSelectIter<R, T> {
        IntoSelectIter {
            distribution_range: Range::new(0.0, 1.0),
            rng: rng,
            total_fitness: roulette_wheel.total_fitness,
            fitnesses: roulette_wheel.fitnesses,
            population: roulette_wheel.population
        }
    }
}

impl<R: Rng, T> Iterator for IntoSelectIter<R, T> {
    type Item = (f32, T);

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.fitnesses.len(), Some(self.fitnesses.len()))
    }

    fn next(&mut self) -> Option<Self::Item> {
        if !self.fitnesses.is_empty() {
            let sample = self.distribution_range.ind_sample(&mut self.rng);
            let mut selection = sample * self.total_fitness;
            let index = self.fitnesses.iter().position(|fit| {
                            selection -= *fit;
                            selection <= 0.0
                        }).unwrap();
            let fitness = self.fitnesses.swap_remove(index);
            let individual = self.population.swap_remove(index);
            self.total_fitness -= fitness;
            Some((fitness, individual))
        }
        else { None }
    }
}

#[cfg(test)]
mod tests {
    use rand::SeedableRng;
    use rand::StdRng;
    use {RouletteWheel, SelectIter, IntoSelectIter};

    const SEED: [usize; 4] = [4, 2, 42, 4242];

    #[test]
    fn select_iter_seeded() {
        let rng = StdRng::from_seed(&SEED);

        let fitnesses = [0.1, 0.2, 0.3, 0.4, 0.5];
        let fitnesses = fitnesses.iter().cloned();
        let population = 15..20;
        let rw: RouletteWheel<_> = fitnesses.zip(population).collect();

        let mut iter = SelectIter::from_rng(&rw, rng);

        assert_eq!(iter.next(), Some((0.5, &19)));
        assert_eq!(iter.next(), Some((0.3, &17)));
        assert_eq!(iter.next(), Some((0.4, &18)));
        assert_eq!(iter.next(), Some((0.2, &16)));
        assert_eq!(iter.next(), Some((0.1, &15)));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn into_select_iter_seeded() {
        let rng = StdRng::from_seed(&SEED);

        let fitnesses = [0.1, 0.2, 0.3, 0.4, 0.5];
        let fitnesses = fitnesses.iter().cloned();
        let population = 15..20;
        let rw: RouletteWheel<_> = fitnesses.zip(population).collect();

        let mut iter = IntoSelectIter::from_rng(rw, rng);

        assert_eq!(iter.next(), Some((0.5, 19)));
        assert_eq!(iter.next(), Some((0.3, 17)));
        assert_eq!(iter.next(), Some((0.4, 18)));
        assert_eq!(iter.next(), Some((0.2, 16)));
        assert_eq!(iter.next(), Some((0.1, 15)));
        assert_eq!(iter.next(), None);
    }
}
<commit_msg>tests: Add a .len() test<commit_after>//! A Little implementation of the roulette-wheel principle, `RouletteWheel<T>`.
//! https://wikipedia.org/wiki/Fitness_proportionate_selection
//!
//! ![Fitness proportionate selection](https://upload.wikimedia.org/wikipedia/commons/2/2a/Fitness_proportionate_selection_example.png)
//!
//! # Usage
//!
//! # Examples
//!
//! You can explicitly create a `RouletteWheel<T>` with `new()`:
//!
//! ```
//! use roulette_wheel::RouletteWheel;
//!
//! let rw: RouletteWheel<u8> = RouletteWheel::new();
//! ```
//!
//! You can `push` values in the roulette-wheel (which will grow the wheel as needed):
//! ```
//! use roulette_wheel::RouletteWheel;
//!
//! let mut rw = RouletteWheel::new();
//!
//! rw.push(5., 'a');
//! rw.push(1., 'b');
//!
//! // .iter() will not consume the roulette wheel
//! for (fit, &ind) in rw.iter() {
//!    // do things with individuals here
//! }
//!
//! // .into_iter() consume the roulette wheel
//! for (fit, ind) in rw.into_iter() {
//!     // do things with individuals here
//! }
//!
//! ```

extern crate rand;

use std::iter::{FromIterator, Iterator, IntoIterator};
use rand::{Rng, ThreadRng, thread_rng};
use rand::distributions::{Range, IndependentSample};

/// A roulette-wheel container
pub struct RouletteWheel<T> {
    total_fitness: f32,
    fitnesses: Vec<f32>,
    population: Vec<T>
}

impl<T: Clone> Clone for RouletteWheel<T> {
    fn clone(&self) -> RouletteWheel<T> {
        RouletteWheel {
            total_fitness: self.total_fitness,
            fitnesses: self.fitnesses.clone(),
            population: self.population.clone()
        }
    }
}

impl<T> FromIterator<(f32, T)> for RouletteWheel<T> {
    fn from_iter<A>(iter: A) -> Self where A: IntoIterator<Item=(f32, T)> {
        let iter = iter.into_iter();
        let (lower, _) = iter.size_hint();

        let mut total_fitness = 0.0;
        let mut fitnesses = Vec::with_capacity(lower);
        let mut population = Vec::with_capacity(lower);

        for (fitness, individual) in iter {
            total_fitness += fitness;
            fitnesses.push(fitness);
            population.push(individual);
        }

        RouletteWheel {
            total_fitness: total_fitness,
            fitnesses: fitnesses,
            population: population
        }

        // #![feature(field_init_shorthand)]
        // RouletteWheel { total_fitness, fitnesses, population }
    }
}

impl<T> RouletteWheel<T> {
    /// create a new empty random-wheel.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let rw = RouletteWheel::<u8>::new();
    /// ```
    pub fn new() -> RouletteWheel<T> {
        RouletteWheel {
            total_fitness: 0.0,
            fitnesses: Vec::new(),
            population: Vec::new()
        }
    }

    /// Creates an empty RouletteWheel with space for at least n elements.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let rw = RouletteWheel::<u8>::with_capacity(15);
    ///
    /// assert_eq!(rw.len(), 0);
    /// ```
    pub fn with_capacity(cap: usize) -> RouletteWheel<T> {
        RouletteWheel {
            total_fitness: 0.0,
            fitnesses: Vec::with_capacity(cap),
            population: Vec::with_capacity(cap)
        }
    }

    /// Reserves capacity for at least `additional` more elements to be inserted.
    /// The collection may reserve more space to avoid frequent reallocations.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::<u8>::new();
    /// rw.reserve(20);
    ///
    /// assert_eq!(rw.len(), 0);
    /// ```
    pub fn reserve(&mut self, additional: usize) {
        self.fitnesses.reserve(additional);
        self.population.reserve(additional);
    }

    /// Returns the number of elements in the wheel.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let rw: RouletteWheel<_> = [(0.1, 10), (0.2, 15), (0.5, 20)].iter().cloned().collect();
    ///
    /// assert_eq!(rw.len(), 3);
    /// ```
    pub fn len(&self) -> usize {
        self.population.len()
    }

    /// Returns `true` if empty else return `false`.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let empty_rw = RouletteWheel::<u8>::new();
    ///
    /// assert_eq!(empty_rw.is_empty(), true);
    ///
    /// let non_empty_rw: RouletteWheel<_> = [(0.1, 10), (0.2, 15), (0.5, 20)].iter().cloned().collect();
    ///
    /// assert_eq!(non_empty_rw.is_empty(), false);
    /// ```
    pub fn is_empty(&self) -> bool {
        self.population.is_empty()
    }

    /// Remove all elements in this wheel.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw: RouletteWheel<_> = [(0.1, 10), (0.2, 15), (0.5, 20)].iter().cloned().collect();
    ///
    /// assert_eq!(rw.len(), 3);
    ///
    /// rw.clear();
    ///
    /// assert_eq!(rw.len(), 0);
    /// ```
    pub fn clear(&mut self) {
        self.fitnesses.clear();
        self.population.clear();
    }

    /// Add an element associated with a probability.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::new();
    ///
    /// rw.push(1.0, 'r');
    /// rw.push(1.0, 'c');
    /// rw.push(1.0, 'a');
    ///
    /// assert_eq!(rw.len(), 3);
    /// ```
    pub fn push(&mut self, fitness: f32, individual: T) {
        assert!(fitness >= 0.0, "Can't push the less than zero fitness: {:?}", fitness);
        assert!((self.total_fitness + fitness).is_finite(), "Fitnesses sum reached a non-finite value!");
        unsafe { self.unchecked_push(fitness, individual) }
    }

    /// Add an element associated with a probability.
    /// This unsafe function doesn't check for total fitness overflow
    /// nether fitness positivity.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::new();
    ///
    /// unsafe { rw.unchecked_push(1.0, 'r') };
    /// unsafe { rw.unchecked_push(1.0, 'c') };
    /// unsafe { rw.unchecked_push(1.0, 'a') };
    ///
    /// assert_eq!(rw.len(), 3);
    /// ```
    pub unsafe fn unchecked_push(&mut self, fitness: f32, individual: T) {
        self.total_fitness += fitness;
        self.fitnesses.push(fitness);
        self.population.push(individual);
    }

    /// Returns the sum of all individual fitnesses.
    /// # Example
    ///
    /// ```
    /// use roulette_wheel::RouletteWheel;
    ///
    /// let mut rw = RouletteWheel::new();
    ///
    /// rw.push(3.0, 'r');
    /// rw.push(2.0, 'c');
    /// rw.push(1.5, 'a');
    ///
    /// assert_eq!(rw.total_fitness(), 6.5);
    /// ```
    pub fn total_fitness(&self) -> f32 {
        self.total_fitness
    }

    pub fn select_iter(&self) -> SelectIter<ThreadRng, T> {
        SelectIter::<ThreadRng, _>::new(&self)
    }
}

pub struct SelectIter<'a, R: Rng, T: 'a> {
    distribution_range: Range<f32>,
    rng: R,
    total_fitness: f32,
    fitnesses_ids: Vec<(usize, f32)>,
    roulette_wheel: &'a RouletteWheel<T>
}

impl<'a, R: Rng, T> SelectIter<'a, R, T> {
    pub fn new(roulette_wheel: &'a RouletteWheel<T>) -> SelectIter<'a, ThreadRng, T> {
        SelectIter::from_rng(roulette_wheel, thread_rng())
    }

    pub fn from_rng(roulette_wheel: &'a RouletteWheel<T>, rng: R) -> SelectIter<'a, R, T> {
        SelectIter {
            distribution_range: Range::new(0.0, 1.0),
            rng: rng,
            total_fitness: roulette_wheel.total_fitness,
            fitnesses_ids: roulette_wheel.fitnesses.iter().cloned().enumerate().collect(),
            roulette_wheel: roulette_wheel
        }
    }
}

impl<'a, R: Rng, T: 'a> Iterator for SelectIter<'a, R, T> {
    type Item = (f32, &'a T);

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.fitnesses_ids.len(), Some(self.fitnesses_ids.len()))
    }

    fn next(&mut self) -> Option<Self::Item> {
        if !self.fitnesses_ids.is_empty() {
            let sample = self.distribution_range.ind_sample(&mut self.rng);
            let mut selection = sample * self.total_fitness;
            let index = self.fitnesses_ids.iter().position(|&(_, fit)| {
                            selection -= fit;
                            selection <= 0.0
                        }).unwrap();
            let (index, fitness) = self.fitnesses_ids.swap_remove(index);
            self.total_fitness -= fitness;
            Some((fitness, &self.roulette_wheel.population[index]))
        }
        else { None }
    }
}

impl<T> IntoIterator for RouletteWheel<T> {
    type Item = (f32, T);
    type IntoIter = IntoSelectIter<ThreadRng, T>;

    fn into_iter(self) -> IntoSelectIter<ThreadRng, T> {
        IntoSelectIter::<ThreadRng, _>::new(self)
    }
}

pub struct IntoSelectIter<R: Rng, T> {
    distribution_range: Range<f32>,
    rng: R,
    total_fitness: f32,
    fitnesses: Vec<f32>,
    population: Vec<T>
}

impl<R: Rng, T> IntoSelectIter<R, T> {
    pub fn new(roulette_wheel: RouletteWheel<T>) -> IntoSelectIter<ThreadRng, T> {
        IntoSelectIter::from_rng(roulette_wheel, thread_rng())
    }

    pub fn from_rng(roulette_wheel: RouletteWheel<T>, rng: R) -> IntoSelectIter<R, T> {
        IntoSelectIter {
            distribution_range: Range::new(0.0, 1.0),
            rng: rng,
            total_fitness: roulette_wheel.total_fitness,
            fitnesses: roulette_wheel.fitnesses,
            population: roulette_wheel.population
        }
    }
}

impl<R: Rng, T> Iterator for IntoSelectIter<R, T> {
    type Item = (f32, T);

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.fitnesses.len(), Some(self.fitnesses.len()))
    }

    fn next(&mut self) -> Option<Self::Item> {
        if !self.fitnesses.is_empty() {
            let sample = self.distribution_range.ind_sample(&mut self.rng);
            let mut selection = sample * self.total_fitness;
            let index = self.fitnesses.iter().position(|fit| {
                            selection -= *fit;
                            selection <= 0.0
                        }).unwrap();
            let fitness = self.fitnesses.swap_remove(index);
            let individual = self.population.swap_remove(index);
            self.total_fitness -= fitness;
            Some((fitness, individual))
        }
        else { None }
    }
}

#[cfg(test)]
mod tests {
    use rand::SeedableRng;
    use rand::StdRng;
    use {RouletteWheel, SelectIter, IntoSelectIter};

    const SEED: [usize; 4] = [4, 2, 42, 4242];

    #[test]
    fn select_iter_seeded() {
        let rng = StdRng::from_seed(&SEED);

        let fitnesses = [0.1, 0.2, 0.3, 0.4, 0.5];
        let fitnesses = fitnesses.iter().cloned();
        let population = 15..20;
        let rw: RouletteWheel<_> = fitnesses.zip(population).collect();

        let mut iter = SelectIter::from_rng(&rw, rng);

        assert_eq!(iter.next(), Some((0.5, &19)));
        assert_eq!(iter.next(), Some((0.3, &17)));
        assert_eq!(iter.next(), Some((0.4, &18)));
        assert_eq!(iter.next(), Some((0.2, &16)));
        assert_eq!(iter.next(), Some((0.1, &15)));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn into_select_iter_seeded() {
        let rng = StdRng::from_seed(&SEED);

        let fitnesses = [0.1, 0.2, 0.3, 0.4, 0.5];
        let fitnesses = fitnesses.iter().cloned();
        let population = 15..20;
        let rw: RouletteWheel<_> = fitnesses.zip(population).collect();

        let mut iter = IntoSelectIter::from_rng(rw, rng);

        assert_eq!(iter.next(), Some((0.5, 19)));
        assert_eq!(iter.next(), Some((0.3, 17)));
        assert_eq!(iter.next(), Some((0.4, 18)));
        assert_eq!(iter.next(), Some((0.2, 16)));
        assert_eq!(iter.next(), Some((0.1, 15)));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_len() {
        let mut rw = RouletteWheel::<u8>::new();

        assert_eq!(rw.len(), 0);

        rw.push(0.1, 1);
        rw.push(0.1, 1);
        rw.push(0.1, 1);
        rw.push(0.1, 1);

        assert_eq!(rw.len(), 4);
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(all(test, feature = "unstable"), feature(test))]
#![cfg_attr(all(test, feature = "unstable"), feature(iter_arith))]

#[cfg(all(test, feature = "unstable"))] extern crate test;

#[macro_use]
extern crate either;

extern crate num;
extern crate rand;

use either::{Either, Left, Right};
use num::Float;
use rand::{Rng, XorShiftRng, weak_rng};
use std::mem::replace;

/// A set of possible criteria to decide when to stop the optimization.
pub enum StopCriterion {
    /// Stop if a specified number of function evaluations has been reached.
    MaxFunctionEvaluations(usize),

    /// Stop if a specified number of iterations has been reached.
    MaxIterations(usize),

    /// Stop if a certain target value range has been reached.
    ///
    /// The first value represents the desired **target value**, while the second one
    /// specifies an **absolute error** which is additionally applied to both directions.
    TargetValue(f64, f64)
}

/// Provide means to adjust the main paramters driving the algorithm.
pub struct Config {
    foods: usize,
    max_trials: usize,
    max_evals: Option<usize>,
    max_iters: Option<usize>,
    value_range: Option<(f64, f64)>
}

impl Config {
    /// Create a new `Config` which describes the behavior of the ABC algorithm.
    ///
    /// By default, this configuration provides sane defaults, i.e., it uses 20 food
    /// sources, a limit of 100 trials per food source
    pub fn new() -> Config {
        Config {
            foods: 20,
            max_trials: 100,
            max_evals: None,
            max_iters: None,
            value_range: None
        }
    }

    /// Set the number of food sources.
    pub fn foods(&mut self, foods: usize) -> &mut Config {
        assert!(foods > 0, "There must be at least 1 food source");
        self.foods = foods;
        self
    }

    /// Set the number of maximum trials per food source, bevore abondining the
    /// source.
    pub fn max_trials(&mut self, max_trials: usize) -> &mut Config {
        self.max_trials = max_trials;
        self
    }

    /// Enable the supplied stop criterion for future optimizations.
    pub fn stop(&mut self, criterion: StopCriterion) -> &mut Config {
        match criterion {
            MaxFunctionEvaluations(v) => self.max_evals = Some(v),
            MaxIterations(v) => self.max_iters = Some(v),
            TargetValue(v, e) => self.value_range = Some((v - e, v + e))
        }
        self
    }

    /// Optimize the given objective function by using the ABC algorithm with
    /// the specified parameters.
    pub fn optimize<F>(&self, dimensions: usize, low: f64, high: f64, objective: F)
        -> Solution
        where //X: Float,
              //Y: NumCast + Zero + PartialOrd + Clone,
              F: Fn(&[f64]) -> f64
    {
        assert!(dimensions > 0, "`dimensions` must be greater than zero");
        assert!(low.is_finite(), "`low` must be a finite number");
        assert!(high.is_finite(), "`high` must be a finite number");
        assert!(low < high, "`low` must be less than `high`");

        Optimizer::new(self, dimensions, low, high, &objective).run().left().unwrap()
    }
}

struct Optimizer<'a> {
    rng: XorShiftRng,
    config: &'a Config,
    dimensions: usize,
    low: f64,
    high: f64,
    objective: &'a Fn(&[f64]) -> f64,
    candidates: Vec<Candidate>,
    optimum: Option<Candidate>,
    evals: usize,
    iter: usize
}

impl<'a> Optimizer<'a> {
    fn new(config: &'a Config, dimensions: usize, low: f64, high: f64,
        objective: &'a Fn(&[f64]) -> f64) -> Optimizer<'a>
    {
        Optimizer {
            rng: weak_rng(),
            config: config,
            dimensions: dimensions,
            low: low,
            high: high,
            objective: objective,
            candidates: Vec::with_capacity(config.foods),
            optimum: None,
            evals: 0,
            iter: 0
        }
    }

    fn run(mut self) -> Either<Solution, ()> {
        self = try_right!(self.initialize());

        loop {
            self = try_right!(self.send_employed_bees());
            self = try_right!(self.send_onlooker_bees());
            self = try_right!(self.send_scout_bees());

            if self.config.max_iters.map_or(false, |max_iters| max_iters == self.iter + 1) {
                return Left(self.finish());
            }

            self.iter += 1;
        }
    }

    fn initialize(mut self) -> Either<Solution, Self> {
        for _ in 0..self.config.foods {
            self = try_right!(self.new_candidate(None));
        }

        Right(self)
    }

    fn send_employed_bees(mut self) -> Either<Solution, Self> {
        for i in 0..self.config.foods {
            self = try_right!(self.mutate_candidate(i));
        }

        // calculating probabilities
        let max_fitness = self.candidates.iter().fold(self.candidates[0].fitness,
            |acc, candidate| acc.max(candidate.fitness));

        for candidate in &mut self.candidates {
            candidate.probability = 0.9 * candidate.fitness / max_fitness + 0.1;
        }

        Right(self)
    }

    fn send_onlooker_bees(mut self) -> Either<Solution, Self> {
        let mut s = 0;
        let mut t = 0;

        while t < self.config.foods {
            if self.rng.gen::<f64>() < self.candidates[s].probability {
                t += 1;

                self = try_right!(self.mutate_candidate(s));
            }

            s += 1;
            if s == self.config.foods {
                s = 0;
            }
        }

        Right(self)
    }

    fn send_scout_bees(mut self) -> Either<Solution, Self> {
        let max_trials = self.candidates.iter().map(|c| c.trials).max().unwrap();

        if max_trials >= self.config.max_trials {
            for i in 0..self.config.foods {
                if self.candidates[i].trials == max_trials {
                    self = try_right!(self.new_candidate(Some(i)));
                }
            }
        }

        Right(self)
    }

    fn new_candidate(mut self, idx: Option<usize>) -> Either<Solution, Self> {
        let xs: Vec<_> = (0..self.dimensions).map(|_| {
            let x = self.rng.gen_range(self.low, self.high);
            assert!(x.is_finite(), "New random value is not finite!");
            x
            }).collect();
        let y = (self.objective)(&xs);
        let fitness = calc_fitness(y);

        self.evals += 1;

        let candidate = Candidate {
            xs: xs,
            y: y,
            fitness: fitness,
            probability: 0.0,
            trials: 0
        };

        let i = match idx {
            Some(i) => {
                self.candidates[i] = candidate;
                i
            },
            None => {
                self.candidates.push(candidate);
                self.candidates.len() - 1
            }
        };

        self.evaluate_candidate(i)
    }

    fn mutate_candidate(mut self, index: usize) -> Either<Solution, Self> {
        let rand_value = self.rng.gen_range(0, self.dimensions);
        let rand_candidate = (0..).map(|_| self.rng.gen_range(0, self.config.foods))
            .filter(|&other| other != index)
            .next().unwrap();

        let new_value = (self.candidates[index].xs[rand_value] +
            (self.candidates[index].xs[rand_value] - self.candidates[rand_candidate].xs[rand_value]) *
            (self.rng.gen::<f64>() - 0.5) * 2.0)
            .min(self.high).max(self.low);

        assert!(new_value.is_finite(), "Mutated value is not finite!");

        let old_value = replace(&mut self.candidates[index].xs[rand_value], new_value);

        let y = (self.objective)(&(self.candidates[index].xs));
        let fitness = calc_fitness(y);

        self.evals += 1;

        {
            let candidate = &mut self.candidates[index];
            if candidate.fitness < fitness {
                candidate.y = y;
                candidate.fitness = fitness;
                candidate.trials = 0;
            } else {
                candidate.xs[rand_value] = old_value;
                candidate.trials += 1;
            }
        }

        self.evaluate_candidate(index)
    }

    fn evaluate_candidate(mut self, idx: usize) -> Either<Solution, Self> {
        {
            let candidate = &self.candidates[idx];

            // upldate global optimum, if necessary
            if self.optimum.as_ref().map_or(true, |optimum| candidate.y < optimum.y) {
                self.optimum = Some(candidate.clone());
            }
        }

        // abort if the maximum
        if self.config.max_evals.map_or(false, |max_evals| max_evals == self.evals) ||
            self.config.value_range.map_or(false, |(l, h)| self.optimum.as_ref().map_or(false,
                |optimum| l <= optimum.y && optimum.y <= h))
        {
            return Left(self.finish());
        }

        Right(self)
    }

    fn finish(self) -> Solution {
        let optimum = self.optimum.unwrap();

        Solution {
            xs: optimum.xs,
            y: optimum.y,
            fitness: optimum.fitness,
            evals: self.evals,
            iters: self.iter + 1
        }
    }
}

#[derive(Clone, Debug)]
struct Candidate {
    xs: Vec<f64>,
    y: f64,
    fitness: f64,
    probability: f64,
    trials: usize
}

#[derive(Debug)]
pub struct Solution {
    pub xs: Vec<f64>,
    pub y: f64,
    pub fitness: f64,
    pub evals: usize,
    pub iters: usize
}

#[inline(always)]
fn calc_fitness(y: f64) -> f64 {
    if y >= 0.0 {
        1.0 / (1.0 + y)
    } else {
        1.0 + y.abs()
    }
}

// re-exports
pub use StopCriterion::*;

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use test::Bencher;

    use super::*;
    use super::problems::*;

    macro_rules! benches {
        ($($name: ident, $d: expr, $l: expr, $u: expr, $y: expr, $e: expr, $xs: expr, $f: ident;)*) => {
            $(
                #[bench]
                fn $name(b: &mut Bencher) {
                    let optimum = $xs;
                    assert!($d == optimum.len());

                    let mut abc = Config::new();
                    abc.stop(TargetValue($y, $e));

                    b.iter(|| {
                        let solution = abc.optimize($d, $l, $u, &$f);

                        for i in 0..$d {
                            assert!((solution.xs[i] - optimum[i]).abs() < 1.0e-5);
                        }

                        solution
                    })
                }
            )*
        }
    }

    benches! {
        sphere_01d, 1, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 1], sphere;
        sphere_02d, 2, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 2], sphere;
        sphere_04d, 4, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 4], sphere;
        sphere_08d, 8, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 8], sphere;
        sphere_16d, 16, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 16], sphere;

        ackley_01d, 1, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 1], ackley;
        ackley_02d, 2, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 2], ackley;
        ackley_04d, 4, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 4], ackley;
        ackley_08d, 8, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 8], ackley;
        ackley_16d, 16, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 16], ackley;

        booth_2d, 2, -10.0, 10.0, 0.0, 1.0e-10, vec![1.0, 3.0], booth;
    }
}

#[cfg(all(test, feature = "unstable"))]
mod problems {
    use std::f64::consts::PI;

    // bowl-shaped
    pub fn sphere(xs: &[f64]) -> f64 {
        xs.iter().map(|v| v * v).sum()
    }

    // many local minima
    pub fn ackley(xs: &[f64]) -> f64 {
        let a = 20.0;
        let b = 0.2;
        let c = 2.0 * PI;

        -a * (-b * (xs.iter().map(|&x| x.powi(2)).sum::<f64>() / xs.len() as f64).sqrt()).exp() -
            (xs.iter().map(|&x| (c*x).cos()).sum::<f64>() / xs.len() as f64).exp() + a + 1.0f64.exp()
    }

    // plate-shaped
    pub fn booth(xs: &[f64]) -> f64 {
        assert!(xs.len() == 2);

        (xs[0] + 2.0*xs[1] - 7.0).powi(2) + (2.0*xs[0] + xs[1] - 5.0).powi(2)
    }
}
<commit_msg>Add documentation and remove obsolete code<commit_after>#![cfg_attr(all(test, feature = "unstable"), feature(test))]
#![cfg_attr(all(test, feature = "unstable"), feature(iter_arith))]

#[cfg(all(test, feature = "unstable"))] extern crate test;

#[macro_use]
extern crate either;

extern crate num;
extern crate rand;

use either::{Either, Left, Right};
use num::Float;
use rand::{Rng, XorShiftRng, weak_rng};
use std::mem::replace;

/// A set of possible criteria to decide when to stop the optimization.
pub enum StopCriterion {
    /// Stop if a specified number of function evaluations has been reached.
    MaxFunctionEvaluations(usize),

    /// Stop if a specified number of iterations has been reached.
    MaxIterations(usize),

    /// Stop if a certain target value range has been reached.
    ///
    /// The first value represents the desired **target value**, while the second one
    /// specifies an **absolute error** which is additionally applied to both directions.
    TargetValue(f64, f64)
}

/// Groups information relevant to a solution.
#[derive(Clone, Debug)]
pub struct Solution {
    /// The set of parameters supplied to the objective function.
    pub xs: Vec<f64>,

    /// Target value for the set of parameters.
    pub y: f64,

    /// Number of function evaluations.
    pub evals: usize,

    /// Number of iterations.
    pub iters: usize
}

/// Provide means to adjust the main paramters driving the algorithm.
pub struct Config {
    foods: usize,
    max_trials: usize,
    max_evals: Option<usize>,
    max_iters: Option<usize>,
    value_range: Option<(f64, f64)>
}

impl Config {
    /// Create a new `Config` which describes the behavior of the ABC algorithm.
    ///
    /// By default, this configuration provides sane defaults, i.e., it uses 20 food
    /// sources, a limit of 100 trials per food source
    pub fn new() -> Config {
        Config {
            foods: 20,
            max_trials: 100,
            max_evals: None,
            max_iters: None,
            value_range: None
        }
    }

    /// Set the number of food sources.
    pub fn foods(&mut self, foods: usize) -> &mut Config {
        assert!(foods > 0, "There must be at least 1 food source");
        self.foods = foods;
        self
    }

    /// Set the number of maximum trials per food source, bevore abondining the
    /// source.
    pub fn max_trials(&mut self, max_trials: usize) -> &mut Config {
        self.max_trials = max_trials;
        self
    }

    /// Enable the supplied stop criterion for future optimizations.
    pub fn stop(&mut self, criterion: StopCriterion) -> &mut Config {
        match criterion {
            MaxFunctionEvaluations(v) => self.max_evals = Some(v),
            MaxIterations(v) => self.max_iters = Some(v),
            TargetValue(v, e) => self.value_range = Some((v - e, v + e))
        }
        self
    }

    /// Optimize the given objective function by using the ABC algorithm with
    /// the specified parameters.
    pub fn optimize<F>(&self, dimensions: usize, low: f64, high: f64, objective: F) -> Solution
        where F: Fn(&[f64]) -> f64
    {
        assert!(dimensions > 0, "`dimensions` must be greater than zero");
        assert!(low.is_finite(), "`low` must be a finite number");
        assert!(high.is_finite(), "`high` must be a finite number");
        assert!(low < high, "`low` must be less than `high`");

        Optimizer::new(self, dimensions, low, high, &objective).run().left().unwrap()
    }
}

struct Optimizer<'a> {
    rng: XorShiftRng,
    config: &'a Config,
    dimensions: usize,
    low: f64,
    high: f64,
    objective: &'a Fn(&[f64]) -> f64,
    candidates: Vec<Candidate>,
    optimum: Option<Candidate>,
    evals: usize,
    iter: usize
}

impl<'a> Optimizer<'a> {
    fn new(config: &'a Config, dimensions: usize, low: f64, high: f64,
        objective: &'a Fn(&[f64]) -> f64) -> Optimizer<'a>
    {
        Optimizer {
            rng: weak_rng(),
            config: config,
            dimensions: dimensions,
            low: low,
            high: high,
            objective: objective,
            candidates: Vec::with_capacity(config.foods),
            optimum: None,
            evals: 0,
            iter: 0
        }
    }

    fn run(mut self) -> Either<Solution, ()> {
        self = try_right!(self.initialize());

        loop {
            self = try_right!(self.send_employed_bees());
            self = try_right!(self.send_onlooker_bees());
            self = try_right!(self.send_scout_bees());

            if self.config.max_iters.map_or(false, |max_iters| max_iters == self.iter + 1) {
                return Left(self.finish());
            }

            self.iter += 1;
        }
    }

    fn initialize(mut self) -> Either<Solution, Self> {
        for _ in 0..self.config.foods {
            self = try_right!(self.new_candidate(None));
        }

        Right(self)
    }

    fn send_employed_bees(mut self) -> Either<Solution, Self> {
        for i in 0..self.config.foods {
            self = try_right!(self.mutate_candidate(i));
        }

        // calculating probabilities
        let max_fitness = self.candidates.iter().fold(self.candidates[0].fitness,
            |acc, candidate| acc.max(candidate.fitness));

        for candidate in &mut self.candidates {
            candidate.probability = 0.9 * candidate.fitness / max_fitness + 0.1;
        }

        Right(self)
    }

    fn send_onlooker_bees(mut self) -> Either<Solution, Self> {
        let mut s = 0;
        let mut t = 0;

        while t < self.config.foods {
            if self.rng.gen::<f64>() < self.candidates[s].probability {
                t += 1;

                self = try_right!(self.mutate_candidate(s));
            }

            s += 1;
            if s == self.config.foods {
                s = 0;
            }
        }

        Right(self)
    }

    fn send_scout_bees(mut self) -> Either<Solution, Self> {
        let max_trials = self.candidates.iter().map(|c| c.trials).max().unwrap();

        if max_trials >= self.config.max_trials {
            for i in 0..self.config.foods {
                if self.candidates[i].trials == max_trials {
                    self = try_right!(self.new_candidate(Some(i)));
                }
            }
        }

        Right(self)
    }

    fn new_candidate(mut self, idx: Option<usize>) -> Either<Solution, Self> {
        let xs: Vec<_> = (0..self.dimensions).map(|_| {
            let x = self.rng.gen_range(self.low, self.high);
            assert!(x.is_finite(), "New random value is not finite!");
            x
            }).collect();
        let y = (self.objective)(&xs);
        let fitness = calc_fitness(y);

        self.evals += 1;

        let candidate = Candidate {
            xs: xs,
            y: y,
            fitness: fitness,
            probability: 0.0,
            trials: 0
        };

        let i = match idx {
            Some(i) => {
                self.candidates[i] = candidate;
                i
            },
            None => {
                self.candidates.push(candidate);
                self.candidates.len() - 1
            }
        };

        self.evaluate_candidate(i)
    }

    fn mutate_candidate(mut self, index: usize) -> Either<Solution, Self> {
        let rand_value = self.rng.gen_range(0, self.dimensions);
        let rand_candidate = (0..).map(|_| self.rng.gen_range(0, self.config.foods))
            .filter(|&other| other != index)
            .next().unwrap();

        let new_value = (self.candidates[index].xs[rand_value] +
            (self.candidates[index].xs[rand_value] - self.candidates[rand_candidate].xs[rand_value]) *
            (self.rng.gen::<f64>() - 0.5) * 2.0)
            .min(self.high).max(self.low);

        assert!(new_value.is_finite(), "Mutated value is not finite!");

        let old_value = replace(&mut self.candidates[index].xs[rand_value], new_value);

        let y = (self.objective)(&(self.candidates[index].xs));
        let fitness = calc_fitness(y);

        self.evals += 1;

        {
            let candidate = &mut self.candidates[index];
            if candidate.fitness < fitness {
                candidate.y = y;
                candidate.fitness = fitness;
                candidate.trials = 0;
            } else {
                candidate.xs[rand_value] = old_value;
                candidate.trials += 1;
            }
        }

        self.evaluate_candidate(index)
    }

    fn evaluate_candidate(mut self, idx: usize) -> Either<Solution, Self> {
        {
            let candidate = &self.candidates[idx];

            // upldate global optimum, if necessary
            if self.optimum.as_ref().map_or(true, |optimum| candidate.y < optimum.y) {
                self.optimum = Some(candidate.clone());
            }
        }

        // abort if the maximum
        if self.config.max_evals.map_or(false, |max_evals| max_evals == self.evals) ||
            self.config.value_range.map_or(false, |(l, h)| self.optimum.as_ref().map_or(false,
                |optimum| l <= optimum.y && optimum.y <= h))
        {
            return Left(self.finish());
        }

        Right(self)
    }

    fn finish(self) -> Solution {
        let optimum = self.optimum.unwrap();

        Solution {
            xs: optimum.xs,
            y: optimum.y,
            evals: self.evals,
            iters: self.iter + 1
        }
    }
}

#[derive(Clone, Debug)]
struct Candidate {
    xs: Vec<f64>,
    y: f64,
    fitness: f64,
    probability: f64,
    trials: usize
}

#[inline(always)]
fn calc_fitness(y: f64) -> f64 {
    if y >= 0.0 {
        1.0 / (1.0 + y)
    } else {
        1.0 + y.abs()
    }
}

// re-exports
pub use StopCriterion::*;

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use test::Bencher;

    use super::*;
    use super::problems::*;

    macro_rules! benches {
        ($($name: ident, $d: expr, $l: expr, $u: expr, $y: expr, $e: expr, $xs: expr, $f: ident;)*) => {
            $(
                #[bench]
                fn $name(b: &mut Bencher) {
                    let optimum = $xs;
                    assert!($d == optimum.len());

                    let mut abc = Config::new();
                    abc.stop(TargetValue($y, $e));

                    b.iter(|| {
                        let solution = abc.optimize($d, $l, $u, &$f);

                        for i in 0..$d {
                            assert!((solution.xs[i] - optimum[i]).abs() < 1.0e-5);
                        }

                        solution
                    })
                }
            )*
        }
    }

    benches! {
        sphere_01d, 1, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 1], sphere;
        sphere_02d, 2, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 2], sphere;
        sphere_04d, 4, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 4], sphere;
        sphere_08d, 8, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 8], sphere;
        sphere_16d, 16, -5.12, 5.12, 0.0, 1.0e-10, vec![0.0f64; 16], sphere;

        ackley_01d, 1, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 1], ackley;
        ackley_02d, 2, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 2], ackley;
        ackley_04d, 4, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 4], ackley;
        ackley_08d, 8, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 8], ackley;
        ackley_16d, 16, -32.768, 32.768, 0.0, 1.0e-10, vec![0.0f64; 16], ackley;

        booth_2d, 2, -10.0, 10.0, 0.0, 1.0e-10, vec![1.0, 3.0], booth;
    }
}

#[cfg(all(test, feature = "unstable"))]
mod problems {
    use std::f64::consts::PI;

    // bowl-shaped
    pub fn sphere(xs: &[f64]) -> f64 {
        xs.iter().map(|v| v * v).sum()
    }

    // many local minima
    pub fn ackley(xs: &[f64]) -> f64 {
        let a = 20.0;
        let b = 0.2;
        let c = 2.0 * PI;

        -a * (-b * (xs.iter().map(|&x| x.powi(2)).sum::<f64>() / xs.len() as f64).sqrt()).exp() -
            (xs.iter().map(|&x| (c*x).cos()).sum::<f64>() / xs.len() as f64).exp() + a + 1.0f64.exp()
    }

    // plate-shaped
    pub fn booth(xs: &[f64]) -> f64 {
        assert!(xs.len() == 2);

        (xs[0] + 2.0*xs[1] - 7.0).powi(2) + (2.0*xs[0] + xs[1] - 5.0).powi(2)
    }
}
<|endoftext|>


<commit_before>#![feature(plugin)]
extern crate time;
extern crate curl;
extern crate rustc_serialize;
extern crate openssl;
extern crate url;
#[macro_use]
extern crate log;

use url::Url;
use std::str;
use std::io::prelude::*;

use curl::easy::{Easy, List};

use openssl::crypto::{hmac, hash};
use rustc_serialize::base64::{ToBase64, STANDARD};
use rustc_serialize::hex::ToHex;

pub const LONG_DATE: &'static str = "%Y%m%dT%H%M%SZ";
const SHORT_DATE: &'static str = "%Y%m%d";
const AMZ_EXPIRE: &'static str = "604800";
const AMZ_ALGO: &'static str = "AWS4-HMAC-SHA256";
const AMZ_REGION: &'static str = "eu-west-1";
const AMZ_SERVICE: &'static str = "s3";
const AMZ_PAYLOAD: &'static str = "UNSIGNED-PAYLOAD";
const AMZ_REQ_VER: &'static str = "aws4_request";

pub struct Bucket {
    name: String,
    region: Option<String>,
    access_key: String,
    secret_key: String,
    proto: String,
}

macro_rules! build_headers {
    ($list:ident, $($x:expr),+) => (
      $($list.append($x).unwrap())+
    )
}

macro_rules! unwrap_get {
  ($get:ident) => (
      match $get {
      Some(x) => {
        if x.len() > 0 {
          x
        } else {
          panic!("These are not the droids you are looking for.")
        }
      },
      None => unreachable!()
    }
  )
}

macro_rules! headers {
  ($transfer:ident, $headers:ident) => (
    $transfer.header_function(|header| {
      $headers.push(str::from_utf8(header).unwrap().to_string());
      true
    }).unwrap();
  )
}

enum Command<'a> {
  Put {
    content: &'a [u8]
  },
  Get,
  List {
    prefix: &'a str,
    delimiter: &'a str
  }
}

impl Bucket {
  pub fn new(name: String,
             region: Option<String>,
             access_key: String,
             secret_key: String,
             proto: &str) -> Bucket {
      Bucket {
          name: name,
          region: region,
          access_key: access_key,
          secret_key: secret_key,
          proto: proto.to_string(),
      }
  }

  fn execute(&self,
                 cmd: Command,
                 path: &str ) -> (String, Option<Vec<u8>>) {

    let content_type = "text/plain";

    let path = if path.starts_with("/") {&path[1..]} else {path};
    debug!("s3_path: {}", path);

    let host = self.host();
    debug!("host: {}", host);

    let date = time::now_utc().rfc822z().to_string();
    debug!("date: {}", date);

    let cmd_string = match cmd {
      Command::Put { content } => "PUT",
      _ => "GET"
    };

    debug!("command: {}", cmd_string);
    let a = self.auth(&cmd_string, &date, path, "", &content_type);
    debug!("auth: {}", a);

    // TODO implement delimiter into request for List
    let url = match cmd {
      Command::List { prefix, delimiter } => format!("{}://{}/?prefix={}", self.proto, host, prefix),
      _ => format!("{}://{}/{}", self.proto, host, path)
    };
    debug!("url: {}", url);

    let mut handle = Easy::new();
    let _ = handle.url(&url[..]);

    let mut list = List::new();

    let l = match cmd {
      Command::Put { content } => content.len(),
      _ => 0
    };

    build_headers!(list,
            &format!("Host: {}", &host),
            &format!("Date: {}", &date),
            &format!("Authorization: {}", &a),
            &format!("Content-Type: {}", &content_type),
            &format!("Content-Length: {}", l));

    for l in list.iter(){
      debug!("{:?}", String::from_utf8_lossy(l));
    }

    match cmd {
      Command::Put { content } => {
        handle.put(true).unwrap();
        let _ = handle.post_field_size(l as u64);
      }
      _ => {}
    }

    handle.http_headers(list).unwrap();

    let result = match cmd {
                  Command::Put { content } => self.put(&mut handle, content),
                  Command::Get => self.get(&mut handle),
                  Command::List { delimiter, prefix }=> self.list(&mut handle)
                };
    (url, result)
  }

  pub fn put(&self, handle: &mut Easy, mut content: &[u8]) -> Option<Vec<u8>> {
    let mut headers = Vec::new();
    {
      let mut transfer = handle.transfer();

      transfer.read_function(|buf| {
          Ok(content.read(buf).unwrap())
      }).unwrap();

      headers!(transfer, headers);

      transfer.perform().unwrap();
    }
    debug!("recieved headers: {:?}", headers);
    debug!("response: {}", handle.response_code().unwrap());
    None
  }

  pub fn get(&self, handle: &mut Easy) -> Option<Vec<u8>> {
    let mut headers = Vec::new();
    let mut dst = Vec::new();

    {
      let mut transfer = handle.transfer();
      transfer.write_function(|data| {
          dst.extend_from_slice(data);
          Ok(data.len())
      }).unwrap();

      headers!(transfer, headers);

      transfer.perform().unwrap();
    }
    debug!("recieved headers: {:?}", headers);
    Some(dst)
  }

  pub fn list(&self, handle: &mut Easy) -> Option<Vec<u8>> {
    let mut dst = Vec::new();
    let mut headers = Vec::new();
    {
      let mut transfer = handle.transfer();
      transfer.write_function(|data| {
          dst.extend_from_slice(data);
          Ok(data.len())
      }).unwrap();

      headers!(transfer, headers);

      transfer.perform().unwrap();
    }
    Some(dst)
  }

    // TODO implement DELETE Command

    // pub fn delete<'a, 'b>(&self, handle: &'a mut http::Handle, path: &str)
    //                       -> http::Request<'a, 'b> {
    //     let path = if path.starts_with("/") {&path[1..]} else {path};
    //     let host = self.host();
    //     let date = time::now().rfc822z().to_string();
    //     let auth = self.auth("DELETE", &date, path, "", "");
    //     let url = format!("{}://{}/{}", self.proto, host, path);
    //     handle.delete(&url[..])
    //           .header("Host", &host)
    //           .header("Date", &date)
    //           .header("Authorization", &auth)
    // }

  pub fn host(&self) -> String {
      format!("{}.s3{}.amazonaws.com", self.name,
              match self.region {
                  Some(ref r) => format!("-{}", r),
                  None => String::new(),
              })
  }

  fn auth(&self, verb: &str, date: &str, path: &str,
          md5: &str, content_type: &str) -> String {
      let string = format!("{verb}\n{md5}\n{ty}\n{date}\n{headers}{resource}",
                           verb = verb,
                           md5 = md5,
                           ty = content_type,
                           date = date,
                           headers = "",
                           resource = format!("/{}/{}", self.name, path));
      let signature = {
          let mut hmac = hmac::HMAC::new(hash::Type::SHA1, self.secret_key.as_bytes());
          let _ = hmac.write_all(string.as_bytes());
          hmac.finish().to_base64(STANDARD)
      };
      format!("AWS {}:{}", self.access_key, signature)
  }
}

pub fn get_s3(bucket: &Bucket, s3_path: Option<&str>) -> Vec<u8> {
      let path = match s3_path {
        Some(x) => x,
        None => "/"
      };
      let (_, result) = bucket.execute(Command::Get, path);
      unwrap_get!(result)
    }

// TODO prefix + delimiter support, default delimiter is / ATM
pub fn list_s3(bucket: &Bucket, path: &str, prefix: &str, delimiter: &str) -> Vec<u8> {
      let (_, result) = bucket.execute(Command::List { prefix: prefix, delimiter: delimiter }, path);
      unwrap_get!(result)
    }

fn sign(key: &[u8], msg: &[u8]) -> Vec<u8> {
    let mut hmac = hmac::HMAC::new(hash::Type::SHA256, key);
    let _ = hmac.write_all(msg);
    hmac.finish()
  }

fn get_signature_key(key: &str, date_stamp: &str, region_name: &str, service_name: &str) -> Vec<u8> {
  let kdate = sign(format!("AWS4{}", key).as_bytes(), date_stamp.as_bytes());
  let kregion = sign(kdate.as_slice(), region_name.as_bytes());
  let kservice = sign(kregion.as_slice(), service_name.as_bytes());
  sign(&kservice, AMZ_REQ_VER.as_bytes())
}

pub fn put_s3(bucket: &Bucket,
              s3_path: &str,
              output: &[u8]) -> String {

  let (url, _) = bucket.execute(Command::Put { content: output }, &s3_path);
  let t = time::now();
  let method = &"GET";
  let amzdate = match t.strftime(LONG_DATE) {
    Ok(x) => x.to_string(),
    Err(e) => panic!("{:?}", e)
  };
  let datestamp = match t.strftime(SHORT_DATE) {
    Ok(x) => x.to_string(),
    Err(e) => panic!("{:?}", e)
  };
  let url: Url = match url.parse() {
    Ok(x) => x,
    Err(e) => panic!("{:?}", e)
  };
  let canonical_uri = &url.path();
  let host = match url.host_str() {
    Some(x) => x,
    None => panic!("Unable to extract host string from url: {:?}", url)
  };
  let endpoint = format!("https://{}", host);
  let canonical_headers = format!("host:{}\n", host);
  let signed_headers = &"host";
  let payload_hash = AMZ_PAYLOAD;
  let credential_scope = format!("{}/{}/{}/{}",
                                 datestamp,
                                 AMZ_REGION,
                                 AMZ_SERVICE,
                                 AMZ_REQ_VER);
  let canonical_querystring = format!("{}={}&{}={}%2F{}&{}={}&{}={}&{}={}",
                                &"X-Amz-Algorithm", AMZ_ALGO,
                                &"X-Amz-Credential", bucket.access_key, credential_scope.replace("/", "%2F"),
                                &"X-Amz-Date", &amzdate,
                                &"X-Amz-Expires", AMZ_EXPIRE,
                                &"X-Amz-SignedHeaders", signed_headers);
  let canonical_request = format!("{}\n{}\n{}\n{}\n{}\n{}",
                                  method,
                                  &canonical_uri,
                                  canonical_querystring,
                                  canonical_headers,
                                  signed_headers,
                                  payload_hash);
  let string_to_sign = format!("{}\n{}\n{}\n{}",
                           AMZ_ALGO,
                           amzdate,
                           credential_scope,
                           hash::hash(hash::Type::SHA256, canonical_request.as_bytes()).to_hex());
  let signing_key = get_signature_key(&bucket.secret_key, &datestamp, AMZ_REGION, AMZ_SERVICE);
  let signature = sign(&signing_key, &string_to_sign.as_bytes()).to_hex();
  let canonical_querystring = format!("{}&{}={}",
                                      canonical_querystring,
                                      &"X-Amz-Signature", signature);
  format!("{}{}?{}",
                endpoint,
                canonical_uri,
                canonical_querystring)
}

<commit_msg>Update lib.rs<commit_after>extern crate time;
extern crate curl;
extern crate rustc_serialize;
extern crate openssl;
extern crate url;
#[macro_use]
extern crate log;

use url::Url;
use std::str;
use std::io::prelude::*;

use curl::easy::{Easy, List};

use openssl::crypto::{hmac, hash};
use rustc_serialize::base64::{ToBase64, STANDARD};
use rustc_serialize::hex::ToHex;

pub const LONG_DATE: &'static str = "%Y%m%dT%H%M%SZ";
const SHORT_DATE: &'static str = "%Y%m%d";
const AMZ_EXPIRE: &'static str = "604800";
const AMZ_ALGO: &'static str = "AWS4-HMAC-SHA256";
const AMZ_REGION: &'static str = "eu-west-1";
const AMZ_SERVICE: &'static str = "s3";
const AMZ_PAYLOAD: &'static str = "UNSIGNED-PAYLOAD";
const AMZ_REQ_VER: &'static str = "aws4_request";

pub struct Bucket {
    name: String,
    region: Option<String>,
    access_key: String,
    secret_key: String,
    proto: String,
}

macro_rules! build_headers {
    ($list:ident, $($x:expr),+) => (
      $($list.append($x).unwrap())+
    )
}

macro_rules! unwrap_get {
  ($get:ident) => (
      match $get {
      Some(x) => {
        if x.len() > 0 {
          x
        } else {
          panic!("These are not the droids you are looking for.")
        }
      },
      None => unreachable!()
    }
  )
}

macro_rules! headers {
  ($transfer:ident, $headers:ident) => (
    $transfer.header_function(|header| {
      $headers.push(str::from_utf8(header).unwrap().to_string());
      true
    }).unwrap();
  )
}

enum Command<'a> {
  Put {
    content: &'a [u8]
  },
  Get,
  List {
    prefix: &'a str,
    delimiter: &'a str
  }
}

impl Bucket {
  pub fn new(name: String,
             region: Option<String>,
             access_key: String,
             secret_key: String,
             proto: &str) -> Bucket {
      Bucket {
          name: name,
          region: region,
          access_key: access_key,
          secret_key: secret_key,
          proto: proto.to_string(),
      }
  }

  fn execute(&self,
                 cmd: Command,
                 path: &str ) -> (String, Option<Vec<u8>>) {

    let content_type = "text/plain";

    let path = if path.starts_with("/") {&path[1..]} else {path};
    debug!("s3_path: {}", path);

    let host = self.host();
    debug!("host: {}", host);

    let date = time::now_utc().rfc822z().to_string();
    debug!("date: {}", date);

    let cmd_string = match cmd {
      Command::Put { content } => "PUT",
      _ => "GET"
    };

    debug!("command: {}", cmd_string);
    let a = self.auth(&cmd_string, &date, path, "", &content_type);
    debug!("auth: {}", a);

    // TODO implement delimiter into request for List
    let url = match cmd {
      Command::List { prefix, delimiter } => format!("{}://{}/?prefix={}", self.proto, host, prefix),
      _ => format!("{}://{}/{}", self.proto, host, path)
    };
    debug!("url: {}", url);

    let mut handle = Easy::new();
    let _ = handle.url(&url[..]);

    let mut list = List::new();

    let l = match cmd {
      Command::Put { content } => content.len(),
      _ => 0
    };

    build_headers!(list,
            &format!("Host: {}", &host),
            &format!("Date: {}", &date),
            &format!("Authorization: {}", &a),
            &format!("Content-Type: {}", &content_type),
            &format!("Content-Length: {}", l));

    for l in list.iter(){
      debug!("{:?}", String::from_utf8_lossy(l));
    }

    match cmd {
      Command::Put { content } => {
        handle.put(true).unwrap();
        let _ = handle.post_field_size(l as u64);
      }
      _ => {}
    }

    handle.http_headers(list).unwrap();

    let result = match cmd {
                  Command::Put { content } => self.put(&mut handle, content),
                  Command::Get => self.get(&mut handle),
                  Command::List { delimiter, prefix }=> self.list(&mut handle)
                };
    (url, result)
  }

  pub fn put(&self, handle: &mut Easy, mut content: &[u8]) -> Option<Vec<u8>> {
    let mut headers = Vec::new();
    {
      let mut transfer = handle.transfer();

      transfer.read_function(|buf| {
          Ok(content.read(buf).unwrap())
      }).unwrap();

      headers!(transfer, headers);

      transfer.perform().unwrap();
    }
    debug!("recieved headers: {:?}", headers);
    debug!("response: {}", handle.response_code().unwrap());
    None
  }

  pub fn get(&self, handle: &mut Easy) -> Option<Vec<u8>> {
    let mut headers = Vec::new();
    let mut dst = Vec::new();

    {
      let mut transfer = handle.transfer();
      transfer.write_function(|data| {
          dst.extend_from_slice(data);
          Ok(data.len())
      }).unwrap();

      headers!(transfer, headers);

      transfer.perform().unwrap();
    }
    debug!("recieved headers: {:?}", headers);
    Some(dst)
  }

  pub fn list(&self, handle: &mut Easy) -> Option<Vec<u8>> {
    let mut dst = Vec::new();
    let mut headers = Vec::new();
    {
      let mut transfer = handle.transfer();
      transfer.write_function(|data| {
          dst.extend_from_slice(data);
          Ok(data.len())
      }).unwrap();

      headers!(transfer, headers);

      transfer.perform().unwrap();
    }
    Some(dst)
  }

    // TODO implement DELETE Command

    // pub fn delete<'a, 'b>(&self, handle: &'a mut http::Handle, path: &str)
    //                       -> http::Request<'a, 'b> {
    //     let path = if path.starts_with("/") {&path[1..]} else {path};
    //     let host = self.host();
    //     let date = time::now().rfc822z().to_string();
    //     let auth = self.auth("DELETE", &date, path, "", "");
    //     let url = format!("{}://{}/{}", self.proto, host, path);
    //     handle.delete(&url[..])
    //           .header("Host", &host)
    //           .header("Date", &date)
    //           .header("Authorization", &auth)
    // }

  pub fn host(&self) -> String {
      format!("{}.s3{}.amazonaws.com", self.name,
              match self.region {
                  Some(ref r) => format!("-{}", r),
                  None => String::new(),
              })
  }

  fn auth(&self, verb: &str, date: &str, path: &str,
          md5: &str, content_type: &str) -> String {
      let string = format!("{verb}\n{md5}\n{ty}\n{date}\n{headers}{resource}",
                           verb = verb,
                           md5 = md5,
                           ty = content_type,
                           date = date,
                           headers = "",
                           resource = format!("/{}/{}", self.name, path));
      let signature = {
          let mut hmac = hmac::HMAC::new(hash::Type::SHA1, self.secret_key.as_bytes());
          let _ = hmac.write_all(string.as_bytes());
          hmac.finish().to_base64(STANDARD)
      };
      format!("AWS {}:{}", self.access_key, signature)
  }
}

pub fn get_s3(bucket: &Bucket, s3_path: Option<&str>) -> Vec<u8> {
      let path = match s3_path {
        Some(x) => x,
        None => "/"
      };
      let (_, result) = bucket.execute(Command::Get, path);
      unwrap_get!(result)
    }

// TODO prefix + delimiter support, default delimiter is / ATM
pub fn list_s3(bucket: &Bucket, path: &str, prefix: &str, delimiter: &str) -> Vec<u8> {
      let (_, result) = bucket.execute(Command::List { prefix: prefix, delimiter: delimiter }, path);
      unwrap_get!(result)
    }

fn sign(key: &[u8], msg: &[u8]) -> Vec<u8> {
    let mut hmac = hmac::HMAC::new(hash::Type::SHA256, key);
    let _ = hmac.write_all(msg);
    hmac.finish()
  }

fn get_signature_key(key: &str, date_stamp: &str, region_name: &str, service_name: &str) -> Vec<u8> {
  let kdate = sign(format!("AWS4{}", key).as_bytes(), date_stamp.as_bytes());
  let kregion = sign(kdate.as_slice(), region_name.as_bytes());
  let kservice = sign(kregion.as_slice(), service_name.as_bytes());
  sign(&kservice, AMZ_REQ_VER.as_bytes())
}

pub fn put_s3(bucket: &Bucket,
              s3_path: &str,
              output: &[u8]) -> String {

  let (url, _) = bucket.execute(Command::Put { content: output }, &s3_path);
  let t = time::now();
  let method = &"GET";
  let amzdate = match t.strftime(LONG_DATE) {
    Ok(x) => x.to_string(),
    Err(e) => panic!("{:?}", e)
  };
  let datestamp = match t.strftime(SHORT_DATE) {
    Ok(x) => x.to_string(),
    Err(e) => panic!("{:?}", e)
  };
  let url: Url = match url.parse() {
    Ok(x) => x,
    Err(e) => panic!("{:?}", e)
  };
  let canonical_uri = &url.path();
  let host = match url.host_str() {
    Some(x) => x,
    None => panic!("Unable to extract host string from url: {:?}", url)
  };
  let endpoint = format!("https://{}", host);
  let canonical_headers = format!("host:{}\n", host);
  let signed_headers = &"host";
  let payload_hash = AMZ_PAYLOAD;
  let credential_scope = format!("{}/{}/{}/{}",
                                 datestamp,
                                 AMZ_REGION,
                                 AMZ_SERVICE,
                                 AMZ_REQ_VER);
  let canonical_querystring = format!("{}={}&{}={}%2F{}&{}={}&{}={}&{}={}",
                                &"X-Amz-Algorithm", AMZ_ALGO,
                                &"X-Amz-Credential", bucket.access_key, credential_scope.replace("/", "%2F"),
                                &"X-Amz-Date", &amzdate,
                                &"X-Amz-Expires", AMZ_EXPIRE,
                                &"X-Amz-SignedHeaders", signed_headers);
  let canonical_request = format!("{}\n{}\n{}\n{}\n{}\n{}",
                                  method,
                                  &canonical_uri,
                                  canonical_querystring,
                                  canonical_headers,
                                  signed_headers,
                                  payload_hash);
  let string_to_sign = format!("{}\n{}\n{}\n{}",
                           AMZ_ALGO,
                           amzdate,
                           credential_scope,
                           hash::hash(hash::Type::SHA256, canonical_request.as_bytes()).to_hex());
  let signing_key = get_signature_key(&bucket.secret_key, &datestamp, AMZ_REGION, AMZ_SERVICE);
  let signature = sign(&signing_key, &string_to_sign.as_bytes()).to_hex();
  let canonical_querystring = format!("{}&{}={}",
                                      canonical_querystring,
                                      &"X-Amz-Signature", signature);
  format!("{}{}?{}",
                endpoint,
                canonical_uri,
                canonical_querystring)
}

<|endoftext|>