Loop.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Data.Massiv.Core.Loop
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Core.Loop (
  loop,
  loopF,
  nextMaybeF,
  loopA,
  loopA_,
  loopM,
  loopM_,
  iloopM,
  iloopA_,
  loopNextM,
  loopNextA_,
  loopDeepM,
  splitLinearly,
  splitLinearlyM,
  splitLinearlyM_,
  splitLinearlyWith_,
  splitLinearlyWithM_,
  splitLinearlyWithStartAtM_,
  splitLinearlyWithStatefulM_,
  iterLinearST_,
  iterLinearAccST_,
  iterLinearAccST,
  splitNumChunks,
  stepStartAdjust,

  -- * Experimental
  splitWorkWithFactorST,
  scheduleMassivWork,
  withMassivScheduler_,
) where

import Control.Monad (void, when)
import Control.Monad.IO.Unlift (MonadUnliftIO (..))
import Control.Monad.Primitive
import Control.Monad.ST (ST)
import Control.Scheduler (
  Comp (..),
  Scheduler,
  SchedulerWS,
  numWorkers,
  scheduleWork,
  scheduleWorkState_,
  scheduleWork_,
  trivialScheduler_,
  unwrapSchedulerWS,
  withScheduler_,
 )
import Control.Scheduler.Global (globalScheduler, withGlobalScheduler_)
import Data.Coerce
import Data.Functor.Identity

