{-# LANGUAGE GeneralizedNewtypeDeriving, TypeFamilies, 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, evalIntMapT, execIntMapT, execIntMapT_, evalIntMapT_, IntMapM, execIntMapM, evalIntMapM,  execIntMapM_, evalIntMapM_) where

import Control.Monad.Array.Class

import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Writer.Class
import Control.Monad.Trans

import qualified Data.IntMap as IM
import Data.IntMap(IntMap)

import Control.Monad

-- | 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 'askSize' operation returns the number of associations in the 'IntMap'.
newtype IntMapT e m a = IntMapT {runIMapT :: ReaderT e (StateT (IntMap e) m) a} deriving (Monad, MonadFix, MonadPlus, MonadIO, MonadWriter w)

-- | Basic monad version of 'IntMapT'.
newtype IntMapM e a = IntMapM {runIMapM :: ReaderT e (State (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 (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 -> 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 (IntMap e)
execIntMapT_ = execIntMapT emptyElement

evalIntMapM_ :: IntMapM e a -> a
evalIntMapM_ = evalIntMapM emptyElement

execIntMapM_ :: IntMapM e a -> IntMap e
execIntMapM_ = execIntMapM emptyElement

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

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

instance Monad m => MonadArray (IntMapT e m) where
	type ArrayElem (IntMapT e m) = e
	readAt i = IntMapT $ gets (IM.lookup i) >>= maybe ask return
	writeAt i x = IntMapT $ modify (IM.insert i x)
	askSize = IntMapT $ gets IM.size
	ensureSize _ = return ()
	askElems = IntMapT $ gets IM.elems
	askAssocs = IntMapT $ gets IM.toList

instance MonadArray (IntMapM e) where
	type ArrayElem (IntMapM e) = e
	readAt i = IntMapM $ gets (IM.lookup i) >>= maybe ask return
	writeAt i x = IntMapM $ modify (IM.insert i x)
	askSize = IntMapM $ gets IM.size
	ensureSize _ = return ()
	askElems = IntMapM $ gets IM.elems
	askAssocs = IntMapM $ gets IM.toList

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 =<<)