<commit_before>#![feature(const_fn)]
#![feature(catch_panic)] // For testing

use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering};

/// A cell holding values protected by an atomic lock.
///
/// This cell is designed to be instantiated as a global variable, to
/// hold a single value and to distribute it to threads as needed.  It
/// was initially designed to distribute clones of `Sender` to
/// hundreds of clients across dozens of modules without having to
/// pass these senders as argument through hundreds of intermediate
/// functions.
///
/// # Example
///
/// ```
/// #![feature(const_fn)]
/// use std::thread;
/// use std::sync::mpsc::{channel, Sender};
/// use std::ops::Add;
///
/// use atomic_cell::StaticCell;
///
/// // Sender cannot be defined as a `static` for two reasons:
/// // - it does not implement `Sync`;
/// // - it has a destructor.
/// // So let's implement it as a `StaticCell`.
/// static CELL: StaticCell<Sender<u32>> = StaticCell::new();
///
/// fn main() {
///   let (sender, receiver) = channel();
///   let _guard = CELL.init(sender);
///   // From this point and until `_guard` is dropped, `CELL` owns `sender`.
///
///   for i in 0..10 {
///     thread::spawn(move || {
///       // Any thread can now access clones of `sender`.
///       let sender = CELL.get().unwrap();
///       sender.send(i).unwrap();
///     });
///   }
///
///   // Make sure that all data was received properly.
///   assert_eq!(receiver.iter().take(10).fold(0, Add::add), 45);
///
///   // When we leave `main()`, `_guard` will go out of scope and
///   // will be dropped. This will cause `sender` to be dropped.
/// }
/// ```
///
///
/// # Performance
///
/// This cell is optimized for low contention. If there is no
/// contention, each call to `get` is resolved as a single atomic
/// read. In presence of high contention on a single cell, though,
/// performance may degrade considerably.
///
///
///
/// # Resource-safety
///
/// Unlike other kinds of static holders, values in `StaticCell` are
/// scoped and will be dropped, once the guard is dropped.
///
///
///
pub struct StaticCell<T> where T: Clone {
    internal: InternalAtomicCell<T>,
    initialized: AtomicBool
}

impl<T> StaticCell<T> where T: Clone {
    ///
    /// Create an empty cell.
    ///
    /// Call `init` to fill the cell.
    ///
    pub const fn new() -> Self {
        StaticCell {
            initialized: AtomicBool::new(false),
            internal: InternalAtomicCell::const_new()
        }
    }

    /// Initialize the cell.
    ///
    /// This methods returns a guard, which will drop `value` once it is dropped itself.
    /// Once this is done, `self` will return to being an empty cell. It will, however,
    /// remain initialized and cannot ever be initialized again.
    ///
    /// # Panics
    ///
    /// This method panicks if the cell is already initialized.
    pub fn init<'a>(&'a self, value: T) -> CleanGuard<'a> {
        let initialized = self.initialized.compare_and_swap(false, true, Ordering::Relaxed);
        if initialized {
            panic!("StaticCell is already initialized.");
        }
        self.internal.set(value);
        return CleanGuard {
            item: self
        }
    }

    /// Get a clone of the value held by the cell.
    ///
    /// Returns `None` if the cell is empty, either because it is not
    /// initialized or because the guard has been dropped.
    ///
    /// Receiving `None` is almost always a programming error, so
    /// client code is encouraged to `unwrap()` immediately.
    ///
    /// # Panics
    ///
    /// This method panicks if the call to `value.clone()` causes a
    /// panic.  However, the cell remains usable.
    pub fn get(&self) -> Option<Box<T>> {
        self.internal.get()
    }
}

trait CleanMeUp {
    fn clean(&self);
}

/// A guard used to drop the value held by a `StaticCell` at a
/// deterministic point in code.
///
/// Once the CleanGuard returned by `init()` is dropped, the value
/// held by the cell is also dropped.
pub struct CleanGuard<'a> {
    item: &'a CleanMeUp
}

impl<T> CleanMeUp for StaticCell<T> where T: Clone {
    ///
    /// Perform any cleanup.
    ///
    fn clean(&self) {
        self.internal.unset();
    }
}

impl<'a> Drop for CleanGuard<'a> {
    fn drop(&mut self) {
        self.item.clean();
    }
}

/// A cell holding values protected by an atomic lock.
///
/// This cell is designed to be instantiated as a local variable, to
/// hold a single value and to distribute it to threads as needed.
///
/// This cell cannot be allocated as a global variable. If you need a
/// global variable, use `StaticAtomicCell`.
///
///
/// # Performance
///
/// This cell is optimized for low contention. If there is no
/// contention, each call to `get` is resolved as a single atomic
/// read. In presence of high contention on a single cell, though,
/// performance may degrade considerably.
///
pub struct AtomicCell<T> where T: Clone {
    internal: InternalAtomicCell<T>
}

impl<T> AtomicCell<T> where T: Clone {
    ///
    /// Create a new empty cell.
    ///
    /// Use `set` or `swap` to add contents.
    ///
    pub fn new() -> Self {
        AtomicCell {
            internal: InternalAtomicCell::new()
        }
    }

    ///
    /// Set the contents of the cell.
    ///
    /// `value` will be dropped either when the cell is dropped, or
    /// when `set` is called once again. Property of `value` is
    /// transferred to the client if `swap` is called.
    ///
    /// If the cell already held some contents, drop these contents.
    ///
    pub fn set(&self, value: T) {
        self.internal.set(value)
    }

    /// Get a clone of the value held by the cell.
    ///
    /// Returns `None` if the cell is empty.
    ///
    /// # Panics
    ///
    /// A panic during the call to `value.clone()` will propagate and
    /// cause a panic in the cell.  However, the cell will remain
    /// usable.
    pub fn get(&self) -> Option<Box<T>> {
        self.internal.get()
    }

    ///
    /// Empty the cell manually.
    ///
    /// If the cell was empty, this is a noop. Otherwise, the previous
    /// value held is dropped.
    ///
    pub fn unset(&self) {
        self.internal.unset()
    }

    ///
    /// Swap the value held by the cell with a new value.
    ///
    pub fn swap(&self, value: Option<Box<T>>) -> Option<Box<T>> {
        self.internal.swap(value)
    }
}

impl<T> Drop for AtomicCell<T> where T: Clone {
    ///
    /// Drop any content present in the cell.
    ///
    fn drop(&mut self) {
        self.unset();
    }
}

/// A cell holding values protected by an atomic lock.
///
/// This cell is designed to hold values that implement `Clone` and to
/// distribute them to threads as needed. It may be instantiated as
/// a `static` value.
///
/// This cell is optimized for low contention.
///
/// This version does NOT guarantee that the values it holds are
/// dropped.
///
impl<T> InternalAtomicCell<T> where T: Clone {
    ///
    /// Create an empty cell.
    ///
    pub const fn const_new() -> Self {
        InternalAtomicCell {
            ptr: AtomicPtr::new(std::ptr::null_mut()),
            lock: AtomicBool::new(true),
        }
    }
    pub fn new() -> Self {
        InternalAtomicCell {
            ptr: AtomicPtr::new(std::ptr::null_mut()),
            lock: AtomicBool::new(true),
        }
    }

    ///
    /// Put a value in the cell.
    ///
    /// If there was a previous value, it is dropped.
    ///
    /// Argument `value` will **not** be dropped automatically. You should
    /// call `set`, `unset` or `swap` to ensure that it is dropped.
    ///
    /// See also `swap`.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    pub fn set(&self, value: T) {
        self.swap(Some(Box::new(value)));
    }

    ///
    /// Remove whichever value is in the cell.
    ///
    /// If there was a previous value, it is dropped.
    ///
    /// Argument `value` will **not** be dropped automatically. You should
    /// call `set`, `unset` or `swap` to ensure that it is dropped.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    pub fn unset(&self) {
        self.swap(None);
    }

    ///
    /// Get a clone of the current value in the cell.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    /// # Panics
    ///
    /// If the `clone` method panics, this will cause a panic. The cell
    /// will, however, remain usable.
    ///
    pub fn get(&self) -> Option<Box<T>> {
        self.with_lock(|| {
            let ptr = self.ptr.load(Ordering::Relaxed);

            if ptr.is_null() {
                return None;
            }

            // Don't cast back to Box, as this would cause us to
            // `drop` the value in case of panic.
            let clone = unsafe { (*ptr).clone() };
            // If `clone` panics, `with_lock` will ensure that the
            // lock is released.

            Some(Box::new(clone))
        })
    }

    ///
    /// Replace the value currently in the cell with another value.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    pub fn swap(&self, value: Option<Box<T>>) -> Option<Box<T>> {
        let new_ptr = ptr_from_opt(value);
        // We are now the owner of `value` as a pointer. From this
        // point, we are in charge of dropping it manually.
        self.with_lock(|| {
            let old_ptr = self.ptr.swap(new_ptr, Ordering::Relaxed);
            opt_from_ptr(old_ptr)
        })
    }

    ///
    /// Acquire the lock.
    ///
    /// We use an atomic spinlock.
    ///
    /// # Panics
    ///
    /// If `f` panics.
    ///
    fn with_lock<F, R>(&self, f: F) -> R where F: FnOnce() -> R {
        loop {
            // Attempt to acquire the lock.
            // This data structure is designed for low contention, so we can use
            // an atomic spinlock.
            let owning = self.lock.compare_and_swap(/*must be available*/true,
                                                    /*mark as unavailable*/false, Ordering::Relaxed); // FIXME: Check ordering.
            if owning {
                // We are now the owner of the lock.

                // Make sure that we eventually release the lock.
                let _guard = GuardLock::new(&self.lock);

                let result = f();

                return result;
            }
        }
    }
}

fn ptr_from_opt<T>(value: Option<Box<T>>) -> *mut T {
    match value {
        None => std::ptr::null_mut(),
        Some(b) => Box::into_raw(b)
    }
}

fn opt_from_ptr<T>(ptr: *mut T) -> Option<Box<T>> {
    if ptr.is_null() {
        None
    } else {
        unsafe { Some(Box::from_raw(ptr)) }
    }
}

struct InternalAtomicCell<T> where T: Clone {
    ptr: AtomicPtr<T>,
    lock: AtomicBool,
}

unsafe impl<T> Sync for InternalAtomicCell<T> where T: Clone + Send {
}

// A guard used to ensure that we release a lock, even in case of
// panic.
struct GuardLock<'a> {
    lock: &'a AtomicBool
}
impl<'a> GuardLock<'a> {
    fn new(lock: &AtomicBool) -> GuardLock {
        GuardLock {
            lock: lock
        }
    }
}
impl<'a> Drop for GuardLock<'a> {
    fn drop(&mut self) {
        self.lock.swap(true, Ordering::Relaxed);
    }
}

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

    use std::ops::Add;
    use std::sync::mpsc::{channel, Sender};
    use std::thread;
    use std::sync::atomic::{AtomicBool, Ordering};

    // Test that we can allocate a cell and use it to distribute value
    // to several threads.
    static CELL: StaticCell<Sender<u32>> = StaticCell::new();
    #[test]
    fn test_channels() {
        let (tx, rx) = channel();
        let _guard = CELL.init(tx);
        for i in 0..10 {
            thread::spawn(move || {
                let tx = CELL.get().unwrap();
                tx.send(i).unwrap();
            });
        }
        assert_eq!(rx.iter().take(10).fold(0, Add::add), 45);
    }

    // Test that `get()` on an empty cell returns `None`.
    static CELL2: StaticCell<u32> = StaticCell::new();
    #[test]
    fn test_empty() {
        for _ in 0..10 {
            thread::spawn(move || {
                assert_eq!(CELL2.get(), None);
            });
        }
        assert_eq!(CELL2.get(), None);
    }

    // Test that a panic does not make the cell unusable.
    static mut should_panic: AtomicBool = AtomicBool::new(false);
    struct Panicky {
        foo: u32
    }
    impl Clone for Panicky {
        fn clone(&self) -> Self {
            unsafe {
                if should_panic.load(Ordering::Relaxed) {
                    panic!("I have panicked, as expected");
                }
            }
            Panicky {
                foo: self.foo
            }
        }
    }

    static CELL_PANIC: StaticCell<Panicky> = StaticCell::new();
    #[test]
    fn test_panic() {
        let original = Panicky { foo: 500 };
        let _guard = CELL_PANIC.init(original.clone());

        // Now, cause a panic.
        unsafe {
            should_panic.swap(true, Ordering::Relaxed);
        }
        let panic = thread::catch_panic(|| {
            // This should panic.
            CELL_PANIC.get()
        });
        assert!(panic.is_err());

        // Now stop the panic.
        unsafe {
            should_panic.swap(false, Ordering::Relaxed);
        }

        // This shouldn't panic. Moreover, we should find
        // our original value.
        let result = CELL_PANIC.get().unwrap();
        assert_eq!(result.foo, original.foo);
    }

    static CELL_INIT: StaticCell<u32> = StaticCell::new();
    #[test]
    fn test_init() {
        {
            let mut _guard = CELL_INIT.init(0);
            CELL_INIT.get().unwrap();
        }
        assert_eq!(CELL_INIT.get(), None);
    }
}
<commit_msg>Checked ordering<commit_after>#![feature(const_fn)]
#![feature(catch_panic)] // For testing

