{-# 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 :: forall a.
Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a
loop Int
initial Int -> Bool
condition Int -> Int
increment a
initAcc Int -> a -> a
f =
  forall a. Identity a -> a
runIdentity (forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
initial Int -> Bool
condition Int -> Int
increment a
initAcc (coerce :: forall a b. Coercible a b => a -> b
coerce Int -> a -> a
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 :: forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM !Int
initial Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> a -> m a
f =
  Int -> a -> m a
go Int
initial a
initAcc
  where
    go :: Int -> a -> m a
go !Int
step !a
acc
      | Int -> Bool
condition Int
step = Int -> a -> m a
f Int
step a
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m a
go (Int -> Int
increment Int
step)
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
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 :: forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM !Int
istart !Int
initIx Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> Int -> a -> m a
f = Int -> Int -> a -> m a
go Int
istart Int
initIx a
initAcc
  where
    go :: Int -> Int -> a -> m a
go !Int
i !Int
step !a
acc
      | Int -> Bool
condition Int
step = Int -> Int -> a -> m a
f Int
i Int
step a
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> a -> m a
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int
increment Int
step)
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
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_ :: forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ !Int
initial Int -> Bool
condition Int -> Int
increment Int -> m a
f = Int -> m ()
go Int
initial
  where
    go :: Int -> m ()
go !Int
step
      | Int -> Bool
condition Int
step = Int -> m a
f Int
step forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
go (Int -> Int
increment Int
step)
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
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_ :: forall (f :: * -> *) a.
Applicative f =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f a)
-> f ()
iloopA_ !Int
istart !Int
initIx Int -> Bool
condition Int -> Int
increment Int -> Int -> f a
f = Int -> Int -> f ()
go Int
istart Int
initIx
  where
    go :: Int -> Int -> f ()
go !Int
i !Int
step
      | Int -> Bool
condition Int
step = Int -> Int -> f a
f Int
i Int
step forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> Int -> f ()
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int
increment Int
step)
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
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_ :: forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ !Int
initial Int -> Bool
condition Int -> Int
increment Int -> Int -> f a
f = Int -> f ()
go Int
initial
  where
    go :: Int -> f ()
go !Int
step
      | Int -> Bool
condition Int
step =
          let !next :: Int
next = Int -> Int
increment Int
step
           in Int -> Int -> f a
f Int
step Int
next forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> f ()
go Int
next
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
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 :: forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
loopNextM !Int
initial Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> Int -> a -> m a
f = Int -> a -> m a
go Int
initial a
initAcc
  where
    go :: Int -> a -> m a
go !Int
step !a
acc
      | Int -> Bool
condition Int
step =
          let !next :: Int
next = Int -> Int
increment Int
step
           in Int -> Int -> a -> m a
f Int
step Int
next a
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m a
go Int
next
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
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_ :: forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ !Int
initial Int -> Bool
condition Int -> Int
increment Int -> f a
f =
  forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
