<commit_before>use Configuration;
use deque;
use deque::{Worker, Stealer, Stolen};
use job::{JobRef, StackJob};
use latch::{Latch, CountLatch, LockLatch};
#[allow(unused_imports)]
use log::Event::*;
use rand::{self, Rng};
use sleep::Sleep;
use std::cell::{Cell, UnsafeCell};
use std::env;
use std::str::FromStr;
use std::sync::{Arc, Mutex, Once, ONCE_INIT};
use std::thread;
use std::mem;
use std::u32;
use std::usize;
use unwind;
use util::leak;
use num_cpus;

/// ////////////////////////////////////////////////////////////////////////

pub struct Registry {
    thread_infos: Vec<ThreadInfo>,
    state: Mutex<RegistryState>,
    sleep: Sleep,
    job_uninjector: Stealer<JobRef>,

    // When this latch reaches 0, it means that all work on this
    // registry must be complete. This is ensured in the following ways:
    //
    // - if this is the global registry, there is a ref-count that never
    //   gets released.
    // - if this is a user-created thread-pool, then so long as the thread-pool
    //   exists, it holds a reference.
    // - when we inject a "blocking job" into the registry with `ThreadPool::install()`,
    //   no adjustment is needed; the `ThreadPool` holds the reference, and since we won't
    //   return until the blocking job is complete, that ref will continue to be held.
    // - when `join()` or `scope()` is invoked, similarly, no adjustments are needed.
    //   These are always owned by some other job (e.g., one injected by `ThreadPool::install()`)
    //   and that job will keep the pool alive.
    terminate_latch: CountLatch,
}

struct RegistryState {
    job_injector: Worker<JobRef>,
}

/// ////////////////////////////////////////////////////////////////////////
/// Initialization

static mut THE_REGISTRY: Option<&'static Arc<Registry>> = None;
static THE_REGISTRY_SET: Once = ONCE_INIT;

/// Starts the worker threads (if that has not already happened). If
/// initialization has not already occurred, use the default
/// configuration.
pub fn global_registry() -> &'static Arc<Registry> {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(Configuration::new()) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Starts the worker threads (if that has not already happened) with
/// the given configuration.
pub fn get_registry_with_config(config: Configuration) -> &'static Registry {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(config) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Initializes the global registry with the given configuration.
/// Meant to be called from within the `THE_REGISTRY_SET` once
/// function. Declared `unsafe` because it writes to `THE_REGISTRY` in
/// an unsynchronized fashion.
unsafe fn init_registry(config: Configuration) {
    let registry = leak(Arc::new(Registry::new(config)));
    THE_REGISTRY = Some(registry);
}

impl Registry {
    pub fn new(mut configuration: Configuration) -> Arc<Registry> {
        let limit_value = match configuration.num_threads() {
            Some(value) => value,
            None => {
                match env::var("RAYON_RS_NUM_CPUS") {
                    Ok(s) => usize::from_str(&s).expect("invalid value for RAYON_RS_NUM_CPUS"),
                    Err(_) => num_cpus::get(),
                }
            }
        };

        let (inj_worker, inj_stealer) = deque::new();
        let (workers, stealers): (Vec<_>, Vec<_>) = (0..limit_value).map(|_| deque::new()).unzip();

        let registry = Arc::new(Registry {
            thread_infos: stealers.into_iter()
                .map(|s| ThreadInfo::new(s))
                .collect(),
            state: Mutex::new(RegistryState::new(inj_worker)),
            sleep: Sleep::new(),
            job_uninjector: inj_stealer,
            terminate_latch: CountLatch::new(),
        });

        for (index, worker) in workers.into_iter().enumerate() {
            let registry = registry.clone();
            if let Some(name) = configuration.base_thread_name(index) {
                let _ = thread::Builder::new()
                    .name(name)
                    .spawn(move || unsafe { main_loop(worker, registry, index) });
            } else {
                thread::spawn(move || unsafe { main_loop(worker, registry, index) });
            }
        }

        registry
    }

    /// Returns an opaque identifier for this registry.
    pub fn id(&self) -> RegistryId {
        // We can rely on `self` not to change since we only ever create
        // registries that are boxed up in an `Arc` (see `new()` above).
        RegistryId { addr: self as *const Self as usize }
    }

    pub fn num_threads(&self) -> usize {
        self.thread_infos.len()
    }

    /// Waits for the worker threads to get up and running.  This is
    /// meant to be used for benchmarking purposes, primarily, so that
    /// you can get more consistent numbers by having everything
    /// "ready to go".
    pub fn wait_until_primed(&self) {
        for info in &self.thread_infos {
            info.primed.wait();
        }
    }

    /// Waits for the worker threads to stop. This is used for testing
    /// -- so we can check that termination actually works.
    #[cfg(test)]
    pub fn wait_until_stopped(&self) {
        for info in &self.thread_infos {
            info.stopped.wait();
        }
    }

    /// ////////////////////////////////////////////////////////////////////////
    /// MAIN LOOP
    ///
    /// So long as all of the worker threads are hanging out in their
    /// top-level loop, there is no work to be done.

    pub unsafe fn inject(&self, injected_jobs: &[JobRef]) {
        log!(InjectJobs { count: injected_jobs.len() });
        {
            let state = self.state.lock().unwrap();

            // It should not be possible for `state.terminate` to be true
            // here. It is only set to true when the user creates (and
            // drops) a `ThreadPool`; and, in that case, they cannot be
            // calling `inject()` later, since they dropped their
            // `ThreadPool`.
            assert!(!self.terminate_latch.probe(), "inject() sees state.terminate as true");

            for &job_ref in injected_jobs {
                state.job_injector.push(job_ref);
            }
        }
        self.sleep.tickle(usize::MAX);
    }

    fn pop_injected_job(&self, worker_index: usize) -> Option<JobRef> {
        loop {
            match self.job_uninjector.steal() {
                Stolen::Empty => return None,
                Stolen::Abort => (), // retry
                Stolen::Data(v) => {
                    log!(UninjectedWork { worker: worker_index });
                    return Some(v);
                }
            }
        }
    }

    /// Signals that the thread-pool which owns this registry has been
    /// dropped. The worker threads will gradually terminate, once any
    /// extant work is completed.
    pub fn terminate(&self) {
        self.terminate_latch.set();
        self.sleep.tickle(usize::MAX);
    }
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct RegistryId {
    addr: usize
}

impl RegistryState {
    pub fn new(job_injector: Worker<JobRef>) -> RegistryState {
        RegistryState {
            job_injector: job_injector,
        }
    }
}

struct ThreadInfo {
    /// Latch set once thread has started and we are entering into the
    /// main loop. Used to wait for worker threads to become primed,
    /// primarily of interest for benchmarking.
    primed: LockLatch,

    /// Latch is set once worker thread has completed. Used to wait
    /// until workers have stopped; only used for tests.
    stopped: LockLatch,

