<commit_before>use crate::enter;
use crate::unpark_mutex::UnparkMutex;
use futures_core::future::Future;
use futures_core::task::{Context, Poll};
use futures_task::{FutureObj, Spawn, SpawnError};
use futures_task::{ArcWake, waker_ref};
use futures_util::future::FutureExt;
use std::cmp;
use std::fmt;
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;

/// A general-purpose thread pool for scheduling tasks that poll futures to
/// completion.
///
/// The thread pool multiplexes any number of tasks onto a fixed number of
/// worker threads.
///
/// This type is a clonable handle to the threadpool itself.
/// Cloning it will only create a new reference, not a new threadpool.
///
/// This type is only available when the `thread-pool` feature of this
/// library is activated.
pub struct ThreadPool {
    state: Arc<PoolState>,
}

/// Thread pool configuration object.
///
/// This type is only available when the `thread-pool` feature of this
/// library is activated.
pub struct ThreadPoolBuilder {
    pool_size: usize,
    stack_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<dyn Fn(usize) + Send + Sync>>,
    before_stop: Option<Arc<dyn Fn(usize) + Send + Sync>>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for ThreadPool {}

struct PoolState {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
}

impl fmt::Debug for ThreadPool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ThreadPool")
            .field("size", &self.state.size)
            .finish()
    }
}

impl fmt::Debug for ThreadPoolBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ThreadPoolBuilder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .finish()
    }
}

enum Message {
    Run(Task),
    Close,
}

impl ThreadPool {
    /// Creates a new thread pool with the default configuration.
    ///
    /// See documentation for the methods in
    /// [`ThreadPoolBuilder`](ThreadPoolBuilder) for details on the default
    /// configuration.
    pub fn new() -> Result<ThreadPool, io::Error> {
        ThreadPoolBuilder::new().create()
    }

    /// Create a default thread pool configuration, which can then be customized.
    ///
    /// See documentation for the methods in
    /// [`ThreadPoolBuilder`](ThreadPoolBuilder) for details on the default
    /// configuration.
    pub fn builder() -> ThreadPoolBuilder {
        ThreadPoolBuilder::new()
    }

    /// Spawns a future that will be run to completion.
    ///
    /// > **Note**: This method is similar to `Spawn::spawn_obj`, except that
    /// >           it is guaranteed to always succeed.
    pub fn spawn_obj_ok(&self, future: FutureObj<'static, ()>) {
        let task = Task {
            future,
            wake_handle: Arc::new(WakeHandle {
                exec: self.clone(),
                mutex: UnparkMutex::new(),
            }),
            exec: self.clone(),
        };
        self.state.send(Message::Run(task));
    }

    /// Spawns a task that polls the given future with output `()` to
    /// completion.
    ///
    /// ```
    /// use futures::executor::ThreadPool;
    ///
    /// let pool = ThreadPool::new().unwrap();
    ///
    /// let future = async { /* ... */ };
    /// pool.spawn_ok(future);
    /// ```
    ///
    /// > **Note**: This method is similar to `SpawnExt::spawn`, except that
    /// >           it is guaranteed to always succeed.
    pub fn spawn_ok<Fut>(&self, future: Fut)
    where
        Fut: Future<Output = ()> + Send + 'static,
    {
        self.spawn_obj_ok(FutureObj::new(Box::new(future)))
    }
}

impl Spawn for ThreadPool {
    fn spawn_obj(
        &self,
        future: FutureObj<'static, ()>,
    ) -> Result<(), SpawnError> {
        self.spawn_obj_ok(future);
        Ok(())
    }
}

impl PoolState {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self,
            idx: usize,
            after_start: Option<Arc<dyn Fn(usize) + Send + Sync>>,
            before_stop: Option<Arc<dyn Fn(usize) + Send + Sync>>) {
        let _scope = enter().unwrap();
        if let Some(after_start) = after_start {
            after_start(idx);
        }
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(task) => task.run(),
                Message::Close => break,
            }
        }
        if let Some(before_stop) = before_stop {
            before_stop(idx);
        }
    }
}

impl Clone for ThreadPool {
    fn clone(&self) -> ThreadPool {
        self.state.cnt.fetch_add(1, Ordering::Relaxed);
        ThreadPool { state: self.state.clone() }
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        if self.state.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.state.size {
                self.state.send(Message::Close);
            }
        }
    }
}

impl ThreadPoolBuilder {
    /// Create a default thread pool configuration.
    ///
    /// See the other methods on this type for details on the defaults.
    pub fn new() -> ThreadPoolBuilder {
        ThreadPoolBuilder {
            pool_size: cmp::max(1, num_cpus::get()),
            stack_size: 0,
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future ThreadPool
    ///
    /// The size of a thread pool is the number of worker threads spawned. By
    /// default, this is equal to the number of CPU cores.
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set stack size of threads in the pool.
    ///
    /// By default, worker threads use Rust's standard stack size.
    pub fn stack_size(&mut self, stack_size: usize) -> &mut Self {
        self.stack_size = stack_size;
        self
    }

    /// Set thread name prefix of a future ThreadPool.
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    ///
    /// By default, worker threads are assigned Rust's standard thread name.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute the closure `f` immediately after each worker thread is started,
    /// but before running any tasks on it.
    ///
    /// This hook is intended for bookkeeping and monitoring.
    /// The closure `f` will be dropped after the `builder` is dropped
    /// and all worker threads in the pool have executed it.
    ///
    /// The closure provided will receive an index corresponding to the worker
    /// thread it's running on.
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn(usize) + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute closure `f` just prior to shutting down each worker thread.
    ///
    /// This hook is intended for bookkeeping and monitoring.
    /// The closure `f` will be dropped after the `builder` is droppped
    /// and all threads in the pool have executed it.
    ///
    /// The closure provided will receive an index corresponding to the worker
    /// thread it's running on.
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn(usize) + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create a [`ThreadPool`](ThreadPool) with the given configuration.
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> Result<ThreadPool, io::Error> {
        let (tx, rx) = mpsc::channel();
        let pool = ThreadPool {
            state: Arc::new(PoolState {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let state = pool.state.clone();
            let after_start = self.after_start.clone();
            let before_stop = self.before_stop.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            if self.stack_size > 0 {
                thread_builder = thread_builder.stack_size(self.stack_size);
            }
            thread_builder.spawn(move || state.work(counter, after_start, before_stop))?;
        }
        Ok(pool)
    }
}

impl Default for ThreadPoolBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// A task responsible for polling a future to completion.
struct Task {
    future: FutureObj<'static, ()>,
    exec: ThreadPool,
    wake_handle: Arc<WakeHandle>,
}

struct WakeHandle {
    mutex: UnparkMutex<Task>,
    exec: ThreadPool,
}

impl Task {
    /// Actually run the task (invoking `poll` on the future) on the current
    /// thread.
    fn run(self) {
        let Task { mut future, wake_handle, mut exec } = self;
        let waker = waker_ref(&wake_handle);
        let mut cx = Context::from_waker(&waker);

        // Safety: The ownership of this `Task` object is evidence that
        // we are in the `POLLING`/`REPOLL` state for the mutex.
        unsafe {
            wake_handle.mutex.start_poll();

            loop {
                let res = future.poll_unpin(&mut cx);
                match res {
                    Poll::Pending => {}
                    Poll::Ready(()) => return wake_handle.mutex.complete(),
                }
                let task = Task {
                    future,
                    wake_handle: wake_handle.clone(),
                    exec,
                };
                match wake_handle.mutex.wait(task) {
                    Ok(()) => return, // we've waited
                    Err(task) => { // someone's notified us
                        future = task.future;
                        exec = task.exec;
                    }
                }
            }
        }
    }
}

impl fmt::Debug for Task {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Task")
            .field("contents", &"...")
            .finish()
    }
}

impl ArcWake for WakeHandle {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        match arc_self.mutex.notify() {
            Ok(task) => arc_self.exec.state.send(Message::Run(task)),
            Err(()) => {}
        }
    }
}

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

    #[test]
    fn test_drop_after_start() {
        let (tx, rx) = mpsc::sync_channel(2);
        let _cpu_pool = ThreadPoolBuilder::new()
            .pool_size(2)
            .after_start(move |_| tx.send(1).unwrap()).create().unwrap();

        // After ThreadPoolBuilder is deconstructed, the tx should be droped
        // so that we can use rx as an iterator.
        let count = rx.into_iter().count();
        assert_eq!(count, 2);
    }
}
<commit_msg>remove check for create and add check to pool size<commit_after>use crate::enter;
use crate::unpark_mutex::UnparkMutex;
use futures_core::future::Future;
use futures_core::task::{Context, Poll};
use futures_task::{FutureObj, Spawn, SpawnError};
use futures_task::{ArcWake, waker_ref};
use futures_util::future::FutureExt;
use std::cmp;
use std::fmt;
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;

/// A general-purpose thread pool for scheduling tasks that poll futures to
/// completion.
///
/// The thread pool multiplexes any number of tasks onto a fixed number of
/// worker threads.
///
/// This type is a clonable handle to the threadpool itself.
/// Cloning it will only create a new reference, not a new threadpool.
///
/// This type is only available when the `thread-pool` feature of this
/// library is activated.
pub struct ThreadPool {
    state: Arc<PoolState>,
}

/// Thread pool configuration object.
///
/// This type is only available when the `thread-pool` feature of this
/// library is activated.
pub struct ThreadPoolBuilder {
    pool_size: usize,
    stack_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<dyn Fn(usize) + Send + Sync>>,
    before_stop: Option<Arc<dyn Fn(usize) + Send + Sync>>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for ThreadPool {}

struct PoolState {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
}

impl fmt::Debug for ThreadPool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ThreadPool")
            .field("size", &self.state.size)
            .finish()
    }
}

impl fmt::Debug for ThreadPoolBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ThreadPoolBuilder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .finish()
    }
}

enum Message {
    Run(Task),
    Close,
}

