{-# LANGUAGE BangPatterns #-}
module Control.Concurrent.WVar
(
WVar
, newWVar
, takeWVar
, tryTakeWVar
, putWVar
, readWVar
, readFreshWVar
, tryReadFreshWVar
, WCached(..)
, WTicket
, cacheWVar
, recacheWCached
, readWTicket
, takeWCached
, tryTakeWCached
, putWCached
, tryPutWCached
, readWCached
, readFreshWCached
, tryReadFreshWCached
)
where
import Control.Concurrent.MVar (MVar)
import qualified Control.Concurrent.MVar as MVar
import qualified Control.Monad as M
import qualified Data.Atomics as Atm
import Data.IORef (IORef)
import qualified Data.IORef as Ref
newtype WVar a = WVar (IORef (WContent a))
deriving Eq
{-# INLINE newWVar #-}
newWVar :: a -> IO (WVar a)
newWVar !a = WVar <$> Ref.newIORef WContent
{ wvalue = a
, wstate = Fresh
}
{-# INLINE takeWVar #-}
takeWVar :: WVar a -> IO a
takeWVar wv = do
wc <- cacheWVar wv
wt1 <- takeWCached wc
return $ readWTicket wt1
{-# INLINE tryTakeWVar #-}
tryTakeWVar :: WVar a -> IO (Bool, a)
tryTakeWVar wv = cacheWVar wv >>= go
where
go wc = do
(suc, wt1) <- tryTakeWCached wc
case (suc, wstate (Atm.peekTicket wt1)) of
(False, Fresh) -> go $ WCached wv wt1
_ -> return (suc, readWTicket wt1)
{-# INLINE putWVar #-}
putWVar :: WVar a -> a -> IO ()
putWVar wv a = do
wc <- cacheWVar wv
M.void $ putWCached wc a
{-# INLINE readWVar #-}
readWVar :: WVar a -> IO a
readWVar (WVar ref) = wvalue <$> Ref.readIORef ref
{-# INLINE readFreshWVar #-}
readFreshWVar :: WVar a -> IO a
readFreshWVar wv = do
wc <- cacheWVar wv
readWTicket <$> readFreshWCached wc
{-# INLINE tryReadFreshWVar #-}
tryReadFreshWVar :: WVar a -> IO (Bool, a)
tryReadFreshWVar wv = do
wc <- cacheWVar wv
(suc, wt1) <- tryReadFreshWCached wc
return (suc, readWTicket wt1)
data WCached a = WCached
{ cachedVar :: {-# UNPACK #-} !(WVar a)
, cachedTicket :: WTicket a
} deriving Eq
instance Show a => Show (WCached a) where
show (WCached _ wt) = "WCached " ++ show (readWTicket wt)
type WTicket a = Atm.Ticket (WContent a)
data WContent a = WContent
{ wvalue :: !a
, wstate :: !(WState a)
}
instance Show a => Show (WContent a) where
show wcnt = concat [show (wstate wcnt), " ", show (wvalue wcnt)]
data WState a =
Fresh
| Updating
| Waiting {-# UNPACK #-} !(MVar (WTicket a))
instance Show (WState a) where
show Fresh = "Fresh"
show Updating = "Updating"
show (Waiting _) = "Waiting"
{-# INLINE cacheWVar #-}
cacheWVar :: WVar a -> IO (WCached a)
cacheWVar wv@(WVar ref) = do
wt <- Atm.readForCAS ref
return WCached { cachedVar = wv, cachedTicket = wt }
recacheWCached :: WCached a -> IO (WCached a)
recacheWCached = cacheWVar . cachedVar
{-# INLINE readWTicket #-}
readWTicket :: WTicket a -> a
readWTicket = wvalue . Atm.peekTicket
{-# INLINE takeWCached #-}
takeWCached :: WCached a -> IO (WTicket a)
takeWCached wc@(WCached wv@(WVar ref) wt0) = do
let wcnt0 = Atm.peekTicket wt0
(suc, wt2) <- case wstate wcnt0 of
Fresh -> do
(suc, wt1) <- Atm.casIORef ref wt0 wcnt0 { wstate = Updating }
if suc
then return (True, wt1)
else return (False, wt1)
Updating -> do
wt1 <- beginWaitWCached wc
return (False, wt1)
Waiting mv -> do
wt1 <- MVar.readMVar mv
return (False, wt1)
if suc
then return wt2
else takeWCached $ WCached wv wt2
{-# INLINE tryTakeWCached #-}
tryTakeWCached :: WCached a -> IO (Bool, WTicket a)
tryTakeWCached (WCached (WVar ref) wt0) = do
let wcnt0 = Atm.peekTicket wt0
case wstate wcnt0 of
Fresh -> Atm.casIORef ref wt0 wcnt0 { wstate = Updating }
Updating -> return (False, wt0)
Waiting _ -> return (False, wt0)
{-# INLINE putWCached #-}
putWCached :: WCached a -> a -> IO (WTicket a)
putWCached wc0 !a = do
(suc, wt1) <- tryPutWCached wc0 a
if suc
then return wt1
else do
let wc1 = wc0 { cachedTicket = wt1 }
putWCached wc1 a
{-# INLINE tryPutWCached #-}
tryPutWCached :: WCached a -> a -> IO (Bool, WTicket a)
tryPutWCached (WCached (WVar ref) wt0) !a = do
let !wcnt1 = WContent { wvalue = a, wstate = Fresh }
(suc, wt1) <- Atm.casIORef ref wt0 wcnt1
M.when suc $ case wstate $ Atm.peekTicket wt0 of
Waiting mv -> MVar.putMVar mv wt1
_ -> return ()
return (suc, wt1)
{-# INLINE readWCached #-}
readWCached :: WCached a -> a
readWCached (WCached _ wt0) = readWTicket wt0
{-# INLINE readFreshWCached #-}
readFreshWCached :: WCached a -> IO (WTicket a)
readFreshWCached wc@(WCached wv wt0) = do
let wcnt0 = Atm.peekTicket wt0
(suc, wt2) <- case wstate wcnt0 of
Fresh -> return (True, wt0)
Updating -> do
wt1 <- beginWaitWCached wc
return (False, wt1)
Waiting mv -> do
wt1 <- MVar.readMVar mv
return (True, wt1)
if suc
then return wt2
else readFreshWCached $ WCached wv wt2
{-# INLINE tryReadFreshWCached #-}
tryReadFreshWCached :: WCached a -> IO (Bool, WTicket a)
tryReadFreshWCached (WCached _ wt0) = case wstate $ Atm.peekTicket wt0 of
Fresh -> return (True, wt0)
_ -> return (False, wt0)
{-# INLINE beginWaitWCached #-}
beginWaitWCached :: WCached a -> IO (WTicket a)
beginWaitWCached (WCached (WVar ref) wt0) = do
mv <- MVar.newEmptyMVar
let wcnt0 = Atm.peekTicket wt0
wcnt1 = wcnt0 { wstate = Waiting mv }
snd <$> Atm.casIORef ref wt0 wcnt1