ArrayT.hs

{-# LANGUAGE GeneralizedNewtypeDeriving, UnboxedTuples, MagicHash, OverlappingInstances, RankNTypes, FlexibleInstances, MultiParamTypeClasses, UndecidableInstances #-}

-- | A monad that cleanly generalizes out implementation details of array manipulation in a monad.  In general, this is likely to be the most efficient array transformer implementation made available in this library, but if improperly used, elements of this implementation may lead to segfaults.
module Control.Monad.Array.ArrayT (ArrayM, ArrayT, runArrayM, runArrayMIO, runArrayM_, runArrayMIO_, runArrayT, runArrayT_) where

import GHC.Exts
import GHC.ST(ST(..))

import Prelude hiding (getContents)
import Control.Monad.ST
import Control.Monad.Fix
import Control.Monad.Array.Class
import Control.Monad
import Control.Monad.Trans
import Control.Monad.RWS.Class
import Control.Monad.State
import Control.Monad.ST.Class

data MArr s e = MArr {-# UNPACK #-} !Int e (MutableArray# s e)

-- | Monad controlling safe access to an underlying array.
type ArrayM s e = ArrayT s e (ST s)
-- | Monad transformer that safely grants the underlying monad access to a mutable array.
newtype ArrayT s e m a = ArrayT {runArrT :: StateT (MArr s e) m a} deriving (Monad, MonadTrans, MonadFix, MonadST s, MonadReader r, MonadWriter w)

instance MonadState t m => MonadState t (ArrayT s e m) where
	get = lift get
	put = lift . put

runArrayM :: Int -> e -> (forall s . ArrayM s e a) -> a
runArrayM n d m = runST $ runArrayT n d m

runArrayMIO :: Int -> e -> ArrayM RealWorld e a -> IO a
runArrayMIO n d m = stToIO $ runArrayT n d m

runArrayM_ :: Int -> (forall s . ArrayM s e a) -> a
runArrayM_ n = runArrayM n emptyElement

runArrayMIO_ :: Int -> ArrayM RealWorld e a -> IO a
runArrayMIO_ n = runArrayMIO n emptyElement

runArrayT :: (MonadST s m, Monad m) => Int -> e -> ArrayT s e m a -> m a
runArrayT n d m = liftST (newMArr n d) >>= evalStateT (runArrT m)

runArrayT_ :: (MonadST s m, Monad m) => Int -> ArrayT s e m a -> m a
runArrayT_ n = runArrayT n emptyElement

emptyElement = error "Undefined array element"

instance (MonadST s m, Monad m) => MonadArray e (ArrayT s e m) where
	{-# INLINE unsafeReadAt #-}
	{-# INLINE unsafeWriteAt #-}
	{-# INLINE getSize #-}
	{-# INLINE resize #-}
	unsafeReadAt i = ArrayT $ 	do	arr <- get
						liftST $ readMArr arr i
	unsafeWriteAt i x = ArrayT $ 	do	arr <- get
						liftST $ writeMArr arr i x
	getSize = ArrayT $ 	do	MArr n _ _ <- get
					return n
	resize n' = ArrayT $ 	do	a@(MArr n d _) <- get
					a' <- liftST $ newMArr n' d
					liftST $ mapM_ (\ i -> readMArr a i >>= writeMArr a' i) [0..n-1]
					put a'

newMArr :: Int -> e -> ST s (MArr s e)
newMArr (I# n) d = ST $ \ s -> case newArray# n d s of (# s', arr' #) -> (# s', MArr (I# n) d arr' #)

readMArr :: MArr s e -> Int -> ST s e
readMArr (MArr n _ arr) i@(I# i#) = ST $ readArray# arr i#

writeMArr :: MArr s e -> Int -> e -> ST s ()
writeMArr (MArr n _ arr) i@(I# i#) x = ST $ \ s -> (# writeArray# arr i# x s, () #)