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

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

extern crate crossbeam;
extern crate rand;

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

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

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

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

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

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

struct Job {
    func: JobFn,
    sender: Sender<JobExitStatus>,
    scope_counter: *const AtomicUsize,
}

unsafe impl Send for Job {}

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

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

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

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

            unused.push(i);
        }

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

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

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

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

        let (tx, rx) = channel();

        let entry = &jobs[next];

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

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

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

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

        jobs.reserve(new_len);

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

            self.unused.push(i);
        }
    }

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

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

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

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

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

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

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

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

        guard.pop_back()
    }

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

        guard.pop_front()
    }

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

        guard.push_back(job);
    }
}

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

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

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

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

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

        f(&s);

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

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

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

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

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

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

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

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

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

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

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

}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            _ => {}
        }
        worker.run_next_job();
    }
}

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

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

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

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

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

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

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

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

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

            handles.push(handle);
        }

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

    }

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

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

            for i in v {
                assert_eq!(i, 1);
            }
        }).wait().unwrap();
    }
}<commit_msg>Switched TreiberStack with SegQueue<commit_after>#![feature(catch_panic, fnbox)]

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

extern crate crossbeam;
extern crate rand;

use crossbeam::mem::CachePadded;
use crossbeam::sync::SegQueue;

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

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

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

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

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

struct Job {
    func: JobFn,
    sender: Sender<JobExitStatus>,
    scope_counter: *const AtomicUsize,
}

unsafe impl Send for Job {}

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

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

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

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

            unused.push(i);
        }

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

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

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

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

        let (tx, rx) = channel();

        let entry = &jobs[next];

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

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

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

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

        jobs.reserve(new_len);

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

            self.unused.push(i);
        }
    }

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

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

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

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

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

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

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

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

        guard.pop_back()
    }

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

        guard.pop_front()
    }

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

        guard.push_back(job);
    }
}

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

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

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

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

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

        f(&s);

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

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

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

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

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

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

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

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

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

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

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

}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            _ => {}
        }
        worker.run_next_job();
    }
}

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

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

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

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

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

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

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

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

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

            handles.push(handle);
        }

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

    }

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

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

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