use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering};

/// A cell holding values protected by an atomic lock.
///
/// This cell is designed to be instantiated as a global variable, to
/// hold a single value and to distribute it to threads as needed.  It
/// was initially designed to distribute clones of `Sender` to
/// hundreds of clients across dozens of modules without having to
/// pass these senders as argument through hundreds of intermediate
/// functions.
///
/// # Example
///
/// ```
/// #![feature(const_fn)]
/// use std::thread;
/// use std::sync::mpsc::{channel, Sender};
/// use std::ops::Add;
///
/// use atomic_cell::StaticCell;
///
/// // Sender cannot be defined as a `static` for two reasons:
/// // - it does not implement `Sync`;
/// // - it has a destructor.
/// // So let's implement it as a `StaticCell`.
/// static CELL: StaticCell<Sender<u32>> = StaticCell::new();
///
/// fn main() {
///   let (sender, receiver) = channel();
///   let _guard = CELL.init(sender);
///   // From this point and until `_guard` is dropped, `CELL` owns `sender`.
///
///   for i in 0..10 {
///     thread::spawn(move || {
///       // Any thread can now access clones of `sender`.
///       let sender = CELL.get().unwrap();
///       sender.send(i).unwrap();
///     });
///   }
///
///   // Make sure that all data was received properly.
///   assert_eq!(receiver.iter().take(10).fold(0, Add::add), 45);
///
///   // When we leave `main()`, `_guard` will go out of scope and
///   // will be dropped. This will cause `sender` to be dropped.
/// }
/// ```
///
///
/// # Performance
///
/// This cell is optimized for low contention. If there is no
/// contention, each call to `get` is resolved as a single atomic
/// read. In presence of high contention on a single cell, though,
/// performance may degrade considerably.
///
///
///
/// # Resource-safety
///
/// Unlike other kinds of static holders, values in `StaticCell` are
/// scoped and will be dropped, once the guard is dropped.
///
///
///
pub struct StaticCell<T> where T: Clone {
    internal: InternalAtomicCell<T>,
    initialized: AtomicBool
}

impl<T> StaticCell<T> where T: Clone {
    ///
    /// Create an empty cell.
    ///
    /// Call `init` to fill the cell.
    ///
    pub const fn new() -> Self {
        StaticCell {
            initialized: AtomicBool::new(false),
            internal: InternalAtomicCell::const_new()
        }
    }

    /// Initialize the cell.
    ///
    /// This methods returns a guard, which will drop `value` once it is dropped itself.
    /// Once this is done, `self` will return to being an empty cell. It will, however,
    /// remain initialized and cannot ever be initialized again.
    ///
    /// # Panics
    ///
    /// This method panicks if the cell is already initialized.
    pub fn init<'a>(&'a self, value: T) -> CleanGuard<'a> {
        let initialized = self.initialized.compare_and_swap(false, true, Ordering::SeqCst);
        if initialized {
            panic!("StaticCell is already initialized.");
        }
        self.internal.set(value);
        return CleanGuard {
            item: self
        }
    }

    /// Get a clone of the value held by the cell.
    ///
    /// Returns `None` if the cell is empty, either because it is not
    /// initialized or because the guard has been dropped.
    ///
    /// Receiving `None` is almost always a programming error, so
    /// client code is encouraged to `unwrap()` immediately.
    ///
    /// # Panics
    ///
    /// This method panicks if the call to `value.clone()` causes a
    /// panic.  However, the cell remains usable.
    pub fn get(&self) -> Option<Box<T>> {
        self.internal.get()
    }
}

trait CleanMeUp {
    fn clean(&self);
}

/// A guard used to drop the value held by a `StaticCell` at a
/// deterministic point in code.
///
/// Once the CleanGuard returned by `init()` is dropped, the value
/// held by the cell is also dropped.
pub struct CleanGuard<'a> {
    item: &'a CleanMeUp
}

impl<T> CleanMeUp for StaticCell<T> where T: Clone {
    ///
    /// Perform any cleanup.
    ///
    fn clean(&self) {
        self.internal.unset();
    }
}

impl<'a> Drop for CleanGuard<'a> {
    fn drop(&mut self) {
        self.item.clean();
    }
}

/// A cell holding values protected by an atomic lock.
///
/// This cell is designed to be instantiated as a local variable, to
/// hold a single value and to distribute it to threads as needed.
///
/// This cell cannot be allocated as a global variable. If you need a
/// global variable, use `StaticAtomicCell`.
///
///
/// # Performance
///
/// This cell is optimized for low contention. If there is no
/// contention, each call to `get` is resolved as a single atomic
/// read. In presence of high contention on a single cell, though,
/// performance may degrade considerably.
///
pub struct AtomicCell<T> where T: Clone {
    internal: InternalAtomicCell<T>
}

impl<T> AtomicCell<T> where T: Clone {
    ///
    /// Create a new empty cell.
    ///
    /// Use `set` or `swap` to add contents.
    ///
    pub fn new() -> Self {
        AtomicCell {
            internal: InternalAtomicCell::new()
        }
    }

    ///
    /// Set the contents of the cell.
    ///
    /// `value` will be dropped either when the cell is dropped, or
    /// when `set` is called once again. Property of `value` is
    /// transferred to the client if `swap` is called.
    ///
    /// If the cell already held some contents, drop these contents.
    ///
    pub fn set(&self, value: T) {
        self.internal.set(value)
    }

    /// Get a clone of the value held by the cell.
    ///
    /// Returns `None` if the cell is empty.
    ///
    /// # Panics
    ///
    /// A panic during the call to `value.clone()` will propagate and
    /// cause a panic in the cell.  However, the cell will remain
    /// usable.
    pub fn get(&self) -> Option<Box<T>> {
        self.internal.get()
    }

    ///
    /// Empty the cell manually.
    ///
    /// If the cell was empty, this is a noop. Otherwise, the previous
    /// value held is dropped.
    ///
    pub fn unset(&self) {
        self.internal.unset()
    }

    ///
    /// Swap the value held by the cell with a new value.
    ///
    pub fn swap(&self, value: Option<Box<T>>) -> Option<Box<T>> {
        self.internal.swap(value)
    }
}

impl<T> Drop for AtomicCell<T> where T: Clone {
    ///
    /// Drop any content present in the cell.
    ///
    fn drop(&mut self) {
        self.unset();
    }
}

/// A cell holding values protected by an atomic lock.
///
/// This cell is designed to hold values that implement `Clone` and to
/// distribute them to threads as needed. It may be instantiated as
/// a `static` value.
///
/// This cell is optimized for low contention.
///
/// This version does NOT guarantee that the values it holds are
/// dropped.
///
impl<T> InternalAtomicCell<T> where T: Clone {
    ///
    /// Create an empty cell.
    ///
    pub const fn const_new() -> Self {
        InternalAtomicCell {
            ptr: AtomicPtr::new(std::ptr::null_mut()),
            lock: AtomicBool::new(true),
        }
    }
    pub fn new() -> Self {
        InternalAtomicCell {
            ptr: AtomicPtr::new(std::ptr::null_mut()),
            lock: AtomicBool::new(true),
        }
    }

    ///
    /// Put a value in the cell.
    ///
    /// If there was a previous value, it is dropped.
    ///
    /// Argument `value` will **not** be dropped automatically. You should
    /// call `set`, `unset` or `swap` to ensure that it is dropped.
    ///
    /// See also `swap`.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    pub fn set(&self, value: T) {
        self.swap(Some(Box::new(value)));
    }

    ///
    /// Remove whichever value is in the cell.
    ///
    /// If there was a previous value, it is dropped.
    ///
    /// Argument `value` will **not** be dropped automatically. You should
    /// call `set`, `unset` or `swap` to ensure that it is dropped.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    pub fn unset(&self) {
        self.swap(None);
    }

    ///
    /// Get a clone of the current value in the cell.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    /// # Panics
    ///
    /// If the `clone` method panics, this will cause a panic. The cell
    /// will, however, remain usable.
    ///
    pub fn get(&self) -> Option<Box<T>> {
        self.with_lock(|| {
            let ptr = self.ptr.load(Ordering::Relaxed);

            if ptr.is_null() {
                return None;
            }

            // Don't cast back to Box, as this would cause us to
            // `drop` the value in case of panic.
            let clone = unsafe { (*ptr).clone() };
            // If `clone` panics, `with_lock` will ensure that the
            // lock is released.

            Some(Box::new(clone))
        })
    }

    ///
    /// Replace the value currently in the cell with another value.
    ///
    /// # Performance
    ///
    /// This method requires acquiring an atomic lock. If contention
    /// is low, it should be very fast. However, in case of high
    /// contention, performance may degrade considerably.
    ///
    pub fn swap(&self, value: Option<Box<T>>) -> Option<Box<T>> {
        let new_ptr = ptr_from_opt(value);
        // We are now the owner of `value` as a pointer. From this
        // point, we are in charge of dropping it manually.
        self.with_lock(|| {
            // Since we are protected by the outer lock, a relaxed
            // order is sufficient here.
            let old_ptr = self.ptr.swap(new_ptr, Ordering::Relaxed);
            opt_from_ptr(old_ptr)
        })
    }

    ///
    /// Acquire the lock.
    ///
    /// We use an atomic spinlock.
    ///
    /// # Panics
    ///
    /// If `f` panics.
    ///
    fn with_lock<F, R>(&self, f: F) -> R where F: FnOnce() -> R {
        loop {
            // Attempt to acquire the lock.
            // This data structure is designed for low contention, so we can use
            // an atomic spinlock.
            let owning = self.lock.compare_and_swap(/*must be available*/true,
                                                    /*mark as unavailable*/false,
                                                    Ordering::SeqCst);
            if owning {
                // We are now the owner of the lock.

                // Make sure that we eventually release the lock.
                let _guard = GuardLock::new(&self.lock);

                let result = f();

                return result;
            }
        }
    }
}

fn ptr_from_opt<T>(value: Option<Box<T>>) -> *mut T {
    match value {
        None => std::ptr::null_mut(),
        Some(b) => Box::into_raw(b)
    }
}

fn opt_from_ptr<T>(ptr: *mut T) -> Option<Box<T>> {
    if ptr.is_null() {
        None
    } else {
        unsafe { Some(Box::from_raw(ptr)) }
    }
}

struct InternalAtomicCell<T> where T: Clone {
    ptr: AtomicPtr<T>,
    lock: AtomicBool,
}

unsafe impl<T> Sync for InternalAtomicCell<T> where T: Clone + Send {
}

