content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE MonoLocalBinds #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- |
  -- Module      : Data.Massiv.Array.Ops.Map
  -- Copyright   : (c) Alexey Kuleshevich 2018-2022
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  module Data.Massiv.Array.Ops.Map (
    map,
    imap,
  
    -- ** Traversing
  
    -- *** Applicative
    traverseA,
    traverseA_,
    itraverseA,
    itraverseA_,
    sequenceA,
    sequenceA_,
  
    -- *** PrimMonad
    traversePrim,
    itraversePrim,
  
    -- ** Monadic mapping
  
    -- *** Sequential
    mapM,
    forM,
    imapM,
    iforM,
    mapM_,
    forM_,
    imapM_,
    iforM_,
  
    -- *** Parallelizable
    mapIO,
    mapWS,
    mapIO_,
    imapIO,
    imapWS,
    imapIO_,
    forIO,
    forWS,
    forIO_,
    iforIO,
    iforWS,
    iforIO_,
    imapSchedulerM_,
    iforSchedulerM_,
    iterArrayLinearM_,
    iterArrayLinearWithSetM_,
    iterArrayLinearWithStrideM_,
  
    -- ** Zipping
    zip,
    zip3,
    zip4,
    unzip,
    unzip3,
    unzip4,
    zipWith,
    zipWith3,
    zipWith4,
    izipWith,
    izipWith3,
    izipWith4,
  
    -- *** Applicative
    zipWithA,
    izipWithA,
    zipWith3A,
    izipWith3A,
  ) where
  
  import Control.Monad (void)
  import Control.Monad.Primitive
  import Control.Scheduler
  import Data.Coerce
  import Data.Massiv.Array.Delayed.Pull
  import Data.Massiv.Array.Manifest.List
  import Data.Massiv.Array.Mutable
  import Data.Massiv.Array.Ops.Construct (makeArrayA, makeArrayLinearA)
  import Data.Massiv.Core.Common
  import Data.Traversable (traverse)
  import Prelude hiding (
    map,
    mapM,
    mapM_,
    sequenceA,
    traverse,
    unzip,
    unzip3,
    zip,
    zip3,
    zipWith,
    zipWith3,
   )
  
  --------------------------------------------------------------------------------
  -- map -------------------------------------------------------------------------
  --------------------------------------------------------------------------------
  
  -- | Map a function over an array
  --
  -- @since 0.1.0
  map :: (Index ix, Source r e') => (e' -> e) -> Array r ix e' -> Array D ix e
  map f = fmap f . delay
  {-# INLINE map #-}
  
  --------------------------------------------------------------------------------
  -- zip -------------------------------------------------------------------------
  --------------------------------------------------------------------------------
  
  -- | Zip two arrays
  --
  -- @since 0.1.0
  zip
    :: (Index ix, Source r1 e1, Source r2 e2)
    => Array r1 ix e1
    -> Array r2 ix e2
    -> Array D ix (e1, e2)
  zip = zipWith (,)
  {-# INLINE zip #-}
  
  -- | Zip three arrays
  --
  -- @since 0.1.0
  zip3
    :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3)
    => Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> Array D ix (e1, e2, e3)
  zip3 = zipWith3 (,,)
  {-# INLINE zip3 #-}
  
  -- | Zip four arrays
  --
  -- @since 0.5.4
  zip4
    :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4)
    => Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> Array r4 ix e4
    -> Array D ix (e1, e2, e3, e4)
  zip4 = zipWith4 (,,,)
  {-# INLINE zip4 #-}
  
  -- | Unzip two arrays
  --
  -- @since 0.1.0
  unzip :: (Index ix, Source r (e1, e2)) => Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2)
  unzip arr = (map fst arr, map snd arr)
  {-# INLINE unzip #-}
  
  -- | Unzip three arrays
  --
  -- @since 0.1.0
  unzip3
    :: (Index ix, Source r (e1, e2, e3))
    => Array r ix (e1, e2, e3)
    -> (Array D ix e1, Array D ix e2, Array D ix e3)
  unzip3 arr = (map (\(e, _, _) -> e) arr, map (\(_, e, _) -> e) arr, map (\(_, _, e) -> e) arr)
  {-# INLINE unzip3 #-}
  
  -- | Unzip four arrays
  --
  -- @since 0.5.4
  unzip4
    :: (Index ix, Source r (e1, e2, e3, e4))
    => Array r ix (e1, e2, e3, e4)
    -> (Array D ix e1, Array D ix e2, Array D ix e3, Array D ix e4)
  unzip4 arr =
    ( map (\(e, _, _, _) -> e) arr
    , map (\(_, e, _, _) -> e) arr
    , map (\(_, _, e, _) -> e) arr
    , map (\(_, _, _, e) -> e) arr
    )
  {-# INLINE unzip4 #-}
  
  --------------------------------------------------------------------------------
  -- zipWith ---------------------------------------------------------------------
  --------------------------------------------------------------------------------
  
  -- | Zip two arrays with a function. Resulting array will be an intersection of
  -- source arrays in case their dimensions do not match.
  zipWith
    :: (Index ix, Source r1 e1, Source r2 e2)
    => (e1 -> e2 -> e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array D ix e
  zipWith f arr1 arr2 = DArray comp sz prefIndex
    where
      sz = SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))
      comp = getComp arr1 <> getComp arr2
      prefIndex = PrefIndex (\ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix))
  -- Somehow checking for size equality destroys performance
  --  | PrefIndexLinear gi1 <- unsafePrefIndex arr1,
  --    PrefIndexLinear gi2 <- unsafePrefIndex arr2,
  --    size arr1 == size arr2 =
  --      PrefIndexLinear (\i -> f (gi1 i) (gi2 i))
  --  | otherwise = PrefIndex (\ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix))
  {-# INLINE zipWith #-}
  
  -- | Just like `zipWith`, except with an index aware function.
  izipWith
    :: (Index ix, Source r1 e1, Source r2 e2)
    => (ix -> e1 -> e2 -> e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array D ix e
  izipWith f arr1 arr2 =
    DArray
      (getComp arr1 <> getComp arr2)
      (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))))
      (PrefIndex (\ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)))
  {-# INLINE izipWith #-}
  
  -- | Just like `zipWith`, except zip three arrays with a function.
  zipWith3
    :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3)
    => (e1 -> e2 -> e3 -> e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> Array D ix e
  zipWith3 f arr1 arr2 arr3 = izipWith3 (\_ e1 e2 e3 -> f e1 e2 e3) arr1 arr2 arr3
  -- See note on zipWith
  --  | sz1 == size arr2 && sz1 == size arr3
  --  , PrefIndexLinear gi1 <- unsafePrefIndex arr1
  --  , PrefIndexLinear gi2 <- unsafePrefIndex arr2
  --  , PrefIndexLinear gi3 <- unsafePrefIndex arr3 =
  --    makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i))
  --  | otherwise = izipWith3 (\_ e1 e2 e3 -> f e1 e2 e3) arr1 arr2 arr3
  -- where
  --   comp = getComp arr1 <> getComp arr2 <> getComp arr3
  --   sz1 = size arr1
  {-# INLINE zipWith3 #-}
  
  -- | Just like `zipWith3`, except with an index aware function.
  izipWith3
    :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3)
    => (ix -> e1 -> e2 -> e3 -> e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> Array D ix e
  izipWith3 f arr1 arr2 arr3 =
    DArray
      (getComp arr1 <> getComp arr2 <> getComp arr3)
      ( SafeSz
          ( liftIndex2
              min
              (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))
              (coerce (size arr3))
          )
      )
      (PrefIndex $ \ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix))
  {-# INLINE izipWith3 #-}
  
  -- | Just like `zipWith`, except zip four arrays with a function.
  --
  -- @since 0.5.4
  zipWith4
    :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4)
    => (e1 -> e2 -> e3 -> e4 -> e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> Array r4 ix e4
    -> Array D ix e
  zipWith4 f arr1 arr2 arr3 arr4 =
    izipWith4 (\_ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4
  -- See note on zipWith
  --  | sz1 == size arr2 && sz1 == size arr3 && sz1 == size arr4
  --  , PrefIndexLinear gi1 <- unsafePrefIndex arr1
  --  , PrefIndexLinear gi2 <- unsafePrefIndex arr2
  --  , PrefIndexLinear gi3 <- unsafePrefIndex arr3
  --  , PrefIndexLinear gi4 <- unsafePrefIndex arr4 =
  --    makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i) (gi4 i))
  --  | otherwise = izipWith4 (\ _ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4
  --  where
  --    comp = getComp arr1 <> getComp arr2 <> getComp arr3 <> getComp arr4
  --    sz1 = size arr1
  {-# INLINE zipWith4 #-}
  
  -- | Just like `zipWith4`, except with an index aware function.
  --
  -- @since 0.5.4
  izipWith4
    :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4)
    => (ix -> e1 -> e2 -> e3 -> e4 -> e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> Array r4 ix e4
    -> Array D ix e
  izipWith4 f arr1 arr2 arr3 arr4 =
    makeArray
      (getComp arr1 <> getComp arr2 <> getComp arr3 <> getComp arr4)
      ( SafeSz
          ( liftIndex2
              min
              ( liftIndex2
                  min
                  (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))
                  (coerce (size arr3))
              )
              (coerce (size arr4))
          )
      )
      ( \ !ix ->
          f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix) (unsafeIndex arr4 ix)
      )
  {-# INLINE izipWith4 #-}
  
  -- | Similar to `zipWith`, except does it sequentially and using the `Applicative`. Note that
  -- resulting array has Manifest representation.
  --
  -- @since 0.3.0
  zipWithA
    :: (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix)
    => (e1 -> e2 -> f e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> f (Array r ix e)
  zipWithA f arr1 arr2
    | sz1 == size arr2
    , PrefIndexLinear gi1 <- unsafePrefIndex arr1
    , PrefIndexLinear gi2 <- unsafePrefIndex arr2 =
        setComp (getComp arr1 <> getComp arr2) <$> makeArrayLinearA sz1 (\ !i -> f (gi1 i) (gi2 i))
    | otherwise = izipWithA (const f) arr1 arr2
    where
      !sz1 = size arr1
  {-# INLINE zipWithA #-}
  
  -- | Similar to `zipWith`, except does it sequentially and using the `Applicative`. Note that
  -- resulting array has Manifest representation.
  --
  -- @since 0.3.0
  izipWithA
    :: (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix)
    => (ix -> e1 -> e2 -> f e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> f (Array r ix e)
  izipWithA f arr1 arr2 =
    setComp (getComp arr1 <> getComp arr2)
      <$> makeArrayA
        (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))))
        (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix))
  {-# INLINE izipWithA #-}
  
  -- | Same as `zipWithA`, but for three arrays.
  --
  -- @since 0.3.0
  zipWith3A
    :: (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix)
    => (e1 -> e2 -> e3 -> f e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> f (Array r ix e)
  zipWith3A f = izipWith3A (const f)
  {-# INLINE zipWith3A #-}
  
  -- | Same as `izipWithA`, but for three arrays.
  --
  -- @since 0.3.0
  izipWith3A
    :: (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix)
    => (ix -> e1 -> e2 -> e3 -> f e)
    -> Array r1 ix e1
    -> Array r2 ix e2
    -> Array r3 ix e3
    -> f (Array r ix e)
  izipWith3A f arr1 arr2 arr3 =
    setComp (getComp arr1 <> getComp arr2 <> getComp arr3)
      <$> makeArrayA sz (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix))
    where
      sz =
        SafeSz $
          liftIndex2 min (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) (coerce (size arr3))
  {-# INLINE izipWith3A #-}
  
  --------------------------------------------------------------------------------
  -- traverse --------------------------------------------------------------------
  --------------------------------------------------------------------------------
  
  -- | Traverse with an `Applicative` action over an array sequentially.
  --
  -- /Note/ - using `traversePrim` instead will always be significantly faster, roughly
  -- about 30 times faster in practice.
  --
  -- @since 0.2.6
  traverseA
    :: forall r ix e r' a f
     . (Source r' a, Manifest r e, Index ix, Applicative f)
    => (a -> f e)
    -> Array r' ix a
    -> f (Array r ix e)
  traverseA f arr =
    unsafeResize (size arr) . fromList (getComp arr) <$> traverse f (toList arr)
  {-# INLINE traverseA #-}
  
  -- | Traverse sequentially over a source array, while discarding the result.
  --
  -- @since 0.3.0
  traverseA_
    :: forall r ix e a f
     . (Index ix, Source r e, Applicative f)
    => (e -> f a)
    -> Array r ix e
    -> f ()
  traverseA_ f arr =
    case unsafePrefIndex arr of
      PrefIndex gix -> iterA_ zeroIndex (unSz sz) oneIndex (<) (f . gix)
      PrefIndexLinear gi -> loopA_ 0 (< totalElem sz) (+ 1) (f . gi)
    where
      sz = size arr
  {-# INLINE traverseA_ #-}
  
  -- | Sequence actions in a source array.
  --
  -- @since 0.3.0
  sequenceA
    :: forall r ix e r' f
     . (Source r' (f e), Manifest r e, Index ix, Applicative f)
    => Array r' ix (f e)
    -> f (Array r ix e)
  sequenceA = traverseA id
  {-# INLINE sequenceA #-}
  
  -- | Sequence actions in a source array, while discarding the result.
  --
  -- @since 0.3.0
  sequenceA_
    :: forall r ix e f
     . (Index ix, Source r (f e), Applicative f)
    => Array r ix (f e)
    -> f ()
  sequenceA_ = traverseA_ id
  {-# INLINE sequenceA_ #-}
  
  -- | Traverse with an `Applicative` index aware action over an array sequentially.
  --
  -- @since 0.2.6
  itraverseA
    :: forall r ix e r' a f
     . (Source r' a, Manifest r e, Index ix, Applicative f)
    => (ix -> a -> f e)
    -> Array r' ix a
    -> f (Array r ix e)
  itraverseA f arr =
    setComp (getComp arr) <$> makeArrayA (size arr) (\ !ix -> f ix (unsafeIndex arr ix))
  {-# INLINE itraverseA #-}
  
  -- | Traverse with an `Applicative` index aware action over an array sequentially.
  --
  -- @since 0.2.6
  itraverseA_
    :: forall r ix e a f
     . (Source r a, Index ix, Applicative f)
    => (ix -> a -> f e)
    -> Array r ix a
    -> f ()
  itraverseA_ f arr =
    case unsafePrefIndex arr of
      PrefIndex gix ->
        iterA_ zeroIndex (unSz sz) oneIndex (<) (\ !ix -> f ix (gix ix))
      PrefIndexLinear gi ->
        iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix -> f ix (gi i)
    where
      sz = size arr
  {-# INLINE itraverseA_ #-}
  
  -- | Traverse sequentially within `PrimMonad` over an array with an action.
  --
  -- @since 0.3.0
  traversePrim
    :: forall r ix b r' a m
     . (Source r' a, Manifest r b, Index ix, PrimMonad m)
    => (a -> m b)
    -> Array r' ix a
    -> m (Array r ix b)
  traversePrim f arr = do
    let sz = size arr
    marr <- unsafeNew sz
    case unsafePrefIndex arr of
      PrefIndex gix ->
        iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix ->
          f (gix ix) >>= unsafeLinearWrite marr i
      PrefIndexLinear gi ->
        loopA_ 0 (< totalElem sz) (+ 1) $ \i ->
          f (gi i) >>= unsafeLinearWrite marr i
    unsafeFreeze (getComp arr) marr
  {-# INLINE traversePrim #-}
  
  -- | Same as `traversePrim`, but traverse with index aware action.
  --
  -- @since 0.3.0
  itraversePrim
    :: forall r ix b r' a m
     . (Source r' a, Manifest r b, Index ix, PrimMonad m)
    => (ix -> a -> m b)
    -> Array r' ix a
    -> m (Array r ix b)
  itraversePrim f arr = do
    let sz = size arr
    marr <- unsafeNew sz
    case unsafePrefIndex arr of
      PrefIndex gix ->
        iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix ->
          f ix (gix ix) >>= unsafeLinearWrite marr i
      PrefIndexLinear gi ->
        iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix ->
          f ix (gi i) >>= unsafeLinearWrite marr i
    unsafeFreeze (getComp arr) marr
  {-# INLINE itraversePrim #-}
  
  --------------------------------------------------------------------------------
  -- mapM ------------------------------------------------------------------------
  --------------------------------------------------------------------------------
  
  -- | Map a monadic action over an array sequentially.
  --
  -- @since 0.2.6
  mapM
    :: forall r ix b r' a m
     . (Source r' a, Manifest r b, Index ix, Monad m)
    => (a -> m b)
    -- ^ Mapping action
    -> Array r' ix a
    -- ^ Source array
    -> m (Array r ix b)
  mapM = traverseA
  {-# INLINE mapM #-}
  
  -- | Same as `mapM` except with arguments flipped.
  --
  -- @since 0.2.6
  forM
    :: forall r ix b r' a m
     . (Source r' a, Manifest r b, Index ix, Monad m)
    => Array r' ix a
    -> (a -> m b)
    -> m (Array r ix b)
  forM = flip traverseA
  {-# INLINE forM #-}
  
  -- | Map an index aware monadic action over an array sequentially.
  --
  -- @since 0.2.6
  imapM
    :: forall r ix b r' a m
     . (Source r' a, Manifest r b, Index ix, Monad m)
    => (ix -> a -> m b)
    -> Array r' ix a
    -> m (Array r ix b)
  imapM = itraverseA
  {-# INLINE imapM #-}
  
  -- | Same as `forM`, except with an index aware action.
  --
  -- @since 0.5.1
  iforM
    :: forall r ix b r' a m
     . (Source r' a, Manifest r b, Index ix, Monad m)
    => Array r' ix a
    -> (ix -> a -> m b)
    -> m (Array r ix b)
  iforM = flip itraverseA
  {-# INLINE iforM #-}
  
  -- | Map a monadic function over an array sequentially, while discarding the result.
  --
  -- ==== __Examples__
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> rangeStepM Par (Ix1 10) 12 60 >>= A.mapM_ print
  -- 10
  -- 22
  -- 34
  -- 46
  -- 58
  --
  -- @since 0.1.0
  mapM_ :: (Source r a, Index ix, Monad m) => (a -> m b) -> Array r ix a -> m ()
  mapM_ = traverseA_
  {-# INLINE mapM_ #-}
  
  -- | Just like `mapM_`, except with flipped arguments.
  --
  -- ==== __Examples__
  --
  -- Here is a common way of iterating N times using a for loop in an imperative
  -- language with mutation being an obvious side effect:
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> import Data.IORef
  -- >>> ref <- newIORef 0 :: IO (IORef Int)
  -- >>> A.forM_ (range Seq (Ix1 0) 1000) $ \ i -> modifyIORef' ref (+i)
  -- >>> readIORef ref
  -- 499500
  forM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (a -> m b) -> m ()
  forM_ = flip traverseA_
  {-# INLINE forM_ #-}
  
  -- | Map a monadic index aware function over an array sequentially, while discarding the result.
  --
  -- ==== __Examples__
  --
  -- >>> import Data.Massiv.Array
  -- >>> imapM_ (curry print) $ range Seq (Ix1 10) 15
  -- (0,10)
  -- (1,11)
  -- (2,12)
  -- (3,13)
  -- (4,14)
  --
  -- @since 0.1.0
  imapM_ :: (Index ix, Source r a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m ()
  imapM_ = itraverseA_
  {-# INLINE imapM_ #-}
  
  -- | Just like `imapM_`, except with flipped arguments.
  iforM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (ix -> a -> m b) -> m ()
  iforM_ = flip itraverseA_
  {-# INLINE iforM_ #-}
  
  -- | Map an `IO` action over an `Array`. Underlying computation strategy is respected and will be
  -- parallelized when requested. Unfortunately no fusion is possible and new array will be create
  -- upon each call.
  --
  -- @since 0.2.6
  mapIO
    :: forall r ix b r' a m
     . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m)
    => (a -> m b)
    -> Array r' ix a
    -> m (Array r ix b)
  mapIO action = imapIO (const action)
  {-# INLINE mapIO #-}
  
  -- | Similar to `mapIO`, but ignores the result of mapping action and does not
  -- create a resulting array, therefore it is faster. Use this instead of `mapIO`
  -- when result is irrelevant. Most importantly it will follow the iteration
  -- logic outlined by the supplied array.
  --
  -- @since 0.2.6
  mapIO_
    :: forall r ix e a m
     . (Load r ix e, MonadUnliftIO m)
    => (e -> m a)
    -> Array r ix e
    -> m ()
  mapIO_ action arr =
    withRunInIO $ \run ->
      withMassivScheduler_ (getComp arr) $ \scheduler ->
        iterArrayLinearM_ scheduler arr (\_ -> void . run . action)
  {-# INLINE mapIO_ #-}
  
  -- | Same as `mapIO_`, but map an index aware action instead.
  --
  -- @since 0.2.6
  imapIO_
    :: forall r ix e a m
     . (Load r ix e, MonadUnliftIO m)
    => (ix -> e -> m a)
    -> Array r ix e
    -> m ()
  imapIO_ action arr =
    withRunInIO $ \run ->
      withMassivScheduler_ (getComp arr) $ \scheduler ->
        let sz = outerSize arr
         in -- It is ok to use outerSize in context of DS and L. Former is 1-dim,
            -- so sz is never evaluated and for the latter outerSize has to be
            -- called regardless how this function is implemented.
            iterArrayLinearM_ scheduler arr (\i -> void . run . action (fromLinearIndex sz i))
  {-# INLINE imapIO_ #-}
  
  -- | Same as `mapIO` but map an index aware action instead. Respects computation strategy.
  --
  -- @since 0.2.6
  imapIO
    :: forall r ix b r' a m
     . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m)
    => (ix -> a -> m b)
    -> Array r' ix a
    -> m (Array r ix b)
  imapIO action arr = do
    let sz = size arr
    withRunInIO $ \run -> do
      marr <- unsafeNew sz
      withMassivScheduler_ (getComp arr) $ \scheduler ->
        iterArrayLinearM_ scheduler arr $ \ !i e ->
          run (action (fromLinearIndex sz i) e) >>= unsafeLinearWrite marr i
      unsafeFreeze (getComp arr) marr
  {-# INLINE imapIO #-}
  
  -- | Same as `mapIO` but with arguments flipped.
  --
  -- @since 0.2.6
  forIO
    :: forall r ix b r' a m
     . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m)
    => Array r' ix a
    -> (a -> m b)
    -> m (Array r ix b)
  forIO = flip mapIO
  {-# INLINE forIO #-}
  
  -- | Same as `imapIO`, but ignores the inner computation strategy and uses
  -- stateful workers during computation instead. Use
  -- `Control.Scheduler.initWorkerStates` for the `WorkerStates` initialization.
  --
  -- @since 0.3.4
  imapWS
    :: forall r ix b r' a s m
     . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m)
    => WorkerStates s
    -> (ix -> a -> s -> m b)
    -> Array r' ix a
    -> m (Array r ix b)
  imapWS states f arr = generateArrayWS states (size arr) (\ix s -> f ix (unsafeIndex arr ix) s)
  {-# INLINE imapWS #-}
  
  -- | Same as `imapWS`, but without the index.
  --
  -- @since 0.3.4
  mapWS
    :: forall r ix b r' a s m
     . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m)
    => WorkerStates s
    -> (a -> s -> m b)
    -> Array r' ix a
    -> m (Array r ix b)
  mapWS states f = imapWS states (\_ -> f)
  {-# INLINE mapWS #-}
  
  -- | Same as `imapWS`, but with source array and mapping action arguments flipped.
  --
  -- @since 0.3.4
  iforWS
    :: forall r ix b r' a s m
     . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m)
    => WorkerStates s
    -> Array r' ix a
    -> (ix -> a -> s -> m b)
    -> m (Array r ix b)
  iforWS states f arr = imapWS states arr f
  {-# INLINE iforWS #-}
  
  -- | Same as `iforWS`, but without the index.
  --
  -- @since 0.3.4
  forWS
    :: forall r ix b r' a s m
     . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m)
    => WorkerStates s
    -> Array r' ix a
    -> (a -> s -> m b)
    -> m (Array r ix b)
  forWS states arr f = imapWS states (\_ -> f) arr
  {-# INLINE forWS #-}
  
  -- | Same as `mapIO_` but with arguments flipped.
  --
  -- ==== __Example__
  --
  -- This is the same example as in `forM_`, with important difference that accumulator `ref` will be
  -- modified concurrently by as many threads as there are capabilities.
  --
  -- >>> import Data.Massiv.Array
  -- >>> import Data.IORef
  -- >>> ref <- newIORef 0 :: IO (IORef Int)
  -- >>> forIO_ (range Par (Ix1 0) 1000) $ \ i -> atomicModifyIORef' ref (\v -> (v+i, ()))
  -- >>> readIORef ref
  -- 499500
  --
  -- @since 0.2.6
  forIO_ :: (Load r ix e, MonadUnliftIO m) => Array r ix e -> (e -> m a) -> m ()
  forIO_ = flip mapIO_
  {-# INLINE forIO_ #-}
  
  -- | Same as `imapIO` but with arguments flipped.
  --
  -- @since 0.2.6
  iforIO
    :: forall r ix b r' a m
     . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m)
    => Array r' ix a
    -> (ix -> a -> m b)
    -> m (Array r ix b)
  iforIO = flip imapIO
  {-# INLINE iforIO #-}
  
  -- | Same as `imapIO_` but with arguments flipped.
  --
  -- @since 0.2.6
  iforIO_
    :: forall r ix e a m
     . (Load r ix e, MonadUnliftIO m)
    => Array r ix e
    -> (ix -> e -> m a)
    -> m ()
  iforIO_ = flip imapIO_
  {-# INLINE iforIO_ #-}
  
  iterArrayLinearM_
    :: forall r ix e m s
     . (Load r ix e, MonadPrimBase s m)
    => Scheduler s ()
    -> Array r ix e
    -- ^ Array that is being loaded
    -> (Int -> e -> m ())
    -- ^ Function that writes an element into target array
    -> m ()
  iterArrayLinearM_ scheduler arr f =
    stToPrim $ iterArrayLinearST_ scheduler arr (\i -> primToPrim . f i)
  {-# INLINE iterArrayLinearM_ #-}
  
  iterArrayLinearWithSetM_
    :: forall r ix e m s
     . (Load r ix e, MonadPrimBase s m)
    => Scheduler s ()
    -> Array r ix e
    -- ^ Array that is being loaded
    -> (Int -> e -> m ())
    -- ^ Function that writes an element into target array
    -> (Ix1 -> Sz1 -> e -> m ())
    -- ^ Function that efficiently sets a region of an array
    -- to the supplied value target array
    -> m ()
  iterArrayLinearWithSetM_ scheduler arr f set =
    stToPrim $
      iterArrayLinearWithSetST_ scheduler arr (\i -> primToPrim . f i) (\i n -> primToPrim . set i n)
  {-# INLINE iterArrayLinearWithSetM_ #-}
  
  iterArrayLinearWithStrideM_
    :: forall r ix e m s
     . (StrideLoad r ix e, MonadPrimBase s m)
    => Scheduler s ()
    -> Stride ix
    -- ^ Stride to use
    -> Sz ix
    -- ^ Size of the target array affected by the stride.
    -> Array r ix e
    -- ^ Array that is being loaded
    -> (Int -> e -> m ())
    -- ^ Function that writes an element into target array
    -> m ()
  iterArrayLinearWithStrideM_ scheduler stride sz arr f =
    stToPrim $ iterArrayLinearWithStrideST_ scheduler stride sz arr (\i -> primToPrim . f i)
  {-# INLINE iterArrayLinearWithStrideM_ #-}
  
  -- iterArrayM_ ::
  --      Scheduler s ()
  --   -> Array r ix e -- ^ Array that is being loaded
  --   -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array
  --   -> ST s ()
  -- iterArrayM_ scheduler arr uWrite
  
  -- Deprecated
  
  -- | Same as `imapM_`, but will use the supplied scheduler.
  --
  -- @since 0.3.1
  imapSchedulerM_
    :: (Index ix, Source r e, MonadPrimBase s m)
    => Scheduler s ()
    -> (ix -> e -> m a)
    -> Array r ix e
    -> m ()
  imapSchedulerM_ scheduler action arr = do
    let sz = size arr
    splitLinearlyWith_
      scheduler
      (totalElem sz)
      (unsafeLinearIndex arr)
      (\i -> void . action (fromLinearIndex sz i))
  {-# INLINE imapSchedulerM_ #-}
  
  -- | Same as `imapM_`, but will use the supplied scheduler.
  --
  -- @since 0.3.1
  iforSchedulerM_
    :: (Index ix, Source r e, MonadPrimBase s m)
    => Scheduler s ()
    -> Array r ix e
    -> (ix -> e -> m a)
    -> m ()
  iforSchedulerM_ scheduler arr action = imapSchedulerM_ scheduler action arr
  {-# INLINE iforSchedulerM_ #-}
  
  -- -- | Load an array into memory.
  -- --
  -- -- @since 0.3.0
  -- loadArrayM
  --   :: Scheduler s ()
  --   -> Array r ix e -- ^ Array that is being loaded
  --   -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array
  --   -> ST s ()
  -- loadArrayM scheduler arr uWrite =
  --   loadArrayWithSetM scheduler arr uWrite $ \offset sz e ->
  --     loopM_ offset (< (offset + unSz sz)) (+1) (`uWrite` e)
  -- {-# INLINE loadArrayM #-}
  
  -- -- | Load an array into memory, just like `loadArrayM`. Except it also accepts a
  -- -- function that is potentially optimized for setting many cells in a region to the same
  -- -- value
  -- --
  -- -- @since 0.5.8
  -- loadArrayWithSetM
  --   :: Scheduler s ()
  --   -> Array r ix e -- ^ Array that is being loaded
  --   -> (Ix1 -> e -> ST s ()) -- ^ Function that writes an element into target array
  --   -> (Ix1 -> Sz1 -> e -> ST s ()) -- ^ Function that efficiently sets a region of an array
  --                                   -- to the supplied value target array
  --   -> ST s ()
  -- loadArrayWithSetM scheduler arr uWrite _ = loadArrayM scheduler arr uWrite
  -- {-# INLINE loadArrayWithSetM #-}
  
  -- iterArrayLinearWithStrideST
  --   :: Scheduler s ()
  --   -> Stride ix -- ^ Stride to use
  --   -> Sz ix -- ^ Size of the target array affected by the stride.
  --   -> Array r ix e -- ^ Array that is being loaded
  --   -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array
  --   -> ST s ()
  