{-# LANGUAGE BangPatterns #-}

-- |
-- Module    : Z3.RLock
-- Copyright  : (c) 2010-2011 Bas van Dijk & Roel van Dijk
-- License   : BSD3
-- Maintainer: Iago Abal <mail@iagoabal.eu>,
--             David Castro <david.castro.dcp@gmail.com>
--
-- A minimal implementation of a re-entrant lock, adapted from https://github.com/basvandijk/concurrent-extra
--
--

module Z3.RLock where

import qualified Z3.Lock as Lock

import Z3.Lock(Lock)

import Control.Concurrent( ThreadId, myThreadId, MVar, putMVar, takeMVar, newMVar )
import Control.Exception( mask_, bracket_ )
import Control.Applicative( liftA2 )

{-| A reentrant lock is in one of two states: \"locked\" or \"unlocked\". When
the lock is in the \"locked\" state it has two additional properties:

* Its /owner/: the thread that acquired the lock.

* Its /acquired count/: how many times its owner acquired the lock.
-}
newtype RLock = RLock {RLock -> MVar (State, Lock)
un :: MVar (State, Lock)}
    deriving RLock -> RLock -> Bool
(RLock -> RLock -> Bool) -> (RLock -> RLock -> Bool) -> Eq RLock
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RLock -> RLock -> Bool
== :: RLock -> RLock -> Bool
$c/= :: RLock -> RLock -> Bool
/= :: RLock -> RLock -> Bool
Eq
{-| The state of an 'RLock'.

* 'Nothing' indicates an \"unlocked\" state.

* @'Just' (tid, n)@ indicates a \"locked\" state where the thread identified by
@tid@ acquired the lock @n@ times.
-}
type State = Maybe (ThreadId, Integer)

-- | Create a reentrant lock in the \"unlocked\" state.
new :: IO RLock
new :: IO RLock
new = do Lock
lock <- IO Lock
Lock.new
         MVar (State, Lock) -> RLock
RLock (MVar (State, Lock) -> RLock)
-> IO (MVar (State, Lock)) -> IO RLock
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (State, Lock) -> IO (MVar (State, Lock))
forall a. a -> IO (MVar a)
newMVar (State
forall a. Maybe a
Nothing, Lock
lock)

{-|
Acquires the 'RLock'. Blocks if another thread has acquired the 'RLock'.
@acquire@ behaves as follows:

* When the state is \"unlocked\", @acquire@ changes the state to \"locked\"
with the current thread as owner and an acquired count of 1.

* When the state is \"locked\" and the current thread owns the lock @acquire@
only increments the acquired count.

* When the state is \"locked\" and the current thread does not own the lock
@acquire@ /blocks/ until the owner releases the lock. If the thread that called
@acquire@ is woken upon release of the lock it will take ownership and change
the state to \"locked\" with an acquired count of 1.

There are two further important properties of @acquire@:

* @acquire@ is single-wakeup. That is, if there are multiple threads blocked on
@acquire@, and the lock is released, only one thread will be woken up. The
runtime guarantees that the woken thread completes its @acquire@ operation.

* When multiple threads are blocked on @acquire@ they are woken up in FIFO
order. This is useful for providing fairness properties of abstractions built
using locks. (Note that this differs from the Python implementation where the
wake-up order is undefined.)
-}
acquire :: RLock -> IO ()
acquire :: RLock -> IO ()
acquire (RLock MVar (State, Lock)
mv) = do
  ThreadId
myTID <- IO ThreadId
myThreadId
  IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ let acq :: IO ()
acq = do t :: (State, Lock)
t@(State
mb, Lock
lock) <- MVar (State, Lock) -> IO (State, Lock)
forall a. MVar a -> IO a
takeMVar MVar (State, Lock)
mv
                       case State
mb of
                         State
Nothing          -> do Lock -> IO ()
Lock.acquire Lock
lock
                                                MVar (State, Lock) -> (State, Lock) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (State, Lock)
mv ((ThreadId, Integer) -> State
forall a. a -> Maybe a
Just (ThreadId
myTID, Integer
1), Lock
lock)
                         Just (ThreadId
tid, Integer
n)
                           | ThreadId
myTID ThreadId -> ThreadId -> Bool
forall a. Eq a => a -> a -> Bool
== ThreadId
tid -> let !sn :: Integer
sn = Integer -> Integer
forall a. Enum a => a -> a
succ Integer
n
                                             in MVar (State, Lock) -> (State, Lock) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (State, Lock)
mv ((ThreadId, Integer) -> State
forall a. a -> Maybe a
Just (ThreadId
tid, Integer
sn), Lock
lock)
                           | Bool
otherwise    -> do MVar (State, Lock) -> (State, Lock) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (State, Lock)
mv (State, Lock)
t
                                                Lock -> IO ()
Lock.wait Lock
lock
                                                IO ()
acq
          in IO ()
acq

{-| @release@ decrements the acquired count. When a lock is released with an
acquired count of 1 its state is changed to \"unlocked\".

Note that it is both an error to release a lock in the \"unlocked\" state and to
release a lock that is not owned by the current thread.

If there are any threads blocked on 'acquire' the thread that first called
@acquire@ will be woken up.
-}
release :: RLock -> IO ()
release :: RLock -> IO ()
release (RLock MVar (State, Lock)
mv) = do
  ThreadId
myTID <- IO ThreadId
myThreadId
  IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    t :: (State, Lock)
t@(State
mb, Lock
lock) <- MVar (State, Lock) -> IO (State, Lock)
forall a. MVar a -> IO a
takeMVar MVar (State, Lock)
mv
    let err :: [Char] -> IO b
err [Char]
msg = do MVar (State, Lock) -> (State, Lock) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (State, Lock)
mv (State, Lock)
t
                     [Char] -> IO b
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO b) -> [Char] -> IO b
forall a b. (a -> b) -> a -> b
$ [Char]
"Z3.RLock.release: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
msg
    case State
mb of
      State
Nothing -> [Char] -> IO ()
forall {b}. [Char] -> IO b
err [Char]
"Can't release an unacquired RLock!"
      Just (ThreadId
tid, Integer
n)
        | ThreadId
myTID ThreadId -> ThreadId -> Bool
forall a. Eq a => a -> a -> Bool
== ThreadId
tid -> if Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1
                          then do Lock -> IO ()
Lock.release Lock
lock
                                  MVar (State, Lock) -> (State, Lock) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (State, Lock)
mv (State
forall a. Maybe a
Nothing, Lock
lock)
                          else let !pn :: Integer
pn = Integer -> Integer
forall a. Enum a => a -> a
pred Integer
n
                               in MVar (State, Lock) -> (State, Lock) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (State, Lock)
mv ((ThreadId, Integer) -> State
forall a. a -> Maybe a
Just (ThreadId
tid, Integer
pn), Lock
lock)
        | Bool
otherwise -> [Char] -> IO ()
forall {b}. [Char] -> IO b
err [Char]
"Calling thread does not own the RLock!"

{-| A convenience function which first acquires the lock and then
performs the computation. When the computation terminates, whether
normally or by raising an exception, the lock is released.

Note that: @with = 'liftA2' 'bracket_' 'acquire' 'release'@.
-}
with :: RLock -> IO a -> IO a
with :: forall a. RLock -> IO a -> IO a
with = (IO () -> IO () -> IO a -> IO a)
-> (RLock -> IO ()) -> (RLock -> IO ()) -> RLock -> IO a -> IO a
forall a b c.
(a -> b -> c) -> (RLock -> a) -> (RLock -> b) -> RLock -> c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ RLock -> IO ()
acquire RLock -> IO ()
release