<commit_before>//! Lettre is a mailer written in Rust. It provides a simple email builder and several transports.
//!
//! ## Overview
//!
//! This mailer is divided into:
//!
//! * An `email` part: builds the email message
//! * A `transport` part: contains the available transports for your emails. To be sendable, the
//!   emails have to implement `SendableEmail`.
//!
//! ## Creating messages
//!
//! The `email` part builds email messages. For now, it does not support attachment nor alternative bodies.
//! An email is built using an `EmailBuilder`. The simplest email could be:
//!
//! ```rust
//! use lettre::email::EmailBuilder;
//!
//! // Create an email
//! let email = EmailBuilder::new()
//!     // Addresses can be specified by the couple (email, alias)
//!     .to(("user@example.org", "Firstname Lastname"))
//!     // ... or by an address only
//!     .from("user@example.com")
//!     .subject("Hi, Hello world")
//!     .body("Hello world.")
//!     .build();
//!
//! assert!(email.is_ok());
//! ```
//!
//! When the `build` method is called, the `EmailBuilder` will add the missing headers (like `Message-ID` or `Date`) and
//! check for missing necessary ones (like `From` or `To`). It will then generate an `Email` that can
//! be sent.
//!
//! Below is a more complete example, not using method chaining:
//!
//! ```rust
//! use lettre::email::EmailBuilder;
//!
//! let mut builder = EmailBuilder::new();
//! builder = builder.to(("user@example.org", "Alias name"));
//! builder = builder.cc(("user@example.net", "Alias name"));
//! builder = builder.from("no-reply@example.com");
//! builder = builder.from("no-reply@example.eu");
//! builder = builder.sender("no-reply@example.com");
//! builder = builder.subject("Hello world");
//! builder = builder.body("Hi, Hello world.");
//! builder = builder.reply_to("contact@example.com");
//! builder = builder.add_header(("X-Custom-Header", "my header"));
//!
//! let email = builder.build();
//! assert!(email.is_ok());
//! ```
//!
//! See the `EmailBuilder` documentation for a complete list of methods.
//!
//! ## Sending messages
//!
//! The following sections describe the available transport methods to treat emails.
//!
//! * The `SmtpTransport` uses the SMTP protocol to send the message over the network. It is
//!   the prefered way of sending emails.
//! * The `FileTransport` creates a file containing the email content to be sent. It can be used for debugging
//!   or if you want to keep all sent emails.
//! * The `StubTransport` is useful for debugging, and only prints the content of the email in the logs.
//!
//! ### SMTP transport
//!
//! This SMTP follows [RFC
//! 5321](https://tools.ietf.org/html/rfc5321), but is still
//! a work in progress. It is designed to efficiently send emails from an
//! application to a
//! relay email server, as it relies as much as possible on the relay server
//! for sanity and RFC
//! compliance checks.
//!
//! It implements the following extensions:
//!
//! * 8BITMIME ([RFC 6152](https://tools.ietf.org/html/rfc6152))
//! * AUTH ([RFC 4954](http://tools.ietf.org/html/rfc4954)) with PLAIN and
//! CRAM-MD5 mechanisms
//! * STARTTLS ([RFC 2487](http://tools.ietf.org/html/rfc2487))
//! * SMTPUTF8 ([RFC 6531](http://tools.ietf.org/html/rfc6531))
//!
//! #### Simple example
//!
//! This is the most basic example of usage:
//!
//! ```rust
//! use lettre::transport::smtp::{SmtpTransport, SmtpTransportBuilder};
//! use lettre::email::EmailBuilder;
//! use lettre::transport::EmailTransport;
//!
//! let email = EmailBuilder::new()
//!                     .to("root@localhost")
//!                     .from("user@localhost")
//!                     .body("Hello World!")
//!                     .subject("Hello")
//!                     .build()
//!                     .unwrap();
//!
//! // Open a local connection on port 25
//! let mut mailer =
//! SmtpTransportBuilder::localhost().unwrap().build();
//! // Send the email
//! let result = mailer.send(email);
//!
//! assert!(result.is_ok());
//! ```
//!
//! #### Complete example
//!
//! ```rust,no_run
//! use lettre::email::EmailBuilder;
//! use lettre::transport::smtp::{SecurityLevel, SmtpTransport,
//! SmtpTransportBuilder};
//! use lettre::transport::smtp::authentication::Mechanism;
//! use lettre::transport::smtp::SUBMISSION_PORT;
//! use lettre::transport::EmailTransport;
//!
//! let email = EmailBuilder::new()
//!                     .to("root@localhost")
//!                     .from("user@localhost")
//!                     .body("Hello World!")
//!                     .subject("Hello")
//!                     .build()
//!                     .unwrap();
//!
//! // Connect to a remote server on a custom port
//! let mut mailer = SmtpTransportBuilder::new(("server.tld",
//! SUBMISSION_PORT)).unwrap()
//!     // Set the name sent during EHLO/HELO, default is `localhost`
//!     .hello_name("my.hostname.tld")
//!     // Add credentials for authentication
//!     .credentials("username", "password")
//!     // Specify a TLS security level. You can also specify an SslContext with
//!     // .ssl_context(SslContext::Ssl23)
//!     .security_level(SecurityLevel::AlwaysEncrypt)
//!     // Enable SMTPUTF8 is the server supports it
//!     .smtp_utf8(true)
//!     // Configure accepted authetication mechanisms
//!     .authentication_mechanisms(vec![Mechanism::CramMd5])
//!     // Enable connection reuse
//!     .connection_reuse(true).build();
//!
//! let result_1 = mailer.send(email.clone());
//! assert!(result_1.is_ok());
//!
//! // The second email will use the same connection
//! let result_2 = mailer.send(email);
//! assert!(result_2.is_ok());
//!
//! // Explicitely close the SMTP transaction as we enabled connection reuse
//! mailer.close();
//! ```
//!
//! #### Lower level
//!
//! You can also send commands, here is a simple email transaction without
//! error handling:
//!
//! ```rust
//! use lettre::transport::smtp::SMTP_PORT;
//! use lettre::transport::smtp::client::Client;
//! use lettre::transport::smtp::client::net::NetworkStream;
//!
//! let mut email_client: Client<NetworkStream> = Client::new();
//! let _ = email_client.connect(&("localhost", SMTP_PORT), None);
//! let _ = email_client.ehlo("my_hostname");
//! let _ = email_client.mail("user@example.com", None);
//! let _ = email_client.rcpt("user@example.org");
//! let _ = email_client.data();
//! let _ = email_client.message("Test email");
//! let _ = email_client.quit();
//! ```
//!
//! ### Stub transport
//!
//! The stub transport only logs message envelope and drops the content. It can be useful for
//! testing purposes.
//!
//! ```rust
//! use lettre::transport::stub::StubEmailTransport;
//! use lettre::transport::EmailTransport;
//! use lettre::email::EmailBuilder;
//!
//! let email = EmailBuilder::new()
//!                     .to("root@localhost")
//!                     .from("user@localhost")
//!                     .body("Hello World!")
//!                     .subject("Hello")
//!                     .build()
//!                     .unwrap();
//!
//! let mut sender = StubEmailTransport;
//! let result = sender.send(email);
//! assert!(result.is_ok());
//! ```
//!
//! Will log the line:
//!
//! ```text
//! b7c211bc-9811-45ce-8cd9-68eab575d695: from=<user@localhost> to=<root@localhost>
//! ```
//!
//! ### File transport
//!
//! The file transport writes the emails to the given directory. The name of the file will be
//! `message_id.txt`.
//! It can be useful for testing purposes, or if you want to keep track of sent messages.
//!
//! ```rust
//! use std::env::temp_dir;
//!
//! use lettre::transport::file::FileEmailTransport;
//! use lettre::transport::EmailTransport;
//! use lettre::email::{EmailBuilder, SendableEmail};
//!
//! // Write to the local temp directory
//! let mut sender = FileEmailTransport::new(temp_dir());
//! let email = EmailBuilder::new()
//!                 .to("root@localhost")
//!                 .from("user@localhost")
//!                 .body("Hello World!")
//!                 .subject("Hello")
//!                 .build()
//!                 .unwrap();
//!
//! let result = sender.send(email);
//! assert!(result.is_ok());
//! ```
//! Example result in `/tmp/b7c211bc-9811-45ce-8cd9-68eab575d695.txt`:
//!
//! ```text
//! b7c211bc-9811-45ce-8cd9-68eab575d695: from=<user@localhost> to=<root@localhost>
//! To: <root@localhost>
//! From: <user@localhost>
//! Subject: Hello
//! Date: Sat, 31 Oct 2015 13:42:19 +0100
//! Message-ID: <b7c211bc-9811-45ce-8cd9-68eab575d695.lettre@localhost>
//!
//! Hello World!
//! ```

#![deny(missing_docs, unsafe_code, unstable_features)]

#[macro_use]
extern crate log;
extern crate rustc_serialize;
extern crate crypto;
extern crate time;
extern crate uuid;
extern crate email as email_format;
extern crate bufstream;
extern crate openssl;

