<commit_before>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

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

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

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

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here becuase `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

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

    fn work(&self) {
        self.after_start.as_ref().map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        self.before_stop.as_ref().map(|fun| fun());
    }
}

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

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

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// 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.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
                after_start: self.after_start.clone(),
                before_stop: self.before_stop.clone(),
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.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));
            }
            thread_builder.spawn(move || inner.work()).unwrap();
        }

        return pool
    }
}
<commit_msg>futures_cpupool: Add `derive(Debug)` for `CpuFuture`<commit_after>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

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

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

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

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
#[derive(Debug)]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here becuase `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

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

    fn work(&self) {
        self.after_start.as_ref().map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        self.before_stop.as_ref().map(|fun| fun());
    }
}

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

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

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// 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.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
                after_start: self.after_start.clone(),
                before_stop: self.before_stop.clone(),
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.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));
            }
            thread_builder.spawn(move || inner.work()).unwrap();
        }

        return pool
    }
}
<|endoftext|>


<commit_before>#[deriving(Eq)]
pub struct Sym(@str);

#[deriving(Eq)]
pub struct RecordName(@DottedName, @str);

#[deriving(Eq)]
pub struct DottedName(@[Sym]);

#[deriving(Eq)]
pub enum Lit {
    IntegerLiteral(@str),
    FloatingLiteral(@str),
    StringLiteral(@str),
    BytesLiteral(@str)
}

pub enum LetPat {
    // a(b, c, ...)
    FunctionPattern(DottedName, @[Pat]),

    // a.b.c
    DottedPattern(DottedName),

    // ...
    BasicPattern(Pat)
}

pub enum Pat {
    // _
    AnyPattern,

    // ...
    ManyPattern,

    // A(a, b, c, ...)
    RecordPattern(RecordName, @[Pat]),

    // a | b
    DisjunctivePattern(@Pat, @Pat),

    // a & b
    ConjunctivePattern(@Pat, @Pat),

    // [a, b, c, ...]
    ListPattern(@[Pat]),

    // a as a'
    BoundPattern(Sym, @Pat),

    // some kind of literal
    LiteralPattern(Lit),

    // a
    SymbolPattern(Sym),

    // ()
    UnitPattern
}

pub struct Exp {
    exp: BareExp,
    lineno: uint
}

pub enum BareExp {
    // if a then b [ else c ]
    IfThenElseExpression(@Exp, @Exp, option::Option<@Exp>),

    // while a do b
    WhileDoExpression(@Exp, @Exp),

    // match a with case b' -> b case c' -> c
    MatchWithExpression(@Exp, @[(Pat, Exp)]),

    // a OP b
    BinaryExpression(@Exp, @str, @Exp),

    // some kind of literal
    LiteralExpression(Lit),

    // { a; b; c } or { a; b; c; } if boolean field is true
    BlockExpression(@[Stmt], bool),

    // a(b, c, d, ...)
    CallExpression(@Exp, @[Exp]),

    // [a, b, c, d]
    ListExpression(@[Exp]),

    // fn (a) -> b
    LambdaExpression(@[Pat], @Exp),

    // a.b
    RecordAccessExpression(@Exp, Sym),

    // a:b
    RecordFunctionBindExpression(@Exp, Sym),

    // a
    SymbolExpression(Sym),

    // a = b
    AssignmentExpression(@Exp, @Exp),

    // ()
    UnitExpression
}

pub struct Stmt {
    stmt: BareStmt,
    lineno: uint
}

pub enum BareStmt {
    // let a = b
    LetBindingStatement(LetPat, Exp),

    // ...
    ExpressionStatement(Exp)
}

pub fn mk_expression_statement(exp: Exp) -> Stmt {
    Stmt {
        stmt: ExpressionStatement(exp),
        lineno: exp.lineno
    }
}

pub struct ImportDeclaration {
    module: DottedName,
    qualified: bool,
    lineno: uint
}

pub struct RecordDeclaration {
    name: RecordName,
    slots: @[Sym],
    lineno: uint
}

pub struct Program {
    imports: @[ImportDeclaration],
    records: @[RecordDeclaration],
    body: @[Stmt]
}
<commit_msg>implement pattern specificities<commit_after>#[deriving(Eq)]
pub struct Sym(@str);

#[deriving(Eq)]
pub struct RecordName(@DottedName, @str);

#[deriving(Eq)]
pub struct DottedName(@[Sym]);

#[deriving(Eq)]
pub enum Lit {
    IntegerLiteral(@str),
    FloatingLiteral(@str),
    StringLiteral(@str),
    BytesLiteral(@str)
}

pub enum LetPat {
    // a(b, c, ...)
    FunctionPattern(DottedName, @[Pat]),

    // a.b.c
    DottedPattern(DottedName),

    // ...
    BasicPattern(Pat)
}

pub enum Pat {
    // _
    AnyPattern,

    // ...
    ManyPattern,

    // A(a, b, c, ...)
    RecordPattern(RecordName, @[Pat]),

    // a | b
    DisjunctivePattern(@Pat, @Pat),

    // a & b
    ConjunctivePattern(@Pat, @Pat),

    // [a, b, c, ...]
    ListPattern(@[Pat]),

    // a as a'
    BoundPattern(Sym, @Pat),

    // some kind of literal
    LiteralPattern(Lit),

    // a
    SymbolPattern(Sym),

    // ()
    UnitPattern
}

pub impl Pat {
    fn specificity(&self) -> uint {
        match *self {
            AnyPattern => 0,
            ManyPattern => 0,
            RecordPattern(_, pats) => pats.map(|pat| pat.specificity()).foldr(0, |x, acc| x + acc),
            DisjunctivePattern(l, r) => uint::min(l.specificity(), r.specificity()),
            ConjunctivePattern(l, r) => uint::max(l.specificity(), r.specificity()),
            ListPattern(pats) => pats.map(|pat| pat.specificity()).foldr(0, |x, acc| x + acc),
            BoundPattern(_, pat) => pat.specificity(),
            LiteralPattern(_) => 1,
            SymbolPattern(_) => 1,
            UnitPattern => 1
        }
    }
}

pub struct Exp {
    exp: BareExp,
    lineno: uint
}

pub enum BareExp {
    // if a then b [ else c ]
    IfThenElseExpression(@Exp, @Exp, option::Option<@Exp>),

    // while a do b
    WhileDoExpression(@Exp, @Exp),

    // match a with case b' -> b case c' -> c
    MatchWithExpression(@Exp, @[(Pat, Exp)]),

    // a OP b
    BinaryExpression(@Exp, @str, @Exp),

    // some kind of literal
    LiteralExpression(Lit),

    // { a; b; c } or { a; b; c; } if boolean field is true
    BlockExpression(@[Stmt], bool),

    // a(b, c, d, ...)
    CallExpression(@Exp, @[Exp]),

    // [a, b, c, d]
    ListExpression(@[Exp]),

    // fn (a) -> b
    LambdaExpression(@[Pat], @Exp),

    // a.b
    RecordAccessExpression(@Exp, Sym),

    // a:b
    RecordFunctionBindExpression(@Exp, Sym),

    // a
    SymbolExpression(Sym),

    // a = b
    AssignmentExpression(@Exp, @Exp),

    // ()
    UnitExpression
}

pub struct Stmt {
    stmt: BareStmt,
    lineno: uint
}

pub enum BareStmt {
    // let a = b
    LetBindingStatement(LetPat, Exp),

    // ...
    ExpressionStatement(Exp)
}