impl ThreadPool {
    /// Creates a new thread pool with the default configuration.
    ///
    /// See documentation for the methods in
    /// [`ThreadPoolBuilder`](ThreadPoolBuilder) for details on the default
    /// configuration.
    pub fn new() -> Result<ThreadPool, io::Error> {
        ThreadPoolBuilder::new().create()
    }

    /// Create a default thread pool configuration, which can then be customized.
    ///
    /// See documentation for the methods in
    /// [`ThreadPoolBuilder`](ThreadPoolBuilder) for details on the default
    /// configuration.
    pub fn builder() -> ThreadPoolBuilder {
        ThreadPoolBuilder::new()
    }

    /// Spawns a future that will be run to completion.
    ///
    /// > **Note**: This method is similar to `Spawn::spawn_obj`, except that
    /// >           it is guaranteed to always succeed.
    pub fn spawn_obj_ok(&self, future: FutureObj<'static, ()>) {
        let task = Task {
            future,
            wake_handle: Arc::new(WakeHandle {
                exec: self.clone(),
                mutex: UnparkMutex::new(),
            }),
            exec: self.clone(),
        };
        self.state.send(Message::Run(task));
    }

    /// Spawns a task that polls the given future with output `()` to
    /// completion.
    ///
    /// ```
    /// use futures::executor::ThreadPool;
    ///
    /// let pool = ThreadPool::new().unwrap();
    ///
    /// let future = async { /* ... */ };
    /// pool.spawn_ok(future);
    /// ```
    ///
    /// > **Note**: This method is similar to `SpawnExt::spawn`, except that
    /// >           it is guaranteed to always succeed.
    pub fn spawn_ok<Fut>(&self, future: Fut)
    where
        Fut: Future<Output = ()> + Send + 'static,
    {
        self.spawn_obj_ok(FutureObj::new(Box::new(future)))
    }
}

impl Spawn for ThreadPool {
    fn spawn_obj(
        &self,
        future: FutureObj<'static, ()>,
    ) -> Result<(), SpawnError> {
        self.spawn_obj_ok(future);
        Ok(())
    }
}

impl PoolState {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self,
            idx: usize,
            after_start: Option<Arc<dyn Fn(usize) + Send + Sync>>,
            before_stop: Option<Arc<dyn Fn(usize) + Send + Sync>>) {
        let _scope = enter().unwrap();
        if let Some(after_start) = after_start {
            after_start(idx);
        }
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(task) => task.run(),
                Message::Close => break,
            }
        }
        if let Some(before_stop) = before_stop {
            before_stop(idx);
        }
    }
}

impl Clone for ThreadPool {
    fn clone(&self) -> ThreadPool {
        self.state.cnt.fetch_add(1, Ordering::Relaxed);
        ThreadPool { state: self.state.clone() }
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        if self.state.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.state.size {
                self.state.send(Message::Close);
            }
        }
    }
}

impl ThreadPoolBuilder {
    /// Create a default thread pool configuration.
    ///
    /// See the other methods on this type for details on the defaults.
    pub fn new() -> Self {
        Self {
            pool_size: cmp::max(1, num_cpus::get()),
            stack_size: 0,
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future ThreadPool
    ///
    /// The size of a thread pool is the number of worker threads spawned. By
    /// default, this is equal to the number of CPU cores.
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        assert!(size > 0);
        self.pool_size = size;
        self
    }

    /// Set stack size of threads in the pool.
    ///
    /// By default, worker threads use Rust's standard stack size.
    pub fn stack_size(&mut self, stack_size: usize) -> &mut Self {
        self.stack_size = stack_size;
        self
    }

    /// Set thread name prefix of a future ThreadPool.
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    ///
    /// By default, worker threads are assigned Rust's standard thread name.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute the closure `f` immediately after each worker thread is started,
    /// but before running any tasks on it.
    ///
    /// This hook is intended for bookkeeping and monitoring.
    /// The closure `f` will be dropped after the `builder` is dropped
    /// and all worker threads in the pool have executed it.
    ///
    /// The closure provided will receive an index corresponding to the worker
    /// thread it's running on.
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn(usize) + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute closure `f` just prior to shutting down each worker thread.
    ///
    /// This hook is intended for bookkeeping and monitoring.
    /// The closure `f` will be dropped after the `builder` is droppped
    /// and all threads in the pool have executed it.
    ///
    /// The closure provided will receive an index corresponding to the worker
    /// thread it's running on.
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn(usize) + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create a [`ThreadPool`](ThreadPool) with the given configuration.
    pub fn create(&mut self) -> Result<ThreadPool, io::Error> {
        let (tx, rx) = mpsc::channel();
        let pool = ThreadPool {
            state: Arc::new(PoolState {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
            }),
        };

        for counter in 0..self.pool_size {
            let state = pool.state.clone();
            let after_start = self.after_start.clone();
            let before_stop = self.before_stop.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            if self.stack_size > 0 {
                thread_builder = thread_builder.stack_size(self.stack_size);
            }
            thread_builder.spawn(move || state.work(counter, after_start, before_stop))?;
        }
        Ok(pool)
    }
}

impl Default for ThreadPoolBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// A task responsible for polling a future to completion.
struct Task {
    future: FutureObj<'static, ()>,
    exec: ThreadPool,
    wake_handle: Arc<WakeHandle>,
}

struct WakeHandle {
    mutex: UnparkMutex<Task>,
    exec: ThreadPool,
}

impl Task {
    /// Actually run the task (invoking `poll` on the future) on the current
    /// thread.
    fn run(self) {
        let Task { mut future, wake_handle, mut exec } = self;
        let waker = waker_ref(&wake_handle);
        let mut cx = Context::from_waker(&waker);

        // Safety: The ownership of this `Task` object is evidence that
        // we are in the `POLLING`/`REPOLL` state for the mutex.
        unsafe {
            wake_handle.mutex.start_poll();

            loop {
                let res = future.poll_unpin(&mut cx);
                match res {
                    Poll::Pending => {}
                    Poll::Ready(()) => return wake_handle.mutex.complete(),
                }
                let task = Task {
                    future,
                    wake_handle: wake_handle.clone(),
                    exec,
                };
                match wake_handle.mutex.wait(task) {
                    Ok(()) => return, // we've waited
                    Err(task) => { // someone's notified us
                        future = task.future;
                        exec = task.exec;
                    }
                }
            }
        }
    }
}

impl fmt::Debug for Task {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Task")
            .field("contents", &"...")
            .finish()
    }
}

impl ArcWake for WakeHandle {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        match arc_self.mutex.notify() {
            Ok(task) => arc_self.exec.state.send(Message::Run(task)),
            Err(()) => {}
        }
    }
}

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

    #[test]
    fn test_drop_after_start() {
        let (tx, rx) = mpsc::sync_channel(2);
        let _cpu_pool = ThreadPoolBuilder::new()
            .pool_size(2)
            .after_start(move |_| tx.send(1).unwrap()).create().unwrap();

        // After ThreadPoolBuilder is deconstructed, the tx should be droped
        // so that we can use rx as an iterator.
        let count = rx.into_iter().count();
        assert_eq!(count, 2);
    }
}
<|endoftext|>


<commit_before>use futures::stable::block_on_stable;
use futures::executor::{block_on, ThreadPool};
use futures::prelude::*;

#[async]
fn foo() -> Result<i32, i32> {
    Ok(1)
}

#[async]
fn bar(x: &i32) -> Result<i32, i32> {
    Ok(*x)
}

#[async]
fn baz(x: i32) -> Result<i32, i32> {
    await!(bar(&x))
}

#[async(boxed)]
fn boxed(x: i32) -> Result<i32, i32> {
    Ok(x)
}

#[async(boxed)]
fn boxed_borrow(x: &i32) -> Result<i32, i32> {
    Ok(*x)
}

#[async(boxed, send)]
fn boxed_send(x: i32) -> Result<i32, i32> {
    Ok(x)
}

#[async(boxed, send)]
fn boxed_send_borrow(x: &i32) -> Result<i32, i32> {
    Ok(*x)
}

#[async(boxed, send)]
fn spawnable() -> Result<(), Never> {
    Ok(())
}

#[async_stream(item = u64)]
fn _stream1() -> Result<(), i32> {
    fn integer() -> u64 { 1 }
    let x = &integer();
    stream_yield!(0);
    stream_yield!(*x);
    Ok(())
}

#[async_stream(boxed, item = u64)]
fn _stream_boxed() -> Result<(), i32> {
    fn integer() -> u64 { 1 }
    let x = &integer();
    stream_yield!(0);
    stream_yield!(*x);
    Ok(())
}

#[async]
pub fn uses_async_for() -> Result<Vec<u64>, i32> {
    let mut v = vec![];
    #[async]
    for i in _stream1() {
        v.push(i);
    }
    Ok(v)
}

#[test]
fn main() {
    assert_eq!(block_on_stable(foo()), Ok(1));
    assert_eq!(block_on_stable(bar(&1)), Ok(1));
    assert_eq!(block_on_stable(baz(17)), Ok(17));
    assert_eq!(block_on(boxed(17)), Ok(17));
    assert_eq!(block_on(boxed_send(17)), Ok(17));
    assert_eq!(block_on(boxed_borrow(&17)), Ok(17));
    assert_eq!(block_on(boxed_send_borrow(&17)), Ok(17));
    assert_eq!(block_on_stable(uses_async_for()), Ok(vec![0, 1]));
}

#[test]
fn run_pinned_future_in_thread_pool() {
    let mut pool = ThreadPool::new().unwrap();
    pool.spawn_pinned(spawnable()).unwrap();
}
<commit_msg>Add test cases for async_block variants<commit_after>use futures::stable::block_on_stable;
use futures::executor::{block_on, ThreadPool};
use futures::prelude::*;

#[async]
fn foo() -> Result<i32, i32> {
    Ok(1)
}

#[async]
fn bar(x: &i32) -> Result<i32, i32> {
    Ok(*x)
}

#[async]
fn baz(x: i32) -> Result<i32, i32> {
    await!(bar(&x))
}

#[async(boxed)]
fn boxed(x: i32) -> Result<i32, i32> {
    Ok(x)
}