pub mod transport;
pub mod email;
<commit_msg>Doc formatting<commit_after>//! Lettre is a mailer written in Rust. It provides a simple email builder and several transports.
//!
//! ## Overview
//!
//! This mailer is divided into:
//!
//! * An `email` part: builds the email message
//! * A `transport` part: contains the available transports for your emails. To be sendable, the
//!   emails have to implement `SendableEmail`.
//!
//! ## Creating messages
//!
//! The `email` part builds email messages. For now, it does not support attachment nor alternative
//! bodies. An email is built using an `EmailBuilder`. The simplest email could be:
//!
//! ```rust
//! use lettre::email::EmailBuilder;
//!
//! // Create an email
//! let email = EmailBuilder::new()
//!     // Addresses can be specified by the couple (email, alias)
//!     .to(("user@example.org", "Firstname Lastname"))
//!     // ... or by an address only
//!     .from("user@example.com")
//!     .subject("Hi, Hello world")
//!     .body("Hello world.")
//!     .build();
//!
//! assert!(email.is_ok());
//! ```
//!
//! When the `build` method is called, the `EmailBuilder` will add the missing headers (like
//! `Message-ID` or `Date`) and check for missing necessary ones (like `From` or `To`). It will
//! then generate an `Email` that can be sent.
//!
//! Below is a more complete example, not using method chaining:
//!
//! ```rust
//! use lettre::email::EmailBuilder;
//!
//! let mut builder = EmailBuilder::new();
//! builder = builder.to(("user@example.org", "Alias name"));
//! builder = builder.cc(("user@example.net", "Alias name"));
//! builder = builder.from("no-reply@example.com");
//! builder = builder.from("no-reply@example.eu");
//! builder = builder.sender("no-reply@example.com");
//! builder = builder.subject("Hello world");
//! builder = builder.body("Hi, Hello world.");
//! builder = builder.reply_to("contact@example.com");
//! builder = builder.add_header(("X-Custom-Header", "my header"));
//!
//! let email = builder.build();
//! assert!(email.is_ok());
//! ```
//!
//! See the `EmailBuilder` documentation for a complete list of methods.
//!
//! ## Sending messages
//!
//! The following sections describe the available transport methods to treat emails.
//!
//! * The `SmtpTransport` uses the SMTP protocol to send the message over the network. It is
//!   the prefered way of sending emails.
//! * The `FileTransport` creates a file containing the email content to be sent. It can be used
//!   for debugging or if you want to keep all sent emails.
//! * The `StubTransport` is useful for debugging, and only prints the content of the email in the
//!   logs.
//!
//! ### SMTP transport
//!
//! This SMTP follows [RFC
//! 5321](https://tools.ietf.org/html/rfc5321), but is still
//! a work in progress. It is designed to efficiently send emails from an
//! application to a
//! relay email server, as it relies as much as possible on the relay server
//! for sanity and RFC
//! compliance checks.
//!
//! It implements the following extensions:
//!
//! * 8BITMIME ([RFC 6152](https://tools.ietf.org/html/rfc6152))
//! * AUTH ([RFC 4954](http://tools.ietf.org/html/rfc4954)) with PLAIN and
//! CRAM-MD5 mechanisms
//! * STARTTLS ([RFC 2487](http://tools.ietf.org/html/rfc2487))
//! * SMTPUTF8 ([RFC 6531](http://tools.ietf.org/html/rfc6531))
//!
//! #### Simple example
//!
//! This is the most basic example of usage:
//!
//! ```rust
//! use lettre::transport::smtp::{SmtpTransport, SmtpTransportBuilder};
//! use lettre::email::EmailBuilder;
//! use lettre::transport::EmailTransport;
//!
//! let email = EmailBuilder::new()
//!                     .to("root@localhost")
//!                     .from("user@localhost")
//!                     .body("Hello World!")
//!                     .subject("Hello")
//!                     .build()
//!                     .unwrap();
//!
//! // Open a local connection on port 25
//! let mut mailer =
//! SmtpTransportBuilder::localhost().unwrap().build();
//! // Send the email
//! let result = mailer.send(email);
//!
//! assert!(result.is_ok());
//! ```
//!
//! #### Complete example
//!
//! ```rust,no_run
//! use lettre::email::EmailBuilder;
//! use lettre::transport::smtp::{SecurityLevel, SmtpTransport,
//! SmtpTransportBuilder};
//! use lettre::transport::smtp::authentication::Mechanism;
//! use lettre::transport::smtp::SUBMISSION_PORT;
//! use lettre::transport::EmailTransport;
//!
//! let email = EmailBuilder::new()
//!                     .to("root@localhost")
//!                     .from("user@localhost")
//!                     .body("Hello World!")
//!                     .subject("Hello")
//!                     .build()
//!                     .unwrap();
//!
//! // Connect to a remote server on a custom port
//! let mut mailer = SmtpTransportBuilder::new(("server.tld",
//! SUBMISSION_PORT)).unwrap()
//!     // Set the name sent during EHLO/HELO, default is `localhost`
//!     .hello_name("my.hostname.tld")
//!     // Add credentials for authentication
//!     .credentials("username", "password")
//!     // Specify a TLS security level. You can also specify an SslContext with
//!     // .ssl_context(SslContext::Ssl23)
//!     .security_level(SecurityLevel::AlwaysEncrypt)
//!     // Enable SMTPUTF8 is the server supports it
//!     .smtp_utf8(true)
//!     // Configure accepted authetication mechanisms
//!     .authentication_mechanisms(vec![Mechanism::CramMd5])
//!     // Enable connection reuse
//!     .connection_reuse(true).build();
//!
//! let result_1 = mailer.send(email.clone());
//! assert!(result_1.is_ok());
//!
//! // The second email will use the same connection
//! let result_2 = mailer.send(email);
//! assert!(result_2.is_ok());
//!
//! // Explicitely close the SMTP transaction as we enabled connection reuse
//! mailer.close();
//! ```
//!
//! #### Lower level
//!
//! You can also send commands, here is a simple email transaction without
//! error handling:
//!
//! ```rust
//! use lettre::transport::smtp::SMTP_PORT;
//! use lettre::transport::smtp::client::Client;
//! use lettre::transport::smtp::client::net::NetworkStream;
//!
//! let mut email_client: Client<NetworkStream> = Client::new();
//! let _ = email_client.connect(&("localhost", SMTP_PORT), None);
//! let _ = email_client.ehlo("my_hostname");
//! let _ = email_client.mail("user@example.com", None);
//! let _ = email_client.rcpt("user@example.org");
//! let _ = email_client.data();
//! let _ = email_client.message("Test email");
//! let _ = email_client.quit();
//! ```
//!
//! ### Stub transport
//!
//! The stub transport only logs message envelope and drops the content. It can be useful for
//! testing purposes.
//!
//! ```rust
//! use lettre::transport::stub::StubEmailTransport;
//! use lettre::transport::EmailTransport;
//! use lettre::email::EmailBuilder;
//!
//! let email = EmailBuilder::new()
//!                     .to("root@localhost")
//!                     .from("user@localhost")
//!                     .body("Hello World!")
//!                     .subject("Hello")
//!                     .build()
//!                     .unwrap();
//!
//! let mut sender = StubEmailTransport;
//! let result = sender.send(email);
//! assert!(result.is_ok());
//! ```
//!
//! Will log the line:
//!
//! ```text
//! b7c211bc-9811-45ce-8cd9-68eab575d695: from=<user@localhost> to=<root@localhost>
//! ```
//!
//! ### File transport
//!
//! The file transport writes the emails to the given directory. The name of the file will be
//! `message_id.txt`.
//! It can be useful for testing purposes, or if you want to keep track of sent messages.
//!
//! ```rust
//! use std::env::temp_dir;
//!
//! use lettre::transport::file::FileEmailTransport;
//! use lettre::transport::EmailTransport;
//! use lettre::email::{EmailBuilder, SendableEmail};
//!
//! // Write to the local temp directory
//! let mut sender = FileEmailTransport::new(temp_dir());
//! let email = EmailBuilder::new()
//!                 .to("root@localhost")
//!                 .from("user@localhost")
//!                 .body("Hello World!")
//!                 .subject("Hello")
//!                 .build()
//!                 .unwrap();
//!
//! let result = sender.send(email);
//! assert!(result.is_ok());
//! ```
//! Example result in `/tmp/b7c211bc-9811-45ce-8cd9-68eab575d695.txt`:
//!
//! ```text
//! b7c211bc-9811-45ce-8cd9-68eab575d695: from=<user@localhost> to=<root@localhost>
//! To: <root@localhost>
//! From: <user@localhost>
//! Subject: Hello
//! Date: Sat, 31 Oct 2015 13:42:19 +0100
//! Message-ID: <b7c211bc-9811-45ce-8cd9-68eab575d695.lettre@localhost>
//!
//! Hello World!
//! ```

