{-# OPTIONS -fglasgow-exts -cpp #-}

module Control.Concurrent.STM.TArray (
    TArray
) where

import Control.Monad (replicateM)
import Control.Monad.ST (unsafeIOToST)
import Data.Array (Array, listArray)
import Data.Array.Base (arrEleBottom, unsafeAt, MArray(..), HasBounds(..))
import Data.Array.ST (runSTArray)
import Data.Ix (rangeSize)
import GHC.Conc (STM, TVar,
#if __GLASGOW_HASKELL__ >= 650
    newTVarIO,
#else
    newTVar,
#endif
    readTVar, writeTVar)

newtype TArray i e = TArray (Array i (TVar e))

instance MArray TArray e STM where
#if __GLASGOW_HASKELL__ >= 650
    newArray b e = return $ TArray $ runSTArray $ do
        a <- newArray_ b
        let n = rangeSize b
        let fill i | i == n = return ()
                   | otherwise = do
                       t <- unsafeIOToST (newTVarIO e)
                       unsafeWrite a i t
                       fill (i + 1)
        fill 0
        return a
#else
    newArray b e = do
        a <- replicateM (rangeSize b) (newTVar e)
        return $ TArray (listArray b a)
#endif
    newArray_ b = newArray b arrEleBottom
    unsafeRead (TArray a) i = readTVar $ unsafeAt a i
    unsafeWrite (TArray a) i e = writeTVar (unsafeAt a i) e

instance HasBounds TArray where
    bounds (TArray a) = bounds a