    /// the "stealer" half of the worker's deque
    stealer: Stealer<JobRef>,
}

impl ThreadInfo {
    fn new(stealer: Stealer<JobRef>) -> ThreadInfo {
        ThreadInfo {
            primed: LockLatch::new(),
            stopped: LockLatch::new(),
            stealer: stealer,
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// WorkerThread identifiers

pub struct WorkerThread {
    worker: Worker<JobRef>,
    index: usize,

    /// A weak random number generator.
    rng: UnsafeCell<rand::XorShiftRng>,

    registry: Arc<Registry>,
}

// This is a bit sketchy, but basically: the WorkerThread is
// allocated on the stack of the worker on entry and stored into this
// thread local variable. So it will remain valid at least until the
// worker is fully unwound. Using an unsafe pointer avoids the need
// for a RefCell<T> etc.
thread_local! {
    static WORKER_THREAD_STATE: Cell<*const WorkerThread> =
        Cell::new(0 as *const WorkerThread)
}

impl WorkerThread {
    /// Gets the `WorkerThread` index for the current thread; returns
    /// NULL if this is not a worker thread. This pointer is valid
    /// anywhere on the current thread.
    #[inline]
    pub unsafe fn current() -> *const WorkerThread {
        WORKER_THREAD_STATE.with(|t| t.get())
    }

    /// Sets `self` as the worker thread index for the current thread.
    /// This is done during worker thread startup.
    unsafe fn set_current(thread: *const WorkerThread) {
        WORKER_THREAD_STATE.with(|t| {
            assert!(t.get().is_null());
            t.set(thread);
        });
    }

    /// Returns the registry that owns this worker thread.
    pub fn registry(&self) -> &Arc<Registry> {
        &self.registry
    }

    /// Our index amongst the worker threads (ranges from `0..self.num_threads()`).
    #[inline]
    pub fn index(&self) -> usize {
        self.index
    }

    #[inline]
    pub unsafe fn push(&self, job: JobRef) {
        self.worker.push(job);
        self.registry.sleep.tickle(self.index);
    }

    /// Pop `job` from top of stack, returning `false` if it has been
    /// stolen.
    #[inline]
    pub unsafe fn pop(&self) -> Option<JobRef> {
        self.worker.pop()
    }

    /// Wait until the latch is set. Try to keep busy by popping and
    /// stealing tasks as necessary.
    #[inline]
    pub unsafe fn wait_until<L: Latch>(&self, latch: &L) {
        log!(WaitUntil { worker: self.index });
        if !latch.probe() {
            self.wait_until_cold(latch);
        }
    }

    #[cold]
    unsafe fn wait_until_cold<L: Latch>(&self, latch: &L) {
        // the code below should swallow all panics and hence never
        // unwind; but if something does wrong, we want to abort,
        // because otherwise other code in rayon may assume that the
        // latch has been signaled, and that can lead to random memory
        // accesses, which would be *very bad*
        let abort_guard = unwind::AbortIfPanic;

        let mut yields = 0;
        while !latch.probe() {
            // Try to find some work to do. We give preference first
            // to things in our local deque, then in other workers
            // deques, and finally to injected jobs from the
            // outside. The idea is to finish what we started before
            // we take on something new.
            if let Some(job) = self.pop()
                                   .or_else(|| self.steal())
                                   .or_else(|| self.registry.pop_injected_job(self.index)) {
                yields = self.registry.sleep.work_found(self.index, yields);
                self.execute(job);
            } else {
                yields = self.registry.sleep.no_work_found(self.index, yields);
            }
        }

        // If we were sleepy, we are not anymore. We "found work" --
        // whatever the surrounding thread was doing before it had to
        // wait.
        self.registry.sleep.work_found(self.index, yields);

        log!(LatchSet { worker: self.index });
        mem::forget(abort_guard); // successful execution, do not abort
    }

    pub unsafe fn execute(&self, job: JobRef) {
        job.execute();

        // Subtle: executing this job will have `set()` some of its
        // latches.  This may mean that a sleepy (or sleeping) worker
        // can now make progress. So we have to tickle them to let
        // them know.
        self.registry.sleep.tickle(self.index);
    }

    /// Try to steal a single job and return it.
    ///
    /// This should only be done as a last resort, when there is no
    /// local work to do.
    unsafe fn steal(&self) -> Option<JobRef> {
        // we only steal when we don't have any work to do locally
        debug_assert!(self.worker.pop().is_none());

        // otherwise, try to steal
        let num_threads = self.registry.thread_infos.len();
        if num_threads <= 1 {
            return None;
        }
        assert!(num_threads < (u32::MAX as usize),
                "we do not support more than u32::MAX worker threads");

        let start = {
            // OK to use this UnsafeCell because (a) this data is
            // confined to current thread, as WorkerThread is not Send
            // nor Sync and (b) rand crate will not call back into
            // this method.
            let rng = &mut *self.rng.get();
            rng.next_u32() % num_threads as u32
        } as usize;
        (start .. num_threads)
            .chain(0 .. start)
            .filter(|&i| i != self.index)
            .filter_map(|victim_index| {
                let victim = &self.registry.thread_infos[victim_index];
                loop {
                    match victim.stealer.steal() {
                        Stolen::Empty => return None,
                        Stolen::Abort => (), // retry
                        Stolen::Data(v) => {
                            log!(StoleWork { worker: self.index, victim: victim_index });
                            return Some(v);
                        }
                    }
                }
            })
            .next()
    }
}

/// ////////////////////////////////////////////////////////////////////////

unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usize) {
    let worker_thread = WorkerThread {
        worker: worker,
        index: index,
        rng: UnsafeCell::new(rand::weak_rng()),
        registry: registry.clone(),
    };
    WorkerThread::set_current(&worker_thread);

    // let registry know we are ready to do work
    registry.thread_infos[index].primed.set();

    // Worker threads should not panic. If they do, just abort, as the
    // internal state of the threadpool is corrupted. Note that if
    // **user code** panics, we should catch that and redirect.
    let abort_guard = unwind::AbortIfPanic;

    worker_thread.wait_until(&registry.terminate_latch);

    // Should not be any work left in our queue.
    debug_assert!(worker_thread.pop().is_none());

    // let registry know we are done
    registry.thread_infos[index].stopped.set();

    // Normal termination, do not abort.
    mem::forget(abort_guard);
}

/// If already in a worker-thread, just execute `op`.  Otherwise,
/// execute `op` in the default thread-pool. Either way, block until
/// `op` completes and return its return value. If `op` panics, that
/// panic will be propagated as well.
pub fn in_worker<OP, R>(op: OP) -> R
    where OP: FnOnce(&WorkerThread) -> R + Send, R: Send
{
    unsafe {
        let owner_thread = WorkerThread::current();
        if !owner_thread.is_null() {
            // Perfectly valid to give them a `&T`: this is the
            // current thread, so we know the data structure won't be
            // invalidated until we return.
            return op(&*owner_thread);
        } else {
            return in_worker_cold(op);
        }
    }
}

#[cold]
unsafe fn in_worker_cold<OP, R>(op: OP) -> R
    where OP: FnOnce(&WorkerThread) -> R + Send, R: Send
{
    // never run from a worker thread; just shifts over into worker threads
    debug_assert!(WorkerThread::current().is_null());
    let job = StackJob::new(|| in_worker(op), LockLatch::new());
    global_registry().inject(&[job.as_job_ref()]);
    job.latch.wait();
    job.into_result()
}
<commit_msg>Make thread spawning DRYer<commit_after>use Configuration;
use deque;
use deque::{Worker, Stealer, Stolen};
use job::{JobRef, StackJob};
use latch::{Latch, CountLatch, LockLatch};
#[allow(unused_imports)]
use log::Event::*;
use rand::{self, Rng};
use sleep::Sleep;
use std::cell::{Cell, UnsafeCell};
use std::env;
use std::str::FromStr;
use std::sync::{Arc, Mutex, Once, ONCE_INIT};
use std::thread;
use std::mem;
use std::u32;
use std::usize;
use unwind;
use util::leak;
use num_cpus;

/// ////////////////////////////////////////////////////////////////////////

pub struct Registry {
    thread_infos: Vec<ThreadInfo>,
    state: Mutex<RegistryState>,
    sleep: Sleep,
    job_uninjector: Stealer<JobRef>,

    // When this latch reaches 0, it means that all work on this
    // registry must be complete. This is ensured in the following ways:
    //
    // - if this is the global registry, there is a ref-count that never
    //   gets released.
    // - if this is a user-created thread-pool, then so long as the thread-pool
    //   exists, it holds a reference.
    // - when we inject a "blocking job" into the registry with `ThreadPool::install()`,
    //   no adjustment is needed; the `ThreadPool` holds the reference, and since we won't
    //   return until the blocking job is complete, that ref will continue to be held.
    // - when `join()` or `scope()` is invoked, similarly, no adjustments are needed.
    //   These are always owned by some other job (e.g., one injected by `ThreadPool::install()`)
    //   and that job will keep the pool alive.
    terminate_latch: CountLatch,
}

struct RegistryState {
    job_injector: Worker<JobRef>,
}

/// ////////////////////////////////////////////////////////////////////////
/// Initialization

static mut THE_REGISTRY: Option<&'static Arc<Registry>> = None;
static THE_REGISTRY_SET: Once = ONCE_INIT;

/// Starts the worker threads (if that has not already happened). If
/// initialization has not already occurred, use the default
/// configuration.
pub fn global_registry() -> &'static Arc<Registry> {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(Configuration::new()) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Starts the worker threads (if that has not already happened) with
/// the given configuration.
pub fn get_registry_with_config(config: Configuration) -> &'static Registry {
    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(config) });
    unsafe { THE_REGISTRY.unwrap() }
}

/// Initializes the global registry with the given configuration.
/// Meant to be called from within the `THE_REGISTRY_SET` once
/// function. Declared `unsafe` because it writes to `THE_REGISTRY` in
/// an unsynchronized fashion.
unsafe fn init_registry(config: Configuration) {
    let registry = leak(Arc::new(Registry::new(config)));
    THE_REGISTRY = Some(registry);
}

impl Registry {
    pub fn new(mut configuration: Configuration) -> Arc<Registry> {
        let limit_value = match configuration.num_threads() {
            Some(value) => value,
            None => {
                match env::var("RAYON_RS_NUM_CPUS") {
                    Ok(s) => usize::from_str(&s).expect("invalid value for RAYON_RS_NUM_CPUS"),
                    Err(_) => num_cpus::get(),
                }
            }
        };

        let (inj_worker, inj_stealer) = deque::new();
        let (workers, stealers): (Vec<_>, Vec<_>) = (0..limit_value).map(|_| deque::new()).unzip();

        let registry = Arc::new(Registry {
            thread_infos: stealers.into_iter()
                .map(|s| ThreadInfo::new(s))
                .collect(),
            state: Mutex::new(RegistryState::new(inj_worker)),
            sleep: Sleep::new(),
            job_uninjector: inj_stealer,
            terminate_latch: CountLatch::new(),
        });

        for (index, worker) in workers.into_iter().enumerate() {
            let registry = registry.clone();
            let mut b = thread::Builder::new();
            if let Some(name) = configuration.base_thread_name(index) {
                b = b.name(name);
            }
            b.spawn(move || unsafe { main_loop(worker, registry, index) });
        }

        registry
    }

    /// Returns an opaque identifier for this registry.
    pub fn id(&self) -> RegistryId {
        // We can rely on `self` not to change since we only ever create
        // registries that are boxed up in an `Arc` (see `new()` above).
        RegistryId { addr: self as *const Self as usize }
    }

    pub fn num_threads(&self) -> usize {
        self.thread_infos.len()
    }

    /// Waits for the worker threads to get up and running.  This is
    /// meant to be used for benchmarking purposes, primarily, so that
    /// you can get more consistent numbers by having everything
    /// "ready to go".
    pub fn wait_until_primed(&self) {
        for info in &self.thread_infos {
            info.primed.wait();
        }
    }

    /// Waits for the worker threads to stop. This is used for testing
    /// -- so we can check that termination actually works.
    #[cfg(test)]
    pub fn wait_until_stopped(&self) {
        for info in &self.thread_infos {
            info.stopped.wait();
        }
    }

    /// ////////////////////////////////////////////////////////////////////////
    /// MAIN LOOP
    ///
    /// So long as all of the worker threads are hanging out in their
    /// top-level loop, there is no work to be done.

    pub unsafe fn inject(&self, injected_jobs: &[JobRef]) {
        log!(InjectJobs { count: injected_jobs.len() });
        {
            let state = self.state.lock().unwrap();

            // It should not be possible for `state.terminate` to be true
            // here. It is only set to true when the user creates (and
            // drops) a `ThreadPool`; and, in that case, they cannot be
            // calling `inject()` later, since they dropped their
            // `ThreadPool`.
            assert!(!self.terminate_latch.probe(), "inject() sees state.terminate as true");

            for &job_ref in injected_jobs {
                state.job_injector.push(job_ref);
            }
        }
        self.sleep.tickle(usize::MAX);
    }

    fn pop_injected_job(&self, worker_index: usize) -> Option<JobRef> {
        loop {
            match self.job_uninjector.steal() {
                Stolen::Empty => return None,
                Stolen::Abort => (), // retry
                Stolen::Data(v) => {
                    log!(UninjectedWork { worker: worker_index });
                    return Some(v);
                }
            }
        }
    }

    /// Signals that the thread-pool which owns this registry has been
    /// dropped. The worker threads will gradually terminate, once any
    /// extant work is completed.
    pub fn terminate(&self) {
        self.terminate_latch.set();
        self.sleep.tickle(usize::MAX);
    }
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct RegistryId {
    addr: usize
}

impl RegistryState {
    pub fn new(job_injector: Worker<JobRef>) -> RegistryState {
        RegistryState {
            job_injector: job_injector,
        }
    }
}

struct ThreadInfo {
    /// Latch set once thread has started and we are entering into the
    /// main loop. Used to wait for worker threads to become primed,
    /// primarily of interest for benchmarking.
    primed: LockLatch,

    /// Latch is set once worker thread has completed. Used to wait
    /// until workers have stopped; only used for tests.
    stopped: LockLatch,