#![deny(missing_docs, unsafe_code, unstable_features)]

#[macro_use]
extern crate log;
extern crate rustc_serialize;
extern crate crypto;
extern crate time;
extern crate uuid;
extern crate email as email_format;
extern crate bufstream;
extern crate openssl;

pub mod transport;
pub mod email;
<|endoftext|>


<commit_before>//! Manipulations and data types that represent polynomial.

#![warn(bad_style)]
#![warn(missing_docs)]
#![warn(trivial_casts)]
#![warn(trivial_numeric_casts)]
#![warn(unused)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

extern crate num_traits;

use std::{cmp, fmt};
use std::ops::{Add, Mul, Neg, Sub};
use num_traits::{Zero, One};

/// Polynomial expression
#[derive(Eq, PartialEq, Clone, Debug)]
pub struct Polynomial<T> {
    data: Vec<T>,
}

impl<T: Zero> Polynomial<T> {
    /// Creates new `Polynomial`.
    ///
    /// ```rust
    /// use polynomial::Polynomial;
    /// let poly = Polynomial::new(vec![1, 2, 3]);
    /// assert_eq!("1+2*x+3*x^2", poly.pretty("x"));
    /// ```
    #[inline]
    pub fn new(mut data: Vec<T>) -> Polynomial<T> {
        while let Some(true) = data.last().map(|x| x.is_zero()) {
            let _ = data.pop();
        }
        Polynomial { data: data }
    }
}

impl<T: Zero + One + Clone> Polynomial<T> {
    /// Evaluates the polynomial value.
    ///
    /// ```rust
    /// use polynomial::Polynomial;
    /// let poly = Polynomial::new(vec![1, 2, 3]);
    /// assert_eq!(1, poly.eval(0));
    /// assert_eq!(6, poly.eval(1));
    /// assert_eq!(17, poly.eval(2));
    /// ```
    #[inline]
    pub fn eval(&self, x: T) -> T {
        let mut sum: T = Zero::zero();
        let mut x_n: T = One::one();
        for n in self.data.iter() {
            sum = sum + n.clone() * x_n.clone();
            x_n = x_n * x.clone();
        }
        sum
    }
}

impl<T> Polynomial<T> {
    /// Gets the slice of internal data.
    #[inline]
    pub fn data(&self) -> &[T] {
        &self.data
    }
}

impl<T> Polynomial<T>
    where T: Zero + One + Eq + Neg<Output = T> + Ord + fmt::Display + Clone
{
    /// Pretty prints the polynomial.
    pub fn pretty(&self, x: &str) -> String {
        if self.is_zero() {
            return "0".to_string();
        }

        let one = One::one();
        let mut s = Vec::new();
        for (i, n) in self.data.iter().enumerate() {
            // output n*x^i / -n*x^i
            if n.is_zero() {
                continue;
            }

            let term = if i.is_zero() {
                n.to_string()
            } else if i == 1 {
                if (*n) == one {
                    x.to_string()
                } else if (*n) == -one.clone() {
                    format!("-{}", x)
                } else {
                    format!("{}*{}", n.to_string(), x)
                }
            } else {
                if (*n) == one {
                    format!("{}^{}", x, i)
                } else if (*n) == -one.clone() {
                    format!("-{}^{}", x, i)
                } else {
                    format!("{}*{}^{}", n.to_string(), x, i)
                }
            };

            if s.len() > 0 && (*n) > Zero::zero() {
                s.push("+".to_string());
            }
            s.push(term);
        }

        s.concat()
    }
}

impl<'a, T> Neg for Polynomial<T>
    where T: Neg + Zero + Clone,
          <T as Neg>::Output: Zero
{
    type Output = Polynomial<<T as Neg>::Output>;

    #[inline]
    fn neg(self) -> Polynomial<<T as Neg>::Output> {
        -&self
    }
}