#[async(boxed)]
fn boxed_borrow(x: &i32) -> Result<i32, i32> {
    Ok(*x)
}

#[async(boxed, send)]
fn boxed_send(x: i32) -> Result<i32, i32> {
    Ok(x)
}

#[async(boxed, send)]
fn boxed_send_borrow(x: &i32) -> Result<i32, i32> {
    Ok(*x)
}

#[async(boxed, send)]
fn spawnable() -> Result<(), Never> {
    Ok(())
}

fn baz_block(x: i32) -> impl StableFuture<Item = i32, Error = i32> {
    async_block! {
        await!(bar(&x))
    }
}

#[async_stream(item = u64)]
fn stream1() -> Result<(), i32> {
    fn integer() -> u64 { 1 }
    let x = &integer();
    stream_yield!(0);
    stream_yield!(*x);
    Ok(())
}

fn stream1_block() -> impl StableStream<Item = u64, Error = i32> {
    async_stream_block! {
        #[async]
        for item in stream1() {
            stream_yield!(item)
        }
        Ok(())
    }
}

#[async_stream(boxed, item = u64)]
fn _stream_boxed() -> Result<(), i32> {
    fn integer() -> u64 { 1 }
    let x = &integer();
    stream_yield!(0);
    stream_yield!(*x);
    Ok(())
}

#[async]
pub fn uses_async_for() -> Result<Vec<u64>, i32> {
    let mut v = vec![];
    #[async]
    for i in stream1() {
        v.push(i);
    }
    Ok(v)
}

#[async]
pub fn uses_async_for_block() -> Result<Vec<u64>, i32> {
    let mut v = vec![];
    #[async]
    for i in stream1_block() {
        v.push(i);
    }
    Ok(v)
}

#[test]
fn main() {
    assert_eq!(block_on_stable(foo()), Ok(1));
    assert_eq!(block_on_stable(bar(&1)), Ok(1));
    assert_eq!(block_on_stable(baz(17)), Ok(17));
    assert_eq!(block_on(boxed(17)), Ok(17));
    assert_eq!(block_on(boxed_send(17)), Ok(17));
    assert_eq!(block_on(boxed_borrow(&17)), Ok(17));
    assert_eq!(block_on(boxed_send_borrow(&17)), Ok(17));
    assert_eq!(block_on_stable(baz_block(18)), Ok(18));
    assert_eq!(block_on_stable(uses_async_for()), Ok(vec![0, 1]));
    assert_eq!(block_on_stable(uses_async_for_block()), Ok(vec![0, 1]));
}

#[test]
fn run_pinned_future_in_thread_pool() {
    let mut pool = ThreadPool::new().unwrap();
    pool.spawn_pinned(spawnable()).unwrap();
}
<|endoftext|>


<commit_before>use test_utils::tested_by;
use ra_db::SourceDatabase;
use ra_syntax::{
    AstNode, SyntaxNode, TextUnit, TextRange,
    SyntaxKind::FN_DEF,
    ast::{self, ArgListOwner},
    algo::find_node_at_offset,
};
use hir::Docs;

use crate::{FilePosition, CallInfo, db::RootDatabase};

/// Computes parameter information for the given call expression.
pub(crate) fn call_info(db: &RootDatabase, position: FilePosition) -> Option<CallInfo> {
    let file = db.parse(position.file_id);
    let syntax = file.syntax();

    // Find the calling expression and it's NameRef
    let calling_node = FnCallNode::with_node(syntax, position.offset)?;
    let name_ref = calling_node.name_ref()?;

    // Resolve the function's NameRef (NOTE: this isn't entirely accurate).
    let file_symbols = crate::symbol_index::index_resolve(db, name_ref);
    let symbol = file_symbols.into_iter().find(|it| it.ptr.kind() == FN_DEF)?;
    let fn_file = db.parse(symbol.file_id);
    let fn_def = symbol.ptr.to_node(&fn_file);
    let fn_def = ast::FnDef::cast(fn_def).unwrap();
    let function = hir::source_binder::function_from_source(db, symbol.file_id, fn_def)?;

    let mut call_info = CallInfo::new(db, function, fn_def)?;
    // If we have a calling expression let's find which argument we are on
    let num_params = call_info.parameters.len();
    let has_self = fn_def.param_list().and_then(|l| l.self_param()).is_some();

    if num_params == 1 {
        if !has_self {
            call_info.active_parameter = Some(0);
        }
    } else if num_params > 1 {
        // Count how many parameters into the call we are.
        // TODO: This is best effort for now and should be fixed at some point.
        // It may be better to see where we are in the arg_list and then check
        // where offset is in that list (or beyond).
        // Revisit this after we get documentation comments in.
        if let Some(ref arg_list) = calling_node.arg_list() {
            let arg_list_range = arg_list.syntax().range();
            if !arg_list_range.contains_inclusive(position.offset) {
                tested_by!(call_info_bad_offset);
                return None;
            }
            let start = arg_list_range.start();

            let range_search = TextRange::from_to(start, position.offset);
            let mut commas: usize =
                arg_list.syntax().text().slice(range_search).to_string().matches(',').count();

            // If we have a method call eat the first param since it's just self.
            if has_self {
                commas += 1;
            }

            call_info.active_parameter = Some(commas);
        }
    }

    Some(call_info)
}

enum FnCallNode<'a> {
    CallExpr(&'a ast::CallExpr),
    MethodCallExpr(&'a ast::MethodCallExpr),
}

impl<'a> FnCallNode<'a> {
    pub fn with_node(syntax: &'a SyntaxNode, offset: TextUnit) -> Option<FnCallNode<'a>> {
        if let Some(expr) = find_node_at_offset::<ast::CallExpr>(syntax, offset) {
            return Some(FnCallNode::CallExpr(expr));
        }
        if let Some(expr) = find_node_at_offset::<ast::MethodCallExpr>(syntax, offset) {
            return Some(FnCallNode::MethodCallExpr(expr));
        }
        None
    }

    pub fn name_ref(&self) -> Option<&'a ast::NameRef> {
        match *self {
            FnCallNode::CallExpr(call_expr) => Some(match call_expr.expr()?.kind() {
                ast::ExprKind::PathExpr(path_expr) => path_expr.path()?.segment()?.name_ref()?,
                _ => return None,
            }),

            FnCallNode::MethodCallExpr(call_expr) => {
                call_expr.syntax().children().filter_map(ast::NameRef::cast).nth(0)
            }
        }
    }

    pub fn arg_list(&self) -> Option<&'a ast::ArgList> {
        match *self {
            FnCallNode::CallExpr(expr) => expr.arg_list(),
            FnCallNode::MethodCallExpr(expr) => expr.arg_list(),
        }
    }
}

impl CallInfo {
    fn new(db: &RootDatabase, function: hir::Function, node: &ast::FnDef) -> Option<Self> {
        let label = crate::completion::function_label(node)?;
        let doc = function.docs(db);

        Some(CallInfo { parameters: param_list(node), label, doc, active_parameter: None })
    }
}

fn param_list(node: &ast::FnDef) -> Vec<String> {
    let mut res = vec![];
    if let Some(param_list) = node.param_list() {
        if let Some(self_param) = param_list.self_param() {
            res.push(self_param.syntax().text().to_string())
        }

        // Maybe use param.pat here? See if we can just extract the name?
        //res.extend(param_list.params().map(|p| p.syntax().text().to_string()));
        res.extend(
            param_list.params().filter_map(|p| p.pat()).map(|pat| pat.syntax().text().to_string()),
        );
    }
    res
}

#[cfg(test)]
mod tests {
    use test_utils::covers;

    use crate::mock_analysis::single_file_with_position;

    use super::*;

    fn call_info(text: &str) -> CallInfo {
        let (analysis, position) = single_file_with_position(text);
        analysis.call_info(position).unwrap().unwrap()
    }

    #[test]
    fn test_fn_signature_two_args_first() {
        let info = call_info(
            r#"fn foo(x: u32, y: u32) -> u32 {x + y}
fn bar() { foo(<|>3, ); }"#,
        );

        assert_eq!(info.parameters, vec!("x".to_string(), "y".to_string()));
        assert_eq!(info.active_parameter, Some(0));
    }

    #[test]
    fn test_fn_signature_two_args_second() {
        let info = call_info(
            r#"fn foo(x: u32, y: u32) -> u32 {x + y}
fn bar() { foo(3, <|>); }"#,
        );

        assert_eq!(info.parameters, vec!("x".to_string(), "y".to_string()));
        assert_eq!(info.active_parameter, Some(1));
    }

    #[test]
    fn test_fn_signature_for_impl() {
        let info = call_info(
            r#"struct F; impl F { pub fn new() { F{}} }
fn bar() {let _ : F = F::new(<|>);}"#,
        );

        assert_eq!(info.parameters, Vec::<String>::new());
        assert_eq!(info.active_parameter, None);
    }

    #[test]
    fn test_fn_signature_for_method_self() {
        let info = call_info(
            r#"struct F;
impl F {
    pub fn new() -> F{
        F{}
    }

    pub fn do_it(&self) {}
}

fn bar() {
    let f : F = F::new();
    f.do_it(<|>);
}"#,
        );

        assert_eq!(info.parameters, vec!["&self".to_string()]);
        assert_eq!(info.active_parameter, None);
    }

    #[test]
    fn test_fn_signature_for_method_with_arg() {
        let info = call_info(
            r#"struct F;
impl F {
    pub fn new() -> F{
        F{}
    }

    pub fn do_it(&self, x: i32) {}
}

fn bar() {
    let f : F = F::new();
    f.do_it(<|>);
}"#,
        );

        assert_eq!(info.parameters, vec!["&self".to_string(), "x".to_string()]);
        assert_eq!(info.active_parameter, Some(1));
    }

    #[test]
    fn test_fn_signature_with_docs_simple() {
        let info = call_info(
            r#"
/// test
// non-doc-comment
fn foo(j: u32) -> u32 {
    j
}

fn bar() {
    let _ = foo(<|>);
}
"#,
        );

        assert_eq!(info.parameters, vec!["j".to_string()]);
        assert_eq!(info.active_parameter, Some(0));
        assert_eq!(info.label, "fn foo(j: u32) -> u32".to_string());
        assert_eq!(info.doc.map(|it| it.into()), Some("test".to_string()));
    }

