-- | 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 =
  Chan a -> IORef Int -> IORef (Maybe a) -> SizedChan a
forall a. Chan a -> IORef Int -> IORef (Maybe a) -> SizedChan a
SizedChan
    (Chan a -> IORef Int -> IORef (Maybe a) -> SizedChan a)
-> IO (Chan a) -> IO (IORef Int -> IORef (Maybe a) -> SizedChan a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Chan a)
forall a. IO (Chan a)
newChan
    IO (IORef Int -> IORef (Maybe a) -> SizedChan a)
-> IO (IORef Int) -> IO (IORef (Maybe a) -> SizedChan a)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
    IO (IORef (Maybe a) -> SizedChan a)
-> IO (IORef (Maybe a)) -> IO (SizedChan a)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe a -> IO (IORef (Maybe a))
forall a. a -> IO (IORef a)
newIORef Maybe a
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
  Chan a -> a -> IO ()
forall a. Chan a -> a -> IO ()
writeChan Chan a
chan a
val
  IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef Int -> Int
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 <- IORef (Maybe a) -> IO (Maybe a)
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
      IORef (Maybe a) -> Maybe a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef Maybe a
forall a. Maybe a
Nothing
      IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef Int -> Int
forall a. Enum a => a -> a
pred
      a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
val
    -- else read from the channel
    Maybe a
Nothing -> do
      a
val <- Chan a -> IO a
forall a. Chan a -> IO a
readChan Chan a
chan
      IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef Int -> Int
forall a. Enum a => a -> a
pred
      a -> IO a
forall a. a -> IO a
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 <- IORef (Maybe a) -> IO (Maybe a)
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
      IORef (Maybe a) -> Maybe a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef Maybe a
forall a. Maybe a
Nothing
      IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef Int -> Int
forall a. Enum a => a -> a
pred
      Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
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 <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
sizeIORef
      if Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
        else do
          a
val <- Chan a -> IO a
forall a. Chan a -> IO a
readChan Chan a
chan
          IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
sizeIORef Int -> Int
forall a. Enum a => a -> a
pred
          Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
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 <- IORef (Maybe a) -> IO (Maybe a)
forall a. IORef a -> IO a
readIORef IORef (Maybe a)
peekedIORef
  case Maybe a
peeked of
    -- return the peeked value
    Just a
val -> a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
val
    -- read from the channel instead
    Maybe a
Nothing -> do 
      a
val <- Chan a -> IO a
forall a. Chan a -> IO a
readChan Chan a
chan
      IORef (Maybe a) -> Maybe a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef (a -> Maybe a
forall a. a -> Maybe a
Just a
val)
      a -> IO a
forall a. a -> IO a
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 <- IORef (Maybe a) -> IO (Maybe a)
forall a. IORef a -> IO a
readIORef IORef (Maybe a)
peekedIORef
  case Maybe a
peeked of
    -- return the peeked value
    Just a
val -> Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
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 <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
sizeIORef
      if Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
        else do
          a
val <- Chan a -> IO a
forall a. Chan a -> IO a
readChan Chan a
chan
          IORef (Maybe a) -> Maybe a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe a)
peekedIORef (a -> Maybe a
forall a. a -> Maybe a
Just a
val)
          Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
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)
_) = IORef Int -> IO Int
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 <- SizedChan a -> IO Int
forall a. SizedChan a -> IO Int
measureSizedChan SizedChan a
chan
  Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0