module Control.Concurrent.Bag.TaskBufferSTM
( TaskBufferSTM (..)
, BufferType (..)
, SplitFunction
, takeFirst
, splitVertical
, splitHalf
, newChanBufferSTM
, newStackBufferSTM )
where
import Control.Concurrent.STM
( STM
, TChan
, newTChan
, writeTChan
, readTChan
, tryReadTChan
, isEmptyTChan
, unGetTChan
, retry )
import Control.Concurrent.STM.TStack
import Control.Monad ( liftM )
import Data.Maybe ( isNothing, fromJust )
import Control.Concurrent.Bag.BufferType
data TaskBufferSTM a = TaskBufferSTM {
writeBufferSTM :: a -> STM ()
, unGetBufferSTM :: a -> STM ()
, readBufferSTM :: STM a
, tryReadBufferSTM :: STM (Maybe a)
, isEmptyBufferSTM :: STM Bool
}
newChanBufferSTM :: STM (TaskBufferSTM r)
newChanBufferSTM = do
c <- newTChan
return $ TaskBufferSTM (writeTChan c) (unGetTChan c) (readTChan c) (tryReadTChan c) (isEmptyTChan c)
newStackBufferSTM :: STM (TaskBufferSTM r)
newStackBufferSTM = do
s <- newTStack
return $ TaskBufferSTM (writeTStack s) (writeTStack s) (readTStack s) (tryReadTStack s) (isEmptyTStack s)
type SplitFunction r = TaskBufferSTM (IO (Maybe r))
-> TaskBufferSTM (IO (Maybe r))
-> STM (IO (Maybe r))
takeFirst :: SplitFunction r
takeFirst _ from = readBufferSTM from
splitVertical :: SplitFunction r
splitVertical to from = do
first <- readBufferSTM from
splitRest to from
return first
where
splitRest to from = do
first <- tryReadBufferSTM from
second <- tryReadBufferSTM from
case (first, second) of
(Nothing, _) -> return ()
(Just f, Nothing) -> do
unGetBufferSTM from f
(Just f, Just s) -> do
splitRest to from
unGetBufferSTM to s
unGetBufferSTM from f
splitHalf :: SplitFunction r
splitHalf to from = do
splitRest to from 0
first <- readBufferSTM to
return first
where
splitRest to from n = do
first <- tryReadBufferSTM from
case first of
Nothing -> return n
Just f -> do
c <- splitRest to from (n+1)
if c > 0
then do
unGetBufferSTM to f
return (c2)
else do
unGetBufferSTM from f
return (c2)