{-# OPTIONS_GHC -fglasgow-exts #-}
module MyST (ST, runST
            ,STRef, newSTRef, readSTRef, writeSTRef
            ,STArray
            ,module Data.Array.MArray
            ) where

import Monad
import Data.IORef
import Data.Array.Base (unsafeRead, unsafeWrite)
import Data.Array.MArray
import Data.Array.IO (IOArray)
import System.IO.Unsafe (unsafePerformIO)

--- ST ---

newtype ST s a = ST {unST :: IO a}

runST :: (forall s. ST s a) -> a
runST = unsafePerformIO . unST

instance Monad (ST s) where
    return = ST . return
    ST a >>= f = ST (a >>= unST . f)

--- STRef ---

newtype STRef s a = STRef {unSTRef :: IORef a}

newSTRef :: a -> ST s (STRef s a)
newSTRef = ST . liftM STRef . newIORef

readSTRef :: STRef s a -> ST s a
readSTRef = ST . readIORef . unSTRef

writeSTRef :: STRef s a -> a -> ST s ()
writeSTRef r = ST . writeIORef (unSTRef r)

--- STArray ---

newtype STArray s i e = STArray {unSTArray :: IOArray i e}

instance HasBounds (STArray s) where
    bounds = bounds . unSTArray

instance MArray (STArray s) e (ST s) where
    newArray b = ST . liftM STArray . newArray b
    unsafeRead a = ST . unsafeRead (unSTArray a)
    unsafeWrite a i = ST . unsafeWrite (unSTArray a) i