-- | Efficient loop with an accumulator
--
-- @since 0.1.0
loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a
loop initial condition increment initAcc f =
  runIdentity (loopM initial condition increment initAcc (coerce f))
{-# INLINE loop #-}

-- | Efficient monadic loop with an accumulator
--
-- >>> loopM 1 (< 20) (+ 2) [] (\i a -> Just (i:a))
-- Just [19,17,15,13,11,9,7,5,3,1]
--
-- @since 0.1.0
loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM !initial condition increment !initAcc f =
  go initial initAcc
  where
    go !step !acc
      | condition step = f step acc >>= go (increment step)
      | otherwise = pure acc
{-# INLINE loopM #-}

-- | Efficient monadic loop with an accumulator and extra linear index incremented by 1.
--
-- >>> iloopM 100 1 (< 20) (+ 2) [] (\i ix a -> Just ((i, ix) : a))
-- Just [(109,19),(108,17),(107,15),(106,13),(105,11),(104,9),(103,7),(102,5),(101,3),(100,1)]
--
-- @since 1.0.2
iloopM
  :: Monad m => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a
iloopM !istart !initIx condition increment !initAcc f = go istart initIx initAcc
  where
    go !i !step !acc
      | condition step = f i step acc >>= go (i + 1) (increment step)
      | otherwise = pure acc
{-# INLINE iloopM #-}

-- | Efficient monadic loop. Result of each iteration is discarded.
--
-- @since 0.1.0
loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ !initial condition increment f = go initial
  where
    go !step
      | condition step = f step >> go (increment step)
      | otherwise = pure ()
-- loopF initial condition increment (pure ()) (\i ma -> f i >> ma)
{-# INLINE loopM_ #-}
{-# DEPRECATED loopM_ "In favor of `loopA_`" #-}

-- | Efficient monadic loop with extra linear index incremented by 1.
--
-- >>> iloopA_ 100 1 (< 10) (+ 2) (\i ix -> print (i, ix))
-- (100,1)
-- (101,3)
-- (102,5)
-- (103,7)
-- (104,9)
--
-- @since 1.0.2
iloopA_
  :: Applicative f => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
iloopA_ !istart !initIx condition increment f = go istart initIx
  where
    go !i !step
      | condition step = f i step *> go (i + 1) (increment step)
      | otherwise = pure ()
{-# INLINE iloopA_ #-}

-- | Similar to `loopM_` except the action accepts not only the value for current step,
-- but also for the next one as well.
--
-- @since 1.0.2
loopNextA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ !initial condition increment f = go initial
  where
    go !step
      | condition step =
          let !next = increment step
           in f step next *> go next
      | otherwise = pure ()
{-# INLINE loopNextA_ #-}

-- | Similar to `loopM_` except the action accepts not only the value for current step,
-- but also for the next one as well.
--
-- @since 1.0.2
loopNextM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a
loopNextM !initial condition increment !initAcc f = go initial initAcc
  where
    go !step !acc
      | condition step =
          let !next = increment step
           in f step next acc >>= go next
      | otherwise = pure acc
{-# INLINE loopNextM #-}

-- | Efficient Applicative loop. Result of each iteration is discarded.
--
-- > loopA_ initial cond incr f === loopA initial cond incr (pure ()) (\i -> id <$ f i)
--
-- @since 1.0.2
loopA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ !initial condition increment f =
  loopF initial condition increment (pure ()) (\i ma -> f i *> ma)
{-# INLINE loopA_ #-}

-- | Applicative loop. Use monadic `loopM` when possible, since it will be more efficient.
--
-- @since 0.3.0
loopA :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> f b -> (Int -> f (b -> b)) -> f b
loopA !initial condition increment lastAction f =
  loopF initial condition increment lastAction (\i ma -> f i <*> ma)
{-# INLINE loopA #-}

loopF :: Int -> (Int -> Bool) -> (Int -> Int) -> f a -> (Int -> f a -> f a) -> f a
loopF !initial condition increment lastAction f = go initial
  where
    go !step
      | condition step = f step (go (increment step))
      | otherwise = lastAction
{-# INLINE loopF #-}

nextMaybeF :: Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF !cur condition increment f =
  let !i = increment cur
   in f $! if condition i then Just i else Nothing
{-# INLINE nextMaybeF #-}

-- | Similar to `loopM`, but way less efficient monadic loop with an accumulator that reverses
-- the direction of action application. eg:
--
-- >>> loopDeepM 1 (< 20) (+ 2) [] (\i a -> Just (i:a))
-- Just [1,3,5,7,9,11,13,15,17,19]
--
-- Equivalent to:
--
-- >>> loopM 19 (>= 1) (subtract 2) [] (\i a -> Just (i:a))
-- Just [1,3,5,7,9,11,13,15,17,19]
--
-- @since 0.1.0
loopDeepM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM !initial condition increment !initAcc f =
  loopF initial condition increment (pure initAcc) (\i ma -> ma >>= f i)
{-# INLINE loopDeepM #-}

-- | Divide length in chunks and apply a function to the computed results
--
-- @since 0.2.1
splitLinearly
  :: Int
  -- ^ Number of chunks
  -> Int
  -- ^ Total length
  -> (Int -> Int -> a)
  -- ^ Function that accepts a chunk length and slack start index
  -> a
splitLinearly numChunks totalLength action = action chunkLength slackStart
  where
    !chunkLength = totalLength `quot` numChunks
    !slackStart = chunkLength * numChunks
{-# INLINE splitLinearly #-}

-- | Iterator that expects an action that accepts starting linear index as well as the ending
--
-- @since 0.5.7
splitLinearlyM_
  :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ scheduler totalLength action =
  splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do
    loopNextA_ 0 (< slackStart) (+ chunkLength) $ \start next ->
      scheduleWork_ scheduler $ action start next
    when (slackStart < totalLength) $
      scheduleWork_ scheduler $
        action slackStart totalLength
{-# INLINE splitLinearlyM_ #-}

-- | Iterator that expects an action that accepts starting linear index as well as the ending
--
-- @since 1.0.2
splitLinearlyM
  :: MonadPrimBase s m => Scheduler s a -> Int -> (Int -> Int -> m a) -> m ()
splitLinearlyM scheduler totalLength action =
  splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do
    loopNextA_ 0 (< slackStart) (+ chunkLength) $ \start next ->
      scheduleWork scheduler (action start next)
    when (slackStart < totalLength) $
      scheduleWork scheduler (action slackStart totalLength)
{-# INLINE splitLinearlyM #-}

-- | Iterator that can be used to split computation amongst different workers. For monadic
-- generator see `splitLinearlyWithM_`.
--
-- @since 0.2.1
splitLinearlyWith_
  :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ scheduler totalLength index =
  splitLinearlyWithM_ scheduler totalLength (pure . index)
{-# INLINE splitLinearlyWith_ #-}

-- | Iterator that can be used to split computation jobs
--
-- @since 0.2.6
splitLinearlyWithM_
  :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ scheduler totalLength make write =
  splitLinearlyM_ scheduler totalLength go
  where
    go start end = loopM_ start (< end) (+ 1) $ \k -> make k >>= write k
    {-# INLINE go #-}
{-# INLINE splitLinearlyWithM_ #-}

-- | Iterator that can be used to split computation jobs
--
-- @since 0.3.0
splitLinearlyWithStartAtM_
  :: MonadPrimBase s m => Scheduler s () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ scheduler startAt totalLength make write =
  splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do
    loopM_ startAt (< (slackStart + startAt)) (+ chunkLength) $ \ !start ->
      scheduleWork_ scheduler $
        loopM_ start (< (start + chunkLength)) (+ 1) $
          \ !k -> make k >>= write k
    when (slackStart < totalLength) $
      scheduleWork_ scheduler $
        loopM_ (slackStart + startAt) (< (totalLength + startAt)) (+ 1) $
          \ !k -> make k >>= write k
{-# INLINE splitLinearlyWithStartAtM_ #-}

-- | Iterator that can be used to split computation jobs, while using a stateful scheduler.
--
-- @since 0.3.4
splitLinearlyWithStatefulM_
  :: MonadUnliftIO m
  => SchedulerWS ws ()
  -> Int
  -- ^ Total linear length
  -> (Int -> ws -> m b)
  -- ^ Element producing action
  -> (Int -> b -> m c)
  -- ^ Element storing action
  -> m ()
splitLinearlyWithStatefulM_ schedulerWS totalLength make store =
  let nWorkers = numWorkers (unwrapSchedulerWS schedulerWS)
   in withRunInIO $ \run ->
        splitLinearly nWorkers totalLength $ \chunkLength slackStart -> do
          loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start ->
            scheduleWorkState_ schedulerWS $ \s ->
              loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k ->
                run (make k s >>= store k)
          scheduleWorkState_ schedulerWS $ \s ->
            loopM_ slackStart (< totalLength) (+ 1) $ \ !k ->
              run (make k s >>= store k)
{-# INLINE splitLinearlyWithStatefulM_ #-}

-- | This is a major helper function for fair splitting and parallelization of
-- work with ability to use some arbitrary accumulator and splittable seed
--
-- @since 1.0.2
splitWorkWithFactorST
  :: Int
  -- ^ Multiplying factor to be applied to number of workers for number
  -- of jobs to schedule. Higher the factor, more jobs will be
  -- scheduled. Only positive values are valid.
  -> Scheduler s a
  -> Int
  -- ^ Starting index
  -> Int
  -- ^ Stepping value. Can be negative, but must not be zero.
  -> Int
  -- ^ Total number of steps to be taken
  -> b
  -- ^ Initial value for an accumulator
  -> (b -> ST s (b, b))
  -- ^ An action to split accumulator for multiple threads
  -> (Int -> Int -> Int -> Int -> b -> ST s a)
  -- ^ A job to be scheduled. Accepts:
  --
  -- * Chunk index start
  -- * Chunk length
  -- * Chunk start index adjusted for supplied start and stepping value
  -- * Chunk stop index adjusted for supplied start and stepping value
  -> ST s b
splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc action = do
  let !(chunkLength, slackStart) = splitNumChunks fact (numWorkers scheduler) totalLength
  slackAcc <-
    loopM 0 (< slackStart) (+ chunkLength) initAcc $ \ !chunkStart !acc -> do
      (accCur, accNext) <- splitAcc acc
      scheduleMassivWork scheduler $ do
        let !chunkStartAdj = start + chunkStart * step
            !chunkStopAdj = chunkStartAdj + chunkLength * step
        action chunkStart chunkLength chunkStartAdj chunkStopAdj accCur
      pure accNext
  let !slackLength = totalLength - slackStart
  if slackLength > 0
    then do
      (curAcc, nextAcc) <- splitAcc slackAcc
      scheduleMassivWork scheduler $ do
        let !slackStartAdj = start + slackStart * step
            !slackStopAdj = slackStartAdj + slackLength * step
        action slackStart slackLength slackStartAdj slackStopAdj curAcc
      pure nextAcc
    else pure slackAcc
{-# INLINE splitWorkWithFactorST #-}

-- | Linear iterator that supports multiplying factor
--
-- @since 1.0.2
iterLinearST_
  :: Int
  -> Scheduler s ()
  -> Int
  -> Int
  -> Int
  -> (Int -> ST s a)
  -> ST s ()
iterLinearST_ fact scheduler start step n action = do
  let totalLength = (n - start) `quot` step
  splitWorkWithFactorST fact scheduler start step totalLength () (\_ -> pure ((), ())) $
    \_ _ chunkStartAdj chunkStopAdj _ ->
      loopA_ chunkStartAdj (< chunkStopAdj) (+ step) action
{-# INLINE iterLinearST_ #-}

-- | Linear iterator that supports multiplying factor and accumulator, but the results are discarded.
--
-- @since 1.0.2
iterLinearAccST_
  :: Int
  -> Scheduler s ()
  -> Int
  -> Int
  -> Int
  -> a
  -> (a -> ST s (a, a))
  -> (Int -> a -> ST s a)
  -> ST s ()
iterLinearAccST_ fact scheduler start step n initAcc splitAcc action = do
  let totalLength = (n - start) `quot` step
  void $
    splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc $
      \_ _ chunkStartAdj chunkStopAdj accCur ->
        void $ loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action
{-# INLINE iterLinearAccST_ #-}

-- | Linear iterator that supports multiplying factor and accumulator. Results
-- of actions are stored in the scheduler.
--
-- @since 1.0.2
iterLinearAccST
  :: Int
  -> Scheduler s a
  -> Int
  -> Int
  -- ^ Step. Must be non-zero
  -> Int
  -> a
  -> (a -> ST s (a, a))
  -> (Int -> a -> ST s a)
  -> ST s a
iterLinearAccST fact scheduler start step n initAcc splitAcc action = do
  let totalLength = (n - start) `quot` step
  splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc $
    \_ _ chunkStartAdj chunkStopAdj accCur ->
      loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action
{-# INLINE iterLinearAccST #-}

-- | Helper for figuring out the chunk length and slack start
splitNumChunks :: Int -> Int -> Int -> (Int, Int)
splitNumChunks fact nw totalLength =
  let maxNumChunks = nw * max 1 fact
      !numChunks
        | nw == 1 || totalLength <= 0 = 1 -- Optimize for Seq and avoid `quot` by 0.
        | totalLength <= nw = totalLength
        | totalLength >= maxNumChunks = maxNumChunks
        | otherwise = nw
      !chunkLength = totalLength `quot` numChunks
      !slackStart = chunkLength * numChunks
   in (chunkLength, slackStart)

-- | Helper for adjusting stride of a chunk
stepStartAdjust :: Int -> Int -> Int
stepStartAdjust step ix = ix + ((step - (ix `mod` step)) `mod` step)
{-# INLINE stepStartAdjust #-}

-- | Internal version of a `scheduleWork` that will be replaced by
-- `scheduleWork_` by the compiler whenever action produces `()`
scheduleMassivWork :: PrimBase m => Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork = scheduleWork
{-# INLINE [0] scheduleMassivWork #-}

{-# RULES
"scheduleWork/scheduleWork_/ST" forall (scheduler :: Scheduler s ()) (action :: ST s ()).
  scheduleMassivWork scheduler action =
    scheduleWork_ scheduler action
"scheduleWork/scheduleWork_/IO" forall (scheduler :: Scheduler RealWorld ()) (action :: IO ()).
  scheduleMassivWork scheduler action =
    scheduleWork_ scheduler action
  #-}

-- | Selects an optimal scheduler for the supplied strategy, but it works only in `IO`
--
-- @since 1.0.0
withMassivScheduler_ :: Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ comp f =
  case comp of
    Par -> withGlobalScheduler_ globalScheduler f
    Seq -> f trivialScheduler_
    _ -> withScheduler_ comp f
{-# INLINE withMassivScheduler_ #-}