pub fn mk_expression_statement(exp: Exp) -> Stmt {
    Stmt {
        stmt: ExpressionStatement(exp),
        lineno: exp.lineno
    }
}

pub struct ImportDeclaration {
    module: DottedName,
    qualified: bool,
    lineno: uint
}

pub struct RecordDeclaration {
    name: RecordName,
    slots: @[Sym],
    lineno: uint
}

pub struct Program {
    imports: @[ImportDeclaration],
    records: @[RecordDeclaration],
    body: @[Stmt]
}
<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::fs;
use std::io;
use std::path::Path;
use std::process;
use std::result;

use errno_result;

use libc::{c_long, pid_t, syscall, CLONE_NEWPID, CLONE_NEWUSER, SIGCHLD};

use syscall_defines::linux::LinuxSyscall::SYS_clone;

/// Controls what namespace `clone_process` will have. See NAMESPACES(7).
#[repr(u32)]
pub enum CloneNamespace {
    /// The new process will inherit the namespace from the old process.
    Inherit = 0,
    /// The new process with be in a new user and PID namespace.
    NewUserPid = CLONE_NEWUSER as u32 | CLONE_NEWPID as u32,
}

#[derive(Debug)]
pub enum CloneError {
    /// There was an error trying to iterate this process's threads.
    IterateTasks(io::Error),
    /// There are multiple threads running. The `usize` indicates how many threads.
    Multithreaded(usize),
    /// There was an error while cloning.
    Sys(::Error),
}

unsafe fn do_clone(flags: i32) -> ::Result<pid_t> {
    // Forking is unsafe, this function must be unsafe as there is no way to guarantee safety
    // without more context about the state of the program.
    let pid = syscall(SYS_clone as c_long, flags | SIGCHLD as i32, 0);
    if pid < 0 {
        errno_result()
    } else {
        Ok(pid as pid_t)
    }
}

fn count_dir_entries<P: AsRef<Path>>(path: P) -> io::Result<usize> {
    Ok(fs::read_dir(path)?.count())
}

/// Clones this process and calls a closure in the new process.
///
/// After `post_clone_cb` returns or panics, the new process exits. Similar to how a `fork` syscall
/// works, the new process is the same as the current process with the exception of the namespace
/// controlled with the `ns` argument.
///
/// # Arguments
/// * `ns` - What namespace the new process will have (see NAMESPACES(7)).
/// * `post_clone_cb` - Callback to run in the new process
pub fn clone_process<F>(ns: CloneNamespace, post_clone_cb: F) -> result::Result<pid_t, CloneError>
where
    F: FnOnce(),
{
    match count_dir_entries("/proc/self/task") {
        Ok(1) => {}
        Ok(thread_count) => {
            // Test cfg gets a free pass on this because tests generally have multiple independent
            // test threads going.
            let _ = thread_count;
            #[cfg(not(test))]
            return Err(CloneError::Multithreaded(thread_count));
        }
        Err(e) => return Err(CloneError::IterateTasks(e)),
    }
    // Forking is considered unsafe in mutlithreaded programs, but we just checked for other threads
    // in this process. We also only allow valid flags from CloneNamespace and check the return
    // result for errors. We also never let the cloned process return from this function.
    let ret = unsafe { do_clone(ns as i32) }.map_err(CloneError::Sys)?;
    if ret == 0 {
        struct ExitGuard;
        impl Drop for ExitGuard {
            fn drop(&mut self) {
                process::exit(101);
            }
        }
        // Prevents a panic in post_clone_cb from bypassing the process::exit.
        #[allow(unused_variables)]
        let exit_guard = ExitGuard {};
        post_clone_cb();
        // ! Never returns
        process::exit(0);
    }

    Ok(ret)
}

#[cfg(test)]
mod tests {
    use super::*;
    use libc;
    use {getpid, EventFd};

    fn wait_process() -> libc::c_int {
        let mut status: libc::c_int = 0;
        unsafe {
            libc::wait(&mut status as *mut libc::c_int);
            libc::WEXITSTATUS(status)
        }
    }

    #[test]
    fn pid_diff() {
        let evt_fd = EventFd::new().expect("failed to create EventFd");
        let evt_fd_fork = evt_fd.try_clone().expect("failed to clone EventFd");
        let pid = getpid();
        clone_process(CloneNamespace::Inherit, || {
            // checks that this is a genuine fork with a new PID
            if pid != getpid() {
                evt_fd_fork.write(1).unwrap()
            } else {
                evt_fd_fork.write(2).unwrap()
            }
        })
        .expect("failed to clone");
        assert_eq!(evt_fd.read(), Ok(1));
    }

    #[test]
    fn panic_safe() {
        let pid = getpid();
        assert_ne!(pid, 0);

        clone_process(CloneNamespace::Inherit, || {
            assert!(false);
        })
        .expect("failed to clone");

        // This should never happen;
        if pid != getpid() {
            process::exit(2);
        }

        let status = wait_process();
        assert!(status == 101 || status == 0);
    }
}
<commit_msg>sys_util: wait on specific pid in panic_safe test<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::fs;
use std::io;
use std::path::Path;
use std::process;
use std::result;

use errno_result;

use libc::{c_long, pid_t, syscall, CLONE_NEWPID, CLONE_NEWUSER, SIGCHLD};

use syscall_defines::linux::LinuxSyscall::SYS_clone;

/// Controls what namespace `clone_process` will have. See NAMESPACES(7).
#[repr(u32)]
pub enum CloneNamespace {
    /// The new process will inherit the namespace from the old process.
    Inherit = 0,
    /// The new process with be in a new user and PID namespace.
    NewUserPid = CLONE_NEWUSER as u32 | CLONE_NEWPID as u32,
}

#[derive(Debug)]
pub enum CloneError {
    /// There was an error trying to iterate this process's threads.
    IterateTasks(io::Error),
    /// There are multiple threads running. The `usize` indicates how many threads.
    Multithreaded(usize),
    /// There was an error while cloning.
    Sys(::Error),
}

unsafe fn do_clone(flags: i32) -> ::Result<pid_t> {
    // Forking is unsafe, this function must be unsafe as there is no way to guarantee safety
    // without more context about the state of the program.
    let pid = syscall(SYS_clone as c_long, flags | SIGCHLD as i32, 0);
    if pid < 0 {
        errno_result()
    } else {
        Ok(pid as pid_t)
    }
}

fn count_dir_entries<P: AsRef<Path>>(path: P) -> io::Result<usize> {
    Ok(fs::read_dir(path)?.count())
}

/// Clones this process and calls a closure in the new process.
///
/// After `post_clone_cb` returns or panics, the new process exits. Similar to how a `fork` syscall
/// works, the new process is the same as the current process with the exception of the namespace
/// controlled with the `ns` argument.
///
/// # Arguments
/// * `ns` - What namespace the new process will have (see NAMESPACES(7)).
/// * `post_clone_cb` - Callback to run in the new process
pub fn clone_process<F>(ns: CloneNamespace, post_clone_cb: F) -> result::Result<pid_t, CloneError>
where
    F: FnOnce(),
{
    match count_dir_entries("/proc/self/task") {
        Ok(1) => {}
        Ok(thread_count) => {
            // Test cfg gets a free pass on this because tests generally have multiple independent
            // test threads going.
            let _ = thread_count;
            #[cfg(not(test))]
            return Err(CloneError::Multithreaded(thread_count));
        }
        Err(e) => return Err(CloneError::IterateTasks(e)),
    }
    // Forking is considered unsafe in mutlithreaded programs, but we just checked for other threads
    // in this process. We also only allow valid flags from CloneNamespace and check the return
    // result for errors. We also never let the cloned process return from this function.
    let ret = unsafe { do_clone(ns as i32) }.map_err(CloneError::Sys)?;
    if ret == 0 {
        struct ExitGuard;
        impl Drop for ExitGuard {
            fn drop(&mut self) {
                process::exit(101);
            }
        }
        // Prevents a panic in post_clone_cb from bypassing the process::exit.
        #[allow(unused_variables)]
        let exit_guard = ExitGuard {};
        post_clone_cb();
        // ! Never returns
        process::exit(0);
    }

    Ok(ret)
}