    #[test]
    fn test_fn_signature_with_docs() {
        let info = call_info(
            r#"
/// Adds one to the number given.
///
/// # Examples
///
/// ```
/// let five = 5;
///
/// assert_eq!(6, my_crate::add_one(5));
/// ```
pub fn add_one(x: i32) -> i32 {
    x + 1
}

pub fn do() {
    add_one(<|>
}"#,
        );

        assert_eq!(info.parameters, vec!["x".to_string()]);
        assert_eq!(info.active_parameter, Some(0));
        assert_eq!(info.label, "pub fn add_one(x: i32) -> i32".to_string());
        assert_eq!(
            info.doc.map(|it| it.into()),
            Some(
                r#"Adds one to the number given.

# Examples

```
let five = 5;

assert_eq!(6, my_crate::add_one(5));
```"#
                    .to_string()
            )
        );
    }

    #[test]
    fn test_fn_signature_with_docs_impl() {
        let info = call_info(
            r#"
struct addr;
impl addr {
    /// Adds one to the number given.
    ///
    /// # Examples
    ///
    /// ```
    /// let five = 5;
    ///
    /// assert_eq!(6, my_crate::add_one(5));
    /// ```
    pub fn add_one(x: i32) -> i32 {
        x + 1
    }
}

pub fn do_it() {
    addr {};
    addr::add_one(<|>);
}"#,
        );

        assert_eq!(info.parameters, vec!["x".to_string()]);
        assert_eq!(info.active_parameter, Some(0));
        assert_eq!(info.label, "pub fn add_one(x: i32) -> i32".to_string());
        assert_eq!(
            info.doc.map(|it| it.into()),
            Some(
                r#"Adds one to the number given.

# Examples

```
let five = 5;

assert_eq!(6, my_crate::add_one(5));
```"#
                    .to_string()
            )
        );
    }

    #[test]
    fn test_fn_signature_with_docs_from_actix() {
        let info = call_info(
            r#"
pub trait WriteHandler<E>
where
    Self: Actor,
    Self::Context: ActorContext,
{
    /// Method is called when writer emits error.
    ///
    /// If this method returns `ErrorAction::Continue` writer processing
    /// continues otherwise stream processing stops.
    fn error(&mut self, err: E, ctx: &mut Self::Context) -> Running {
        Running::Stop
    }

    /// Method is called when writer finishes.
    ///
    /// By default this method stops actor's `Context`.
    fn finished(&mut self, ctx: &mut Self::Context) {
        ctx.stop()
    }
}

pub fn foo() {
    WriteHandler r;
    r.finished(<|>);
}

"#,
        );

        assert_eq!(info.parameters, vec!["&mut self".to_string(), "ctx".to_string()]);
        assert_eq!(info.active_parameter, Some(1));
        assert_eq!(
            info.doc.map(|it| it.into()),
            Some(
                r#"Method is called when writer finishes.

By default this method stops actor's `Context`."#
                    .to_string()
            )
        );
    }

    #[test]
    fn call_info_bad_offset() {
        covers!(call_info_bad_offset);
        let (analysis, position) = single_file_with_position(
            r#"fn foo(x: u32, y: u32) -> u32 {x + y}
               fn bar() { foo <|> (3, ); }"#,
        );
        let call_info = analysis.call_info(position).unwrap();
        assert!(call_info.is_none());
    }
}
<commit_msg>Be precise about the argument list<commit_after>use test_utils::tested_by;
use ra_db::SourceDatabase;
use ra_syntax::{
    AstNode, SyntaxNode, TextUnit,
    SyntaxKind::FN_DEF,
    ast::{self, ArgListOwner},
    algo::find_node_at_offset,
};
use hir::Docs;

use crate::{FilePosition, CallInfo, db::RootDatabase};

/// Computes parameter information for the given call expression.
pub(crate) fn call_info(db: &RootDatabase, position: FilePosition) -> Option<CallInfo> {
    let file = db.parse(position.file_id);
    let syntax = file.syntax();

    // Find the calling expression and it's NameRef
    let calling_node = FnCallNode::with_node(syntax, position.offset)?;
    let name_ref = calling_node.name_ref()?;

    // Resolve the function's NameRef (NOTE: this isn't entirely accurate).
    let file_symbols = crate::symbol_index::index_resolve(db, name_ref);
    let symbol = file_symbols.into_iter().find(|it| it.ptr.kind() == FN_DEF)?;
    let fn_file = db.parse(symbol.file_id);
    let fn_def = symbol.ptr.to_node(&fn_file);
    let fn_def = ast::FnDef::cast(fn_def).unwrap();
    let function = hir::source_binder::function_from_source(db, symbol.file_id, fn_def)?;

    let mut call_info = CallInfo::new(db, function, fn_def)?;
    // If we have a calling expression let's find which argument we are on
    let num_params = call_info.parameters.len();
    let has_self = fn_def.param_list().and_then(|l| l.self_param()).is_some();

    if num_params == 1 {
        if !has_self {
            call_info.active_parameter = Some(0);
        }
    } else if num_params > 1 {
        // Count how many parameters into the call we are.
        if let Some(ref arg_list) = calling_node.arg_list() {
            let arg_list_range = arg_list.syntax().range();
            if !arg_list_range.contains_inclusive(position.offset) {
                tested_by!(call_info_bad_offset);
                return None;
            }

            let param = arg_list
                .args()
                .position(|arg| arg.syntax().range().contains(position.offset))
                .or(Some(num_params - 1))
                .unwrap();

            call_info.active_parameter = Some(param);
        }
    }

    Some(call_info)
}

enum FnCallNode<'a> {
    CallExpr(&'a ast::CallExpr),
    MethodCallExpr(&'a ast::MethodCallExpr),
}

impl<'a> FnCallNode<'a> {
    pub fn with_node(syntax: &'a SyntaxNode, offset: TextUnit) -> Option<FnCallNode<'a>> {
        if let Some(expr) = find_node_at_offset::<ast::CallExpr>(syntax, offset) {
            return Some(FnCallNode::CallExpr(expr));
        }
        if let Some(expr) = find_node_at_offset::<ast::MethodCallExpr>(syntax, offset) {
            return Some(FnCallNode::MethodCallExpr(expr));
        }
        None
    }

    pub fn name_ref(&self) -> Option<&'a ast::NameRef> {
        match *self {
            FnCallNode::CallExpr(call_expr) => Some(match call_expr.expr()?.kind() {
                ast::ExprKind::PathExpr(path_expr) => path_expr.path()?.segment()?.name_ref()?,
                _ => return None,
            }),

            FnCallNode::MethodCallExpr(call_expr) => {
                call_expr.syntax().children().filter_map(ast::NameRef::cast).nth(0)
            }
        }
    }

    pub fn arg_list(&self) -> Option<&'a ast::ArgList> {
        match *self {
            FnCallNode::CallExpr(expr) => expr.arg_list(),
            FnCallNode::MethodCallExpr(expr) => expr.arg_list(),
        }
    }
}

impl CallInfo {
    fn new(db: &RootDatabase, function: hir::Function, node: &ast::FnDef) -> Option<Self> {
        let label = crate::completion::function_label(node)?;
        let doc = function.docs(db);

        Some(CallInfo { parameters: param_list(node), label, doc, active_parameter: None })
    }
}

fn param_list(node: &ast::FnDef) -> Vec<String> {
    let mut res = vec![];
    if let Some(param_list) = node.param_list() {
        if let Some(self_param) = param_list.self_param() {
            res.push(self_param.syntax().text().to_string())
        }

        // Maybe use param.pat here? See if we can just extract the name?
        //res.extend(param_list.params().map(|p| p.syntax().text().to_string()));
        res.extend(
            param_list.params().filter_map(|p| p.pat()).map(|pat| pat.syntax().text().to_string()),
        );
    }
    res
}

#[cfg(test)]
mod tests {
    use test_utils::covers;

    use crate::mock_analysis::single_file_with_position;

    use super::*;

    fn call_info(text: &str) -> CallInfo {
        let (analysis, position) = single_file_with_position(text);
        analysis.call_info(position).unwrap().unwrap()
    }

    #[test]
    fn test_fn_signature_two_args_first() {
        let info = call_info(
            r#"fn foo(x: u32, y: u32) -> u32 {x + y}
fn bar() { foo(<|>3, ); }"#,
        );

        assert_eq!(info.parameters, vec!("x".to_string(), "y".to_string()));
        assert_eq!(info.active_parameter, Some(0));
    }

    #[test]
    fn test_fn_signature_two_args_second() {
        let info = call_info(
            r#"fn foo(x: u32, y: u32) -> u32 {x + y}
fn bar() { foo(3, <|>); }"#,
        );

        assert_eq!(info.parameters, vec!("x".to_string(), "y".to_string()));
        assert_eq!(info.active_parameter, Some(1));
    }

    #[test]
    fn test_fn_signature_for_impl() {
        let info = call_info(
            r#"struct F; impl F { pub fn new() { F{}} }
fn bar() {let _ : F = F::new(<|>);}"#,
        );

        assert_eq!(info.parameters, Vec::<String>::new());
        assert_eq!(info.active_parameter, None);
    }

    #[test]
    fn test_fn_signature_for_method_self() {
        let info = call_info(
            r#"struct F;
impl F {
    pub fn new() -> F{
        F{}
    }

    pub fn do_it(&self) {}
}

fn bar() {
    let f : F = F::new();
    f.do_it(<|>);
}"#,
        );

        assert_eq!(info.parameters, vec!["&self".to_string()]);
        assert_eq!(info.active_parameter, None);
    }

    #[test]
    fn test_fn_signature_for_method_with_arg() {
        let info = call_info(
            r#"struct F;
impl F {
    pub fn new() -> F{
        F{}
    }

    pub fn do_it(&self, x: i32) {}
}

fn bar() {
    let f : F = F::new();
    f.do_it(<|>);
}"#,
        );

        assert_eq!(info.parameters, vec!["&self".to_string(), "x".to_string()]);
        assert_eq!(info.active_parameter, Some(1));
    }

