content: {-# 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_ #-}
  