// A guard used to ensure that we release a lock, even in case of
// panic.
struct GuardLock<'a> {
    lock: &'a AtomicBool
}
impl<'a> GuardLock<'a> {
    fn new(lock: &AtomicBool) -> GuardLock {
        GuardLock {
            lock: lock
        }
    }
}
impl<'a> Drop for GuardLock<'a> {
    fn drop(&mut self) {
        self.lock.swap(true, Ordering::Relaxed);
    }
}

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

    use std::ops::Add;
    use std::sync::mpsc::{channel, Sender};
    use std::thread;
    use std::sync::atomic::{AtomicBool, Ordering};

    // Test that we can allocate a cell and use it to distribute value
    // to several threads.
    static CELL: StaticCell<Sender<u32>> = StaticCell::new();
    #[test]
    fn test_channels() {
        let (tx, rx) = channel();
        let _guard = CELL.init(tx);
        for i in 0..10 {
            thread::spawn(move || {
                let tx = CELL.get().unwrap();
                tx.send(i).unwrap();
            });
        }
        assert_eq!(rx.iter().take(10).fold(0, Add::add), 45);
    }

    // Test that `get()` on an empty cell returns `None`.
    static CELL2: StaticCell<u32> = StaticCell::new();
    #[test]
    fn test_empty() {
        for _ in 0..10 {
            thread::spawn(move || {
                assert_eq!(CELL2.get(), None);
            });
        }
        assert_eq!(CELL2.get(), None);
    }

    // Test that a panic does not make the cell unusable.
    static mut should_panic: AtomicBool = AtomicBool::new(false);
    struct Panicky {
        foo: u32
    }
    impl Clone for Panicky {
        fn clone(&self) -> Self {
            unsafe {
                if should_panic.load(Ordering::Relaxed) {
                    panic!("I have panicked, as expected");
                }
            }
            Panicky {
                foo: self.foo
            }
        }
    }

    static CELL_PANIC: StaticCell<Panicky> = StaticCell::new();
    #[test]
    fn test_panic() {
        let original = Panicky { foo: 500 };
        let _guard = CELL_PANIC.init(original.clone());

        // Now, cause a panic.
        unsafe {
            should_panic.swap(true, Ordering::Relaxed);
        }
        let panic = thread::catch_panic(|| {
            // This should panic.
            CELL_PANIC.get()
        });
        assert!(panic.is_err());

        // Now stop the panic.
        unsafe {
            should_panic.swap(false, Ordering::Relaxed);
        }

        // This shouldn't panic. Moreover, we should find
        // our original value.
        let result = CELL_PANIC.get().unwrap();
        assert_eq!(result.foo, original.foo);
    }

    static CELL_INIT: StaticCell<u32> = StaticCell::new();
    #[test]
    fn test_init() {
        {
            let mut _guard = CELL_INIT.init(0);
            CELL_INIT.get().unwrap();
        }
        assert_eq!(CELL_INIT.get(), None);
    }
}
<|endoftext|>


<commit_before>//! # The Redox OS Kernel, version 2
//!
//! The Redox OS Kernel is a microkernel that supports `x86_64` systems and
//! provides Unix-like syscalls for primarily Rust applications

//#![deny(warnings)]
#![cfg_attr(feature = "clippy", allow(if_same_then_else))]
#![cfg_attr(feature = "clippy", allow(inline_always))]
#![cfg_attr(feature = "clippy", allow(many_single_char_names))]
#![cfg_attr(feature = "clippy", allow(module_inception))]
#![cfg_attr(feature = "clippy", allow(new_without_default))]
#![cfg_attr(feature = "clippy", allow(not_unsafe_ptr_arg_deref))]
#![cfg_attr(feature = "clippy", allow(or_fun_call))]
#![cfg_attr(feature = "clippy", allow(too_many_arguments))]
#![feature(alloc)]
#![feature(allocator_api)]
#![feature(asm)]
#![feature(concat_idents)]
#![feature(const_fn)]
#![feature(core_intrinsics)]
#![feature(integer_atomics)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(never_type)]
#![feature(ptr_internals)]
#![feature(thread_local)]
#![no_std]

pub extern crate x86;

#[macro_use]
extern crate alloc;

#[macro_use]
extern crate bitflags;
extern crate goblin;
extern crate linked_list_allocator;
extern crate rustc_demangle;
extern crate spin;
#[cfg(feature = "slab")]
extern crate slab_allocator;

use alloc::vec::Vec;
use core::sync::atomic::{AtomicUsize, Ordering};

use scheme::{FileHandle, SchemeNamespace};

pub use consts::*;

#[macro_use]
/// Shared data structures
pub mod common;

/// Architecture-dependent stuff
#[macro_use]
pub mod arch;
pub use arch::*;

/// Constants like memory locations
pub mod consts;

/// Heap allocators
pub mod allocator;

/// ACPI table parsing
#[cfg(feature = "acpi")]
mod acpi;

/// Context management
pub mod context;

/// Architecture-independent devices
pub mod devices;

/// ELF file parsing
#[cfg(not(feature="doc"))]
pub mod elf;

/// Event handling
pub mod event;

/// External functions
pub mod externs;

/// Logging
pub mod log;

/// Memory management
pub mod memory;

/// Panic
#[cfg(not(any(feature="doc", test)))]
pub mod panic;

/// Schemes, filesystem handlers
pub mod scheme;

/// Synchronization primitives
pub mod sync;

/// Syscall handlers
pub mod syscall;

/// Time
pub mod time;

/// Tests
#[cfg(test)]
pub mod tests;

#[global_allocator]
static ALLOCATOR: allocator::Allocator = allocator::Allocator;

/// A unique number that identifies the current CPU - used for scheduling
#[thread_local]
static CPU_ID: AtomicUsize = AtomicUsize::new(0);

/// Get the current CPU's scheduling ID
#[inline(always)]
pub fn cpu_id() -> usize {
    CPU_ID.load(Ordering::Relaxed)
}

/// The count of all CPUs that can have work scheduled
static CPU_COUNT : AtomicUsize = AtomicUsize::new(0);

/// Get the number of CPUs currently active
#[inline(always)]
pub fn cpu_count() -> usize {
    CPU_COUNT.load(Ordering::Relaxed)
}

static mut INIT_ENV: &[u8] = &[];

/// Initialize userspace by running the initfs:bin/init process
/// This function will also set the CWD to initfs:bin and open debug: as stdio
pub extern fn userspace_init() {
    let path = b"initfs:/bin/init";
    let env = unsafe { INIT_ENV };

    assert_eq!(syscall::chdir(b"initfs:"), Ok(0));

    assert_eq!(syscall::open(b"debug:", syscall::flag::O_RDONLY).map(FileHandle::into), Ok(0));
    assert_eq!(syscall::open(b"debug:", syscall::flag::O_WRONLY).map(FileHandle::into), Ok(1));
    assert_eq!(syscall::open(b"debug:", syscall::flag::O_WRONLY).map(FileHandle::into), Ok(2));

    let fd = syscall::open(path, syscall::flag::O_RDONLY).expect("failed to open init");

    let mut args = Vec::new();
    args.push(path.to_vec().into_boxed_slice());

    let mut vars = Vec::new();
    for var in env.split(|b| *b == b'\n') {
        if ! var.is_empty() {
            vars.push(var.to_vec().into_boxed_slice());
        }
    }

    syscall::fexec_kernel(fd, args.into_boxed_slice(), vars.into_boxed_slice(), None).expect("failed to execute init");

    panic!("init returned");
}

/// This is the kernel entry point for the primary CPU. The arch crate is responsible for calling this
pub fn kmain(cpus: usize, env: &'static [u8]) -> ! {
    CPU_ID.store(0, Ordering::SeqCst);
    CPU_COUNT.store(cpus, Ordering::SeqCst);
    unsafe { INIT_ENV = env };

    //Initialize the first context, stored in kernel/src/context/mod.rs
    context::init();

    let pid = syscall::getpid();
    println!("BSP: {:?} {}", pid, cpus);
    println!("Env: {:?}", ::core::str::from_utf8(env));

    match context::contexts_mut().spawn(userspace_init) {
        Ok(context_lock) => {
            let mut context = context_lock.write();
            context.rns = SchemeNamespace::from(1);
            context.ens = SchemeNamespace::from(1);
            context.status = context::Status::Runnable;
        },
        Err(err) => {
            panic!("failed to spawn userspace_init: {:?}", err);
        }
    }

    loop {
        unsafe {
            interrupt::disable();
            if context::switch() {
                interrupt::enable_and_nop();
            } else {
                // Enable interrupts, then halt CPU (to save power) until the next interrupt is actually fired.
                interrupt::enable_and_halt();
            }
        }
    }
}

/// This is the main kernel entry point for secondary CPUs
#[allow(unreachable_code, unused_variables)]
pub fn kmain_ap(id: usize) -> ! {
    CPU_ID.store(id, Ordering::SeqCst);

    if cfg!(feature = "multi_core") {
        context::init();

        let pid = syscall::getpid();
        println!("AP {}: {:?}", id, pid);

        loop {
            unsafe {
                interrupt::disable();
                if context::switch() {
                    interrupt::enable_and_nop();
                } else {
                    // Enable interrupts, then halt CPU (to save power) until the next interrupt is actually fired.
                    interrupt::enable_and_halt();
                }
            }
        }
    } else {
        println!("AP {}: Disabled", id);

        loop {
            unsafe {
                interrupt::disable();
                interrupt::halt();
            }
        }
    }
}

/// Allow exception handlers to send signal to arch-independant kernel
#[no_mangle]
pub extern fn ksignal(signal: usize) {
    println!("SIGNAL {}, CPU {}, PID {:?}", signal, cpu_id(), context::context_id());
    {
        let contexts = context::contexts();
        if let Some(context_lock) = contexts.current() {
            let context = context_lock.read();
            println!("NAME {}", unsafe { ::core::str::from_utf8_unchecked(&context.name.lock()) });
        }
    }
    syscall::exit(signal & 0x7F);
}
<commit_msg>Do not allow unreachable patterns (often a logic issue)<commit_after>//! # The Redox OS Kernel, version 2
//!
//! The Redox OS Kernel is a microkernel that supports `x86_64` systems and
//! provides Unix-like syscalls for primarily Rust applications

//#![deny(warnings)]
#![cfg_attr(feature = "clippy", allow(if_same_then_else))]
#![cfg_attr(feature = "clippy", allow(inline_always))]
#![cfg_attr(feature = "clippy", allow(many_single_char_names))]
#![cfg_attr(feature = "clippy", allow(module_inception))]
#![cfg_attr(feature = "clippy", allow(new_without_default))]
#![cfg_attr(feature = "clippy", allow(not_unsafe_ptr_arg_deref))]
#![cfg_attr(feature = "clippy", allow(or_fun_call))]
#![cfg_attr(feature = "clippy", allow(too_many_arguments))]
#![deny(unreachable_patterns)]
#![feature(alloc)]
#![feature(allocator_api)]
#![feature(asm)]
#![feature(concat_idents)]
#![feature(const_fn)]
#![feature(core_intrinsics)]
#![feature(integer_atomics)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(never_type)]
#![feature(ptr_internals)]
#![feature(thread_local)]
#![no_std]

pub extern crate x86;

#[macro_use]
extern crate alloc;

#[macro_use]
extern crate bitflags;
extern crate goblin;
extern crate linked_list_allocator;
extern crate rustc_demangle;
extern crate spin;
#[cfg(feature = "slab")]
extern crate slab_allocator;

use alloc::vec::Vec;
use core::sync::atomic::{AtomicUsize, Ordering};

use scheme::{FileHandle, SchemeNamespace};

pub use consts::*;

#[macro_use]
/// Shared data structures
pub mod common;

/// Architecture-dependent stuff
#[macro_use]
pub mod arch;
pub use arch::*;

/// Constants like memory locations
pub mod consts;

/// Heap allocators
pub mod allocator;

/// ACPI table parsing
#[cfg(feature = "acpi")]
mod acpi;

/// Context management
pub mod context;

/// Architecture-independent devices
pub mod devices;

/// ELF file parsing
#[cfg(not(feature="doc"))]
pub mod elf;

/// Event handling
pub mod event;

/// External functions
pub mod externs;

/// Logging
pub mod log;

/// Memory management
pub mod memory;

/// Panic
#[cfg(not(any(feature="doc", test)))]
pub mod panic;

/// Schemes, filesystem handlers
pub mod scheme;