    /// the "stealer" half of the worker's deque
    stealer: Stealer<JobRef>,
}

impl ThreadInfo {
    fn new(stealer: Stealer<JobRef>) -> ThreadInfo {
        ThreadInfo {
            primed: LockLatch::new(),
            stopped: LockLatch::new(),
            stealer: stealer,
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// WorkerThread identifiers

pub struct WorkerThread {
    worker: Worker<JobRef>,
    index: usize,

    /// A weak random number generator.
    rng: UnsafeCell<rand::XorShiftRng>,

    registry: Arc<Registry>,
}

// This is a bit sketchy, but basically: the WorkerThread is
// allocated on the stack of the worker on entry and stored into this
// thread local variable. So it will remain valid at least until the
// worker is fully unwound. Using an unsafe pointer avoids the need
// for a RefCell<T> etc.
thread_local! {
    static WORKER_THREAD_STATE: Cell<*const WorkerThread> =
        Cell::new(0 as *const WorkerThread)
}

impl WorkerThread {
    /// Gets the `WorkerThread` index for the current thread; returns
    /// NULL if this is not a worker thread. This pointer is valid
    /// anywhere on the current thread.
    #[inline]
    pub unsafe fn current() -> *const WorkerThread {
        WORKER_THREAD_STATE.with(|t| t.get())
    }

    /// Sets `self` as the worker thread index for the current thread.
    /// This is done during worker thread startup.
    unsafe fn set_current(thread: *const WorkerThread) {
        WORKER_THREAD_STATE.with(|t| {
            assert!(t.get().is_null());
            t.set(thread);
        });
    }

    /// Returns the registry that owns this worker thread.
    pub fn registry(&self) -> &Arc<Registry> {
        &self.registry
    }

    /// Our index amongst the worker threads (ranges from `0..self.num_threads()`).
    #[inline]
    pub fn index(&self) -> usize {
        self.index
    }

    #[inline]
    pub unsafe fn push(&self, job: JobRef) {
        self.worker.push(job);
        self.registry.sleep.tickle(self.index);
    }

    /// Pop `job` from top of stack, returning `false` if it has been
    /// stolen.
    #[inline]
    pub unsafe fn pop(&self) -> Option<JobRef> {
        self.worker.pop()
    }

    /// Wait until the latch is set. Try to keep busy by popping and
    /// stealing tasks as necessary.
    #[inline]
    pub unsafe fn wait_until<L: Latch>(&self, latch: &L) {
        log!(WaitUntil { worker: self.index });
        if !latch.probe() {
            self.wait_until_cold(latch);
        }
    }

    #[cold]
    unsafe fn wait_until_cold<L: Latch>(&self, latch: &L) {
        // the code below should swallow all panics and hence never
        // unwind; but if something does wrong, we want to abort,
        // because otherwise other code in rayon may assume that the
        // latch has been signaled, and that can lead to random memory
        // accesses, which would be *very bad*
        let abort_guard = unwind::AbortIfPanic;

        let mut yields = 0;
        while !latch.probe() {
            // Try to find some work to do. We give preference first
            // to things in our local deque, then in other workers
            // deques, and finally to injected jobs from the
            // outside. The idea is to finish what we started before
            // we take on something new.
            if let Some(job) = self.pop()
                                   .or_else(|| self.steal())
                                   .or_else(|| self.registry.pop_injected_job(self.index)) {
                yields = self.registry.sleep.work_found(self.index, yields);
                self.execute(job);
            } else {
                yields = self.registry.sleep.no_work_found(self.index, yields);
            }
        }

        // If we were sleepy, we are not anymore. We "found work" --
        // whatever the surrounding thread was doing before it had to
        // wait.
        self.registry.sleep.work_found(self.index, yields);

        log!(LatchSet { worker: self.index });
        mem::forget(abort_guard); // successful execution, do not abort
    }

    pub unsafe fn execute(&self, job: JobRef) {
        job.execute();

        // Subtle: executing this job will have `set()` some of its
        // latches.  This may mean that a sleepy (or sleeping) worker
        // can now make progress. So we have to tickle them to let
        // them know.
        self.registry.sleep.tickle(self.index);
    }

    /// Try to steal a single job and return it.
    ///
    /// This should only be done as a last resort, when there is no
    /// local work to do.
    unsafe fn steal(&self) -> Option<JobRef> {
        // we only steal when we don't have any work to do locally
        debug_assert!(self.worker.pop().is_none());

        // otherwise, try to steal
        let num_threads = self.registry.thread_infos.len();
        if num_threads <= 1 {
            return None;
        }
        assert!(num_threads < (u32::MAX as usize),
                "we do not support more than u32::MAX worker threads");

        let start = {
            // OK to use this UnsafeCell because (a) this data is
            // confined to current thread, as WorkerThread is not Send
            // nor Sync and (b) rand crate will not call back into
            // this method.
            let rng = &mut *self.rng.get();
            rng.next_u32() % num_threads as u32
        } as usize;
        (start .. num_threads)
            .chain(0 .. start)
            .filter(|&i| i != self.index)
            .filter_map(|victim_index| {
                let victim = &self.registry.thread_infos[victim_index];
                loop {
                    match victim.stealer.steal() {
                        Stolen::Empty => return None,
                        Stolen::Abort => (), // retry
                        Stolen::Data(v) => {
                            log!(StoleWork { worker: self.index, victim: victim_index });
                            return Some(v);
                        }
                    }
                }
            })
            .next()
    }
}

/// ////////////////////////////////////////////////////////////////////////

unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usize) {
    let worker_thread = WorkerThread {
        worker: worker,
        index: index,
        rng: UnsafeCell::new(rand::weak_rng()),
        registry: registry.clone(),
    };
    WorkerThread::set_current(&worker_thread);

    // let registry know we are ready to do work
    registry.thread_infos[index].primed.set();

    // Worker threads should not panic. If they do, just abort, as the
    // internal state of the threadpool is corrupted. Note that if
    // **user code** panics, we should catch that and redirect.
    let abort_guard = unwind::AbortIfPanic;

    worker_thread.wait_until(&registry.terminate_latch);

    // Should not be any work left in our queue.
    debug_assert!(worker_thread.pop().is_none());

    // let registry know we are done
    registry.thread_infos[index].stopped.set();

    // Normal termination, do not abort.
    mem::forget(abort_guard);
}

/// If already in a worker-thread, just execute `op`.  Otherwise,
/// execute `op` in the default thread-pool. Either way, block until
/// `op` completes and return its return value. If `op` panics, that
/// panic will be propagated as well.
pub fn in_worker<OP, R>(op: OP) -> R
    where OP: FnOnce(&WorkerThread) -> R + Send, R: Send
{
    unsafe {
        let owner_thread = WorkerThread::current();
        if !owner_thread.is_null() {
            // Perfectly valid to give them a `&T`: this is the
            // current thread, so we know the data structure won't be
            // invalidated until we return.
            return op(&*owner_thread);
        } else {
            return in_worker_cold(op);
        }
    }
}

#[cold]
unsafe fn in_worker_cold<OP, R>(op: OP) -> R
    where OP: FnOnce(&WorkerThread) -> R + Send, R: Send
{
    // never run from a worker thread; just shifts over into worker threads
    debug_assert!(WorkerThread::current().is_null());
    let job = StackJob::new(|| in_worker(op), LockLatch::new());
    global_registry().inject(&[job.as_job_ref()]);
    job.latch.wait();
    job.into_result()
}
<|endoftext|>


<commit_before>/// The query resolver that operates on the AST and the TOML object

use toml::Value;
use tokenizer::Token;
use error::*;

fn resolve(toml: &mut Value, tokens: Token) -> Result<&mut Value> {
    let ident = tokens.identifier().unwrap().clone();
    Err(Error::from(ErrorKind::IdentifierNotFoundInDocument(ident)))
}

#[cfg(test)]
mod test {
    use toml::from_str as toml_from_str;
    use toml::Value;
    use tokenizer::*;
    use error::*;
    use super::resolve;

    macro_rules! do_resolve {
        ( $toml:ident => $query:expr ) => {
            resolve(&mut $toml, tokenize_with_seperator(&String::from($query), '.').unwrap())
        }
    }

    #[test]
    fn test_resolve_empty_toml_simple_query() {
        let mut toml = toml_from_str("").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_err());
        let result = result.unwrap_err();

        let errkind = result.kind();
        assert!(is_match!(errkind, &ErrorKind::IdentifierNotFoundInDocument { .. }));
    }

    #[test]
    fn test_resolve_present_bool() {
        let mut toml = toml_from_str("example = true").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Boolean(true)));
    }

    #[test]
    fn test_resolve_present_integer() {
        let mut toml = toml_from_str("example = 1").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Integer(1)));
    }

    #[test]
    fn test_resolve_present_float() {
        let mut toml = toml_from_str("example = 1.0").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Float(1.0)));
    }

    #[test]
    fn test_resolve_present_string() {
        let mut toml = toml_from_str("example = 'string'").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::String(_)));
        match result {
            &mut Value::String(ref s) => assert_eq!("string", s),
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_resolve_present_array_bools() {
        let mut toml = toml_from_str("example = [ true, false ]").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Array(_)));
        match result {
            &mut Value::Array(ref ary) => {
                assert_eq!(ary[0], Value::Boolean(true));
                assert_eq!(ary[1], Value::Boolean(false));
            },
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_resolve_present_array_integers() {
        let mut toml = toml_from_str("example = [ 1, 1337 ]").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Array(_)));
        match result {
            &mut Value::Array(ref ary) => {
                assert_eq!(ary[0], Value::Integer(1));
                assert_eq!(ary[1], Value::Integer(1337));
            },
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_resolve_present_array_floats() {
        let mut toml = toml_from_str("example = [ 1.0, 133.25 ]").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Array(_)));
        match result {
            &mut Value::Array(ref ary) => {
                assert_eq!(ary[0], Value::Float(1.0));
                assert_eq!(ary[1], Value::Float(133.25));
            },
            _ => panic!("What just happened?"),
        }
    }

}
<commit_msg>Impl resolve to fix test_resolve_present_bool<commit_after>/// The query resolver that operates on the AST and the TOML object

use toml::Value;
use tokenizer::Token;
use error::*;

fn resolve(toml: &mut Value, tokens: Token) -> Result<&mut Value> {
    let ident = tokens.identifier().unwrap().clone();

    match toml {
        &mut Value::Table(ref mut t) => {
            match t.get_mut(&ident) {
                None    => Err(Error::from(ErrorKind::IdentifierNotFoundInDocument(ident))),
                Some(s) => Ok(s),
            }
        },
        _ => {
            Err(Error::from(ErrorKind::IdentifierNotFoundInDocument(ident)))
        }
    }
}

#[cfg(test)]
mod test {
    use toml::from_str as toml_from_str;
    use toml::Value;
    use tokenizer::*;
    use error::*;
    use super::resolve;

    macro_rules! do_resolve {
        ( $toml:ident => $query:expr ) => {
            resolve(&mut $toml, tokenize_with_seperator(&String::from($query), '.').unwrap())
        }
    }

    #[test]
    fn test_resolve_empty_toml_simple_query() {
        let mut toml = toml_from_str("").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_err());
        let result = result.unwrap_err();

