-- |
-- Module       : Control.Monad.Trans.Loop
-- Copyright    : (c) Joseph Adams 2012
-- License      : BSD3
-- Maintainer   : joeyadams3.14159@gmail.com
--

{-# LANGUAGE Rank2Types #-}

-- Needed for the MonadBase instance
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}

module Control.Monad.Trans.Loop (
    -- * The LoopT monad transformer
    LoopT(..),
    stepLoopT,

    -- * continue and exit
    continue,
    exit,
    continueWith,
    exitWith,

    -- * Looping constructs
    foreach,
    while,
    doWhile,
    once,
    repeatLoopT,
    iterateLoopT,

    -- * Lifting other operations
    liftLocalLoopT,
) where

import Control.Applicative          (Applicative(pure, (<*>)))
import Control.Monad.Base           (MonadBase(liftBase), liftBaseDefault)
import Control.Monad.IO.Class       (MonadIO(liftIO))
import Control.Monad.Trans.Class    (MonadTrans(lift))

-- | 'LoopT' is a monad transformer for the loop body.  It provides two
-- capabilities:
--
--  * 'continue' to the next iteration.
--
--  * 'exit' the whole loop.
newtype LoopT c e m a = LoopT
    { forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT :: forall r.     -- This universal quantification forces the
                                -- LoopT computation to call one of the
                                -- following continuations.
                  (c -> m r)    -- continue
               -> (e -> m r)    -- exit
               -> (a -> m r)    -- return a value
               -> m r
    }

instance Functor (LoopT c e m) where
    fmap :: forall a b. (a -> b) -> LoopT c e m a -> LoopT c e m b
fmap a -> b
f LoopT c e m a
m = (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
 -> LoopT c e m b)