#[cfg(test)]
mod tests {
    use super::*;
    use libc;
    use {getpid, EventFd};

    fn wait_process(pid: libc::pid_t) -> ::Result<libc::c_int> {
        let mut status: libc::c_int = 0;
        unsafe {
            if libc::waitpid(pid, &mut status as *mut libc::c_int, 0) < 0 {
                errno_result()
            } else {
                Ok(libc::WEXITSTATUS(status))
            }
        }
    }

    #[test]
    fn pid_diff() {
        let evt_fd = EventFd::new().expect("failed to create EventFd");
        let evt_fd_fork = evt_fd.try_clone().expect("failed to clone EventFd");
        let pid = getpid();
        clone_process(CloneNamespace::Inherit, || {
            // checks that this is a genuine fork with a new PID
            if pid != getpid() {
                evt_fd_fork.write(1).unwrap()
            } else {
                evt_fd_fork.write(2).unwrap()
            }
        })
        .expect("failed to clone");
        assert_eq!(evt_fd.read(), Ok(1));
    }

    #[test]
    fn panic_safe() {
        let pid = getpid();
        assert_ne!(pid, 0);

        let clone_pid = clone_process(CloneNamespace::Inherit, || {
            assert!(false);
        })
        .expect("failed to clone");

        // This should never happen;
        if pid != getpid() {
            process::exit(2);
        }

        let status = wait_process(clone_pid).expect("wait_process failed");
        assert!(status == 101 || status == 0);
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::fmt::{self, Formatter};
use std::hash;
use std::hash::Hash;
use std::path::Path;
use std::sync::Arc;

use semver;
use serde::de;
use serde::ser;

use core::interning::InternedString;
use core::source::SourceId;
use util::{CargoResult, ToSemver};

/// Identifier for a specific version of a package in a specific source.
#[derive(Clone)]
pub struct PackageId {
    inner: Arc<PackageIdInner>,
}

#[derive(PartialEq, PartialOrd, Eq, Ord)]
struct PackageIdInner {
    name: InternedString,
    version: semver::Version,
    source_id: SourceId,
}

impl ser::Serialize for PackageId {
    fn serialize<S>(&self, s: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        s.collect_str(&format_args!(
            "{} {} ({})",
            self.inner.name,
            self.inner.version,
            self.inner.source_id.to_url()
        ))
    }
}

impl<'de> de::Deserialize<'de> for PackageId {
    fn deserialize<D>(d: D) -> Result<PackageId, D::Error>
    where
        D: de::Deserializer<'de>,
    {
        let string = String::deserialize(d)?;
        let mut s = string.splitn(3, ' ');
        let name = s.next().unwrap();
        let version = match s.next() {
            Some(s) => s,
            None => return Err(de::Error::custom("invalid serialized PackageId")),
        };
        let version = semver::Version::parse(version).map_err(de::Error::custom)?;
        let url = match s.next() {
            Some(s) => s,
            None => return Err(de::Error::custom("invalid serialized PackageId")),
        };
        let url = if url.starts_with('(') && url.ends_with(')') {
            &url[1..url.len() - 1]
        } else {
            return Err(de::Error::custom("invalid serialized PackageId"));
        };
        let source_id = SourceId::from_url(url).map_err(de::Error::custom)?;

        Ok(PackageId {
            inner: Arc::new(PackageIdInner {
                name: InternedString::new(name),
                version,
                source_id,
            }),
        })
    }
}

impl Hash for PackageId {
    fn hash<S: hash::Hasher>(&self, state: &mut S) {
        self.inner.name.hash(state);
        self.inner.version.hash(state);
        self.inner.source_id.hash(state);
    }
}

impl PartialEq for PackageId {
    fn eq(&self, other: &PackageId) -> bool {
        (*self.inner).eq(&*other.inner)
    }
}
impl PartialOrd for PackageId {
    fn partial_cmp(&self, other: &PackageId) -> Option<Ordering> {
        (*self.inner).partial_cmp(&*other.inner)
    }
}
impl Eq for PackageId {}
impl Ord for PackageId {
    fn cmp(&self, other: &PackageId) -> Ordering {
        (*self.inner).cmp(&*other.inner)
    }
}

impl PackageId {
    pub fn new<T: ToSemver>(name: &str, version: T, sid: SourceId) -> CargoResult<PackageId> {
        let v = version.to_semver()?;
        Ok(PackageId {
            inner: Arc::new(PackageIdInner {
                name: InternedString::new(name),
                version: v,
                source_id: sid,
            }),
        })
    }

    pub fn name(&self) -> InternedString {
        self.inner.name
    }
    pub fn version(&self) -> &semver::Version {
        &self.inner.version
    }
    pub fn source_id(&self) -> SourceId {
        self.inner.source_id
    }

    pub fn with_precise(&self, precise: Option<String>) -> PackageId {
        PackageId {
            inner: Arc::new(PackageIdInner {
                name: self.inner.name,
                version: self.inner.version.clone(),
                source_id: self.inner.source_id.with_precise(precise),
            }),
        }
    }

    pub fn with_source_id(&self, source: SourceId) -> PackageId {
        PackageId {
            inner: Arc::new(PackageIdInner {
                name: self.inner.name,
                version: self.inner.version.clone(),
                source_id: source,
            }),
        }
    }

    pub fn stable_hash<'a>(&'a self, workspace: &'a Path) -> PackageIdStableHash<'a> {
        PackageIdStableHash(self, workspace)
    }
}

pub struct PackageIdStableHash<'a>(&'a PackageId, &'a Path);

impl<'a> Hash for PackageIdStableHash<'a> {
    fn hash<S: hash::Hasher>(&self, state: &mut S) {
        self.0.inner.name.hash(state);
        self.0.inner.version.hash(state);
        self.0.inner.source_id.stable_hash(self.1, state);
    }
}

impl fmt::Display for PackageId {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{} v{}", self.inner.name, self.inner.version)?;

        if !self.inner.source_id.is_default_registry() {
            write!(f, " ({})", self.inner.source_id)?;
        }

        Ok(())
    }
}

impl fmt::Debug for PackageId {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("PackageId")
            .field("name", &self.inner.name)
            .field("version", &self.inner.version.to_string())
            .field("source", &self.inner.source_id.to_string())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::PackageId;
    use core::source::SourceId;
    use sources::CRATES_IO_INDEX;
    use util::ToUrl;