        let errkind = result.kind();
        assert!(is_match!(errkind, &ErrorKind::IdentifierNotFoundInDocument { .. }));
    }

    #[test]
    fn test_resolve_present_bool() {
        let mut toml = toml_from_str("example = true").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Boolean(true)));
    }

    #[test]
    fn test_resolve_present_integer() {
        let mut toml = toml_from_str("example = 1").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Integer(1)));
    }

    #[test]
    fn test_resolve_present_float() {
        let mut toml = toml_from_str("example = 1.0").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Float(1.0)));
    }

    #[test]
    fn test_resolve_present_string() {
        let mut toml = toml_from_str("example = 'string'").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::String(_)));
        match result {
            &mut Value::String(ref s) => assert_eq!("string", s),
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_resolve_present_array_bools() {
        let mut toml = toml_from_str("example = [ true, false ]").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Array(_)));
        match result {
            &mut Value::Array(ref ary) => {
                assert_eq!(ary[0], Value::Boolean(true));
                assert_eq!(ary[1], Value::Boolean(false));
            },
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_resolve_present_array_integers() {
        let mut toml = toml_from_str("example = [ 1, 1337 ]").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Array(_)));
        match result {
            &mut Value::Array(ref ary) => {
                assert_eq!(ary[0], Value::Integer(1));
                assert_eq!(ary[1], Value::Integer(1337));
            },
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_resolve_present_array_floats() {
        let mut toml = toml_from_str("example = [ 1.0, 133.25 ]").unwrap();
        let result = do_resolve!(toml => "example");

        assert!(result.is_ok());
        let result = result.unwrap();

        assert!(is_match!(result, &mut Value::Array(_)));
        match result {
            &mut Value::Array(ref ary) => {
                assert_eq!(ary[0], Value::Float(1.0));
                assert_eq!(ary[1], Value::Float(133.25));
            },
            _ => panic!("What just happened?"),
        }
    }

}
<|endoftext|>


<commit_before>// Copyright (c) 2016
// Jeff Nettleton
//
// Licensed under the MIT license (http://opensource.org/licenses/MIT). This
// file may not be copied, modified, or distributed except according to those
// terms

use std::collections::HashMap;
use chrono::UTC;
use rustc_serialize::{json, Encodable};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

/// A trait that converts data from the handler function to a u8 slice.
pub trait ToOutput {
    fn to_output(&self) -> &[u8];
}

impl ToOutput for str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for &'static str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for String {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for Vec<u8> {
    fn to_output(&self) -> &[u8] {
        self.as_slice()
    }
}

/// This struct reprsents the response to an HTTP client.
#[derive(Debug)]
pub struct Response {
    status:     u16,
    cmsg:       String,
    ctype:      String,
    headers:    HashMap<String, String>,
    payload:    Vec<u8>,
}

impl Response {
    /// Create a new, empty Response.
    pub fn new() -> Response {
        let mut res = Response {
            status:     200,
            cmsg:       String::from("OK"),
            ctype:      String::from("text/plain"),
            headers:    HashMap::new(),
            payload:    Vec::with_capacity(2048),
        };

        let now = UTC::now().format("%a, %d %b %Y, %H:%M:%S %Z").to_string();

        res.add_header("Connection", "close");
        res.add_header("Server", &format!("canteen/{}", VERSION));
        res.add_header("Date", &now);

        res
    }

    /// Creates a Response with a JSON body
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// use canteen::Response;
    ///
    /// #[derive(RustcEncodable)]
    /// struct Foo {
    ///     item: i32,
    /// }
    ///
    /// let foo = Foo { item: 12345 };
    /// let res = Response::as_json(&foo);
    /// ```
    pub fn as_json<T: Encodable>(data: &T) -> Response {
        let mut res = Response::new();

        res.set_content_type("application/json");
        res.append(json::encode(data).unwrap());

        res
    }

    /// Gets the HTTP message for a given status.
    fn get_http_message(status: u16) -> String {
        let msg = match status {
            100 => "Continue",
            101 => "Switching Protocols",
            200 => "OK",
            201 => "Created",
            202 => "Accepted",
            203 => "Non-Authoritative Information",
            204 => "No Content",
            205 => "Reset Content",
            206 => "Partial Content",
            300 => "Multiple Choices",
            301 => "Moved Permanently",
            302 => "Found",
            303 => "See Other",
            304 => "Not Modified",
            305 => "Use Proxy",
            307 => "Temporary Redirect",
            400 => "Bad Request",
            401 => "Unauthorized",
            402 => "Payment Required",
            403 => "Forbidden",
            404 => "Not Found",
            405 => "Method Not Allowed",
            406 => "Not Acceptable",
            407 => "Proxy Authentication Required",
            408 => "Request Time Out",
            409 => "Conflict",
            410 => "Gone",
            411 => "Length Required",
            412 => "Precondition Failed",
            413 => "Request Entity Too Large",
            414 => "Request-URI Too Large",
            415 => "Unsupported Media Type",
            416 => "Requested Range Not Satisfiable",
            417 => "Expectation Failed",
            500 => "Internal Server Error",
            501 => "Not Implemented",
            502 => "Bad Gateway",
            503 => "Service Unavailable",
            504 => "Gateway Time-out",
            505 => "HTTP Version Not Supported",
            _   => "OK",
        };

        String::from(msg)
    }

    /// Sets the response status for the HTTP response.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// res.set_status(200);
    /// ```
    pub fn set_status(&mut self, status: u16) {
        self.status = status;
        self.cmsg = Response::get_http_message(status);
    }

    /// Sets the Content-Type header for the HTTP response.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// res.set_content_type("text/html");
    /// ```
    pub fn set_content_type(&mut self, ctype: &str) {
        self.ctype = String::from(ctype);
    }

    /// Adds a header to the HTTP response.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// res.add_header("Content-Type", "text/html");
    /// ```
    pub fn add_header(&mut self, key: &str, value: &str) {
        if !self.headers.contains_key(key) {
            self.headers.insert(String::from(key), String::from(value));
        }
    }

    /// Appends data to the body of the HTTP response. The trait ToOutput must
    /// be implemented for the type passed.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// let data = "{ message: \"Hello, world!\" }";
    /// res.append(data);
    /// ```
    pub fn append<T: ToOutput>(&mut self, payload: T) {
        self.payload.extend(payload.to_output().into_iter());
    }

    /// Returns a byte array containing the full contents of the HTTP response,
    /// for use by the Canteen struct.
    pub fn gen_output(&self) -> Vec<u8> {
        let mut output: Vec<u8> = Vec::with_capacity(self.payload.len() + 500);
        let mut inter = String::new();

        inter.push_str(&format!("HTTP/1.1 {} {}\r\n", self.status, self.cmsg));

        for (key, value) in &self.headers {
            inter.push_str(&format!("{}: {}\r\n", key, value));
        }

        inter.push_str(&format!("Content-Type: {}\r\n", self.ctype));
        inter.push_str(&format!("Content-Length: {}\r\n", self.payload.len()));
        inter.push_str("\r\n");

        output.extend(inter.as_bytes());
        output.extend(self.payload.iter());

        output
    }
}

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

    #[derive(RustcEncodable)]
    struct Foo {
        item: i32,
    }

    #[test]
    fn test_response_as_json() {
        let foo = Foo { item: 12345 };
        let res_j = Response::as_json(&foo);
        let mut res_r = Response::new();

        res_r.set_content_type("application/json");
        res_r.append(json::encode(&foo).unwrap());

        assert_eq!(res_r.gen_output(), res_j.gen_output());
    }

    #[test]
    fn test_response_http_message() {
        assert_eq!("OK", Response::get_http_message(200));
    }

    #[test]
    fn test_tooutput_trait_static_str() {
        let ar: [u8; 3] = [97, 98, 99];

        assert_eq!(ar, "abc".to_output());
    }

    #[test]
    fn test_tooutput_trait_str() {
        let ar: [u8; 3] = [97, 98, 99];
        let st = "abc";

        assert_eq!(ar, st.to_output());
    }

    #[test]
    fn test_tooutput_trait_string() {
        let ar: [u8; 3] = [97, 98, 99];
        let st = String::from("abc");

        assert_eq!(ar, st.to_output());
    }

    #[test]
    fn test_tooutput_trait_vec() {
        let ar: [u8; 5] = [1, 2, 3, 4, 5];
        let vc: Vec<u8> = vec![1, 2, 3, 4, 5];

        assert_eq!(ar, vc.to_output());
    }
}
<commit_msg>use std::collections::BTreeMap instead of HashMap so that headers are ordered consistently<commit_after>// Copyright (c) 2016
// Jeff Nettleton
//
// Licensed under the MIT license (http://opensource.org/licenses/MIT). This
// file may not be copied, modified, or distributed except according to those
// terms

use std::collections::BTreeMap;
use chrono::UTC;
use rustc_serialize::{json, Encodable};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

/// A trait that converts data from the handler function to a u8 slice.
pub trait ToOutput {
    fn to_output(&self) -> &[u8];
}

impl ToOutput for str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for &'static str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for String {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for Vec<u8> {
    fn to_output(&self) -> &[u8] {
        self.as_slice()
    }
}

/// This struct reprsents the response to an HTTP client.
#[derive(Debug)]
pub struct Response {
    status:     u16,
    cmsg:       String,
    ctype:      String,
    headers:    BTreeMap<String, String>,
    payload:    Vec<u8>,
}

impl Response {
    /// Create a new, empty Response.
    pub fn new() -> Response {
        let mut res = Response {
            status:     200,
            cmsg:       String::from("OK"),
            ctype:      String::from("text/plain"),
            headers:    BTreeMap::new(),
            payload:    Vec::with_capacity(2048),
        };

        let now = UTC::now().format("%a, %d %b %Y, %H:%M:%S %Z").to_string();

        res.add_header("Connection", "close");
        res.add_header("Server", &format!("canteen/{}", VERSION));
        res.add_header("Date", &now);

        res
    }

    /// Creates a Response with a JSON body
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// use canteen::Response;
    ///
    /// #[derive(RustcEncodable)]
    /// struct Foo {
    ///     item: i32,
    /// }
    ///
    /// let foo = Foo { item: 12345 };
    /// let res = Response::as_json(&foo);
    /// ```
    pub fn as_json<T: Encodable>(data: &T) -> Response {
        let mut res = Response::new();

        res.set_content_type("application/json");
        res.append(json::encode(data).unwrap());

        res
    }

    /// Gets the HTTP message for a given status.
    fn get_http_message(status: u16) -> String {
        let msg = match status {
            100 => "Continue",
            101 => "Switching Protocols",
            200 => "OK",
            201 => "Created",
            202 => "Accepted",
            203 => "Non-Authoritative Information",
            204 => "No Content",
            205 => "Reset Content",
            206 => "Partial Content",
            300 => "Multiple Choices",
            301 => "Moved Permanently",
            302 => "Found",
            303 => "See Other",
            304 => "Not Modified",
            305 => "Use Proxy",
            307 => "Temporary Redirect",
            400 => "Bad Request",
            401 => "Unauthorized",
            402 => "Payment Required",
            403 => "Forbidden",
            404 => "Not Found",
            405 => "Method Not Allowed",
            406 => "Not Acceptable",
            407 => "Proxy Authentication Required",
            408 => "Request Time Out",
            409 => "Conflict",
            410 => "Gone",
            411 => "Length Required",
            412 => "Precondition Failed",
            413 => "Request Entity Too Large",
            414 => "Request-URI Too Large",
            415 => "Unsupported Media Type",
            416 => "Requested Range Not Satisfiable",
            417 => "Expectation Failed",
            500 => "Internal Server Error",
            501 => "Not Implemented",
            502 => "Bad Gateway",
            503 => "Service Unavailable",
            504 => "Gateway Time-out",
            505 => "HTTP Version Not Supported",
            _   => "OK",
        };

        String::from(msg)
    }