impl<'a, T> Neg for &'a Polynomial<T>
    where T: Neg + Zero + Clone,
          <T as Neg>::Output: Zero
{
    type Output = Polynomial<<T as Neg>::Output>;

    #[inline]
    fn neg(self) -> Polynomial<<T as Neg>::Output> {
        Polynomial::new(self.data.iter().map(|x| -x.clone()).collect())
    }
}

macro_rules! forward_val_val_binop {
    (impl $imp:ident, $method:ident) => {
        impl<Lhs, Rhs> $imp<Polynomial<Rhs>> for Polynomial<Lhs>
            where Lhs: Zero + $imp<Rhs> + Clone, Rhs: Zero + Clone,
                  <Lhs as $imp<Rhs>>::Output: Zero
        {
            type Output = Polynomial<<Lhs as $imp<Rhs>>::Output>;

            #[inline]
            fn $method(self, other: Polynomial<Rhs>) -> Polynomial<<Lhs as $imp<Rhs>>::Output> {
                (&self).$method(&other)
            }
        }
    }
}

macro_rules! forward_ref_val_binop {
    (impl $imp:ident, $method:ident) => {
        impl<'a, Lhs, Rhs> $imp<Polynomial<Rhs>> for &'a Polynomial<Lhs>
            where Lhs: Zero + $imp<Rhs> + Clone, Rhs: Zero + Clone,
                  <Lhs as $imp<Rhs>>::Output: Zero
        {
            type Output = Polynomial<<Lhs as $imp<Rhs>>::Output>;

            #[inline]
            fn $method(self, other: Polynomial<Rhs>) -> Polynomial<<Lhs as $imp<Rhs>>::Output> {
                self.$method(&other)
            }
        }
    }
}

