<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
// If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

extern crate libc;
extern crate simple_slab;
#[macro_use] extern crate bitflags;

use std::ops::Drop;
use std::sync::Mutex;
use std::io::{self, Error};
use std::os::unix::io::{RawFd, AsRawFd};
use std::time::{Duration, Instant};

use simple_slab::Slab;
pub use interest::Interest;

mod interest;

bitflags! {
    flags ControlOptions: i32 {
        const EPOLL_CTL_ADD = libc::EPOLL_CTL_ADD,
        const EPOLL_CTL_MOD = libc::EPOLL_CTL_MOD,
        const EPOLL_CTL_DEL = libc::EPOLL_CTL_DEL
    }
}

bitflags! {
    pub flags Events: u32 {
        /// Sets the Edge Triggered behavior for the associated file descriptor.
        ///
        /// The default behavior for epoll is Level Triggered.
        const EPOLLET      = libc::EPOLLET as u32,
        /// The associated file is available for read operations.
        const EPOLLIN      = libc::EPOLLIN as u32,
        /// Error condition happened on the associated file descriptor.
        ///
        /// `wait` will always wait for this event; is not necessary to set it in events.
        const EPOLLERR     = libc::EPOLLERR as u32,
        /// Hang up happened on the associated file descriptor.
        ///
        /// `wait` will always wait for this event; it is not necessary to set it in events.
        /// Note that when reading from a channel such as a pipe or a stream socket, this event
        /// merely indicates that the peer closed its end of the channel. Subsequent reads from
        /// the channel will return 0 (end of file) only after all outstanding data in the
        /// channel has been consumed.
        const EPOLLHUP     = libc::EPOLLHUP as u32,
        /// The associated file is available for write operations.
        const EPOLLOUT     = libc::EPOLLOUT as u32,
        /// There is urgent data available for read operations.
        const EPOLLPRI     = libc::EPOLLPRI as u32,
        /// Stream socket peer closed connection, or shut down writing half of connection.
        ///
        /// This flag is especially useful for writing simple code to detect peer shutdown when
        /// using Edge Triggered monitoring.
        const EPOLLRDHUP   = libc::EPOLLRDHUP as u32,
        /// If `EPOLLONESHOT` and `EPOLLET` are clear and the process has the `CAP_BLOCK_SUSPEND`
        /// capability, ensure that the system does not enter "suspend" or "hibernate" while this
        /// event is pending or being processed.
        ///
        /// The event is considered as being "processed" from the time when it is returned by
        /// a call to `wait` until the next call to `wait` on the same `EpollInstance`
        /// descriptor, the closure of that file descriptor, the removal of the event file
        /// descriptor with `EPOLL_CTL_DEL`, or the clearing of `EPOLLWAKEUP` for the event file
        /// descriptor with `EPOLL_CTL_MOD`.
        const EPOLLWAKEUP  = libc::EPOLLWAKEUP as u32,
        /// Sets the one-shot behavior for the associated file descriptor.
        ///
        /// This means that after an event is pulled out with `wait` the associated file
        /// descriptor is internally disabled and no other events will be reported by the epoll
        /// interface.  The user must call `ctl` with `EPOLL_CTL_MOD` to rearm the file
        /// descriptor with a new event mask.
        const EPOLLONESHOT = libc::EPOLLONESHOT as u32
    }
}

/// Thread safe abstraction around the returned `fd` from `libc::epoll_create(1)`
pub struct EpollInstance {
    fd: libc::c_int,
    interest_mutex: Mutex<Slab<Interest>>,
    events: u64,
    wait: Duration
}

impl EpollInstance {

    /// Creates a new `EpollInstance`.
    ///
    /// ## Notes
    /// * `FD_CLOEXEC` flag is set on the underlying fd returned.
    pub fn new() -> io::Result<EpollInstance> {
        let epfd = unsafe {
            let fd = try!(cvt(libc::epoll_create(1)));
            let mut flags = try!(cvt(libc::fcntl(fd, libc::F_GETFD)));
            flags |= libc::FD_CLOEXEC;
            try!(cvt(libc::fcntl(fd, libc::F_SETFD, flags)));
            fd
        };

        Ok(EpollInstance {
            fd: epfd,
            interest_mutex: Mutex::new(Slab::<Interest>::new()),
            events: 0,
            wait: Duration::new(0, 0)
        })
    }

    /// Register an initial `Interest` with this instance.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn add_interest(&mut self, interest: Interest) -> io::Result<()> {
        let mut event_mask = libc::epoll_event {
            events: interest.events().bits() as u32,
            u64: interest.data()
        };
        try!(ctl(self.fd,
                 EPOLL_CTL_ADD.bits(),
                 interest.as_raw_fd(),
                 &mut event_mask as *mut libc::epoll_event));

        let mut list = self.interest_mutex.lock().unwrap();
        (*list).insert(interest);

        Ok(())
    }

    /// Modify the original `Interest`, identified by its `RawFd`, to the passed
    /// interest's events and data fields.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn mod_interest(&mut self, interest: &Interest) -> io::Result<()> {
        let mut event_mask = libc::epoll_event {
            events: interest.events().bits() as u32,
            u64: interest.data()
        };
        try!(ctl(self.fd,
                 EPOLL_CTL_MOD.bits(),
                 interest.as_raw_fd(),
                 &mut event_mask as *mut libc::epoll_event));

        let mut list = self.interest_mutex.lock().unwrap();
        for ref mut _interest in (*list).iter_mut() {
            if _interest.as_raw_fd() == interest.as_raw_fd() {
                _interest.set_events(interest.events());
                _interest.set_data(interest.data());
                break;
            }
        }

        Ok(())
    }

    /// Remove the passed `Interest`, identified by its `RawFd`, from this instance.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn del_interest(&mut self, interest: &Interest) -> io::Result<()> {
        // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required a non-null
        // pointer in event, even though this argument is ignored.
        let mut event_mask = libc::epoll_event {
            events: 0u32,
            u64: 0u64
        };
        try!(ctl(self.fd,
                 EPOLL_CTL_DEL.bits(),
                 interest.as_raw_fd(),
                 &mut event_mask as *mut libc::epoll_event));

        let mut offset: usize = 0;
        let mut list = self.interest_mutex.lock().unwrap();
        for ref mut _interest in (*list).iter() {
            if _interest.as_raw_fd() == interest.as_raw_fd() {
                break;
            }
            offset += 1
        }
        (*list).remove(offset);

        Ok(())
    }

    /// Waits for events on this instance for at most `timeout` milliseconds and returns at most
    /// `max_returned` `Interests`.
    ///
    /// ## Notes
    ///
    /// * If `timeout` is negative, it will block until an event is received.
    /// * `max_returned` must be greater than zero.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn wait(&mut self, timeout: i32, max_returned: usize) -> io::Result<Vec<Interest>> {
        let timeout = if timeout < -1 { -1 } else { timeout };
        let mut ret_buf = Vec::<Interest>::with_capacity(max_returned);
        let mut buf = Vec::<libc::epoll_event>::with_capacity(max_returned);

        let start = Instant::now();
        unsafe {
            let num_events = try!(cvt(libc::epoll_wait(self.fd,
                                                       buf.as_mut_ptr(),
                                                       max_returned as i32,
                                                       timeout))) as usize;
            buf.set_len(num_events);
        };

        self.events += buf.len() as u64;
        self.wait += start.elapsed();

        let mut list = self.interest_mutex.lock().unwrap();
        for ref event in buf.iter() {
            for ref mut interest in (*list).iter_mut() {
                if interest.data() == event.u64 {
                    (*interest.events_mut()) = Events::from_bits(event.events).unwrap();
                    ret_buf.push(interest.clone());
                    break;
                }
            }
        }

        Ok((ret_buf))
    }

    /// Returns the total number of events reported.
    pub fn events(&self) -> u64 {
        self.events
    }

    /// Returns the total time spent in a blocked, waiting state.
    pub fn wait_time(&self) -> Duration {
        self.wait
    }
}

impl AsRawFd for EpollInstance {
    fn as_raw_fd(&self) -> RawFd { self.fd }
}

impl Drop for EpollInstance {
    fn drop(&mut self) {
        unsafe { libc::close(self.fd) };
    }
}

unsafe impl Send for EpollInstance {}
unsafe impl Sync for EpollInstance {}

fn ctl(epfd: libc::c_int,
       op: libc::c_int,
       fd: libc::c_int,
       event: *mut libc::epoll_event)
       -> io::Result<()>
{
    unsafe { try!(cvt(libc::epoll_ctl(epfd, op, fd, event))) };
    Ok(())
}

fn cvt(result: libc::c_int) -> io::Result<libc::c_int> {
    if result < 0 {
        Err(Error::last_os_error())
    } else {
        Ok(result)
    }
}
<commit_msg>Added free standing fns for any non-mutex wants.<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
// If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

extern crate libc;
extern crate simple_slab;
#[macro_use] extern crate bitflags;

use std::ops::Drop;
use std::sync::Mutex;
use std::io::{self, Error};
use std::os::unix::io::{RawFd, AsRawFd};
use std::time::{Duration, Instant};

use simple_slab::Slab;
pub use interest::Interest;

mod interest;

bitflags! {
    pub flags ControlOptions: i32 {
        const EPOLL_CTL_ADD = libc::EPOLL_CTL_ADD,
        const EPOLL_CTL_MOD = libc::EPOLL_CTL_MOD,
        const EPOLL_CTL_DEL = libc::EPOLL_CTL_DEL
    }
}

bitflags! {
    pub flags Events: u32 {
        /// Sets the Edge Triggered behavior for the associated file descriptor.
        ///
        /// The default behavior for epoll is Level Triggered.
        const EPOLLET      = libc::EPOLLET as u32,
        /// The associated file is available for read operations.
        const EPOLLIN      = libc::EPOLLIN as u32,
        /// Error condition happened on the associated file descriptor.
        ///
        /// `wait` will always wait for this event; is not necessary to set it in events.
        const EPOLLERR     = libc::EPOLLERR as u32,
        /// Hang up happened on the associated file descriptor.
        ///
        /// `wait` will always wait for this event; it is not necessary to set it in events.
        /// Note that when reading from a channel such as a pipe or a stream socket, this event
        /// merely indicates that the peer closed its end of the channel. Subsequent reads from
        /// the channel will return 0 (end of file) only after all outstanding data in the
        /// channel has been consumed.
        const EPOLLHUP     = libc::EPOLLHUP as u32,
        /// The associated file is available for write operations.
        const EPOLLOUT     = libc::EPOLLOUT as u32,
        /// There is urgent data available for read operations.
        const EPOLLPRI     = libc::EPOLLPRI as u32,
        /// Stream socket peer closed connection, or shut down writing half of connection.
        ///
        /// This flag is especially useful for writing simple code to detect peer shutdown when
        /// using Edge Triggered monitoring.
        const EPOLLRDHUP   = libc::EPOLLRDHUP as u32,
        /// If `EPOLLONESHOT` and `EPOLLET` are clear and the process has the `CAP_BLOCK_SUSPEND`
        /// capability, ensure that the system does not enter "suspend" or "hibernate" while this
        /// event is pending or being processed.
        ///
        /// The event is considered as being "processed" from the time when it is returned by
        /// a call to `wait` until the next call to `wait` on the same `EpollInstance`
        /// descriptor, the closure of that file descriptor, the removal of the event file
        /// descriptor with `EPOLL_CTL_DEL`, or the clearing of `EPOLLWAKEUP` for the event file
        /// descriptor with `EPOLL_CTL_MOD`.
        const EPOLLWAKEUP  = libc::EPOLLWAKEUP as u32,
        /// Sets the one-shot behavior for the associated file descriptor.
        ///
        /// This means that after an event is pulled out with `wait` the associated file
        /// descriptor is internally disabled and no other events will be reported by the epoll
        /// interface.  The user must call `ctl` with `EPOLL_CTL_MOD` to rearm the file
        /// descriptor with a new event mask.
        const EPOLLONESHOT = libc::EPOLLONESHOT as u32
    }
}