    #[test]
    fn invalid_version_handled_nicely() {
        let loc = CRATES_IO_INDEX.to_url().unwrap();
        let repo = SourceId::for_registry(&loc).unwrap();

        assert!(PackageId::new("foo", "1.0", repo).is_err());
        assert!(PackageId::new("foo", "1", repo).is_err());
        assert!(PackageId::new("foo", "bar", repo).is_err());
        assert!(PackageId::new("foo", "", repo).is_err());
    }
}
<commit_msg>Test PackageId Debug & Display<commit_after>use std::cmp::Ordering;
use std::fmt::{self, Formatter};
use std::hash;
use std::hash::Hash;
use std::path::Path;
use std::sync::Arc;

use semver;
use serde::de;
use serde::ser;

use core::interning::InternedString;
use core::source::SourceId;
use util::{CargoResult, ToSemver};

/// Identifier for a specific version of a package in a specific source.
#[derive(Clone)]
pub struct PackageId {
    inner: Arc<PackageIdInner>,
}

#[derive(PartialEq, PartialOrd, Eq, Ord)]
struct PackageIdInner {
    name: InternedString,
    version: semver::Version,
    source_id: SourceId,
}

impl ser::Serialize for PackageId {
    fn serialize<S>(&self, s: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        s.collect_str(&format_args!(
            "{} {} ({})",
            self.inner.name,
            self.inner.version,
            self.inner.source_id.to_url()
        ))
    }
}

impl<'de> de::Deserialize<'de> for PackageId {
    fn deserialize<D>(d: D) -> Result<PackageId, D::Error>
    where
        D: de::Deserializer<'de>,
    {
        let string = String::deserialize(d)?;
        let mut s = string.splitn(3, ' ');
        let name = s.next().unwrap();
        let version = match s.next() {
            Some(s) => s,
            None => return Err(de::Error::custom("invalid serialized PackageId")),
        };
        let version = semver::Version::parse(version).map_err(de::Error::custom)?;
        let url = match s.next() {
            Some(s) => s,
            None => return Err(de::Error::custom("invalid serialized PackageId")),
        };
        let url = if url.starts_with('(') && url.ends_with(')') {
            &url[1..url.len() - 1]
        } else {
            return Err(de::Error::custom("invalid serialized PackageId"));
        };
        let source_id = SourceId::from_url(url).map_err(de::Error::custom)?;

        Ok(PackageId {
            inner: Arc::new(PackageIdInner {
                name: InternedString::new(name),
                version,
                source_id,
            }),
        })
    }
}

impl Hash for PackageId {
    fn hash<S: hash::Hasher>(&self, state: &mut S) {
        self.inner.name.hash(state);
        self.inner.version.hash(state);
        self.inner.source_id.hash(state);
    }
}

impl PartialEq for PackageId {
    fn eq(&self, other: &PackageId) -> bool {
        (*self.inner).eq(&*other.inner)
    }
}
impl PartialOrd for PackageId {
    fn partial_cmp(&self, other: &PackageId) -> Option<Ordering> {
        (*self.inner).partial_cmp(&*other.inner)
    }
}
impl Eq for PackageId {}
impl Ord for PackageId {
    fn cmp(&self, other: &PackageId) -> Ordering {
        (*self.inner).cmp(&*other.inner)
    }
}

impl PackageId {
    pub fn new<T: ToSemver>(name: &str, version: T, sid: SourceId) -> CargoResult<PackageId> {
        let v = version.to_semver()?;
        Ok(PackageId {
            inner: Arc::new(PackageIdInner {
                name: InternedString::new(name),
                version: v,
                source_id: sid,
            }),
        })
    }

    pub fn name(&self) -> InternedString {
        self.inner.name
    }
    pub fn version(&self) -> &semver::Version {
        &self.inner.version
    }
    pub fn source_id(&self) -> SourceId {
        self.inner.source_id
    }

    pub fn with_precise(&self, precise: Option<String>) -> PackageId {
        PackageId {
            inner: Arc::new(PackageIdInner {
                name: self.inner.name,
                version: self.inner.version.clone(),
                source_id: self.inner.source_id.with_precise(precise),
            }),
        }
    }

    pub fn with_source_id(&self, source: SourceId) -> PackageId {
        PackageId {
            inner: Arc::new(PackageIdInner {
                name: self.inner.name,
                version: self.inner.version.clone(),
                source_id: source,
            }),
        }
    }

    pub fn stable_hash<'a>(&'a self, workspace: &'a Path) -> PackageIdStableHash<'a> {
        PackageIdStableHash(self, workspace)
    }
}

pub struct PackageIdStableHash<'a>(&'a PackageId, &'a Path);

impl<'a> Hash for PackageIdStableHash<'a> {
    fn hash<S: hash::Hasher>(&self, state: &mut S) {
        self.0.inner.name.hash(state);
        self.0.inner.version.hash(state);
        self.0.inner.source_id.stable_hash(self.1, state);
    }
}

impl fmt::Display for PackageId {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{} v{}", self.inner.name, self.inner.version)?;

        if !self.inner.source_id.is_default_registry() {
            write!(f, " ({})", self.inner.source_id)?;
        }

        Ok(())
    }
}

impl fmt::Debug for PackageId {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("PackageId")
            .field("name", &self.inner.name)
            .field("version", &self.inner.version.to_string())
            .field("source", &self.inner.source_id.to_string())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::PackageId;
    use core::source::SourceId;
    use sources::CRATES_IO_INDEX;
    use util::ToUrl;

    #[test]
    fn invalid_version_handled_nicely() {
        let loc = CRATES_IO_INDEX.to_url().unwrap();
        let repo = SourceId::for_registry(&loc).unwrap();

        assert!(PackageId::new("foo", "1.0", repo).is_err());
        assert!(PackageId::new("foo", "1", repo).is_err());
        assert!(PackageId::new("foo", "bar", repo).is_err());
        assert!(PackageId::new("foo", "", repo).is_err());
    }

    #[test]
    fn debug() {
        let loc = CRATES_IO_INDEX.to_url().unwrap();
        let pkg_id = PackageId::new("foo", "1.0.0", SourceId::for_registry(&loc).unwrap()).unwrap();
        assert_eq!(r#"PackageId { name: "foo", version: "1.0.0", source: "registry `https://github.com/rust-lang/crates.io-index`" }"#, format!("{:?}", pkg_id));

        let pretty = r#"
PackageId {
    name: "foo",
    version: "1.0.0",
    source: "registry `https://github.com/rust-lang/crates.io-index`"
}
"#.trim();
        assert_eq!(pretty, format!("{:#?}", pkg_id));
    }

    #[test]
    fn display() {
        let loc = CRATES_IO_INDEX.to_url().unwrap();
        let pkg_id = PackageId::new("foo", "1.0.0", SourceId::for_registry(&loc).unwrap()).unwrap();
        assert_eq!("foo v1.0.0", pkg_id.to_string());
    }
}
<|endoftext|>


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

use crate::core::package::PackageSet;
use crate::core::{Dependency, Package, PackageId, Summary};
use crate::util::{CargoResult, Config};

mod source_id;

pub use self::source_id::{GitReference, SourceId};

/// Something that finds and downloads remote packages based on names and versions.
pub trait Source {
    /// Returns the `SourceId` corresponding to this source.
    fn source_id(&self) -> SourceId;

    /// Returns the replaced `SourceId` corresponding to this source.
    fn replaced_source_id(&self) -> SourceId {
        self.source_id()
    }

    /// Returns whether or not this source will return summaries with
    /// checksums listed.
    fn supports_checksums(&self) -> bool;

    /// Returns whether or not this source will return summaries with
    /// the `precise` field in the source id listed.
    fn requires_precise(&self) -> bool;

    /// Attempts to find the packages that match a dependency request.
    fn query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()>;

    /// Attempts to find the packages that are close to a dependency request.
    /// Each source gets to define what `close` means for it.
    /// Path/Git sources may return all dependencies that are at that URI,
    /// whereas an `Index` source may return dependencies that have the same canonicalization.
    fn fuzzy_query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()>;

    fn query_vec(&mut self, dep: &Dependency) -> CargoResult<Vec<Summary>> {
        let mut ret = Vec::new();
        self.query(dep, &mut |s| ret.push(s))?;
        Ok(ret)
    }