/// Synchronization primitives
pub mod sync;

/// Syscall handlers
pub mod syscall;

/// Time
pub mod time;

/// Tests
#[cfg(test)]
pub mod tests;

#[global_allocator]
static ALLOCATOR: allocator::Allocator = allocator::Allocator;

/// A unique number that identifies the current CPU - used for scheduling
#[thread_local]
static CPU_ID: AtomicUsize = AtomicUsize::new(0);

/// Get the current CPU's scheduling ID
#[inline(always)]
pub fn cpu_id() -> usize {
    CPU_ID.load(Ordering::Relaxed)
}

/// The count of all CPUs that can have work scheduled
static CPU_COUNT : AtomicUsize = AtomicUsize::new(0);

/// Get the number of CPUs currently active
#[inline(always)]
pub fn cpu_count() -> usize {
    CPU_COUNT.load(Ordering::Relaxed)
}

static mut INIT_ENV: &[u8] = &[];

/// Initialize userspace by running the initfs:bin/init process
/// This function will also set the CWD to initfs:bin and open debug: as stdio
pub extern fn userspace_init() {
    let path = b"initfs:/bin/init";
    let env = unsafe { INIT_ENV };

    assert_eq!(syscall::chdir(b"initfs:"), Ok(0));

    assert_eq!(syscall::open(b"debug:", syscall::flag::O_RDONLY).map(FileHandle::into), Ok(0));
    assert_eq!(syscall::open(b"debug:", syscall::flag::O_WRONLY).map(FileHandle::into), Ok(1));
    assert_eq!(syscall::open(b"debug:", syscall::flag::O_WRONLY).map(FileHandle::into), Ok(2));

    let fd = syscall::open(path, syscall::flag::O_RDONLY).expect("failed to open init");

    let mut args = Vec::new();
    args.push(path.to_vec().into_boxed_slice());

    let mut vars = Vec::new();
    for var in env.split(|b| *b == b'\n') {
        if ! var.is_empty() {
            vars.push(var.to_vec().into_boxed_slice());
        }
    }

    syscall::fexec_kernel(fd, args.into_boxed_slice(), vars.into_boxed_slice(), None).expect("failed to execute init");

    panic!("init returned");
}

/// This is the kernel entry point for the primary CPU. The arch crate is responsible for calling this
pub fn kmain(cpus: usize, env: &'static [u8]) -> ! {
    CPU_ID.store(0, Ordering::SeqCst);
    CPU_COUNT.store(cpus, Ordering::SeqCst);
    unsafe { INIT_ENV = env };

    //Initialize the first context, stored in kernel/src/context/mod.rs
    context::init();

    let pid = syscall::getpid();
    println!("BSP: {:?} {}", pid, cpus);
    println!("Env: {:?}", ::core::str::from_utf8(env));

    match context::contexts_mut().spawn(userspace_init) {
        Ok(context_lock) => {
            let mut context = context_lock.write();
            context.rns = SchemeNamespace::from(1);
            context.ens = SchemeNamespace::from(1);
            context.status = context::Status::Runnable;
        },
        Err(err) => {
            panic!("failed to spawn userspace_init: {:?}", err);
        }
    }

    loop {
        unsafe {
            interrupt::disable();
            if context::switch() {
                interrupt::enable_and_nop();
            } else {
                // Enable interrupts, then halt CPU (to save power) until the next interrupt is actually fired.
                interrupt::enable_and_halt();
            }
        }
    }
}

/// This is the main kernel entry point for secondary CPUs
#[allow(unreachable_code, unused_variables)]
pub fn kmain_ap(id: usize) -> ! {
    CPU_ID.store(id, Ordering::SeqCst);

    if cfg!(feature = "multi_core") {
        context::init();

        let pid = syscall::getpid();
        println!("AP {}: {:?}", id, pid);

        loop {
            unsafe {
                interrupt::disable();
                if context::switch() {
                    interrupt::enable_and_nop();
                } else {
                    // Enable interrupts, then halt CPU (to save power) until the next interrupt is actually fired.
                    interrupt::enable_and_halt();
                }
            }
        }
    } else {
        println!("AP {}: Disabled", id);

        loop {
            unsafe {
                interrupt::disable();
                interrupt::halt();
            }
        }
    }
}

/// Allow exception handlers to send signal to arch-independant kernel
#[no_mangle]
pub extern fn ksignal(signal: usize) {
    println!("SIGNAL {}, CPU {}, PID {:?}", signal, cpu_id(), context::context_id());
    {
        let contexts = context::contexts();
        if let Some(context_lock) = contexts.current() {
            let context = context_lock.read();
            println!("NAME {}", unsafe { ::core::str::from_utf8_unchecked(&context.name.lock()) });
        }
    }
    syscall::exit(signal & 0x7F);
}
<|endoftext|>


<commit_before>extern crate toml;

// Import traits
use std::io::prelude::Read;

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}

pub fn test() {
    println!("Hello from lib.rs");
}

#[derive(Debug)]
enum Step {
    Echo {name: String, string: String},
    Run,
    Shell,
    SetEnv,
}

fn parse_toml(filename: &str) -> Vec<Step> {

    println!("Reading toml file {}", filename);

    // Read toml file content
    let mut input = String::new();
    std::fs::File::open(&filename).and_then(|mut f| {
        f.read_to_string(&mut input)
    }).unwrap();

    // Pase toml file
    let mut parser = toml::Parser::new(&input);

    let toml = parser.parse().unwrap_or_else(|| {
        for err in &parser.errors {
            let (loline, locol) = parser.to_linecol(err.lo);
            let (hiline, hicol) = parser.to_linecol(err.hi);
            println!("{}:{}:{}-{}:{} error: {}",
                     filename, loline, locol, hiline, hicol, err.desc);
        }
        panic!("Error parsing toml file.");
    });

    println!("toml: {:#?}", toml);
    println!("toml: {:?}", toml.len());

    // Parse "on" table
    // unimplemented!();

    // Parse "variables" table
    // unimplemented!();

    // Parse "steps" table
    let mut steps = Vec::<Step>::with_capacity(10);

    for step in toml.get("step").unwrap().as_slice().unwrap() {
        let action = step.as_table().unwrap().get("action").unwrap().as_str().unwrap();
        let name = step.as_table().unwrap().get("name").unwrap().as_str().unwrap().to_string();
        let details = step.as_table().unwrap().get("details").unwrap().as_table().unwrap();
        let step = match action {
            "echo" => {
                let string = details.get("string").unwrap().as_str().unwrap().to_string();
                Step::Echo {name: name, string: string}
            },
            // "shell" => {
            // },
            // "run" => {
            // },
            // "setenv" => {
            // },
            _ => {
                println!("Unknown action '{}' in toml file", action);
                unimplemented!();
            },
        };
        println!("Step: {:?}", step);

        steps.push(step);
    }

    steps
}

pub fn execute_steps(filename: &str) {

    let steps = parse_toml(filename);

    println!("Steps: {:#?}", steps);
}
<commit_msg>New enums OnSuccess and OnFailure<commit_after>extern crate toml;

// Import traits
use std::io::prelude::Read;

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}

pub fn test() {
    println!("Hello from lib.rs");
}

#[derive(Debug)]
enum OnSuccess {
    Continue,
    Warn {message: String},
    Abort {message: String},
}

#[derive(Debug)]
enum OnFailure {
    Continue,
    Warn {message: String},
    Abort {message: String},
}

#[derive(Debug)]
enum Step {
    Echo {name: String, string: String},
    Run,
    Shell,
    SetEnv,
}

fn parse_toml(filename: &str) -> Vec<Step> {

    println!("Reading toml file {}", filename);

    // Read toml file content
    let mut input = String::new();
    std::fs::File::open(&filename).and_then(|mut f| {
        f.read_to_string(&mut input)
    }).unwrap();

    // Pase toml file
    let mut parser = toml::Parser::new(&input);

    let toml = parser.parse().unwrap_or_else(|| {
        for err in &parser.errors {
            let (loline, locol) = parser.to_linecol(err.lo);
            let (hiline, hicol) = parser.to_linecol(err.hi);
            println!("{}:{}:{}-{}:{} error: {}",
                     filename, loline, locol, hiline, hicol, err.desc);
        }
        panic!("Error parsing toml file.");
    });

    println!("toml: {:#?}", toml);
    println!("toml: {:?}", toml.len());

    // Parse "on" table
    // unimplemented!();

    // Parse "variables" table
    // unimplemented!();

    // Parse "steps" table
    let mut steps = Vec::<Step>::with_capacity(10);

    for step in toml.get("step").unwrap().as_slice().unwrap() {
        let action = step.as_table().unwrap().get("action").unwrap().as_str().unwrap();
        let name = step.as_table().unwrap().get("name").unwrap().as_str().unwrap().to_string();
        let details = step.as_table().unwrap().get("details").unwrap().as_table().unwrap();
        let step = match action {
            "echo" => {
                let string = details.get("string").unwrap().as_str().unwrap().to_string();
                Step::Echo {name: name, string: string}
            },
            // "shell" => {
            // },
            // "run" => {
            // },
            // "setenv" => {
            // },
            _ => {
                println!("Unknown action '{}' in toml file", action);
                unimplemented!();
            },
        };
        println!("Step: {:?}", step);

        steps.push(step);
    }

    steps
}

pub fn execute_steps(filename: &str) {

    let steps = parse_toml(filename);

    println!("Steps: {:#?}", steps);
}
<|endoftext|>


<commit_before>//! Capture a bitmap image of a display. The resulting screenshot is stored in
//! the `Screenshot` type, which varies per platform.
//!
//! # Platform-specific details
//!
//! Despite OS X's CoreGraphics documentation, the bitmap returned has its
//! origin at the top left corner. It uses ARGB pixels.
//!
//! The Windows GDI bitmap has its coordinate origin at the bottom left. We
//! attempt to undo this by reordering the rows. Windows also uses ARGB pixels.
//!
//! TODO Linux support. Contributions welcome.

#![allow(unstable, unused_assignments)]

extern crate libc;

use std::intrinsics::{size_of, offset};
pub use ffi::{get_screenshot};

#[derive(Copy)]
pub struct Pixel {
	pub a: u8,
	pub r: u8,
	pub g: u8,
	pub b: u8,
}

/// An image buffer containing the screenshot.
pub struct Screenshot {
	data: Vec<u8>,
	height: usize,
	width: usize,
	row_len: usize, // Might be superfluous
	pixel_width: usize,
}

impl Screenshot {
	/// Height of image in pixels.
	pub fn height(&self) -> usize { self.height }

	/// Width of image in pixels.
	pub fn width(&self) -> usize { self.width }

	/// Number of bytes in one row of bitmap.
	pub fn row_len(&self) -> usize { self.row_len }

	/// Width of pixel in bytes.
	pub fn pixel_width(&self) -> usize { self.pixel_width }

	/// Raw bitmap.
	pub unsafe fn raw_data(&self) -> *const u8 {
		&self.data[0] as *const u8
	}

	/// Number of bytes in bitmap
	pub fn raw_len(&self) -> usize {
		self.data.len() * unsafe {size_of::<u8>()}
	}

	/// Gets pixel at (row, col)
	pub fn get_pixel(&self, row: usize, col: usize) -> Pixel {
		let idx = (row*self.row_len() + col*self.pixel_width()) as isize;
		unsafe {
			let data = &self.data[0] as *const u8;
			if idx as usize > self.raw_len() { panic!("Bounds overflow"); }

			Pixel {
				a: *offset(data, idx+3),
				r: *offset(data, idx+2),
				g: *offset(data, idx+1),
				b: *offset(data, idx),
			}
		}
	}
}

pub type ScreenResult = Result<Screenshot, &'static str>;

#[cfg(target_os = "macos")]
mod ffi {
	#![allow(non_upper_case_globals, dead_code)]

	use libc;
	use ::Screenshot;
	use ::ScreenResult;

	type CFIndex = libc::c_long;
	type CFDataRef = *const u8; // *const CFData

	#[cfg(target_arch = "x86")]
	type CGFloat = libc::c_float;
	#[cfg(target_arch = "x86_64")]
	type CGFloat = libc::c_double;
	type CGError = libc::int32_t;