    /// Sets the response status for the HTTP response.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// res.set_status(200);
    /// ```
    pub fn set_status(&mut self, status: u16) {
        self.status = status;
        self.cmsg = Response::get_http_message(status);
    }

    /// Sets the Content-Type header for the HTTP response.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// res.set_content_type("text/html");
    /// ```
    pub fn set_content_type(&mut self, ctype: &str) {
        self.ctype = String::from(ctype);
    }

    /// Adds a header to the HTTP response.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// res.add_header("Content-Type", "text/html");
    /// ```
    pub fn add_header(&mut self, key: &str, value: &str) {
        if !self.headers.contains_key(key) {
            self.headers.insert(String::from(key), String::from(value));
        }
    }

    /// Appends data to the body of the HTTP response. The trait ToOutput must
    /// be implemented for the type passed.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use canteen::Response;
    ///
    /// let mut res = Response::new();
    /// let data = "{ message: \"Hello, world!\" }";
    /// res.append(data);
    /// ```
    pub fn append<T: ToOutput>(&mut self, payload: T) {
        self.payload.extend(payload.to_output().into_iter());
    }

    /// Returns a byte array containing the full contents of the HTTP response,
    /// for use by the Canteen struct.
    pub fn gen_output(&self) -> Vec<u8> {
        let mut output: Vec<u8> = Vec::with_capacity(self.payload.len() + 500);
        let mut inter = String::new();

        inter.push_str(&format!("HTTP/1.1 {} {}\r\n", self.status, self.cmsg));

        for (key, value) in &self.headers {
            inter.push_str(&format!("{}: {}\r\n", key, value));
        }

        inter.push_str(&format!("Content-Type: {}\r\n", self.ctype));
        inter.push_str(&format!("Content-Length: {}\r\n", self.payload.len()));
        inter.push_str("\r\n");

        output.extend(inter.as_bytes());
        output.extend(self.payload.iter());

        output
    }
}

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

    #[derive(RustcEncodable)]
    struct Foo {
        item: i32,
    }

    #[test]
    fn test_response_as_json() {
        let foo = Foo { item: 12345 };
        let res_j = Response::as_json(&foo);
        let mut res_r = Response::new();

        res_r.set_content_type("application/json");
        res_r.append(json::encode(&foo).unwrap());

        assert_eq!(res_r.gen_output(), res_j.gen_output());
    }

    #[test]
    fn test_response_http_message() {
        assert_eq!("OK", Response::get_http_message(200));
    }

    #[test]
    fn test_tooutput_trait_static_str() {
        let ar: [u8; 3] = [97, 98, 99];

        assert_eq!(ar, "abc".to_output());
    }

    #[test]
    fn test_tooutput_trait_str() {
        let ar: [u8; 3] = [97, 98, 99];
        let st = "abc";

        assert_eq!(ar, st.to_output());
    }

    #[test]
    fn test_tooutput_trait_string() {
        let ar: [u8; 3] = [97, 98, 99];
        let st = String::from("abc");

        assert_eq!(ar, st.to_output());
    }

    #[test]
    fn test_tooutput_trait_vec() {
        let ar: [u8; 5] = [1, 2, 3, 4, 5];
        let vc: Vec<u8> = vec![1, 2, 3, 4, 5];

        assert_eq!(ar, vc.to_output());
    }
}
<|endoftext|>


<commit_before>use crate::{return_variant, ErrorCode, ReturnVariant};

use core::mem::transmute;

/// The response type from the [`command`](crate::Syscalls::command) syscall.
/// Can represent a success or a failure with or without associated data.
///
/// After a syscall is made, registers `r1`-`r3` contain the output as
/// described by [TRD 104][trd-104]. Some syscalls only return success/failure,
/// while others provide associated data. This is done by placing the _return
/// variant_ in `r0`, which specifies how the output registers should be
/// interpreted. For syscalls other than `command`, the possible output
/// variants are fixed; you always know which variants are expected given the
/// syscall class.
///
/// However, the `command` syscall is flexible - there must be one success
/// variant and one failure variant for a given driver/command ID, but
/// which variants those are, and what data is expected, cannot be checked
/// statically. Capsules and userspace APIs must agree on the expected
/// variants for success and failure.
///
/// # Example
///
/// This uses the [`to_result`] method to implicitly check variants and convert
/// to a `Result`.
///
/// ```ignore
/// let res: Result<(u32, u32), ErrorCode> = Syscalls::command(314, 1, 1, 2).to_result();
/// match res {
///     Ok((val1, val2)) => {
///         // Success with associated data in val1, val2.
///     }
///     Err(ErrorCode::BadRVal) => {
///         // Incorrect return variant! We may choose to handle this
///         // explicitly or propagate upwards without branching.
///     }
///     Err(ec) => {
///         // The driver returned an error (or it doesn't exist).
///     }
/// }
/// ```
///
/// This uses the `get_*` methods to check the variants explicitly and extract
/// the associated data.
///
/// ```ignore
/// let command_return = Syscalls::command(314, 1, 1, 2);
/// if let Some((val1, val2)) = command_return.get_success_2_u32() {
///     // If there was a success, there is an associated data (u32, u32).
/// } else if let Some(error_code) = command_return.get_failure() {
///     // If there was a failure, there's no associated data and we only
///     // have an error code.
/// } else {
///     // Incorrect return variant! If this occurs, your capsule and userspace
///     // API do not agree on what the return variants should be.
///     // An application may want to panic in this case to catch this early.
/// }
/// ```
///
/// [trd-104]: https://github.com/tock/tock/blob/master/doc/reference/trd104-syscalls.md#32-return-values
#[must_use = "this `CommandReturn` may represent an error, which should be handled"]
#[derive(Clone, Copy, Debug)]
pub struct CommandReturn {
    return_variant: ReturnVariant,

    // Safety invariant on r1: If return_variant is failure variant, r1 must be
    // a valid ErrorCode.
    r1: u32,
    r2: u32,
    r3: u32,
}

impl CommandReturn {
    /// # Safety
    /// If return_variant is a failure variant, r1 must be a valid ErrorCode.
    pub unsafe fn new(return_variant: ReturnVariant, r1: u32, r2: u32, r3: u32) -> Self {
        CommandReturn {
            return_variant,
            r1,
            r2,
            r3,
        }
    }

    /// Returns true if this CommandReturn is of type Failure. Note that this
    /// does not return true for other failure types, such as Failure with u32.
    pub fn is_failure(&self) -> bool {
        self.return_variant == return_variant::FAILURE
    }

    /// Returns true if this CommandReturn is of type Failure with u32.
    pub fn is_failure_u32(&self) -> bool {
        self.return_variant == return_variant::FAILURE_U32
    }

    /// Returns true if this CommandReturn is of type Failure with 2 u32.
    pub fn is_failure_2_u32(&self) -> bool {
        self.return_variant == return_variant::FAILURE_2_U32
    }

    /// Returns true if this CommandReturn is of type Failure with u64.
    pub fn is_failure_u64(&self) -> bool {
        self.return_variant == return_variant::FAILURE_U64
    }

    /// Returns true if this CommandReturn is of type Success. Note that this
    /// does not return true for other success types, such as Success with u32.
    pub fn is_success(&self) -> bool {
        self.return_variant == return_variant::SUCCESS
    }

    /// Returns true if this CommandReturn is of type Success with u32.
    pub fn is_success_u32(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_U32
    }

    /// Returns true if this CommandReturn is of type Success with 2 u32.
    pub fn is_success_2_u32(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_2_U32
    }

    /// Returns true if this CommandReturn is of type Success with u64.
    pub fn is_success_u64(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_U64
    }

    /// Returns true if this CommandReturn is of type Success with 3 u32.
    pub fn is_success_3_u32(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_3_U32
    }

    /// Returns true if this CommandReturn is of type Success with u32 and u64.
    pub fn is_success_u32_u64(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_U32_U64
    }

    /// Returns the error code if this CommandReturn is of type Failure.
    pub fn get_failure(&self) -> Option<ErrorCode> {
        if !self.is_failure() {
            return None;
        }
        Some(unsafe { transmute(self.r1 as u16) })
    }

    /// Returns the error code and value if this CommandReturn is of type
    /// Failure with u32.
    pub fn get_failure_u32(&self) -> Option<(ErrorCode, u32)> {
        if !self.is_failure_u32() {
            return None;
        }
        Some((unsafe { transmute(self.r1 as u16) }, self.r2))
    }

    /// Returns the error code and return values if this CommandReturn is of
    /// type Failure with 2 u32.
    pub fn get_failure_2_u32(&self) -> Option<(ErrorCode, u32, u32)> {
        if !self.is_failure_2_u32() {
            return None;
        }
        Some((unsafe { transmute(self.r1 as u16) }, self.r2, self.r3))
    }

    /// Returns the error code and return value if this CommandReturn is of type
    /// Failure with u64.
    pub fn get_failure_u64(&self) -> Option<(ErrorCode, u64)> {
        if !self.is_failure_u64() {
            return None;
        }
        Some((
            unsafe { transmute(self.r1 as u16) },
            self.r2 as u64 + ((self.r3 as u64) << 32),
        ))
    }

    /// Returns the value if this CommandReturn is of type Success with u32.
    pub fn get_success_u32(&self) -> Option<u32> {
        if !self.is_success_u32() {
            return None;
        }
        Some(self.r1)
    }

    /// Returns the values if this CommandReturn is of type Success with 2 u32.
    pub fn get_success_2_u32(&self) -> Option<(u32, u32)> {
        if !self.is_success_2_u32() {
            return None;
        }
        Some((self.r1, self.r2))
    }

    /// Returns the value if this CommandReturn is of type Success with u64.
    pub fn get_success_u64(&self) -> Option<u64> {
        if !self.is_success_u64() {
            return None;
        }
        Some(self.r1 as u64 + ((self.r2 as u64) << 32))
    }

    /// Returns the values if this CommandReturn is of type Success with 3 u32.
    pub fn get_success_3_u32(&self) -> Option<(u32, u32, u32)> {
        if !self.is_success_3_u32() {
            return None;
        }
        Some((self.r1, self.r2, self.r3))
    }

    /// Returns the values if this CommandReturn is of type Success with u32 and
    /// u64.
    pub fn get_success_u32_u64(&self) -> Option<(u32, u64)> {
        if !self.is_success_u32_u64() {
            return None;
        }
        Some((self.r1, self.r2 as u64 + ((self.r3 as u64) << 32)))
    }