/// Thread safe abstraction around the returned `fd` from `libc::epoll_create(1)`
pub struct EpollInstance {
    fd: libc::c_int,
    interest_mutex: Mutex<Slab<Interest>>,
    events: u64,
    wait: Duration
}

impl EpollInstance {

    /// Creates a new `EpollInstance`.
    ///
    /// ## Notes
    /// * `FD_CLOEXEC` flag is set on the underlying fd returned.
    pub fn new() -> io::Result<EpollInstance> {
        Ok(EpollInstance {
            fd: try!(create(true)),
            interest_mutex: Mutex::new(Slab::<Interest>::new()),
            events: 0,
            wait: Duration::new(0, 0)
        })
    }

    /// Register an initial `Interest` with this instance.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn add_interest(&mut self, interest: Interest) -> io::Result<()> {
        let mut event_mask = libc::epoll_event {
            events: interest.events().bits() as u32,
            u64: interest.data()
        };

        try!(ctl(self.fd,
                 EPOLL_CTL_ADD,
                 interest.as_raw_fd(),
                 &mut event_mask as *mut libc::epoll_event));

        let mut list = self.interest_mutex.lock().unwrap();
        (*list).insert(interest);

        Ok(())
    }

    /// Modify the original `Interest`, identified by its `RawFd`, to the passed
    /// interest's events and data fields.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn mod_interest(&mut self, interest: &Interest) -> io::Result<()> {
        let mut event_mask = libc::epoll_event {
            events: interest.events().bits() as u32,
            u64: interest.data()
        };

        try!(ctl(self.fd,
                 EPOLL_CTL_MOD,
                 interest.as_raw_fd(),
                 &mut event_mask as *mut libc::epoll_event));

        let mut list = self.interest_mutex.lock().unwrap();
        for ref mut _interest in (*list).iter_mut() {
            if _interest.as_raw_fd() == interest.as_raw_fd() {
                _interest.set_events(interest.events());
                _interest.set_data(interest.data());
                break;
            }
        }

        Ok(())
    }

    /// Remove the passed `Interest`, identified by its `RawFd`, from this instance.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn del_interest(&mut self, interest: &Interest) -> io::Result<()> {
        // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required a non-null
        // pointer in event, even though this argument is ignored.
        let mut event_mask = libc::epoll_event {
            events: 0u32,
            u64: 0u64
        };

        try!(ctl(self.fd,
                 EPOLL_CTL_DEL,
                 interest.as_raw_fd(),
                 &mut event_mask as *mut libc::epoll_event));

        let mut offset: usize = 0;
        let mut list = self.interest_mutex.lock().unwrap();
        for ref mut _interest in (*list).iter() {
            if _interest.as_raw_fd() == interest.as_raw_fd() {
                break;
            }
            offset += 1
        }
        (*list).remove(offset);

        Ok(())
    }

    /// Waits for events on this instance for at most `timeout` milliseconds and returns at most
    /// `max_returned` `Interests`.
    ///
    /// ## Notes
    ///
    /// * If `timeout` is negative, it will block until an event is received.
    /// * `max_returned` must be greater than zero.
    ///
    /// ## Panics
    ///
    /// Panics if the interior Mutex has been poisoned.
    pub fn wait(&mut self, timeout: i32, max_returned: usize) -> io::Result<Vec<Interest>> {
        let start = Instant::now();
        let buf = try!(wait(self.fd, timeout, max_returned));
        self.events += buf.len() as u64;
        self.wait += start.elapsed();

        let mut ret_buf = Vec::<Interest>::with_capacity(buf.len());
        let mut list = self.interest_mutex.lock().unwrap();
        for ref event in buf.iter() {
            for ref mut interest in (*list).iter_mut() {
                if interest.data() == event.u64 {
                    (*interest.events_mut()) = Events::from_bits(event.events).unwrap();
                    ret_buf.push(interest.clone());
                    break;
                }
            }
        }

        Ok(ret_buf)
    }

    /// Returns the total number of events reported.
    pub fn events(&self) -> u64 {
        self.events
    }

    /// Returns the total time spent in a blocked, waiting state.
    pub fn wait_time(&self) -> Duration {
        self.wait
    }
}

impl AsRawFd for EpollInstance {
    fn as_raw_fd(&self) -> RawFd { self.fd }
}

impl Drop for EpollInstance {
    fn drop(&mut self) {
        unsafe { libc::close(self.fd) };
    }
}

unsafe impl Send for EpollInstance {}
unsafe impl Sync for EpollInstance {}

/// Creates a new epoll file descriptor
///
/// If `cloexec` is true, `FD_CLOEXEC` will be set on the returned file descriptor.
///
/// ## Notes
///
/// * `epoll_create(1)` is the underlying syscall.
pub fn create(cloexec: bool) -> io::Result<RawFd> {
    let epfd = unsafe {
        let fd = try!(cvt(libc::epoll_create(1)));

        if cloexec {
            let mut flags = try!(cvt(libc::fcntl(fd, libc::F_GETFD)));
            flags |= libc::FD_CLOEXEC;
            try!(cvt(libc::fcntl(fd, libc::F_SETFD, flags)));
        }
        fd
    };

    Ok(epfd)
}

/// Calls the `epoll_ctl` syscall with the passed arguments.
pub fn ctl(epfd: RawFd,
           op: ControlOptions,
           fd: RawFd,
           event: *mut libc::epoll_event)
           -> io::Result<()>
{
    unsafe { try!(cvt(libc::epoll_ctl(epfd, op.bits, fd, event))) };
    Ok(())
}

/// Calls the `epoll_wait` syscall with the passed arguments.
///
/// ## Notes
///
/// * If `timeout` is negative, it will block until an event is received.
/// * `max_returned` must be greater than zero.
pub fn wait(epfd: RawFd,
            timeout: i32,
            max_events: usize)
            -> io::Result<Vec<libc::epoll_event>>
{
    let timeout = if timeout < -1 { -1 } else { timeout };
    let mut buf = Vec::<libc::epoll_event>::with_capacity(max_events);

    unsafe {
        let num_events = try!(cvt(libc::epoll_wait(epfd,
                                                   buf.as_mut_ptr(),
                                                   max_events as i32,
                                                   timeout))) as usize;
        buf.set_len(num_events);
    };

    Ok(buf)
}

fn cvt(result: libc::c_int) -> io::Result<libc::c_int> {
    if result < 0 {
        Err(Error::last_os_error())
    } else {
        Ok(result)
    }
}
<|endoftext|>


<commit_before>//! `primal` puts raw power into prime numbers.
//!
//! This crates includes
//!
//! - optimised prime sieves
//! - checking for primality
//! - enumerating primes
//! - factorising numbers
//! - estimating upper and lower bounds for π(*n*) (the number of primes
//!   below *n*) and *pk* (the kth prime)
//!
//! This uses a state-of-the-art cache-friendly Sieve of Eratosthenes
//! to enumerate the primes up to some fixed bound (in a memory
//! efficient manner), and then allows this cached information to be
//! used for things like enumerating and counting primes.
//!
//! `primal` takes around 2.8 seconds and less than 3MB of RAM to
//! count the exact number of primes below 1010 (455052511)
//! on my laptop (i7-3517U).
//!
//! [*Source*](http://github.com/huonw/primal)
//!
//! # Using this library
//!
//! Just add the following to your [`Cargo.toml`](http://crates.io/):
//!
//! ```toml
//! [dependencies]
//! primal = "0.2"
//! ```
//!
//! # Examples
//!
//! ## "Indexing" Primes
//!
//! Let's find the 10001st prime. The easiest way is to enumerate the
//! primes, and find the 10001st:
//!
//! ```rust
//! // (.nth is zero indexed.)
//! let p = primal::Primes::all().nth(10001 - 1).unwrap();
//! println!("The 10001st prime is {}", p); // 104743
//! ```
//!
//! This takes around 400 microseconds on my computer, which seems
//! nice and quick, but, `Primes` is flexible at the cost of
//! performance: we can make it faster. The `StreamingSieve` type
//! offers a specialised `nth_prime` function:
//!
//! ```rust
//! let p = primal::StreamingSieve::nth_prime(10001);
//! println!("The 10001st prime is {}", p); // 104743
//! ```
//!
//! This runs in only 10 microseconds! `StreamingSieve` is extremely
//! efficient and uses very little memory. It is the best way to solve
//! this task with `primal`.
//!
//! Since that was so easy, let's now make the problem bigger and
//! harder: find the sum of the 100,000th, 200,000th, 300,000th, ...,
//! 10,000,000th primes (100 in total).
//!
//! We could call `StreamingSieve::nth_prime` repeatedly:
//!
//! ```rust,no_run
//! // the primes we want to find
//! let ns = (1..100 + 1).map(|x| x * 100_000).collect::<Vec<_>>();
//!
//! // search and sum them up
//! let sum = ns.iter()
//!             .map(|n| primal::StreamingSieve::nth_prime(*n))
//!             .fold(0, |a, b| a + b);
//! println!("the sum is {}", sum);
//! ```
//!
//! This takes around 1.6s seconds to print `the sum is 8795091674`;
//! not so speedy. Each call to `nth_prime` is individually fast (400
//! microseconds for 100,000 to 40 milliseconds for 10,000,000) but
//! they add up to something bad. Every one is starting from the start
//! and redoing work that previous calls have done... wouldn't it be
//! nice if we could just do the computation for 10,000,000 and reuse
//! that for the smaller ones?
//!
//! The `Sieve` type is a wrapper around `StreamingSieve` that
//! caches information, allowing repeated queries to be answered
//! efficiently.
//!
//! There's one hitch: `Sieve` requires a limit to know how far to
//! sieve: we need some way to find an upper bound to be guaranteed to
//! be at least as large as all our primes. We could guess that, say,
//! 1010 will be large enough and use that, but that's a
//! huge overestimate (spoilers: the 10,000,000th prime is around
//! 2×108). We could also try filtering with
//! exponentially larger upper bounds until we find one that works
//! (e.g. doubling each time), or, we could just take a shortcut and
//! use deeper mathematics via
//! [`estimate_nth_prime`](fn.estimate_nth_prime.html).
//!
//! ```rust
//! // the primes we want to find
//! let ns = (1..100 + 1).map(|x| x * 100_000).collect::<Vec<_>>();
//!
//! // find our upper bound
//! let (_lo, hi) = primal::estimate_nth_prime(10_000_000);
//!
//! // find the primes up to this upper bound
//! let sieve = primal::Sieve::new(hi as usize);
//!
//! // now we can efficiently sum them up
//! let sum = ns.iter()
//!             .map(|n| sieve.nth_prime(*n))
//!             .fold(0, |a, b| a + b);
//! println!("the sum is {}", sum);
//! ```
//!
//! This takes around 40 milliseconds, and gives the same output: much
//! better!
//!
//! (By the way, the version using 1010 as the bound
//! instead of the more accurate estimate still only takes ~3
//! seconds.)
//!
//! ## Counting Primes
//!
//! Another problem: count the number of primes below 1 million. This
//! is evaluating the [prime-counting function
//! π](https://en.wikipedia.org/wiki/Prime-counting_function),
//! i.e. π(106).
//!
//! As above, there's a few ways to attack this: the iterator, and the
//! sieves.
//!
//! ```rust
//! const LIMIT: usize = 1_000_000;
//!
//! // iterator
//! let count = primal::Primes::all().take_while(|p| *p < LIMIT).count();
//! println!("there are {} primes below 1 million", count); // 78498
//!
//! // sieves
//! let sieve = primal::Sieve::new(LIMIT);
//! let count = sieve.prime_pi(LIMIT);
//! println!("there are {} primes below 1 million", count);
//!
//! let count = primal::StreamingSieve::prime_pi(LIMIT);
//! println!("there are {} primes below 1 million", count);
//! ```
//!
//! `StreamingSieve` is fastest (380 microseconds) followed by `Sieve`
//! (400) with `Primes` bringing up the rear at 1300 microseconds. Of
//! course, repeated queries will be faster with `Sieve` than with
//! `StreamingSieve`, but that flexibility comes at the cost of extra
//! memory use.
//!
//! If an approximation is all that is required, `estimate_prime_pi`
//! provides close upper and lower bounds:
//!
//! ```rust
//! let (lo, hi) = primal::estimate_prime_pi(1_000_000);
//! println!("there are between {} and {} primes below 1 million", lo, hi);
//! // 78304, 78573
//! ```
//!
//! ## Searching Primes
//!
//! Now for something where `Primes` might be useful: find the first
//! prime where the binary expansion (not including trailing zeros)
//! ends like `00..001` with at least 27 zeros. This condition is
//! checked by:
//!
//! ```rust
//! fn check(p: usize) -> bool {
//!     p > 1 && (p / 2).trailing_zeros() >= 27
//! }
//! ```
//!
//! I have no idea how large the prime might be: I know it's
//! guaranteed to be at *least* 227 + 1 + 1, but not an
//! upper limit.
//!
//! The `Primes` iterator works perfectly for this:
//!
//! ```rust
//! # fn check(p: usize) -> bool {  p > 1 && (p / 2).trailing_zeros() >= 5 } // 27 is too slow
//! let p = primal::Primes::all().find(|p| check(*p)).unwrap();
//! println!("the prime is {}", p);
//! # assert_eq!(p, 193);
//! ```
//!
//! It takes about 3.1 seconds for my computer to spit out 3,221,225,473.
//!
//! Using a sieve is a little trickier: one approach is to start with
//! some estimated upper bound (like double the absolute lower bound),
//! look for a valid prime. If one isn't found, double the upper bound
//! and start again. The `primes_from` method allows for saving a
//! little bit of work: we can start iterating from an arbitrary point
//! in the sequence, such as the lower bound.
//!
//! ```rust
//! # fn check(p: usize) -> bool {  p > 1 && (p / 2).trailing_zeros() >= 5 } // 27 is too slow
//! let p;
//! let mut lower_bound = 1 << (27 + 1);
//! # let mut lower_bound = 1 << (5 + 1);
//! loop {
//!     // our upper bound is double the lower bound
//!     let sieve = primal::Sieve::new(lower_bound * 2);
//!     if let Some(p_) = sieve.primes_from(lower_bound).find(|p| check(*p)) {
//!         p = p_;
//!         break
//!     }
//!     lower_bound *= 2;
//! }
//! println!("the prime is {}", p);
//! # assert_eq!(p, 193);
//! ```
//!
//! This takes around 3.5 seconds to print the same number. Slower
//! than the iterator!
//!
//! I was just using this silly condition as an example of something
//! that doesn't have an obvious upper bound, rather than a problem
//! that is hard to do fast. There's a much faster way to tackle it,
//! by inverting the problem: construct numbers that satisfy `check`,
//! and check the primality of those.
//!
//! The numbers that satisfy `check` are `k * (1 << (27 + 1)) + 1` for
//! `k >= 1`, so the only hard bit is testing primality. Fortunately,
//! `primal` offers the `is_prime` function which is an efficient way
//! to do primality tests, even of very large numbers.
//!
//! ```rust
//! let mut p = 0;
//! for k in 1.. {
//!     p = k * (1 << (27 + 1)) + 1;
//!     if primal::is_prime(p) { break }
//! }
//! println!("the prime is {}", p);
//! # assert_eq!(p, 3_221_225_473);
//! ```
//!
//! This takes 6 microseconds: more than 500,000×
//! faster than the iterator!

