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_ #-}