Loop monad for Haskell

Loop.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

module Loop where

import Prelude hiding (break)

import Control.Monad (ap)
import Control.Monad.State.Lazy (MonadState(..))
import Control.Monad.Trans (MonadTrans(..))

data Loop a r = Break    a
              | Continue a
              | Value    r

instance Functor (Loop a) where
    fmap :: (r -> s) -> Loop a r -> Loop a s
    fmap _ (Break    t) = Break t
    fmap _ (Continue t) = Continue t
    fmap f (Value    x) = Value (f x)

newtype LoopT a m r = LoopT { runLoopT :: m (Loop a r) }

instance Functor m => Functor (LoopT a m) where
    fmap :: (r -> s) -> LoopT a m r -> LoopT a m s
    fmap f (LoopT mx) = LoopT (fmap (fmap f) mx)

instance Monad m => Applicative (LoopT a m) where
    pure :: r -> LoopT a m r
    pure = LoopT . pure . Value

    (<*>) :: LoopT a m (r -> s) -> LoopT a m r -> LoopT a m s
    (<*>) = ap

instance Monad m => Monad (LoopT a m) where
    (>>=) :: LoopT a m r -> (r -> LoopT a m s) -> LoopT a m s
    (LoopT mx) >>= f = LoopT $ do x <- mx
                                  case x of
                                       Break t    -> return (Break t)
                                       Continue t -> return (Continue t)
                                       Value y    -> runLoopT (f y)

instance MonadTrans (LoopT a) where
    lift :: Monad m => m r -> LoopT t m r
    lift = LoopT . fmap Value

instance MonadState s m => MonadState s (LoopT a m) where
    get :: LoopT a m s
    get = lift get

    put :: s -> LoopT a m ()
    put = lift . put

break :: Monad m => a -> LoopT a m к
break = LoopT . return . Break

continue :: Monad m => a -> LoopT a m r
continue = LoopT . return . Continue

inCycle :: Functor m => LoopT a m a -> LoopT a m a
inCycle (LoopT mx) = LoopT (fmap go mx)
    where
        go :: Loop a a -> Loop a a
        go (Break    x) = Break x
        go (Continue x) = Value x
        go (Value x)    = Value x

afterCycle :: Functor m => LoopT a m a -> m a
afterCycle (LoopT mx) = fmap go mx
    where
        go :: Loop a a -> a
        go (Break    x) = x
        go (Continue x) = x
        go (Value    x) = x

Usage_DfsHopcroft.hs

{-# INLINE dfsEdge #-}
dfsEdge :: Int -> a -> b -> LoopT Bool (State (HKState a b (Set.Set a))) Bool
dfsEdge dst v u = do (HKS d m s) <- get
                     case u `Map.lookup` pairOfRight m of
                          Nothing -> do modify (addEdgeTo v u)
                                        break True
                          Just w  -> do when (w `Set.member` s) (continue False)
                                        let dw = fromJust (w `Map.lookup` d)
                                        when (dw /= dst + 1) (continue False)
                                        z <- lift (dfsVertex w)
                                        when (not z) (continue False)
                                        modify (addEdgeTo v u)
                                        break True

dfsVertex :: a -> HKDfsMonad a b Bool
dfsVertex v = do (HKS d m s) <- get
                 let dist = fromJust (v `Map.lookup` d)
                 put (HKS d m (Set.insert v s))
                 afterCycle $ or <$> mapM (inCycle . dfsEdge dist v) (neighbours v)

dfs :: HKDfsMonad a b ()
dfs = do (HKS _ m _) <- get
         mapM_ dfsVertex [ v | v <- leftVertexList g, not (leftCovered v m) ]