#![cfg_attr(all(test, feature = "unstable"), feature(test, step_by))]

extern crate primal_estimate;
extern crate primal_check;
extern crate primal_sieve;

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

pub use primal_estimate::prime_pi as estimate_prime_pi;
pub use primal_estimate::nth_prime as estimate_nth_prime;
pub use primal_check::miller_rabin as is_prime;
pub use primal_check::{as_perfect_power, as_prime_power};

pub use primal_sieve::{StreamingSieve, Sieve, SievePrimes, Primes};

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

    use super::{Sieve, is_prime};
    use self::test::Bencher;

    const N: usize = 1_000_000;
    const STEP: usize = 101;
    #[bench]
    fn bench_miller_rabin_tests(b: &mut Bencher) {
        b.iter(|| {
            (1..N).step_by(STEP)
                .filter(|&n| is_prime(n as u64)).count()
        })
    }
    #[bench]
    fn bench_sieve_tests(b: &mut Bencher) {
        b.iter(|| {
            let sieve = Sieve::new(1_000_000);
            (1..N).step_by(STEP)
                .filter(|&n| sieve.is_prime(n)).count()
        })
    }
}
<commit_msg>Auto merge of #2 - wackywendell:doc-asserts, r=huonw<commit_after>//! `primal` puts raw power into prime numbers.
//!
//! This crates includes
//!
//! - optimised prime sieves
//! - checking for primality
//! - enumerating primes
//! - factorising numbers
//! - estimating upper and lower bounds for π(*n*) (the number of primes
//!   below *n*) and *pk* (the kth prime)
//!
//! This uses a state-of-the-art cache-friendly Sieve of Eratosthenes
//! to enumerate the primes up to some fixed bound (in a memory
//! efficient manner), and then allows this cached information to be
//! used for things like enumerating and counting primes.
//!
//! `primal` takes around 2.8 seconds and less than 3MB of RAM to
//! count the exact number of primes below 1010 (455052511)
//! on my laptop (i7-3517U).
//!
//! [*Source*](http://github.com/huonw/primal)
//!
//! # Using this library
//!
//! Just add the following to your [`Cargo.toml`](http://crates.io/):
//!
//! ```toml
//! [dependencies]
//! primal = "0.2"
//! ```
//!
//! # Examples
//!
//! ## "Indexing" Primes
//!
//! Let's find the 10001st prime. The easiest way is to enumerate the
//! primes, and find the 10001st:
//!
//! ```rust
//! // (.nth is zero indexed.)
//! let p = primal::Primes::all().nth(10001 - 1).unwrap();
//! println!("The 10001st prime is {}", p); // 104743
//! # assert_eq!(p, 104743);
//! ```
//!
//! This takes around 400 microseconds on my computer, which seems
//! nice and quick, but, `Primes` is flexible at the cost of
//! performance: we can make it faster. The `StreamingSieve` type
//! offers a specialised `nth_prime` function:
//!
//! ```rust
//! let p = primal::StreamingSieve::nth_prime(10001);
//! println!("The 10001st prime is {}", p); // 104743
//! # assert_eq!(p, 104743);
//! ```
//!
//! This runs in only 10 microseconds! `StreamingSieve` is extremely
//! efficient and uses very little memory. It is the best way to solve
//! this task with `primal`.
//!
//! Since that was so easy, let's now make the problem bigger and
//! harder: find the sum of the 100,000th, 200,000th, 300,000th, ...,
//! 10,000,000th primes (100 in total).
//!
//! We could call `StreamingSieve::nth_prime` repeatedly:
//!
//! ```rust,no_run
//! // the primes we want to find
//! let ns = (1..100 + 1).map(|x| x * 100_000).collect::<Vec<_>>();
//!
//! // search and sum them up
//! let sum = ns.iter()
//!             .map(|n| primal::StreamingSieve::nth_prime(*n))
//!             .fold(0, |a, b| a + b);
//! println!("the sum is {}", sum);
//! ```
//!
//! This takes around 1.6s seconds to print `the sum is 8795091674`;
//! not so speedy. Each call to `nth_prime` is individually fast (400
//! microseconds for 100,000 to 40 milliseconds for 10,000,000) but
//! they add up to something bad. Every one is starting from the start
//! and redoing work that previous calls have done... wouldn't it be
//! nice if we could just do the computation for 10,000,000 and reuse
//! that for the smaller ones?
//!
//! The `Sieve` type is a wrapper around `StreamingSieve` that
//! caches information, allowing repeated queries to be answered
//! efficiently.
//!
//! There's one hitch: `Sieve` requires a limit to know how far to
//! sieve: we need some way to find an upper bound to be guaranteed to
//! be at least as large as all our primes. We could guess that, say,
//! 1010 will be large enough and use that, but that's a
//! huge overestimate (spoilers: the 10,000,000th prime is around
//! 2×108). We could also try filtering with
//! exponentially larger upper bounds until we find one that works
//! (e.g. doubling each time), or, we could just take a shortcut and
//! use deeper mathematics via
//! [`estimate_nth_prime`](fn.estimate_nth_prime.html).
//!
//! ```rust
//! // the primes we want to find
//! let ns = (1..100 + 1).map(|x| x * 100_000).collect::<Vec<_>>();
//!
//! // find our upper bound
//! let (_lo, hi) = primal::estimate_nth_prime(10_000_000);
//!
//! // find the primes up to this upper bound
//! let sieve = primal::Sieve::new(hi as usize);
//!
//! // now we can efficiently sum them up
//! let sum = ns.iter()
//!             .map(|n| sieve.nth_prime(*n))
//!             .fold(0, |a, b| a + b);
//! println!("the sum is {}", sum);
//! # assert_eq!(sum, 8795091674);
//! ```
//!
//! This takes around 40 milliseconds, and gives the same output: much
//! better!
//!
//! (By the way, the version using 1010 as the bound
//! instead of the more accurate estimate still only takes ~3
//! seconds.)
//!
//! ## Counting Primes
//!
//! Another problem: count the number of primes below 1 million. This
//! is evaluating the [prime-counting function
//! π](https://en.wikipedia.org/wiki/Prime-counting_function),
//! i.e. π(106).
//!
//! As above, there's a few ways to attack this: the iterator, and the
//! sieves.
//!
//! ```rust
//! const LIMIT: usize = 1_000_000;
//!
//! // iterator
//! let count = primal::Primes::all().take_while(|p| *p < LIMIT).count();
//! println!("there are {} primes below 1 million", count); // 78498
//! # assert_eq!(count, 78498);
//!
//! // sieves
//! let sieve = primal::Sieve::new(LIMIT);
//! let count = sieve.prime_pi(LIMIT);
//! println!("there are {} primes below 1 million", count);
//! # assert_eq!(count, 78498);
//!
//! let count = primal::StreamingSieve::prime_pi(LIMIT);
//! println!("there are {} primes below 1 million", count);
//! # assert_eq!(count, 78498);
//! ```
//!
//! `StreamingSieve` is fastest (380 microseconds) followed by `Sieve`
//! (400) with `Primes` bringing up the rear at 1300 microseconds. Of
//! course, repeated queries will be faster with `Sieve` than with
//! `StreamingSieve`, but that flexibility comes at the cost of extra
//! memory use.
//!
//! If an approximation is all that is required, `estimate_prime_pi`
//! provides close upper and lower bounds:
//!
//! ```rust
//! let (lo, hi) = primal::estimate_prime_pi(1_000_000);
//! println!("there are between {} and {} primes below 1 million", lo, hi);
//! // 78304, 78573
//! # assert_eq!(lo, 78304);
//! # assert_eq!(hi, 78573);
//! ```
//!
//! ## Searching Primes
//!
//! Now for something where `Primes` might be useful: find the first
//! prime where the binary expansion (not including trailing zeros)
//! ends like `00..001` with at least 27 zeros. This condition is
//! checked by:
//!
//! ```rust
//! fn check(p: usize) -> bool {
//!     p > 1 && (p / 2).trailing_zeros() >= 27
//! }
//! ```
//!
//! I have no idea how large the prime might be: I know it's
//! guaranteed to be at *least* 227 + 1 + 1, but not an
//! upper limit.
//!
//! The `Primes` iterator works perfectly for this:
//!
//! ```rust
//! # fn check(p: usize) -> bool {  p > 1 && (p / 2).trailing_zeros() >= 5 } // 27 is too slow
//! let p = primal::Primes::all().find(|p| check(*p)).unwrap();
//! println!("the prime is {}", p);
//! # assert_eq!(p, 193);
//! ```
//!
//! It takes about 3.1 seconds for my computer to spit out 3,221,225,473.
//!
//! Using a sieve is a little trickier: one approach is to start with
//! some estimated upper bound (like double the absolute lower bound),
//! look for a valid prime. If one isn't found, double the upper bound
//! and start again. The `primes_from` method allows for saving a
//! little bit of work: we can start iterating from an arbitrary point
//! in the sequence, such as the lower bound.
//!
//! ```rust
//! # fn check(p: usize) -> bool {  p > 1 && (p / 2).trailing_zeros() >= 5 } // 27 is too slow
//! let p;
//! let mut lower_bound = 1 << (27 + 1);
//! # let mut lower_bound = 1 << (5 + 1);
//! loop {
//!     // our upper bound is double the lower bound
//!     let sieve = primal::Sieve::new(lower_bound * 2);
//!     if let Some(p_) = sieve.primes_from(lower_bound).find(|p| check(*p)) {
//!         p = p_;
//!         break
//!     }
//!     lower_bound *= 2;
//! }
//! println!("the prime is {}", p);
//! # assert_eq!(p, 193);
//! ```
//!
//! This takes around 3.5 seconds to print the same number. Slower
//! than the iterator!
//!
//! I was just using this silly condition as an example of something
//! that doesn't have an obvious upper bound, rather than a problem
//! that is hard to do fast. There's a much faster way to tackle it,
//! by inverting the problem: construct numbers that satisfy `check`,
//! and check the primality of those.
//!
//! The numbers that satisfy `check` are `k * (1 << (27 + 1)) + 1` for
//! `k >= 1`, so the only hard bit is testing primality. Fortunately,
//! `primal` offers the `is_prime` function which is an efficient way
//! to do primality tests, even of very large numbers.
//!
//! ```rust
//! let mut p = 0;
//! for k in 1.. {
//!     p = k * (1 << (27 + 1)) + 1;
//!     if primal::is_prime(p) { break }
//! }
//! println!("the prime is {}", p);
//! # assert_eq!(p, 3_221_225_473);
//! ```
//!
//! This takes 6 microseconds: more than 500,000×
//! faster than the iterator!