    #[test]
    fn test_fn_signature_with_docs_simple() {
        let info = call_info(
            r#"
/// test
// non-doc-comment
fn foo(j: u32) -> u32 {
    j
}

fn bar() {
    let _ = foo(<|>);
}
"#,
        );

        assert_eq!(info.parameters, vec!["j".to_string()]);
        assert_eq!(info.active_parameter, Some(0));
        assert_eq!(info.label, "fn foo(j: u32) -> u32".to_string());
        assert_eq!(info.doc.map(|it| it.into()), Some("test".to_string()));
    }

    #[test]
    fn test_fn_signature_with_docs() {
        let info = call_info(
            r#"
/// Adds one to the number given.
///
/// # Examples
///
/// ```
/// let five = 5;
///
/// assert_eq!(6, my_crate::add_one(5));
/// ```
pub fn add_one(x: i32) -> i32 {
    x + 1
}

pub fn do() {
    add_one(<|>
}"#,
        );

        assert_eq!(info.parameters, vec!["x".to_string()]);
        assert_eq!(info.active_parameter, Some(0));
        assert_eq!(info.label, "pub fn add_one(x: i32) -> i32".to_string());
        assert_eq!(
            info.doc.map(|it| it.into()),
            Some(
                r#"Adds one to the number given.

# Examples

```
let five = 5;

assert_eq!(6, my_crate::add_one(5));
```"#
                    .to_string()
            )
        );
    }

    #[test]
    fn test_fn_signature_with_docs_impl() {
        let info = call_info(
            r#"
struct addr;
impl addr {
    /// Adds one to the number given.
    ///
    /// # Examples
    ///
    /// ```
    /// let five = 5;
    ///
    /// assert_eq!(6, my_crate::add_one(5));
    /// ```
    pub fn add_one(x: i32) -> i32 {
        x + 1
    }
}

pub fn do_it() {
    addr {};
    addr::add_one(<|>);
}"#,
        );

        assert_eq!(info.parameters, vec!["x".to_string()]);
        assert_eq!(info.active_parameter, Some(0));
        assert_eq!(info.label, "pub fn add_one(x: i32) -> i32".to_string());
        assert_eq!(
            info.doc.map(|it| it.into()),
            Some(
                r#"Adds one to the number given.

# Examples

```
let five = 5;

assert_eq!(6, my_crate::add_one(5));
```"#
                    .to_string()
            )
        );
    }

    #[test]
    fn test_fn_signature_with_docs_from_actix() {
        let info = call_info(
            r#"
pub trait WriteHandler<E>
where
    Self: Actor,
    Self::Context: ActorContext,
{
    /// Method is called when writer emits error.
    ///
    /// If this method returns `ErrorAction::Continue` writer processing
    /// continues otherwise stream processing stops.
    fn error(&mut self, err: E, ctx: &mut Self::Context) -> Running {
        Running::Stop
    }

    /// Method is called when writer finishes.
    ///
    /// By default this method stops actor's `Context`.
    fn finished(&mut self, ctx: &mut Self::Context) {
        ctx.stop()
    }
}

pub fn foo() {
    WriteHandler r;
    r.finished(<|>);
}

"#,
        );

        assert_eq!(info.parameters, vec!["&mut self".to_string(), "ctx".to_string()]);
        assert_eq!(info.active_parameter, Some(1));
        assert_eq!(
            info.doc.map(|it| it.into()),
            Some(
                r#"Method is called when writer finishes.

By default this method stops actor's `Context`."#
                    .to_string()
            )
        );
    }

    #[test]
    fn call_info_bad_offset() {
        covers!(call_info_bad_offset);
        let (analysis, position) = single_file_with_position(
            r#"fn foo(x: u32, y: u32) -> u32 {x + y}
               fn bar() { foo <|> (3, ); }"#,
        );
        let call_info = analysis.call_info(position).unwrap();
        assert!(call_info.is_none());
    }
}
<|endoftext|>


<commit_before>#![feature(conservative_impl_trait)]

extern crate websocket;
extern crate futures;
extern crate futures_cpupool;
extern crate tokio_core;
extern crate serde;
extern crate serde_json;
#[macro_use]
extern crate serde_derive;

use websocket::message::OwnedMessage;
use websocket::server::InvalidConnection;
use websocket::async::Server;

use tokio_core::reactor::{Handle, Core};

use futures::{Future, Sink, Stream};
use futures::future::{self, Loop};
use futures::sync::mpsc;
use futures_cpupool::CpuPool;

use std::sync::{RwLock, Arc};
use std::thread;
use std::rc::Rc;
use std::fmt::Debug;
use std::time::Duration;
use std::ops::Deref;
use std::collections::HashMap;
use std::cell::RefCell;

type Id = u32;

fn main() {
    let mut core = Core::new().expect("Failed to create Tokio event loop");
    let handle = core.handle();
    let server = Server::bind("localhost:8081", &handle).expect("Failed to create server");
    let pool = Rc::new(CpuPool::new_num_cpus());
    let connections = Arc::new(RwLock::new(HashMap::new()));
    let state = Arc::new(RwLock::new(State::new()));
    let (read_channel_out, read_channel_in) = mpsc::unbounded();
    let conn_id = Rc::new(RefCell::new(Counter::new()));

    let connections_inner = connections.clone();
    let conn_id_inner = conn_id.clone();
    let connection_handler = server.incoming()
        // we don't wanna save the stream if it drops
        .map_err(|InvalidConnection { error, .. }| error)
        .for_each(move |(upgrade, addr)| {
            let connections_inner = connections_inner.clone();
            println!("Got a connection from: {}", addr);
            let channel = read_channel_out.clone();
            let handle_inner = handle.clone();
            let conn_id = conn_id_inner.clone();
            let f = upgrade
                .use_protocol("rust-websocket")
                .accept()
                .and_then(move |(framed, _)| {
                    let (sink, stream) = framed.split();
                    let f = channel.send(stream);
                    spawn_future(f, "Senk stream to connection pool", &handle_inner);
                    let id = conn_id
                        .borrow_mut()
                        .next()
                        .expect("maximum amount of ids reached");
                    connections_inner.write().unwrap().insert(id, sink);
                    Ok(())
                });
            spawn_future(f, "Handle new connection", &handle);
            Ok(())
        })
        .map_err(|_| ());

    let state_read = state.clone();
    let remote = core.remote();
    let read_handler = pool.spawn_fn(|| {
        read_channel_in.for_each(move |stream| {
            let state_read = state_read.clone();
            remote.spawn(|_| {
                stream.for_each(move |msg| {
                                    handle_incoming(&mut state_read.write().unwrap(), &msg);
                                    Ok(())
                                })
                    .map_err(|_| ())
            });
            Ok(())
        })
    });

    let (write_channel_out, write_channel_in) = mpsc::unbounded();
    let connections_inner = connections.clone();
    let write_handler = pool.spawn_fn(move || {
        let connections = connections_inner.clone();
        write_channel_in.for_each(move |(id, state): (Id, Arc<RwLock<State>>)| {
                let connections = connections.clone();
                let mut connections = connections.write().unwrap();
                let sink = connections.remove(&id).expect("Tried to send to invalid client id");

                let msg = serde_json::to_string(state.read().unwrap().deref()).unwrap();
                println!("Sending message: {}", msg);
                let sink = sink.send(OwnedMessage::Text(msg))
                    .wait()
                    .expect("Error while sending to client");
                connections.insert(id, sink);
                Ok(())
            })
            .map_err(|_| ())
    });

    let remote = core.remote();
    let game_loop = pool.spawn_fn(move || {
        future::loop_fn(write_channel_out, move |write_channel_out| {
            thread::sleep(Duration::from_millis(100));
            let connections = connections.clone();
            let state = state.clone();
            let write_channel_out_inner = write_channel_out.clone();
            remote.spawn(move |handle| {
                for (id, _) in connections.read().unwrap().iter() {
                    let f = write_channel_out_inner.clone().send((*id, state.clone()));
                    spawn_future(f, "Send message to client", handle);
                }
                Ok(())
            });

            // insert your terminating condition here
            match 1 {
                1 => Ok(Loop::Continue(write_channel_out)),
                2 => Ok(Loop::Break(())),
                _ => Err(()),
            }
        })
    });

    let handlers =
        game_loop.select2(connection_handler.select2(read_handler.select(write_handler)));
    core.run(handlers).map_err(|_| println!("Error while running core loop")).unwrap();
}

fn spawn_future<F, I, E>(f: F, desc: &'static str, handle: &Handle)
    where F: Future<Item = I, Error = E> + 'static,
          E: Debug
{
    handle.spawn(f.map_err(move |e| println!("Error in {}: '{:?}'", desc, e))
                     .map(move |_| println!("{}: Finished.", desc)));
}

fn handle_incoming(_: &mut State, msg: &OwnedMessage) {
    if let OwnedMessage::Text(ref txt) = *msg {
        println!("Received message: {}", txt);
        //state.msg = txt.clone();
    }
}

struct Counter {
    count: Id,
}
impl Counter {
    fn new() -> Self {
        Counter { count: 0 }
    }
}

impl Iterator for Counter {
    type Item = Id;

    fn next(&mut self) -> Option<Id> {
        if self.count != Id::max_value() {
            self.count += 1;
            Some(self.count)
        } else {
            None
        }
    }
}

#[derive(Serialize, Deserialize)]
struct State {
    positions: Vec<Pos>,
}

impl State {
    fn new() -> Self {
        Self { positions: Vec::new() }
    }
}

#[derive(Serialize, Deserialize)]
struct Pos {
    x: i32,
    y: i32,
}

/*
impl Pos {
    fn new(x: i32, y: i32) -> Self {
        Self { x, y }
    }

    fn add_x(&mut self, x: i32) -> &mut Self {
        self.x += x;
        self
    }

    fn add_y(&mut self, y: i32) -> &mut Self {
        self.y += y;
        self
    }

    fn add(&mut self, x: i32, y: i32) -> &mut Self {
        self.x += x;
        self.y += y;
        self
    }
}
*/
<commit_msg>Core: Let readers know which ID sent the data<commit_after>#![feature(conservative_impl_trait)]

extern crate websocket;
extern crate futures;
extern crate futures_cpupool;
extern crate tokio_core;
extern crate serde;
extern crate serde_json;
#[macro_use]
extern crate serde_derive;

use websocket::message::OwnedMessage;
use websocket::server::InvalidConnection;
use websocket::async::Server;

use tokio_core::reactor::{Handle, Core};

use futures::{Future, Sink, Stream};
use futures::future::{self, Loop};
use futures::sync::mpsc;
use futures_cpupool::CpuPool;

use std::sync::{RwLock, Arc};
use std::thread;
use std::rc::Rc;
use std::fmt::Debug;
use std::time::Duration;
use std::ops::Deref;
use std::collections::HashMap;
use std::cell::RefCell;

type Id = u32;

fn main() {
    let mut core = Core::new().expect("Failed to create Tokio event loop");
    let handle = core.handle();
    let server = Server::bind("localhost:8081", &handle).expect("Failed to create server");
    let pool = Rc::new(CpuPool::new_num_cpus());
    let connections = Arc::new(RwLock::new(HashMap::new()));
    let state = Arc::new(RwLock::new(State::new()));
    let (read_channel_out, read_channel_in) = mpsc::unbounded();
    let conn_id = Rc::new(RefCell::new(Counter::new()));

    let connections_inner = connections.clone();
    let conn_id_inner = conn_id.clone();
    let connection_handler = server.incoming()
        // we don't wanna save the stream if it drops
        .map_err(|InvalidConnection { error, .. }| error)
        .for_each(move |(upgrade, addr)| {
            let connections_inner = connections_inner.clone();
            println!("Got a connection from: {}", addr);
            let channel = read_channel_out.clone();
            let handle_inner = handle.clone();
            let conn_id = conn_id_inner.clone();
            let f = upgrade
                .use_protocol("rust-websocket")
                .accept()
                .and_then(move |(framed, _)| {
                    let id = conn_id
                        .borrow_mut()
                        .next()
                        .expect("maximum amount of ids reached");
                    let (sink, stream) = framed.split();
                    let f = channel.send((id, stream));
                    spawn_future(f, "Senk stream to connection pool", &handle_inner);
                    connections_inner.write().unwrap().insert(id, sink);
                    Ok(())
                });
            spawn_future(f, "Handle new connection", &handle);
            Ok(())
        })
        .map_err(|_| ());

    let state_read = state.clone();
    let remote = core.remote();
    let read_handler = pool.spawn_fn(|| {
        read_channel_in.for_each(move |(id, stream)| {
            let state_read = state_read.clone();
            remote.spawn(move |_| {
                stream.for_each(move |msg| {
                                    handle_incoming(id, &msg, &mut state_read.write().unwrap());
                                    Ok(())
                                })
                    .map_err(|_| ())
            });
            Ok(())
        })
    });

    let (write_channel_out, write_channel_in) = mpsc::unbounded();
    let connections_inner = connections.clone();
    let write_handler = pool.spawn_fn(move || {
        let connections = connections_inner.clone();
        write_channel_in.for_each(move |(id, state): (Id, Arc<RwLock<State>>)| {
                let connections = connections.clone();
                let mut connections = connections.write().unwrap();

                let sink = connections.remove(&id).expect("Tried to send to invalid client id");

                let msg = serde_json::to_string(state.read().unwrap().deref()).unwrap();
                println!("Sending message: {}", msg);
                let sink = sink.send(OwnedMessage::Text(msg))
                    .wait()
                    .expect("Error while sending to client");

                connections.insert(id, sink);
                Ok(())
            })
            .map_err(|_| ())
    });

    let remote = core.remote();
    let game_loop = pool.spawn_fn(move || {
        future::loop_fn(write_channel_out, move |write_channel_out| {
            thread::sleep(Duration::from_millis(100));
            let connections = connections.clone();
            let state = state.clone();
            let write_channel_out_inner = write_channel_out.clone();
            remote.spawn(move |handle| {
                for (id, _) in connections.read().unwrap().iter() {
                    let f = write_channel_out_inner.clone().send((*id, state.clone()));
                    spawn_future(f, "Send message to client", handle);
                }
                Ok(())
            });

            // insert your terminating condition here
            match 1 {
                1 => Ok(Loop::Continue(write_channel_out)),
                2 => Ok(Loop::Break(())),
                _ => Err(()),
            }
        })
    });

    let handlers =
        game_loop.select2(connection_handler.select2(read_handler.select(write_handler)));
    core.run(handlers).map_err(|_| println!("Error while running core loop")).unwrap();
}

fn spawn_future<F, I, E>(f: F, desc: &'static str, handle: &Handle)
    where F: Future<Item = I, Error = E> + 'static,
          E: Debug
{
    handle.spawn(f.map_err(move |e| println!("Error in {}: '{:?}'", desc, e))
                     .map(move |_| println!("{}: Finished.", desc)));
}

fn handle_incoming(_:u32, msg: &OwnedMessage,  _: &mut State) {
    if let OwnedMessage::Text(ref txt) = *msg {
        println!("Received message: {}", txt);
        //state.msg = txt.clone();
    }
}

struct Counter {
    count: Id,
}
impl Counter {
    fn new() -> Self {
        Counter { count: 0 }
    }
}

impl Iterator for Counter {
    type Item = Id;

    fn next(&mut self) -> Option<Id> {
        if self.count != Id::max_value() {
            self.count += 1;
            Some(self.count)
        } else {
            None
        }
    }
}

#[derive(Serialize, Deserialize)]
struct State {
    positions: Vec<Pos>,
}

impl State {
    fn new() -> Self {
        Self { positions: Vec::new() }
    }
}

#[derive(Serialize, Deserialize)]
struct Pos {
    x: i32,
    y: i32,
}

/*
impl Pos {
    fn new(x: i32, y: i32) -> Self {
        Self { x, y }
    }

    fn add_x(&mut self, x: i32) -> &mut Self {
        self.x += x;
        self
    }

    fn add_y(&mut self, y: i32) -> &mut Self {
        self.y += y;
        self
    }

    fn add(&mut self, x: i32, y: i32) -> &mut Self {
        self.x += x;
        self.y += y;
        self
    }
}
*/
<|endoftext|>


<commit_before>use super::*;

use crate::ok_or_shutdown;
use crate::state_helper::{pause_on_failure, save_state, LockedState};

impl TaskHandler {
    /// See if we can start a new queued task.
    pub fn spawn_new(&mut self) {
        let cloned_state_mutex = self.state.clone();
        let mut state = cloned_state_mutex.lock().unwrap();
        // Get the next task id that can be started
        if let Some(id) = self.get_next_task_id(&state) {
            self.start_process(id, &mut state);
        }
    }

    /// Search and return the next task that can be started.
    /// Precondition for a task to be started:
    /// - is in Queued state
    /// - There are free slots in the task's group
    /// - The group is running
    /// - has all its dependencies in `Done` state
    pub fn get_next_task_id(&mut self, state: &LockedState) -> Option<usize> {
        // Check how many tasks are running in each group
        let mut running_tasks_per_group: HashMap<String, usize> = HashMap::new();

        // Create a default group for tasks without an explicit group
        running_tasks_per_group.insert("default".into(), 0);

        // Walk through all tasks and save the number of running tasks by group
        for (_, task) in state.tasks.iter() {
            // We are only interested in currently running tasks.
            if !matches!(task.status, TaskStatus::Running | TaskStatus::Paused) {
                continue;
            }

            match running_tasks_per_group.get_mut(&task.group) {
                Some(count) => {
                    *count += 1;
                }
                None => {
                    running_tasks_per_group.insert(task.group.clone(), 1);
                }
            }
        }

        state
            .tasks
            .iter()
            .filter(|(_, task)| task.status == TaskStatus::Queued)
            .filter(|(_, task)| {
                // The task is assigned to a group.
                // First let's check if the group is running. If it isn't, simply return false.
                if !matches!(state.groups.get(&task.group), Some(GroupStatus::Running)) {
                    return false;
                }

                // If there's no running task for the group yet, we can safely return true
                //
                // If there are running tasks for this group, we have to ensure that there are
                // fewer running tasks than allowed for this group.
                match running_tasks_per_group.get(&task.group) {
                    None => true,
                    Some(count) => match state.settings.daemon.groups.get(&task.group) {
                        Some(allowed) => count < allowed,
                        None => {
                            error!(
                                "Got task with unknown group {}. Please report this!",
                                &task.group
                            );
                            false
                        }
                    },
                }
            })
            .find(|(_, task)| {
                // Check whether all dependencies for this task are fulfilled.
                task.dependencies
                    .iter()
                    .flat_map(|id| state.tasks.get(id))
                    .all(|task| matches!(task.status, TaskStatus::Done(TaskResult::Success)))
            })
            .map(|(id, _)| *id)
    }

    /// Actually spawn a new sub process
    /// The output of subprocesses is piped into a seperate file for easier access
    pub fn start_process(&mut self, task_id: usize, state: &mut LockedState) {
        // Check if the task exists and can actually be spawned. Otherwise do an early return.
        match state.tasks.get(&task_id) {
            Some(task) => {
                if !matches!(
                    &task.status,
                    TaskStatus::Stashed { .. } | TaskStatus::Queued | TaskStatus::Paused
                ) {
                    info!("Tried to start task with status: {}", task.status);
                    return;
                }
            }
            None => {
                info!("Tried to start non-existing task: {}", task_id);
                return;
            }
        };

        // Try to get the log files to which the output of the process
        // will be written to. Error if this doesn't work!
        let (stdout_log, stderr_log) = match create_log_file_handles(task_id, &self.pueue_directory)
        {
            Ok((out, err)) => (out, err),
            Err(err) => {
                error!("Failed to create child log files: {:?}", err);
                return;
            }
        };

        // Get all necessary info for starting the task
        let (command, path, envs) = {
            let task = state.tasks.get(&task_id).unwrap();
            (task.command.clone(), task.path.clone(), task.envs.clone())
        };

        // Spawn the actual subprocess
        let mut command = compile_shell_command(&command);

        let spawned_command = command
            .current_dir(path)
            .stdin(Stdio::piped())
            .envs(envs)
            .stdout(Stdio::from(stdout_log))
            .stderr(Stdio::from(stderr_log))
            .spawn();

        // Check if the task managed to spawn
        let child = match spawned_command {
            Ok(child) => child,
            Err(err) => {
                let error = format!("Failed to spawn child {} with err: {:?}", task_id, err);
                error!("{}", error);
                clean_log_handles(task_id, &self.pueue_directory);

                // Update all necessary fields on the task.
                let group = {
                    let task = state.tasks.get_mut(&task_id).unwrap();
                    task.status = TaskStatus::Done(TaskResult::FailedToSpawn(error));
                    task.start = Some(Local::now());
                    task.end = Some(Local::now());
                    self.spawn_callback(task);

                    task.group.clone()
                };

                pause_on_failure(state, group);
                ok_or_shutdown!(self, save_state(state));
                return;
            }
        };
        self.children.insert(task_id, child);

        let task = state.tasks.get_mut(&task_id).unwrap();

        task.start = Some(Local::now());
        task.status = TaskStatus::Running;

        info!("Started task: {}", task.command);
        ok_or_shutdown!(self, save_state(state));
    }
}
<commit_msg>feature: Multi-spawn tasks per loop<commit_after>use super::*;

use crate::ok_or_shutdown;
use crate::state_helper::{pause_on_failure, save_state, LockedState};

impl TaskHandler {
    /// See if we can start a new queued task.
    pub fn spawn_new(&mut self) {
        let cloned_state_mutex = self.state.clone();
        let mut state = cloned_state_mutex.lock().unwrap();
        // Check whether a new task can be started.
        // Spawn tasks until we no longer have free slots available.
        while let Some(id) = self.get_next_task_id(&state) {
            self.start_process(id, &mut state);
        }
    }

    /// Search and return the next task that can be started.
    /// Precondition for a task to be started:
    /// - is in Queued state
    /// - There are free slots in the task's group
    /// - The group is running
    /// - has all its dependencies in `Done` state
    pub fn get_next_task_id(&mut self, state: &LockedState) -> Option<usize> {
        // Check how many tasks are running in each group
        let mut running_tasks_per_group: HashMap<String, usize> = HashMap::new();

        // Create a default group for tasks without an explicit group
        running_tasks_per_group.insert("default".into(), 0);

        // Walk through all tasks and save the number of running tasks by group
        for (_, task) in state.tasks.iter() {
            // We are only interested in currently running tasks.
            if !matches!(task.status, TaskStatus::Running | TaskStatus::Paused) {
                continue;
            }

            match running_tasks_per_group.get_mut(&task.group) {
                Some(count) => {
                    *count += 1;
                }
                None => {
                    running_tasks_per_group.insert(task.group.clone(), 1);
                }
            }
        }

        state
            .tasks
            .iter()
            .filter(|(_, task)| task.status == TaskStatus::Queued)
            .filter(|(_, task)| {
                // The task is assigned to a group.
                // First let's check if the group is running. If it isn't, simply return false.
                if !matches!(state.groups.get(&task.group), Some(GroupStatus::Running)) {
                    return false;
                }

                // If there's no running task for the group yet, we can safely return true
                //
                // If there are running tasks for this group, we have to ensure that there are
                // fewer running tasks than allowed for this group.
                match running_tasks_per_group.get(&task.group) {
                    None => true,
                    Some(count) => match state.settings.daemon.groups.get(&task.group) {
                        Some(allowed) => count < allowed,
                        None => {
                            error!(
                                "Got task with unknown group {}. Please report this!",
                                &task.group
                            );
                            false
                        }
                    },
                }
            })
            .find(|(_, task)| {
                // Check whether all dependencies for this task are fulfilled.
                task.dependencies
                    .iter()
                    .flat_map(|id| state.tasks.get(id))
                    .all(|task| matches!(task.status, TaskStatus::Done(TaskResult::Success)))
            })
            .map(|(id, _)| *id)
    }

    /// Actually spawn a new sub process
    /// The output of subprocesses is piped into a seperate file for easier access
    pub fn start_process(&mut self, task_id: usize, state: &mut LockedState) {
        // Check if the task exists and can actually be spawned. Otherwise do an early return.
        match state.tasks.get(&task_id) {
            Some(task) => {
                if !matches!(
                    &task.status,
                    TaskStatus::Stashed { .. } | TaskStatus::Queued | TaskStatus::Paused
                ) {
                    info!("Tried to start task with status: {}", task.status);
                    return;
                }
            }
            None => {
                info!("Tried to start non-existing task: {}", task_id);
                return;
            }
        };

        // Try to get the log files to which the output of the process
        // will be written to. Error if this doesn't work!
        let (stdout_log, stderr_log) = match create_log_file_handles(task_id, &self.pueue_directory)
        {
            Ok((out, err)) => (out, err),
            Err(err) => {
                error!("Failed to create child log files: {:?}", err);
                return;
            }
        };

        // Get all necessary info for starting the task
        let (command, path, envs) = {
            let task = state.tasks.get(&task_id).unwrap();
            (task.command.clone(), task.path.clone(), task.envs.clone())
        };

        // Spawn the actual subprocess
        let mut command = compile_shell_command(&command);

        let spawned_command = command
            .current_dir(path)
            .stdin(Stdio::piped())
            .envs(envs)
            .stdout(Stdio::from(stdout_log))
            .stderr(Stdio::from(stderr_log))
            .spawn();

        // Check if the task managed to spawn
        let child = match spawned_command {
            Ok(child) => child,
            Err(err) => {
                let error = format!("Failed to spawn child {} with err: {:?}", task_id, err);
                error!("{}", error);
                clean_log_handles(task_id, &self.pueue_directory);

                // Update all necessary fields on the task.
                let group = {
                    let task = state.tasks.get_mut(&task_id).unwrap();
                    task.status = TaskStatus::Done(TaskResult::FailedToSpawn(error));
                    task.start = Some(Local::now());
                    task.end = Some(Local::now());
                    self.spawn_callback(task);

                    task.group.clone()
                };

                pause_on_failure(state, group);
                ok_or_shutdown!(self, save_state(state));
                return;
            }
        };
        self.children.insert(task_id, child);

        let task = state.tasks.get_mut(&task_id).unwrap();

        task.start = Some(Local::now());
        task.status = TaskStatus::Running;

        info!("Started task: {}", task.command);
        ok_or_shutdown!(self, save_state(state));
    }
}
<|endoftext|>


<commit_before>//! Simplistic and relatively unoptimised handling of basic
//! tasks around primes:
//!
//! - 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 basic Sieve of Eratosthenes to enumerate the primes up to
//! some fixed bound (in a relatively memory efficient manner), and then
//! allows this cached information to be used for things like enumerating
//! the primes, and factorisation via trial division.
//!
//! (Despite the name, it can sieve the primes up to 109 in
//! about 5 seconds.)
//!
//! [*Source*](http://github.com/huonw/primal)
//!
//! # Example
//!
//! Let's find the 10001st prime. The basic idea is to enumerate the
//! primes and then take the 10001st in that list.
//!
//! Unfortunately, `Primes::sieve` takes an upper bound, and it gives
//! us no information beyond this; so we really need some way to find
//! an upper bound to be guaranteed to include the 10001st prime. If
//! we had an a priori number we could just use that, but we don't
//! (for the purposes of this example, anyway). Hence, we can either
//! try filtering with exponentially larger upper bounds until we find
//! one that works (e.g. doubling each time), or just take a shortcut
//! and use deeper mathematics via
//! [`estimate_nth_prime`](fn.estimate_nth_prime.html).
//!
//! ```rust
//! // find our upper bound
//! let (_lo, hi) = primal::estimate_nth_prime(10001);
//!
//! // find the primes up to this upper bound
//! let sieve = primal::Primes::sieve(hi as usize);
//!
//! // (.nth is zero indexed.)
//! match sieve.primes().nth(10001 - 1) {
//!     Some(p) => println!("The 10001st prime is {}", p), // 104743
//!     None => unreachable!(),
//! }
//! ```
//!
//! # Using this library
//!
//! Just add the following to your [`Cargo.toml`](http://crates.io/):
//!
//! ```toml
//! [dependencies.primal]
//! git = "https://github.com/huonw/primal"
//! ```

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

extern crate num as num_;
extern crate primal_estimate;
extern crate primal_bit;
extern crate primal_check;
extern crate primal_smallsieve;
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 fast_sieve::Sieve;
pub use primal_check::miller_rabin as is_prime;
pub use primal_check::{as_perfect_power, as_prime_power};
pub use primal_smallsieve::{Primes, PrimeIterator};

pub use primal_sieve::StreamingSieve;

/// (prime, exponent) pairs storing the prime factorisation of a
/// number.
pub type Factors = Vec<(usize, usize)>;

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

    use super::{Primes, 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 = Primes::sieve(1_000_000);
            (1..N).step_by(STEP)
                .filter(|&n| sieve.is_prime(n)).count()
        })
    }
}
<commit_msg>Export non-streaming sieve.<commit_after>//! Simplistic and relatively unoptimised handling of basic
//! tasks around primes:
//!
//! - 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 basic Sieve of Eratosthenes to enumerate the primes up to
//! some fixed bound (in a relatively memory efficient manner), and then
//! allows this cached information to be used for things like enumerating
//! the primes, and factorisation via trial division.
//!
//! (Despite the name, it can sieve the primes up to 109 in
//! about 5 seconds.)
//!
//! [*Source*](http://github.com/huonw/primal)
//!
//! # Example
//!
//! Let's find the 10001st prime. The basic idea is to enumerate the
//! primes and then take the 10001st in that list.
//!
//! Unfortunately, `Primes::sieve` takes an upper bound, and it gives
//! us no information beyond this; so we really need some way to find
//! an upper bound to be guaranteed to include the 10001st prime. If
//! we had an a priori number we could just use that, but we don't
//! (for the purposes of this example, anyway). Hence, we can either
//! try filtering with exponentially larger upper bounds until we find
//! one that works (e.g. doubling each time), or just take a shortcut
//! and use deeper mathematics via
//! [`estimate_nth_prime`](fn.estimate_nth_prime.html).
//!
//! ```rust
//! // find our upper bound
//! let (_lo, hi) = primal::estimate_nth_prime(10001);
//!
//! // find the primes up to this upper bound
//! let sieve = primal::Primes::sieve(hi as usize);
//!
//! // (.nth is zero indexed.)
//! match sieve.primes().nth(10001 - 1) {
//!     Some(p) => println!("The 10001st prime is {}", p), // 104743
//!     None => unreachable!(),
//! }
//! ```
//!
//! # Using this library
//!
//! Just add the following to your [`Cargo.toml`](http://crates.io/):
//!
//! ```toml
//! [dependencies.primal]
//! git = "https://github.com/huonw/primal"
//! ```

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

extern crate num as num_;
extern crate primal_estimate;
extern crate primal_bit;
extern crate primal_check;
extern crate primal_smallsieve;
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 fast_sieve::Sieve;
pub use primal_check::miller_rabin as is_prime;
pub use primal_check::{as_perfect_power, as_prime_power};
pub use primal_smallsieve::{Primes, PrimeIterator};

pub use primal_sieve::{StreamingSieve, Sieve};

/// (prime, exponent) pairs storing the prime factorisation of a
/// number.
pub type Factors = Vec<(usize, usize)>;

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