    /// Returns the register values used to create this command.
    pub fn raw_values(&self) -> (ReturnVariant, u32, u32, u32) {
        (self.return_variant, self.r1, self.r2, self.r3)
    }

    /// Returns the return variant of this command.
    pub fn return_variant(&self) -> ReturnVariant {
        self.return_variant
    }

    /// Interprets this `CommandReturn` as a `Result`, checking the success and
    /// failure variants, as well as extracting the relevant data.
    ///
    /// If neither the success or failure variants match what is required by
    /// `T` and `E`, this function will return `Err(ErrorCode::BadRVal)`.
    /// If `E` contains non-`ErrorCode` data in this case, the data will be 0.
    ///
    /// It is recommended to use type ascription or `::<>` to make the types
    /// for `T` and `E` explicit at call-site.
    pub fn to_result<T, E>(self) -> Result<T, E>
    where
        T: SuccessData,
        E: FailureData,
    {
        let (return_variant, r1, mut r2, mut r3) = self.raw_values();
        if return_variant == T::RETURN_VARIANT {
            return Ok(T::from_raw_values(r1, r2, r3));
        }
        let ec: ErrorCode = if return_variant == E::RETURN_VARIANT {
            // Safety: E::RETURN_VARIANT must be a failure variant, and
            // failure variants must contain a valid ErrorCode in r1.
            unsafe { core::mem::transmute(r1 as u16) }
        } else {
            r2 = 0;
            r3 = 0;
            ErrorCode::BadRVal
        };
        Err(E::from_raw_values(ec, r2, r3))
    }
}

mod sealed {
    pub trait Sealed {}
}

/// Output from a successful `command` syscall.
pub trait SuccessData: sealed::Sealed {
    /// The return variant for this success type, stored in `r0` after
    /// performing a `command` syscall.
    const RETURN_VARIANT: ReturnVariant;

    /// Constructs the success data given the raw register values.
    fn from_raw_values(r1: u32, r2: u32, r3: u32) -> Self;
}

impl sealed::Sealed for () {}
impl SuccessData for () {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS;

    fn from_raw_values(_r1: u32, _r2: u32, _r3: u32) -> Self {}
}
impl sealed::Sealed for u32 {}
impl SuccessData for u32 {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_U32;

    fn from_raw_values(r1: u32, _r2: u32, _r3: u32) -> Self {
        r1
    }
}
impl sealed::Sealed for u64 {}
impl SuccessData for u64 {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_U64;

    fn from_raw_values(r1: u32, r2: u32, _r3: u32) -> Self {
        r1 as u64 | ((r2 as u64) << 32)
    }
}
impl sealed::Sealed for (u32, u32) {}
impl SuccessData for (u32, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_2_U32;

    fn from_raw_values(r1: u32, r2: u32, _r3: u32) -> Self {
        (r1, r2)
    }
}
impl sealed::Sealed for (u32, u64) {}
impl SuccessData for (u32, u64) {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_U32_U64;

    fn from_raw_values(r1: u32, r2: u32, r3: u32) -> Self {
        (r1, r2 as u64 | ((r3 as u64) << 32))
    }
}
impl sealed::Sealed for (u32, u32, u32) {}
impl SuccessData for (u32, u32, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_3_U32;

    fn from_raw_values(r1: u32, r2: u32, r3: u32) -> Self {
        (r1, r2, r3)
    }
}

/// Output from a failed `command` syscall.
pub unsafe trait FailureData: sealed::Sealed {
    /// The return variant for this failure type, stored in `r0` after
    /// performing a `command` syscall.
    ///
    /// # Safety
    /// This must represent a failure variant, such that `r1` will always be
    /// a valid [`ErrorCode`].
    const RETURN_VARIANT: ReturnVariant;

    /// Constructs the error data given the raw register values.
    fn from_raw_values(r1: ErrorCode, r2: u32, r3: u32) -> Self;
}

impl sealed::Sealed for ErrorCode {}
unsafe impl FailureData for ErrorCode {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE;

    fn from_raw_values(r1: ErrorCode, _r2: u32, _r3: u32) -> Self {
        r1
    }
}
impl sealed::Sealed for (ErrorCode, u32) {}
unsafe impl FailureData for (ErrorCode, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE_U32;

    fn from_raw_values(r1: ErrorCode, r2: u32, _r3: u32) -> Self {
        (r1, r2)
    }
}
impl sealed::Sealed for (ErrorCode, u32, u32) {}
unsafe impl FailureData for (ErrorCode, u32, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE_2_U32;

    fn from_raw_values(r1: ErrorCode, r2: u32, r3: u32) -> Self {
        (r1, r2, r3)
    }
}
impl sealed::Sealed for (ErrorCode, u64) {}
unsafe impl FailureData for (ErrorCode, u64) {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE_U64;

    fn from_raw_values(r1: ErrorCode, r2: u32, r3: u32) -> Self {
        (r1, r2 as u64 | ((r3 as u64) << 32))
    }
}
<commit_msg>Document trait sealing<commit_after>use crate::{return_variant, ErrorCode, ReturnVariant};

use core::mem::transmute;

/// The response type from the [`command`](crate::Syscalls::command) syscall.
/// Can represent a success or a failure with or without associated data.
///
/// After a syscall is made, registers `r1`-`r3` contain the output as
/// described by [TRD 104][trd-104]. Some syscalls only return success/failure,
/// while others provide associated data. This is done by placing the _return
/// variant_ in `r0`, which specifies how the output registers should be
/// interpreted. For syscalls other than `command`, the possible output
/// variants are fixed; you always know which variants are expected given the
/// syscall class.
///
/// However, the `command` syscall is flexible - there must be one success
/// variant and one failure variant for a given driver/command ID, but
/// which variants those are, and what data is expected, cannot be checked
/// statically. Capsules and userspace APIs must agree on the expected
/// variants for success and failure.
///
/// # Example
///
/// This uses the [`to_result`] method to implicitly check variants and convert
/// to a `Result`.
///
/// ```ignore
/// let res: Result<(u32, u32), ErrorCode> = Syscalls::command(314, 1, 1, 2).to_result();
/// match res {
///     Ok((val1, val2)) => {
///         // Success with associated data in val1, val2.
///     }
///     Err(ErrorCode::BadRVal) => {
///         // Incorrect return variant! We may choose to handle this
///         // explicitly or propagate upwards without branching.
///     }
///     Err(ec) => {
///         // The driver returned an error (or it doesn't exist).
///     }
/// }
/// ```
///
/// This uses the `get_*` methods to check the variants explicitly and extract
/// the associated data.
///
/// ```ignore
/// let command_return = Syscalls::command(314, 1, 1, 2);
/// if let Some((val1, val2)) = command_return.get_success_2_u32() {
///     // If there was a success, there is an associated data (u32, u32).
/// } else if let Some(error_code) = command_return.get_failure() {
///     // If there was a failure, there's no associated data and we only
///     // have an error code.
/// } else {
///     // Incorrect return variant! If this occurs, your capsule and userspace
///     // API do not agree on what the return variants should be.
///     // An application may want to panic in this case to catch this early.
/// }
/// ```
///
/// [trd-104]: https://github.com/tock/tock/blob/master/doc/reference/trd104-syscalls.md#32-return-values
#[must_use = "this `CommandReturn` may represent an error, which should be handled"]
#[derive(Clone, Copy, Debug)]
pub struct CommandReturn {
    return_variant: ReturnVariant,

    // Safety invariant on r1: If return_variant is failure variant, r1 must be
    // a valid ErrorCode.
    r1: u32,
    r2: u32,
    r3: u32,
}

impl CommandReturn {
    /// # Safety
    /// If return_variant is a failure variant, r1 must be a valid ErrorCode.
    pub unsafe fn new(return_variant: ReturnVariant, r1: u32, r2: u32, r3: u32) -> Self {
        CommandReturn {
            return_variant,
            r1,
            r2,
            r3,
        }
    }

    /// Returns true if this CommandReturn is of type Failure. Note that this
    /// does not return true for other failure types, such as Failure with u32.
    pub fn is_failure(&self) -> bool {
        self.return_variant == return_variant::FAILURE
    }

    /// Returns true if this CommandReturn is of type Failure with u32.
    pub fn is_failure_u32(&self) -> bool {
        self.return_variant == return_variant::FAILURE_U32
    }

    /// Returns true if this CommandReturn is of type Failure with 2 u32.
    pub fn is_failure_2_u32(&self) -> bool {
        self.return_variant == return_variant::FAILURE_2_U32
    }

    /// Returns true if this CommandReturn is of type Failure with u64.
    pub fn is_failure_u64(&self) -> bool {
        self.return_variant == return_variant::FAILURE_U64
    }

    /// Returns true if this CommandReturn is of type Success. Note that this
    /// does not return true for other success types, such as Success with u32.
    pub fn is_success(&self) -> bool {
        self.return_variant == return_variant::SUCCESS
    }

    /// Returns true if this CommandReturn is of type Success with u32.
    pub fn is_success_u32(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_U32
    }

    /// Returns true if this CommandReturn is of type Success with 2 u32.
    pub fn is_success_2_u32(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_2_U32
    }

    /// Returns true if this CommandReturn is of type Success with u64.
    pub fn is_success_u64(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_U64
    }

    /// Returns true if this CommandReturn is of type Success with 3 u32.
    pub fn is_success_3_u32(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_3_U32
    }

    /// Returns true if this CommandReturn is of type Success with u32 and u64.
    pub fn is_success_u32_u64(&self) -> bool {
        self.return_variant == return_variant::SUCCESS_U32_U64
    }

    /// Returns the error code if this CommandReturn is of type Failure.
    pub fn get_failure(&self) -> Option<ErrorCode> {
        if !self.is_failure() {
            return None;
        }
        Some(unsafe { transmute(self.r1 as u16) })
    }

    /// Returns the error code and value if this CommandReturn is of type
    /// Failure with u32.
    pub fn get_failure_u32(&self) -> Option<(ErrorCode, u32)> {
        if !self.is_failure_u32() {
            return None;
        }
        Some((unsafe { transmute(self.r1 as u16) }, self.r2))
    }

    /// Returns the error code and return values if this CommandReturn is of
    /// type Failure with 2 u32.
    pub fn get_failure_2_u32(&self) -> Option<(ErrorCode, u32, u32)> {
        if !self.is_failure_2_u32() {
            return None;
        }
        Some((unsafe { transmute(self.r1 as u16) }, self.r2, self.r3))
    }

    /// Returns the error code and return value if this CommandReturn is of type
    /// Failure with u64.
    pub fn get_failure_u64(&self) -> Option<(ErrorCode, u64)> {
        if !self.is_failure_u64() {
            return None;
        }
        Some((
            unsafe { transmute(self.r1 as u16) },
            self.r2 as u64 + ((self.r3 as u64) << 32),
        ))
    }