	type CGDirectDisplayID = libc::uint32_t;
	type CGDisplayCount = libc::uint32_t;
	type CGImageRef = *mut u8; // *mut CGImage
	type CGDataProviderRef = *mut u8; // *mut CGDataProvider

	const kCGErrorSuccess: CGError = 0;
	const kCGErrorFailure: CGError = 1000;
	const CGDisplayNoErr: CGError = kCGErrorSuccess;

	#[link(name = "CoreGraphics", kind = "framework")]
	extern "C" {
		fn CGGetActiveDisplayList(max_displays: libc::uint32_t,
	                              active_displays: *mut CGDirectDisplayID,
	                              display_count: *mut CGDisplayCount) -> CGError;
		fn CGDisplayCreateImage(displayID: CGDirectDisplayID) -> CGImageRef;
		fn CGImageRelease(image: CGImageRef);

		fn CGImageGetBitsPerComponent(image: CGImageRef) -> libc::size_t;
		fn CGImageGetBitsPerPixel(image: CGImageRef) -> libc::size_t;
		fn CGImageGetBytesPerRow(image: CGImageRef) -> libc::size_t;
		fn CGImageGetDataProvider(image: CGImageRef) -> CGDataProviderRef;
		fn CGImageGetHeight(image: CGImageRef) -> libc::size_t;
		fn CGImageGetWidth (image: CGImageRef) -> libc::size_t;

		fn CGDataProviderCopyData(provider: CGDataProviderRef) -> CFDataRef;
	}
	#[link(name = "CoreFoundation", kind = "framework")]
	extern "C" {
		fn CFDataGetLength (theData: CFDataRef) -> CFIndex;
		fn CFDataGetBytePtr(theData: CFDataRef) -> *const u8;
		fn CFRelease(cf: *const libc::c_void);
	}

	/// Get a screenshot of the requested display.
	pub fn get_screenshot(screen: usize) -> ScreenResult {
		let mut count: CGDisplayCount = 0;
		let mut err = CGDisplayNoErr;
		unsafe {
			err = CGGetActiveDisplayList(0, 0 as *mut CGDirectDisplayID, &mut count);
		};
		if err != CGDisplayNoErr {
			return Err("Error getting list of displays.");
		}

		let mut disps: Vec<CGDisplayCount> = Vec::with_capacity(count as usize);
		unsafe {
			disps.set_len(count as usize);
			err = CGGetActiveDisplayList(disps.len() as libc::uint32_t,
				&mut disps[0] as *mut CGDirectDisplayID,
				&mut count);
		};

		if err != CGDisplayNoErr {
			return Err("Error loading the display.");
		}

		let disp_id = disps[screen];

		unsafe {
			// Get screenshot of display #disp_id
			let cg_img = CGDisplayCreateImage(disp_id);

			// Get info about image
			let width = CGImageGetWidth(cg_img) as usize;
			let height = CGImageGetHeight(cg_img) as usize;
			let row_len = CGImageGetBytesPerRow(cg_img) as usize;
			let pixel_bits = CGImageGetBitsPerPixel(cg_img) as usize;
			if pixel_bits % 8 != 0 {
				return Err("Pixels aren't integral bytes.");
			}

			// Copy image into a Vec buffer
			let cf_data = CGDataProviderCopyData(CGImageGetDataProvider(cg_img));
			let raw_len = CFDataGetLength(cf_data) as usize;
			if width*height*pixel_bits != raw_len*8 {
				return Err("Image size is inconsistent with W*H*D.");
			}
			let data = Vec::<u8>::from_raw_buf(CFDataGetBytePtr(cf_data), raw_len);

			// Release native objects
			CGImageRelease(cg_img);
			CFRelease(cf_data as *const libc::c_void);

			Ok(Screenshot {
				data: data,
				height: height,
				width: width,
				row_len: row_len,
				pixel_width: pixel_bits/8
			})
		}
	}
}

#[cfg(target_os = "windows")]
mod ffi {
	#![allow(non_snake_case)]

	use libc::{c_int, c_uint, c_long, c_void};
	use std::intrinsics::{size_of};

	use ::Screenshot;
	use ::ScreenResult;

	type PVOID = *mut c_void;
	type LPVOID = *mut c_void;
	type WORD = u16; // c_uint;
	type DWORD = u32; // c_ulong;
	type BOOL = c_int;
	type BYTE = u8;
	type UINT = c_uint;
	type LONG = c_long;
	type HANDLE = PVOID;
	type HWND = HANDLE;
	type HDC = HANDLE;
	type HBITMAP = HANDLE;
	type HGDIOBJ = HANDLE;

	type LPBITMAPINFO = PVOID; // Hack

	const NULL: *mut c_void = 0us as *mut c_void;
	const HGDI_ERROR: *mut c_void = -1is as *mut c_void;
	const SM_CXSCREEN: c_int = 0;
	const SM_CYSCREEN: c_int = 1;

	/// TODO verify value
	const SRCCOPY: u32 = 0x00CC0020;
	const CAPTUREBLT: u32 = 0x40000000;
	const DIB_RGB_COLORS: UINT = 0;
	const BI_RGB: DWORD = 0;

	#[repr(C)]
	struct BITMAPINFOHEADER {
		biSize: DWORD,
		biWidth: LONG,
		biHeight: LONG,
		biPlanes: WORD,
		biBitCount: WORD,
		biCompression: DWORD,
		biSizeImage: DWORD,
		biXPelsPerMeter: LONG,
		biYPelsPerMeter: LONG,
		biClrUsed: DWORD,
		biClrImportant: DWORD,
	}

	#[repr(C)]
	struct RGBQUAD {
		rgbBlue: BYTE,
		rgbGreen: BYTE,
		rgbRed: BYTE,
		rgbReserved: BYTE,
	}

	/// WARNING variable sized struct
	#[repr(C)]
	struct BITMAPINFO {
		bmiHeader: BITMAPINFOHEADER,
		bmiColors: [RGBQUAD; 1],
	}

	#[link(name = "user32")]
	extern "system" {
		fn GetSystemMetrics(m: c_int) -> c_int;
	}

	#[link(name = "gdi32")]
	extern "system" {
		fn GetDesktopWindow() -> HWND;
		fn GetDC(hWnd: HWND) -> HDC;
		fn CreateCompatibleDC(hdc: HDC) -> HDC;
		fn CreateCompatibleBitmap(hdc: HDC, nWidth: c_int, nHeight: c_int) -> HBITMAP;
		fn SelectObject(hdc: HDC, hgdiobj: HGDIOBJ) -> HGDIOBJ;
		fn BitBlt(hdcDest: HDC, nXDest: c_int, nYDest: c_int, nWidth: c_int, nHeight: c_int,
                  hdcSrc: HDC, nXSrc: c_int, nYSrc: c_int, dwRop: DWORD) -> BOOL;
		fn GetDIBits(hdc: HDC, hbmp: HBITMAP, uStartScan: UINT, cScanLines: UINT,
					 lpvBits: LPVOID, lpbi: LPBITMAPINFO, uUsage: UINT) -> c_int;

		fn DeleteObject(hObject: HGDIOBJ) -> BOOL;
		fn ReleaseDC(hWnd: HWND, hDC: HDC) -> c_int;
		fn DeleteDC(hdc: HDC) -> BOOL;
	}

	/// Reorder rows in bitmap, last to first.
	/// TODO rewrite functionally
	fn flip_rows(data: Vec, height: usize, row_len: usize) -> Vec {
		let mut new_data = Vec::with_capacity(data.len());
		unsafe {new_data.set_len(data.len())};
		for row_i in range(0, height) {
			for byte_i in range(0, row_len) {
				let old_idx = (height-row_i-1)*row_len + byte_i;
				let new_idx = row_i*row_len + byte_i;
				new_data[new_idx] = data[old_idx];
			}
		}
		new_data
	}

	/// TODO don't ignore screen number
	pub fn get_screenshot(_screen: usize) -> ScreenResult {
		unsafe {
			let width = GetSystemMetrics(SM_CXSCREEN);
			let height = GetSystemMetrics(SM_CYSCREEN);

			let h_wnd_screen = GetDesktopWindow();
			let h_dc_screen = GetDC(h_wnd_screen);

			// Create a Windows Bitmap, and copy the bits into it
			let h_dc = CreateCompatibleDC(h_dc_screen);
			if h_dc == NULL { return Err("Can't get a Windows display.");}

			let h_bmp = CreateCompatibleBitmap(h_dc_screen, width, height);
			if h_bmp == NULL { return Err("Can't create a Windows buffer");}

			let res = SelectObject(h_dc, h_bmp);
			if res == NULL || res == HGDI_ERROR {
				return Err("Can't select Windows buffer.");
			}

			let res = BitBlt(h_dc, 0, 0, width, height, h_dc_screen, 0, 0, SRCCOPY|CAPTUREBLT);
			if res == 0 { return Err("Failed to copy screen to Windows buffer");}

			// Get image info
			let pixel_width: usize = 4; // FIXME
			let mut bmi = BITMAPINFO {
				bmiHeader: BITMAPINFOHEADER {
					biSize: size_of::<BITMAPINFOHEADER>() as DWORD,
					biWidth: width as LONG,
					biHeight: height as LONG,
					biPlanes: 1,
					biBitCount: 8*pixel_width as WORD,
					biCompression: BI_RGB,
					biSizeImage: (width * height * pixel_width as c_int) as DWORD,
					biXPelsPerMeter: 0,
					biYPelsPerMeter: 0,
					biClrUsed: 0,
					biClrImportant: 0,
				},
				bmiColors: [RGBQUAD {
					rgbBlue: 0,
					rgbGreen: 0,
					rgbRed: 0,
					rgbReserved: 0
				}],
			};

			// Create a Vec for image
			let size: usize = (width*height) as usize * pixel_width;
			let mut data: Vec<u8> = Vec::with_capacity(size);
			data.set_len(size);

			// copy bits into Vec
			GetDIBits(h_dc, h_bmp, 0, height as DWORD,
				&mut data[0] as *mut u8 as *mut c_void,
				&mut bmi as *mut BITMAPINFO as *mut c_void,
				DIB_RGB_COLORS);

			// Release native image buffers
			ReleaseDC(h_wnd_screen, h_dc_screen); // don't need screen anymore
			DeleteDC(h_dc);
			DeleteObject(h_bmp);

			let data = flip_rows(data, height as usize, width as usize*pixel_width);

			Ok(Screenshot {
				data: data,
				height: height as usize,
				width: width as usize,
				row_len: width as usize*pixel_width,
				pixel_width: pixel_width,
			})
		}
	}
}

#[test]
fn test_get_screenshot() {
	let s: Screenshot = get_screenshot(0).unwrap();
	println!("width: {}\n height: {}\npixel width: {}\n bytes: {}",
		s.width(), s.height(), s.pixel_width(), s.raw_len());
}
<commit_msg>Condense unsafe code, fix windows build.<commit_after>//! Capture a bitmap image of a display. The resulting screenshot is stored in
//! the `Screenshot` type, which varies per platform.
//!
//! # Platform-specific details
//!
//! Despite OS X's CoreGraphics documentation, the bitmap returned has its
//! origin at the top left corner. It uses ARGB pixels.
//!
//! The Windows GDI bitmap has its coordinate origin at the bottom left. We
//! attempt to undo this by reordering the rows. Windows also uses ARGB pixels.
//!
//! TODO Linux support. Contributions welcome.

#![allow(unstable, unused_assignments)]

extern crate libc;

use std::intrinsics::{size_of, offset};
pub use ffi::{get_screenshot};

#[derive(Copy)]
pub struct Pixel {
	pub a: u8,
	pub r: u8,
	pub g: u8,
	pub b: u8,
}

/// An image buffer containing the screenshot.
pub struct Screenshot {
	data: Vec<u8>,
	height: usize,
	width: usize,
	row_len: usize, // Might be superfluous
	pixel_width: usize,
}

impl Screenshot {
	/// Height of image in pixels.
	pub fn height(&self) -> usize { self.height }