#![cfg_attr(all(test, feature = "unstable"), feature(test, step_by))]

extern crate primal_estimate;
extern crate primal_check;
extern crate primal_sieve;

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

pub use primal_estimate::prime_pi as estimate_prime_pi;
pub use primal_estimate::nth_prime as estimate_nth_prime;
pub use primal_check::miller_rabin as is_prime;
pub use primal_check::{as_perfect_power, as_prime_power};

pub use primal_sieve::{StreamingSieve, Sieve, SievePrimes, Primes};

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

    use super::{Sieve, is_prime};
    use self::test::Bencher;

    const N: usize = 1_000_000;
    const STEP: usize = 101;
    #[bench]
    fn bench_miller_rabin_tests(b: &mut Bencher) {
        b.iter(|| {
            (1..N).step_by(STEP)
                .filter(|&n| is_prime(n as u64)).count()
        })
    }
    #[bench]
    fn bench_sieve_tests(b: &mut Bencher) {
        b.iter(|| {
            let sieve = Sieve::new(1_000_000);
            (1..N).step_by(STEP)
                .filter(|&n| sieve.is_prime(n)).count()
        })
    }
}
<|endoftext|>


<commit_before>#![feature(catch_panic, fnbox)]

//! Work-stealing job-queue implementation.
//! Inspired by the Molecular Matters blog posts.

extern crate crossbeam;
extern crate rand;

use crossbeam::mem::CachePadded;
use crossbeam::sync::TreiberStack;

use std::any::Any;
use std::boxed::FnBox;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use std::thread;

use ::rand::thread_rng;
use ::rand::distributions::{IndependentSample, Range};

/// Messages that can be sent to workers.
enum WorkerMessage {
    Stop,
}

type Job = Box<FnBox(&Worker) + 'static + Send>;

/// A job's id in the pool.
pub struct JobId {
    idx: usize,
}

struct PoolEntry {
    job: Job,
    sender: Sender<JobExitStatus>,
    associated_counter: *const AtomicUsize,
}

unsafe impl Send for PoolEntry {}

enum JobExitStatus {
    Success,
    Panic(Box<Any>),
}

/// A mostly lock-free pool allocator of jobs.
/// It only locks when doubling the size of the vector.
/// Eventually, I may implement a lock-free vector based off of
/// Stroustroup-et-al's paper on the subject, but for now, I've
/// decided to "keep it simple, stupid".
/// However, this should be mostly wait free unless *many* jobs
/// are being submitted.
struct JobPool {
    // we ensure with the lock-free stack of free indices that
    // an index is only used by one thread at a time.
    // mutable, disjoint borrows are therefore acceptable
    // as long as they occur within the umbrella of a read
    // lock and will not be interfered with by doubling of the
    // vector.
    jobs: RwLock<Vec<RefCell<Option<CachePadded<PoolEntry>>>>>,
    unused: TreiberStack<usize>,
}

impl JobPool {
    /// Initialize this pool with the given capacity.
    fn with_capacity(size: usize) -> Self {
        let mut v = Vec::with_capacity(size);
        let unused = TreiberStack::new();

        for i in (0..size).rev() {
            v.push(RefCell::new(None));

            unused.push(i);
        }

        JobPool {
            jobs: RwLock::new(v),
            unused: unused,
        }
    }

    /// Submit a job to be stored in the pool.
    fn submit(&self, job: Job) -> (JobId, Receiver<JobExitStatus>) {
        self.submit_with_counter(job, ::std::ptr::null())
    }

    /// Submit a job to be stored in the pool.
    fn submit_with_counter(&self, job: Job, counter: *const AtomicUsize) -> (JobId, Receiver<JobExitStatus>) {
        let next = if let Some(next) = self.unused.pop() {
            next
        } else {
            self.double();
            self.unused.pop().unwrap()
        };

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let (tx, rx) = channel();

        let entry = &jobs[next];
        *entry.borrow_mut() = Some(CachePadded::new(PoolEntry {
            job: job,
            sender: tx,
            associated_counter: counter,
        }));

        (JobId { idx: next, }, rx)
    }

    // double the vector size
    fn double(&self) {
        let mut jobs = match self.jobs.write() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let old_len = jobs.len();
        let new_len = old_len * 2;

        jobs.reserve(new_len);

        for i in (old_len..new_len).rev() {
            jobs.push(RefCell::new(None));

            self.unused.push(i);
        }
    }

    /// Get the job given by id, passing the worker
    /// as a parameter.
    fn get_job(&self, id: JobId) -> PoolEntry {
        use std::ptr;

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        // take the job out before pushing the index back onto
        // the unused stack.
        let entry = jobs[id.idx].borrow_mut().take().unwrap();

        // CachePadded doesn't run destructors, so this is ok.
        let job = unsafe { ptr::read(&*entry) };
        self.unused.push(id.idx);
        job
    }
}

unsafe impl Send for JobPool {}
unsafe impl Sync for JobPool {}

/// A thread-safe double-ended queue for jobs.
/// Right now, it's just a dumb locking wrapper of VecDeque.
/// It stores indices of jobs in the global job pool.
struct JobQueue {
    inner: Mutex<VecDeque<JobId>>,
}

impl JobQueue {
    fn new() -> Self {
        JobQueue {
            inner: Mutex::new(VecDeque::new()),
        }
    }

    // pop a job from the private end of the queue.
    // only the owner of this queue should call this.
    fn pop(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_back()
    }

    // steal a job from the public end of the queue.
    fn steal(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_front()
    }

    // push another job to the queue.
    fn push(&self, job: JobId) {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.push_back(job);
    }
}

unsafe impl Send for JobQueue {}
unsafe impl Sync for JobQueue {}

/// A worker executes jobs.
pub struct Worker {
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
    siblings: Box<[Arc<JobQueue>]>,
}

impl Worker {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: Some(self),
        }
    }

    fn run_next_job(&self) {
        let id = if let Some(id) = self.queue.pop() {
            Some(id)
        } else {
            let mut rng = thread_rng();

            let idx = Range::new(0, self.siblings.len()).ind_sample(&mut rng);
            let other_queue: &Arc<JobQueue> = &self.siblings[idx];

            // make sure we don't steal from our own queue.
            if &*self.queue as *const JobQueue != &**other_queue as *const JobQueue {
                other_queue.steal()
            } else {
                None
            }
        };

        if let Some(id) = id {
            self.run_job(id);
        }
    }

    fn run_job(&self, id: JobId) {
        use std::mem;
        let entry = self.pool.get_job(id);
        let PoolEntry { job, sender, associated_counter: counter } = entry;

        // is there a better way to do this than to box again?
        let unbounded_fn: Box<FnBox() + Send> = Box::new(move || job.call_box((self,)));
        let bounded_fn: Box<FnBox() + 'static + Send> = unsafe { mem::transmute(unbounded_fn) };

        let exit_status = match thread::catch_panic(|| bounded_fn()) {
            Ok(_) => { JobExitStatus::Success }
            Err(e) => { JobExitStatus::Panic(e) }
        };

        // decrement the associated counter and send the exit status
        if !counter.is_null() {
             unsafe { (*counter).fetch_sub(1, Ordering::SeqCst); }
        }
        drop(sender.send(exit_status)); // message doesn't _have_ to arrive
    }

    fn run_to_completion(self) {
        while let Some(job) = self.queue.pop() {
            self.run_job(job);
        }
    }

}

/// A JoinHandle can be used to manually wait for
/// a job to be completed.
pub struct JoinHandle<'a> {
    notifier: Receiver<JobExitStatus>,
    worker: Option<&'a Worker>,
}

impl<'a> JoinHandle<'a> {
    pub fn wait(self) -> Result<(), Box<Any>> {
        loop {
            match self.notifier.try_recv() {
                Ok(exit_status) => {
                    return match exit_status {
                        JobExitStatus::Success => Ok(()),
                        JobExitStatus::Panic(err) => Err(err),
                    }
                }

                // sender _never_ hangs up before the job is done.
                Err(TryRecvError::Disconnected) => unreachable!(),

                // nothing yet, just do some more work.
                _ => {
                    if let &Some(ref worker) = &self.worker {
                        worker.run_next_job()
                    }
                }
            }
        }
    }
}

pub struct JobSystem {
    handles: Vec<Sender<WorkerMessage>>,
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
}

impl JobSystem {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: None,
        }
    }
}

impl Drop for JobSystem {
    fn drop(&mut self) {
        for handle in &self.handles {
            match handle.send(WorkerMessage::Stop) {
                _ => {}
            }
        }
    }
}