    /// Performs any network operations required to get the entire list of all names,
    /// versions and dependencies of packages managed by the `Source`.
    fn update(&mut self) -> CargoResult<()>;

    /// Fetches the full package for each name and version specified.
    fn download(&mut self, package: PackageId) -> CargoResult<MaybePackage>;

    fn download_now(self: Box<Self>, package: PackageId, config: &Config) -> CargoResult<Package>
    where
        Self: std::marker::Sized,
    {
        let mut sources = SourceMap::new();
        sources.insert(self);
        let pkg_set = PackageSet::new(&[package], sources, config)?;
        let pkg = pkg_set.get_one(package)?;
        Ok(Package::clone(pkg))
    }

    fn finish_download(&mut self, package: PackageId, contents: Vec<u8>) -> CargoResult<Package>;

    /// Generates a unique string which represents the fingerprint of the
    /// current state of the source.
    ///
    /// This fingerprint is used to determine the "fresheness" of the source
    /// later on. It must be guaranteed that the fingerprint of a source is
    /// constant if and only if the output product will remain constant.
    ///
    /// The `pkg` argument is the package which this fingerprint should only be
    /// interested in for when this source may contain multiple packages.
    fn fingerprint(&self, pkg: &Package) -> CargoResult<String>;

    /// If this source supports it, verifies the source of the package
    /// specified.
    ///
    /// Note that the source may also have performed other checksum-based
    /// verification during the `download` step, but this is intended to be run
    /// just before a crate is compiled so it may perform more expensive checks
    /// which may not be cacheable.
    fn verify(&self, _pkg: PackageId) -> CargoResult<()> {
        Ok(())
    }

    /// Describes this source in a human readable fashion, used for display in
    /// resolver error messages currently.
    fn describe(&self) -> String;

    /// Returns whether a source is being replaced by another here.
    fn is_replaced(&self) -> bool {
        false
    }

    /// Add a number of crates that should be whitelisted for showing up during
    /// queries, even if they are yanked. Currently only applies to registry
    /// sources.
    fn add_to_yanked_whitelist(&mut self, pkgs: &[PackageId]);

    /// Query if a package is yanked. Only registry sources can mark packages
    /// as yanked. This ignores the yanked whitelist.
    fn is_yanked(&mut self, _pkg: PackageId) -> CargoResult<bool>;
}

pub enum MaybePackage {
    Ready(Package),
    Download { url: String, descriptor: String },
}

impl<'a, T: Source + ?Sized + 'a> Source for Box<T> {
    /// Forwards to `Source::source_id`.
    fn source_id(&self) -> SourceId {
        (**self).source_id()
    }

    /// Forwards to `Source::replaced_source_id`.
    fn replaced_source_id(&self) -> SourceId {
        (**self).replaced_source_id()
    }

    /// Forwards to `Source::supports_checksums`.
    fn supports_checksums(&self) -> bool {
        (**self).supports_checksums()
    }

    /// Forwards to `Source::requires_precise`.
    fn requires_precise(&self) -> bool {
        (**self).requires_precise()
    }

    /// Forwards to `Source::query`.
    fn query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).query(dep, f)
    }

    /// Forwards to `Source::query`.
    fn fuzzy_query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).fuzzy_query(dep, f)
    }

    /// Forwards to `Source::update`.
    fn update(&mut self) -> CargoResult<()> {
        (**self).update()
    }

    /// Forwards to `Source::download`.
    fn download(&mut self, id: PackageId) -> CargoResult<MaybePackage> {
        (**self).download(id)
    }

    fn finish_download(&mut self, id: PackageId, data: Vec<u8>) -> CargoResult<Package> {
        (**self).finish_download(id, data)
    }

    /// Forwards to `Source::fingerprint`.
    fn fingerprint(&self, pkg: &Package) -> CargoResult<String> {
        (**self).fingerprint(pkg)
    }

    /// Forwards to `Source::verify`.
    fn verify(&self, pkg: PackageId) -> CargoResult<()> {
        (**self).verify(pkg)
    }

    fn describe(&self) -> String {
        (**self).describe()
    }

    fn is_replaced(&self) -> bool {
        (**self).is_replaced()
    }

    fn add_to_yanked_whitelist(&mut self, pkgs: &[PackageId]) {
        (**self).add_to_yanked_whitelist(pkgs);
    }

    fn is_yanked(&mut self, pkg: PackageId) -> CargoResult<bool> {
        (**self).is_yanked(pkg)
    }
}

impl<'a, T: Source + ?Sized + 'a> Source for &'a mut T {
    fn source_id(&self) -> SourceId {
        (**self).source_id()
    }

    fn replaced_source_id(&self) -> SourceId {
        (**self).replaced_source_id()
    }

    fn supports_checksums(&self) -> bool {
        (**self).supports_checksums()
    }

    fn requires_precise(&self) -> bool {
        (**self).requires_precise()
    }

    fn query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).query(dep, f)
    }

    fn fuzzy_query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).fuzzy_query(dep, f)
    }

    fn update(&mut self) -> CargoResult<()> {
        (**self).update()
    }

    fn download(&mut self, id: PackageId) -> CargoResult<MaybePackage> {
        (**self).download(id)
    }

    fn finish_download(&mut self, id: PackageId, data: Vec<u8>) -> CargoResult<Package> {
        (**self).finish_download(id, data)
    }

    fn fingerprint(&self, pkg: &Package) -> CargoResult<String> {
        (**self).fingerprint(pkg)
    }

    fn verify(&self, pkg: PackageId) -> CargoResult<()> {
        (**self).verify(pkg)
    }

    fn describe(&self) -> String {
        (**self).describe()
    }

    fn is_replaced(&self) -> bool {
        (**self).is_replaced()
    }

    fn add_to_yanked_whitelist(&mut self, pkgs: &[PackageId]) {
        (**self).add_to_yanked_whitelist(pkgs);
    }

    fn is_yanked(&mut self, pkg: PackageId) -> CargoResult<bool> {
        (**self).is_yanked(pkg)
    }
}

/// A `HashMap` of `SourceId` -> `Box<Source>`.
#[derive(Default)]
pub struct SourceMap<'src> {
    map: HashMap<SourceId, Box<dyn Source + 'src>>,
}

// `impl Debug` on source requires specialization, if even desirable at all.
impl<'src> fmt::Debug for SourceMap<'src> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SourceMap ")?;
        f.debug_set().entries(self.map.keys()).finish()
    }
}