	/// Width of image in pixels.
	pub fn width(&self) -> usize { self.width }

	/// Number of bytes in one row of bitmap.
	pub fn row_len(&self) -> usize { self.row_len }

	/// Width of pixel in bytes.
	pub fn pixel_width(&self) -> usize { self.pixel_width }

	/// Raw bitmap.
	pub unsafe fn raw_data(&self) -> *const u8 {
		&self.data[0] as *const u8
	}

	/// Number of bytes in bitmap
	pub fn raw_len(&self) -> usize {
		self.data.len() * unsafe {size_of::<u8>()}
	}

	/// Gets pixel at (row, col)
	pub fn get_pixel(&self, row: usize, col: usize) -> Pixel {
		let idx = (row*self.row_len() + col*self.pixel_width()) as isize;
		unsafe {
			let data = &self.data[0] as *const u8;
			if idx as usize > self.raw_len() { panic!("Bounds overflow"); }

			Pixel {
				a: *offset(data, idx+3),
				r: *offset(data, idx+2),
				g: *offset(data, idx+1),
				b: *offset(data, idx),
			}
		}
	}
}

pub type ScreenResult = Result<Screenshot, &'static str>;

#[cfg(target_os = "macos")]
mod ffi {
	#![allow(non_upper_case_globals, dead_code)]

	use libc;
	use ::Screenshot;
	use ::ScreenResult;

	type CFIndex = libc::c_long;
	type CFDataRef = *const u8; // *const CFData

	#[cfg(target_arch = "x86")]
	type CGFloat = libc::c_float;
	#[cfg(target_arch = "x86_64")]
	type CGFloat = libc::c_double;
	type CGError = libc::int32_t;

	type CGDirectDisplayID = libc::uint32_t;
	type CGDisplayCount = libc::uint32_t;
	type CGImageRef = *mut u8; // *mut CGImage
	type CGDataProviderRef = *mut u8; // *mut CGDataProvider

	const kCGErrorSuccess: CGError = 0;
	const kCGErrorFailure: CGError = 1000;
	const CGDisplayNoErr: CGError = kCGErrorSuccess;

	#[link(name = "CoreGraphics", kind = "framework")]
	extern "C" {
		fn CGGetActiveDisplayList(max_displays: libc::uint32_t,
	                              active_displays: *mut CGDirectDisplayID,
	                              display_count: *mut CGDisplayCount) -> CGError;
		fn CGDisplayCreateImage(displayID: CGDirectDisplayID) -> CGImageRef;
		fn CGImageRelease(image: CGImageRef);

		fn CGImageGetBitsPerComponent(image: CGImageRef) -> libc::size_t;
		fn CGImageGetBitsPerPixel(image: CGImageRef) -> libc::size_t;
		fn CGImageGetBytesPerRow(image: CGImageRef) -> libc::size_t;
		fn CGImageGetDataProvider(image: CGImageRef) -> CGDataProviderRef;
		fn CGImageGetHeight(image: CGImageRef) -> libc::size_t;
		fn CGImageGetWidth (image: CGImageRef) -> libc::size_t;

		fn CGDataProviderCopyData(provider: CGDataProviderRef) -> CFDataRef;
	}
	#[link(name = "CoreFoundation", kind = "framework")]
	extern "C" {
		fn CFDataGetLength (theData: CFDataRef) -> CFIndex;
		fn CFDataGetBytePtr(theData: CFDataRef) -> *const u8;
		fn CFRelease(cf: *const libc::c_void);
	}

	/// Get a screenshot of the requested display.
	pub fn get_screenshot(screen: usize) -> ScreenResult {
		unsafe {
			// Get number of displays
			let mut count: CGDisplayCount = 0;
			let mut err = CGDisplayNoErr;
			err = CGGetActiveDisplayList(0, 0 as *mut CGDirectDisplayID, &mut count);
			if err != CGDisplayNoErr {
				return Err("Error getting number of displays.");
			}

			// Get list of displays
			let mut disps: Vec<CGDisplayCount> = Vec::with_capacity(count as usize);
			disps.set_len(count as usize);
			err = CGGetActiveDisplayList(disps.len() as libc::uint32_t,
				&mut disps[0] as *mut CGDirectDisplayID,
				&mut count);
			if err != CGDisplayNoErr {
				return Err("Error getting list of displays.");
			}

			// Get screenshot of requested display
			let disp_id = disps[screen];
			let cg_img = CGDisplayCreateImage(disp_id);

			// Get info about image
			let width = CGImageGetWidth(cg_img) as usize;
			let height = CGImageGetHeight(cg_img) as usize;
			let row_len = CGImageGetBytesPerRow(cg_img) as usize;
			let pixel_bits = CGImageGetBitsPerPixel(cg_img) as usize;
			if pixel_bits % 8 != 0 {
				return Err("Pixels aren't integral bytes.");
			}

			// Copy image into a Vec buffer
			let cf_data = CGDataProviderCopyData(CGImageGetDataProvider(cg_img));
			let raw_len = CFDataGetLength(cf_data) as usize;
			if width*height*pixel_bits != raw_len*8 {
				return Err("Image size is inconsistent with W*H*D.");
			}
			let data = Vec::<u8>::from_raw_buf(CFDataGetBytePtr(cf_data), raw_len);

			// Release native objects
			CGImageRelease(cg_img);
			CFRelease(cf_data as *const libc::c_void);

			Ok(Screenshot {
				data: data,
				height: height,
				width: width,
				row_len: row_len,
				pixel_width: pixel_bits/8
			})
		}
	}
}

#[cfg(target_os = "windows")]
mod ffi {
	#![allow(non_snake_case, dead_code)]

	use libc::{c_int, c_uint, c_long, c_void};
	use std::intrinsics::{size_of};

	use ::Screenshot;
	use ::ScreenResult;

	type PVOID = *mut c_void;
	type LPVOID = *mut c_void;
	type WORD = u16; // c_uint;
	type DWORD = u32; // c_ulong;
	type BOOL = c_int;
	type BYTE = u8;
	type UINT = c_uint;
	type LONG = c_long;
	type LPARAM = c_long;

	#[repr(C)]
	struct RECT {
		left: LONG,
		top: LONG,
		right: LONG, // immediately outside rect
		bottom: LONG, // immediately outside rect
	}
	type LPCRECT = *const RECT;
	type LPRECT = *mut RECT;

	type HANDLE = PVOID;
	type HMONITOR = HANDLE;
	type HWND = HANDLE;
	type HDC = HANDLE;
	#[repr(C)]
	struct MONITORINFO {
		cbSize: DWORD,
		rcMonitor: RECT,
		rcWork: RECT,
		dwFlags: DWORD,
	}
	type LPMONITORINFO = *mut MONITORINFO;
	type MONITORENUMPROC = fn(HMONITOR, HDC, LPRECT, LPARAM) -> BOOL;

	type HBITMAP = HANDLE;
	type HGDIOBJ = HANDLE;
	type LPBITMAPINFO = PVOID; // Hack

	const NULL: *mut c_void = 0us as *mut c_void;
	const HGDI_ERROR: *mut c_void = -1is as *mut c_void;
	const SM_CXSCREEN: c_int = 0;
	const SM_CYSCREEN: c_int = 1;

	/// TODO verify value
	const SRCCOPY: u32 = 0x00CC0020;
	const CAPTUREBLT: u32 = 0x40000000;
	const DIB_RGB_COLORS: UINT = 0;
	const BI_RGB: DWORD = 0;

	#[repr(C)]
	struct BITMAPINFOHEADER {
		biSize: DWORD,
		biWidth: LONG,
		biHeight: LONG,
		biPlanes: WORD,
		biBitCount: WORD,
		biCompression: DWORD,
		biSizeImage: DWORD,
		biXPelsPerMeter: LONG,
		biYPelsPerMeter: LONG,
		biClrUsed: DWORD,
		biClrImportant: DWORD,
	}

	#[repr(C)]
	struct RGBQUAD {
		rgbBlue: BYTE,
		rgbGreen: BYTE,
		rgbRed: BYTE,
		rgbReserved: BYTE,
	}

	/// WARNING variable sized struct
	#[repr(C)]
	struct BITMAPINFO {
		bmiHeader: BITMAPINFOHEADER,
		bmiColors: [RGBQUAD; 1],
	}

	#[link(name = "user32")]
	extern "system" {
		fn GetSystemMetrics(m: c_int) -> c_int;
		fn EnumDisplayMonitors(hdc: HDC, lprcClip: LPCRECT,
							   lpfnEnum: MONITORENUMPROC, dwData: LPARAM) -> BOOL;
		fn GetMonitorInfo(hMonitor: HMONITOR, lpmi: LPMONITORINFO) -> BOOL;
	}

	#[link(name = "gdi32")]
	extern "system" {
		fn GetDesktopWindow() -> HWND;
		fn GetDC(hWnd: HWND) -> HDC;
		fn CreateCompatibleDC(hdc: HDC) -> HDC;
		fn CreateCompatibleBitmap(hdc: HDC, nWidth: c_int, nHeight: c_int) -> HBITMAP;
		fn SelectObject(hdc: HDC, hgdiobj: HGDIOBJ) -> HGDIOBJ;
		fn BitBlt(hdcDest: HDC, nXDest: c_int, nYDest: c_int, nWidth: c_int, nHeight: c_int,
                  hdcSrc: HDC, nXSrc: c_int, nYSrc: c_int, dwRop: DWORD) -> BOOL;
		fn GetDIBits(hdc: HDC, hbmp: HBITMAP, uStartScan: UINT, cScanLines: UINT,
					 lpvBits: LPVOID, lpbi: LPBITMAPINFO, uUsage: UINT) -> c_int;

		fn DeleteObject(hObject: HGDIOBJ) -> BOOL;
		fn ReleaseDC(hWnd: HWND, hDC: HDC) -> c_int;
		fn DeleteDC(hdc: HDC) -> BOOL;
	}

	/// Reorder rows in bitmap, last to first.
	/// TODO rewrite functionally
	fn flip_rows(data: Vec<u8>, height: usize, row_len: usize) -> Vec<u8> {
		let mut new_data = Vec::with_capacity(data.len());
		unsafe {new_data.set_len(data.len())};
		for row_i in range(0, height) {
			for byte_i in range(0, row_len) {
				let old_idx = (height-row_i-1)*row_len + byte_i;
				let new_idx = row_i*row_len + byte_i;
				new_data[new_idx] = data[old_idx];
			}
		}
		new_data
	}

	/// TODO Support multiple screens
	/// This may never happen, given the horrific quality of Win32 APIs
	pub fn get_screenshot(_screen: usize) -> ScreenResult {
		unsafe {
			// Enumerate monitors, getting a handle and DC for requested monitor.
			// loljk, because doing that on Windows is worse than death
			let h_wnd_screen = GetDesktopWindow();
			let h_dc_screen = GetDC(h_wnd_screen);
			let width = GetSystemMetrics(SM_CXSCREEN);
			let height = GetSystemMetrics(SM_CYSCREEN);

			// Create a Windows Bitmap, and copy the bits into it
			let h_dc = CreateCompatibleDC(h_dc_screen);
			if h_dc == NULL { return Err("Can't get a Windows display.");}

			let h_bmp = CreateCompatibleBitmap(h_dc_screen, width, height);
			if h_bmp == NULL { return Err("Can't create a Windows buffer");}

			let res = SelectObject(h_dc, h_bmp);
			if res == NULL || res == HGDI_ERROR {
				return Err("Can't select Windows buffer.");
			}

			let res = BitBlt(h_dc, 0, 0, width, height, h_dc_screen, 0, 0, SRCCOPY|CAPTUREBLT);
			if res == 0 { return Err("Failed to copy screen to Windows buffer");}

			// Get image info
			let pixel_width: usize = 4; // FIXME
			let mut bmi = BITMAPINFO {
				bmiHeader: BITMAPINFOHEADER {
					biSize: size_of::<BITMAPINFOHEADER>() as DWORD,
					biWidth: width as LONG,
					biHeight: height as LONG,
					biPlanes: 1,
					biBitCount: 8*pixel_width as WORD,
					biCompression: BI_RGB,
					biSizeImage: (width * height * pixel_width as c_int) as DWORD,
					biXPelsPerMeter: 0,
					biYPelsPerMeter: 0,
					biClrUsed: 0,
					biClrImportant: 0,
				},
				bmiColors: [RGBQUAD {
					rgbBlue: 0,
					rgbGreen: 0,
					rgbRed: 0,
					rgbReserved: 0
				}],
			};

			// Create a Vec for image
			let size: usize = (width*height) as usize * pixel_width;
			let mut data: Vec<u8> = Vec::with_capacity(size);
			data.set_len(size);

			// copy bits into Vec
			GetDIBits(h_dc, h_bmp, 0, height as DWORD,
				&mut data[0] as *mut u8 as *mut c_void,
				&mut bmi as *mut BITMAPINFO as *mut c_void,
				DIB_RGB_COLORS);

			// Release native image buffers
			ReleaseDC(h_wnd_screen, h_dc_screen); // don't need screen anymore
			DeleteDC(h_dc);
			DeleteObject(h_bmp);

			let data = flip_rows(data, height as usize, width as usize*pixel_width);

			Ok(Screenshot {
				data: data,
				height: height as usize,
				width: width as usize,
				row_len: width as usize*pixel_width,
				pixel_width: pixel_width,
			})
		}
	}
}

