{-# LANGUAGE TypeFamilies #-}
module PrioritySync.Internal.RoomCore
    (Room,
     newRoom,
     Claim,
     ClaimMode(..),
     claimedRoom,
     claimedThread,
     userData,
     approve,
     claim_,
     inUse,
     isEmpty)
    where

import PrioritySync.Internal.UserData
import Data.Unique
import Data.Set as Set
import Data.Map as Map
import GHC.Conc
import Control.Monad
import Control.Exception

-- | A resource pool, parameterized against arbitrary user data.
data Room u = Room {
    room_user_data :: (u,Unique),
    room_occupants :: (TVar (Set ThreadId)),
    room_is_empty :: (TVar Bool) }

type instance UserData (Room u) = u

instance Eq (Room u) where
    (==) r1 r2 = snd (room_user_data r1) == snd (room_user_data r2)

instance Ord (Room u) where
    compare r1 r2 = compare (snd $ room_user_data r1) (snd $ room_user_data r2)

-- | A 'Claim', or attempt to acquire or release a 'Room'.
data Claim u = Claim (Room u) ThreadId ClaimMode (TVar Bool)

type instance UserData (Claim u) = u

data ClaimMode = Acquire | Release deriving (Eq)

-- | Create a new Room with some arbitrary user data.
newRoom :: u -> IO (Room u)
newRoom u = return Room `ap` liftM ((,) u) newUnique `ap` newTVarIO Set.empty `ap` newTVarIO True

-- | Get the user data associated with a 'Room'.
userData :: Room u -> u
userData = fst . room_user_data

-- | Whether a Claim is to acquire or release a room.
claimMode :: Claim u -> ClaimMode
claimMode (Claim _ _ b _) = b

-- | Get the 'Room' target of a 'Claim'.
claimedRoom :: Claim u -> Room u
claimedRoom (Claim m _ _ _) = m

-- | Get the thread attempting a 'Claim'.
claimedThread :: Claim u -> ThreadId
claimedThread (Claim _ t _ _) = t

-- | Approve a claim.  This actually acquires or releases a 'Room'.
approve :: Claim u -> STM ()
approve (Claim r me want claim_var) =
    do claim_state <- readTVar claim_var
       when (claim_state == False) $ 
           do writeTVar claim_var True
              writeTVar (room_occupants r) . (case want of Acquire -> Set.insert; Release -> Set.delete) me =<< readTVar (room_occupants r)
              watchRoom r

-- | Set the 'room_is_empty' flag.
watchRoom :: Room u -> STM ()
watchRoom r =
    do is_empty <- readTVar (room_is_empty r)
       occus <- readTVar (room_occupants r)
       when (is_empty && (not $ Set.null occus)) $ writeTVar (room_is_empty r) True
       when (not is_empty && Set.null occus) $ writeTVar (room_is_empty r) False

-- | Acquire and/or release some rooms for the duration of a critical section.
--
-- * which 'Room's to 'Acquire', and later release, or 'Release' and later reacquire for the duration of the critical section.
--
-- * a transaction to 'approve' all entering 'Claim's
--
-- * a transaction to 'approve' all exiting 'Claim's
--
-- * a transaction to run one or more times if and only if this thread is waiting for approval
--
-- * the critical section
--
-- A separate 'Claim' is generated each time a Room needs to be acquired.  The critical
-- section will not enter until every claim has been 'approve'd.
--
-- When the critical section exits, an inverse group of claims will be generated, and the critical
-- section will not exit until those claims have been 'approve'd.
--
-- It is guaranteed that when and only when all 'Claim's have been 'approve'd, the waiting thread will enter
-- (or exit) the critical section.  The lock on each 'Room' is acquired when it's 'Claim' is 'approve'd,
-- not when the critical section is entered.
--
-- 'Claim's may be 'approve'd from any transaction, even from another thread.
--
claim_ :: Map (Room u) ClaimMode -> ([Claim u] -> STM ()) -> ([Claim u] -> STM ()) -> STM () -> IO a -> IO a
claim_ entering_rooms_map approveEnteringSTM approveExitingSTM waitingSTM actionIO =
    do let entering_rooms = Map.toList entering_rooms_map
       me <- myThreadId
       -- transition: generate and request (but do not wait for) approval for all claims
       let transition rooms approveSTM = atomically $
               (\claims -> approveSTM (Prelude.filter ((== Acquire) . claimMode) claims) >> return claims) =<< -- request approval of Acquire claims.
                   (mapM $ \c -> when (claimMode c == Release) (approve c) >> return c) =<< --auto-approve Release claims.
                   (mapM $ \(m,want) -> liftM (Claim m me want) (newTVar False)) =<< -- build claims
                   filterM (\(m,want) -> liftM ((/= (want == Acquire)) . Set.member me) $ inUse m) rooms -- get the difference between the rooms we want and the rooms we have
       -- confirm: wait for all claims to be approved
       let confirm claims = forM_ claims $ \(Claim _ _ _ claim_var) ->
               do claim_state <- readTVar claim_var
                  unless claim_state retry
       -- confirm or perform the user's wait action, return True iff we have confirmed
       let confirmWithWaitAction claims=
               do done <- atomically $ (confirm claims >> return True) `orElse` (waitingSTM >> return False)
                  unless done $ confirmWithWaitAction claims
       -- when entering we signal the claims (transition) and then wait on the approval of the claims (confirm)
       -- then we generate a list of inverse claims for when it comes time to exit the critical section
       -- If an exception is thrown from the approve*STM or waitingSTM, this would leave dangling 
       -- room locks, therefore, we force the rooms into their former state, ignoring any 
       -- constraints that might have been placed on them.  On the theory that it's better to violate 
       -- constraints than to leave dangling locks.
       let transitionAndConfirm approveSTM rooms = flip finally (transition rooms $ mapM_ approve) $
               do claims <- transition rooms approveSTM
                  confirmWithWaitAction claims
                  return $ Prelude.map (\(Claim m _ want _) -> (m,case want of Acquire -> Release; Release -> Acquire)) claims
       bracket (transitionAndConfirm approveEnteringSTM entering_rooms)
               (transitionAndConfirm approveExitingSTM)
               (const actionIO)

-- | Get all 'ThreadId's that are currently claiming this 'Room'.
inUse :: Room u -> STM (Set ThreadId)
inUse = readTVar . room_occupants

-- | True iff a Room is empty.
isEmpty :: Room u -> STM Bool
isEmpty = readTVar . room_is_empty