    /// Returns the value if this CommandReturn is of type Success with u32.
    pub fn get_success_u32(&self) -> Option<u32> {
        if !self.is_success_u32() {
            return None;
        }
        Some(self.r1)
    }

    /// Returns the values if this CommandReturn is of type Success with 2 u32.
    pub fn get_success_2_u32(&self) -> Option<(u32, u32)> {
        if !self.is_success_2_u32() {
            return None;
        }
        Some((self.r1, self.r2))
    }

    /// Returns the value if this CommandReturn is of type Success with u64.
    pub fn get_success_u64(&self) -> Option<u64> {
        if !self.is_success_u64() {
            return None;
        }
        Some(self.r1 as u64 + ((self.r2 as u64) << 32))
    }

    /// Returns the values if this CommandReturn is of type Success with 3 u32.
    pub fn get_success_3_u32(&self) -> Option<(u32, u32, u32)> {
        if !self.is_success_3_u32() {
            return None;
        }
        Some((self.r1, self.r2, self.r3))
    }

    /// Returns the values if this CommandReturn is of type Success with u32 and
    /// u64.
    pub fn get_success_u32_u64(&self) -> Option<(u32, u64)> {
        if !self.is_success_u32_u64() {
            return None;
        }
        Some((self.r1, self.r2 as u64 + ((self.r3 as u64) << 32)))
    }

    /// Returns the register values used to create this command.
    pub fn raw_values(&self) -> (ReturnVariant, u32, u32, u32) {
        (self.return_variant, self.r1, self.r2, self.r3)
    }

    /// Returns the return variant of this command.
    pub fn return_variant(&self) -> ReturnVariant {
        self.return_variant
    }

    /// Interprets this `CommandReturn` as a `Result`, checking the success and
    /// failure variants, as well as extracting the relevant data.
    ///
    /// If neither the success or failure variants match what is required by
    /// `T` and `E`, this function will return `Err(ErrorCode::BadRVal)`.
    /// If `E` contains non-`ErrorCode` data in this case, the data will be 0.
    ///
    /// It is recommended to use type ascription or `::<>` to make the types
    /// for `T` and `E` explicit at call-site.
    pub fn to_result<T, E>(self) -> Result<T, E>
    where
        T: SuccessData,
        E: FailureData,
    {
        let (return_variant, r1, mut r2, mut r3) = self.raw_values();
        if return_variant == T::RETURN_VARIANT {
            return Ok(T::from_raw_values(r1, r2, r3));
        }
        let ec: ErrorCode = if return_variant == E::RETURN_VARIANT {
            // Safety: E::RETURN_VARIANT must be a failure variant, and
            // failure variants must contain a valid ErrorCode in r1.
            unsafe { core::mem::transmute(r1 as u16) }
        } else {
            r2 = 0;
            r3 = 0;
            ErrorCode::BadRVal
        };
        Err(E::from_raw_values(ec, r2, r3))
    }
}

mod sealed {
    pub trait Sealed {}
}

/// Output from a successful `command` syscall.
///
/// This trait is [sealed], meaning foreign implementations cannot be defined,
/// even though it can be referenced by foreign crates.
///
/// [sealed]: https://rust-lang.github.io/api-guidelines/future-proofing.html#sealed-traits-protect-against-downstream-implementations-c-sealed
pub trait SuccessData: sealed::Sealed {
    /// The return variant for this success type, stored in `r0` after
    /// performing a `command` syscall.
    const RETURN_VARIANT: ReturnVariant;

    /// Constructs the success data given the raw register values.
    fn from_raw_values(r1: u32, r2: u32, r3: u32) -> Self;
}

impl sealed::Sealed for () {}
impl SuccessData for () {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS;

    fn from_raw_values(_r1: u32, _r2: u32, _r3: u32) -> Self {}
}
impl sealed::Sealed for u32 {}
impl SuccessData for u32 {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_U32;

    fn from_raw_values(r1: u32, _r2: u32, _r3: u32) -> Self {
        r1
    }
}
impl sealed::Sealed for u64 {}
impl SuccessData for u64 {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_U64;

    fn from_raw_values(r1: u32, r2: u32, _r3: u32) -> Self {
        r1 as u64 | ((r2 as u64) << 32)
    }
}
impl sealed::Sealed for (u32, u32) {}
impl SuccessData for (u32, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_2_U32;

    fn from_raw_values(r1: u32, r2: u32, _r3: u32) -> Self {
        (r1, r2)
    }
}
impl sealed::Sealed for (u32, u64) {}
impl SuccessData for (u32, u64) {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_U32_U64;

    fn from_raw_values(r1: u32, r2: u32, r3: u32) -> Self {
        (r1, r2 as u64 | ((r3 as u64) << 32))
    }
}
impl sealed::Sealed for (u32, u32, u32) {}
impl SuccessData for (u32, u32, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::SUCCESS_3_U32;

    fn from_raw_values(r1: u32, r2: u32, r3: u32) -> Self {
        (r1, r2, r3)
    }
}

/// Output from a failed `command` syscall.
///
/// This trait is [sealed], meaning foreign implementations cannot be defined,
/// even though it can be referenced by foreign crates.
///
/// [sealed]: https://rust-lang.github.io/api-guidelines/future-proofing.html#sealed-traits-protect-against-downstream-implementations-c-sealed
pub unsafe trait FailureData: sealed::Sealed {
    /// The return variant for this failure type, stored in `r0` after
    /// performing a `command` syscall.
    ///
    /// # Safety
    /// This must represent a failure variant, such that `r1` will always be
    /// a valid [`ErrorCode`].
    const RETURN_VARIANT: ReturnVariant;

    /// Constructs the error data given the raw register values.
    fn from_raw_values(r1: ErrorCode, r2: u32, r3: u32) -> Self;
}

impl sealed::Sealed for ErrorCode {}
unsafe impl FailureData for ErrorCode {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE;

    fn from_raw_values(r1: ErrorCode, _r2: u32, _r3: u32) -> Self {
        r1
    }
}
impl sealed::Sealed for (ErrorCode, u32) {}
unsafe impl FailureData for (ErrorCode, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE_U32;

    fn from_raw_values(r1: ErrorCode, r2: u32, _r3: u32) -> Self {
        (r1, r2)
    }
}
impl sealed::Sealed for (ErrorCode, u32, u32) {}
unsafe impl FailureData for (ErrorCode, u32, u32) {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE_2_U32;

    fn from_raw_values(r1: ErrorCode, r2: u32, r3: u32) -> Self {
        (r1, r2, r3)
    }
}
impl sealed::Sealed for (ErrorCode, u64) {}
unsafe impl FailureData for (ErrorCode, u64) {
    const RETURN_VARIANT: ReturnVariant = return_variant::FAILURE_U64;

    fn from_raw_values(r1: ErrorCode, r2: u32, r3: u32) -> Self {
        (r1, r2 as u64 | ((r3 as u64) << 32))
    }
}
<|endoftext|>


<commit_before>/*
 * diosix microkernel 'menchi'
 *
 * Handle interrupts for x86 systems
 *
 * Maintainer: Chris Williams (diosix.org)
 *
 */

use ::hardware::swis;
use ::hardware::pic;
use ::hardware::exceptions;
use errors::KernelInternalError;

const MAX_IDT_ENTRY: usize = 255;

/* define flags in the IDT entries */
const INT_GATE_PRESENT: u8 = 1 << 7; /* gate is present */
const INT_GATE_USER_OK: u8 = 3 << 5; /* allow ring 3 to trigger this gate */
const INT_GATE_TYPE:    u8 = 14;     /* full-fat 32-bit/64-bit interrupt gate */

extern
{
    static kernel_cs: u64; /* kernel code selector */
    static boot_idtr: u64; /* pointer to boot IDTR */
    static mut boot_idt: [idt_entry; MAX_IDT_ENTRY + 1]; /* table of 256 IDT entries */
}

/* when an interrupt happens, we take a snapshot of the
 * running thread's registers and stack them along with
 * the interrupt number and error code. this struct
 * can be used to access the stacked information. */
#[repr(C, packed)]
pub struct interrupted_thread_registers
{
    ds: u64,

    r15: u64, r14: u64, r13: u64, r12: u64, r11: u64,
    r10: u64,  r9: u64,  r8: u64, rdi: u64, rsi: u64,
    rbp: u64, rdx: u64, rcx: u64, rbx: u64, rax: u64,

    interrupt_number: u64,
    error_code: u64,
    rip: u64, cs: u64, flags: u64, rsp: u64, ss: u64
}

/* a 64-bit mode IDT entry */
#[repr(C, packed)]
struct idt_entry
{
    offset_low: u16,        /* bits 0 to 15 of handler address */
    gdt_select: u16,        /* GDT code selector for the handler code */
    reserved_zero_byte: u8, /* must be zero */
    flags: u8,              /* type and attribute bits */
    offset_middle: u16,     /* bits 16 to 31 of handler address */
    offset_high: u32,       /* bits 32 to 63 of handler address */
    reserved_zero_word: u32 /* must be zero */
}

/* init()
 *
 * Initialize the interrupt system with basic exception
 * and interrupt handling.
 * <= returns error code if a failure happens
 *
 */
pub fn init() -> Result<(), KernelInternalError>
{
    try!(exceptions::init());
    try!(pic::init());
    try!(swis::init());

    kprintln!("[x86] using boot interrupt table at {:p} (idtr: {:p})", unsafe{&boot_idt}, &boot_idtr);

    /* load a pointer to the IDT into the CPU and
     * enable interrupts. */
    unsafe
    {
        asm!("lidt (%rax)" : : "{rax}"(&boot_idtr));
        asm!("sti");
    }

    kprintln!("[x86] interrupts and exceptions enabled");

    Ok(())
}

/* set_boot_idt_gate
 *
 * Set the interrupt vector entry for the boot IDT.
 * This is set up early on before the rest of the system
 * is initialized.
 * => vector = vector number to set up (0 to 255)
 *    handler = pointer to low-level interrupt handler
 * <= returns error on failure
 */
pub fn set_boot_idt_gate(vector: usize, handler: unsafe extern "C" fn()) -> Result<(), KernelInternalError>
{
    /* bail out if vector isn't sane */
    if vector > MAX_IDT_ENTRY
    {
        kprintln!("[x86] BUG! set_idt_gate() called with bad vector {}", vector);
        return Err(KernelInternalError::BadIndex);
    }

    let handler_addr = handler as u64;

    /* this bit is unsafe because we're fiddling with a global mutable variable.
     * but this function should only be called by the boot processor core
     * during system startup, so therefore no races. */
    unsafe
    {
        let entry           = &mut boot_idt[vector];
        entry.offset_low    = (handler_addr & 0xffff) as u16;   /* use lowest 16 bits of handler address */
        entry.gdt_select    = kernel_cs as u16;                 /* use kernel's code selector for handler */
        entry.reserved_zero_byte = 0;
        entry.flags         = INT_GATE_PRESENT | INT_GATE_TYPE; /* present, interrupt gate, only kernel or hw can trigger */
        entry.offset_middle = ((handler_addr & 0xffff0000) >> 16) as u16; /* bits 16-31 of the handler address */
        entry.offset_high   = ((handler_addr & 0xffffffff00000000) >> 32) as u32; /* top 32 bits of handler address */
        entry.reserved_zero_word = 0;
    }

    Ok(())
}

