/* mutex.rs
 *
 * Developed by Tim Walls <tim.walls@snowgoons.com>
 * Copyright (c) All Rights Reserved, Tim Walls
 */

// Imports ===================================================================
use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};
use avr_oxide::concurrency::{interrupt, scheduler, TryLockError, TryLockResult};
use avr_oxide::concurrency::thread;
use avr_oxide::concurrency::util::{ThreadSet, ThreadId};
use avr_oxide::OxideResult;
use avr_oxide::OxideResult::{Ok,Err};

// Declarations ==============================================================

struct SimpleThreadMutex {
  waiting_threads: ThreadSet,
  locked_by: Option<ThreadId>
}

/// A mutual exclusion primitive useful for protecting shared data
///
/// This mutex will block threads waiting for the lock to become available. The
/// mutex can also be statically initialized or created via a [`new`]
/// constructor. Each mutex has a type parameter which represents the data that
/// it is protecting. The data can only be accessed through the RAII guards
/// returned from [`lock`] and [`try_lock`], which guarantees that the data is only
/// ever accessed when the mutex is locked.
///
/// [`new`]: Mutex::new
/// [`lock`]: Mutex::lock
/// [`try_lock`]: Mutex::try_lock
pub struct Mutex<T: ?Sized> {
  lock: UnsafeCell<SimpleThreadMutex>,
  // Because T is unsized, this needs to always be the last field in the struct
  data: UnsafeCell<T>,
}

unsafe impl<T: ?Sized + Send> Send for Mutex<T> {}
unsafe impl<T: ?Sized + Send> Sync for Mutex<T> {}

/// An RAII implementation of a "scoped lock" of a mutex. When this structure is
/// dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be accessed through this guard via its
/// [`Deref`] and [`DerefMut`] implementations.
///
/// This structure is created by the [`lock`] and [`try_lock`] methods on
/// [`Mutex`].
///
/// [`lock`]: Mutex::lock
/// [`try_lock`]: Mutex::try_lock
#[must_use = "if unused the Mutex will immediately unlock"]
pub struct MutexGuard<'a, T: ?Sized + 'a> {
  mutex: &'a Mutex<T>
}

impl<T: ?Sized> !Send for MutexGuard<'_, T> {}
unsafe impl<T: ?Sized + Sync> Sync for MutexGuard<'_, T> {}


// Code ======================================================================
impl SimpleThreadMutex {
  fn unlocked() -> SimpleThreadMutex {
    SimpleThreadMutex {
      waiting_threads: ThreadSet::new(),
      locked_by: None
    }
  }

  fn try_lock(&mut self) -> OxideResult<(),TryLockError> {
    interrupt::isolated(|isotoken|{
      match self.locked_by {
        None => {
          self.locked_by.replace(scheduler::current_thread_id(isotoken));
          Ok(())
        },
        Some(thread_id) => {
          if thread_id == scheduler::current_thread_id(isotoken) {
            // lolwut?  I already have this lock.  Odd.  Probably shouldn't happen.
            Ok(())
          } else {
            Err(TryLockError::WouldBlock)
          }
        }
      }
    })
  }

  fn lock(&mut self) {
    loop {
      // We will try spinning a few times to get the lock, if the lock is
      // quickly going to be released then it's quicker than going down
      // the road of context switches
      for _i in 0..5 {
        if self.try_lock().is_ok() {
          return;
        }
      }

      // OK, that didn't work, so we should go on the wait list
      if interrupt::isolated(|isotoken|{
        match self.locked_by {
          // Still one last chance that we could acquire it now
          None => {
            self.locked_by.replace(scheduler::current_thread_id(isotoken));
            true
          },
          Some(_other_thread_id) => {
            // It's definitely another thread, 'cos if it was my thread
            // the try_lock would have succeeded (and there's no way *my*
            // thread could have acquired it between that try_lock and now
            // unless something supernatural is happening)
            self.waiting_threads.add_current_thread(isotoken);
            scheduler::set_current_thread_state(isotoken, scheduler::ThreadState::BlockedOnMutex);
            false
          }
        }
      }) {
        return
      } else {
        // OK, at this point everybody knows I am waiting - now I should yield
        thread::yield_now();
      }

      // At this point, I'm running again.  Maybe I can get the lock this
      // time?  Who knows...  We'll go round the loop and find out though.
    }
  }

  /// Release all threads that are waiting on this lock.  Which one,
  /// if any, will get the lock is up to luck and the scheduler.
  fn unlock_and_release_waiting(&mut self) {
    interrupt::isolated(|isotoken|{
      self.locked_by.take();
      scheduler::release_all_threads_and_clear(isotoken, &mut self.waiting_threads);
    });
  }
}

impl<T> Mutex<T> {
  /// Creates a new mutex in an unlocked state ready for use.
  ///
  pub fn new(t: T) -> Mutex<T> {
    Mutex {
      lock: UnsafeCell::new(SimpleThreadMutex::unlocked()),
      data: UnsafeCell::new(t),
    }
  }
}