/// Create a new `JobSystem` with `n_threads` worker threads.
/// This will fail if `n_threads` is less than 1.
pub fn make_pool(n_threads: usize) -> Option<JobSystem> {
    if n_threads < 1 {
        return None;
    }

    let job_pool = Arc::new(JobPool::with_capacity(4096));

    // make a queue for every thread, including this one.
    let queues: Vec<_> = (0..n_threads+1)
        .map(|_| Arc::new(JobQueue::new())).collect();

    let queues = queues.into_boxed_slice();
    let mut handles = Vec::new();

    // Spawn a worker thread for every queue but the last one,
    // since that belongs to the thread this is spawned on.
    for queue in queues.iter().cloned().take(n_threads) {
        let (tx, rx) = channel();
        let siblings = queues.clone();
        let pool = job_pool.clone();

        thread::spawn(move || {
            let worker = Worker {
                pool: pool,
                queue: queue,
                siblings: siblings,
            };

            worker_main(worker, rx);
        });
        handles.push(tx);
    }

    Some(JobSystem {
        handles: handles,
        pool: job_pool,
        queue: queues.last().unwrap().clone(),
    })
}

// the main thread for each worker
fn worker_main(worker: Worker, rx: Receiver<WorkerMessage>) {
    loop {
        // check for message
        match rx.try_recv() {
            // time to stop!
            Ok(WorkerMessage::Stop) => {
                worker.run_to_completion();
                break;
            }

            // somehow the JobSystem got dropped without disconnecting.
            Err(TryRecvError::Disconnected) => {
                worker.run_to_completion();
                break;
            }

            _ => {}
        }
        // run a job
        worker.run_next_job();
    }
}

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

    #[test]
    #[should_panic]
    fn no_zero_thread_pool() {
        make_pool(0).unwrap();
    }

    #[test]
    fn basic_spawning() {
        // four worker threads, including this one.
        let pool = make_pool(4).unwrap();

        let mut handles = Vec::new();
        for _ in 0..100 {
            // empty job
            let handle = pool.submit(|_| {});
            handles.push(handle);
        }

        for handle in handles {
            // make sure every job runs and
            // no job panics.
            handle.wait().unwrap();
        }
    }

    #[test]
    fn child_jobs() {
        let pool = make_pool(4).unwrap();

        pool.submit(|worker| {
            for _ in 0..100 {
                // empty job
                worker.submit(|_| {});
            };
        }).wait().unwrap();
    }

    #[test]
    fn panic_in_the_jobqueue() {
        let pool = make_pool(4).unwrap();

        let mut handles = Vec::new();
        for i in 0..100 {
            // 4 threads, 20 panics.
            let handle = pool.submit(move |_| match i % 5 {
                0 => panic!(i),
                _ => {}
            });

            handles.push(handle);
        }

        for handle in handles {
            match handle.wait() {
                Ok(_) => {},
                Err(b) => {
                    let panic_val = b.downcast_ref::<i32>().unwrap();
                    assert_eq!(panic_val % 5, 0);
                }
            }
        }

    }

/*    #[test]
    fn scoping() {
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();
        this_worker.submit(|worker| {
            let mut v = vec![0; 256];

            worker.scope(|scope| {
                for chunk in v.chunks_mut(32) {
                    scope.submit(move |_| {
                        for i in chunk { *i += 1 }
                    });
                }
            });

            for i in v {
                assert_eq!(i, 1);
            }
        }).wait();
    }*/
}<commit_msg>Scoping implemented again<commit_after>#![feature(catch_panic, fnbox)]

//! Work-stealing job-queue implementation.
//! Inspired by the Molecular Matters blog posts.

extern crate crossbeam;
extern crate rand;

use crossbeam::mem::CachePadded;
use crossbeam::sync::TreiberStack;

use std::any::Any;
use std::boxed::FnBox;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use std::thread;

use ::rand::thread_rng;
use ::rand::distributions::{IndependentSample, Range};

/// Messages that can be sent to workers.
enum WorkerMessage {
    Stop,
}

type JobFn = Box<FnBox(&Worker) + 'static + Send>;

/// A job's id in the pool.
pub struct JobId {
    idx: usize,
}

struct Job {
    func: JobFn,
    sender: Sender<JobExitStatus>,
    // counter for this job's associated scope.
    // null if there isn't one.
    scope_counter: *const AtomicUsize,
}

unsafe impl Send for Job {}

enum JobExitStatus {
    Success,
    Panic(Box<Any>),
}

/// A mostly lock-free pool allocator of jobs.
/// It only locks when doubling the size of the vector.
/// Eventually, I may implement a lock-free vector based off of
/// Stroustroup-et-al's paper on the subject, but for now, I've
/// decided to "keep it simple, stupid".
/// However, this should be mostly wait free unless *many* jobs
/// are being submitted.
struct JobPool {
    // we ensure with the lock-free stack of free indices that
    // an index is only used by one thread at a time.
    // mutable, disjoint borrows are therefore acceptable
    // as long as they occur within the umbrella of a read
    // lock and will not be interfered with by doubling of the
    // vector.
    jobs: RwLock<Vec<RefCell<Option<CachePadded<Job>>>>>,
    unused: TreiberStack<usize>,
}

impl JobPool {
    /// Initialize this pool with the given capacity.
    fn with_capacity(size: usize) -> Self {
        let mut v = Vec::with_capacity(size);
        let unused = TreiberStack::new();

        for i in (0..size).rev() {
            v.push(RefCell::new(None));

            unused.push(i);
        }

        JobPool {
            jobs: RwLock::new(v),
            unused: unused,
        }
    }

    /// Submit a job to be stored in the pool.
    fn submit(&self, job: JobFn) -> (JobId, Receiver<JobExitStatus>) {
        self.submit_with_counter(job, ::std::ptr::null())
    }

    /// Submit a job to be stored in the pool with a pointer to a scope counter.
    fn submit_with_counter(&self, job: JobFn, counter: *const AtomicUsize) -> (JobId, Receiver<JobExitStatus>) {
        let next = if let Some(next) = self.unused.pop() {
            next
        } else {
            self.double();
            self.unused.pop().unwrap()
        };

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let (tx, rx) = channel();

        let entry = &jobs[next];

        *entry.borrow_mut() = Some(CachePadded::new(Job {
            func: job,
            sender: tx,
            scope_counter: counter,
        }));

        (JobId { idx: next, }, rx)
    }

    // double the vector size
    fn double(&self) {
        let mut jobs = match self.jobs.write() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let old_len = jobs.len();
        let new_len = old_len * 2;

        jobs.reserve(new_len);

        for i in (old_len..new_len).rev() {
            jobs.push(RefCell::new(None));

            self.unused.push(i);
        }
    }

    /// Get the job given by id, passing the worker
    /// as a parameter.
    fn get_job(&self, id: JobId) -> Job {
        use std::ptr;

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        // take the job out before pushing the index back onto
        // the unused stack.
        let entry = jobs[id.idx].borrow_mut().take().unwrap();

        // CachePadded doesn't run destructors, so this is ok.
        let job = unsafe { ptr::read(&*entry) };
        self.unused.push(id.idx);
        job
    }
}

unsafe impl Send for JobPool {}
unsafe impl Sync for JobPool {}

/// A thread-safe double-ended queue for jobs.
/// Right now, it's just a dumb locking wrapper of VecDeque.
/// It stores indices of jobs in the global job pool.
struct JobQueue {
    inner: Mutex<VecDeque<JobId>>,
}

impl JobQueue {
    fn new() -> Self {
        JobQueue {
            inner: Mutex::new(VecDeque::new()),
        }
    }

    // pop a job from the private end of the queue.
    // only the owner of this queue should call this.
    fn pop(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_back()
    }

    // steal a job from the public end of the queue.
    fn steal(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_front()
    }

    // push another job to the queue.
    fn push(&self, job: JobId) {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.push_back(job);
    }
}

unsafe impl Send for JobQueue {}
unsafe impl Sync for JobQueue {}

/// A worker executes jobs.
pub struct Worker {
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
    siblings: Box<[Arc<JobQueue>]>,
}

impl Worker {
    /// Submit a job to be run.
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: Some(self),
        }
    }

    /// Create a scope to run jobs that access local stack data.
    /// This will block until child jobs are finished.
    pub fn scope<'a, 'b, F: FnOnce(&Scope<'a, 'b>)>(&'a self, f: F) {
        let s = Scope {
            counter: AtomicUsize::new(0),
            pool: &*self.pool,
            queue: &*self.queue,
            worker: Some(self),
            _marker: PhantomData,
        };

        f(&s);

        // the move incurred by dropping s actually
        // invalidates all the pointers to its counter.
        // so this should preferably be done out of the
        // destructor.
        s.wait_all();
    }

    fn run_next_job(&self) {
        let id = if let Some(id) = self.queue.pop() {
            Some(id)
        } else {
            let mut rng = thread_rng();

            let idx = Range::new(0, self.siblings.len()).ind_sample(&mut rng);
            let other_queue: &Arc<JobQueue> = &self.siblings[idx];

            // make sure we dn't steal from our own queue.
            if &*self.queue as *const JobQueue != &**other_queue as *const JobQueue {
                other_queue.steal()
            } else {
                None
            }
        };

        if let Some(id) = id {
            self.run_job(id);
        }
    }

    fn run_job(&self, id: JobId) {
        use std::mem;
        let job = self.pool.get_job(id);
        let Job { func, sender, scope_counter } = job;

        // is there a better way to do this than to box again?
        let unbounded_fn: Box<FnBox() + Send> = Box::new(move || func.call_box((self,)));
        let bounded_fn: Box<FnBox() + 'static + Send> = unsafe { mem::transmute(unbounded_fn) };

        let exit_status = match thread::catch_panic(|| bounded_fn()) {
            Ok(_) => { JobExitStatus::Success }
            Err(e) => { JobExitStatus::Panic(e) }
        };

        // decrement the scope counter and send the exit status
        if !scope_counter.is_null() {
            let counter: &AtomicUsize = unsafe { (&*scope_counter) };
            counter.fetch_sub(1, Ordering::SeqCst);
        }

        drop(sender.send(exit_status)); // message doesn't _have_ to arrive
    }

    fn run_to_completion(self) {
        while let Some(job) = self.queue.pop() {
            self.run_job(job);
        }
    }

}

/// A scope for submitting jobs.
pub struct Scope<'a, 'b> {
    counter: AtomicUsize,
    pool: &'a JobPool,
    queue: &'a JobQueue,
    worker: Option<&'a Worker>,
    _marker: PhantomData<RefCell<&'b ()>>,
}

impl<'a, 'b> Scope<'a, 'b> {
    pub fn submit<F: FnOnce(&Worker) + 'b + Send>(&'a self, f: F) -> JoinHandle<'a> {
        use std::mem;

        let unbounded_fn: Box<FnBox(&Worker) + 'b + Send> = Box::new(f);
        let bounded_fn: Box<FnBox(&Worker) + 'static + Send> = unsafe { mem::transmute(unbounded_fn) };

        // increment the counter before submitting the job just in
        // case it gets grabbed really quickly
        let (id, notifier) = self.pool.submit_with_counter(bounded_fn, &self.counter as *const AtomicUsize);
        self.counter.fetch_add(1, Ordering::SeqCst);
        self.queue.push(id);
        JoinHandle {
            notifier: notifier,
            worker: None,
        }
    }

    fn wait_all(&self) {
        // busy wait until they're all done.
        while self.counter.load(Ordering::SeqCst) > 0 {
            if let &Some(ref worker) = &self.worker {
                worker.run_next_job();
            }
        }
    }
}

/// A JoinHandle can be used to manually wait for
/// a job to be completed.
pub struct JoinHandle<'a> {
    notifier: Receiver<JobExitStatus>,
    worker: Option<&'a Worker>,
}

impl<'a> JoinHandle<'a> {
    pub fn wait(self) -> Result<(), Box<Any>> {
        loop {
            match self.notifier.try_recv() {
                Ok(exit_status) => {
                    return match exit_status {
                        JobExitStatus::Success => Ok(()),
                        JobExitStatus::Panic(err) => Err(err),
                    }
                }

                // sender _never_ hangs up before the job is done.
                Err(TryRecvError::Disconnected) => unreachable!(),

                // nothing yet, just do some more work.
                _ => {
                    if let &Some(ref worker) = &self.worker {
                        worker.run_next_job()
                    }
                }
            }
        }
    }
}