    use super::{Primes, 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 = Primes::sieve(1_000_000);
            (1..N).step_by(STEP)
                .filter(|&n| sieve.is_prime(n)).count()
        })
    }
}
<|endoftext|>


<commit_before>#![feature(convert,custom_derive,plugin)]
#![plugin(serde_macros)]
extern crate hyper;
extern crate serde;
extern crate serde_json;

use std::io::Read;

use hyper::Client;
use hyper::client::IntoUrl;
use hyper::client::Response;
use hyper::net::HttpsConnector;
use hyper::net::Openssl;
use serde_json::Value;

pub struct LxdServer {
    url: String,
    cert_location: String,
    key_location: String,
}

impl LxdServer {

    pub fn new(url: &str, cert_location: &str, key_location: &str) -> LxdServer {
        LxdServer {
            url: url.to_string(),
            cert_location: cert_location.to_string(),
            key_location: key_location.to_string()
        }
    }

    fn get_client(&self) -> Client {
        let connector = HttpsConnector::new(
            Openssl::with_cert_and_key(&self.cert_location, &self.key_location).unwrap());
        Client::with_connector(connector)
    }

    fn get(&self, relative_url: &str) -> Response {
        let client = self.get_client();
        let url = match (self.url.clone() + relative_url).into_url() {
            Err(why) => panic!("{:?}", why),
            Ok(url) => url,
        };
        let unsent_response = client.get(url);
        match unsent_response.send() {
            Err(why) => panic!("{:?}", why),
            Ok(response) => response,
        }
    }

