{-# LANGUAGE DeriveDataTypeable #-}
-- | An MChan is a combination of a Map and a TChan.
module Network.Hermes.MChan(
  MChan ,newMChan, newMChanIO, readMChan, writeMChan, writeMChan'
  ,existsMChan, ensureMChan, deleteMChan
  ) where

import Prelude hiding(lookup)

import Control.Concurrent.STM
import Control.Applicative((<$>))
import Data.Typeable
import Data.Map

import Network.Hermes.Misc(modifyTVar)

-- | MChan is an abstract type representing a keyed, unbounded FIFO channel
newtype MChan k v = MChan (TVar (Map k (TChan v)))
                      deriving(Typeable)

-- | Builds and returns a new instance of MChan
newMChan :: Ord k => STM (MChan k v)
newMChan = MChan <$> newTVar empty

-- | IO version of 'newMChan'. This is useful for creating top-level
-- 'MChan's using 'System.IO.Unsafe.unsafePerformIO', because using
-- 'atomically' inside 'System.IO.Unsafe.unsafePerformIO' isn't safe.
newMChanIO :: Ord k => IO (MChan k v)
newMChanIO = MChan <$> newTVarIO empty

-- | Read the next value from an MChan. If the specified key doesn't
-- exist (or is removed while waiting), it returns Nothing.
readMChan :: Ord k => MChan k v -> k -> STM (Maybe v)
readMChan (MChan var) k = do
  chan <- lookup k <$> readTVar var
  case chan of
    Nothing -> return Nothing
    Just chan' -> Just <$> readTChan chan'

-- | Write a value to an MChan. Returns false and discards the value
-- if the specified key doesn't exist.
writeMChan :: Ord k => MChan k v -> k -> v -> STM Bool
writeMChan (MChan var) k v = do
  chan <- lookup k <$> readTVar var
  case chan of
    Nothing -> return False
    Just chan' -> writeTChan chan' v >> return True

-- | Write a value to an MChan, creating the key if it doesn't exist.
writeMChan' :: Ord k => MChan k v -> k -> v -> STM ()
writeMChan' mchan@(MChan var) k v = do
  chan <- lookup k <$> readTVar var
  case chan of
    Nothing -> ensureMChan mchan k >> writeMChan' mchan k v
    Just chan' -> writeTChan chan' v

-- | Checks whether the key exists
existsMChan :: Ord k => MChan k v -> k -> STM Bool
existsMChan (MChan var) k = member k <$> readTVar var

-- | Creates the key if it doesn't already exist
ensureMChan :: Ord k => MChan k v -> k -> STM ()
ensureMChan (MChan var) k = do
  chan <- newTChan
  modifyTVar var (insertWith (const id) k chan)

-- | Delete a key from an MChan
deleteMChan :: Ord k => MChan k v -> k -> STM ()
deleteMChan (MChan var) k = modifyTVar var (delete k)