Control.Monad.Trans.Dijkstra

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE UndecidableInstances #-}

module Control.Monad.Trans.Dijkstra
( Dijkstra , dijkstra , runDijkstra
, fromState , toState
, evalDijkstra , execDijkstra , mapDijkstra , withDijkstra
, DijkstraT(..)
, fromStateT , toStateT
, evalDijkstraT , execDijkstraT , mapDijkstraT , withDijkstraT
) where

import Control.Applicative (Applicative(pure, (<*>)))
import Control.Monad (Monad(return, (>>=)), (>=>), liftM)
import Control.Monad.Reader (MonadReader(ask, local, reader))
import Control.Monad.RWS (MonadRWS)
import Control.Monad.State (MonadState(get, put, state), StateT(..), State(..))
import Control.Monad.Trans (MonadTrans(lift))
import Control.Monad.Writer (MonadWriter(listen, pass, tell, writer))
import Data.Function ((.))
import Data.Functor (Functor(fmap))
import Data.Functor.Identity (Identity(..))
import Data.Tuple (fst, snd)

type Dijkstra s = DijkstraT s Identity

dijkstra :: (forall r. ((a, s) -> r) -> s -> r) -> Dijkstra s a
dijkstra f = DijkstraT (\asr -> runIdentity . f (Identity . asr))

runDijkstra :: Dijkstra s a -> forall r. ((a, s) -> r) -> s -> r
runDijkstra m asr = runIdentity . runDijkstraT m (Identity . asr)

fromState :: State s a -> Dijkstra s a
fromState = fromStateT

toState :: Dijkstra s a -> State s a
toState = toStateT

evalDijkstra :: Dijkstra s a -> s -> a
evalDijkstra m = runIdentity . evalDijkstraT m

execDijkstra :: Dijkstra s a -> s -> s
execDijkstra m = runIdentity . execDijkstraT m

mapDijkstra :: ((a, s) -> (b, s)) -> Dijkstra s a -> Dijkstra s b
mapDijkstra = mapDijkstraT . liftM

withDijkstra :: (s -> s) -> Dijkstra s a -> Dijkstra s a
withDijkstra = withDijkstraT

newtype DijkstraT s m a = DijkstraT { runDijkstraT :: forall r. ((a, s) -> m r) -> s -> m r }

fromStateT :: (Monad m) => StateT s m a -> DijkstraT s m a
fromStateT m = DijkstraT (runStateT m >=>)

toStateT :: (Monad m) => DijkstraT s m a -> StateT s m a
toStateT m = StateT (runDijkstraT m return)

evalDijkstraT :: (Monad m) => DijkstraT s m a -> s -> m a
-- evalDijkstraT = evalStateT . toStateT
-- evalDijkstraT m = evalStateT (toStateT m)
-- evalDijkstraT m = evalStateT (StateT (runDijkstraT m return))
-- evalDijkstraT m s = liftM fst (runStateT (StateT (runDijkstraT m return)) s)
-- evalDijkstraT m s = liftM fst (runDijkstraT m return s)
evalDijkstraT m = runDijkstraT m (return . fst)

execDijkstraT :: (Monad m) => DijkstraT s m a -> s -> m s
-- execDijkstraT = execStateT . toStateT
-- execDijkstraT m = execStateT (toStateT m)
-- execDijkstraT m = execStateT (StateT (runDijkstraT m return))
-- execDijkstraT m s = liftM snd (runStateT (StateT (runDijkstraT m return)) s)
-- execDijkstraT m s = liftM snd (runDijkstraT m return s)
execDijkstraT m = runDijkstraT m (return . snd)

mapDijkstraT :: (Monad m, Monad n) => (m (a, s) -> n (b, s)) -> DijkstraT s m a -> DijkstraT s n b
-- mapDijkstraT f = fromStateT . mapStateT f . toStateT
-- mapDijkstraT f m = fromStateT (mapStateT f (toStateT m))
-- mapDijkstraT f m = DijkstraT (runStateT (mapStateT f (StateT (runDijkstraT m return))) >=>)
-- mapDijkstraT f m = DijkstraT (runStateT (StateT (f . runStateT (StateT (runDijkstraT m return)))) >=>)
mapDijkstraT f m = DijkstraT ((f . runDijkstraT m return) >=>)

withDijkstraT :: (Monad m) => (s -> s) -> DijkstraT s m a -> DijkstraT s m a
-- withDijkstraT f = fromStateT . withStateT f . toStateT
-- withDijkstraT f m = fromStateT (withStateT f (toStateT m))
-- withDijkstraT f m = DijkstraT (runStateT (withStateT f (StateT (runDijkstraT m return))) >=>)
-- withDijkstraT f m = DijkstraT (runStateT (StateT (runStateT (StateT (runDijkstraT m return)) . f)) >=>)
withDijkstraT f m = DijkstraT ((runDijkstraT m return . f) >=>)

instance Functor (DijkstraT s m) where
        fmap ab m = DijkstraT (\bsmr -> runDijkstraT m (\ ~(a, s) -> bsmr (ab a, s)))

instance Applicative (DijkstraT s m) where
        pure a = DijkstraT (\asmr s -> asmr (a, s))
        mab <*> ma = DijkstraT (\bsmr -> runDijkstraT mab (\ ~(ab, s) -> runDijkstraT ma (\ ~(a, s') -> bsmr (ab a, s')) s))

instance Monad (DijkstraT s m) where
        return a = DijkstraT (\asmr s -> asmr (a, s))
        m >>= k = DijkstraT (\bsmr -> runDijkstraT m (\ ~(a, s) -> runDijkstraT (k a) bsmr s))

instance MonadTrans (DijkstraT s) where
        lift ma = DijkstraT (\asmr s -> ma >>= \a -> asmr (a, s))

instance (MonadReader r m) => MonadReader r (DijkstraT s m) where
        ask = lift ask
        local = mapDijkstraT . local
        reader = lift . reader

instance (MonadWriter w m) => MonadWriter w (DijkstraT s m) where
        listen = mapDijkstraT (liftM swap_ . listen)
        pass = mapDijkstraT (pass . liftM swap_)
        tell = lift . tell
        writer = lift . writer

instance MonadState s (DijkstraT s m) where
        get = DijkstraT (\ssmr s -> ssmr (s, s))
        put s = DijkstraT (\xsmr _ -> xsmr ((), s))
        state sas = DijkstraT (. sas)

instance (MonadReader r m, MonadWriter w m) => MonadRWS r w s (DijkstraT s m)

swap_ :: ((a, b), c) -> ((a, c), b)
swap_   ~((a, b), c) =  ((a, c), b)
{-# INLINABLE swap_ #-}