impl<'src> SourceMap<'src> {
    /// Creates an empty map.
    pub fn new() -> SourceMap<'src> {
        SourceMap {
            map: HashMap::new(),
        }
    }

    /// Like `HashMap::contains_key`.
    pub fn contains(&self, id: SourceId) -> bool {
        self.map.contains_key(&id)
    }

    /// Like `HashMap::get`.
    pub fn get(&self, id: SourceId) -> Option<&(dyn Source + 'src)> {
        self.map.get(&id).map(|s| s.as_ref())
    }

    /// Like `HashMap::get_mut`.
    pub fn get_mut(&mut self, id: SourceId) -> Option<&mut (dyn Source + 'src)> {
        self.map.get_mut(&id).map(|s| s.as_mut())
    }

    /// Like `HashMap::get`, but first calculates the `SourceId` from a `PackageId`.
    pub fn get_by_package_id(&self, pkg_id: PackageId) -> Option<&(dyn Source + 'src)> {
        self.get(pkg_id.source_id())
    }

    /// Like `HashMap::insert`, but derives the `SourceId` key from the `Source`.
    pub fn insert(&mut self, source: Box<dyn Source + 'src>) {
        let id = source.source_id();
        self.map.insert(id, source);
    }

    /// Like `HashMap::is_empty`.
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Like `HashMap::len`.
    pub fn len(&self) -> usize {
        self.map.len()
    }

    /// Like `HashMap::values`.
    pub fn sources<'a>(&'a self) -> impl Iterator<Item = &'a Box<dyn Source + 'src>> {
        self.map.values()
    }

    /// Like `HashMap::iter_mut`.
    pub fn sources_mut<'a>(
        &'a mut self,
    ) -> impl Iterator<Item = (&'a SourceId, &'a mut (dyn Source + 'src))> {
        self.map.iter_mut().map(|(a, b)| (a, &mut **b))
    }

    /// Merge the given map into self.
    pub fn add_source_map(&mut self, other: SourceMap<'src>) {
        for (key, value) in other.map {
            self.map.entry(key).or_insert(value);
        }
    }
}
<commit_msg>Remove some dead code in SourceMap<commit_after>use std::collections::hash_map::HashMap;
use std::fmt;

use crate::core::package::PackageSet;
use crate::core::{Dependency, Package, PackageId, Summary};
use crate::util::{CargoResult, Config};

mod source_id;

pub use self::source_id::{GitReference, SourceId};

/// Something that finds and downloads remote packages based on names and versions.
pub trait Source {
    /// Returns the `SourceId` corresponding to this source.
    fn source_id(&self) -> SourceId;

    /// Returns the replaced `SourceId` corresponding to this source.
    fn replaced_source_id(&self) -> SourceId {
        self.source_id()
    }

    /// Returns whether or not this source will return summaries with
    /// checksums listed.
    fn supports_checksums(&self) -> bool;

    /// Returns whether or not this source will return summaries with
    /// the `precise` field in the source id listed.
    fn requires_precise(&self) -> bool;

    /// Attempts to find the packages that match a dependency request.
    fn query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()>;

    /// Attempts to find the packages that are close to a dependency request.
    /// Each source gets to define what `close` means for it.
    /// Path/Git sources may return all dependencies that are at that URI,
    /// whereas an `Index` source may return dependencies that have the same canonicalization.
    fn fuzzy_query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()>;

    fn query_vec(&mut self, dep: &Dependency) -> CargoResult<Vec<Summary>> {
        let mut ret = Vec::new();
        self.query(dep, &mut |s| ret.push(s))?;
        Ok(ret)
    }

    /// Performs any network operations required to get the entire list of all names,
    /// versions and dependencies of packages managed by the `Source`.
    fn update(&mut self) -> CargoResult<()>;

    /// Fetches the full package for each name and version specified.
    fn download(&mut self, package: PackageId) -> CargoResult<MaybePackage>;

    fn download_now(self: Box<Self>, package: PackageId, config: &Config) -> CargoResult<Package>
    where
        Self: std::marker::Sized,
    {
        let mut sources = SourceMap::new();
        sources.insert(self);
        let pkg_set = PackageSet::new(&[package], sources, config)?;
        let pkg = pkg_set.get_one(package)?;
        Ok(Package::clone(pkg))
    }

    fn finish_download(&mut self, package: PackageId, contents: Vec<u8>) -> CargoResult<Package>;

    /// Generates a unique string which represents the fingerprint of the
    /// current state of the source.
    ///
    /// This fingerprint is used to determine the "fresheness" of the source
    /// later on. It must be guaranteed that the fingerprint of a source is
    /// constant if and only if the output product will remain constant.
    ///
    /// The `pkg` argument is the package which this fingerprint should only be
    /// interested in for when this source may contain multiple packages.
    fn fingerprint(&self, pkg: &Package) -> CargoResult<String>;

    /// If this source supports it, verifies the source of the package
    /// specified.
    ///
    /// Note that the source may also have performed other checksum-based
    /// verification during the `download` step, but this is intended to be run
    /// just before a crate is compiled so it may perform more expensive checks
    /// which may not be cacheable.
    fn verify(&self, _pkg: PackageId) -> CargoResult<()> {
        Ok(())
    }

    /// Describes this source in a human readable fashion, used for display in
    /// resolver error messages currently.
    fn describe(&self) -> String;

    /// Returns whether a source is being replaced by another here.
    fn is_replaced(&self) -> bool {
        false
    }

    /// Add a number of crates that should be whitelisted for showing up during
    /// queries, even if they are yanked. Currently only applies to registry
    /// sources.
    fn add_to_yanked_whitelist(&mut self, pkgs: &[PackageId]);

    /// Query if a package is yanked. Only registry sources can mark packages
    /// as yanked. This ignores the yanked whitelist.
    fn is_yanked(&mut self, _pkg: PackageId) -> CargoResult<bool>;
}

pub enum MaybePackage {
    Ready(Package),
    Download { url: String, descriptor: String },
}

impl<'a, T: Source + ?Sized + 'a> Source for Box<T> {
    /// Forwards to `Source::source_id`.
    fn source_id(&self) -> SourceId {
        (**self).source_id()
    }

    /// Forwards to `Source::replaced_source_id`.
    fn replaced_source_id(&self) -> SourceId {
        (**self).replaced_source_id()
    }

    /// Forwards to `Source::supports_checksums`.
    fn supports_checksums(&self) -> bool {
        (**self).supports_checksums()
    }

    /// Forwards to `Source::requires_precise`.
    fn requires_precise(&self) -> bool {
        (**self).requires_precise()
    }

    /// Forwards to `Source::query`.
    fn query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).query(dep, f)
    }

    /// Forwards to `Source::query`.
    fn fuzzy_query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).fuzzy_query(dep, f)
    }

    /// Forwards to `Source::update`.
    fn update(&mut self) -> CargoResult<()> {
        (**self).update()
    }

    /// Forwards to `Source::download`.
    fn download(&mut self, id: PackageId) -> CargoResult<MaybePackage> {
        (**self).download(id)
    }

    fn finish_download(&mut self, id: PackageId, data: Vec<u8>) -> CargoResult<Package> {
        (**self).finish_download(id, data)
    }

    /// Forwards to `Source::fingerprint`.
    fn fingerprint(&self, pkg: &Package) -> CargoResult<String> {
        (**self).fingerprint(pkg)
    }

    /// Forwards to `Source::verify`.
    fn verify(&self, pkg: PackageId) -> CargoResult<()> {
        (**self).verify(pkg)
    }

    fn describe(&self) -> String {
        (**self).describe()
    }

    fn is_replaced(&self) -> bool {
        (**self).is_replaced()
    }

    fn add_to_yanked_whitelist(&mut self, pkgs: &[PackageId]) {
        (**self).add_to_yanked_whitelist(pkgs);
    }

    fn is_yanked(&mut self, pkg: PackageId) -> CargoResult<bool> {
        (**self).is_yanked(pkg)
    }
}