pub struct JobSystem {
    handles: Vec<Sender<WorkerMessage>>,
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
}

impl JobSystem {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: None,
        }
    }
}

impl Drop for JobSystem {
    fn drop(&mut self) {
        for handle in &self.handles {
            match handle.send(WorkerMessage::Stop) {
                _ => {}
            }
        }
    }
}

/// Create a new `JobSystem` with `n_threads` worker threads.
/// This will fail if `n_threads` is less than 1.
pub fn make_pool(n_threads: usize) -> Option<JobSystem> {
    if n_threads < 1 {
        return None;
    }

    let job_pool = Arc::new(JobPool::with_capacity(4096));

    // make a queue for every thread, including this one.
    let queues: Vec<_> = (0..n_threads+1)
        .map(|_| Arc::new(JobQueue::new())).collect();

    let queues = queues.into_boxed_slice();
    let mut handles = Vec::new();

    // Spawn a worker thread for every queue but the last one,
    // since that belongs to the thread this is spawned on.
    for queue in queues.iter().cloned().take(n_threads) {
        let (tx, rx) = channel();
        let siblings = queues.clone();
        let pool = job_pool.clone();

        thread::spawn(move || {
            let worker = Worker {
                pool: pool,
                queue: queue,
                siblings: siblings,
            };

            worker_main(worker, rx);
        });
        handles.push(tx);
    }

    Some(JobSystem {
        handles: handles,
        pool: job_pool,
        queue: queues.last().unwrap().clone(),
    })
}

// the main thread for each worker
fn worker_main(worker: Worker, rx: Receiver<WorkerMessage>) {
    loop {
        // check for message
        match rx.try_recv() {
            // time to stop!
            Ok(WorkerMessage::Stop) => {
                worker.run_to_completion();
                break;
            }

            // somehow the JobSystem got dropped without disconnecting.
            Err(TryRecvError::Disconnected) => {
                worker.run_to_completion();
                break;
            }

            _ => {}
        }
        // run a job
        worker.run_next_job();
    }
}

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

    #[test]
    #[should_panic]
    fn no_zero_thread_pool() {
        make_pool(0).unwrap();
    }

    #[test]
    fn basic_spawning() {
        // four worker threads, including this one.
        let pool = make_pool(4).unwrap();

        let mut handles = Vec::new();
        for _ in 0..100 {
            // empty job
            let handle = pool.submit(|_| {});
            handles.push(handle);
        }

        for handle in handles {
            // make sure every job runs and
            // no job panics.
            handle.wait().unwrap();
        }
    }

    #[test]
    fn child_jobs() {
        let pool = make_pool(4).unwrap();

        pool.submit(|worker| {
            for _ in 0..100 {
                // empty job
                worker.submit(|_| {});
            };
        }).wait().unwrap();
    }

    #[test]
    fn panic_in_the_jobqueue() {
        let pool = make_pool(4).unwrap();

        let mut handles = Vec::new();
        for i in 0..100 {
            // 4 threads, 20 panics.
            let handle = pool.submit(move |_| match i % 5 {
                0 => panic!(i),
                _ => {}
            });

            handles.push(handle);
        }

        for handle in handles {
            match handle.wait() {
                Ok(_) => {},
                Err(b) => {
                    let panic_val = b.downcast_ref::<i32>().unwrap();
                    assert_eq!(panic_val % 5, 0);
                }
            }
        }

    }

    #[test]
    fn scoping() {
        let pool = make_pool(4).unwrap();
        pool.submit(|worker| {
            let mut v = vec![0; 256];

            worker.scope(|scope| {
                for chunk in v.chunks_mut(32) {
                    scope.submit(move |_| {
                        for i in chunk { *i += 1 }
                    });
                }
            });

            for i in v {
                assert_eq!(i, 1);
            }
        }).wait().unwrap();
    }
}<|endoftext|>


<commit_before>extern crate librsync_sys as raw;

use std::io::{self, Read};
use std::marker::PhantomData;
use std::ops::Deref;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum MagicNumber {
    Delta,
    MD4,
    Blake2,
}

#[derive(Debug)]
pub enum Error {
    Done,
    Blocked,
    Running,
    TestSkipped,
    Io(io::Error),
    Syntax,
    Mem,
    InputEnded,
    BadMagic,
    Unimplemented,
    Corrupt,
    Internal,
    Param,
}

pub type Result<T> = std::result::Result<T, Error>;

pub struct Signature<R: Read> {
    old: R,
    job: Job,
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
    // true when input has ended
    done: bool,
}

struct Job(*mut raw::rs_job_t);

// Wrapper around rs_buffers_t.
struct Buffers<'a> {
    inner: raw::rs_buffers_t,
    _phantom: PhantomData<&'a u8>,
}

impl<R: Read> Signature<R> {
    pub fn new(old: R,
               new_block_len: usize,
               strong_len: usize,
               sig_magic: MagicNumber)
               -> Result<Self> {
        let job = unsafe { raw::rs_sig_begin(new_block_len, strong_len, sig_magic.as_raw()) };
        if job.is_null() {
            return Err(Error::BadMagic);
        }
        Ok(Signature {
            old: old,
            job: Job(job),
            buf: vec![0; raw::RS_DEFAULT_BLOCK_LEN],
            pos: 0,
            cap: 0,
            done: false,
        })
    }

    pub fn into_inner(self) -> R {
        self.old
    }
}

impl<R: Read> Read for Signature<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut out_pos = 0;
        let mut out_cap = buf.len();

        loop {
            if self.cap == 0 && !self.done {
                // need to read some data from input
                let read = try!(self.old.read(&mut self.buf));
                if read == 0 {
                    self.done = true;
                }
                self.pos = 0;
                self.cap = read;
            }

            // work
            let mut buffers = Buffers::new(&self.buf[self.pos..self.pos + self.cap],
                                           &mut buf[out_pos..],
                                           self.done);
            let res = unsafe { raw::rs_job_iter(*self.job, buffers.as_raw()) };
            if res != raw::RS_DONE && res != raw::RS_BLOCKED {
                let err = Error::from_raw(res);
                return Err(other_io_err(format!("Error processing signature: \'{:?}\'", err)));
            }

            // update buffer cap and pos
            let read = self.cap - buffers.available_input();
            self.pos += read;
            self.cap -= read;
            // determine written
            let written = out_cap - buffers.available_output();
            out_pos += written;
            out_cap -= written;
            if out_cap == 0 || self.done {
                return Ok(written);
            }
        }
    }
}

impl Error {
    fn from_raw(val: raw::rs_result) -> Option<Self> {
        match val {
            raw::RS_DONE => Some(Error::Done),
            raw::RS_BLOCKED => Some(Error::Blocked),
            raw::RS_RUNNING => Some(Error::Running),
            raw::RS_TEST_SKIPPED => Some(Error::TestSkipped),
            raw::RS_IO_ERROR => Some(Error::Io(other_io_err("Unknown IO error"))),
            raw::RS_SYNTAX_ERROR => Some(Error::Syntax),
            raw::RS_MEM_ERROR => Some(Error::Mem),
            raw::RS_INPUT_ENDED => Some(Error::InputEnded),
            raw::RS_BAD_MAGIC => Some(Error::BadMagic),
            raw::RS_UNIMPLEMENTED => Some(Error::Unimplemented),
            raw::RS_CORRUPT => Some(Error::Corrupt),
            raw::RS_INTERNAL_ERROR => Some(Error::Internal),
            raw::RS_PARAM_ERROR => Some(Error::Param),
            _ => None,
        }
    }
}

impl MagicNumber {
    pub fn from_raw(raw: raw::rs_magic_number) -> Option<Self> {
        match raw {
            raw::RS_DELTA_MAGIC => Some(MagicNumber::Delta),
            raw::RS_MD4_SIG_MAGIC => Some(MagicNumber::MD4),
            raw::RS_BLAKE2_SIG_MAGIC => Some(MagicNumber::Blake2),
            _ => None,
        }
    }

    fn as_raw(&self) -> raw::rs_magic_number {
        match *self {
            MagicNumber::Delta => raw::RS_DELTA_MAGIC,
            MagicNumber::MD4 => raw::RS_MD4_SIG_MAGIC,
            MagicNumber::Blake2 => raw::RS_BLAKE2_SIG_MAGIC,
        }
    }
}

impl Deref for Job {
    type Target = *mut raw::rs_job_t;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Drop for Job {
    fn drop(&mut self) {
        unsafe {
            if !self.0.is_null() {
                raw::rs_job_free(self.0);
            }
        }
    }
}

impl<'a> Buffers<'a> {
    pub fn new(in_buf: &'a [u8], out_buf: &'a mut [u8], eof_in: bool) -> Self {
        Buffers {
            inner: raw::rs_buffers_t {
                next_in: in_buf.as_ptr() as *const i8,
                avail_in: in_buf.len(),
                eof_in: if eof_in { 1 } else { 0 },
                next_out: out_buf.as_mut_ptr() as *mut i8,
                avail_out: out_buf.len(),
            },
            _phantom: PhantomData,
        }
    }

    pub fn as_raw(&mut self) -> *mut raw::rs_buffers_t {
        &mut self.inner
    }

    pub fn available_input(&self) -> usize {
        self.inner.avail_in as usize
    }

    pub fn available_output(&self) -> usize {
        self.inner.avail_out as usize
    }
}

fn other_io_err<T: AsRef<str>>(msg: T) -> io::Error {
    io::Error::new(io::ErrorKind::Other, msg.as_ref())
}

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

    #[test]
    fn signature() {
        let data = "this is a string to be tested";
        let cursor = Cursor::new(data);
        let _sig = Signature::new(cursor, 10, 5, MagicNumber::MD4).unwrap();
    }

    #[test]
    fn signature_invalid_magic() {
        let data = "this is a string to be tested";
        let cursor = Cursor::new(data);
        let sig = Signature::new(cursor, 10, 5, MagicNumber::Delta);
        assert!(sig.is_err());
    }
}
<commit_msg>Fixed bug and added test<commit_after>extern crate librsync_sys as raw;

use std::io::{self, Read};
use std::marker::PhantomData;
use std::ops::Deref;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum MagicNumber {
    Delta,
    MD4,
    Blake2,
}

#[derive(Debug)]
pub enum Error {
    Done,
    Blocked,
    Running,
    TestSkipped,
    Io(io::Error),
    Syntax,
    Mem,
    InputEnded,
    BadMagic,
    Unimplemented,
    Corrupt,
    Internal,
    Param,
}

pub type Result<T> = std::result::Result<T, Error>;

pub struct Signature<R: Read> {
    old: R,
    job: Job,
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
    // true when input has ended
    done: bool,
}

struct Job(*mut raw::rs_job_t);

// Wrapper around rs_buffers_t.
struct Buffers<'a> {
    inner: raw::rs_buffers_t,
    _phantom: PhantomData<&'a u8>,
}

impl<R: Read> Signature<R> {
    pub fn new(old: R,
               new_block_len: usize,
               strong_len: usize,
               sig_magic: MagicNumber)
               -> Result<Self> {
        let job = unsafe { raw::rs_sig_begin(new_block_len, strong_len, sig_magic.as_raw()) };
        if job.is_null() {
            return Err(Error::BadMagic);
        }
        Ok(Signature {
            old: old,
            job: Job(job),
            buf: vec![0; raw::RS_DEFAULT_BLOCK_LEN],
            pos: 0,
            cap: 0,
            done: false,
        })
    }

    pub fn into_inner(self) -> R {
        self.old
    }
}

