{-# LANGUAGE CPP #-}
module LockedChan
  ( LockedChan
  , cloneLChan
  , newLockedChan
  , peekLChan
  , unlockChan 
  , writeLChan )
 where


import Control.Monad.STM
import Control.Concurrent.STM

data LockedChan a = LockedChan
    { lock :: TVar Bool
    , chan :: TChan a
    }

unlockChan :: LockedChan a -> IO (TChan a)
unlockChan c = do
    waslocked <- atomically $ swapTVar (lock c) False
    if waslocked
        then return (chan c) 
        else error "Attempt to read unlocked channel"

writeLChan :: LockedChan a -> a -> STM ()
writeLChan c a = writeTChan (chan c) a

peekLChan :: LockedChan a -> STM a
peekLChan c = peekTChan (chan c)

newLockedChan :: STM (LockedChan a)
newLockedChan = do
    lock <- newTVar True
    chan <- newTChan
    return $ LockedChan lock chan

cloneLChan :: LockedChan a -> IO (LockedChan a)
cloneLChan c = do
    mchan <- atomically $ do
        locked <- readTVar (lock c)
        if locked
            then fmap Just $ do
                    c2 <- cloneTChan (chan c)
                    l2 <- newTVar True
                    return $ LockedChan l2 c2
            else return Nothing
    return $ maybe (error "Attempt to clone unlocked channel")
                   id
                   mchan

#if MIN_VERSION_stm(2,4,0)
#else
-- |Clone a 'TChan': similar to dupTChan, but the cloned channel starts with the
-- same content available as the original channel.
--
-- Terrible inefficient implementation provided to build against older libraries.
cloneTChan :: TChan a -> STM (TChan a)
cloneTChan chan = do
    contents <- chanContents' chan
    chan2 <- dupTChan chan
    mapM_ (writeTChan chan) contents
    return chan2
 where
    chanContents' chan = do
        b <- isEmptyTChan chan
        if b then return [] else do
        x <- readTChan chan
        xs <- chanContents' chan
        return (x:xs)
#endif