<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)]
#![deny(missing_debug_implementations)]

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,
    stack_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,
}

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

impl fmt::Debug for Builder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Builder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .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 because `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, after_start: Option<Arc<Fn() + Send + Sync>>, before_stop: Option<Arc<Fn() + Send + Sync>>) {
        after_start.map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        before_stop.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(),
            stack_size: 0,
            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 stack size of threads in the pool.
    pub fn stack_size(&mut self, stack_size: usize) -> &mut Self {
        self.stack_size = stack_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.
    /// The `f` will be deconstructed after the `builder` is deconstructed
    /// and all threads in the pool has executed it.
    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.
    /// The `f` will be deconstructed after the `builder` is deconstructed
    /// and all threads in the pool has executed it.
    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,
            }),
        };
        assert!(self.pool_size > 0);

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

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

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

        // After Builder is deconstructed, the tx should be droped
        // so that we can use rx as an iterator.
        let count = rx.into_iter().count();
        assert_eq!(count, 2);
    }
}
<commit_msg>Better error message when polling CpuFuture twice<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)]
#![deny(missing_debug_implementations)]

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,
    stack_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,
}

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

impl fmt::Debug for Builder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Builder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .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 because `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, after_start: Option<Arc<Fn() + Send + Sync>>, before_stop: Option<Arc<Fn() + Send + Sync>>) {
        after_start.map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        before_stop.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("cannot poll CpuFuture twice") {
            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(),
            stack_size: 0,
            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 stack size of threads in the pool.
    pub fn stack_size(&mut self, stack_size: usize) -> &mut Self {
        self.stack_size = stack_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.
    /// The `f` will be deconstructed after the `builder` is deconstructed
    /// and all threads in the pool has executed it.
    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.
    /// The `f` will be deconstructed after the `builder` is deconstructed
    /// and all threads in the pool has executed it.
    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,
            }),
        };
        assert!(self.pool_size > 0);

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

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

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

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


<commit_before>use fs::fields as f;
use output::colours::Colours;
use output::cell::{TextCell, DisplayWidth};
use ansi_term::Style;

impl f::Permissions {
    pub fn render(&self, colours: &Colours, file_type: f::Type, xattrs: bool) -> TextCell {
       let bit = |bit, chr: &'static str, style: Style| {
           if bit { style.paint(chr) } else { colours.punctuation.paint("-") }
       };

       let type_char = match file_type {
           f::Type::File        => colours.filetypes.normal.paint("."),
           f::Type::Directory   => colours.filetypes.directory.paint("d"),
           f::Type::Pipe        => colours.filetypes.pipe.paint("|"),
           f::Type::Link        => colours.filetypes.symlink.paint("l"),
           f::Type::CharDevice  => colours.filetypes.device.paint("c"),
           f::Type::BlockDevice => colours.filetypes.device.paint("b"),
           f::Type::Socket      => colours.filetypes.socket.paint("s"),
           f::Type::Special     => colours.filetypes.special.paint("?"),
       };

       let x_colour = if file_type.is_regular_file() { colours.perms.user_execute_file }
                                                else { colours.perms.user_execute_other };

       let mut chars = vec![
           type_char,
           bit(self.user_read,     "r", colours.perms.user_read),
           bit(self.user_write,    "w", colours.perms.user_write),
           bit(self.user_execute,  "x", x_colour),
           bit(self.group_read,    "r", colours.perms.group_read),
           bit(self.group_write,   "w", colours.perms.group_write),
           bit(self.group_execute, "x", colours.perms.group_execute),
           bit(self.other_read,    "r", colours.perms.other_read),
           bit(self.other_write,   "w", colours.perms.other_write),
           bit(self.other_execute, "x", colours.perms.other_execute),
       ];

       if xattrs {
           chars.push(colours.perms.attribute.paint("@"));
       }

       // As these are all ASCII characters, we can guarantee that they’re
       // all going to be one character wide, and don’t need to compute the
       // cell’s display width.
       let width = DisplayWidth::from(chars.len());

       TextCell {
           contents: chars.into(),
           width:    width,
       }
   }
}

#[cfg(test)]
#[allow(unused_results)]
pub mod test {
    use output::details::Details;

    use fs::fields as f;
    use output::cell::TextCell;

    use users::{User, Group};
    use users::mock::MockUsers;
    use users::os::unix::GroupExt;
    use ansi_term::Colour::*;

    #[test]
    fn named() {
        let mut details = Details::default();
        details.colours.users.group_not_yours = Fixed(101).normal();

        let mut users = MockUsers::with_current_uid(1000);
        users.add_group(Group::new(100, "folk"));

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(101).normal(), "folk");
        assert_eq!(expected, group.render(&details.colours, &users))
    }

    #[test]
    fn unnamed() {
        let mut details = Details::default();
        details.colours.users.group_not_yours = Fixed(87).normal();

        let users = MockUsers::with_current_uid(1000);

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(87).normal(), "100");
        assert_eq!(expected, group.render(&details.colours, &users));
    }

    #[test]
    fn primary() {
        let mut details = Details::default();
        details.colours.users.group_yours = Fixed(64).normal();

        let mut users = MockUsers::with_current_uid(2);
        users.add_user(User::new(2, "eve", 100));
        users.add_group(Group::new(100, "folk"));

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(64).normal(), "folk");
        assert_eq!(expected, group.render(&details.colours, &users))
    }

    #[test]
    fn secondary() {
        let mut details = Details::default();
        details.colours.users.group_yours = Fixed(31).normal();

        let mut users = MockUsers::with_current_uid(2);
        users.add_user(User::new(2, "eve", 666));

        let test_group = Group::new(100, "folk").add_member("eve");
        users.add_group(test_group);

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(31).normal(), "folk");
        assert_eq!(expected, group.render(&details.colours, &users))
    }

    #[test]
    fn overflow() {
        let mut details = Details::default();
        details.colours.users.group_not_yours = Blue.underline();

        let group = f::Group(2_147_483_648);
        let expected = TextCell::paint_str(Blue.underline(), "2147483648");
        assert_eq!(expected, group.render(&details.colours, &MockUsers::with_current_uid(0)));
    }
}
<commit_msg>Extract file type render function<commit_after>use fs::fields as f;
use output::colours::Colours;
use output::cell::{TextCell, DisplayWidth};
use ansi_term::{ANSIString, Style};

impl f::Permissions {
    pub fn render(&self, colours: &Colours, file_type: f::Type, xattrs: bool) -> TextCell {
       let bit = |bit, chr: &'static str, style: Style| {
           if bit { style.paint(chr) } else { colours.punctuation.paint("-") }
       };

       let x_colour = if file_type.is_regular_file() { colours.perms.user_execute_file }
                                                else { colours.perms.user_execute_other };

       let mut chars = vec![
           file_type.render(colours),
           bit(self.user_read,     "r", colours.perms.user_read),
           bit(self.user_write,    "w", colours.perms.user_write),
           bit(self.user_execute,  "x", x_colour),
           bit(self.group_read,    "r", colours.perms.group_read),
           bit(self.group_write,   "w", colours.perms.group_write),
           bit(self.group_execute, "x", colours.perms.group_execute),
           bit(self.other_read,    "r", colours.perms.other_read),
           bit(self.other_write,   "w", colours.perms.other_write),
           bit(self.other_execute, "x", colours.perms.other_execute),
       ];

       if xattrs {
           chars.push(colours.perms.attribute.paint("@"));
       }

       // As these are all ASCII characters, we can guarantee that they’re
       // all going to be one character wide, and don’t need to compute the
       // cell’s display width.
       let width = DisplayWidth::from(chars.len());

       TextCell {
           contents: chars.into(),
           width:    width,
       }
   }
}

impl f::Type {
    pub fn render(&self, colours: &Colours) -> ANSIString<'static> {
        match *self {
            f::Type::File        => colours.filetypes.normal.paint("."),
            f::Type::Directory   => colours.filetypes.directory.paint("d"),
            f::Type::Pipe        => colours.filetypes.pipe.paint("|"),
            f::Type::Link        => colours.filetypes.symlink.paint("l"),
            f::Type::CharDevice  => colours.filetypes.device.paint("c"),
            f::Type::BlockDevice => colours.filetypes.device.paint("b"),
            f::Type::Socket      => colours.filetypes.socket.paint("s"),
            f::Type::Special     => colours.filetypes.special.paint("?"),
        }
    }
}