#[test]
fn test_get_screenshot() {
	let s: Screenshot = get_screenshot(0).unwrap();
	println!("width: {}\n height: {}\npixel width: {}\n bytes: {}",
		s.width(), s.height(), s.pixel_width(), s.raw_len());
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate toml;
extern crate ldap3;
extern crate serde;
#[macro_use]
extern crate serde_derive;

mod pam;
mod config;

use std::time::Duration;
use std::collections::BTreeSet;
use ldap3::{LdapConn, LdapConnBuilder, Scope, SearchEntry};

// Re-export the PAM callbacks
pub use pam::callbacks::*;
use pam::{PamResultCode, get_user, set_user};
use config::Config;

#[allow(dead_code)]
fn spare() {
    println!("");
    panic!("");
}

pub fn acct_mgmt(pamh: pam::PamHandleT, args: Vec<String>, silent: bool) -> PamResultCode {

    let user = match get_user(pamh) {
        Ok(u) => u,
        Err(_) => return PamResultCode::PAM_AUTH_ERR,
    };

    if args.len() != 1 {
        return PamResultCode::PAM_SERVICE_ERR;
    }

    let config = match Config::load(&args[0]) {
        Ok(c) => c,
        Err(e) => {
            if !silent {
                println!("ERROR: {:?}", e);
            }
            return PamResultCode::PAM_SERVICE_ERR;
        }
    };

    let lconn = match ldap_connect(&config.ldap) {
        Ok(c) => c,
        Err(_) => return PamResultCode::PAM_SERVICE_ERR,
    };

    let groups = match get_user_groups(&lconn, &config.ldap, &user) {
        Ok(g) => g,
        Err(e) => return e,
    };

    let mut ret = PamResultCode::PAM_AUTH_ERR;
    for (group, mapped_user) in config.mappings.iter() {
        if groups.contains(group) {
            if !silent {
                println!("Mapping {} -> {}", user, mapped_user);
            }
            match set_user(pamh, mapped_user.clone()) {
                Ok(_) => {
                    ret = PamResultCode::PAM_SUCCESS;
                    break;
                }
                Err(_) => {
                    ret = PamResultCode::PAM_SERVICE_ERR;
                    break;
                }
            }
        };
    }
    lconn.unbind().ok();
    ret
}

fn extract_ldap_servers<'a>(uri: &'a String) -> Vec<&'a str> {
    uri.split(',').collect::<Vec<_>>()
}

#[test]
fn test_extract_ldap_servers() {
    assert_eq!(
        extract_ldap_servers(&String::from("asdf,qwer")),
        vec!["asdf", "qwer"]
    );
}

fn ldap_connect(ldap: &config::LdapConfig) -> Result<LdapConn, ()> {
    let servers = extract_ldap_servers(&ldap.uri);

    for server in servers {
        println!("Trying to connect to {}", server);
        let lconn = match LdapConnBuilder::<LdapConn>::new()
            .with_conn_timeout(Duration::from_secs(ldap.conn_timeout))
            .connect(server) {
            Ok(c) => {
                println!("Connected to {}", server);
                c
            }
            Err(_) => {
                println!("Failed to connect to {}, trying next.", server);
                continue;
            }
        };
        match lconn
            .with_timeout(Duration::from_secs(ldap.op_timeout))
            .simple_bind(&ldap.user, &ldap.pass) {
            Ok(_) => {
                println!("Simple Bind Succeeded");
                return Ok(lconn);
            }
            Err(_) => {
                println!("Simple Bind Failed, trying next.");
                continue;
            }
        }
    }
    println!("Server list exhausted.");
    Err(())
}

fn get_user_groups(
    lconn: &LdapConn,
    config: &config::LdapConfig,
    user: &String,
) -> Result<BTreeSet<String>, PamResultCode> {
    lconn
        .with_timeout(Duration::from_secs(config.op_timeout))
        .search(
            &config.user_base_dn,
            Scope::Subtree,
            &format!("({}={})", config.uid_attribute, user),
            vec![config.group_attribute.clone()],
        )
        .map(|result_tuple| {
            result_tuple
                .0
                .into_iter()
                .map(|r| SearchEntry::construct(r))
                .flat_map(|e| e.attrs)
                .filter_map(|e| if e.0 == config.group_attribute {
                    Some(e.1)
                } else {
                    None
                })
                .flat_map(|e| e)
                .filter(|e| {
                    e.to_lowercase().ends_with(
                        &config.group_base_dn.to_lowercase(),
                    )
                })
                .filter_map(|e| if let Some(end) = e.find(",") {
                    Some(e[3..end].to_owned())
                } else {
                    None
                })
                .collect::<BTreeSet<_>>()
        })
        .map_err(|_| PamResultCode::PAM_SERVICE_ERR)
}
<commit_msg>Adapt LDAP code for ldap3-0.5<commit_after>extern crate libc;
extern crate toml;
extern crate ldap3;
extern crate serde;
#[macro_use]
extern crate serde_derive;

mod pam;
mod config;

use std::time::Duration;
use std::collections::BTreeSet;
use ldap3::{LdapConn, LdapConnBuilder, Scope, SearchEntry};
use ldap3::result::SearchResult;
use ldap3::ldap_escape;

// Re-export the PAM callbacks
pub use pam::callbacks::*;
use pam::{PamResultCode, get_user, set_user};
use config::Config;

#[allow(dead_code)]
fn spare() {
    println!("");
    panic!("");
}

pub fn acct_mgmt(pamh: pam::PamHandleT, args: Vec<String>, silent: bool) -> PamResultCode {

    let user = match get_user(pamh) {
        Ok(u) => u,
        Err(_) => return PamResultCode::PAM_AUTH_ERR,
    };

    if args.len() != 1 {
        return PamResultCode::PAM_SERVICE_ERR;
    }

    let config = match Config::load(&args[0]) {
        Ok(c) => c,
        Err(e) => {
            if !silent {
                println!("ERROR: {:?}", e);
            }
            return PamResultCode::PAM_SERVICE_ERR;
        }
    };

    let lconn = match ldap_connect(&config.ldap) {
        Ok(c) => c,
        Err(_) => return PamResultCode::PAM_SERVICE_ERR,
    };

    let groups = match get_user_groups(&lconn, &config.ldap, &user) {
        Ok(g) => g,
        Err(e) => return e,
    };

    let mut ret = PamResultCode::PAM_AUTH_ERR;
    for (group, mapped_user) in config.mappings.iter() {
        if groups.contains(group) {
            if !silent {
                println!("Mapping {} -> {}", user, mapped_user);
            }
            match set_user(pamh, mapped_user.clone()) {
                Ok(_) => {
                    ret = PamResultCode::PAM_SUCCESS;
                    break;
                }
                Err(_) => {
                    ret = PamResultCode::PAM_SERVICE_ERR;
                    break;
                }
            }
        };
    }
    lconn.unbind().ok();
    ret
}

fn extract_ldap_servers<'a>(uri: &'a String) -> Vec<&'a str> {
    uri.split(',').collect::<Vec<_>>()
}

#[test]
fn test_extract_ldap_servers() {
    assert_eq!(
        extract_ldap_servers(&String::from("asdf,qwer")),
        vec!["asdf", "qwer"]
    );
}

fn ldap_connect(ldap: &config::LdapConfig) -> Result<LdapConn, ()> {
    let servers = extract_ldap_servers(&ldap.uri);

    for server in servers {
        println!("Trying to connect to {}", server);
        let lconn = match LdapConnBuilder::<LdapConn>::new()
            .with_conn_timeout(Duration::from_secs(ldap.conn_timeout))
            .connect(server) {
            Ok(c) => {
                println!("Connected to {}", server);
                c
            }
            Err(_) => {
                println!("Failed to connect to {}, trying next.", server);
                continue;
            }
        };
        match lconn
            .with_timeout(Duration::from_secs(ldap.op_timeout))
            .simple_bind(&ldap.user, &ldap.pass) {
            Ok(_) => {
                println!("Simple Bind Succeeded");
                return Ok(lconn);
            }
            Err(_) => {
                println!("Simple Bind Failed, trying next.");
                continue;
            }
        }
    }
    println!("Server list exhausted.");
    Err(())
}

fn get_user_groups(
    lconn: &LdapConn,
    config: &config::LdapConfig,
    user: &String,
) -> Result<BTreeSet<String>, PamResultCode> {
    lconn
        .with_timeout(Duration::from_secs(config.op_timeout))
        .search(
            &config.user_base_dn,
            Scope::Subtree,
            &format!("({}={})", config.uid_attribute, ldap_escape(user.as_ref())),
            vec![config.group_attribute.clone()],
        )
        .and_then(SearchResult::success)
        .map(|(rset, _)| {
            rset
                .into_iter()
                .map(|r| SearchEntry::construct(r))
                .flat_map(|e| e.attrs)
                .filter_map(|e| if e.0 == config.group_attribute {
                    Some(e.1)
                } else {
                    None
                })
                .flat_map(|e| e)
                .filter(|e| {
                    e.to_lowercase().ends_with(
                        &config.group_base_dn.to_lowercase(),
                    )
                })
                .filter_map(|e| if let Some(end) = e.find(",") {
                    Some(e[3..end].to_owned())
                } else {
                    None
                })
                .collect::<BTreeSet<_>>()
        })
        .map_err(|_| PamResultCode::PAM_SERVICE_ERR)
}
<|endoftext|>


<commit_before>use std::io::{Error as IoError, Read};
use std::fs::File;
use std::num::ParseIntError;
use std::path::Path;
use std::time::Duration;

// TODO: ADD DOCS!!!

const UTF8_BOM: &'static str = "\u{feff}";

#[derive(Debug)]
enum State {
    Pos,
    Time,
    Text,
}

#[derive(Debug)]
pub struct Subtitle {
    pub pos: usize,
    pub start_time: Duration,
    pub end_time: Duration,
    pub text: String,
}

