{-# LANGUAGE BangPatterns #-}
module Data.Massiv.Core.Iterator
( loop
, loopA_
, loopM
, loopM_
, loopDeepM
, splitLinearly
, splitLinearlyM_
, splitLinearlyWith_
, splitLinearlyWithM_
, splitLinearlyWithStartAtM_
, splitLinearlyWithStatefulM_
) where
import Control.Scheduler
import Control.Monad
loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a
loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a
loop !Int
init' Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> a -> a
f = Int -> a -> a
go Int
init' a
initAcc
where
go :: Int -> a -> a
go !Int
step !a
acc
| Int -> Bool
condition Int
step = Int -> a -> a
go (Int -> Int
increment Int
step) (Int -> a -> a
f Int
step a
acc)
| Bool
otherwise = a
acc
{-# INLINE loop #-}
loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM :: Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM !Int
init' Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> a -> m a
f = Int -> a -> m a
go Int
init' 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 m a -> (a -> m a) -> m a
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 = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
acc
{-# INLINE loopM #-}
loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ :: Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ !Int
init' Int -> Bool
condition Int -> Int
increment Int -> m a
f = Int -> m ()
go Int
init'
where
go :: Int -> m ()
go !Int
step
| Int -> Bool
condition Int
step = Int -> m a
f Int
step m a -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
go (Int -> Int
increment Int
step)
| Bool
otherwise = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE loopM_ #-}
loopNextM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> m a) -> m ()
loopNextM_ :: Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> m a) -> m ()
loopNextM_ !Int
init' Int -> Bool
condition Int -> Int
increment Int -> Int -> m a
f = Int -> m ()
go Int
init'
where
go :: Int -> m ()
go Int
step =
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int -> Bool
condition Int
step) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
let !next :: Int
next = Int -> Int
increment Int
step
in Int -> Int -> m a
f Int
step Int
next m a -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
go Int
next
{-# INLINE loopNextM_ #-}
loopA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ :: Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ !Int
init' Int -> Bool
condition Int -> Int
increment Int -> f a
f = Int -> f ()
go Int
init'
where
go :: Int -> f ()
go !Int
step
| Int -> Bool
condition Int
step = Int -> f a
f Int
step f a -> f () -> f ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> f ()
go (Int -> Int
increment Int
step)
| Bool
otherwise = () -> f ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE loopA_ #-}
loopDeepM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM :: Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM !Int
init' Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> a -> m a
f = Int -> a -> m a
go Int
init' a
initAcc
where
go :: Int -> a -> m a
go !Int
step !a
acc
| Int -> Bool
condition Int
step = Int -> a -> m a
go (Int -> Int
increment Int
step) a
acc m a -> (a -> m a) -> m a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m a
f Int
step
| Bool
otherwise = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
acc
{-# INLINE loopDeepM #-}
splitLinearly :: Int
-> Int
-> (Int -> Int -> a)
-> a
splitLinearly :: 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 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
numChunks
!slackStart :: Int
slackStart = Int
chunkLength Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
numChunks
{-# INLINE splitLinearly #-}
splitLinearlyM_ ::
Monad m => Scheduler m () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ :: Scheduler m () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler m ()
scheduler Int
totalLength Int -> Int -> m ()
action =
Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler m () -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers Scheduler m ()
scheduler) Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
Int
-> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> m a) -> m ()
loopNextM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ Int
start Int
next ->
Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m ()
action Int
start Int
next
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m ()
action Int
slackStart Int
totalLength
{-# INLINE splitLinearlyM_ #-}
splitLinearlyWith_ ::
Monad m => Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ :: Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler m ()
scheduler Int
totalLength Int -> b
index =
Scheduler m () -> Int -> (Int -> m b) -> (Int -> b -> m ()) -> m ()
forall (m :: * -> *) b c.
Monad m =>
Scheduler m () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ Scheduler m ()
scheduler Int
totalLength (b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> m b) -> (Int -> b) -> Int -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> b
index)
{-# INLINE splitLinearlyWith_ #-}
splitLinearlyWithM_ ::
Monad m => Scheduler m () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ :: Scheduler m () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ Scheduler m ()
scheduler Int
totalLength Int -> m b
make Int -> b -> m c
write =
Scheduler m () -> Int -> (Int -> Int -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler m ()
scheduler Int
totalLength Int -> Int -> m ()
go
where
go :: Int -> Int -> m ()
go Int
start Int
end = Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m c) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m c) -> m ()) -> (Int -> m c) -> m ()
forall a b. (a -> b) -> a -> b
$ \ Int
k -> Int -> m b
make Int
k m b -> (b -> m c) -> m c
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
k
{-# INLINE go #-}
{-# INLINE splitLinearlyWithM_ #-}
splitLinearlyWithStartAtM_ ::
Monad m => Scheduler m () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ :: Scheduler m ()
-> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ Scheduler m ()
scheduler Int
startAt Int
totalLength Int -> m b
make Int -> b -> m c
write =
Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler m () -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers Scheduler m ()
scheduler) Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
slackStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m c) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m c) -> m ()) -> (Int -> m c) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k -> Int -> m b
make Int
k m b -> (b -> m c) -> m c
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
k
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m c) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ (Int
slackStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt) (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
totalLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m c) -> m ()) -> (Int -> m c) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k -> Int -> m b
make Int
k m b -> (b -> m c) -> m c
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
k
{-# INLINE splitLinearlyWithStartAtM_ #-}
splitLinearlyWithStatefulM_ ::
Monad m
=> SchedulerWS s m ()
-> Int
-> (Int -> s -> m b)
-> (Int -> b -> m c)
-> m ()
splitLinearlyWithStatefulM_ :: SchedulerWS s m ()
-> Int -> (Int -> s -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStatefulM_ SchedulerWS s m ()
schedulerWS Int
totalLength Int -> s -> m b
make Int -> b -> m c
store =
let nWorkers :: Int
nWorkers = Scheduler m () -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers (SchedulerWS s m () -> Scheduler m ()
forall s (m :: * -> *) a. SchedulerWS s m a -> Scheduler m a
unwrapSchedulerWS SchedulerWS s m ()
schedulerWS)
in Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly Int
nWorkers Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
SchedulerWS s m () -> (s -> m ()) -> m ()
forall s (m :: * -> *). SchedulerWS s m () -> (s -> m ()) -> m ()
scheduleWorkState_ SchedulerWS s m ()
schedulerWS ((s -> m ()) -> m ()) -> (s -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \s
s ->
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m c) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m c) -> m ()) -> (Int -> m c) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
Int -> s -> m b
make Int
k s
s m b -> (b -> m c) -> m c
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
store Int
k
SchedulerWS s m () -> (s -> m ()) -> m ()
forall s (m :: * -> *). SchedulerWS s m () -> (s -> m ()) -> m ()
scheduleWorkState_ SchedulerWS s m ()
schedulerWS ((s -> m ()) -> m ()) -> (s -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \s
s ->
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m c) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
slackStart (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m c) -> m ()) -> (Int -> m c) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
Int -> s -> m b
make Int
k s
s m b -> (b -> m c) -> m c
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
store Int
k
{-# INLINE splitLinearlyWithStatefulM_ #-}