initial Int -> Bool
condition Int -> Int
increment (forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (\Int
i f ()
ma -> Int -> f a
f Int
i forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> f ()
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 :: forall (f :: * -> *) b.
Applicative f =>
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f b
-> (Int -> f (b -> b))
-> f b
loopA !Int
initial Int -> Bool
condition Int -> Int
increment f b
lastAction Int -> f (b -> b)
f =
  forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
initial Int -> Bool
condition Int -> Int
increment f b
lastAction (\Int
i f b
ma -> Int -> f (b -> b)
f Int
i forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f b
ma)
{-# INLINE loopA #-}

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

nextMaybeF :: Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF :: forall (f :: * -> *) a.
Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF !Int
cur Int -> Bool
condition Int -> Int
increment Maybe Int -> f a
f =
  let !i :: Int
i = Int -> Int
increment Int
cur
   in Maybe Int -> f a
f forall a b. (a -> b) -> a -> b
$! if Int -> Bool
condition Int
i then forall a. a -> Maybe a
Just Int
i else forall a. Maybe a
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 :: forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM !Int
initial Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> a -> m a
f =
  forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
initial Int -> Bool
condition Int -> Int
increment (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
initAcc) (\Int
i m a
ma -> m a
ma forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m a
f Int
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 :: forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly Int
numChunks Int
totalLength Int -> Int -> a
action = Int -> Int -> a
action Int
chunkLength Int
slackStart
  where
    !chunkLength :: Int
chunkLength = Int
totalLength forall a. Integral a => a -> a -> a
`quot` Int
numChunks
    !slackStart :: Int
slackStart = Int
chunkLength forall a. Num a => a -> a -> a
* Int
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_ :: forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler s ()
scheduler Int
totalLength Int -> Int -> m ()
action =
  forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
    forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \Int
start Int
next ->
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Int -> Int -> m ()
action Int
start Int
next
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
        Int -> Int -> m ()
action Int
slackStart Int
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 :: forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> Int -> (Int -> Int -> m a) -> m ()
splitLinearlyM Scheduler s a
scheduler Int
totalLength Int -> Int -> m a
action =
  forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
    forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \Int
start Int
next ->
      forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s a
scheduler (Int -> Int -> m a
action Int
start Int
next)
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
      forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s a
scheduler (Int -> Int -> m a
action Int
slackStart Int
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_ :: forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler Int
totalLength Int -> b
index =
  forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ Scheduler s ()
scheduler Int
totalLength (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> b
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_ :: forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ Scheduler s ()
scheduler Int
totalLength Int -> m b
make Int -> b -> m c
write =
  forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler s ()
scheduler Int
totalLength Int -> Int -> m ()
go
  where
    go :: Int -> Int -> m ()
go Int
start Int
end = forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (forall a. Ord a => a -> a -> Bool
< Int
end) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \Int
k -> Int -> m b
make Int
k forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
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_ :: forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s ()
-> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ Scheduler s ()
scheduler Int
startAt Int
totalLength Int -> m b
make Int -> b -> m c
write =
  forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
    forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAt (forall a. Ord a => a -> a -> Bool
< (Int
slackStart forall a. Num a => a -> a -> a
+ Int
startAt)) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (forall a. Ord a => a -> a -> Bool
< (Int
start forall a. Num a => a -> a -> a
+ Int
chunkLength)) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$
          \ !Int
k -> Int -> m b
make Int
k forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
k
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ (Int
slackStart forall a. Num a => a -> a -> a
+ Int
startAt) (forall a. Ord a => a -> a -> Bool
< (Int
totalLength forall a. Num a => a -> a -> a
+ Int
startAt)) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$
          \ !Int
k -> Int -> m b
make Int
k forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
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_ :: forall (m :: * -> *) ws b c.
MonadUnliftIO m =>
SchedulerWS ws ()
-> Int -> (Int -> ws -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStatefulM_ SchedulerWS ws ()
schedulerWS Int
totalLength Int -> ws -> m b
make Int -> b -> m c
store =
  let nWorkers :: Int
nWorkers = forall s a. Scheduler s a -> Int
numWorkers (forall ws a. SchedulerWS ws a -> Scheduler RealWorld a
unwrapSchedulerWS SchedulerWS ws ()
schedulerWS)
   in forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
        forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly Int
nWorkers Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
          forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
            forall (m :: * -> *) ws.
MonadPrimBase RealWorld m =>
SchedulerWS ws () -> (ws -> m ()) -> m ()
scheduleWorkState_ SchedulerWS ws ()
schedulerWS forall a b. (a -> b) -> a -> b
$ \ws
s ->
              forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (forall a. Ord a => a -> a -> Bool
< (Int
start forall a. Num a => a -> a -> a
+ Int
chunkLength)) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
                forall a. m a -> IO a
run (Int -> ws -> m b
make Int
k ws
s forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
store Int
k)
          forall (m :: * -> *) ws.
MonadPrimBase RealWorld m =>
SchedulerWS ws () -> (ws -> m ()) -> m ()
scheduleWorkState_ SchedulerWS ws ()
schedulerWS forall a b. (a -> b) -> a -> b
$ \ws
s ->
            forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
slackStart (forall a. Ord a => a -> a -> Bool
< Int
totalLength) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
              forall a. m a -> IO a
run (Int -> ws -> m b
make Int
k ws
s forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
store Int
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 :: forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s a
scheduler Int
start Int
step Int
totalLength b
initAcc b -> ST s (b, b)
splitAcc Int -> Int -> Int -> Int -> b -> ST s a
action = do
  let !(Int
chunkLength, Int
slackStart) = Int -> Int -> Int -> (Int, Int)
splitNumChunks Int
fact (forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler) Int
totalLength
  b
slackAcc <-
    forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) b
initAcc forall a b. (a -> b) -> a -> b
$ \ !Int
chunkStart !b
acc -> do
      (b
accCur, b
accNext) <- b -> ST s (b, b)
splitAcc b
acc
      forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork Scheduler s a
scheduler forall a b. (a -> b) -> a -> b
$ do
        let !chunkStartAdj :: Int
chunkStartAdj = Int
start forall a. Num a => a -> a -> a
+ Int
chunkStart forall a. Num a => a -> a -> a
* Int
step
            !chunkStopAdj :: Int
chunkStopAdj = Int
chunkStartAdj forall a. Num a => a -> a -> a
+ Int
chunkLength forall a. Num a => a -> a -> a
* Int
step
        Int -> Int -> Int -> Int -> b -> ST s a
action Int
chunkStart Int
chunkLength Int
chunkStartAdj Int
chunkStopAdj b
accCur
      forall (f :: * -> *) a. Applicative f => a -> f a
pure b
accNext
  let !slackLength :: Int
slackLength = Int
totalLength forall a. Num a => a -> a -> a
- Int
slackStart
  if Int
slackLength forall a. Ord a => a -> a -> Bool
> Int
0
    then do
      (b
curAcc, b
nextAcc) <- b -> ST s (b, b)
splitAcc b
slackAcc
      forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork Scheduler s a
scheduler forall a b. (a -> b) -> a -> b
$ do
        let !slackStartAdj :: Int
slackStartAdj = Int
start forall a. Num a => a -> a -> a
+ Int
slackStart forall a. Num a => a -> a -> a
* Int
step
            !slackStopAdj :: Int
slackStopAdj = Int
slackStartAdj forall a. Num a => a -> a -> a
+ Int
slackLength forall a. Num a => a -> a -> a
* Int
step
        Int -> Int -> Int -> Int -> b -> ST s a
action Int
slackStart Int
slackLength Int
slackStartAdj Int
slackStopAdj b
curAcc
      forall (f :: * -> *) a. Applicative f => a -> f a
pure b
nextAcc
    else forall (f :: * -> *) a. Applicative f => a -> f a
pure b
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_ :: forall s a.
Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> (Int -> ST s a)
-> ST s ()
iterLinearST_ Int
fact Scheduler s ()
scheduler Int
start Int
step Int
n Int -> ST s a
action = do
  let totalLength :: Int
totalLength = (Int
n forall a. Num a => a -> a -> a
- Int
start) forall a. Integral a => a -> a -> a
`quot` Int
step
  forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s ()
scheduler Int
start Int
step Int
totalLength () (\()
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ((), ())) forall a b. (a -> b) -> a -> b
$
    \Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj ()
_ ->
      forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
step) Int -> ST s a
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_ :: forall s a.
Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s ()
iterLinearAccST_ Int
fact Scheduler s ()
scheduler Int
start Int
step Int
n a
initAcc a -> ST s (a, a)
splitAcc Int -> a -> ST s a
action = do
  let totalLength :: Int
totalLength = (Int
n forall a. Num a => a -> a -> a
- Int
start) forall a. Integral a => a -> a -> a
`quot` Int
step
  forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
    forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s ()
scheduler Int
start Int
step Int
totalLength a
initAcc a -> ST s (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$
      \Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj a
accCur ->
        forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
step) a
accCur Int -> a -> ST s a
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 :: forall s a.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s a
iterLinearAccST Int
fact Scheduler s a
scheduler Int
start Int
step Int
n a
initAcc a -> ST s (a, a)
splitAcc Int -> a -> ST s a
action = do
  let totalLength :: Int
totalLength = (Int
n forall a. Num a => a -> a -> a
- Int
start) forall a. Integral a => a -> a -> a
`quot` Int
step
  forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s a
scheduler Int
start Int
step Int
totalLength a
initAcc a -> ST s (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$
    \Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj a
accCur ->
      forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
step) a
accCur Int -> a -> ST s a
action
{-# INLINE iterLinearAccST #-}

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

-- | Helper for adjusting stride of a chunk
stepStartAdjust :: Int -> Int -> Int
stepStartAdjust :: Int -> Int -> Int
stepStartAdjust Int
step Int
ix = Int
ix forall a. Num a => a -> a -> a
+ ((Int
step forall a. Num a => a -> a -> a
- (Int
ix forall a. Integral a => a -> a -> a
`mod` Int
step)) forall a. Integral a => a -> a -> a
`mod` Int
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 :: forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork = forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
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 -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ Comp
comp Scheduler RealWorld () -> IO ()
f =
  case Comp
comp of
    Comp
Par -> forall (m :: * -> *) a.
MonadUnliftIO m =>
GlobalScheduler -> (Scheduler RealWorld () -> m a) -> m ()
withGlobalScheduler_ GlobalScheduler
globalScheduler Scheduler RealWorld () -> IO ()
f
    Comp
Seq -> Scheduler RealWorld () -> IO ()
f forall s. Scheduler s ()
trivialScheduler_
    Comp
_ -> forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m ()
withScheduler_ Comp
comp Scheduler RealWorld () -> IO ()
f
{-# INLINE withMassivScheduler_ #-}