{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DoRec              #-}

-----------------------------------------------------------------------
-- |
-- Module     : Control/Concurrent/Stream.hs
-- Copyright  : (c) Ivan Tomac 2012
-- License    : BSD3
--
-- Maintainer : ivan `dot` tomac `at` google `dot` com
-- Stability  : experimental
--
-- Concurrent streams.
--
-----------------------------------------------------------------------

module Control.Concurrent.Stream
  ( Stream
  , newStream
  , readStream
  ) where

import Control.Applicative
import Control.Arrow
import Control.Concurrent
import Control.Monad
import Data.Foldable (foldrM)
import Data.IORef
import Data.Typeable
import System.IO.Unsafe
import System.Mem.Weak

import Control.Concurrent.Queue.Nonblocking

-- A stream contains an MVar with a list of weak pointers to functions
-- that write to the stream and any derivative streams. When a stream
-- is written to, data is broadcast to all derivative streams.
-- Weak pointers are used so that writers to any streams that have been
-- garbage collected can be disposed of.

data Stream a = Repeat a | Stream (MVar [Weak (a -> IO ())]) (Queue a)
    deriving (Eq, Typeable)

-- | Returns a new 'Stream' and a function to write to it.
newStream :: IO (a -> IO (), Stream a)
newStream = do
    q <- newQ'
    r <- newIORef q
    rec v <- newMVar [w]
        w <- weak v (modifyIORef_ r . flip unsafeWriteQ)
    return (modifyMVar_ v . broadcast, Stream v q)

-- | Returns the next value in the 'Stream' if one is available, along
-- with the rest of the 'Stream'.
-- If the 'Stream' is empty, it returns 'Nothing'.
readStream :: Stream a -> IO (Maybe (a, Stream a))
readStream st @ (Repeat   x) = return $ Just (x, st)
readStream      (Stream v q) = (fmap $ second $ Stream v) <$> readQ q

instance Functor Stream where
    fmap f (Repeat   x) = Repeat $ f x
    fmap f (Stream v q) = unsafePerformIO $ do
        (write, st @ (Stream u _)) <- first (. f) <$> newStream
        -- Go through all the items in the stream and push them through
        -- to the new stream. Then add the writer for the new stream
        -- to the list of writers.
        modifyMVar_ v $ (unfoldM readQ q >>= mapM_ write . fst >>)
                      . addWriter u write
        return st

instance Applicative Stream where
    pure = Repeat

    Repeat   f  <*> st          = fmap f     st
    st          <*> Repeat   x  = fmap ($ x) st
    Stream v qf <*> Stream u qx = unsafePerformIO $ do
        (write, st @ (Stream s _)) <- newStream
        t                          <- newMVar (qf, qx)
        -- Pull values from both streams, applying the function from
        -- the first stream to the value from the second stream, until
        -- one of the streams is out of values.
        let pull (qf', qx') =
                liftA2 (,) <$> readQ qf' <*> readQ qx' >>=
                maybe (return (qf', qx'))
                      (uncurry (>>) . (write . app *** pull)
                                    . transpose)
        -- Adds a writer that ignores its argument and instead pulls
        -- data directly from the queues. Since addWriter always adds
        -- the writer to the very end of the list of writers, the
        -- queues data is pulled from should already have been written
        -- to by the time the writer is called.
            modify = addWriter s (const $ modifyMVar_ t pull)

        modifyMVar_ v modify
        modifyMVar_ u modify
        modifyMVar_ t pull
        return st

addWriter :: k     -> (a -> IO ())
          ->    [Weak (a -> IO ())]
          -> IO [Weak (a -> IO ())]
addWriter k write ws =
    weak k write >>= mapMaybeM alive . (ws ++) . (: [])

broadcast :: a -> [Weak (a -> IO ())] -> IO [Weak (a -> IO ())]
broadcast x = foldrM (liftA2 (fmap . (>>=)) deRefWeak f) []
    where f _ ws  Nothing     =            return      ws
          f w ws (Just write) = write x >> return (w : ws)

transpose :: ((a, b), (c, d)) -> ((a, c), (b, d))
transpose ((x, y), (z, w)) = ((x, z), (y, w))

unfoldM :: Monad m => (a -> m (Maybe (b, a))) -> a -> m ([b], a)
unfoldM f x = f x >>= return ([], x) `maybe` uncurry g
    where g y = liftM (first (y :)) . unfoldM f

mapMaybeM :: Monad m => (a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM f = foldrM g []
    where g x ys = f x >>= return . maybe ys (: ys)

modifyIORef_ :: IORef a -> (a -> IO a) -> IO ()
modifyIORef_ r f = readIORef r >>= f >>= writeIORef r

alive :: Weak a -> IO (Maybe (Weak a))
alive = fmap . fmap . const <*> deRefWeak

weak :: k -> a -> IO (Weak a)
weak = ($ Nothing) . flip . mkWeak