impl<'a, T: Source + ?Sized + 'a> Source for &'a mut T {
    fn source_id(&self) -> SourceId {
        (**self).source_id()
    }

    fn replaced_source_id(&self) -> SourceId {
        (**self).replaced_source_id()
    }

    fn supports_checksums(&self) -> bool {
        (**self).supports_checksums()
    }

    fn requires_precise(&self) -> bool {
        (**self).requires_precise()
    }

    fn query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).query(dep, f)
    }

    fn fuzzy_query(&mut self, dep: &Dependency, f: &mut dyn FnMut(Summary)) -> CargoResult<()> {
        (**self).fuzzy_query(dep, f)
    }

    fn update(&mut self) -> CargoResult<()> {
        (**self).update()
    }

    fn download(&mut self, id: PackageId) -> CargoResult<MaybePackage> {
        (**self).download(id)
    }

    fn finish_download(&mut self, id: PackageId, data: Vec<u8>) -> CargoResult<Package> {
        (**self).finish_download(id, data)
    }

    fn fingerprint(&self, pkg: &Package) -> CargoResult<String> {
        (**self).fingerprint(pkg)
    }

    fn verify(&self, pkg: PackageId) -> CargoResult<()> {
        (**self).verify(pkg)
    }

    fn describe(&self) -> String {
        (**self).describe()
    }

    fn is_replaced(&self) -> bool {
        (**self).is_replaced()
    }

    fn add_to_yanked_whitelist(&mut self, pkgs: &[PackageId]) {
        (**self).add_to_yanked_whitelist(pkgs);
    }

    fn is_yanked(&mut self, pkg: PackageId) -> CargoResult<bool> {
        (**self).is_yanked(pkg)
    }
}

/// A `HashMap` of `SourceId` -> `Box<Source>`.
#[derive(Default)]
pub struct SourceMap<'src> {
    map: HashMap<SourceId, Box<dyn Source + 'src>>,
}

// `impl Debug` on source requires specialization, if even desirable at all.
impl<'src> fmt::Debug for SourceMap<'src> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SourceMap ")?;
        f.debug_set().entries(self.map.keys()).finish()
    }
}

impl<'src> SourceMap<'src> {
    /// Creates an empty map.
    pub fn new() -> SourceMap<'src> {
        SourceMap {
            map: HashMap::new(),
        }
    }

    /// Like `HashMap::get`.
    pub fn get(&self, id: SourceId) -> Option<&(dyn Source + 'src)> {
        self.map.get(&id).map(|s| s.as_ref())
    }

    /// Like `HashMap::get_mut`.
    pub fn get_mut(&mut self, id: SourceId) -> Option<&mut (dyn Source + 'src)> {
        self.map.get_mut(&id).map(|s| s.as_mut())
    }

    /// Like `HashMap::insert`, but derives the `SourceId` key from the `Source`.
    pub fn insert(&mut self, source: Box<dyn Source + 'src>) {
        let id = source.source_id();
        self.map.insert(id, source);
    }

    /// Like `HashMap::len`.
    pub fn len(&self) -> usize {
        self.map.len()
    }

    /// Like `HashMap::iter_mut`.
    pub fn sources_mut<'a>(
        &'a mut self,
    ) -> impl Iterator<Item = (&'a SourceId, &'a mut (dyn Source + 'src))> {
        self.map.iter_mut().map(|(a, b)| (a, &mut **b))
    }

    /// Merge the given map into self.
    pub fn add_source_map(&mut self, other: SourceMap<'src>) {
        for (key, value) in other.map {
            self.map.entry(key).or_insert(value);
        }
    }
}
<|endoftext|>


<commit_before><commit_msg>Add flag for overriding variables<commit_after><|endoftext|>


<commit_before>
use std;
import std::io;
import std::io::{writer_util, reader_util};
import result;

enum state {
    fieldstart(bool),
    infield(uint, uint),
    inescapedfield(uint, uint),
    inquote(uint, uint)
}

type rowreader = {
    delim: char,
    quote: char,
    f : io::reader,
    mutable offset : uint,
    mutable buffers : [@[char]],
    mutable state : state
};

type row = {
    fields : [ fieldtype ]
};

type bufferdescr = {
    escaped: bool,
    buffers: [@[char]],
    start: uint,
    end: uint
};

enum fieldtype {
    emptyfield(),
    bufferfield(bufferdescr)
}

iface rowiter {
    fn readrow() -> result::t<row, str>;
}

iface rowaccess {
    fn len() -> uint;
    fn getchars(uint) -> [char];
    fn getstr(uint) -> str;
}

fn new_reader(+f: io::reader, +delim: char, +quote: char) -> rowreader {
    let r : rowreader = {
        delim: delim,
        quote: quote,
        f: f,
        mutable offset : 0u,
        mutable buffers : [],
        mutable state : fieldstart(false)
    };
    ret r;
}

impl of rowaccess for row {
    fn len() -> uint {
        vec::len(self.fields)
    }
    fn getchars(field: uint) -> [char] {
        fn extract_field(buffers: [@[char]], start: uint, end: uint, &r: [char]) {
            let i = 0u;
            while i < vec::len(buffers) {
                let from = i == 0u ? start : 0u;
                let to = (i == vec::len(buffers) - 1u) ?
                         end : vec::len(*buffers[i]);
                r += vec::slice(*buffers[i], from, to);
                i = i + 1u;
            }
        }
        fn unescape(escaped: [char]) -> [char] {
            let r : [char] = [];
            vec::reserve(r, vec::len(escaped));
            let in_q = false;
            for c in escaped {
                if in_q {
                    assert(c == '"');
                    in_q = false;
                } else {
                    in_q = c == '"';
                    r += [c];
                }
            }
            ret r;
        }
        alt self.fields[field] {
            emptyfield() { ret []; }
            bufferfield(desc) {
                let buf = [];
                extract_field(desc.buffers, desc.start, desc.end, buf);
                if desc.escaped {
                    buf = unescape(buf);
                }
                ret buf;
            }
        };
    }
    fn getstr(field: uint) -> str {
        ret str::from_chars(self.getchars(field));
    }
}

impl of rowiter for rowreader {
    fn readrow() -> result::t<row, str> {
        fn row_from_buf(self: rowreader, &fields: [fieldtype]) -> bool {
            fn new_bufferfield(self: rowreader, escaped: bool, sb: uint, so: uint, eo: uint) -> fieldtype {
                let eb = vec::len(self.buffers);
                let sb = sb, so = so, eo = eo;
                if escaped {
                    so += 1u;
                    if so >= vec::len(*self.buffers[sb]) {
                        sb += 1u;
                        so = vec::len(*self.buffers[sb]) - 1u;
                    }
                    if eo > 0u {
                        eo -= 1u;
                    } else {
                        eb -= 1u;
                        eo = vec::len(*self.buffers[sb]) - 1u;
                    }
                }
                bufferfield({ escaped: escaped, buffers: vec::slice(self.buffers, sb, eb), start: so, end: eo })
            }
            let cbuffer = vec::len(self.buffers) - 1u;
            let buf: @[char] = self.buffers[cbuffer];
            while self.offset < vec::len(*buf) {
                let coffset = self.offset;
                let c : char = buf[coffset];
                self.offset += 1u;
                alt self.state {
                    fieldstart(after_delim) {
                        //io::println(#fmt("fieldstart : after_delim %b", after_delim));
                        if c == self.quote {
                            self.state = inescapedfield(cbuffer, coffset);
                        } else if c == '\n' {
                            if after_delim {
                                fields += [emptyfield];
                            }
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [emptyfield];
                        } else {
                            self.state = infield(cbuffer, coffset);
                        }
                    }
                    infield(b,o) {
                        //io::println(#fmt("field : %u %u", b, o));
                        if c == '\n' {
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                        }
                    }
                    inescapedfield(b, o) {
                        //io::println(#fmt("inescapedfield : %u %u", b, o));
                        if c == self.quote {
                            self.state = inquote(b, o);
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                        }
                    }
                    inquote(b, o) {
                        //io::println(#fmt("inquote : %u %u", b, o));
                        if c == '\n' {
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                            ret true;
                        } else if c == self.quote {
                            self.state = inescapedfield(b, o);
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                        }
                        // swallow odd chars, eg. space between field and "
                    }
                }
            }
            ret false;
        }

        self.state = fieldstart(false);
        let do_read = vec::len(self.buffers) == 0u;
        let fields = [];
        while true {
            if do_read {
                let data: @[char] = @self.f.read_chars(1024u);
                //io::println(#fmt("len %u '%s'", vec::len(*data), str::from_chars(*data)));
                if vec::len(*data) == 0u {
                    ret result::err("EOF");
                }
                self.buffers += [data];
                self.offset = 0u;
            }

            if row_from_buf(self, fields) {
                let r: row = { fields: fields };
                fields = [];
                if vec::len(self.buffers) > 1u {
                    self.buffers = vec::slice(self.buffers, vec::len(self.buffers) - 1u, vec::len(self.buffers));
                }
                ret result::ok(r);
            }
            do_read = true;
        }
        ret result::err("unreachable");
    }
}