pub fn parse(source: &str) -> Result<Vec<Subtitle>, Error> {
    let source = source.trim_left_matches(UTF8_BOM).trim().lines();

    let mut result = Vec::new();

    let mut state = State::Pos;

    let mut pos: Option<usize> = None;
    let mut start_time: Option<Duration> = None;
    let mut end_time: Option<Duration> = None;
    let mut text: Option<String> = None;

    macro_rules! push_subtitle {
        ($pos:ident) => (
            result.push(Subtitle {
                pos: $pos,
                start_time: match start_time {
                    Some(val) => val,
                    None => return Err(Error::MissingStartTime),
                },
                end_time: match end_time {
                    Some(val) => val,
                    None => return Err(Error::MissingEndTime),
                },
                text: match text {
                    Some(val) => val,
                    None => return Err(Error::MissingText),
                }
            });
        )
    }

    for line in source {
        match state {
            State::Pos => {
                if let Some(ps) = pos {
                    push_subtitle!(ps);
                    start_time = None;
                    end_time = None;
                    text = None;
                }
                pos = Some(line.parse::<usize>()?);
                state = State::Time;
            }
            State::Time => {
                let mut parts = line.split("-->");
                start_time = match parts.next() {
                    Some(v) => Some(duration_from_str(v)?),
                    None => None,
                };
                end_time = match parts.next() {
                    Some(v) => Some(duration_from_str(v)?),
                    None => None,
                };
                state = State::Text;
            }
            State::Text => {
                if line.is_empty() {
                    state = State::Pos;
                } else {
                    if let Some(ref mut txt) = text {
                        txt.push('\n');
                        txt.push_str(line);
                    } else {
                        text = Some(line.to_string());
                    }
                }
            }
        }
    }
    if let Some(ps) = pos {
        push_subtitle!(ps);
    }

    Ok(result)
}

pub fn read_from_file<P: AsRef<Path>>(path: P) -> Result<Vec<Subtitle>, Error> {
    let mut file = File::open(path)?;
    let mut buf = String::new();
    file.read_to_string(&mut buf)?;
    parse(&buf)
}

#[derive(Debug)]
pub enum Error {
    Io(IoError),
    ParseInt(ParseIntError),
    MissingText,
    MissingEndTime,
    MissingStartTime,
    ParseTime,
}

// TODO: impl StdError

impl From<ParseIntError> for Error {
    fn from(err: ParseIntError) -> Error {
        Error::ParseInt(err)
    }
}

impl From<IoError> for Error {
    fn from(err: IoError) -> Error {
        Error::Io(err)
    }
}

macro_rules! parse_time_part {
    ($part:expr) => {{
        match $part {
            Some(val) => {
                val.trim().parse::<u64>()?
            },
            None => {
                return Err(Error::ParseTime);
            }
        }
    }}
}

fn duration_from_str(time: &str) -> Result<Duration, Error> {
    let mut time = time.split(",");
    let (hours, mut minutes, mut seconds) = match time.next() {
        Some(val) => {
            let mut parts = val.split(":");
            let result = (
                parse_time_part!(parts.next()),
                parse_time_part!(parts.next()),
                parse_time_part!(parts.next()),
            );
            if parts.next().is_some() {
                return Err(Error::ParseTime);
            }
            result
        },
        None => {
            return Err(Error::ParseTime);
        }
    };
    let mut milliseconds = parse_time_part!(time.next());
    if time.next().is_some() {
        return Err(Error::ParseTime);
    }
    minutes += hours * 60;
    seconds += minutes * 60;
    milliseconds += seconds * 1000;
    Ok(Duration::new(milliseconds, 0))
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use ::{Error, UTF8_BOM, parse};

    #[test]
    fn it_works() {
        let source_without_bom = "1
00:00:58,392 --> 00:01:02,563
The war had all but ground to a halt
in the blink of an eye.

2
00:01:04,565 --> 00:01:08,986
Lucian, the most feared and ruthless
leader ever to rule the Lycan clan...

3
00:01:09,070 --> 00:01:11,656
...had finally been killed.

652
01:53:02,325 --> 01:53:06,162
Soon, Marcus will take the throne.
";
        let source_with_bom = format!("{}{}", UTF8_BOM, source_without_bom);

        fn assert_it_works(data: &str) -> Result<(), Error> {
            let result = parse(&data)?;
            assert_eq!(result.len(), 4);
            assert_eq!(result[0].pos, 1);
            assert_eq!(result[0].start_time, Duration::new(58392, 0));
            assert_eq!(result[0].end_time, Duration::new(62563, 0));
            assert_eq!(result[0].text, "The war had all but ground to a halt\nin the blink of an eye.");
            assert_eq!(result[1].pos, 2);
            assert_eq!(result[1].start_time, Duration::new(64565, 0));
            assert_eq!(result[1].end_time, Duration::new(68986, 0));
            assert_eq!(result[1].text, "Lucian, the most feared and ruthless\nleader ever to rule the Lycan clan...");
            assert_eq!(result[2].pos, 3);
            assert_eq!(result[2].start_time, Duration::new(69070, 0));
            assert_eq!(result[2].end_time, Duration::new(71656, 0));
            assert_eq!(result[2].text, "...had finally been killed.");
            assert_eq!(result[3].pos, 652);
            assert_eq!(result[3].start_time, Duration::new(6782325, 0));
            assert_eq!(result[3].end_time, Duration::new(6786162, 0));
            assert_eq!(result[3].text, "Soon, Marcus will take the throne.");
            Ok(())
        }

        assert_it_works(&source_without_bom).expect("Failed to parse UTF-8 source without BOM");
        assert_it_works(&source_with_bom).expect("Failed to parse UTF-8 source with BOM");
    }

    // TODO: NEED MOAR TESTS!!!

}
<commit_msg>Impl std::error::Error<commit_after>use std::fmt;
use std::error::Error as StdError;
use std::io::{Error as IoError, Read};
use std::fs::File;
use std::num::ParseIntError;
use std::path::Path;
use std::result::Result as StdResult;
use std::time::Duration;

// TODO: ADD DOCS!!!

const UTF8_BOM: &'static str = "\u{feff}";

#[derive(Debug)]
enum State {
    Pos,
    Time,
    Text,
}

#[derive(Debug)]
pub struct Subtitle {
    pub pos: usize,
    pub start_time: Duration,
    pub end_time: Duration,
    pub text: String,
}

pub fn parse(source: &str) -> Result<Vec<Subtitle>> {
    let source = source.trim_left_matches(UTF8_BOM).trim().lines();

    let mut result = Vec::new();

    let mut state = State::Pos;

    let mut pos: Option<usize> = None;
    let mut start_time: Option<Duration> = None;
    let mut end_time: Option<Duration> = None;
    let mut text: Option<String> = None;

    macro_rules! push_subtitle {
        ($pos:ident) => (
            result.push(Subtitle {
                pos: $pos,
                start_time: match start_time {
                    Some(val) => val,
                    None => return Err(Error::MissingStartTime),
                },
                end_time: match end_time {
                    Some(val) => val,
                    None => return Err(Error::MissingEndTime),
                },
                text: match text {
                    Some(val) => val,
                    None => return Err(Error::MissingText),
                }
            });
        )
    }

    for line in source {
        match state {
            State::Pos => {
                if let Some(ps) = pos {
                    push_subtitle!(ps);
                    start_time = None;
                    end_time = None;
                    text = None;
                }
                pos = Some(line.parse::<usize>()?);
                state = State::Time;
            }
            State::Time => {
                let mut parts = line.split("-->");
                start_time = match parts.next() {
                    Some(v) => Some(duration_from_str(v)?),
                    None => None,
                };
                end_time = match parts.next() {
                    Some(v) => Some(duration_from_str(v)?),
                    None => None,
                };
                state = State::Text;
            }
            State::Text => {
                if line.is_empty() {
                    state = State::Pos;
                } else {
                    if let Some(ref mut txt) = text {
                        txt.push('\n');
                        txt.push_str(line);
                    } else {
                        text = Some(line.to_string());
                    }
                }
            }
        }
    }
    if let Some(ps) = pos {
        push_subtitle!(ps);
    }

    Ok(result)
}

pub fn read_from_file<P: AsRef<Path>>(path: P) -> Result<Vec<Subtitle>> {
    let mut file = File::open(path)?;
    let mut buf = String::new();
    file.read_to_string(&mut buf)?;
    parse(&buf)
}

/// Alias for std result
pub type Result<T> = StdResult<T, Error>;

/// Describes all errors that may occur.
#[derive(Debug)]
pub enum Error {
    /// An IO error occured. Contains `std::io::Error`.
    Io(IoError),
    /// An error when parsing an integer. Contains `std::num::ParseIntError`.
    ParseInt(ParseIntError),
    /// Subtitle text is missing
    MissingText,
    /// Subtitle end time is missing
    MissingEndTime,
    /// Subtitle start time is missing
    MissingStartTime,
    /// An error when parsing subtitle time
    ParseTime,
}

impl StdError for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Io(ref err) => err.description(),
            Error::ParseInt(ref err) => err.description(),
            Error::MissingText => "Subtitle text is missing",
            Error::MissingEndTime => "Subtitle end time is missing",
            Error::MissingStartTime => "Subtitle start time is missing",
            Error::ParseTime => "Invalid subtitle time",
        }
    }

    fn cause(&self) -> Option<&StdError> {
        match *self {
            Error::Io(ref err) => Some(err),
            Error::ParseInt(ref err) => Some(err),
            _ => None,
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, out: &mut fmt::Formatter) -> fmt::Result {
        write!(out, "{}", self.description())
    }
}

impl From<ParseIntError> for Error {
    fn from(err: ParseIntError) -> Error {
        Error::ParseInt(err)
    }
}

impl From<IoError> for Error {
    fn from(err: IoError) -> Error {
        Error::Io(err)
    }
}

macro_rules! parse_time_part {
    ($part:expr) => {{
        match $part {
            Some(val) => {
                val.trim().parse::<u64>()?
            },
            None => {
                return Err(Error::ParseTime);
            }
        }
    }}
}

fn duration_from_str(time: &str) -> Result<Duration> {
    let mut time = time.split(",");
    let (hours, mut minutes, mut seconds) = match time.next() {
        Some(val) => {
            let mut parts = val.split(":");
            let result = (
                parse_time_part!(parts.next()),
                parse_time_part!(parts.next()),
                parse_time_part!(parts.next()),
            );
            if parts.next().is_some() {
                return Err(Error::ParseTime);
            }
            result
        },
        None => {
            return Err(Error::ParseTime);
        }
    };
    let mut milliseconds = parse_time_part!(time.next());
    if time.next().is_some() {
        return Err(Error::ParseTime);
    }
    minutes += hours * 60;
    seconds += minutes * 60;
    milliseconds += seconds * 1000;
    Ok(Duration::new(milliseconds, 0))
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use ::{Result, UTF8_BOM, parse};

    #[test]
    fn it_works() {
        let source_without_bom = "1
00:00:58,392 --> 00:01:02,563
The war had all but ground to a halt
in the blink of an eye.

2
00:01:04,565 --> 00:01:08,986
Lucian, the most feared and ruthless
leader ever to rule the Lycan clan...

3
00:01:09,070 --> 00:01:11,656
...had finally been killed.

652
01:53:02,325 --> 01:53:06,162
Soon, Marcus will take the throne.
";
        let source_with_bom = format!("{}{}", UTF8_BOM, source_without_bom);

        fn assert_it_works(data: &str) -> Result<()> {
            let result = parse(&data)?;
            assert_eq!(result.len(), 4);
            assert_eq!(result[0].pos, 1);
            assert_eq!(result[0].start_time, Duration::new(58392, 0));
            assert_eq!(result[0].end_time, Duration::new(62563, 0));
            assert_eq!(result[0].text, "The war had all but ground to a halt\nin the blink of an eye.");
            assert_eq!(result[1].pos, 2);
            assert_eq!(result[1].start_time, Duration::new(64565, 0));
            assert_eq!(result[1].end_time, Duration::new(68986, 0));
            assert_eq!(result[1].text, "Lucian, the most feared and ruthless\nleader ever to rule the Lycan clan...");
            assert_eq!(result[2].pos, 3);
            assert_eq!(result[2].start_time, Duration::new(69070, 0));
            assert_eq!(result[2].end_time, Duration::new(71656, 0));
            assert_eq!(result[2].text, "...had finally been killed.");
            assert_eq!(result[3].pos, 652);
            assert_eq!(result[3].start_time, Duration::new(6782325, 0));
            assert_eq!(result[3].end_time, Duration::new(6786162, 0));
            assert_eq!(result[3].text, "Soon, Marcus will take the throne.");
            Ok(())
        }

        assert_it_works(&source_without_bom).expect("Failed to parse UTF-8 source without BOM");
        assert_it_works(&source_with_bom).expect("Failed to parse UTF-8 source with BOM");
    }

    // TODO: NEED MOAR TESTS!!!

}
<|endoftext|>