macro_rules! forward_val_ref_binop {
    (impl $imp:ident, $method:ident) => {
        impl<'a, Lhs, Rhs> $imp<&'a Polynomial<Rhs>> for Polynomial<Lhs>
            where Lhs: Zero + $imp<Rhs> + Clone, Rhs: Zero + Clone,
                  <Lhs as $imp<Rhs>>::Output: Zero
        {
            type Output = Polynomial<<Lhs as $imp<Rhs>>::Output>;

            #[inline]
            fn $method(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as $imp<Rhs>>::Output> {
                (&self).$method(other)
            }
        }
    }
}

macro_rules! forward_all_binop {
    (impl $imp:ident, $method:ident) => {
        forward_val_val_binop!(impl $imp, $method);
        forward_ref_val_binop!(impl $imp, $method);
        forward_val_ref_binop!(impl $imp, $method);
    };
}

forward_all_binop!(impl Add, add);

impl<'a, 'b, Lhs, Rhs> Add<&'b Polynomial<Rhs>> for &'a Polynomial<Lhs>
    where Lhs: Zero + Add<Rhs> + Clone,
          Rhs: Zero + Clone,
          <Lhs as Add<Rhs>>::Output: Zero
{
    type Output = Polynomial<<Lhs as Add<Rhs>>::Output>;

    fn add(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as Add<Rhs>>::Output> {
        let max_len = cmp::max(self.data.len(), other.data.len());
        let min_len = cmp::min(self.data.len(), other.data.len());

        let mut sum = Vec::with_capacity(max_len);
        for i in 0..min_len {
            sum.push(self.data[i].clone() + other.data[i].clone());
        }

        if self.data.len() <= other.data.len() {
            for i in min_len..max_len {
                sum.push(num_traits::zero::<Lhs>() + other.data[i].clone());
            }
        } else {
            for i in min_len..max_len {
                sum.push(self.data[i].clone() + num_traits::zero::<Rhs>());
            }
        }

        Polynomial::new(sum)
    }
}

forward_all_binop!(impl Sub, sub);

impl<'a, 'b, Lhs, Rhs> Sub<&'b Polynomial<Rhs>> for &'a Polynomial<Lhs>
    where Lhs: Zero + Sub<Rhs> + Clone,
          Rhs: Zero + Clone,
          <Lhs as Sub<Rhs>>::Output: Zero
{
    type Output = Polynomial<<Lhs as Sub<Rhs>>::Output>;

    fn sub(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as Sub<Rhs>>::Output> {
        let min_len = cmp::min(self.data.len(), other.data.len());
        let max_len = cmp::max(self.data.len(), other.data.len());

        let mut sub = Vec::with_capacity(max_len);
        for i in 0..min_len {
            sub.push(self.data[i].clone() - other.data[i].clone());
        }
        if self.data.len() <= other.data.len() {
            for i in min_len..max_len {
                sub.push(num_traits::zero::<Lhs>() - other.data[i].clone())
            }
        } else {
            for i in min_len..max_len {
                sub.push(self.data[i].clone() - num_traits::zero::<Rhs>())
            }
        }
        Polynomial::new(sub)
    }
}

forward_all_binop!(impl Mul, mul);

impl<'a, 'b, Lhs, Rhs> Mul<&'b Polynomial<Rhs>> for &'a Polynomial<Lhs>
    where Lhs: Zero + Mul<Rhs> + Clone,
          Rhs: Zero + Clone,
          <Lhs as Mul<Rhs>>::Output: Zero
{
    type Output = Polynomial<<Lhs as Mul<Rhs>>::Output>;

    fn mul(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as Mul<Rhs>>::Output> {
        if self.is_zero() || other.is_zero() {
            return Polynomial::new(vec![]);
        }

        let slen = self.data.len();
        let olen = other.data.len();
        let prod = (0..slen + olen - 1)
            .map(|i| {
                let mut p = num_traits::zero::<<Lhs as Mul<Rhs>>::Output>();
                let kstart = cmp::max(olen, i + 1) - olen;
                let kend = cmp::min(slen, i + 1);
                for k in kstart..kend {
                    p = p + self.data[k].clone() * other.data[i - k].clone();
                }
                p
            })
            .collect();
        Polynomial::new(prod)
    }
}

impl<T: Zero + Clone> Zero for Polynomial<T> {
    #[inline]
    fn zero() -> Polynomial<T> {
        Polynomial { data: vec![] }
    }
    #[inline]
    fn is_zero(&self) -> bool {
        self.data.is_empty()
    }
}

impl<T: Zero + One + Clone> One for Polynomial<T> {
    #[inline]
    fn one() -> Polynomial<T> {
        Polynomial { data: vec![One::one()] }
    }
}

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

    #[test]
    fn new() {
        fn check(dst: Vec<i32>, src: Vec<i32>) {
            assert_eq!(dst, Polynomial::new(src).data);
        }
        check(vec![1, 2, 3], vec![1, 2, 3]);
        check(vec![1, 2, 3], vec![1, 2, 3, 0, 0]);
        check(vec![], vec![0, 0, 0]);
    }

    #[test]
    fn neg_add_sub() {
        fn check(a: &[i32], b: &[i32], c: &[i32]) {
            fn check_eq(a: &Polynomial<i32>, b: &Polynomial<i32>) {
                assert_eq!(*a, *b);
                assert_eq!(-a, -b);
            }
            fn check_add(sum: &Polynomial<i32>, a: &Polynomial<i32>, b: &Polynomial<i32>) {
                check_eq(sum, &(a + b));
                check_eq(sum, &(b + a));
            }
            fn check_sub(sum: &Polynomial<i32>, a: &Polynomial<i32>, b: &Polynomial<i32>) {
                check_eq(a, &(sum - b));
                check_eq(b, &(sum - a));
            }

            let a = &Polynomial::new(a.to_vec());
            let b = &Polynomial::new(b.to_vec());
            let c = &Polynomial::new(c.to_vec());
            check_add(c, a, b);
            check_add(&(-c), &(-a), &(-b));
            check_sub(c, a, b);
            check_sub(&(-c), &(-a), &(-b));
        }
        check(&[], &[], &[]);
        check(&[], &[1], &[1]);
        check(&[1], &[1], &[2]);
        check(&[1, 0, 1], &[1], &[2, 0, 1]);
        check(&[1, 0, -1], &[-1, 0, 1], &[]);
    }

    #[test]
    fn mul() {
        fn check(a: &[i32], b: &[i32], c: &[i32]) {
            let a = Polynomial::new(a.to_vec());
            let b = Polynomial::new(b.to_vec());
            let c = Polynomial::new(c.to_vec());
            assert_eq!(c, &a * &b);
            assert_eq!(c, &b * &a);
        }
        check(&[], &[], &[]);
        check(&[0, 0], &[], &[]);
        check(&[0, 0], &[1], &[]);
        check(&[1, 0], &[1], &[1]);
        check(&[1, 0, 1], &[1], &[1, 0, 1]);
        check(&[1, 1], &[1, 1], &[1, 2, 1]);
        check(&[1, 1], &[1, 0, 1], &[1, 1, 1, 1]);
        check(&[0, 0, 1], &[0, 0, 1], &[0, 0, 0, 0, 1]);
    }

    #[test]
    fn eval() {
        fn check<F: Fn(i32) -> i32>(pol: &[i32], f: F) {
            for n in -10..10 {
                assert_eq!(f(n), Polynomial::new(pol.to_vec()).eval(n));
            }
        }
        check(&[], |_x| 0);
        check(&[1], |_x| 1);
        check(&[1, 1], |x| x + 1);
        check(&[0, 1], |x| x);
        check(&[10, -10, 10], |x| 10 * x * x - 10 * x + 10);
    }

    #[test]
    fn pretty() {
        fn check(slice: &[i32], s: &str) {
            assert_eq!(s.to_string(), Polynomial::new(slice.to_vec()).pretty("x"));
        }
        check(&[0], "0");
        check(&[1], "1");
        check(&[1, 1], "1+x");
        check(&[1, 1, 1], "1+x+x^2");
        check(&[2, 2, 2], "2+2*x+2*x^2");
        check(&[0, 0, 0, 1], "x^3");
        check(&[0, 0, 0, -1], "-x^3");
        check(&[-1, 0, 0, -1], "-1-x^3");
        check(&[-1, 1, 0, -1], "-1+x-x^3");
        check(&[-1, 1, -1, -1], "-1+x-x^2-x^3");
    }
}
<commit_msg>Run rustfmt<commit_after>//! Manipulations and data types that represent polynomial.

#![warn(bad_style)]
#![warn(missing_docs)]
#![warn(trivial_casts)]
#![warn(trivial_numeric_casts)]
#![warn(unused)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

extern crate num_traits;

use num_traits::{One, Zero};
use std::ops::{Add, Mul, Neg, Sub};
use std::{cmp, fmt};

/// Polynomial expression
#[derive(Eq, PartialEq, Clone, Debug)]
pub struct Polynomial<T> {
    data: Vec<T>,
}

impl<T: Zero> Polynomial<T> {
    /// Creates new `Polynomial`.
    ///
    /// ```rust
    /// use polynomial::Polynomial;
    /// let poly = Polynomial::new(vec![1, 2, 3]);
    /// assert_eq!("1+2*x+3*x^2", poly.pretty("x"));
    /// ```
    #[inline]
    pub fn new(mut data: Vec<T>) -> Polynomial<T> {
        while let Some(true) = data.last().map(|x| x.is_zero()) {
            let _ = data.pop();
        }
        Polynomial { data: data }
    }
}

impl<T: Zero + One + Clone> Polynomial<T> {
    /// Evaluates the polynomial value.
    ///
    /// ```rust
    /// use polynomial::Polynomial;
    /// let poly = Polynomial::new(vec![1, 2, 3]);
    /// assert_eq!(1, poly.eval(0));
    /// assert_eq!(6, poly.eval(1));
    /// assert_eq!(17, poly.eval(2));
    /// ```
    #[inline]
    pub fn eval(&self, x: T) -> T {
        let mut sum: T = Zero::zero();
        let mut x_n: T = One::one();
        for n in self.data.iter() {
            sum = sum + n.clone() * x_n.clone();
            x_n = x_n * x.clone();
        }
        sum
    }
}

impl<T> Polynomial<T> {
    /// Gets the slice of internal data.
    #[inline]
    pub fn data(&self) -> &[T] {
        &self.data
    }
}

impl<T> Polynomial<T>
where
    T: Zero + One + Eq + Neg<Output = T> + Ord + fmt::Display + Clone,
{
    /// Pretty prints the polynomial.
    pub fn pretty(&self, x: &str) -> String {
        if self.is_zero() {
            return "0".to_string();
        }

        let one = One::one();
        let mut s = Vec::new();
        for (i, n) in self.data.iter().enumerate() {
            // output n*x^i / -n*x^i
            if n.is_zero() {
                continue;
            }

            let term = if i.is_zero() {
                n.to_string()
            } else if i == 1 {
                if (*n) == one {
                    x.to_string()
                } else if (*n) == -one.clone() {
                    format!("-{}", x)
                } else {
                    format!("{}*{}", n.to_string(), x)
                }
            } else {
                if (*n) == one {
                    format!("{}^{}", x, i)
                } else if (*n) == -one.clone() {
                    format!("-{}^{}", x, i)
                } else {
                    format!("{}*{}^{}", n.to_string(), x, i)
                }
            };

            if s.len() > 0 && (*n) > Zero::zero() {
                s.push("+".to_string());
            }
            s.push(term);
        }

        s.concat()
    }
}

impl<'a, T> Neg for Polynomial<T>
where
    T: Neg + Zero + Clone,
    <T as Neg>::Output: Zero,
{
    type Output = Polynomial<<T as Neg>::Output>;

    #[inline]
    fn neg(self) -> Polynomial<<T as Neg>::Output> {
        -&self
    }
}

impl<'a, T> Neg for &'a Polynomial<T>
where
    T: Neg + Zero + Clone,
    <T as Neg>::Output: Zero,
{
    type Output = Polynomial<<T as Neg>::Output>;

    #[inline]
    fn neg(self) -> Polynomial<<T as Neg>::Output> {
        Polynomial::new(self.data.iter().map(|x| -x.clone()).collect())
    }
}

macro_rules! forward_val_val_binop {
    (impl $imp:ident, $method:ident) => {
        impl<Lhs, Rhs> $imp<Polynomial<Rhs>> for Polynomial<Lhs>
        where
            Lhs: Zero + $imp<Rhs> + Clone,
            Rhs: Zero + Clone,
            <Lhs as $imp<Rhs>>::Output: Zero,
        {
            type Output = Polynomial<<Lhs as $imp<Rhs>>::Output>;

            #[inline]
            fn $method(self, other: Polynomial<Rhs>) -> Polynomial<<Lhs as $imp<Rhs>>::Output> {
                (&self).$method(&other)
            }
        }
    };
}

macro_rules! forward_ref_val_binop {
    (impl $imp:ident, $method:ident) => {
        impl<'a, Lhs, Rhs> $imp<Polynomial<Rhs>> for &'a Polynomial<Lhs>
        where
            Lhs: Zero + $imp<Rhs> + Clone,
            Rhs: Zero + Clone,
            <Lhs as $imp<Rhs>>::Output: Zero,
        {
            type Output = Polynomial<<Lhs as $imp<Rhs>>::Output>;

            #[inline]
            fn $method(self, other: Polynomial<Rhs>) -> Polynomial<<Lhs as $imp<Rhs>>::Output> {
                self.$method(&other)
            }
        }
    };
}

macro_rules! forward_val_ref_binop {
    (impl $imp:ident, $method:ident) => {
        impl<'a, Lhs, Rhs> $imp<&'a Polynomial<Rhs>> for Polynomial<Lhs>
        where
            Lhs: Zero + $imp<Rhs> + Clone,
            Rhs: Zero + Clone,
            <Lhs as $imp<Rhs>>::Output: Zero,
        {
            type Output = Polynomial<<Lhs as $imp<Rhs>>::Output>;

            #[inline]
            fn $method(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as $imp<Rhs>>::Output> {
                (&self).$method(other)
            }
        }
    };
}

macro_rules! forward_all_binop {
    (impl $imp:ident, $method:ident) => {
        forward_val_val_binop!(impl $imp, $method);
        forward_ref_val_binop!(impl $imp, $method);
        forward_val_ref_binop!(impl $imp, $method);
    };
}

forward_all_binop!(impl Add, add);

impl<'a, 'b, Lhs, Rhs> Add<&'b Polynomial<Rhs>> for &'a Polynomial<Lhs>
where
    Lhs: Zero + Add<Rhs> + Clone,
    Rhs: Zero + Clone,
    <Lhs as Add<Rhs>>::Output: Zero,
{
    type Output = Polynomial<<Lhs as Add<Rhs>>::Output>;

    fn add(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as Add<Rhs>>::Output> {
        let max_len = cmp::max(self.data.len(), other.data.len());
        let min_len = cmp::min(self.data.len(), other.data.len());

        let mut sum = Vec::with_capacity(max_len);
        for i in 0..min_len {
            sum.push(self.data[i].clone() + other.data[i].clone());
        }

        if self.data.len() <= other.data.len() {
            for i in min_len..max_len {
                sum.push(num_traits::zero::<Lhs>() + other.data[i].clone());
            }
        } else {
            for i in min_len..max_len {
                sum.push(self.data[i].clone() + num_traits::zero::<Rhs>());
            }
        }

        Polynomial::new(sum)
    }
}

forward_all_binop!(impl Sub, sub);

impl<'a, 'b, Lhs, Rhs> Sub<&'b Polynomial<Rhs>> for &'a Polynomial<Lhs>
where
    Lhs: Zero + Sub<Rhs> + Clone,
    Rhs: Zero + Clone,
    <Lhs as Sub<Rhs>>::Output: Zero,
{
    type Output = Polynomial<<Lhs as Sub<Rhs>>::Output>;

    fn sub(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as Sub<Rhs>>::Output> {
        let min_len = cmp::min(self.data.len(), other.data.len());
        let max_len = cmp::max(self.data.len(), other.data.len());

        let mut sub = Vec::with_capacity(max_len);
        for i in 0..min_len {
            sub.push(self.data[i].clone() - other.data[i].clone());
        }
        if self.data.len() <= other.data.len() {
            for i in min_len..max_len {
                sub.push(num_traits::zero::<Lhs>() - other.data[i].clone())
            }
        } else {
            for i in min_len..max_len {
                sub.push(self.data[i].clone() - num_traits::zero::<Rhs>())
            }
        }
        Polynomial::new(sub)
    }
}

forward_all_binop!(impl Mul, mul);

impl<'a, 'b, Lhs, Rhs> Mul<&'b Polynomial<Rhs>> for &'a Polynomial<Lhs>
where
    Lhs: Zero + Mul<Rhs> + Clone,
    Rhs: Zero + Clone,
    <Lhs as Mul<Rhs>>::Output: Zero,
{
    type Output = Polynomial<<Lhs as Mul<Rhs>>::Output>;

    fn mul(self, other: &Polynomial<Rhs>) -> Polynomial<<Lhs as Mul<Rhs>>::Output> {
        if self.is_zero() || other.is_zero() {
            return Polynomial::new(vec![]);
        }

        let slen = self.data.len();
        let olen = other.data.len();
        let prod = (0..slen + olen - 1)
            .map(|i| {
                let mut p = num_traits::zero::<<Lhs as Mul<Rhs>>::Output>();
                let kstart = cmp::max(olen, i + 1) - olen;
                let kend = cmp::min(slen, i + 1);
                for k in kstart..kend {
                    p = p + self.data[k].clone() * other.data[i - k].clone();
                }
                p
            })
            .collect();
        Polynomial::new(prod)
    }
}

impl<T: Zero + Clone> Zero for Polynomial<T> {
    #[inline]
    fn zero() -> Polynomial<T> {
        Polynomial { data: vec![] }
    }
    #[inline]
    fn is_zero(&self) -> bool {
        self.data.is_empty()
    }
}

impl<T: Zero + One + Clone> One for Polynomial<T> {
    #[inline]
    fn one() -> Polynomial<T> {
        Polynomial {
            data: vec![One::one()],
        }
    }
}

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

    #[test]
    fn new() {
        fn check(dst: Vec<i32>, src: Vec<i32>) {
            assert_eq!(dst, Polynomial::new(src).data);
        }
        check(vec![1, 2, 3], vec![1, 2, 3]);
        check(vec![1, 2, 3], vec![1, 2, 3, 0, 0]);
        check(vec![], vec![0, 0, 0]);
    }

    #[test]
    fn neg_add_sub() {
        fn check(a: &[i32], b: &[i32], c: &[i32]) {
            fn check_eq(a: &Polynomial<i32>, b: &Polynomial<i32>) {
                assert_eq!(*a, *b);
                assert_eq!(-a, -b);
            }
            fn check_add(sum: &Polynomial<i32>, a: &Polynomial<i32>, b: &Polynomial<i32>) {
                check_eq(sum, &(a + b));
                check_eq(sum, &(b + a));
            }
            fn check_sub(sum: &Polynomial<i32>, a: &Polynomial<i32>, b: &Polynomial<i32>) {
                check_eq(a, &(sum - b));
                check_eq(b, &(sum - a));
            }

            let a = &Polynomial::new(a.to_vec());
            let b = &Polynomial::new(b.to_vec());
            let c = &Polynomial::new(c.to_vec());
            check_add(c, a, b);
            check_add(&(-c), &(-a), &(-b));
            check_sub(c, a, b);
            check_sub(&(-c), &(-a), &(-b));
        }
        check(&[], &[], &[]);
        check(&[], &[1], &[1]);
        check(&[1], &[1], &[2]);
        check(&[1, 0, 1], &[1], &[2, 0, 1]);
        check(&[1, 0, -1], &[-1, 0, 1], &[]);
    }

    #[test]
    fn mul() {
        fn check(a: &[i32], b: &[i32], c: &[i32]) {
            let a = Polynomial::new(a.to_vec());
            let b = Polynomial::new(b.to_vec());
            let c = Polynomial::new(c.to_vec());
            assert_eq!(c, &a * &b);
            assert_eq!(c, &b * &a);
        }
        check(&[], &[], &[]);
        check(&[0, 0], &[], &[]);
        check(&[0, 0], &[1], &[]);
        check(&[1, 0], &[1], &[1]);
        check(&[1, 0, 1], &[1], &[1, 0, 1]);
        check(&[1, 1], &[1, 1], &[1, 2, 1]);
        check(&[1, 1], &[1, 0, 1], &[1, 1, 1, 1]);
        check(&[0, 0, 1], &[0, 0, 1], &[0, 0, 0, 0, 1]);
    }

    #[test]
    fn eval() {
        fn check<F: Fn(i32) -> i32>(pol: &[i32], f: F) {
            for n in -10..10 {
                assert_eq!(f(n), Polynomial::new(pol.to_vec()).eval(n));
            }
        }
        check(&[], |_x| 0);
        check(&[1], |_x| 1);
        check(&[1, 1], |x| x + 1);
        check(&[0, 1], |x| x);
        check(&[10, -10, 10], |x| 10 * x * x - 10 * x + 10);
    }

    #[test]
    fn pretty() {
        fn check(slice: &[i32], s: &str) {
            assert_eq!(s.to_string(), Polynomial::new(slice.to_vec()).pretty("x"));
        }
        check(&[0], "0");
        check(&[1], "1");
        check(&[1, 1], "1+x");
        check(&[1, 1, 1], "1+x+x^2");
        check(&[2, 2, 2], "2+2*x+2*x^2");
        check(&[0, 0, 0, 1], "x^3");
        check(&[0, 0, 0, -1], "-x^3");
        check(&[-1, 0, 0, -1], "-1-x^3");
        check(&[-1, 1, 0, -1], "-1+x-x^3");
        check(&[-1, 1, -1, -1], "-1+x-x^2-x^3");
    }
}
<|endoftext|>