fn main(args : [str])
{
    if (vec::len(args) != 2u) {
        ret;
    }
    let f : io::reader = result::get(io::file_reader(args[1]));
    let reader = new_reader(f, ',', '"');
    while true {
        let res = reader.readrow();
        if result::failure(res) {
            break;
        }
        let row = result::get(res);
        io::println(#fmt("---- ROW %u fields -----", row.len()));
        let i = 0u;
        while i < row.len() {
            io::println(row.getstr(i));
            i = i + 1u;
        }
    }
}

<commit_msg>speed a tiny bit<commit_after>
use std;
import std::io;
import std::io::{writer_util, reader_util};
import result;

enum state {
    fieldstart(bool),
    infield(uint, uint),
    inescapedfield(uint, uint),
    inquote(uint, uint)
}

type rowreader = {
    delim: char,
    quote: char,
    f : io::reader,
    mutable offset : uint,
    mutable buffers : [@[char]],
    mutable state : state
};

type row = {
    fields : [ fieldtype ]
};

type bufferdescr = {
    escaped: bool,
    buffers: [@[char]],
    start: uint,
    end: uint
};

enum fieldtype {
    emptyfield(),
    bufferfield(bufferdescr)
}

iface rowiter {
    fn readrow() -> result::t<row, str>;
}

iface rowaccess {
    fn len() -> uint;
    fn getchars(uint) -> [char];
    fn getstr(uint) -> str;
}

fn new_reader(+f: io::reader, +delim: char, +quote: char) -> rowreader {
    let r : rowreader = {
        delim: delim,
        quote: quote,
        f: f,
        mutable offset : 0u,
        mutable buffers : [],
        mutable state : fieldstart(false)
    };
    ret r;
}

impl of rowaccess for row {
    fn len() -> uint {
        vec::len(self.fields)
    }
    fn getchars(field: uint) -> [char] {
        fn extract_field(buffers: [@[char]], start: uint, end: uint, &r: [char]) {
            let i = 0u;
            while i < vec::len(buffers) {
                let from = i == 0u ? start : 0u;
                let to = (i == vec::len(buffers) - 1u) ?
                         end : vec::len(*buffers[i]);
                r += vec::slice(*buffers[i], from, to);
                i = i + 1u;
            }
        }
        fn unescape(escaped: [char]) -> [char] {
            let r : [char] = [];
            vec::reserve(r, vec::len(escaped));
            let in_q = false;
            for c in escaped {
                if in_q {
                    assert(c == '"');
                    in_q = false;
                } else {
                    in_q = c == '"';
                    r += [c];
                }
            }
            ret r;
        }
        alt self.fields[field] {
            emptyfield() { ret []; }
            bufferfield(desc) {
                let buf = [];
                extract_field(desc.buffers, desc.start, desc.end, buf);
                if desc.escaped {
                    buf = unescape(buf);
                }
                ret buf;
            }
        };
    }
    fn getstr(field: uint) -> str {
        ret str::from_chars(self.getchars(field));
    }
}

impl of rowiter for rowreader {
    fn readrow() -> result::t<row, str> {
        fn row_from_buf(self: rowreader, &fields: [fieldtype]) -> bool {
            fn new_bufferfield(self: rowreader, escaped: bool, sb: uint, so: uint, eo: uint) -> fieldtype {
                let eb = vec::len(self.buffers);
                let sb = sb, so = so, eo = eo;
                if escaped {
                    so += 1u;
                    if so >= vec::len(*self.buffers[sb]) {
                        sb += 1u;
                        so = vec::len(*self.buffers[sb]) - 1u;
                    }
                    if eo > 0u {
                        eo -= 1u;
                    } else {
                        eb -= 1u;
                        eo = vec::len(*self.buffers[sb]) - 1u;
                    }
                }
                bufferfield({ escaped: escaped, buffers: vec::slice(self.buffers, sb, eb), start: so, end: eo })
            }
            let cbuffer = vec::len(self.buffers) - 1u;
            let buf: @[char] = self.buffers[cbuffer];
            while self.offset < vec::len(*buf) {
                let coffset = self.offset;
                let c : char = buf[coffset];
                self.offset += 1u;
                alt self.state {
                    fieldstart(after_delim) {
                        //io::println(#fmt("fieldstart : after_delim %b", after_delim));
                        if c == self.quote {
                            self.state = inescapedfield(cbuffer, coffset);
                        } else if c == '\n' {
                            if after_delim {
                                fields += [emptyfield];
                            }
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [emptyfield];
                        } else {
                            self.state = infield(cbuffer, coffset);
                        }
                    }
                    infield(b,o) {
                        //io::println(#fmt("field : %u %u", b, o));
                        if c == '\n' {
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                        }
                    }
                    inescapedfield(b, o) {
                        //io::println(#fmt("inescapedfield : %u %u", b, o));
                        if c == self.quote {
                            self.state = inquote(b, o);
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                        }
                    }
                    inquote(b, o) {
                        //io::println(#fmt("inquote : %u %u", b, o));
                        if c == '\n' {
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                            ret true;
                        } else if c == self.quote {
                            self.state = inescapedfield(b, o);
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                        }
                        // swallow odd chars, eg. space between field and "
                    }
                }
            }
            ret false;
        }

        self.state = fieldstart(false);
        let do_read = vec::len(self.buffers) == 0u;
        let fields = [];
        while true {
            if do_read {
                let data: @[char] = @self.f.read_chars(4096u);
                //io::println(#fmt("len %u '%s'", vec::len(*data), str::from_chars(*data)));
                if vec::len(*data) == 0u {
                    ret result::err("EOF");
                }
                self.buffers += [data];
                self.offset = 0u;
            }

            if row_from_buf(self, fields) {
                let r: row = { fields: fields };
                fields = [];
                if vec::len(self.buffers) > 1u {
                    self.buffers = vec::slice(self.buffers, vec::len(self.buffers) - 1u, vec::len(self.buffers));
                }
                ret result::ok(r);
            }
            do_read = true;
        }
        ret result::err("unreachable");
    }
}

fn main(args : [str])
{
    if (vec::len(args) != 2u) {
        ret;
    }
    let f : io::reader = result::get(io::file_reader(args[1]));
    let reader = new_reader(f, ',', '"');
    while true {
        let res = reader.readrow();
        if result::failure(res) {
            break;
        }
        let row = result::get(res);
        io::println(#fmt("---- ROW %u fields -----", row.len()));
        let i = 0u;
        while i < row.len() {
            io::println(row.getstr(i));
            i = i + 1u;
        }
    }
}

<|endoftext|>