impl<R: Read> Read for Signature<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut out_pos = 0;
        let mut out_cap = buf.len();

        loop {
            if self.cap == 0 && !self.done {
                // need to read some data from input
                let read = try!(self.old.read(&mut self.buf));
                if read == 0 {
                    self.done = true;
                }
                self.pos = 0;
                self.cap = read;
            }

            // work
            let mut buffers = Buffers::new(&self.buf[self.pos..self.pos + self.cap],
                                           &mut buf[out_pos..],
                                           self.done);
            let res = unsafe { raw::rs_job_iter(*self.job, buffers.as_raw()) };
            if res != raw::RS_DONE && res != raw::RS_BLOCKED {
                let err = Error::from_raw(res);
                return Err(other_io_err(format!("Error processing signature: \'{:?}\'", err)));
            }

            // update buffer cap and pos
            let read = self.cap - buffers.available_input();
            self.pos += read;
            self.cap -= read;
            // determine written
            let written = out_cap - buffers.available_output();
            out_pos += written;
            out_cap -= written;
            if out_cap == 0 || written == 0 {
                return Ok(out_pos);
            }
        }
    }
}

impl Error {
    fn from_raw(val: raw::rs_result) -> Option<Self> {
        match val {
            raw::RS_DONE => Some(Error::Done),
            raw::RS_BLOCKED => Some(Error::Blocked),
            raw::RS_RUNNING => Some(Error::Running),
            raw::RS_TEST_SKIPPED => Some(Error::TestSkipped),
            raw::RS_IO_ERROR => Some(Error::Io(other_io_err("Unknown IO error"))),
            raw::RS_SYNTAX_ERROR => Some(Error::Syntax),
            raw::RS_MEM_ERROR => Some(Error::Mem),
            raw::RS_INPUT_ENDED => Some(Error::InputEnded),
            raw::RS_BAD_MAGIC => Some(Error::BadMagic),
            raw::RS_UNIMPLEMENTED => Some(Error::Unimplemented),
            raw::RS_CORRUPT => Some(Error::Corrupt),
            raw::RS_INTERNAL_ERROR => Some(Error::Internal),
            raw::RS_PARAM_ERROR => Some(Error::Param),
            _ => None,
        }
    }
}

impl MagicNumber {
    pub fn from_raw(raw: raw::rs_magic_number) -> Option<Self> {
        match raw {
            raw::RS_DELTA_MAGIC => Some(MagicNumber::Delta),
            raw::RS_MD4_SIG_MAGIC => Some(MagicNumber::MD4),
            raw::RS_BLAKE2_SIG_MAGIC => Some(MagicNumber::Blake2),
            _ => None,
        }
    }

    fn as_raw(&self) -> raw::rs_magic_number {
        match *self {
            MagicNumber::Delta => raw::RS_DELTA_MAGIC,
            MagicNumber::MD4 => raw::RS_MD4_SIG_MAGIC,
            MagicNumber::Blake2 => raw::RS_BLAKE2_SIG_MAGIC,
        }
    }
}

impl Deref for Job {
    type Target = *mut raw::rs_job_t;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Drop for Job {
    fn drop(&mut self) {
        unsafe {
            if !self.0.is_null() {
                raw::rs_job_free(self.0);
            }
        }
    }
}

impl<'a> Buffers<'a> {
    pub fn new(in_buf: &'a [u8], out_buf: &'a mut [u8], eof_in: bool) -> Self {
        Buffers {
            inner: raw::rs_buffers_t {
                next_in: in_buf.as_ptr() as *const i8,
                avail_in: in_buf.len(),
                eof_in: if eof_in { 1 } else { 0 },
                next_out: out_buf.as_mut_ptr() as *mut i8,
                avail_out: out_buf.len(),
            },
            _phantom: PhantomData,
        }
    }

    pub fn as_raw(&mut self) -> *mut raw::rs_buffers_t {
        &mut self.inner
    }

    pub fn available_input(&self) -> usize {
        self.inner.avail_in as usize
    }

    pub fn available_output(&self) -> usize {
        self.inner.avail_out as usize
    }
}

fn other_io_err<T: AsRef<str>>(msg: T) -> io::Error {
    io::Error::new(io::ErrorKind::Other, msg.as_ref())
}

#[cfg(test)]
mod test {
    use super::*;
    use std::io::{Cursor, Read};

    #[test]
    fn signature() {
        let data = "this is a string to be tested";
        let cursor = Cursor::new(data);
        let mut sig = Signature::new(cursor, 10, 5, MagicNumber::MD4).unwrap();
        let mut signature = Vec::new();
        let read = sig.read_to_end(&mut signature).unwrap();
        assert_eq!(read, signature.len());
        let expected = vec![0x72, 0x73, 0x01, 0x36, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00,
                            0x05, 0x1b, 0x21, 0x04, 0x8b, 0xad, 0x3c, 0xbd, 0x19, 0x09, 0x1d,
                            0x1b, 0x04, 0xf0, 0x9d, 0x1f, 0x64, 0x31, 0xde, 0x15, 0xf4, 0x04,
                            0x87, 0x60, 0x96, 0x19, 0x50, 0x39];
        assert_eq!(signature, expected);
    }

    #[test]
    fn signature_invalid_magic() {
        let data = "this is a string to be tested";
        let cursor = Cursor::new(data);
        let sig = Signature::new(cursor, 10, 5, MagicNumber::Delta);
        assert!(sig.is_err());
    }
}
<|endoftext|>


<commit_before>//! Complex numbers

#![deny(missing_docs, warnings)]
#![feature(macro_rules)]

extern crate onezero;

use std::fmt;
use std::rand::{Rand, Rng};

use onezero::{One, Zero};

mod ffi;
mod math;

pub mod f32;
pub mod f64;

/// A complex number in Cartesian form.
#[repr(C)]
#[deriving(Clone, Copy, PartialEq)]
pub struct Complex<T> {
    /// The real part
    pub re: T,
    /// The imaginary part
    pub im: T,
}

/// Single precision complex number
#[allow(non_camel_case_types)]
pub type c64 = Complex<f32>;
/// Double precision complex number
#[allow(non_camel_case_types)]
pub type c128 = Complex<f64>;

impl<T> Complex<T> {
    /// Create a new complex number
    pub fn new(re: T, im: T) -> Complex<T> {
        Complex {
            im: im,
            re: re,
        }
    }
}

impl<T> Complex<T> where T: Add<T, T> + Clone + Mul<T, T> {
    fn norm_sqr(self) -> T {
        self.re.clone() * self.re + self.im.clone() * self.im
    }
}

impl<T> Add<T, Complex<T>> for Complex<T> where T: Add<T, T> + Clone {
    fn add(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re + rhs,
            im: self.im.clone(),
        }
    }
}

impl<T> Add<Complex<T>, Complex<T>> for Complex<T> where T: Add<T, T> {
    fn add(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self.re + rhs.re,
            im: self.im + rhs.im,
        }
    }
}

impl<T> Add<Complex<T>, Complex<T>> for T where T: Add<T, T> + Clone {
    fn add(self, rhs: Complex<T>) -> Complex<T> {
        rhs + self
    }
}

impl<T> Div<T, Complex<T>> for Complex<T> where T: Clone + Div<T, T> {
    fn div(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re / rhs.clone(),
            im: self.im / rhs,
        }
    }
}

impl<T> Div<Complex<T>, Complex<T>> for Complex<T> where
    T: Add<T, T> + Clone + Div<T, T> + Mul<T, T> + Sub<T, T>
{
    fn div(self, rhs: Complex<T>) -> Complex<T> {
        let den = rhs.clone().norm_sqr();

        Complex {
            re: {
                self.re.clone() * rhs.re.clone() + self.im.clone() * rhs.im.clone()
            } / den.clone(),
            im: (self.im * rhs.re - self.re * rhs.im) / den,
        }
    }
}

impl<T> Div<Complex<T>, Complex<T>> for T where
    T: Add<T, T> + Clone + Div<T, T> + Mul<T, T> + Neg<T>,
{
    fn div(self, rhs: Complex<T>) -> Complex<T> {
        let den = rhs.clone().norm_sqr();

        Complex {
            re: (self.clone() * rhs.re) / den.clone(),
            im: -(self * rhs.im) / den,
        }
    }
}

impl<T> Mul<T, Complex<T>> for Complex<T> where T: Clone + Mul<T, T> {
    fn mul(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re * rhs.clone(),
            im: self.im * rhs,
        }
    }
}

impl<T> Mul<Complex<T>, Complex<T>> for Complex<T> where
    T: Add<T, T> + Clone + Mul<T, T> + Sub<T, T>,
{
    fn mul(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self.re.clone() * rhs.re.clone() - self.im.clone() * rhs.im.clone(),
            im: self.re * rhs.im + self.im * rhs.re,
        }
    }
}

impl<T> Mul<Complex<T>, Complex<T>> for T where T: Clone + Mul<T, T> {
    fn mul(self, rhs: Complex<T>) -> Complex<T> {
        rhs * self
    }
}

impl<T> Neg<Complex<T>> for Complex<T> where T: Neg<T> {
    fn neg(self) -> Complex<T> {
        Complex {
            re: -self.re,
            im: -self.im,
        }
    }
}

impl<T> One for Complex<T> where T: Mul<T, T> + One + Sub<T, T> + Zero {
    fn one() -> Complex<T> {
        Complex {
            re: One::one(),
            im: Zero::zero(),
        }
    }
}

impl<T> Rand for Complex<T> where T: Rand {
    fn rand<R>(rng: &mut R) -> Complex<T> where R: Rng {
        Complex::new(rng.gen(), rng.gen())
    }
}

impl<T> fmt::Show for Complex<T> where T: PartialOrd + fmt::Show + Zero {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.im < Zero::zero() {
            write!(f, "{}-{}i", self.re, self.im)
        } else {
            write!(f, "{}+{}i", self.re, self.im)
        }
    }
}

impl<T> Sub<T, Complex<T>> for Complex<T> where T: Clone + Sub<T, T> {
    fn sub(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re - rhs,
            im: self.im.clone(),
        }
    }
}

impl<T> Sub<Complex<T>, Complex<T>> for Complex<T> where T: Sub<T, T> {
    fn sub(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self.re - rhs.re,
            im: self.im - rhs.im,
        }
    }
}

impl<T> Sub<Complex<T>, Complex<T>> for T where T: Neg<T> + Sub<T, T> {
    fn sub(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self - rhs.re,
            im: -rhs.im,
        }
    }
}

impl<T> Zero for Complex<T> where T: Zero {
    fn zero() -> Complex<T> {
        Complex {
            re: Zero::zero(),
            im: Zero::zero(),
        }
    }
}

/// Mathematical operations on complex numbers
// FIXME (AI) `T` should be an associated output type
pub trait Math<T> {
    /// Computes the complex absolute value (also called norm, modulus or magnitude)
    fn abs(self) -> T;

    /// Computes the arc cosine
    fn acos(self) -> Self;

    /// Computes the arc hyperbolic cosine
    fn acosh(self) -> Self;

    /// Computes the argument (also called the phase angle)
    fn arg(self) -> T;

    /// Computes the arc sine
    fn asin(self) -> Self;

    /// Computes the arc hyperbolic sine
    fn asinh(self) -> Self;

    /// Computes the arc tangent
    fn atan(self) -> Self;

    /// Computes the arc hyperbolic tangent
    fn atanh(self) -> Self;

    /// Computes the arc cosine
    fn cos(self) -> Self;

    /// Computes the arc hyperbolic cosine
    fn cosh(self) -> Self;

    /// Computes the complex base-e exponential
    fn exp(self) -> Self;

    /// Returns the imaginary part
    fn imag(self) -> T;

    /// Computes the complex base-e logarithm
    fn log(self) -> Self;

    /// Computes the complex conjugate
    fn conj(self) -> Self;

    /// Computes the complex power
    fn pow(self, Self) -> Self;

