use core::cell::UnsafeCell;
use core::sync::atomic::{AtomicUsize, Ordering};

use util::cpu_relax;

/// A synchronization primitive which can be used to run a one-time global
/// initialization. Unlike its std equivalent, this is generalized so that The
/// closure returns a value and it is stored. Once therefore acts something like
/// 1a future, too.
///
/// # Examples
///
/// ```
/// #![feature(const_fn)]
/// use spin;
///
/// static START: spin::Once<()> = spin::Once::new();
///
/// START.call_once(|| {
///     // run initialization here
/// });
/// ```
pub struct Once<T> {
    state: AtomicUsize,
    data: UnsafeCell<Option<T>>, // TODO remove option and use mem::uninitialized
}

unsafe impl<T: Sync + Sync> Sync for Once<T> {}
unsafe impl<T: Sync + Sync> Send for Once<T> {}

// Four states that a Once can be in, encoded into the lower bits of `state` in
// the Once structure.
const INCOMPLETE: usize = 0x0;
const RUNNING: usize = 0x1;
const COMPLETE: usize = 0x2;

impl<T> Once<T> {
    /// Creates a new `Once` value.
    pub const fn new() -> Once<T> {
        Once {
            state: AtomicUsize::new(INCOMPLETE),
            data: UnsafeCell::new(None),
        }
    }

    fn force_get<'a>(&'a self) -> &'a T {
        match unsafe { &*self.data.get() }.as_ref() {
            None    => unsafe { ::core::intrinsics::unreachable() },
            Some(p) => p,
        }
    }

    /// Performs an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `call_once` has been called,
    /// and otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling thread if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified). The
    /// returned pointer points to the return value of when of those
    /// initialization closures.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(const_fn)]
    /// use spin;
    ///
    /// static INIT: spin::Once<usize> = spin::Once::new();
    ///
    /// fn get_cached_val() -> usize {
    ///     *INIT.call_once(expensive_computation)
    /// }
    ///
    /// fn expensive_computation() -> usize {
    ///     // ...
    /// # 2
    /// }
    /// ```
    pub fn call_once<'a, F>(&'a self, builder: F) -> &'a T
        where F: FnOnce() -> T
    {
        let mut status = self.state.load(Ordering::SeqCst);

        if status == INCOMPLETE {
            status = self.state.compare_and_swap(INCOMPLETE,
                                                 RUNNING,
                                                 Ordering::SeqCst);
            if status == INCOMPLETE { // We init
                unsafe { *self.data.get() = Some(builder()) };
                status = COMPLETE;
                self.state.store(status, Ordering::SeqCst);
                // This next line is strictly an optomization
                return self.force_get();
            }
        }

        loop {
            match status {
                INCOMPLETE => unreachable!(),
                RUNNING => { // We spin
                    cpu_relax();
                    status = self.state.load(Ordering::SeqCst)
                },
                COMPLETE => return self.force_get(),
                _ => unsafe { ::core::intrinsics::unreachable() },
            }
        }
    }

    /// Returns a pointer iff the `Once` was previously initialized
    pub fn try<'a>(&'a self) -> Option<&'a T> {
        match self.state.load(Ordering::SeqCst) {
            COMPLETE => Some(self.force_get()),
            _        => None,
        }
    }

    /// Like try, but will spin if the `Once` is in the process of being
    /// initialized
    pub fn wait<'a>(&'a self) -> Option<&'a T> {
        loop {
            match self.state.load(Ordering::SeqCst) {
                INCOMPLETE => return None,
                RUNNING    => cpu_relax(), // We spin
                COMPLETE   => return Some(self.force_get()),
                _ => unsafe { ::core::intrinsics::unreachable() },
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::prelude::v1::*;

    use std::sync::mpsc::channel;
    use std::thread;
    use super::Once;

    #[test]
    fn smoke_once() {
        static O: Once<()> = Once::new();
        let mut a = 0;
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
    }

    #[test]
    fn smoke_once_value() {
        static O: Once<usize> = Once::new();
        let a = O.call_once(|| 1);
        assert_eq!(*a, 1);
        let b = O.call_once(|| 2);
        assert_eq!(*b, 1);
    }

    #[test]
    fn stampede_once() {
        static O: Once<()> = Once::new();
        static mut run: bool = false;

        let (tx, rx) = channel();
        for _ in 0..10 {
            let tx = tx.clone();
            thread::spawn(move|| {
                for _ in 0..4 { thread::yield_now() }
                unsafe {
                    O.call_once(|| {
                        assert!(!run);
                        run = true;
                    });
                    assert!(run);
                }
                tx.send(()).unwrap();
            });
        }

        unsafe {
            O.call_once(|| {
                assert!(!run);
                run = true;
            });
            assert!(run);
        }

        for _ in 0..10 {
            rx.recv().unwrap();
        }
    }

    #[test]
    fn try() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.try().is_none());
        INIT.call_once(|| 2);
        assert_eq!(INIT.try().map(|r| *r), Some(2));
    }

    #[test]
    fn try_no_wait() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.try().is_none());
        thread::spawn(move|| {
            INIT.call_once(|| loop { });
        });
        assert!(INIT.try().is_none());
    }


    #[test]
    fn wait() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.wait().is_none());
        INIT.call_once(|| 3);
        assert_eq!(INIT.wait().map(|r| *r), Some(3));
    }
}