-- | Chan with size
module Control.Concurrent.SizedChan (SizedChan, newSizedChan, writeSizedChan, readSizedChan, tryReadSizedChan, peekSizedChan, tryPeekSizedChan, isEmptySizedChan) where

import Control.Concurrent.Chan
import Data.IORef

data SizedChan a = 
  SizedChan 
    (Chan a) -- ^ The channel
    (IORef Int) -- ^ Its size
    (IORef (Maybe a)) -- ^ Peeked payload

-- | Build and returns a new instance of 'SizedChan'.
newSizedChan :: IO (SizedChan a)
newSizedChan :: forall a. IO (SizedChan a)
newSizedChan =
  forall a. Chan a -> IORef Int -> IORef (Maybe a) -> SizedChan a
SizedChan
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IO (Chan a)
newChan
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (IORef a)
newIORef Int
0
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing

-- | Write a value to a 'SizedChan'.
writeSizedChan :: SizedChan a -> a -> IO ()
writeSizedChan :: forall a. SizedChan a -> a -> IO ()
writeSizedChan (SizedChan Chan a
chan IORef Int
sizeIORef IORef (Maybe a)
_) a
val = do
  forall a. Chan a -> a -> IO ()
writeChan Chan a
chan a
val
  forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef forall a. Enum a => a -> a
succ

-- | Read the next value from the 'SizedChan'. Blocks when the channel is empty.
readSizedChan :: SizedChan a -> IO a
readSizedChan :: forall a. SizedChan a -> IO a
readSizedChan (SizedChan Chan a
chan IORef Int
sizeIORef IORef (Maybe a)
peekedIORef) = do
  Maybe a
peeked <- forall a. IORef a -> IO a
readIORef IORef (Maybe a)
peekedIORef
  case Maybe a
peeked of
    -- return and remove the peeked value
    Just a
val -> do
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef forall a. Maybe a
Nothing
      forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef forall a. Enum a => a -> a
pred
      forall (m :: * -> *) a. Monad m => a -> m a
return a
val
    -- else read from the channel
    Maybe a
Nothing -> do
      a
val <- forall a. Chan a -> IO a
readChan Chan a
chan
      forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef forall a. Enum a => a -> a
pred
      forall (m :: * -> *) a. Monad m => a -> m a
return a
val

-- | A version of `readSizedChan` which does not block. Instead it returns Nothing if no value is available.
tryReadSizedChan :: SizedChan a -> IO (Maybe a)
tryReadSizedChan :: forall a. SizedChan a -> IO (Maybe a)
tryReadSizedChan (SizedChan Chan a
chan IORef Int
sizeIORef IORef (Maybe a)
peekedIORef) = do
  Maybe a
peeked <- forall a. IORef a -> IO a
readIORef IORef (Maybe a)
peekedIORef
  case Maybe a
peeked of
    -- return and remove the peeked value
    Just a
val -> do
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef forall a. Maybe a
Nothing
      forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef forall a. Enum a => a -> a
pred
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
val
    -- check the size before reading from the channel, to prevent blocking
    Maybe a
Nothing -> do
      Int
size <- forall a. IORef a -> IO a
readIORef IORef Int
sizeIORef
      if Int
size forall a. Eq a => a -> a -> Bool
== Int
0
        then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
        else do
          a
val <- forall a. Chan a -> IO a
readChan Chan a
chan
          forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef forall a. Enum a => a -> a
pred
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
val

-- | Peek the next value from the 'SizedChan' without removing it. Blocks when the channel is empty.
peekSizedChan :: SizedChan a -> IO a
peekSizedChan :: forall a. SizedChan a -> IO a
peekSizedChan (SizedChan Chan a
chan IORef Int
_ IORef (Maybe a)
peekedIORef) = do
  Maybe a
peeked <- forall a. IORef a -> IO a
readIORef IORef (Maybe a)
peekedIORef
  case Maybe a
peeked of
    -- return the peeked value
    Just a
val -> forall (m :: * -> *) a. Monad m => a -> m a
return a
val
    -- read from the channel instead
    Maybe a
Nothing -> do 
      a
val <- forall a. Chan a -> IO a
readChan Chan a
chan
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef (forall a. a -> Maybe a
Just a
val)
      forall (m :: * -> *) a. Monad m => a -> m a
return a
val

-- | A version of `peekSizedChan` which does not block. Instead it returns Nothing if no value is available.
tryPeekSizedChan :: SizedChan a -> IO (Maybe a)
tryPeekSizedChan :: forall a. SizedChan a -> IO (Maybe a)
tryPeekSizedChan (SizedChan Chan a
chan IORef Int
sizeIORef IORef (Maybe a)
peekedIORef) = do 
  Maybe a
peeked <- forall a. IORef a -> IO a
readIORef IORef (Maybe a)
peekedIORef
  case Maybe a
peeked of
    -- return the peeked value
    Just a
val -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
val
    -- check the size before reading from the channel, to prevent blocking
    Maybe a
Nothing -> do
      Int
size <- forall a. IORef a -> IO a
readIORef IORef Int
sizeIORef
      if Int
size forall a. Eq a => a -> a -> Bool
== Int
0
        then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
        else do
          a
val <- forall a. Chan a -> IO a
readChan Chan a
chan
          forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef (forall a. a -> Maybe a
Just a
val)
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
val

measureSizedChan :: SizedChan a -> IO Int
measureSizedChan :: forall a. SizedChan a -> IO Int
measureSizedChan (SizedChan Chan a
_ IORef Int
sizeIORef IORef (Maybe a)
_) = forall a. IORef a -> IO a
readIORef IORef Int
sizeIORef

isEmptySizedChan :: SizedChan a -> IO Bool
isEmptySizedChan :: forall a. SizedChan a -> IO Bool
isEmptySizedChan SizedChan a
chan = do
  Int
size <- forall a. SizedChan a -> IO Int
measureSizedChan SizedChan a
chan
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Int
size forall a. Eq a => a -> a -> Bool
== Int
0