    /// Computes the Riemann sphere projection
    fn proj(self) -> Self;

    /// Returns the real part
    fn real(self) -> T;

    /// Computes the arc sine
    fn sin(self) -> Self;

    /// Computes the arc hyperbolic sine
    fn sinh(self) -> Self;

    /// Computes the square root
    fn sqrt(self) -> Self;

    /// Computes the arc tangent
    fn tan(self) -> Self;

    /// Computes the arc hyperbolic tangent
    fn tanh(self) -> Self;
}

#[cfg(test)]
mod test {
    macro_rules! test {
        ($($ty:ident),+) => {$(
            mod $ty {
                use Complex;
                use $ty::I;

                #[test]
                fn constructor() {
                    assert_eq!(Complex::new(3., 4.), I * 4. + 3.)
                }
            })+
         }
    }

    test!(f32, f64);
}
<commit_msg>Update to changes in master<commit_after>//! Complex numbers

#![deny(missing_docs, warnings)]
#![feature(macro_rules)]

extern crate onezero;

use std::fmt;
use std::ops::{Add, Sub, Mul, Div, Neg};
use std::rand::{Rand, Rng};

use onezero::{One, Zero};

mod ffi;
mod math;

pub mod f32;
pub mod f64;

/// A complex number in Cartesian form.
#[repr(C)]
#[derive(Clone, Copy, PartialEq)]
pub struct Complex<T> {
    /// The real part
    pub re: T,
    /// The imaginary part
    pub im: T,
}

/// Single precision complex number
#[allow(non_camel_case_types)]
pub type c64 = Complex<f32>;
/// Double precision complex number
#[allow(non_camel_case_types)]
pub type c128 = Complex<f64>;

impl<T> Complex<T> {
    /// Create a new complex number
    pub fn new(re: T, im: T) -> Complex<T> {
        Complex {
            im: im,
            re: re,
        }
    }
}

impl<T> Complex<T> where T: Add<T, T> + Clone + Mul<T, T> {
    fn norm_sqr(self) -> T {
        self.re.clone() * self.re + self.im.clone() * self.im
    }
}

impl<T> Add<T, Complex<T>> for Complex<T> where T: Add<T, T> + Clone {
    fn add(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re + rhs,
            im: self.im.clone(),
        }
    }
}

impl<T> Add<Complex<T>, Complex<T>> for Complex<T> where T: Add<T, T> {
    fn add(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self.re + rhs.re,
            im: self.im + rhs.im,
        }
    }
}

impl<T> Add<Complex<T>, Complex<T>> for T where T: Add<T, T> + Clone {
    fn add(self, rhs: Complex<T>) -> Complex<T> {
        rhs + self
    }
}

impl<T> Div<T, Complex<T>> for Complex<T> where T: Clone + Div<T, T> {
    fn div(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re / rhs.clone(),
            im: self.im / rhs,
        }
    }
}

impl<T> Div<Complex<T>, Complex<T>> for Complex<T> where
    T: Add<T, T> + Clone + Div<T, T> + Mul<T, T> + Sub<T, T>
{
    fn div(self, rhs: Complex<T>) -> Complex<T> {
        let den = rhs.clone().norm_sqr();

        Complex {
            re: {
                self.re.clone() * rhs.re.clone() + self.im.clone() * rhs.im.clone()
            } / den.clone(),
            im: (self.im * rhs.re - self.re * rhs.im) / den,
        }
    }
}

impl<T> Div<Complex<T>, Complex<T>> for T where
    T: Add<T, T> + Clone + Div<T, T> + Mul<T, T> + Neg<T>,
{
    fn div(self, rhs: Complex<T>) -> Complex<T> {
        let den = rhs.clone().norm_sqr();

        Complex {
            re: (self.clone() * rhs.re) / den.clone(),
            im: -(self * rhs.im) / den,
        }
    }
}

impl<T> Mul<T, Complex<T>> for Complex<T> where T: Clone + Mul<T, T> {
    fn mul(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re * rhs.clone(),
            im: self.im * rhs,
        }
    }
}

impl<T> Mul<Complex<T>, Complex<T>> for Complex<T> where
    T: Add<T, T> + Clone + Mul<T, T> + Sub<T, T>,
{
    fn mul(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self.re.clone() * rhs.re.clone() - self.im.clone() * rhs.im.clone(),
            im: self.re * rhs.im + self.im * rhs.re,
        }
    }
}

impl<T> Mul<Complex<T>, Complex<T>> for T where T: Clone + Mul<T, T> {
    fn mul(self, rhs: Complex<T>) -> Complex<T> {
        rhs * self
    }
}

impl<T> Neg<Complex<T>> for Complex<T> where T: Neg<T> {
    fn neg(self) -> Complex<T> {
        Complex {
            re: -self.re,
            im: -self.im,
        }
    }
}

impl<T> One for Complex<T> where T: Mul<T, T> + One + Sub<T, T> + Zero {
    fn one() -> Complex<T> {
        Complex {
            re: One::one(),
            im: Zero::zero(),
        }
    }
}

impl<T> Rand for Complex<T> where T: Rand {
    fn rand<R>(rng: &mut R) -> Complex<T> where R: Rng {
        Complex::new(rng.gen(), rng.gen())
    }
}

impl<T> fmt::Show for Complex<T> where T: PartialOrd + fmt::Show + Zero {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.im < Zero::zero() {
            write!(f, "{}-{}i", self.re, self.im)
        } else {
            write!(f, "{}+{}i", self.re, self.im)
        }
    }
}

impl<T> Sub<T, Complex<T>> for Complex<T> where T: Clone + Sub<T, T> {
    fn sub(self, rhs: T) -> Complex<T> {
        Complex {
            re: self.re - rhs,
            im: self.im.clone(),
        }
    }
}

impl<T> Sub<Complex<T>, Complex<T>> for Complex<T> where T: Sub<T, T> {
    fn sub(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self.re - rhs.re,
            im: self.im - rhs.im,
        }
    }
}

impl<T> Sub<Complex<T>, Complex<T>> for T where T: Neg<T> + Sub<T, T> {
    fn sub(self, rhs: Complex<T>) -> Complex<T> {
        Complex {
            re: self - rhs.re,
            im: -rhs.im,
        }
    }
}

impl<T> Zero for Complex<T> where T: Zero {
    fn zero() -> Complex<T> {
        Complex {
            re: Zero::zero(),
            im: Zero::zero(),
        }
    }
}

/// Mathematical operations on complex numbers
// FIXME (AI) `T` should be an associated output type
pub trait Math<T> {
    /// Computes the complex absolute value (also called norm, modulus or magnitude)
    fn abs(self) -> T;

    /// Computes the arc cosine
    fn acos(self) -> Self;

    /// Computes the arc hyperbolic cosine
    fn acosh(self) -> Self;

    /// Computes the argument (also called the phase angle)
    fn arg(self) -> T;

    /// Computes the arc sine
    fn asin(self) -> Self;

    /// Computes the arc hyperbolic sine
    fn asinh(self) -> Self;

    /// Computes the arc tangent
    fn atan(self) -> Self;

    /// Computes the arc hyperbolic tangent
    fn atanh(self) -> Self;

    /// Computes the arc cosine
    fn cos(self) -> Self;

    /// Computes the arc hyperbolic cosine
    fn cosh(self) -> Self;

    /// Computes the complex base-e exponential
    fn exp(self) -> Self;

    /// Returns the imaginary part
    fn imag(self) -> T;

    /// Computes the complex base-e logarithm
    fn log(self) -> Self;

    /// Computes the complex conjugate
    fn conj(self) -> Self;

    /// Computes the complex power
    fn pow(self, Self) -> Self;

    /// Computes the Riemann sphere projection
    fn proj(self) -> Self;

    /// Returns the real part
    fn real(self) -> T;

    /// Computes the arc sine
    fn sin(self) -> Self;

    /// Computes the arc hyperbolic sine
    fn sinh(self) -> Self;

    /// Computes the square root
    fn sqrt(self) -> Self;

    /// Computes the arc tangent
    fn tan(self) -> Self;

    /// Computes the arc hyperbolic tangent
    fn tanh(self) -> Self;
}

#[cfg(test)]
mod test {
    macro_rules! test {
        ($($ty:ident),+) => {$(
            mod $ty {
                use Complex;
                use $ty::I;

                #[test]
                fn constructor() {
                    assert_eq!(Complex::new(3., 4.), I * 4. + 3.)
                }
            })+
         }
    }

    test!(f32, f64);
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

pub type EventData<T> = HashMap<String, T>;
//pub type EventCallback<'callback, T> = |& mut EventData<T>|:'callback;
pub type EventCallback<'callback, T> = FnMut(& mut EventData<T>) -> () + 'callback;

pub trait Events<'callback, T> {
  fn on(&mut self, event_name: String, callback: &'callback mut EventCallback<'callback, T>);
  fn off(&mut self, event_name: String);
  fn emit(&mut self, event_name: String, EventData: & mut EventData<T>);
}

pub struct Emitter<'callback, T> {
  events : HashMap<String, Vec<&'callback mut EventCallback<'callback, T>>>
}

impl<'callback, T> Emitter<'callback, T>{
  pub fn new() -> Emitter<'callback, T> {
    Emitter {
      events: HashMap::new()
    }
  }
}

impl<'callback, T> Events<'callback, T> for Emitter<'callback, T> {
  fn on(&mut self, event_name: String, callback: &'callback mut EventCallback<'callback, T>) {
    if self.events.contains_key(&event_name) {
      match self.events.get_mut(&event_name) {
          Some(callbacks) => callbacks.push(callback),
          None => (),
      }
    } else {
      let vec = vec![callback];
      self.events.insert(event_name, vec);
    }

  }
  fn off(&mut self, event_name: String) {
    self.events.remove(&event_name);
  }
  fn emit(&mut self, event_name: String, event_data: & mut EventData<T>) {
    match self.events.get_mut(&event_name) {
      Some(callbacks) => {
        for callback in callbacks.iter_mut() {
            (*callback)(event_data);
        }
      }
      _ => {;}
    }
  }
}
<commit_msg>Fix  for RFC 1214<commit_after>use std::collections::HashMap;

pub type EventData<T> = HashMap<String, T>;
//pub type EventCallback<'callback, T> = |& mut EventData<T>|:'callback;
pub type EventCallback<'callback, T> = FnMut(& mut EventData<T>) -> () + 'callback;

pub trait Events<'callback, T> {
  fn on(&mut self, event_name: String, callback: &'callback mut EventCallback<'callback, T>);
  fn off(&mut self, event_name: String);
  fn emit(&mut self, event_name: String, EventData: & mut EventData<T>);
}

pub struct Emitter<'callback, T: 'callback> {
  events : HashMap<String, Vec<&'callback mut EventCallback<'callback, T>>>
}

impl<'callback, T> Emitter<'callback, T>{
  pub fn new() -> Emitter<'callback, T> {
    Emitter {
      events: HashMap::new()
    }
  }
}

impl<'callback, T> Events<'callback, T> for Emitter<'callback, T> {
  fn on(&mut self, event_name: String, callback: &'callback mut EventCallback<'callback, T>) {
    if self.events.contains_key(&event_name) {
      match self.events.get_mut(&event_name) {
          Some(callbacks) => callbacks.push(callback),
          None => (),
      }
    } else {
      let vec = vec![callback];
      self.events.insert(event_name, vec);
    }

  }
  fn off(&mut self, event_name: String) {
    self.events.remove(&event_name);
  }
  fn emit(&mut self, event_name: String, event_data: & mut EventData<T>) {
    match self.events.get_mut(&event_name) {
      Some(callbacks) => {
        for callback in callbacks.iter_mut() {
            (*callback)(event_data);
        }
      }
      _ => {;}
    }
  }
}
<|endoftext|>