{-# 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