impl<T: ?Sized> Mutex<T> {
  /// Acquires a mutex, blocking the current thread until it is able to do so.
  ///
  /// This function will block the local thread until it is available to acquire
  /// the mutex. Upon returning, the thread is the only thread with the lock
  /// held. An RAII guard is returned to allow scoped unlock of the lock. When
  /// the guard goes out of scope, the mutex will be unlocked.
  ///
  /// The exact behavior on locking a mutex in the thread which already holds
  /// the lock is left unspecified. However, this function will not return on
  /// the second call (it might panic or deadlock, for example).
  ///
  /// # Errors
  ///
  /// If another user of this mutex panicked while holding the mutex, then
  /// this call will return an error once the mutex is acquired.
  ///
  /// # Panics
  ///
  /// This function might panic when called if the lock is already held by
  /// the current thread.
  ///
  /// # Examples
  ///
  /// ```
  /// use std::sync::{Arc, Mutex};
  /// use std::thread;
  ///
  /// let mutex = Arc::new(Mutex::new(0));
  /// let c_mutex = Arc::clone(&mutex);
  ///
  /// thread::spawn(move || {
  ///     *c_mutex.lock().unwrap() = 10;
  /// }).join().expect("thread::spawn failed");
  /// assert_eq!(*mutex.lock().unwrap(), 10);
  /// ```
  pub fn lock(&self) -> MutexGuard<'_, T> {
    unsafe {
      let lock  = &mut *self.lock.get();

      lock.lock();
      MutexGuard::new(self)
    }
  }

  /// Attempts to acquire this lock.
  ///
  /// If the lock could not be acquired at this time, then [`Err`] is returned.
  /// Otherwise, an RAII guard is returned. The lock will be unlocked when the
  /// guard is dropped.
  ///
  /// This function does not block.
  ///
  /// # Errors
  ///
  /// If the mutex could not be acquired because it is already locked, then
  /// this call will return the [`WouldBlock`] error.
  ///
  /// [`WouldBlock`]: TryLockError::WouldBlock
  ///
  /// # Examples
  ///
  /// ```
  /// use std::sync::{Arc, Mutex};
  /// use std::thread;
  ///
  /// let mutex = Arc::new(Mutex::new(0));
  /// let c_mutex = Arc::clone(&mutex);
  ///
  /// thread::spawn(move || {
  ///     let mut lock = c_mutex.try_lock();
  ///     if let Ok(ref mut mutex) = lock {
  ///         **mutex = 10;
  ///     } else {
  ///         println!("try_lock failed");
  ///     }
  /// }).join().expect("thread::spawn failed");
  /// assert_eq!(*mutex.lock().unwrap(), 10);
  /// ```
  pub fn try_lock(&self) -> TryLockResult<MutexGuard<'_, T>> {
    unsafe {
      let lock  = &mut *self.lock.get();

      match lock.try_lock(){
        Ok(_) => {
          Ok(MutexGuard::new(self))
        },
        Err(e) => {
          Err(e)
        }
      }
    }
  }



  /// Immediately drops the guard, and consequently unlocks the mutex.
  ///
  /// This function is equivalent to calling [`drop`] on the guard but is more self-documenting.
  /// Alternately, the guard will be automatically dropped when it goes out of scope.
  ///
  pub fn unlock(guard: MutexGuard<'_, T>) {
    drop(guard);
  }

  /// Consumes this mutex, returning the underlying data.
  ///
  /// # Errors
  ///
  /// If another user of this mutex panicked while holding the mutex, then
  /// this call will return an error instead.
  ///
  /// # Examples
  ///
  /// ```
  /// use std::sync::Mutex;
  ///
  /// let mutex = Mutex::new(0);
  /// assert_eq!(mutex.into_inner().unwrap(), 0);
  /// ```
  pub fn into_inner(self) -> T
  where
    T: Sized
  {
    self.data.into_inner()
  }

  /// Returns a mutable reference to the underlying data.
  ///
  /// Since this call borrows the `Mutex` mutably, no actual locking needs to
  /// take place -- the mutable borrow statically guarantees no locks exist.
  ///
  /// # Errors
  ///
  /// If another user of this mutex panicked while holding the mutex, then
  /// this call will return an error instead.
  ///
  /// # Examples
  ///
  /// ```
  /// use std::sync::Mutex;
  ///
  /// let mut mutex = Mutex::new(0);
  /// *mutex.get_mut().unwrap() = 10;
  /// assert_eq!(*mutex.lock().unwrap(), 10);
  /// ```
  pub fn get_mut(&mut self) -> &mut T {
    self.data.get_mut()
  }
}

impl<T> From<T> for Mutex<T> {
  /// Creates a new mutex in an unlocked state ready for use.
  /// This is equivalent to [`Mutex::new`].
  fn from(t: T) -> Self {
    Mutex::new(t)
  }
}

impl<T: ?Sized + Default> Default for Mutex<T> {
  /// Creates a `Mutex<T>`, with the `Default` value for T.
  fn default() -> Mutex<T> {
    Mutex::new(Default::default())
  }
}

impl<'mutex, T: ?Sized> MutexGuard<'mutex, T> {
  unsafe fn new(lock: &'mutex Mutex<T>) -> MutexGuard<'mutex, T> {
    MutexGuard { mutex: lock }
  }
}

impl<T: ?Sized> Deref for MutexGuard<'_, T> {
  type Target = T;

  fn deref(&self) -> &T {
    unsafe { &*self.mutex.data.get() }
  }
}

impl<T: ?Sized> DerefMut for MutexGuard<'_, T> {
  fn deref_mut(&mut self) -> &mut T {
    unsafe { &mut *self.mutex.data.get() }
  }
}

impl<T: ?Sized> Drop for MutexGuard<'_, T> {
  #[inline]
  fn drop(&mut self) {
    unsafe {
      (&mut *self.mutex.lock.get()).unlock_and_release_waiting();
    }
  }
}

// Tests =====================================================================