content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE ExplicitForAll #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE MonoLocalBinds #-}
  {-# LANGUAGE PatternSynonyms #-}
  
  -- |
  -- Module      : Data.Massiv.Core.Index.Iterator
  -- Copyright   : (c) Alexey Kuleshevich 2021-2022
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  module Data.Massiv.Core.Index.Iterator (
    Iterator (..),
  
    -- * Extra iterator functions
    iterTargetAccST,
    iterTargetAccST_,
    iterTargetFullWithStrideAccST,
    iterTargetFullWithStrideAccST_,
    iterTargetST_,
    iterTargetFullWithStrideST_,
  
    -- * Iterator implementations
    RowMajor (RowMajor),
    defRowMajor,
    RowMajorLinear (RowMajorLinear),
    defRowMajorLinear,
    RowMajorUnbalanced (RowMajorUnbalanced),
    defRowMajorUnbalanced,
  ) where
  
  import Control.Monad
  import Control.Monad.ST
  import Control.Scheduler
  import Data.Massiv.Core.Index.Internal
  import Data.Massiv.Core.Index.Stride
  import Data.Massiv.Core.Loop
  
  class Iterator it where
    {-# MINIMAL (iterTargetM, iterTargetA_, iterTargetWithStrideAccST, iterTargetWithStrideAccST_) #-}
  
    -- | Iterate over a target region using linear index with access to the source
    -- index, which adjusted according to the stride. Use `iterTargetM` if you
    -- need an accumulator.
    --
    -- @since 1.0.2
    iterTargetA_
      :: (Index ix, Applicative f)
      => it
      -> Int
      -- ^ Target linear index start
      -> Sz ix
      -- ^ Target size
      -> ix
      -- ^ Source start index
      -> Stride ix
      -- ^ Source stride
      -> (Ix1 -> ix -> f a)
      -- ^ Action that accepts a linear index of the target and multi-dimensional
      -- index of the source.
      -> f ()
  
    -- | Iterate over a target region using linear index with access to the source
    -- index, which adjusted according to the stride.
    --
    -- @since 1.0.2
    iterTargetM
      :: (Index ix, Monad m)
      => it
      -> Ix1
      -- ^ Target linear index start
      -> Sz ix
      -- ^ Target size
      -> ix
      -- ^ Source start index
      -> Stride ix
      -- ^ Source stride
      -> a
      -- ^ Accumulator
      -> (Ix1 -> ix -> a -> m a)
      -- ^ Action that accepts a linear index of the target,
      -- multi-dimensional index of the source and accumulator
      -> m a
  
    iterTargetWithStrideAccST
      :: Index ix
      => it
      -> Scheduler s a
      -- ^ Scheduler to use
      -> Ix1
      -- ^ Target linear start index
      -> Sz ix
      -- ^ Target size
      -> ix
      -- ^ Source start index
      -> Stride ix
      -- ^ Source stride
      -> a
      -- ^ Initial accumulator
      -> (a -> ST s (a, a))
      -- ^ Splitting action that produces new accumulators for separate worker threads.
      -> (Ix1 -> ix -> a -> ST s a)
      -- ^ Action
      -> ST s a
  
    iterTargetWithStrideAccST_
      :: Index ix
      => it
      -> Scheduler s ()
      -- ^ Scheduler to use
      -> Ix1
      -- ^ Target linear start index
      -> Sz ix
      -- ^ Target size
      -> ix
      -- ^ Start
      -> Stride ix
      -- ^ Stride
      -> a
      -- ^ Initial accumulator
      -> (a -> ST s (a, a))
      -- ^ Splitting action that produces new accumulators for separate worker threads.
      -> (Ix1 -> ix -> a -> ST s a)
      -- ^ Action
      -> ST s ()
  
    -- | Iterate over a region with a monadic action and accumulator.
    --
    -- @since 1.0.2
    iterFullM
      :: (Index ix, Monad m)
      => it
      -> ix
      -- ^ Source start index
      -> Sz ix
      -- ^ Source size
      -> a
      -- ^ Accumulator
      -> (ix -> a -> m a)
      -- ^ Action that accepts a linear index of the target,
      -- multi-dimensional index of the source and accumulator
      -> m a
    iterFullM it start sz acc f =
      iterTargetM it 0 sz start oneStride acc (const f)
    {-# INLINE iterFullM #-}
  
    -- | Iterate over a region with an applicative action ignoring the result.
    --
    -- @since 1.0.2
    iterFullA_
      :: (Index ix, Applicative f)
      => it
      -> ix
      -- ^ Source start index
      -> Sz ix
      -- ^ Source size
      -> (ix -> f a)
      -- ^ Action that accepts a linear index of the target,
      -- multi-dimensional index of the source and accumulator
      -> f ()
    iterFullA_ it start sz f =
      iterTargetA_ it 0 sz start oneStride (const f)
    {-# INLINE iterFullA_ #-}
  
    -- | Iterate over a region in a ST monad with access to `Scheduler`.
    iterFullAccST
      :: Index ix
      => it
      -- ^ Scheduler multiplying factor. Must be positive
      -> Scheduler s a
      -- ^ Scheduler to use
      -> ix
      -- ^ Start index
      -> Sz ix
      -- ^ Size
      -> a
      -- ^ Initial accumulator
      -> (a -> ST s (a, a))
      -- ^ Function that splits accumulator for each scheduled job.
      -> (ix -> a -> ST s a)
      -- ^ Action
      -> ST s a
    iterFullAccST it scheduler start sz acc splitAcc f =
      iterTargetAccST it scheduler 0 sz start acc splitAcc (const f)
    {-# INLINE iterFullAccST #-}
  
    iterTargetFullAccST
      :: Index ix
      => it
      -> Scheduler s a
      -- ^ Scheduler to use
      -> Ix1
      -- ^ Target linear start index
      -> Sz ix
      -- ^ Target size
      -> a
      -- ^ Initial accumulator
      -> (a -> ST s (a, a))
      -- ^ Function that splits accumulator for each scheduled job.
      -> (Ix1 -> ix -> a -> ST s a)
      -- ^ Action
      -> ST s a
    iterTargetFullAccST it scheduler iStart sz =
      iterTargetFullWithStrideAccST it scheduler iStart sz oneStride
    {-# INLINE iterTargetFullAccST #-}
  
    iterTargetFullAccST_
      :: Index ix
      => it
      -> Scheduler s ()
      -- ^ Scheduler to use
      -> Ix1
      -- ^ Target linear start index
      -> Sz ix
      -- ^ Target size
      -> a
      -- ^ Initial accumulator
      -> (a -> ST s (a, a))
      -- ^ Function that splits accumulator for each scheduled job.
      -> (Ix1 -> ix -> a -> ST s a)
      -- ^ Action
      -> ST s ()
    iterTargetFullAccST_ it scheduler iStart sz =
      iterTargetFullWithStrideAccST_ it scheduler iStart sz oneStride
    {-# INLINE iterTargetFullAccST_ #-}
  
    iterTargetFullST_
      :: Index ix
      => it
      -> Scheduler s ()
      -- ^ Scheduler to use
      -> Ix1
      -- ^ Target linear start index
      -> Sz ix
      -- ^ Target size
      -> (Ix1 -> ix -> ST s ())
      -- ^ Action
      -> ST s ()
    iterTargetFullST_ it scheduler iStart sz =
      iterTargetST_ it scheduler iStart sz (pureIndex 0)
    {-# INLINE iterTargetFullST_ #-}
  
    -- NOTE: this function does not have to be part of the class, but for some
    -- reason it creates a severe regression when moved outside.
  
    -- | Iterate over a target array with a stride without an accumulator
    iterTargetWithStrideST_
      :: Index ix
      => it
      -> Scheduler s ()
      -- ^ Scheduler to use
      -> Ix1
      -- ^ Target linear start index
      -> Sz ix
      -- ^ Target size
      -> ix
      -- ^ Start
      -> Stride ix
      -- ^ Stride
      -> (Ix1 -> ix -> ST s a)
      -- ^ Action
      -> ST s ()
    iterTargetWithStrideST_ it scheduler i sz ix stride action =
      iterTargetWithStrideAccST_ it scheduler i sz ix stride () noSplit $ \j jx _ ->
        void $ action j jx
    {-# INLINE iterTargetWithStrideST_ #-}
  
  -- | Default iterator that parallelizes work in linear chunks. Supplied factor
  -- will be used to schedule that many jobs per capability.
  --
  -- @since 1.0.2
  newtype RowMajor = RowMajorInternal Int
  
  -- | Default row major iterator with multiplying factor set to @8@.
  defRowMajor :: RowMajor
  defRowMajor = RowMajorInternal 8
  
  pattern RowMajor
    :: Int
    -- ^ Multiplier that will be used to scale number of jobs.
    -> RowMajor
  pattern RowMajor f <- RowMajorInternal f
    where
      RowMajor = RowMajorInternal . max 1
  
  {-# COMPLETE RowMajor #-}
  
  instance Iterator RowMajor where
    iterFullM _ start (Sz sz) = iterM start sz (pureIndex 1) (<)
    {-# INLINE iterFullM #-}
    iterFullA_ _ start (Sz sz) = iterA_ start sz (pureIndex 1) (<)
    {-# INLINE iterFullA_ #-}
    iterFullAccST (RowMajorInternal fact) scheduler startIx =
      iterRowMajorST fact scheduler startIx (pureIndex 1)
    {-# INLINE iterFullAccST #-}
    iterTargetA_ _ i sz start (Stride stride) =
      iterTargetRowMajorA_ 0 i sz start stride
    {-# INLINE iterTargetA_ #-}
    iterTargetM _ i sz start (Stride stride) =
      iterTargetRowMajorAccM 0 i sz start stride
    {-# INLINE iterTargetM #-}
    iterTargetWithStrideAccST (RowMajor fact) scheduler i sz ix (Stride stride) =
      iterTargetRowMajorAccST 0 fact scheduler i sz ix stride
    {-# INLINE iterTargetWithStrideAccST #-}
    iterTargetWithStrideAccST_ (RowMajor fact) scheduler i sz ix (Stride stride) =
      iterTargetRowMajorAccST_ 0 fact scheduler i sz ix stride
    {-# INLINE iterTargetWithStrideAccST_ #-}
  
  newtype RowMajorLinear = RowMajorLinear Int
  
  defRowMajorLinear :: RowMajorLinear
  defRowMajorLinear = RowMajorLinear 8
  
  instance Iterator RowMajorLinear where
    iterTargetM _ iStart sz start (Stride stride) acc action =
      loopM 0 (< totalElem sz) (+ 1) acc $ \i ->
        action (iStart + i) (liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i)))
    {-# INLINE iterTargetM #-}
    iterTargetA_ _ iStart sz start (Stride stride) action =
      loopA_ 0 (< totalElem sz) (+ 1) $ \i ->
        action (iStart + i) (liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i)))
    {-# INLINE iterTargetA_ #-}
    iterTargetFullAccST it scheduler iStart sz acc splitAcc action =
      let !(RowMajorLinear fact) = it
       in iterLinearAccST fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i ->
            action i (fromLinearIndex sz i)
    {-# INLINE iterTargetFullAccST #-}
    iterTargetFullAccST_ it scheduler iStart sz acc splitAcc action =
      let !(RowMajorLinear fact) = it
       in iterLinearAccST_ fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i ->
            action i (fromLinearIndex sz i)
    {-# INLINE iterTargetFullAccST_ #-}
    iterTargetFullST_ it scheduler iStart sz action =
      let !(RowMajorLinear fact) = it
       in iterLinearST_ fact scheduler iStart 1 (totalElem sz) $ \ !i ->
            action i (fromLinearIndex sz i)
    {-# INLINE iterTargetFullST_ #-}
    iterTargetWithStrideAccST it scheduler iStart sz start (Stride stride) acc spliAcc action =
      let RowMajorLinear fact = it
       in iterLinearAccST fact scheduler 0 1 (totalElem sz) acc spliAcc $ \i ->
            action (iStart + i) $
              liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i))
    {-# INLINE iterTargetWithStrideAccST #-}
    iterTargetWithStrideAccST_ it scheduler iStart sz start (Stride stride) acc spliAcc action =
      let RowMajorLinear fact = it
       in iterLinearAccST_ fact scheduler 0 1 (totalElem sz) acc spliAcc $ \i ->
            action (iStart + i) $
              liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i))
    {-# INLINE iterTargetWithStrideAccST_ #-}
  
  -- | Parallelizing unbalanced computation (i.e. computing some elements of the
  -- array is much more expensive then the others) it can be benefitial to
  -- interleave iteration. Perfect example of this would be a ray tracer or the
  -- Mandelbrot set.
  --
  -- iteration without parallelization is equivalent to `RowMajor`
  --
  -- @since 1.0.2
  newtype RowMajorUnbalanced = RowMajorUnbalancedInternal Int
  
  defRowMajorUnbalanced :: RowMajorUnbalanced
  defRowMajorUnbalanced = RowMajorUnbalancedInternal 8
  
  pattern RowMajorUnbalanced
    :: Int
    -- ^ Multiplier that will be used to scale number of jobs.
    -> RowMajorUnbalanced
  pattern RowMajorUnbalanced f <- RowMajorUnbalancedInternal f
    where
      RowMajorUnbalanced = RowMajorUnbalancedInternal . max 1
  
  {-# COMPLETE RowMajorUnbalanced #-}
  
  instance Iterator RowMajorUnbalanced where
    iterFullM (RowMajorUnbalanced fact) = iterFullM (RowMajor fact)
    {-# INLINE iterFullM #-}
    iterFullA_ (RowMajorUnbalanced fact) = iterFullA_ (RowMajor fact)
    {-# INLINE iterFullA_ #-}
    iterTargetM (RowMajorUnbalanced fact) = iterTargetM (RowMajor fact)
    {-# INLINE iterTargetM #-}
    iterTargetA_ (RowMajorUnbalanced fact) = iterTargetA_ (RowMajor fact)
    {-# INLINE iterTargetA_ #-}
    iterTargetWithStrideAccST = iterUnbalancedTargetWithStride loopM
    {-# INLINE iterTargetWithStrideAccST #-}
    iterTargetWithStrideAccST_ it scheduler iStart sz start stride acc splitAcc' action =
      void $
        iterUnbalancedTargetWithStride innerLoop it scheduler iStart sz start stride acc splitAcc' action
      where
        innerLoop initial condition increment initAcc f =
          void $ loopM initial condition increment initAcc f
        {-# INLINE innerLoop #-}
    {-# INLINE iterTargetWithStrideAccST_ #-}
  
  iterUnbalancedTargetWithStride
    :: Index ix
    => (Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b)
    -> RowMajorUnbalanced
    -> Scheduler s b
    -> Int
    -> Sz ix
    -> ix
    -> Stride ix
    -> a
    -> (a -> ST s (a, a))
    -> (Int -> ix -> t)
    -> ST s a
  iterUnbalancedTargetWithStride innerLoop it scheduler iStart sz start stride acc splitAcc action =
    let RowMajorUnbalanced fact = it
        !n = totalElem sz
        !step = min (fact * numWorkers scheduler) n
     in loopM 0 (< step) (+ 1) acc $ \ !istep !a -> do
          (curAcc, nextAcc) <- splitAcc a
          scheduleMassivWork scheduler $
            innerLoop istep (< n) (+ step) curAcc $ \i ->
              action (iStart + i) $
                liftIndex2 (+) start (liftIndex2 (*) (unStride stride) (fromLinearIndex sz i))
          pure nextAcc
  {-# INLINE iterUnbalancedTargetWithStride #-}
  
  noSplit :: Applicative m => () -> m ((), ())
  noSplit _ = pure ((), ())
  
  iterTargetAccST
    :: (Iterator it, Index ix)
    => it
    -> Scheduler s a
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start
    -> a
    -> (a -> ST s (a, a))
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s a
  iterTargetAccST it scheduler iStart sz ix =
    iterTargetWithStrideAccST it scheduler iStart sz ix oneStride
  {-# INLINE iterTargetAccST #-}
  
  iterTargetAccST_
    :: (Iterator it, Index ix)
    => it
    -> Scheduler s ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start
    -> a
    -> (a -> ST s (a, a))
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s ()
  iterTargetAccST_ it scheduler iStart sz ix =
    iterTargetWithStrideAccST_ it scheduler iStart sz ix oneStride
  {-# INLINE iterTargetAccST_ #-}
  
  iterTargetFullWithStrideST_
    :: (Iterator it, Index ix)
    => it
    -> Scheduler s ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> Stride ix
    -- ^ Stride
    -> (Ix1 -> ix -> ST s ())
    -- ^ Action
    -> ST s ()
  iterTargetFullWithStrideST_ it scheduler iStart sz =
    iterTargetWithStrideST_ it scheduler iStart sz (pureIndex 0)
  {-# INLINE iterTargetFullWithStrideST_ #-}
  
  iterTargetST_
    :: (Iterator it, Index ix)
    => it
    -> Scheduler s ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Start
    -> (Ix1 -> ix -> ST s ())
    -- ^ Action
    -> ST s ()
  iterTargetST_ it scheduler iStart sz ix =
    iterTargetWithStrideST_ it scheduler iStart sz ix oneStride
  {-# INLINE iterTargetST_ #-}
  
  iterTargetFullWithStrideAccST
    :: (Iterator it, Index ix)
    => it
    -> Scheduler s a
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> Stride ix
    -- ^ Stride
    -> a
    -> (a -> ST s (a, a))
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s a
  iterTargetFullWithStrideAccST it scheduler iStart sz =
    iterTargetWithStrideAccST it scheduler iStart sz (pureIndex 0)
  {-# INLINE iterTargetFullWithStrideAccST #-}
  
  iterTargetFullWithStrideAccST_
    :: (Iterator it, Index ix)
    => it
    -> Scheduler s ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> Stride ix
    -- ^ Stride
    -> a
    -> (a -> ST s (a, a))
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s ()
  iterTargetFullWithStrideAccST_ it scheduler iStart sz =
    iterTargetWithStrideAccST_ it scheduler iStart sz (pureIndex 0)
  {-# INLINE iterTargetFullWithStrideAccST_ #-}
  