/* enable_gate_user_access
 *
 * Grant access to given gate from userspace, allowing normal
 * applications to trigger the interrupt. Use to convert interrupt
 * gates into SWIs (software interrupts)
 * => vector = interrupt to convert
 * <= returns error on failure
 */
pub fn enable_gate_user_access(vector: usize) -> Result<(), KernelInternalError>
{
    /* bail out if vector isn't sane */
    if vector > MAX_IDT_ENTRY
    {
        kprintln!("[x86] BUG! enable_gate_user_access() called with bad vector {}", vector);
        return Err(KernelInternalError::BadIndex);
    }

    /* this bit is unsafe because we're fiddling with a global mutable variable.
     * see set_boot_idt_gate() */
    unsafe
    {
        let entry = &mut boot_idt[vector];
        /* check this entry is sane - the present bit (b7) must be set to show it's initialized */
        if (entry.flags & INT_GATE_PRESENT) == 0
        {
            kprintln!("[x86] BUG! enable_gate_user_access() called with non-present vector {}", vector);
            return Err(KernelInternalError::BadIndex);
        }

        entry.flags = entry.flags | INT_GATE_USER_OK; /* allow ring-3 to trigger this interrupt */
    }

    Ok(())
}

/* kernel_interrupt_handler
 *
 * Entry point to the kernel from an interrupt/exception
 * => stack = pointer to stack containing interrupted
 *            thread's registers plus interrupt number and
 *            error code.
 */
#[no_mangle]
pub extern "C" fn kernel_interrupt_handler(stack: interrupted_thread_registers)
{
    kprintln!("[x86] interrupt incoming: {}", stack.interrupt_number);
}

<commit_msg>[x86] add more debug info when an early CPU exception triggers<commit_after>/*
 * diosix microkernel 'menchi'
 *
 * Handle interrupts for x86 systems
 *
 * Maintainer: Chris Williams (diosix.org)
 *
 */

use ::hardware::swis;
use ::hardware::pic;
use ::hardware::exceptions;
use errors::KernelInternalError;

const MAX_IDT_ENTRY: usize = 255;

/* define flags in the IDT entries */
const INT_GATE_PRESENT: u8 = 1 << 7; /* gate is present */
const INT_GATE_USER_OK: u8 = 3 << 5; /* allow ring 3 to trigger this gate */
const INT_GATE_TYPE:    u8 = 14;     /* full-fat 32-bit/64-bit interrupt gate */

extern
{
    static kernel_cs: u64; /* kernel code selector */
    static boot_idtr: u64; /* pointer to boot IDTR */
    static mut boot_idt: [idt_entry; MAX_IDT_ENTRY + 1]; /* table of 256 IDT entries */
}

/* when an interrupt happens, we take a snapshot of the
 * running thread's registers and stack them along with
 * the interrupt number and error code. this struct
 * can be used to access the stacked information. */
#[repr(C, packed)]
pub struct interrupted_thread_registers
{
    ds: u64,

    r15: u64, r14: u64, r13: u64, r12: u64, r11: u64,
    r10: u64,  r9: u64,  r8: u64, rdi: u64, rsi: u64,
    rbp: u64, rdx: u64, rcx: u64, rbx: u64, rax: u64,

    interrupt_number: u64,
    error_code: u64,
    rip: u64, cs: u64, flags: u64, rsp: u64, ss: u64
}

/* a 64-bit mode IDT entry */
#[repr(C, packed)]
struct idt_entry
{
    offset_low: u16,        /* bits 0 to 15 of handler address */
    gdt_select: u16,        /* GDT code selector for the handler code */
    reserved_zero_byte: u8, /* must be zero */
    flags: u8,              /* type and attribute bits */
    offset_middle: u16,     /* bits 16 to 31 of handler address */
    offset_high: u32,       /* bits 32 to 63 of handler address */
    reserved_zero_word: u32 /* must be zero */
}

/* init()
 *
 * Initialize the interrupt system with basic exception
 * and interrupt handling.
 * <= returns error code if a failure happens
 *
 */
pub fn init() -> Result<(), KernelInternalError>
{
    try!(exceptions::init());
    try!(pic::init());
    try!(swis::init());

    kprintln!("[x86] using boot interrupt table at {:p} (idtr: {:p})", unsafe{&boot_idt}, &boot_idtr);

    /* load a pointer to the IDT into the CPU and
     * enable interrupts. */
    unsafe
    {
        asm!("lidt (%rax)" : : "{rax}"(&boot_idtr));
        asm!("sti");
    }

    kprintln!("[x86] interrupts and exceptions enabled");

    Ok(())
}

/* set_boot_idt_gate
 *
 * Set the interrupt vector entry for the boot IDT.
 * This is set up early on before the rest of the system
 * is initialized.
 * => vector = vector number to set up (0 to 255)
 *    handler = pointer to low-level interrupt handler
 * <= returns error on failure
 */
pub fn set_boot_idt_gate(vector: usize, handler: unsafe extern "C" fn()) -> Result<(), KernelInternalError>
{
    /* bail out if vector isn't sane */
    if vector > MAX_IDT_ENTRY
    {
        kprintln!("[x86] BUG! set_idt_gate() called with bad vector {}", vector);
        return Err(KernelInternalError::BadIndex);
    }

    let handler_addr = handler as u64;

    /* this bit is unsafe because we're fiddling with a global mutable variable.
     * but this function should only be called by the boot processor core
     * during system startup, so therefore no races. */
    unsafe
    {
        let entry           = &mut boot_idt[vector];
        entry.offset_low    = (handler_addr & 0xffff) as u16;   /* use lowest 16 bits of handler address */
        entry.gdt_select    = kernel_cs as u16;                 /* use kernel's code selector for handler */
        entry.reserved_zero_byte = 0;
        entry.flags         = INT_GATE_PRESENT | INT_GATE_TYPE; /* present, interrupt gate, only kernel or hw can trigger */
        entry.offset_middle = ((handler_addr & 0xffff0000) >> 16) as u16; /* bits 16-31 of the handler address */
        entry.offset_high   = ((handler_addr & 0xffffffff00000000) >> 32) as u32; /* top 32 bits of handler address */
        entry.reserved_zero_word = 0;
    }

    Ok(())
}

/* enable_gate_user_access
 *
 * Grant access to given gate from userspace, allowing normal
 * applications to trigger the interrupt. Use to convert interrupt
 * gates into SWIs (software interrupts)
 * => vector = interrupt to convert
 * <= returns error on failure
 */
pub fn enable_gate_user_access(vector: usize) -> Result<(), KernelInternalError>
{
    /* bail out if vector isn't sane */
    if vector > MAX_IDT_ENTRY
    {
        kprintln!("[x86] BUG! enable_gate_user_access() called with bad vector {}", vector);
        return Err(KernelInternalError::BadIndex);
    }

    /* this bit is unsafe because we're fiddling with a global mutable variable.
     * see set_boot_idt_gate() */
    unsafe
    {
        let entry = &mut boot_idt[vector];
        /* check this entry is sane - the present bit (b7) must be set to show it's initialized */
        if (entry.flags & INT_GATE_PRESENT) == 0
        {
            kprintln!("[x86] BUG! enable_gate_user_access() called with non-present vector {}", vector);
            return Err(KernelInternalError::BadIndex);
        }

        entry.flags = entry.flags | INT_GATE_USER_OK; /* allow ring-3 to trigger this interrupt */
    }

    Ok(())
}

/* kernel_interrupt_handler
 *
 * Entry point to the kernel from an interrupt/exception
 * => stack = pointer to stack containing interrupted
 *            thread's registers plus interrupt number and
 *            error code.
 */
#[no_mangle]
pub extern "C" fn kernel_interrupt_handler(stack: interrupted_thread_registers)
{
    if stack.interrupt_number < 32
    {
        let rip: u64 = stack.rip;
        let cr2: u64;

        unsafe
        {
            asm!("mov %cr2, %rax" : "={rax}"(cr2));
        }

        kprintln!("[x86] CPU exception {}: rip = {:x} fault addr (cr2) = {:x}", stack.interrupt_number, rip, cr2);
        panic!("Unhandled exception");
    }

    kprintln!("[x86] interrupt fired: {}", stack.interrupt_number);
}

<|endoftext|>


<commit_before>//! Utilities not directly related to parsing, lexing, or
//! language-manipulation.

/// Panic with a given message unless an expression evaluates to true.
///
/// ## Examples
/// ```rust
/// # #[macro_use] extern crate cripes;
/// # fn main() {
/// panic_unless!(1 + 1 == 2, "Math is broken.");
/// # }
/// ```
///
/// ```rust,should_panic
/// # #[macro_use] extern crate cripes;
/// # fn main() {
/// panic_unless!(true == false, "I'm broken.");
/// # }
/// ```
#[macro_export]
macro_rules! panic_unless {
    ($condition:expr, $($rest:expr),+) => ({ if ! $condition { panic!($($rest),+); } });
}

pub mod set;
pub mod hash;
pub mod flex_box;
pub mod graph;
pub mod splitvec;
<commit_msg>util: macro `panic_unless_eq` already used in tests, might as well add it<commit_after>//! Utilities not directly related to parsing, lexing, or
//! language-manipulation.

/// Panic with a given message unless an expression evaluates to true.
///
/// ## Examples
/// ```rust
/// # #[macro_use] extern crate cripes;
/// # fn main() {
/// panic_unless!(1 + 1 == 2, "Math is broken.");
/// # }
/// ```
///
/// ```rust,should_panic
/// # #[macro_use] extern crate cripes;
/// # fn main() {
/// panic_unless!(true == false, "I'm broken.");
/// # }
/// ```
#[macro_export]
macro_rules! panic_unless {
    ($condition:expr, $($rest:expr),+) => ({ if ! $condition { panic!($($rest),+); } });
}

/// Panic with a given message unless an expression evaluates to true.
#[macro_export]
macro_rules! panic_unless_eq {
    ($expected: expr, $actual: expr) => ({
        match (&$expected, &$actual) {
            (expected, actual) => if !(*actual == *expected) {
                panic!("unexpected value: `{:?}`, expected `{:?}`", actual, expected)
            }
        }
    });
    ($expected: expr, $actual: expr, $($rest:expr),+) => ({
        match (&$expected, &$actual) {
            (expected, actual) => if !(*actual == *expected) {
                panic!("{}: `{:?}`, expected `{:?}`", format_args!($($rest),+), actual, expected)
            }
        }
    })
}

pub mod set;
pub mod hash;
pub mod flex_box;
pub mod graph;
pub mod splitvec;
<|endoftext|>