{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module Control.Monad.Loop.Internal
( LoopT(..), Loop, loop
, Unroll(..), UnTL, Unrolling(), noUnroll
, cons, continue, continue_, break_, exec_
, iterate, forever, for, unfoldl, while
) where
import Control.Applicative (Applicative(..), (<$>), liftA2)
import Control.Category ((<<<), (>>>))
import Control.Monad (unless)
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import qualified GHC.TypeLits as TL
import Data.Foldable
import Data.Functor.Identity
import Data.Maybe (fromJust, isJust)
import Data.Traversable (Traversable(..))
import Prelude hiding (foldr, iterate)
-- | @LoopT m a@ represents a loop over a base type @m@ that yields a value
-- @a@ at each iteration. It can be used as a monad transformer, but there
-- are actually no restrictions on the type @m@. However, this library only
-- provides functions to execute the loop if @m@ is at least 'Applicative'
-- (for 'exec_'). If @m@ is also 'Foldable', so is @LoopT m@. For any other
-- type, you may use 'runLoopT'.
newtype LoopT m a = LoopT
{ runLoopT :: forall r. (a -> m r -> m r -> m r) -> m r -> m r -> m r }
-- | @Loop@ is a pure loop, without side-effects.
type Loop = LoopT Identity
-- | @loop@ is just an aid to type inference. For loops over a base monad,
-- there are usually other constraints that fix the type, but for pure
-- loops, the compiler often has trouble inferring @Identity@.
loop :: Loop a -> Loop a
{-# INLINE loop #-}
loop = id
instance Functor (LoopT m) where
{-# INLINE fmap #-}
fmap f xs = LoopT $ \yield -> runLoopT xs (yield . f)
instance Applicative (LoopT m) where
{-# INLINE pure #-}
pure a = LoopT $ \yield -> yield a
{-# INLINE (<*>) #-}
fs <*> as = LoopT $ \yield next ->
runLoopT fs (\f -> runLoopT (fmap f as) yield) next
instance Monad (LoopT m) where
{-# INLINE return #-}
return = pure
{-# INLINE (>>=) #-}
as >>= f = LoopT $ \yield next ->
runLoopT as (\a -> runLoopT (f a) yield) next
instance MonadTrans LoopT where
{-# INLINE lift #-}
lift m = LoopT $ \yield next brk -> m >>= \a -> yield a next brk
instance MonadIO m => MonadIO (LoopT m) where
{-# INLINE liftIO #-}
liftIO = lift . liftIO
instance (Applicative m, Foldable m) => Foldable (LoopT m) where
{-# INLINE foldr #-}
foldr f r xs = foldr (<<<) id inner r
where
yield a next _ = (f a <<<) <$> next
inner = runLoopT xs yield (pure id) (pure id)
{-# INLINE foldl' #-}
foldl' f r xs = foldl' (!>>>) id inner r
where
(!>>>) h g = h >>> (g $!)
yield a next _ = (flip f a >>>) <$> next
inner = runLoopT xs yield (pure id) (pure id)
instance (Applicative m, Foldable m) => Traversable (LoopT m) where
{-# INLINE sequenceA #-}
sequenceA = foldr (liftA2 cons) (pure continue_)
cons :: a -> LoopT m a -> LoopT m a
{-# INLINE cons #-}
cons a as = LoopT $ \yield next brk -> yield a (runLoopT as yield next brk) next
-- | Yield a value for this iteration of the loop and skip immediately to
-- the next iteration.
continue :: a -> LoopT m a
{-# INLINE continue #-}
continue a = LoopT $ \yield next -> yield a next
-- | Skip immediately to the next iteration of the loop without yielding
-- a value.
continue_ :: LoopT m a
{-# INLINE continue_ #-}
continue_ = LoopT $ \_ next _ -> next
-- | Skip all the remaining iterations of the immediately-enclosing loop.
break_ :: LoopT m a
{-# INLINE break_ #-}
break_ = LoopT $ \_ _ brk -> brk
-- | Execute a loop, sequencing the effects and discarding the values.
exec_ :: Applicative m => LoopT m a -> m ()
{-# INLINE exec_ #-}
exec_ xs = runLoopT xs (\_ next _ -> next) (pure ()) (pure ())
-- | Iterate forever (or until 'break' is used).
iterate
:: Unrolling (UnTL n)
=> Unroll n -- ^ Unrolling factor
-> a -- ^ Starting value of iterator
-> (a -> a) -- ^ Advance the iterator
-> LoopT m a
{-# INLINE iterate #-}
iterate unroll = \a0 adv -> LoopT $ \yield next _ ->
let go a = unrollIterate (fromTypeLit unroll) a adv yield go next
in go a0
-- | Loop forever without yielding (interesting) values.
forever :: Unrolling (UnTL n) => Unroll n -> LoopT m ()
{-# INLINE forever #-}
forever unroll = iterate unroll () id
-- | Standard @for@ loop.
for
:: Unrolling (UnTL n)
=> Unroll n -- ^ Unrolling factor
-> a -- ^ Starting value of iterator
-> (a -> Bool) -- ^ Termination condition. The loop will terminate the
-- first time this is false. The termination condition
-- is checked at the /start/ of each iteration.
-> (a -> a) -- ^ Advance the iterator
-> LoopT m a
{-# INLINE for #-}
for unroll = \a0 cond adv -> LoopT $ \yield next _ ->
let go a = unrollFor (fromTypeLit unroll) a cond adv yield go next
in if cond a0 then go a0 else next
-- | Unfold a loop from the left.
unfoldl
:: Unrolling (UnTL n)
=> Unroll n -- ^ Unrolling factor
-> (i -> Maybe (i, a)) -- ^ @Just (i, a)@ advances the loop, yielding an
-- @a@. @Nothing@ terminates the loop.
-> i -- ^ Starting value
-> LoopT m a
{-# INLINE unfoldl #-}
unfoldl unroll = \unf i0 ->
fromJust . fmap snd <$> for unroll (unf i0) isJust (>>= unf . fst)
while
:: (Unrolling (UnTL n), Monad m)
=> Unroll n
-> m Bool
-> LoopT m ()
{-# INLINE while #-}
while unroll = \cond -> do
forever unroll
p <- lift cond
unless p break_
-- | Proxy type for GHC's type level literal natural numbers. @n@ is the
-- number of times the loop will be unrolled into its own body.
data Unroll (n :: TL.Nat) = Unroll
data Nat = S !Nat | Z
data UnrollInd (n :: Nat) = UnrollInd
-- | Do not unroll the loop at all.
noUnroll :: Unroll 1
noUnroll = Unroll
predUnroll :: UnrollInd (S n) -> UnrollInd n
predUnroll UnrollInd = UnrollInd
type family UnTL (n :: TL.Nat) :: Nat where
UnTL 1 = S Z
UnTL n = S (UnTL ((TL.-) n 1))
fromTypeLit :: Unroll n -> UnrollInd (UnTL n)
fromTypeLit Unroll = UnrollInd
class Unrolling (n :: Nat) where
unrollFor
:: UnrollInd n
-> a -> (a -> Bool) -> (a -> a) -- for parameters
-> (a -> m r -> m r -> m r) -> (a -> m r) -> m r -> m r -- un-newtyped LoopT
unrollIterate
:: UnrollInd n -- unrolling factor
-> a -> (a -> a) -- iterate parameters
-> (a -> m r -> m r -> m r) -> (a -> m r) -> m r -> m r -- un-newtyped LoopT
instance Unrolling Z where
{-# INLINE unrollFor #-}
unrollFor UnrollInd a _ _ _ next _ = next a
{-# INLINE unrollIterate #-}
unrollIterate UnrollInd a _ _ next _ = next a
instance Unrolling n => Unrolling (S n) where
{-# INLINE unrollFor #-}
unrollFor unroll a cond adv yield next brk =
yield a descend brk
where
a' = adv a
descend | cond a' = unrollFor (predUnroll unroll) a' cond adv yield next brk
| otherwise = brk
{-# INLINE unrollIterate #-}
unrollIterate unroll a adv yield next brk =
yield a descend brk
where
descend = unrollIterate (predUnroll unroll) (adv a) adv yield next brk