#[cfg(test)]
#[allow(unused_results)]
pub mod test {
    use output::details::Details;

    use fs::fields as f;
    use output::cell::TextCell;

    use users::{User, Group};
    use users::mock::MockUsers;
    use users::os::unix::GroupExt;
    use ansi_term::Colour::*;

    #[test]
    fn named() {
        let mut details = Details::default();
        details.colours.users.group_not_yours = Fixed(101).normal();

        let mut users = MockUsers::with_current_uid(1000);
        users.add_group(Group::new(100, "folk"));

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(101).normal(), "folk");
        assert_eq!(expected, group.render(&details.colours, &users))
    }

    #[test]
    fn unnamed() {
        let mut details = Details::default();
        details.colours.users.group_not_yours = Fixed(87).normal();

        let users = MockUsers::with_current_uid(1000);

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(87).normal(), "100");
        assert_eq!(expected, group.render(&details.colours, &users));
    }

    #[test]
    fn primary() {
        let mut details = Details::default();
        details.colours.users.group_yours = Fixed(64).normal();

        let mut users = MockUsers::with_current_uid(2);
        users.add_user(User::new(2, "eve", 100));
        users.add_group(Group::new(100, "folk"));

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(64).normal(), "folk");
        assert_eq!(expected, group.render(&details.colours, &users))
    }

    #[test]
    fn secondary() {
        let mut details = Details::default();
        details.colours.users.group_yours = Fixed(31).normal();

        let mut users = MockUsers::with_current_uid(2);
        users.add_user(User::new(2, "eve", 666));

        let test_group = Group::new(100, "folk").add_member("eve");
        users.add_group(test_group);

        let group = f::Group(100);
        let expected = TextCell::paint_str(Fixed(31).normal(), "folk");
        assert_eq!(expected, group.render(&details.colours, &users))
    }

    #[test]
    fn overflow() {
        let mut details = Details::default();
        details.colours.users.group_not_yours = Blue.underline();

        let group = f::Group(2_147_483_648);
        let expected = TextCell::paint_str(Blue.underline(), "2147483648");
        assert_eq!(expected, group.render(&details.colours, &MockUsers::with_current_uid(0)));
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;
extern crate orbimage;
extern crate orbfont;

pub const ICON_SIZE: i32 = 48;
pub const ICON_SMALL_SIZE: i32 = 32;

#[cfg(target_os = "redox")]
static UI_PATH: &'static str = "/ui";

#[cfg(not(target_os = "redox"))]
static UI_PATH: &'static str = "ui";

#[cfg(not(target_os = "redox"))]
fn wait(status: &mut i32) -> usize {
    extern crate libc;

    use std::io::Error;

    let pid = unsafe { libc::waitpid(0, status as *mut i32, libc::WNOHANG) };
    if pid < 0 {
        panic!("waitpid failed: {}", Error::last_os_error());
    }
    pid as usize
}

#[cfg(target_os = "redox")]
fn wait(status: &mut usize) -> usize {
    extern crate syscall;

    syscall::waitpid(0, status, syscall::WNOHANG).unwrap()
}

fn load_icon(path: &str) -> Image {
    let icon = Image::from_path(path).unwrap_or(Image::default());
    if icon.width() == ICON_SIZE as u32 && icon.height() == ICON_SIZE as u32 {
        icon
    } else {
        icon.resize(ICON_SIZE as u32, ICON_SIZE as u32, orbimage::ResizeType::Lanczos3).unwrap()
    }
}

fn load_icon_small(path: &str) -> Image {
    let icon = Image::from_path(path).unwrap_or(Image::default());
    if icon.width() == ICON_SMALL_SIZE as u32 && icon.height() == ICON_SMALL_SIZE as u32 {
        icon
    } else {
        icon.resize(ICON_SMALL_SIZE as u32, ICON_SMALL_SIZE as u32, orbimage::ResizeType::Lanczos3).unwrap()
    }
}

use std::env;
use std::os::unix::process::ExitStatusExt;
use std::path::Path;
use std::process::{Command, ExitStatus};

use orbclient::{EventOption, Renderer, Window, K_ESC};
use orbimage::Image;
use orbfont::Font;

use package::Package;
use theme::{BAR_COLOR, BAR_HIGHLIGHT_COLOR, TEXT_COLOR, TEXT_HIGHLIGHT_COLOR};

mod package;
mod theme;

fn get_packages() -> Vec<Package> {
    let read_dir = Path::new(&format!("{}/apps/", UI_PATH)).read_dir().expect("failed to read apps directory");

    let mut entries = vec![];
    for dir in read_dir {
        let dir = match dir {
            Ok(x) => x,
            Err(_) => continue,
        };
        let file_name = dir.file_name().to_string_lossy().into_owned();
        if dir.file_type().expect("failed to get file_type").is_file() {
            entries.push(file_name);
        }
    }

    entries.sort();

    let mut packages: Vec<Package> = Vec::new();
    for entry in entries.iter() {
        packages.push(Package::from_path(&format!("{}/apps/{}", UI_PATH, entry)));
    }

    packages
}

fn draw(window: &mut Window, packages: &Vec<Package>, start: &Image, selected: i32){
    window.set(BAR_COLOR);

    let mut x = 0;
    let y = 0;
    let mut i = 0;

    {
        if i == selected {
            window.rect(x as i32, y as i32,
                              start.width() as u32, start.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        start.draw(window, x as i32, y as i32);

        x += start.width() as i32;
        i += 1;
    }

    for package in packages.iter() {
        if i == selected {
            window.rect(x as i32, y as i32,
                              package.icon.width() as u32, package.icon.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        package.icon.draw(window, x as i32, y as i32);

        x += package.icon.width() as i32;
        i += 1;
    }

    window.sync();
}

fn draw_chooser(window: &mut Window, font: &Font, packages: &Vec<Package>, selected: i32){
    let w = window.width();

    window.set(BAR_COLOR);

    let mut y = 0;
    for (i, package) in packages.iter().enumerate() {
        if i as i32 == selected {
            window.rect(0, y, w, ICON_SMALL_SIZE as u32, BAR_HIGHLIGHT_COLOR);
        }

        package.icon_small.draw(window, 0, y);

        font.render(&package.name, 16.0).draw(window, ICON_SMALL_SIZE + 8, y + 8, if i as i32 == selected { TEXT_HIGHLIGHT_COLOR } else { TEXT_COLOR });

        y += ICON_SMALL_SIZE;
    }

    window.sync();
}

fn bar_main() {
    let mut children = Vec::new();

    let packages = get_packages();

    let start = load_icon(&format!("{}/icons/places/start-here.png", UI_PATH));

    let mut logout_package = Package::new();
    logout_package.name = "Logout".to_string();
    logout_package.icon = load_icon(&format!("{}/icons/actions/system-log-out.png", UI_PATH));
    logout_package.icon_small = load_icon_small(&format!("{}/icons/actions/system-log-out.png", UI_PATH));
    logout_package.binary = "exit".to_string();

    let mut start_packages = packages.clone();
    start_packages.push(logout_package);

    let (width, height) = orbclient::get_display_size().expect("launcher: failed to get display size");
    let mut window = Window::new(0, height as i32 - ICON_SIZE, width, ICON_SIZE as u32, "").expect("launcher: failed to open window");

    let mut selected = -1;
    let mut mouse_x = 0;
    let mut mouse_y = 0;
    let mut mouse_left = false;
    let mut last_mouse_left = false;

    draw(&mut window, &packages, &start, selected);
    'running: loop {
        for event in window.events() {
            let redraw = match event.to_option() {
                EventOption::Mouse(mouse_event) => {
                    mouse_x = mouse_event.x;
                    mouse_y = mouse_event.y;
                    true
                },
                EventOption::Button(button_event) => {
                    mouse_left = button_event.left;
                    true
                },
                EventOption::Quit(_) => break 'running,
                _ => false
            };

            if redraw {
                let mut now_selected = -1;

                {
                    let mut x = 0;
                    let y = 0;
                    let mut i = 0;

                    {
                        if mouse_y >= y && mouse_x >= x &&
                           mouse_x < x + start.width() as i32 {
                               now_selected = i;
                        }
                        x += start.width() as i32;
                        i += 1;
                    }

                    for package in packages.iter() {
                        if mouse_y >= y && mouse_x >= x &&
                           mouse_x < x + package.icon.width() as i32 {
                            now_selected = i;
                        }
                        x += package.icon.width() as i32;
                        i += 1;
                    }
                }

                if now_selected != selected {
                    selected = now_selected;
                    draw(&mut window, &packages, &start, selected);
                }

                if ! mouse_left && last_mouse_left {
                    let mut i = 0;

                    if i == selected {
                        let start_h = start_packages.len() as u32 * ICON_SMALL_SIZE as u32;
                        let mut start_window = Window::new(0, height as i32 - ICON_SIZE - start_h as i32, 320, start_h, "").unwrap();
                        let font = Font::find(Some("Sans"), None, None).unwrap();

                        let mut selected = -1;
                        let mut mouse_y = 0;
                        let mut mouse_left = false;
                        let mut last_mouse_left = false;

                        draw_chooser(&mut start_window, &font, &start_packages, selected);
                        'start_choosing: loop {
                            for event in start_window.events() {
                                let redraw = match event.to_option() {
                                    EventOption::Mouse(mouse_event) => {
                                        mouse_y = mouse_event.y;
                                        true
                                    },
                                    EventOption::Button(button_event) => {
                                        mouse_left = button_event.left;
                                        true
                                    },
                                    EventOption::Key(key_event) => {
                                        match key_event.scancode {
                                            K_ESC => break 'start_choosing,
                                            _ => false
                                        }
                                    },
                                    EventOption::Focus(focus_event) => if ! focus_event.focused {
                                        break 'start_choosing;
                                    } else {
                                        false
                                    },
                                    EventOption::Quit(_) => break 'start_choosing,
                                    _ => false
                                };

                                if redraw {
                                    let mut now_selected = -1;

                                    let mut y = 0;
                                    for (j, _package) in start_packages.iter().enumerate() {
                                        if mouse_y >= y && mouse_y < y + ICON_SMALL_SIZE {
                                            now_selected = j as i32;
                                        }
                                        y += ICON_SMALL_SIZE;
                                    }

                                    if now_selected != selected {
                                        selected = now_selected;
                                        draw_chooser(&mut start_window, &font, &start_packages, selected);
                                    }

                                    if ! mouse_left && last_mouse_left {
                                        let mut y = 0;
                                        for package in start_packages.iter() {
                                            if mouse_y >= y && mouse_y < y + ICON_SMALL_SIZE {
                                                if package.binary == "exit" {
                                                    break 'running;
                                                } else {
                                                    match Command::new(&package.binary).spawn() {
                                                        Ok(child) => children.push(child),
                                                        Err(err) => println!("launcher: failed to launch {}: {}", package.binary, err)
                                                    }
                                                }
                                                break 'start_choosing;
                                            }
                                            y += ICON_SMALL_SIZE;
                                        }
                                    }

                                    last_mouse_left = mouse_left;
                                }
                            }
                        }
                    }
                    i += 1;

                    for package in packages.iter() {
                        if i == selected {
                            match Command::new(&package.binary).spawn() {
                                Ok(child) => children.push(child),
                                Err(err) => println!("launcher: failed to launch {}: {}", package.binary, err)
                            }
                        }
                        i += 1;
                    }
                }

                last_mouse_left = mouse_left;
            }
        }
    }

    for mut child in children {
        let pid = child.id();
        match child.kill() {
            Ok(()) => (),
            Err(err) => println!("launcher: failed to kill {}: {}", pid, err),
        }
        match child.wait() {
            Ok(status) => println!("launcher: {} exited with {}", pid, status),
            Err(err) => println!("launcher: failed to wait for {}: {}", pid, err),
        }
    }

    loop {
        let mut status = 0;
        let pid = wait(&mut status);
        if pid == 0 {
            break;
        } else {
            println!("launcher: reaping zombie {}: {}", pid, ExitStatus::from_raw(status as i32));
        }
    }
}

fn chooser_main(paths: env::Args) {
    for ref path in paths.skip(1) {
        let mut packages = get_packages();

        packages.retain(|package| -> bool {
            for accept in package.accepts.iter() {
                if (accept.starts_with('*') && path.ends_with(&accept[1 ..])) ||
                   (accept.ends_with('*') && path.starts_with(&accept[.. accept.len() - 1])) {
                    return true;
                }
            }
            false
        });

        if packages.len() > 1 {
            let mut window = Window::new(-1, -1, 320, packages.len() as u32 * ICON_SMALL_SIZE as u32, path).expect("launcher: failed to open window");
            let font = Font::find(Some("Sans"), None, None).expect("launcher: failed to open font");

            let mut selected = -1;
            let mut mouse_y = 0;
            let mut mouse_left = false;
            let mut last_mouse_left = false;

            draw_chooser(&mut window, &font, &packages, selected);
            'choosing: loop {
                for event in window.events() {
                    let redraw = match event.to_option() {
                        EventOption::Mouse(mouse_event) => {
                            mouse_y = mouse_event.y;
                            true
                        },
                        EventOption::Button(button_event) => {
                            mouse_left = button_event.left;
                            true
                        },
                        EventOption::Quit(_) => break 'choosing,
                        _ => false
                    };

                    if redraw {
                        let mut now_selected = -1;

                        let mut y = 0;
                        for (i, _package) in packages.iter().enumerate() {
                            if mouse_y >= y && mouse_y < y + ICON_SIZE {
                                now_selected = i as i32;
                            }
                            y += ICON_SMALL_SIZE;
                        }

                        if now_selected != selected {
                            selected = now_selected;
                            draw_chooser(&mut window, &font, &packages, selected);
                        }

                        if ! mouse_left && last_mouse_left {
                            let mut y = 0;
                            for package in packages.iter() {
                                if mouse_y >= y && mouse_y < y + ICON_SMALL_SIZE {
                                    if let Err(err) = Command::new(&package.binary).arg(path).spawn() {
                                        println!("launcher: failed to launch {}: {}", package.binary, err);
                                    }
                                    break 'choosing;
                                }
                                y += ICON_SMALL_SIZE;
                            }
                        }

                        last_mouse_left = mouse_left;
                    }
                }
            }
        } else if let Some(package) = packages.get(0) {
            if let Err(err) = Command::new(&package.binary).arg(&path).spawn() {
                println!("launcher: failed to launch '{}': {}", package.binary, err);
            }
        } else {
            println!("launcher: no application found for '{}'", path);
        }
    }
}

fn main() {
    let paths = env::args();
    if paths.len() > 1 {
        chooser_main(paths);
    } else {
        bar_main();
    }
}
<commit_msg>Handle screen size events<commit_after>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;
extern crate orbimage;
extern crate orbfont;

pub const ICON_SIZE: i32 = 48;
pub const ICON_SMALL_SIZE: i32 = 32;

#[cfg(target_os = "redox")]
static UI_PATH: &'static str = "/ui";

#[cfg(not(target_os = "redox"))]
static UI_PATH: &'static str = "ui";

#[cfg(not(target_os = "redox"))]
fn wait(status: &mut i32) -> usize {
    extern crate libc;

    use std::io::Error;

    let pid = unsafe { libc::waitpid(0, status as *mut i32, libc::WNOHANG) };
    if pid < 0 {
        panic!("waitpid failed: {}", Error::last_os_error());
    }
    pid as usize
}

#[cfg(target_os = "redox")]
fn wait(status: &mut usize) -> usize {
    extern crate syscall;

    syscall::waitpid(0, status, syscall::WNOHANG).unwrap()
}

fn load_icon(path: &str) -> Image {
    let icon = Image::from_path(path).unwrap_or(Image::default());
    if icon.width() == ICON_SIZE as u32 && icon.height() == ICON_SIZE as u32 {
        icon
    } else {
        icon.resize(ICON_SIZE as u32, ICON_SIZE as u32, orbimage::ResizeType::Lanczos3).unwrap()
    }
}

fn load_icon_small(path: &str) -> Image {
    let icon = Image::from_path(path).unwrap_or(Image::default());
    if icon.width() == ICON_SMALL_SIZE as u32 && icon.height() == ICON_SMALL_SIZE as u32 {
        icon
    } else {
        icon.resize(ICON_SMALL_SIZE as u32, ICON_SMALL_SIZE as u32, orbimage::ResizeType::Lanczos3).unwrap()
    }
}

use std::env;
use std::os::unix::process::ExitStatusExt;
use std::path::Path;
use std::process::{Command, ExitStatus};

use orbclient::{EventOption, Renderer, Window, K_ESC};
use orbimage::Image;
use orbfont::Font;

use package::Package;
use theme::{BAR_COLOR, BAR_HIGHLIGHT_COLOR, TEXT_COLOR, TEXT_HIGHLIGHT_COLOR};

mod package;
mod theme;

fn get_packages() -> Vec<Package> {
    let read_dir = Path::new(&format!("{}/apps/", UI_PATH)).read_dir().expect("failed to read apps directory");

    let mut entries = vec![];
    for dir in read_dir {
        let dir = match dir {
            Ok(x) => x,
            Err(_) => continue,
        };
        let file_name = dir.file_name().to_string_lossy().into_owned();
        if dir.file_type().expect("failed to get file_type").is_file() {
            entries.push(file_name);
        }
    }

    entries.sort();

    let mut packages: Vec<Package> = Vec::new();
    for entry in entries.iter() {
        packages.push(Package::from_path(&format!("{}/apps/{}", UI_PATH, entry)));
    }

    packages
}

fn draw(window: &mut Window, packages: &Vec<Package>, start: &Image, selected: i32){
    window.set(BAR_COLOR);

    let mut x = 0;
    let y = 0;
    let mut i = 0;

    {
        if i == selected {
            window.rect(x as i32, y as i32,
                              start.width() as u32, start.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        start.draw(window, x as i32, y as i32);

        x += start.width() as i32;
        i += 1;
    }

    for package in packages.iter() {
        if i == selected {
            window.rect(x as i32, y as i32,
                              package.icon.width() as u32, package.icon.height() as u32,
                              BAR_HIGHLIGHT_COLOR);
        }

        package.icon.draw(window, x as i32, y as i32);

        x += package.icon.width() as i32;
        i += 1;
    }

    window.sync();
}

fn draw_chooser(window: &mut Window, font: &Font, packages: &Vec<Package>, selected: i32){
    let w = window.width();

    window.set(BAR_COLOR);

    let mut y = 0;
    for (i, package) in packages.iter().enumerate() {
        if i as i32 == selected {
            window.rect(0, y, w, ICON_SMALL_SIZE as u32, BAR_HIGHLIGHT_COLOR);
        }

        package.icon_small.draw(window, 0, y);

        font.render(&package.name, 16.0).draw(window, ICON_SMALL_SIZE + 8, y + 8, if i as i32 == selected { TEXT_HIGHLIGHT_COLOR } else { TEXT_COLOR });

        y += ICON_SMALL_SIZE;
    }

    window.sync();
}

fn bar_main() {
    let mut children = Vec::new();

    let packages = get_packages();

    let start = load_icon(&format!("{}/icons/places/start-here.png", UI_PATH));

    let mut logout_package = Package::new();
    logout_package.name = "Logout".to_string();
    logout_package.icon = load_icon(&format!("{}/icons/actions/system-log-out.png", UI_PATH));
    logout_package.icon_small = load_icon_small(&format!("{}/icons/actions/system-log-out.png", UI_PATH));
    logout_package.binary = "exit".to_string();

    let mut start_packages = packages.clone();
    start_packages.push(logout_package);

    let (mut width, mut height) = orbclient::get_display_size().expect("launcher: failed to get display size");
    let mut window = Window::new(0, height as i32 - ICON_SIZE, width, ICON_SIZE as u32, "").expect("launcher: failed to open window");

    let mut selected = -1;
    let mut mouse_x = 0;
    let mut mouse_y = 0;
    let mut mouse_left = false;
    let mut last_mouse_left = false;

    draw(&mut window, &packages, &start, selected);
    'running: loop {
        for event in window.events() {
            let redraw = match event.to_option() {
                EventOption::Mouse(mouse_event) => {
                    mouse_x = mouse_event.x;
                    mouse_y = mouse_event.y;
                    true
                },
                EventOption::Button(button_event) => {
                    mouse_left = button_event.left;
                    true
                },
                EventOption::Screen(screen_event) => {
                    width = screen_event.width;
                    height = screen_event.height;
                    window.set_pos(0, height as i32 - ICON_SIZE);
                    window.set_size(width, ICON_SIZE as u32);
                    true
                },
                EventOption::Quit(_) => break 'running,
                _ => false
            };

            if redraw {
                let mut now_selected = -1;

                {
                    let mut x = 0;
                    let y = 0;
                    let mut i = 0;

                    {
                        if mouse_y >= y && mouse_x >= x &&
                           mouse_x < x + start.width() as i32 {
                               now_selected = i;
                        }
                        x += start.width() as i32;
                        i += 1;
                    }

                    for package in packages.iter() {
                        if mouse_y >= y && mouse_x >= x &&
                           mouse_x < x + package.icon.width() as i32 {
                            now_selected = i;
                        }
                        x += package.icon.width() as i32;
                        i += 1;
                    }
                }

                if now_selected != selected {
                    selected = now_selected;
                    draw(&mut window, &packages, &start, selected);
                }

                if ! mouse_left && last_mouse_left {
                    let mut i = 0;

                    if i == selected {
                        let start_h = start_packages.len() as u32 * ICON_SMALL_SIZE as u32;
                        let mut start_window = Window::new(0, height as i32 - ICON_SIZE - start_h as i32, 320, start_h, "").unwrap();
                        let font = Font::find(Some("Sans"), None, None).unwrap();

                        let mut selected = -1;
                        let mut mouse_y = 0;
                        let mut mouse_left = false;
                        let mut last_mouse_left = false;

                        draw_chooser(&mut start_window, &font, &start_packages, selected);
                        'start_choosing: loop {
                            for event in start_window.events() {
                                let redraw = match event.to_option() {
                                    EventOption::Mouse(mouse_event) => {
                                        mouse_y = mouse_event.y;
                                        true
                                    },
                                    EventOption::Button(button_event) => {
                                        mouse_left = button_event.left;
                                        true
                                    },
                                    EventOption::Key(key_event) => {
                                        match key_event.scancode {
                                            K_ESC => break 'start_choosing,
                                            _ => false
                                        }
                                    },
                                    EventOption::Focus(focus_event) => if ! focus_event.focused {
                                        break 'start_choosing;
                                    } else {
                                        false
                                    },
                                    EventOption::Quit(_) => break 'start_choosing,
                                    _ => false
                                };

                                if redraw {
                                    let mut now_selected = -1;

                                    let mut y = 0;
                                    for (j, _package) in start_packages.iter().enumerate() {
                                        if mouse_y >= y && mouse_y < y + ICON_SMALL_SIZE {
                                            now_selected = j as i32;
                                        }
                                        y += ICON_SMALL_SIZE;
                                    }

                                    if now_selected != selected {
                                        selected = now_selected;
                                        draw_chooser(&mut start_window, &font, &start_packages, selected);
                                    }

                                    if ! mouse_left && last_mouse_left {
                                        let mut y = 0;
                                        for package in start_packages.iter() {
                                            if mouse_y >= y && mouse_y < y + ICON_SMALL_SIZE {
                                                if package.binary == "exit" {
                                                    break 'running;
                                                } else {
                                                    match Command::new(&package.binary).spawn() {
                                                        Ok(child) => children.push(child),
                                                        Err(err) => println!("launcher: failed to launch {}: {}", package.binary, err)
                                                    }
                                                }
                                                break 'start_choosing;
                                            }
                                            y += ICON_SMALL_SIZE;
                                        }
                                    }

                                    last_mouse_left = mouse_left;
                                }
                            }
                        }
                    }
                    i += 1;

                    for package in packages.iter() {
                        if i == selected {
                            match Command::new(&package.binary).spawn() {
                                Ok(child) => children.push(child),
                                Err(err) => println!("launcher: failed to launch {}: {}", package.binary, err)
                            }
                        }
                        i += 1;
                    }
                }

                last_mouse_left = mouse_left;
            }
        }
    }

    for mut child in children {
        let pid = child.id();
        match child.kill() {
            Ok(()) => (),
            Err(err) => println!("launcher: failed to kill {}: {}", pid, err),
        }
        match child.wait() {
            Ok(status) => println!("launcher: {} exited with {}", pid, status),
            Err(err) => println!("launcher: failed to wait for {}: {}", pid, err),
        }
    }

    loop {
        let mut status = 0;
        let pid = wait(&mut status);
        if pid == 0 {
            break;
        } else {
            println!("launcher: reaping zombie {}: {}", pid, ExitStatus::from_raw(status as i32));
        }
    }
}

fn chooser_main(paths: env::Args) {
    for ref path in paths.skip(1) {
        let mut packages = get_packages();

        packages.retain(|package| -> bool {
            for accept in package.accepts.iter() {
                if (accept.starts_with('*') && path.ends_with(&accept[1 ..])) ||
                   (accept.ends_with('*') && path.starts_with(&accept[.. accept.len() - 1])) {
                    return true;
                }
            }
            false
        });

        if packages.len() > 1 {
            let mut window = Window::new(-1, -1, 320, packages.len() as u32 * ICON_SMALL_SIZE as u32, path).expect("launcher: failed to open window");
            let font = Font::find(Some("Sans"), None, None).expect("launcher: failed to open font");

            let mut selected = -1;
            let mut mouse_y = 0;
            let mut mouse_left = false;
            let mut last_mouse_left = false;

            draw_chooser(&mut window, &font, &packages, selected);
            'choosing: loop {
                for event in window.events() {
                    let redraw = match event.to_option() {
                        EventOption::Mouse(mouse_event) => {
                            mouse_y = mouse_event.y;
                            true
                        },
                        EventOption::Button(button_event) => {
                            mouse_left = button_event.left;
                            true
                        },
                        EventOption::Quit(_) => break 'choosing,
                        _ => false
                    };

                    if redraw {
                        let mut now_selected = -1;

                        let mut y = 0;
                        for (i, _package) in packages.iter().enumerate() {
                            if mouse_y >= y && mouse_y < y + ICON_SIZE {
                                now_selected = i as i32;
                            }
                            y += ICON_SMALL_SIZE;
                        }

                        if now_selected != selected {
                            selected = now_selected;
                            draw_chooser(&mut window, &font, &packages, selected);
                        }

                        if ! mouse_left && last_mouse_left {
                            let mut y = 0;
                            for package in packages.iter() {
                                if mouse_y >= y && mouse_y < y + ICON_SMALL_SIZE {
                                    if let Err(err) = Command::new(&package.binary).arg(path).spawn() {
                                        println!("launcher: failed to launch {}: {}", package.binary, err);
                                    }
                                    break 'choosing;
                                }
                                y += ICON_SMALL_SIZE;
                            }
                        }

                        last_mouse_left = mouse_left;
                    }
                }
            }
        } else if let Some(package) = packages.get(0) {
            if let Err(err) = Command::new(&package.binary).arg(&path).spawn() {
                println!("launcher: failed to launch '{}': {}", package.binary, err);
            }
        } else {
            println!("launcher: no application found for '{}'", path);
        }
    }
}

fn main() {
    let paths = env::args();
    if paths.len() > 1 {
        chooser_main(paths);
    } else {
        bar_main();
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;

use std::fs::File;
use std::io::{Read,Write};

// [Macro clap::clap_app](https://docs.rs/clap/2.33.0/clap/macro.clap_app.html)

fn gpu_freq_stat(tail: &str) -> String {
    format!("/sys/class/devfreq/gpufreq/{}", tail)
}

fn ddr_freq_stat(tail: &str) -> String {
    format!("/sys/class/devfreq/ddrfreq/{}", tail)
}

fn cpu_freq_stat(id: usize, tail: &str) -> String {
    format!("/sys/devices/system/cpu/cpu{}/cpufreq/{}", id, tail)
}

struct AvailableFrequencies
{
    gpu: Vec<u64>,
    ddr: Vec<u64>,
    cpu_little: Vec<u64>,
    cpu_medium: Vec<u64>,
    cpu_big: Vec<u64>,
}

struct ActualFrequencies
{
    gpu: u64,
    ddr: u64,
    cpu_little: u64,
    cpu_medium: u64,
    cpu_big: u64,
}

impl ActualFrequencies {
    fn cons() -> ActualFrequencies {
        ActualFrequencies { gpu: 0, ddr: 0, cpu_little: 0, cpu_medium: 0, cpu_big: 0, }
    }

    fn init(gpu: u64, ddr: u64, cpu_l: u64, cpu_m: u64, cpu_b: u64) -> ActualFrequencies {
        ActualFrequencies { gpu: gpu, ddr: ddr, cpu_little: cpu_l, cpu_medium: cpu_m, cpu_big: cpu_b, }
    }
}

fn read_sys_record(sys_path: &str) -> String {
    let mut read_buf = String::new();
    if let Ok(mut file) = File::open(sys_path) {
        let sz = file.read_to_string(&mut read_buf).expect("cannot read sysfs");
        assert!(sz > 0);
    }
    read_buf.trim_right().to_string()
}

fn read_sys_as_u64(sys_path: &str) -> u64 {
    let str_buf = read_sys_record(sys_path);
    str_buf.parse().expect("cannot parse to u64")
}

fn read_gpu_frequency() -> u64 {
    read_sys_as_u64(&gpu_freq_stat("cur_freq"))
}

fn main() {
    let matches = clap_app!(gf =>
        (version: "0.1")
        (author: "Alexander Samoilov <alexander.samoilov@gmail.com>")
        (about: "ARM Android cross-compile sample project")
        (@arg AVAILABLE: -a --available "display available frequencies")
        (@arg SHOW: -s --show "show current frequencies")
        (@arg SET_GPU: -g --gpu +takes_value "set GPU frequency")
        (@arg LOCK_FREQ: -l --lock ... +takes_value "lock GPU/DDR/CPU frequencies")
        (@arg VERBOSE: -v --verbose "be verbose")
    ).get_matches();

    if matches.is_present("VERBOSE") {
        println!("gf_freq command line args matches: {:?}", matches);
    }

    if matches.is_present("SHOW") {
        println!("show device current frequencies: gpu: {}", read_gpu_frequency());
    }

    if let Ok(gpu_freq) = value_t!(matches, "SET_GPU", u64) {
        println!("Set GPU frequency to: {}", gpu_freq);
    }

    if let Ok(freq_vec) = values_t!(matches, "LOCK_FREQ", u64) {
        println!("freq_vec: {:?}", freq_vec);
        for v in freq_vec {
            println!("Set frequency to: {}", v);
        }
    }

}
<commit_msg>read available freqs.<commit_after>#[macro_use]
extern crate clap;

use std::fs::File;
use std::io::{Read,Write};

// [Macro clap::clap_app](https://docs.rs/clap/2.33.0/clap/macro.clap_app.html)

fn gpu_freq_stat(tail: &str) -> String {
    format!("/sys/class/devfreq/gpufreq/{}", tail)
}

fn ddr_freq_stat(tail: &str) -> String {
    format!("/sys/class/devfreq/ddrfreq/{}", tail)
}

fn cpu_freq_stat(id: usize, tail: &str) -> String {
    format!("/sys/devices/system/cpu/cpu{}/cpufreq/{}", id, tail)
}

struct AvailableFrequencies
{
    gpu: Vec<u64>,
    ddr: Vec<u64>,
    cpu_little: Vec<u64>,
    cpu_medium: Vec<u64>,
    cpu_big: Vec<u64>,
}

struct ActualFrequencies
{
    gpu: u64,
    ddr: u64,
    cpu_little: u64,
    cpu_medium: u64,
    cpu_big: u64,
}

impl ActualFrequencies {
    fn cons() -> ActualFrequencies {
        ActualFrequencies { gpu: 0, ddr: 0, cpu_little: 0, cpu_medium: 0, cpu_big: 0, }
    }

    fn init(gpu: u64, ddr: u64, cpu_l: u64, cpu_m: u64, cpu_b: u64) -> ActualFrequencies {
        ActualFrequencies { gpu: gpu, ddr: ddr, cpu_little: cpu_l, cpu_medium: cpu_m, cpu_big: cpu_b, }
    }
}

fn read_sys_record(sys_path: &str) -> String {
    let mut read_buf = String::new();
    if let Ok(mut file) = File::open(sys_path) {
        let sz = file.read_to_string(&mut read_buf).expect("cannot read sysfs");
        assert!(sz > 0);
    }
    read_buf.trim().to_string()
}

fn read_sys_as_u64(sys_path: &str) -> u64 {
    let str_buf = read_sys_record(sys_path);
    str_buf.parse().expect("cannot parse to u64")
}

fn read_gpu_frequency() -> u64 {
    read_sys_as_u64(&gpu_freq_stat("cur_freq"))
}

fn get_device_available_frequencies() -> AvailableFrequencies {
    fn get_sys_freqs(sys_path: &str) -> Vec<u64> {
        let mut freqs: Vec<u64> = Vec::new();
        if let Ok(mut file) = File::open(sys_path) {
            let mut read_buf = String::new();
            let sz = file.read_to_string(&mut read_buf).expect("cannot read sysfs");
            assert!(sz > 0);
            let vec: Vec<&str> = read_buf.trim().split(' ').collect();
            freqs = vec.into_iter().map(|x| x.parse().expect("cannot parse to u64")).collect();
        }
        freqs
    }
    AvailableFrequencies { gpu: get_sys_freqs(&gpu_freq_stat("available_frequencies")),
                           ddr: Vec::new(), cpu_little: Vec::new(), cpu_medium: Vec::new(), cpu_big: Vec::new(), }
}

fn main() {
    let matches = clap_app!(gf =>
        (version: "0.1")
        (author: "Alexander Samoilov <alexander.samoilov@gmail.com>")
        (about: "ARM Android cross-compile sample project")
        (@arg AVAILABLE: -a --available "display available frequencies")
        (@arg SHOW: -s --show "show current frequencies")
        (@arg SET_GPU: -g --gpu +takes_value "set GPU frequency")
        (@arg LOCK_FREQ: -l --lock ... +takes_value "lock GPU/DDR/CPU frequencies")
        (@arg VERBOSE: -v --verbose "be verbose")
    ).get_matches();

    if matches.is_present("VERBOSE") {
        println!("gf_freq command line args matches: {:?}", matches);
    }

    if matches.is_present("AVAILABLE") {
        let avail = get_device_available_frequencies();
        println!("GPU available frequencies: {:?}", avail.gpu);
    }

    if matches.is_present("SHOW") {
        println!("show device current frequencies: gpu: {}", read_gpu_frequency());
    }

    if let Ok(gpu_freq) = value_t!(matches, "SET_GPU", u64) {
        println!("Set GPU frequency to: {}", gpu_freq);
    }

    if let Ok(freq_vec) = values_t!(matches, "LOCK_FREQ", u64) {
        println!("freq_vec: {:?}", freq_vec);
        for v in freq_vec {
            println!("Set frequency to: {}", v);
        }
    }

}
<|endoftext|>


<commit_before>use memory::MemSegment;
use super::PaletteIndex;
use super::PaletteSet;
use super::TilePattern;
use super::SCREEN_BUFFER_SIZE;
use super::SCREEN_WIDTH;
use super::SCREEN_HEIGHT;
use super::ppu_reg::PPUReg;
use super::ppu_memory::PPUMemory;
use std::cmp;

const TRANSPARENT: PaletteIndex = PaletteIndex {
    set: PaletteSet::Sprite,
    palette_id: 0,
    color_id: 0,
};

bitflags! {
    flags OAMAttr : u8 {
        const FLIP_VERT = 0b1000_0000,
        const FLIP_HORZ = 0b0100_0000,
        const BEHIND    = 0b0010_0000,
        const PALETTE1  = 0b0000_0010,
        const PALETTE2  = 0b0000_0001,
    }
}

impl OAMAttr {
    fn palette(&self) -> u8 {
        self.bits & 0x03
    }

    fn priority(&self) -> SpritePriority {
        if self.contains(BEHIND) {
            SpritePriority::Background
        } else {
            SpritePriority::Foreground
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct OAMEntry {
    y: u8,
    tile: u8,
    attr: OAMAttr,
    x: u8,
}

impl OAMEntry {
    fn is_on_scanline(&self, scanline: u16) -> bool {
        let y = self.y as u16;
        y <= scanline && scanline < y + 8
    }

    fn build_details(&self,
                     idx: usize,
                     sl: u16,
                     reg: &PPUReg,
                     mem: &mut PPUMemory)
                     -> SpriteDetails {
        let tile_id = self.tile;
        let fine_y_scroll = get_fine_scroll(sl, self.y as u16, self.attr.contains(FLIP_VERT));
        let tile_table = reg.ppuctrl.sprite_table();
        let tile = mem.read_tile_pattern(tile_id, fine_y_scroll, tile_table);
        SpriteDetails {
            idx: idx,
            x: self.x,
            attr: self.attr,
            tile: tile,
        }
    }
}

impl Default for OAMEntry {
    fn default() -> OAMEntry {
        OAMEntry {
            y: 0xFF,
            tile: 0,
            attr: OAMAttr::from_bits_truncate(0),
            x: 0xFF,
        }
    }
}

impl MemSegment for OAMEntry {
    fn read(&mut self, idx: u16) -> u8 {
        match idx {
            0 => self.y,
            1 => self.tile,
            2 => self.attr.bits(),
            3 => self.x,
            x => invalid_address!(x),
        }
    }

    fn write(&mut self, idx: u16, val: u8) {
        match idx {
            0 => self.y = val,
            1 => self.tile = val,
            2 => self.attr = OAMAttr::from_bits_truncate(val),
            3 => self.x = val,
            x => invalid_address!(x),
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct SpriteDetails {
    idx: usize,
    x: u8,
    attr: OAMAttr,
    tile: TilePattern,
}

impl SpriteDetails {
    fn do_get_pixel(&self, x: u16) -> (SpritePriority, PaletteIndex) {
        let fine_x = get_fine_scroll(x, self.x as u16, self.attr.contains(FLIP_HORZ));
        let attr = self.attr;
        let color_id = self.tile.get_color_in_pattern(fine_x as u32);
        let idx = PaletteIndex {
            set: PaletteSet::Sprite,
            palette_id: attr.palette(),
            color_id: color_id,
        };
        return (attr.priority(), idx);
    }

    fn blit(&self,
            pixel_line: &mut [PaletteIndex],
            priority_line: &mut [SpritePriority],
            sprite0_line: &mut [bool],
            segment: &Interval,
            sprite_interval: &Interval) {
        let intersection = segment.intersection(sprite_interval);
        for x in intersection.start..(intersection.end) {
            let (pri, pal) = self.do_get_pixel(x as u16);
            if !pal.is_transparent() {
                pixel_line[x] = pal;
                priority_line[x] = pri;
                sprite0_line[x] = self.idx == 0;
            }
        }
    }
}

impl Default for SpriteDetails {
    fn default() -> SpriteDetails {
        SpriteDetails {
            idx: 0xFF,
            x: 0xFF,
            attr: OAMAttr::empty(),
            tile: Default::default(),
        }
    }
}

#[derive(Debug)]
struct Interval {
    start: usize,
    end: usize,
}

impl Interval {
    fn new(start: usize, end: usize) -> Interval {
        Interval {
            start: start,
            end: end,
        }
    }

    fn intersects_with(&self, other: &Interval) -> bool {
        if self.start >= other.end {
            false
        } else if self.end <= other.start {
            false
        } else {
            true
        }
    }

    fn intersection(&self, other: &Interval) -> Interval {
        Interval {
            start: cmp::max(self.start, other.start),
            end: cmp::min(self.end, other.end),
        }
    }
}

pub struct SpriteRenderer {
    primary_oam: [OAMEntry; 64],
    secondary_oam: [[SpriteDetails; 8]; SCREEN_HEIGHT],

    pixel_buffer: Box<[PaletteIndex]>,
    priority_buffer: Box<[SpritePriority]>,
    sprite0_buffer: Box<[bool]>,
}

impl Default for SpriteRenderer {
    fn default() -> SpriteRenderer {
        SpriteRenderer {
            primary_oam: [Default::default(); 64],
            secondary_oam: [[Default::default(); 8]; SCREEN_HEIGHT],

            pixel_buffer: vec![Default::default(); SCREEN_BUFFER_SIZE].into_boxed_slice(),
            priority_buffer: vec![SpritePriority::Background; SCREEN_BUFFER_SIZE]
                                 .into_boxed_slice(),
            sprite0_buffer: vec![false; SCREEN_BUFFER_SIZE].into_boxed_slice(),
        }
    }
}

fn get_fine_scroll(screen_dist: u16, sprite_dist: u16, flip: bool) -> u16 {
    let scroll = screen_dist - sprite_dist;
    if flip {
        7 - scroll
    } else {
        scroll
    }
}

impl SpriteRenderer {
    pub fn render(&mut self, start: usize, stop: usize) {
        self.clear(start, stop);
        self.draw(start, stop)
    }

    pub fn run_cycle(&mut self, cyc: u16, sl: i16, reg: &mut PPUReg, mem: &mut PPUMemory) {
        match (cyc, sl) {
            (0, sl@0...239) => self.sprite_eval(sl as u16, reg, mem),
            _ => (),
        }
    }

    fn sprite_eval(&mut self, scanline: u16, reg: &PPUReg, mem: &mut PPUMemory) {
        if scanline + 1 >= SCREEN_HEIGHT as u16 {
            return;
        }
        let mut n = 0;
        let secondary_oam_line = &mut self.secondary_oam[scanline as usize + 1];
        *secondary_oam_line = [Default::default(); 8];
        for x in 0..64 {
            let oam = &self.primary_oam[x];
            if oam.is_on_scanline(scanline) {
                secondary_oam_line[n] = oam.build_details(x, scanline, reg, mem);
                n += 1;
                if n == 8 {
                    return;
                }
            }
        }
    }

    fn clear(&mut self, start: usize, stop: usize) {
        for dest in self.pixel_buffer[start..stop].iter_mut() {
            *dest = TRANSPARENT;
        }
        for dest in self.priority_buffer[start..stop].iter_mut() {
            *dest = SpritePriority::Background;
        }
        for dest in self.sprite0_buffer[start..stop].iter_mut() {
            *dest = false;
        }
    }

    fn draw(&mut self, start: usize, stop: usize) {
        let mut current_scanline = start / SCREEN_WIDTH;
        let mut last_scanline_boundary = current_scanline * SCREEN_WIDTH;
        let next_scanline = current_scanline + 1;
        let mut next_scanline_boundary = next_scanline * SCREEN_WIDTH;

        let mut current = start;
        while current < stop {
            let segment_start = current - last_scanline_boundary;
            let segment_end = cmp::min(next_scanline_boundary, stop) - last_scanline_boundary;

            self.render_segment(current_scanline,
                                last_scanline_boundary,
                                next_scanline_boundary,
                                segment_start,
                                segment_end);
            current_scanline += 1;
            last_scanline_boundary = next_scanline_boundary;
            current = next_scanline_boundary;
            next_scanline_boundary += SCREEN_WIDTH;
        }
    }

    fn render_segment(&mut self,
                      scanline: usize,
                      line_start: usize,
                      line_stop: usize,
                      start: usize,
                      stop: usize) {
        let oam_line = &self.secondary_oam[scanline];
        let pixel_line = &mut self.pixel_buffer[line_start..line_stop];
        let priority_line = &mut self.priority_buffer[line_start..line_stop];
        let sprite0_line = &mut self.sprite0_buffer[line_start..line_stop];

        let segment = Interval::new(start, stop - 1);

        for sprite in oam_line.iter().rev() {
            let sprite_interval = Interval::new(sprite.x as usize, sprite.x as usize + 8);
            if segment.intersects_with(&sprite_interval) {
                sprite.blit(pixel_line,
                            priority_line,
                            sprite0_line,
                            &segment,
                            &sprite_interval);
            }
        }
    }

    pub fn buffers(&self) -> (&[PaletteIndex], &[SpritePriority], &[bool]) {
        (&self.pixel_buffer,
         &self.priority_buffer,
         &self.sprite0_buffer)
    }
}

#[derive(Debug, Copy, Clone)]
pub enum SpritePriority {
    Foreground,
    Background,
}

///Reads the primary OAM table.
impl MemSegment for SpriteRenderer {
    fn read(&mut self, idx: u16) -> u8 {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].read(idx % 4)
    }

    fn write(&mut self, idx: u16, val: u8) {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].write(idx % 4, val)
    }
}

#[cfg(tests)]
mod tests {
    use memory::MemSegment;
    use ppu::PPU;
    use super::*;

    #[test]
    fn reading_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        for x in 0u8..63u8 {
            ppu.oam[x as usize] = OAMEntry::new(x, x, x, x);
        }
        ppu.reg.oamaddr = 0;
        assert_eq!(ppu.read(0x2004), 0);
        ppu.reg.oamaddr = 10;
        assert_eq!(ppu.read(0x2004), 2);
    }

    #[test]
    fn writing_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        ppu.reg.oamaddr = 0;
        ppu.write(0x2004, 12);
        assert_eq!(ppu.oam[0].y, 12);
        ppu.reg.oamaddr = 10;
        ppu.write(0x2004, 3);
        assert_eq!(ppu.oam[2].attr.bits(), 3);
    }

    #[test]
    fn test_sprite_on_scanline() {
        let mut ppu = create_test_ppu();
        let mut oam: OAMEntry = Default::default();
        oam.y = 10;

        assert!(!ppu.is_on_scanline(oam, 9));
        for sl in 10..18 {
            assert!(ppu.is_on_scanline(oam, sl));
        }
        assert!(!ppu.is_on_scanline(oam, 18));
    }
}
<commit_msg>Fix bug where some pixels of sprites weren't drawn.<commit_after>use memory::MemSegment;
use super::PaletteIndex;
use super::PaletteSet;
use super::TilePattern;
use super::SCREEN_BUFFER_SIZE;
use super::SCREEN_WIDTH;
use super::SCREEN_HEIGHT;
use super::ppu_reg::PPUReg;
use super::ppu_memory::PPUMemory;
use std::cmp;

const TRANSPARENT: PaletteIndex = PaletteIndex {
    set: PaletteSet::Sprite,
    palette_id: 0,
    color_id: 0,
};

bitflags! {
    flags OAMAttr : u8 {
        const FLIP_VERT = 0b1000_0000,
        const FLIP_HORZ = 0b0100_0000,
        const BEHIND    = 0b0010_0000,
        const PALETTE1  = 0b0000_0010,
        const PALETTE2  = 0b0000_0001,
    }
}

impl OAMAttr {
    fn palette(&self) -> u8 {
        self.bits & 0x03
    }

    fn priority(&self) -> SpritePriority {
        if self.contains(BEHIND) {
            SpritePriority::Background
        } else {
            SpritePriority::Foreground
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct OAMEntry {
    y: u8,
    tile: u8,
    attr: OAMAttr,
    x: u8,
}

impl OAMEntry {
    fn is_on_scanline(&self, scanline: u16) -> bool {
        let y = self.y as u16;
        y <= scanline && scanline < y + 8
    }

    fn build_details(&self,
                     idx: usize,
                     sl: u16,
                     reg: &PPUReg,
                     mem: &mut PPUMemory)
                     -> SpriteDetails {
        let tile_id = self.tile;
        let fine_y_scroll = get_fine_scroll(sl, self.y as u16, self.attr.contains(FLIP_VERT));
        let tile_table = reg.ppuctrl.sprite_table();
        let tile = mem.read_tile_pattern(tile_id, fine_y_scroll, tile_table);
        SpriteDetails {
            idx: idx,
            x: self.x,
            attr: self.attr,
            tile: tile,
        }
    }
}

impl Default for OAMEntry {
    fn default() -> OAMEntry {
        OAMEntry {
            y: 0xFF,
            tile: 0,
            attr: OAMAttr::from_bits_truncate(0),
            x: 0xFF,
        }
    }
}

impl MemSegment for OAMEntry {
    fn read(&mut self, idx: u16) -> u8 {
        match idx {
            0 => self.y,
            1 => self.tile,
            2 => self.attr.bits(),
            3 => self.x,
            x => invalid_address!(x),
        }
    }

    fn write(&mut self, idx: u16, val: u8) {
        match idx {
            0 => self.y = val,
            1 => self.tile = val,
            2 => self.attr = OAMAttr::from_bits_truncate(val),
            3 => self.x = val,
            x => invalid_address!(x),
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct SpriteDetails {
    idx: usize,
    x: u8,
    attr: OAMAttr,
    tile: TilePattern,
}

impl SpriteDetails {
    fn do_get_pixel(&self, x: u16) -> (SpritePriority, PaletteIndex) {
        let fine_x = get_fine_scroll(x, self.x as u16, self.attr.contains(FLIP_HORZ));
        let attr = self.attr;
        let color_id = self.tile.get_color_in_pattern(fine_x as u32);
        let idx = PaletteIndex {
            set: PaletteSet::Sprite,
            palette_id: attr.palette(),
            color_id: color_id,
        };
        return (attr.priority(), idx);
    }

    fn blit(&self,
            pixel_line: &mut [PaletteIndex],
            priority_line: &mut [SpritePriority],
            sprite0_line: &mut [bool],
            segment: &Interval,
            sprite_interval: &Interval) {
        let intersection = segment.intersection(sprite_interval);
        for x in intersection.start..(intersection.end) {
            let (pri, pal) = self.do_get_pixel(x as u16);
            if !pal.is_transparent() {
                pixel_line[x] = pal;
                priority_line[x] = pri;
                sprite0_line[x] = self.idx == 0;
            }
        }
    }
}

impl Default for SpriteDetails {
    fn default() -> SpriteDetails {
        SpriteDetails {
            idx: 0xFF,
            x: 0xFF,
            attr: OAMAttr::empty(),
            tile: Default::default(),
        }
    }
}

#[derive(Debug)]
struct Interval {
    start: usize,
    end: usize,
}

impl Interval {
    fn new(start: usize, end: usize) -> Interval {
        Interval {
            start: start,
            end: end,
        }
    }

    fn intersects_with(&self, other: &Interval) -> bool {
        if self.start >= other.end {
            false
        } else if self.end <= other.start {
            false
        } else {
            true
        }
    }

    fn intersection(&self, other: &Interval) -> Interval {
        Interval {
            start: cmp::max(self.start, other.start),
            end: cmp::min(self.end, other.end),
        }
    }
}

pub struct SpriteRenderer {
    primary_oam: [OAMEntry; 64],
    secondary_oam: [[SpriteDetails; 8]; SCREEN_HEIGHT],

    pixel_buffer: Box<[PaletteIndex]>,
    priority_buffer: Box<[SpritePriority]>,
    sprite0_buffer: Box<[bool]>,
}

impl Default for SpriteRenderer {
    fn default() -> SpriteRenderer {
        SpriteRenderer {
            primary_oam: [Default::default(); 64],
            secondary_oam: [[Default::default(); 8]; SCREEN_HEIGHT],

            pixel_buffer: vec![Default::default(); SCREEN_BUFFER_SIZE].into_boxed_slice(),
            priority_buffer: vec![SpritePriority::Background; SCREEN_BUFFER_SIZE]
                                 .into_boxed_slice(),
            sprite0_buffer: vec![false; SCREEN_BUFFER_SIZE].into_boxed_slice(),
        }
    }
}

fn get_fine_scroll(screen_dist: u16, sprite_dist: u16, flip: bool) -> u16 {
    let scroll = screen_dist - sprite_dist;
    if flip {
        7 - scroll
    } else {
        scroll
    }
}

impl SpriteRenderer {
    pub fn render(&mut self, start: usize, stop: usize) {
        self.clear(start, stop);
        self.draw(start, stop)
    }

    pub fn run_cycle(&mut self, cyc: u16, sl: i16, reg: &mut PPUReg, mem: &mut PPUMemory) {
        match (cyc, sl) {
            (0, sl@0...239) => self.sprite_eval(sl as u16, reg, mem),
            _ => (),
        }
    }

    fn sprite_eval(&mut self, scanline: u16, reg: &PPUReg, mem: &mut PPUMemory) {
        if scanline + 1 >= SCREEN_HEIGHT as u16 {
            return;
        }
        let mut n = 0;
        let secondary_oam_line = &mut self.secondary_oam[scanline as usize + 1];
        *secondary_oam_line = [Default::default(); 8];
        for x in 0..64 {
            let oam = &self.primary_oam[x];
            if oam.is_on_scanline(scanline) {
                secondary_oam_line[n] = oam.build_details(x, scanline, reg, mem);
                n += 1;
                if n == 8 {
                    return;
                }
            }
        }
    }

    fn clear(&mut self, start: usize, stop: usize) {
        for dest in self.pixel_buffer[start..stop].iter_mut() {
            *dest = TRANSPARENT;
        }
        for dest in self.priority_buffer[start..stop].iter_mut() {
            *dest = SpritePriority::Background;
        }
        for dest in self.sprite0_buffer[start..stop].iter_mut() {
            *dest = false;
        }
    }

    fn draw(&mut self, start: usize, stop: usize) {
        let mut current_scanline = start / SCREEN_WIDTH;
        let mut last_scanline_boundary = current_scanline * SCREEN_WIDTH;
        let next_scanline = current_scanline + 1;
        let mut next_scanline_boundary = next_scanline * SCREEN_WIDTH;

        let mut current = start;
        while current < stop {
            let segment_start = current - last_scanline_boundary;
            let segment_end = cmp::min(next_scanline_boundary, stop) - last_scanline_boundary;

            self.render_segment(current_scanline,
                                last_scanline_boundary,
                                next_scanline_boundary,
                                segment_start,
                                segment_end);
            current_scanline += 1;
            last_scanline_boundary = next_scanline_boundary;
            current = next_scanline_boundary;
            next_scanline_boundary += SCREEN_WIDTH;
        }
    }

    fn render_segment(&mut self,
                      scanline: usize,
                      line_start: usize,
                      line_stop: usize,
                      start: usize,
                      stop: usize) {
        let oam_line = &self.secondary_oam[scanline];
        let pixel_line = &mut self.pixel_buffer[line_start..line_stop];
        let priority_line = &mut self.priority_buffer[line_start..line_stop];
        let sprite0_line = &mut self.sprite0_buffer[line_start..line_stop];

        let segment = Interval::new(start, stop);

        for sprite in oam_line.iter().rev() {
            let sprite_interval = Interval::new(sprite.x as usize, sprite.x as usize + 8);
            if segment.intersects_with(&sprite_interval) {
                sprite.blit(pixel_line,
                            priority_line,
                            sprite0_line,
                            &segment,
                            &sprite_interval);
            }
        }
    }

    pub fn buffers(&self) -> (&[PaletteIndex], &[SpritePriority], &[bool]) {
        (&self.pixel_buffer,
         &self.priority_buffer,
         &self.sprite0_buffer)
    }
}

#[derive(Debug, Copy, Clone)]
pub enum SpritePriority {
    Foreground,
    Background,
}

///Reads the primary OAM table.
impl MemSegment for SpriteRenderer {
    fn read(&mut self, idx: u16) -> u8 {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].read(idx % 4)
    }

    fn write(&mut self, idx: u16, val: u8) {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].write(idx % 4, val)
    }
}

#[cfg(tests)]
mod tests {
    use memory::MemSegment;
    use ppu::PPU;
    use super::*;

    #[test]
    fn reading_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        for x in 0u8..63u8 {
            ppu.oam[x as usize] = OAMEntry::new(x, x, x, x);
        }
        ppu.reg.oamaddr = 0;
        assert_eq!(ppu.read(0x2004), 0);
        ppu.reg.oamaddr = 10;
        assert_eq!(ppu.read(0x2004), 2);
    }

    #[test]
    fn writing_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        ppu.reg.oamaddr = 0;
        ppu.write(0x2004, 12);
        assert_eq!(ppu.oam[0].y, 12);
        ppu.reg.oamaddr = 10;
        ppu.write(0x2004, 3);
        assert_eq!(ppu.oam[2].attr.bits(), 3);
    }

    #[test]
    fn test_sprite_on_scanline() {
        let mut ppu = create_test_ppu();
        let mut oam: OAMEntry = Default::default();
        oam.y = 10;

        assert!(!ppu.is_on_scanline(oam, 9));
        for sl in 10..18 {
            assert!(ppu.is_on_scanline(oam, sl));
        }
        assert!(!ppu.is_on_scanline(oam, 18));
    }
}
<|endoftext|>