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

-- | A monad that cleanly generalizes out implementation details of array manipulation in an array transformer.  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.ArrayM (ArrayM, runArrayM, runArrayM_) 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 transformer that safely grants the underlying monad access to a mutable array.
newtype ArrayM s e a = ArrayM {runArrM :: StateT (MArr s e) (ST s) a} deriving (MonadST s, Monad, MonadFix)

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

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

emptyElement = error "Undefined array element"

instance MonadArray e (ArrayM s e) where
	{-# INLINE unsafeReadAt #-}
	{-# INLINE unsafeWriteAt #-}
	{-# INLINE getSize #-}
	{-# INLINE resize #-}
	unsafeReadAt i = ArrayM $ 	do	arr <- get
						lift $ readMArr arr i
	unsafeWriteAt i x = ArrayM $ 	do	arr <- get
						lift $ writeMArr arr i x
	getSize = ArrayM $ 	do	MArr n _ _ <- get
					return n
	resize n' = ArrayM $ 	do	a@(MArr n d _) <- get
					a' <- lift $ newMArr n' d
					lift $ 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, () #)