{-# LANGUAGE BangPatterns, ScopedTypeVariables, GADTs #-}
--
-- Copyright (c) 2009 Alex Mason - http://axman6.homeip.net/blog/
-- BSD licence - http://www.opensource.org/licenses/bsd-license.php
--

-- |AVars are a form of transactional variables. They internally use a tail
-- recursive function to carry the 'state' of the variable, and allow for
-- use in concurrent systems, where actions are quaranteed to happen. They
-- are designed to cope with exceptions thrown by any modifying functions.
module Data.AVar (
    -- * Types
    AVar,
    Transaction(..),
    -- * functions on AVars
    newAVar,
    putMVar,
    modAVar,
    modAVar',
    getAVar,
    condModAVar,
    swapAVar) where
import Control.Concurrent
import Control.Concurrent.MVar
import Control.Concurrent.Chan
import qualified Control.Exception as E
import System

-- * Types

-- |A 'Transaction' describes what should happen to a variable.
-- They are only used internally, and are here just for reference.
data Transaction a =
      Put a
    -- ^puts the a into the variable
    | Get (MVar a)
    -- ^reads the variable
    | Mod (a -> a) (MVar (Maybe E.SomeException))
    -- ^modifies the variable
    | forall b. Mod' (a -> (a,b)) (MVar (Either E.SomeException b))
    -- ^ modifies the variable, returning the b result to the caller
    | Atom (a -> Bool) (a -> a) (a -> a) (MVar (Either E.SomeException Bool))
    -- ^conditionally modifies a variable

    --Swap a (MVar a)

-- |'AVar's are the means through communication with the variable are conducted.
-- They contain a Chan that is 'connected' to the variable, and is read by the
-- variable's 'handler' function.
data AVar a = AVar (Chan (Transaction a)) 

-- * Functions on AVars

-- |'newAVar' creates a new variable. It forks off the 'handler' that does the
-- work for the variable itself and creates a new AVar.
newAVar :: a -> IO (AVar a)
newAVar x = do
    E.evaluate x
    chan <- newChan :: IO (Chan (Transaction a))
    forkIO (handler chan x)
    return (AVar chan)

-- |'handler' stores the state of the variable in an accumulatin parameter.
-- It reads the chan it was forked with, and takes action depending on the
-- Transaction is recieves. Handler is not something to be used outside of
-- an AVar, and is not exported.
handler :: Chan (Transaction a) -> a -> IO b 
handler chan !x = do
  req <- readChan chan
  case req of
    Put a         -> handler chan a
    Get mvar      -> do
        putMVar mvar x
        handler chan x
    
    Mod f mvar    -> do
        let x' = f x
        p <- E.catch (E.evaluate x' >> return Nothing)
                     (\e -> return (Just e))
        putMVar mvar p
        case p of
          Nothing -> handler chan x'
          _       -> handler chan x
      
    Mod' f mvar    -> do
          let y@(a,b) = f x
          p <- E.try (E.evaluate a >> E.evaluate b)
          case p of
              Right _  -> do
                  putMVar mvar (Right b)
                  handler chan a
              (Left e) -> do
                  putMVar mvar (Left e)
                  handler chan x
    
    Atom test y n res -> do
        let t' = test x
            y' = y x
            n' = n x
        tres <- E.try (E.evaluate t')
        case tres of
            rT@(Right True) -> do
                run <- E.try (E.evaluate y')
                case run of
                    Right x' -> putMVar res rT >> handler chan x'
                    Left e   -> putMVar res (Left e) >> handler chan x
            rF@(Right False) -> do
                run <- E.try (E.evaluate n')
                case run of 
                    Right x' -> putMVar res rF >> handler chan x'
                    Left e   -> putMVar res (Left e) >> handler chan x
            Left e           -> putMVar res (Left e) >> handler chan x
        -- if test x
        --     then do
        --         putMVar res True
        --         handler chan (y x)
        --       else do
        --           putMVar res False
        --           handler chan (n x)

-- |'putAVar' replaces the currect value in the variable with the given x
putAVar :: AVar a -> a -> IO ()
putAVar (AVar chan) x = writeChan chan (Put x)

-- |'modAVar' takes a function from a to a, and returns Nothing if nothing went
-- wrong, or Just e, where e is an exception thrown by the function.
modAVar :: AVar a -> (a -> a) -> IO (Maybe E.SomeException)
modAVar (AVar chan) f = do
    res <- newEmptyMVar
    writeChan chan (Mod f res)
    takeMVar res

-- |'modAVar'' is like modAVar, but it modifies the variable, along with
-- returning a result of type b, within an Either e b.
modAVar' :: AVar a -> (a -> (a,b)) -> IO (Either E.SomeException b)
modAVar' (AVar chan) f = do
    res <- newEmptyMVar
    writeChan chan (Mod' f res)
    takeMVar res

-- |'getAVar' reads the current value inside the AVar.
getAVar :: AVar a -> IO a
getAVar (AVar chan)   = do
    res <- newEmptyMVar
    writeChan chan (Get res)
    takeMVar res

-- |'condModAVar' applies the first finction to the current value in the
-- AVar, and if true will modify the value using the second function if
-- it results in True, or the third function if it results in Fasle.
condModAVar :: AVar a
            -> (a -> Bool)
            -> (a -> a)
            -> (a -> a)
            -> IO (Either E.SomeException Bool)
condModAVar (AVar chan) p t f = do
    res <- newEmptyMVar
    writeChan chan (Atom p t f res)
    takeMVar res

-- |'swapAVar' takes a new value, puts it into the AVar, and returns the old value.
swapAVar :: (AVar a) -> a -> IO (Either E.SomeException a)
swapAVar (AVar chan) new = do
    res <- newEmptyMVar
    writeChan chan (Mod' (\old -> (new, old)) res)
    takeMVar res

main = do
    n <- getArgs >>= \xs -> if null xs then return 1000000 else (readIO.head) xs
    var <- newAVar (0 :: Int)
    m <- newEmptyMVar
    forkIO $ test n var m
    takeMVar m
    where test 0 _ m   = putMVar m ()
          test n var m = do
              res <- getAVar var
              putAVar var (res + 1)
              test (n-1) var m