    pub fn list_containers(&self) ->  Vec<Container> {
        let mut response = self.get("/1.0/containers?recursion=1");
        let payload = response_to_value(&mut response);
        let container_values = payload.find("metadata").unwrap().as_array().unwrap();
        container_values.iter().map(|v| { Container::from_json(v) }).collect()
    }
}

#[derive(Deserialize)]
pub struct ContainerIP {
    pub address: String,
    pub host_veth: String,
    pub interface: String,
    pub protocol: String,
}

pub struct ContainerStatus {
    pub status: String,
    pub ips: Vec<ContainerIP>,
}

pub struct Container {
    pub name: String,
    pub status: ContainerStatus,
    pub ephemeral: bool,
    pub snapshot_urls: Vec<String>,
}

impl Container {

    pub fn from_json(json: &Value) -> Container {
        let value_for_path = |path: &[&str]| {
            match json.find_path(&path) {
                Some(value) => { value }
                None => panic!("Couldn't find {} in JSON", path.join("."))
            }
        };
        let get_array_from_json = |path: &[&str]| {
            match value_for_path(&path).as_array() {
                Some(array) => { array }
                None => panic!("Couldn't find array at {}", path.join("."))
            }
        };
        let get_boolean_from_json = |path: &[&str]| {
            match value_for_path(&path).as_boolean() {
                Some(boolean) => { boolean }
                None => panic!("Couldn't find boolean at {}", path.join("."))
            }
        };
        let get_string_from_json = |path: &[&str]| {
            match value_for_path(&path).as_string() {
                Some(string) => { string.to_string() }
                None => panic!("Couldn't find string at {}", path.join("."))
            }
        };
        Container {
            ephemeral: get_boolean_from_json(&["state", "ephemeral"]),
            name: get_string_from_json(&["state", "name"]),
            status: ContainerStatus {
                status: get_string_from_json(&["state", "status", "status"]),
                ips: get_array_from_json(&["state", "status", "ips"]).iter().map(
                    |ip_value| -> ContainerIP {
                        match serde_json::from_value(ip_value.clone()) {
                            Ok(container_ip) => container_ip,
                            Err(_) => panic!("Couldn't parse a ContainerIP for {:?}", ip_value)
                        }}).collect(),
            },
            snapshot_urls: match json.find_path(&["snaps"]).unwrap().as_array() {
                Some(array) => {
                    array.iter().map(|x| {x.as_string().unwrap().to_string()}).collect()
                }
                None => {
                    vec![]
                }
            }
        }
    }
}

