{-#LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}
module Singh.STMJoin where

import Control.Concurrent.STM
import Control.Concurrent

infixl 5 ?
infixl 5 ??
infixl 6 &
infix 2 >>>

class Joinable t1 t2 where
    (&) :: t1 a -> t2 b -> STM (a,b) 

instance Joinable TChan TChan where
    (&) chanA chanB = chanA & (readTChan chanB)
instance Joinable TChan STM where
    (&) chanA stm = do 
        a <- readTChan chanA
        b <- stm
        return (a,b)
instance Joinable STM TChan where
    stm & chan = do
        (y,x) <- chan & stm
        return (x,y)

joinList :: [TChan a] -> STM [a]
joinList = mapM readTChan

class Handle t where
    (>>>) :: t a -> (a -> IO b) -> STM (IO b)
instance Handle STM where
    joinPattern >>> handler = do
        a <- joinPattern
        return (handler a)
instance Handle TChan where
    chan >>> handler = (readTChan chan) >>> handler

spawn :: TChan a -> a -> IO ()
spawn chan val = atomically (writeTChan chan val)

call chan val = do
    response <- newTChanIO
    spawn chan (response,val)
    atomically $ readTChan response

matchSync lst = do
    io <- atomically $ biasedChoice lst
    io

match lst = do
    forkIO (asyncJoinLoop lst)
    return () -- discard thread id
    where
    asyncJoinLoop lst = do
        io <- atomically $ biasedChoice lst
        forkIO io
        asyncJoinLoop lst

biasedChoice (pattern:xs) = 
    pattern `orElse` (biasedChoice xs)
biasedChoice [] = retry

class Guard m g where
    (?) :: m a -> g -> STM a
instance Guard TChan Bool where
    chan ? guard = (readTChan chan) ? guard
instance Guard TChan (STM Bool) where
    chan ? stmGuard = do
        guard <- stmGuard
        chan ? guard
instance Guard STM (STM Bool) where
    stm ? stmGuard = do
        guard <- stmGuard
        stm ? guard
instance Guard STM Bool where
    stm ? guard = if guard 
        then stm
        else retry
class GuardFun m g where
    (??) :: m a -> (a -> g) -> STM a
instance GuardFun TChan Bool where
    chan ?? guardFunc = (readTChan chan) ?? guardFunc
instance GuardFun STM Bool where
    stm ?? guardFunc = do
        value <- stm
        if guardFunc value
            then return value
            else retry

nonBlocking :: STM a -> STM (Maybe a)
nonBlocking joinPattern = (do result <- joinPattern
                              return $ Just result)
                          `orElse`
                          (return Nothing)