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

-- | A module implementing the array abstraction on a purely functional IntMap.  When attempting to debug a complex array-using algorithm, it may sometimes be useful to use a less segfault-prone implementation.  In addition, the execXXX commands allow the final state of the 'IntMap' to be returned.
module Control.Monad.Array.IntMap (IntMapT, IntMapM, evalIntMapT, evalIntMapM, execIntMapT, execIntMapM, execIntMapT_, evalIntMapT_, execIntMapM_, evalIntMapM_) where

import qualified Data.IntMap as IM
import Data.IntMap(IntMap)
import Control.Monad
import Control.Monad.Trans
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Array.Class
import Control.Monad.ST.Class
import Control.Monad.Writer.Class

-- | An array transformer with an 'IntMap' on the back end.  Provides decent performance while retaining a purely functional back end.  /Note:/ resizing operations have no effect, and the 'getSize' operation returns the number of associations in the 'IntMap'.
newtype IntMapT e m a = IntMapT {runIMapT :: ReaderT e (StateT (IM.IntMap e) m) a} deriving (Monad, MonadFix, MonadPlus, MonadIO, MonadST s, MonadWriter w)

-- | Basic monad version of 'IntMapT'.
newtype IntMapM e a = IntMapM {runIMapM :: ReaderT e (State (IM.IntMap e)) a} deriving (Monad, MonadFix)

-- | Evaluates an 'IntMapT' computation with the specified default element.
evalIntMapT :: Monad m => e -> IntMapT e m a -> m a
evalIntMapT d m = evalStateT (runReaderT (runIMapT m) d) IM.empty

-- | Evaluates an 'IntMapT' computation with the specified default element, returning the final 'IntMap'.
execIntMapT :: Monad m => e -> IntMapT e m a -> m (IM.IntMap e)
execIntMapT d m = execStateT (runReaderT (runIMapT m) d) IM.empty

evalIntMapM :: e -> IntMapM e a -> a
evalIntMapM d m = evalState (runReaderT (runIMapM m) d) IM.empty

execIntMapM :: e -> IntMapM e a -> IM.IntMap e
execIntMapM d m = execState (runReaderT (runIMapM m) d) IM.empty

-- | Evaluates an 'IntMapT' computation with no default element specified.
evalIntMapT_ :: Monad m => IntMapT e m a -> m a
evalIntMapT_ = evalIntMapT emptyElement

-- | Evaluates an 'IntMapT' computation with no default element specified, returning the final 'IntMap'.
execIntMapT_ :: Monad m => IntMapT e m a -> m (IM.IntMap e)
execIntMapT_ = execIntMapT emptyElement
evalIntMapM_ = evalIntMapM emptyElement
execIntMapM_ = execIntMapM emptyElement

emptyElement = error "Undefined array element"

instance MonadTrans (IntMapT e) where
	lift = IntMapT . lift . lift

instance Monad m => MonadArray e (IntMapT e m) where
	unsafeReadAt i = IntMapT $ gets (IM.lookup i) >>= maybe ask return
	unsafeWriteAt i x = IntMapT $ modify (IM.insert i x)
	getSize = IntMapT $ gets IM.size
	ensureSize _ = return ()

instance MonadArray e (IntMapM e) where
	unsafeReadAt i = IntMapM $ gets (IM.lookup i) >>= maybe ask return
	unsafeWriteAt i x = IntMapM $ modify (IM.insert i x)
	getSize = IntMapM $ gets IM.size
	ensureSize _ = return ()

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

instance MonadReader r m => MonadReader r (IntMapT e m) where
	ask = lift ask
	local f = (lift . local f . return =<<)