-> (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall a b. (a -> b) -> a -> b
$ \c -> m r
next e -> m r
fin b -> m r
cont -> LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT c e m a
m c -> m r
next e -> m r
fin (b -> m r
cont (b -> m r) -> (a -> b) -> a -> m r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f)

instance Applicative (LoopT c e m) where
    pure :: forall a. a -> LoopT c e m a
pure a
a    = (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
 -> LoopT c e m a)
-> (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall a b. (a -> b) -> a -> b
$ \c -> m r
_    e -> m r
_   a -> m r
cont -> a -> m r
cont a
a
    LoopT c e m (a -> b)
f1 <*> :: forall a b. LoopT c e m (a -> b) -> LoopT c e m a -> LoopT c e m b
<*> LoopT c e m a
f2 = (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
 -> LoopT c e m b)
-> (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall a b. (a -> b) -> a -> b
$ \c -> m r
next e -> m r
fin b -> m r
cont ->
                LoopT c e m (a -> b)
-> forall r. (c -> m r) -> (e -> m r) -> ((a -> b) -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT c e m (a -> b)
f1 c -> m r
next e -> m r
fin (((a -> b) -> m r) -> m r) -> ((a -> b) -> m r) -> m r
forall a b. (a -> b) -> a -> b
$ \a -> b
f ->
                LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT c e m a
f2 c -> m r
next e -> m r
fin (b -> m r
cont (b -> m r) -> (a -> b) -> a -> m r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f)

instance Monad (LoopT c e m) where
    return :: forall a. a -> LoopT c e m a
return a
a = (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
 -> LoopT c e m a)
-> (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall a b. (a -> b) -> a -> b
$ \c -> m r
_    e -> m r
_   a -> m r
cont -> a -> m r
cont a
a
    LoopT c e m a
m >>= :: forall a b. LoopT c e m a -> (a -> LoopT c e m b) -> LoopT c e m b
>>= a -> LoopT c e m b
k  = (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
 -> LoopT c e m b)
-> (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall a b. (a -> b) -> a -> b
$ \c -> m r
next e -> m r
fin b -> m r
cont ->
               LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT c e m a
m c -> m r
next e -> m r
fin ((a -> m r) -> m r) -> (a -> m r) -> m r
forall a b. (a -> b) -> a -> b
$ \a
a ->
               LoopT c e m b
-> forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT (a -> LoopT c e m b
k a
a) c -> m r
next e -> m r
fin b -> m r
cont

instance MonadTrans (LoopT c e) where
    lift :: forall (m :: * -> *) a. Monad m => m a -> LoopT c e m a
lift m a
m = (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
 -> LoopT c e m a)
-> (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall a b. (a -> b) -> a -> b
$ \c -> m r
_ e -> m r
_ a -> m r
cont -> m a
m m a -> (a -> m r) -> m r
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m r
cont

instance MonadIO m => MonadIO (LoopT c e m) where
    liftIO :: forall a. IO a -> LoopT c e m a
liftIO = m a -> LoopT c e m a
forall (m :: * -> *) a. Monad m => m a -> LoopT c e m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> LoopT c e m a) -> (IO a -> m a) -> IO a -> LoopT c e m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

instance MonadBase b m => MonadBase b (LoopT c e m) where
    liftBase :: forall α. b α -> LoopT c e m α
liftBase = b α -> LoopT c e m α
forall (t :: (* -> *) -> * -> *) (b :: * -> *) (m :: * -> *) α.
(MonadTrans t, MonadBase b m) =>
b α -> t m α
liftBaseDefault

-- | Call a loop body, passing it a continuation for the next iteration.
-- This can be used to construct custom looping constructs.  For example,
-- here is the definition of 'foreach':
--
-- >foreach list body = loop list
-- >  where loop []     = return ()
-- >        loop (x:xs) = stepLoopT (body x) (\_ -> loop xs)
stepLoopT :: Monad m => LoopT c e m c -> (c -> m e) -> m e
stepLoopT :: forall (m :: * -> *) c e.
Monad m =>
LoopT c e m c -> (c -> m e) -> m e
stepLoopT LoopT c e m c
body c -> m e
next = LoopT c e m c
-> forall r. (c -> m r) -> (e -> m r) -> (c -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT c e m c
body c -> m e
next e -> m e
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return c -> m e
next

------------------------------------------------------------------------
-- continue and exit


-- | Skip the rest of the loop body and go to the next iteration.
continue :: LoopT () e m a
continue :: forall e (m :: * -> *) a. LoopT () e m a
continue = () -> LoopT () e m a
forall c e (m :: * -> *) a. c -> LoopT c e m a
continueWith ()

-- | Break out of the loop entirely.
exit :: LoopT c () m a
exit :: forall c (m :: * -> *) a. LoopT c () m a
exit = () -> LoopT c () m a
forall e c (m :: * -> *) a. e -> LoopT c e m a
exitWith ()

-- | Like 'continue', but return a value from the loop body.
continueWith :: c -> LoopT c e m a
continueWith :: forall c e (m :: * -> *) a. c -> LoopT c e m a
continueWith c
c = (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
 -> LoopT c e m a)
-> (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall a b. (a -> b) -> a -> b
$ \c -> m r
next e -> m r
_ a -> m r
_ -> c -> m r
next c
c

-- | Like 'exit', but return a value from the loop as a whole.
-- See the documentation of 'iterateLoopT' for an example.
exitWith :: e -> LoopT c e m a
exitWith :: forall e c (m :: * -> *) a. e -> LoopT c e m a
exitWith e
e = (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
 -> LoopT c e m a)
-> (forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
forall a b. (a -> b) -> a -> b
$ \c -> m r
_ e -> m r
fin a -> m r
_ -> e -> m r
fin e
e


------------------------------------------------------------------------
-- Looping constructs


-- | Call the loop body with each item in the list.
--
-- If you do not need to 'continue' or 'exit' the loop, consider using
-- 'Control.Monad.forM_' instead.
foreach :: Monad m => [a] -> (a -> LoopT c () m c) -> m ()
foreach :: forall (m :: * -> *) a c.
Monad m =>
[a] -> (a -> LoopT c () m c) -> m ()
foreach [a]
list a -> LoopT c () m c
body = [a] -> m ()
loop [a]
list
  where loop :: [a] -> m ()
loop []     = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        loop (a
x:[a]
xs) = LoopT c () m c -> (c -> m ()) -> m ()
forall (m :: * -> *) c e.
Monad m =>
LoopT c e m c -> (c -> m e) -> m e
stepLoopT (a -> LoopT c () m c
body a
x) (\c
_ -> [a] -> m ()
loop [a]
xs)

-- | Repeat the loop body while the predicate holds.  Like a @while@ loop in C,
-- the condition is tested first.
while :: Monad m => m Bool -> LoopT c () m c -> m ()
while :: forall (m :: * -> *) c. Monad m => m Bool -> LoopT c () m c -> m ()
while m Bool
cond LoopT c () m c
body = m ()
loop
  where loop :: m ()
loop = do Bool
b <- m Bool
cond
                  if Bool
b then LoopT c () m c -> (c -> m ()) -> m ()
forall (m :: * -> *) c e.
Monad m =>
LoopT c e m c -> (c -> m e) -> m e
stepLoopT LoopT c () m c
body (\c
_ -> m ()
loop)
                       else () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Like a @do while@ loop in C, where the condition is tested after
-- the loop body.
--
-- 'doWhile' returns the result of the last iteration.  This is possible
-- because, unlike 'foreach' and 'while', the loop body is guaranteed to be
-- executed at least once.
doWhile :: Monad m => LoopT a a m a -> m Bool -> m a
doWhile :: forall (m :: * -> *) a. Monad m => LoopT a a m a -> m Bool -> m a
doWhile LoopT a a m a
body m Bool
cond = m a
loop
  where loop :: m a
loop = LoopT a a m a -> (a -> m a) -> m a
forall (m :: * -> *) c e.
Monad m =>
LoopT c e m c -> (c -> m e) -> m e
stepLoopT LoopT a a m a
body ((a -> m a) -> m a) -> (a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \a
a -> do
            Bool
b <- m Bool
cond
            if Bool
b then m a
loop
                 else a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

-- | Execute the loop body once.  This is a convenient way to introduce early
-- exit support to a block of code.
--
-- 'continue' and 'exit' do the same thing inside of 'once'.
once :: Monad m => LoopT a a m a -> m a
once :: forall (m :: * -> *) a. Monad m => LoopT a a m a -> m a
once LoopT a a m a
body = LoopT a a m a
-> forall r. (a -> m r) -> (a -> m r) -> (a -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT a a m a
body a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return

-- | Execute the loop body again and again.  The only way to exit 'repeatLoopT'
-- is to call 'exit' or 'exitWith'.
repeatLoopT :: Monad m => LoopT c e m a -> m e
repeatLoopT :: forall (m :: * -> *) c e a. Monad m => LoopT c e m a -> m e
repeatLoopT LoopT c e m a
body = m e
loop
  where loop :: m e
loop = LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT c e m a
body (\c
_ -> m e
loop) e -> m e
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (\a
_ -> m e
loop)

-- | Call the loop body again and again, passing it the result of the previous
-- iteration each time around.  The only way to exit 'iterateLoopT' is to call
-- 'exit' or 'exitWith'.
--
-- Example:
--
-- >count :: Int -> IO Int
-- >count n = iterateLoopT 0 $ \i ->
-- >    if i < n
-- >        then do
-- >            lift $ print i
-- >            return $ i+1
-- >        else exitWith i
iterateLoopT :: Monad m => c -> (c -> LoopT c e m c) -> m e
iterateLoopT :: forall (m :: * -> *) c e.
Monad m =>
c -> (c -> LoopT c e m c) -> m e
iterateLoopT c
z c -> LoopT c e m c
body = c -> m e
loop c
z
  where loop :: c -> m e
loop c
c = LoopT c e m c -> (c -> m e) -> m e
forall (m :: * -> *) c e.
Monad m =>
LoopT c e m c -> (c -> m e) -> m e
stepLoopT (c -> LoopT c e m c
body c
c) c -> m e
loop


------------------------------------------------------------------------
-- Lifting other operations


-- | Lift a function like 'Control.Monad.Trans.Reader.local' or
-- 'Control.Exception.mask_'.
liftLocalLoopT :: Monad m => (forall a. m a -> m a) -> LoopT c e m b -> LoopT c e m b
liftLocalLoopT :: forall (m :: * -> *) c e b.
Monad m =>
(forall a. m a -> m a) -> LoopT c e m b -> LoopT c e m b
liftLocalLoopT forall a. m a -> m a
f LoopT c e m b
cb = (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall c e (m :: * -> *) a.
(forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r)
-> LoopT c e m a
LoopT ((forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
 -> LoopT c e m b)
-> (forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r)
-> LoopT c e m b
forall a b. (a -> b) -> a -> b
$ \c -> m r
next e -> m r
fin b -> m r
cont -> do
    m r
m <- m (m r) -> m (m r)
forall a. m a -> m a
f (m (m r) -> m (m r)) -> m (m r) -> m (m r)
forall a b. (a -> b) -> a -> b
$ LoopT c e m b
-> forall r. (c -> m r) -> (e -> m r) -> (b -> m r) -> m r
forall c e (m :: * -> *) a.
LoopT c e m a
-> forall r. (c -> m r) -> (e -> m r) -> (a -> m r) -> m r
runLoopT LoopT c e m b
cb (m r -> m (m r)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (m r -> m (m r)) -> (c -> m r) -> c -> m (m r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> m r
next) (m r -> m (m r)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (m r -> m (m r)) -> (e -> m r) -> e -> m (m r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> m r
fin) (m r -> m (m r)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (m r -> m (m r)) -> (b -> m r) -> b -> m (m r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> m r
cont)
    m r
m