{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, GeneralizedNewtypeDeriving #-}

-- | Provides a 'MonadArray' implementation for any 'MArray'.  Examples of when this would be useful include unboxed arrays and array implementations for specialized monads like STM.
module Control.Monad.Array.MArray (MArrayM, liftMArray, evalMArrayM, execMArrayM, evalMArrayM_, execMArrayM_) where

import Data.Array.MArray
import GHC.Arr
import Data.Array.Base
import Prelude hiding (getContents)
import Control.Monad.Array.Class
import Control.Monad
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Trans
import Control.Monad.ST.Class

-- | Provides a monadic wrapper around any 'MArray' implementation.
newtype MArrayM a e m x = MArrayM {runMArrayM :: ReaderT e (StateT (a Int e) m) x} deriving (Monad, MonadFix, MonadPlus, MonadIO, MonadST s)

-- | Executes an 'MArrayM' computation with the specified initial size and default element.
evalMArrayM :: (Monad m, MArray a e m) => Int -> e -> MArrayM a e m x -> m x
evalMArrayM n d m = newArray (0, n-1) d >>= evalStateT (runReaderT (runMArrayM m) d)

-- | Executes an 'MArrayM' computation with the specified initial size and default element, returning the final array.
execMArrayM :: (Monad m, MArray a e m) => Int -> e -> MArrayM a e m x -> m (a Int e)
execMArrayM n d m = newArray (0, n-1) d >>= execStateT (runReaderT (runMArrayM m) d)

-- | Executes an 'MArrayM' computation with the specified initial size and no default element.
evalMArrayM_ :: (Monad m, MArray a e m) => Int -> MArrayM a e m x -> m x
evalMArrayM_ n = evalMArrayM n emptyElement

-- | Executes an 'MArrayM' computation with the specified initial size and no default element, returning the final array.
execMArrayM_ :: (Monad m, MArray a e m) => Int -> MArrayM a e m x -> m (a Int e)
execMArrayM_ n = execMArrayM n emptyElement

emptyElement :: e
emptyElement = error "Undefined array element"

instance (Monad m, MArray a e m) => MonadArray e (MArrayM a e m) where
	unsafeWriteAt i x = MArrayM $ do	arr <- get
						lift2 $ unsafeWrite arr i x
	unsafeReadAt i = MArrayM $ do	arr <- get
					lift2 $ unsafeRead arr i
	getSize = MArrayM (get >>= lift2 . liftM rangeSize . getBounds)
	resize n = do	prevSize <- getSize
			prevConts <- mapM unsafeReadAt [0..prevSize-1]
			def <- MArrayM ask
			arr' <- liftMArray $ newListArray (0, n-1) (prevConts ++ replicate (n - prevSize) def)
			MArrayM $ put arr'

-- | Lifts a computation in the underlying monad to an 'MArrayM' computation on an array in the same monad.
liftMArray :: (Monad m, MArray a e m) => m x -> MArrayM a e m x
liftMArray = MArrayM . lift2

lift2 :: (MonadTrans t1, MonadTrans t2, Monad m, Monad (t2 m)) => m a -> t1 (t2 m) a
lift2 = lift . lift