{-# LANGUAGE BangPatterns #-}

-- | Provides a convenient and fast alternative to the common
-- @forM_ [1..n]@ idiom, which in many cases GHC cannot fuse to efficient
-- code.
--
-- Notes on fast iteration:
--
-- * For `Int`, @(+1)@ is almost twice as fast as `succ` because `succ`
--   does an overflow check.
--
-- * For `Int`, you can get around that while still using `Enum` using
--   @toEnum . (+ 1) . fromEnum@.
--
-- * However, @toEnum . (+ 1) . fromEnum@ is slower than `succ` for
--   `Word32` on 64-bit machines since `toEnum` has to check if the
--   given `Int` exceeds 32 bits.
--
-- * Using @(+1)@ from `Num` is always the fastest way, but it gives
--   no overflow checking.
--
-- * Using `forLoop` you can flexibly pick the way of increasing the value
--   that best fits your needs.
--
-- * The currently recommended replacement for @forM_ [1..n]@ is
--   @forLoop 1 (<= n) (+1)@.
module Control.Loop
  ( forLoop
  , forLoopState
  , forLoopFold
  , numLoop
  , numLoopState
  , numLoopFold
  ) where


-- | @forLoop start cond inc f@: A C-style for loop with starting value,
-- loop condition and incrementor.
forLoop :: (Monad m) => a -> (a -> Bool) -> (a -> a) -> (a -> m ()) -> m ()
forLoop :: a -> (a -> Bool) -> (a -> a) -> (a -> m ()) -> m ()
forLoop a
start a -> Bool
cond a -> a
inc a -> m ()
f = a -> m ()
go a
start
  where
    go :: a -> m ()
go !a
x | a -> Bool
cond a
x    = a -> m ()
f a
x m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> m ()
go (a -> a
inc a
x)
          | Bool
otherwise = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

{-# INLINE forLoop #-}


-- | @forLoopState start cond inc initialState f@: A C-style for loop with
-- starting value, loop condition, incrementor and a state that is threaded
-- through the computation.
forLoopState :: (Monad m) => a -> (a -> Bool) -> (a -> a) -> b -> (b -> a -> m b) -> m b
forLoopState :: a -> (a -> Bool) -> (a -> a) -> b -> (b -> a -> m b) -> m b
forLoopState a
start a -> Bool
cond a -> a
inc b
initialState b -> a -> m b
f = a -> b -> m b
go a
start b
initialState
  where
    go :: a -> b -> m b
go !a
x !b
state | a -> Bool
cond a
x    = b -> a -> m b
f b
state a
x m b -> (b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> b -> m b
go (a -> a
inc a
x)
                 | Bool
otherwise = b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
state

{-# INLINE forLoopState #-}

-- | @forLoopFold start cond inc acc0 f@: A pure fold using a for loop
-- instead of a list for performance.
--
-- Care is taken that @acc0@ not be strictly evaluated if unless done so by @f@.
forLoopFold :: a -> (a -> Bool) -> (a -> a) -> acc -> (acc -> a -> acc) -> acc
forLoopFold :: a -> (a -> Bool) -> (a -> a) -> acc -> (acc -> a -> acc) -> acc
forLoopFold a
start a -> Bool
cond a -> a
inc acc
acc0 acc -> a -> acc
f = acc -> a -> acc
go acc
acc0 a
start
  where
    -- Not using !acc, see:
    --   http://neilmitchell.blogspot.co.uk/2013/08/destroying-performance-with-strictness.html
    go :: acc -> a -> acc
go acc
acc !a
x | a -> Bool
cond a
x    = let acc' :: acc
acc' = acc -> a -> acc
f acc
acc a
x
                             in acc
acc' acc -> acc -> acc
`seq` acc -> a -> acc
go acc
acc' (a -> a
inc a
x)
              | Bool
otherwise = acc
acc

{-# INLINE forLoopFold #-}


-- | @numLoop start end f@: Loops over a contiguous numerical range, including
-- @end@.
--
-- Does nothing when not @start <= end@.
--
-- It uses @(+ 1)@ so for most integer types it has no bounds (overflow) check.
numLoop :: (Num a, Ord a, Monad m) => a -> a -> (a -> m ()) -> m ()
numLoop :: a -> a -> (a -> m ()) -> m ()
numLoop a
start a
end a -> m ()
f = if a
start a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
end then a -> m ()
go a
start else () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  where
    go :: a -> m ()
go !a
x | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
end  = a -> m ()
f a
x
          | Bool
otherwise = a -> m ()
f a
x m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> m ()
go (a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
1)

{-# INLINE numLoop #-}

-- | @numLoopState start end f initialState@: Loops over a contiguous numerical
-- range, including @end@ threading a state through the computation.
--
-- It uses @(+ 1)@ so for most integer types it has no bounds (overflow) check.
numLoopState :: (Num a, Eq a, Monad m) => a -> a -> b -> (b -> a -> m b) -> m b
numLoopState :: a -> a -> b -> (b -> a -> m b) -> m b
numLoopState a
start a
end b
initState b -> a -> m b
f = a -> b -> m b
go a
start b
initState
  where
    go :: a -> b -> m b
go !a
x !b
state | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
end  = b -> a -> m b
f b
state a
x
                 | Bool
otherwise = b -> a -> m b
f b
state a
x m b -> (b -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> b -> m b
go (a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
1)


{-# INLINE numLoopState #-}

-- | @numLoopFold start end acc0 f@: A pure fold over a contiguous numerical
-- range, including @end@.
--
-- It uses @(+ 1)@ so for most integer types it has no bounds (overflow) check.
--
-- Care is taken that @acc0@ not be strictly evaluated if unless done so by @f@.
numLoopFold :: (Num a, Eq a) => a -> a -> acc -> (acc -> a -> acc) -> acc
numLoopFold :: a -> a -> acc -> (acc -> a -> acc) -> acc
numLoopFold a
start a
end acc
acc0 acc -> a -> acc
f = acc -> a -> acc
go acc
acc0 a
start
  where
    go :: acc -> a -> acc
go acc
acc !a
x | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
end  = acc -> a -> acc
f acc
acc a
x
              | Bool
otherwise = let acc' :: acc
acc' = acc -> a -> acc
f acc
acc a
x
                             in acc
acc' acc -> acc -> acc
`seq` acc -> a -> acc
go acc
acc' (a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
1)

{-# INLINE numLoopFold #-}