{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
-- |
-- Module      : Data.Massiv.Core.Iterator
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Core.Iterator
  ( loop
  , loopA_
  , loopM
  , loopM_
  , loopDeepM
  , splitLinearly
  , splitLinearlyM_
  , splitLinearlyWith_
  , splitLinearlyWithM_
  , splitLinearlyWithStartAtM_
  , splitLinearlyWithStatefulM_
  ) where

import Control.Scheduler
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.IO.Unlift

-- | Efficient loop with an accumulator
--
-- @since 0.1.0
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 #-}


-- | 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 :: 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 #-}


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

-- | Similar to `loopM_` axcept the action accepts not only the value for current step,
-- but also for the next one as well.
--
-- @since 0.5.7
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_ #-}


-- | Efficient Applicative loop. Result of each iteration is discarded.
--
-- @since 0.3.0
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_ #-}


-- | Similar to `loopM`, but slightly 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 :: 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 #-}


-- | 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 :: 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 #-}

-- | 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 s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler s ()
scheduler Int
totalLength Int -> Int -> m ()
action =
  Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
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 s () -> m () -> m ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
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 s () -> m () -> m ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m ()
action Int
slackStart Int
totalLength
{-# INLINE splitLinearlyM_ #-}

-- | Interator 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 s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler Int
totalLength Int -> b
index =
  Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m ()) -> m ()
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 (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_ #-}


-- | Interator 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 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 =
  Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
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 = 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_ #-}


-- | Interator 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 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 =
  Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
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 s () -> m () -> m ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
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 s () -> m () -> m ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
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_ #-}



-- | Interator 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 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 = Scheduler RealWorld () -> Int
forall s a. Scheduler s a -> Int
numWorkers (SchedulerWS ws () -> Scheduler RealWorld ()
forall ws a. SchedulerWS ws a -> Scheduler RealWorld a
unwrapSchedulerWS SchedulerWS ws ()
schedulerWS)
   in ((forall a. m a -> IO a) -> IO ()) -> m ()
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO ()) -> m ())
-> ((forall a. m a -> IO a) -> IO ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
      Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly Int
nWorkers Int
totalLength ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
        Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO ()) -> IO ()
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 -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
          SchedulerWS ws () -> (ws -> IO ()) -> IO ()
forall (m :: * -> *) ws.
MonadPrimBase RealWorld m =>
SchedulerWS ws () -> (ws -> m ()) -> m ()
scheduleWorkState_ SchedulerWS ws ()
schedulerWS ((ws -> IO ()) -> IO ()) -> (ws -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ws
s ->
            Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO c) -> IO ()
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 -> IO c) -> IO ()) -> (Int -> IO c) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
              m c -> IO c
forall a. m a -> IO a
run (Int -> ws -> m b
make Int
k ws
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 ws () -> (ws -> IO ()) -> IO ()
forall (m :: * -> *) ws.
MonadPrimBase RealWorld m =>
SchedulerWS ws () -> (ws -> m ()) -> m ()
scheduleWorkState_ SchedulerWS ws ()
schedulerWS ((ws -> IO ()) -> IO ()) -> (ws -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ws
s ->
          Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO c) -> IO ()
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 -> IO c) -> IO ()) -> (Int -> IO c) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
            m c -> IO c
forall a. m a -> IO a
run (Int -> ws -> m b
make Int
k ws
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_ #-}