fn response_to_value(response: &mut Response) -> Value {
    let mut body = String::new();
    response.read_to_string(&mut body).unwrap();
    serde_json::from_str(body.as_str()).unwrap()
}
<commit_msg>Panic with useful message on non-200 GET response.<commit_after>#![feature(convert,custom_derive,plugin)]
#![plugin(serde_macros)]
extern crate hyper;
extern crate serde;
extern crate serde_json;

use std::io::Read;

use hyper::Client;
use hyper::client::IntoUrl;
use hyper::client::Response;
use hyper::net::HttpsConnector;
use hyper::net::Openssl;
use hyper::status::StatusCode;
use serde_json::Value;

pub struct LxdServer {
    url: String,
    cert_location: String,
    key_location: String,
}

impl LxdServer {

    pub fn new(url: &str, cert_location: &str, key_location: &str) -> LxdServer {
        LxdServer {
            url: url.to_string(),
            cert_location: cert_location.to_string(),
            key_location: key_location.to_string()
        }
    }

    fn get_client(&self) -> Client {
        let connector = HttpsConnector::new(
            Openssl::with_cert_and_key(&self.cert_location, &self.key_location).unwrap());
        Client::with_connector(connector)
    }

    fn get(&self, relative_url: &str) -> Response {
        let client = self.get_client();
        let url = match (self.url.clone() + relative_url).into_url() {
            Err(why) => panic!("{:?}", why),
            Ok(url) => url,
        };
        let unsent_response = client.get(url);
        let response = match unsent_response.send() {
            Err(why) => panic!("{:?}", why),
            Ok(response) => response,
        };
        match response.status {
            StatusCode::Ok => response,
            _ => panic!("Response not OK: {}", response.status)
        }
    }

    pub fn list_containers(&self) ->  Vec<Container> {
        let mut response = self.get("/1.0/containers?recursion=1");
        let payload = response_to_value(&mut response);
        let container_values = payload.find("metadata").unwrap().as_array().unwrap();
        container_values.iter().map(|v| { Container::from_json(v) }).collect()
    }
}

#[derive(Deserialize)]
pub struct ContainerIP {
    pub address: String,
    pub host_veth: String,
    pub interface: String,
    pub protocol: String,
}

pub struct ContainerStatus {
    pub status: String,
    pub ips: Vec<ContainerIP>,
}

pub struct Container {
    pub name: String,
    pub status: ContainerStatus,
    pub ephemeral: bool,
    pub snapshot_urls: Vec<String>,
}

impl Container {

    pub fn from_json(json: &Value) -> Container {
        let value_for_path = |path: &[&str]| {
            match json.find_path(&path) {
                Some(value) => { value }
                None => panic!("Couldn't find {} in JSON", path.join("."))
            }
        };
        let get_array_from_json = |path: &[&str]| {
            match value_for_path(&path).as_array() {
                Some(array) => { array }
                None => panic!("Couldn't find array at {}", path.join("."))
            }
        };
        let get_boolean_from_json = |path: &[&str]| {
            match value_for_path(&path).as_boolean() {
                Some(boolean) => { boolean }
                None => panic!("Couldn't find boolean at {}", path.join("."))
            }
        };
        let get_string_from_json = |path: &[&str]| {
            match value_for_path(&path).as_string() {
                Some(string) => { string.to_string() }
                None => panic!("Couldn't find string at {}", path.join("."))
            }
        };
        Container {
            ephemeral: get_boolean_from_json(&["state", "ephemeral"]),
            name: get_string_from_json(&["state", "name"]),
            status: ContainerStatus {
                status: get_string_from_json(&["state", "status", "status"]),
                ips: get_array_from_json(&["state", "status", "ips"]).iter().map(
                    |ip_value| -> ContainerIP {
                        match serde_json::from_value(ip_value.clone()) {
                            Ok(container_ip) => container_ip,
                            Err(_) => panic!("Couldn't parse a ContainerIP for {:?}", ip_value)
                        }}).collect(),
            },
            snapshot_urls: match json.find_path(&["snaps"]).unwrap().as_array() {
                Some(array) => {
                    array.iter().map(|x| {x.as_string().unwrap().to_string()}).collect()
                }
                None => {
                    vec![]
                }
            }
        }
    }
}

fn response_to_value(response: &mut Response) -> Value {
    let mut body = String::new();
    response.read_to_string(&mut body).unwrap();
    serde_json::from_str(body.as_str()).unwrap()
}
<|endoftext|>