{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MonoLocalBinds #-}
-- |
-- 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 Data.Traversable (traverse)
import Data.Massiv.Array.Manifest.List
import Control.Monad (void)
import Control.Monad.Primitive
import Control.Scheduler
import Data.Coerce
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Mutable
import Data.Massiv.Array.Ops.Construct (makeArrayA, makeArrayLinearA)
import Data.Massiv.Core.Common
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 :: (e' -> e) -> Array r ix e' -> Array D ix e
map e' -> e
f = (e' -> e) -> Array D ix e' -> Array D ix e
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap e' -> e
f (Array D ix e' -> Array D ix e)
-> (Array r ix e' -> Array D ix e')
-> Array r ix e'
-> Array D ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e' -> Array D ix e'
forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
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 :: Array r1 ix e1 -> Array r2 ix e2 -> Array D ix (e1, e2)
zip = (e1 -> e2 -> (e1, e2))
-> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix (e1, e2)
forall ix r1 e1 r2 e2 e.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e
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 :: Array r1 ix e1
-> Array r2 ix e2 -> Array r3 ix e3 -> Array D ix (e1, e2, e3)
zip3 = (e1 -> e2 -> e3 -> (e1, e2, e3))
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> Array D ix (e1, e2, e3)
forall ix r1 e1 r2 e2 r3 e3 e.
(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 (,,)
{-# 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 :: Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> Array r4 ix e4
-> Array D ix (e1, e2, e3, e4)
zip4 = (e1 -> e2 -> e3 -> e4 -> (e1, e2, e3, e4))
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> Array r4 ix e4
-> Array D ix (e1, e2, e3, e4)
forall ix r1 e1 r2 e2 r3 e3 r4 e4 e.
(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 (,,,)
{-# 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 :: Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2)
unzip Array r ix (e1, e2)
arr = (((e1, e2) -> e1) -> Array r ix (e1, e2) -> Array D ix e1
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (e1, e2) -> e1
forall a b. (a, b) -> a
fst Array r ix (e1, e2)
arr, ((e1, e2) -> e2) -> Array r ix (e1, e2) -> Array D ix e2
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (e1, e2) -> e2
forall a b. (a, b) -> b
snd Array r ix (e1, e2)
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 :: Array r ix (e1, e2, e3)
-> (Array D ix e1, Array D ix e2, Array D ix e3)
unzip3 Array r ix (e1, e2, e3)
arr = (((e1, e2, e3) -> e1) -> Array r ix (e1, e2, e3) -> Array D ix e1
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (\ (e1
e, e2
_, e3
_) -> e1
e) Array r ix (e1, e2, e3)
arr, ((e1, e2, e3) -> e2) -> Array r ix (e1, e2, e3) -> Array D ix e2
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (\ (e1
_, e2
e, e3
_) -> e2
e) Array r ix (e1, e2, e3)
arr, ((e1, e2, e3) -> e3) -> Array r ix (e1, e2, e3) -> Array D ix e3
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (\ (e1
_, e2
_, e3
e) -> e3
e) Array r ix (e1, e2, e3)
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 :: Array r ix (e1, e2, e3, e4)
-> (Array D ix e1, Array D ix e2, Array D ix e3, Array D ix e4)
unzip4 Array r ix (e1, e2, e3, e4)
arr =
  ( ((e1, e2, e3, e4) -> e1)
-> Array r ix (e1, e2, e3, e4) -> Array D ix e1
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (\(e1
e, e2
_, e3
_, e4
_) -> e1
e) Array r ix (e1, e2, e3, e4)
arr
  , ((e1, e2, e3, e4) -> e2)
-> Array r ix (e1, e2, e3, e4) -> Array D ix e2
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (\(e1
_, e2
e, e3
_, e4
_) -> e2
e) Array r ix (e1, e2, e3, e4)
arr
  , ((e1, e2, e3, e4) -> e3)
-> Array r ix (e1, e2, e3, e4) -> Array D ix e3
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (\(e1
_, e2
_, e3
e, e4
_) -> e3
e) Array r ix (e1, e2, e3, e4)
arr
  , ((e1, e2, e3, e4) -> e4)
-> Array r ix (e1, e2, e3, e4) -> Array D ix e4
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (\(e1
_, e2
_, e3
_, e4
e) -> e4
e) Array r ix (e1, e2, e3, e4)
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 :: (e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e
zipWith e1 -> e2 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 = Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray Comp
comp Sz ix
sz PrefIndex ix e
prefIndex
  where
    sz :: Sz ix
sz = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
coerce (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (Sz ix -> ix
coerce (Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2)))
    comp :: Comp
comp = Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2
    prefIndex :: PrefIndex ix e
prefIndex = (ix -> e) -> PrefIndex ix e
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (\ix
ix -> e1 -> e2 -> e
f (Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
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 :: (ix -> e1 -> e2 -> e)
-> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e
izipWith ix -> e1 -> e2 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 =
  Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray
    (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2)
    (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
coerce (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (Sz ix -> ix
coerce (Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2))))
    ((ix -> e) -> PrefIndex ix e
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (\ix
ix -> ix -> e1 -> e2 -> e
f ix
ix (Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
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 :: (e1 -> e2 -> e3 -> e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> Array D ix e
zipWith3 e1 -> e2 -> e3 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 = (ix -> e1 -> e2 -> e3 -> e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> Array D ix e
forall ix r1 e1 r2 e2 r3 e3 e.
(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 (\ix
_ e1
e1 e2
e2 e3
e3 -> e1 -> e2 -> e3 -> e
f e1
e1 e2
e2 e3
e3) Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
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 :: (ix -> e1 -> e2 -> e3 -> e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> Array D ix e
izipWith3 ix -> e1 -> e2 -> e3 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 =
  Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray
    (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r3 ix e3 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r3 ix e3
arr3)
    (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz
       ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2
          Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
          ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
coerce (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (Sz ix -> ix
coerce (Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2)))
          (Sz ix -> ix
coerce (Array r3 ix e3 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r3 ix e3
arr3))))
    ((ix -> e) -> PrefIndex ix e
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex ((ix -> e) -> PrefIndex ix e) -> (ix -> e) -> PrefIndex ix e
forall a b. (a -> b) -> a -> b
$ \ !ix
ix -> ix -> e1 -> e2 -> e3 -> e
f ix
ix (Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (Array r3 ix e3 -> ix -> e3
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r3 ix e3
arr3 ix
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 :: (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 e1 -> e2 -> e3 -> e4 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 Array r4 ix e4
arr4 =
  (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
forall ix r1 e1 r2 e2 r3 e3 r4 e4 e.
(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 (\ ix
_ e1
e1 e2
e2 e3
e3 e4
e4 -> e1 -> e2 -> e3 -> e4 -> e
f e1
e1 e2
e2 e3
e3 e4
e4) Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 Array r4 ix e4
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 :: (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 ix -> e1 -> e2 -> e3 -> e4 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 Array r4 ix e4
arr4 =
  Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray
    (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r3 ix e3 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r3 ix e3
arr3 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r4 ix e4 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r4 ix e4
arr4)
    (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz
       ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2
          Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
          ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2
             Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
             ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
coerce (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (Sz ix -> ix
coerce (Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2)))
             (Sz ix -> ix
coerce (Array r3 ix e3 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r3 ix e3
arr3)))
          (Sz ix -> ix
coerce (Array r4 ix e4 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r4 ix e4
arr4))))
    (\ !ix
ix ->
       ix -> e1 -> e2 -> e3 -> e4 -> e
f ix
ix (Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (Array r3 ix e3 -> ix -> e3
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r3 ix e3
arr3 ix
ix) (Array r4 ix e4 -> ix -> e4
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r4 ix e4
arr4 ix
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 :: (e1 -> e2 -> f e)
-> Array r1 ix e1 -> Array r2 ix e2 -> f (Array r ix e)
zipWithA e1 -> e2 -> f e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2
  | Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2
  , PrefIndexLinear Int -> e1
gi1 <- Array r1 ix e1 -> PrefIndex ix e1
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r1 ix e1
arr1
  , PrefIndexLinear Int -> e2
gi2 <- Array r2 ix e2 -> PrefIndex ix e2
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r2 ix e2
arr2 =
    Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2) (Array r ix e -> Array r ix e)
-> f (Array r ix e) -> f (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> (Int -> f e) -> f (Array r ix e)
forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (Int -> f e) -> f (Array r ix e)
makeArrayLinearA Sz ix
sz1 (\ !Int
i -> e1 -> e2 -> f e
f (Int -> e1
gi1 Int
i) (Int -> e2
gi2 Int
i))
  | Bool
otherwise = (ix -> e1 -> e2 -> f e)
-> Array r1 ix e1 -> Array r2 ix e2 -> f (Array r ix e)
forall r1 e1 r2 e2 (f :: * -> *) r e ix.
(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 ((e1 -> e2 -> f e) -> ix -> e1 -> e2 -> f e
forall a b. a -> b -> a
const e1 -> e2 -> f e
f) Array r1 ix e1
arr1 Array r2 ix e2
arr2
  where
    !sz1 :: Sz ix
sz1 = Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1
{-# INLINE zipWithA #-}

-- | Similar to `zipWith`, except does it sequentiall 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 :: (ix -> e1 -> e2 -> f e)
-> Array r1 ix e1 -> Array r2 ix e2 -> f (Array r ix e)
izipWithA ix -> e1 -> e2 -> f e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 =
  Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2) (Array r ix e -> Array r ix e)
-> f (Array r ix e) -> f (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  Sz ix -> (ix -> f e) -> f (Array r ix e)
forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA
    (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
coerce (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (Sz ix -> ix
coerce (Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2))))
    (\ !ix
ix -> ix -> e1 -> e2 -> f e
f ix
ix (Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
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 :: (e1 -> e2 -> e3 -> f e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> f (Array r ix e)
zipWith3A e1 -> e2 -> e3 -> f e
f = (ix -> e1 -> e2 -> e3 -> f e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> f (Array r ix e)
forall r1 e1 r2 e2 r3 e3 (f :: * -> *) r e ix.
(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 ((e1 -> e2 -> e3 -> f e) -> ix -> e1 -> e2 -> e3 -> f e
forall a b. a -> b -> a
const e1 -> e2 -> e3 -> f e
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 :: (ix -> e1 -> e2 -> e3 -> f e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> f (Array r ix e)
izipWith3A ix -> e1 -> e2 -> e3 -> f e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 =
  Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r3 ix e3 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r3 ix e3
arr3) (Array r ix e -> Array r ix e)
-> f (Array r ix e) -> f (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  Sz ix -> (ix -> f e) -> f (Array r ix e)
forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA Sz ix
sz (\ !ix
ix -> ix -> e1 -> e2 -> e3 -> f e
f ix
ix (Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (Array r3 ix e3 -> ix -> e3
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r3 ix e3
arr3 ix
ix))
  where
    sz :: Sz ix
sz =
      ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> Sz ix) -> ix -> Sz ix
forall a b. (a -> b) -> a -> b
$
      (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
coerce (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (Sz ix -> ix
coerce (Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2))) (Sz ix -> ix
coerce (Array r3 ix e3 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r3 ix e3
arr3))
{-# INLINE izipWith3A #-}


--------------------------------------------------------------------------------
-- traverse --------------------------------------------------------------------
--------------------------------------------------------------------------------

-- | Traverse with an `Applicative` action over an array sequentially.
--
-- /Note/ - using `traversePrim` will always be faster, althought not always possible.
--
-- @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 :: (a -> f e) -> Array r' ix a -> f (Array r ix e)
traverseA a -> f e
f Array r' ix a
arr =
  Sz ix -> Array r Int e -> Array r ix e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize (Array r' ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr) (Array r Int e -> Array r ix e)
-> ([e] -> Array r Int e) -> [e] -> Array r ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> [e] -> Array r Int e
forall r e. Manifest r e => Comp -> [e] -> Vector r e
fromList (Array r' ix a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) ([e] -> Array r ix e) -> f [e] -> f (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f e) -> [a] -> f [e]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f e
f (Array r' ix a -> [a]
forall ix r e. (Index ix, Source r e) => Array r ix e -> [e]
toList Array r' ix a
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_ :: (e -> f a) -> Array r ix e -> f ()
traverseA_ e -> f a
f Array r ix e
arr =
  case Array r ix e -> PrefIndex ix e
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
    PrefIndex ix -> e
gix -> ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ ix
forall ix. Index ix => ix
zeroIndex (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) ix
forall ix. Index ix => ix
oneIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) (e -> f a
f (e -> f a) -> (ix -> e) -> ix -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
gix)
    PrefIndexLinear Int -> e
gi -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (e -> f a
f (e -> f a) -> (Int -> e) -> Int -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
gi)
  where
    sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
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 :: Array r' ix (f e) -> f (Array r ix e)
sequenceA = (f e -> f e) -> Array r' ix (f e) -> f (Array r ix e)
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 e -> f e
forall a. a -> a
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_ :: Array r ix (f e) -> f ()
sequenceA_ = (f e -> f e) -> Array r ix (f e) -> f ()
forall r ix e a (f :: * -> *).
(Index ix, Source r e, Applicative f) =>
(e -> f a) -> Array r ix e -> f ()
traverseA_ f e -> f e
forall a. a -> a
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 :: (ix -> a -> f e) -> Array r' ix a -> f (Array r ix e)
itraverseA ix -> a -> f e
f Array r' ix a
arr =
  Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (Array r' ix a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) (Array r ix e -> Array r ix e)
-> f (Array r ix e) -> f (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> (ix -> f e) -> f (Array r ix e)
forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA (Array r' ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr) (\ !ix
ix -> ix -> a -> f e
f ix
ix (Array r' ix a -> ix -> a
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r' ix a
arr ix
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_ :: (ix -> a -> f e) -> Array r ix a -> f ()
itraverseA_ ix -> a -> f e
f Array r ix a
arr =
  case Array r ix a -> PrefIndex ix a
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix a
arr of
    PrefIndex ix -> a
gix ->
      ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f e) -> f ()
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ ix
forall ix. Index ix => ix
zeroIndex (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) ix
forall ix. Index ix => ix
oneIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) (\ !ix
ix -> ix -> a -> f e
f ix
ix (ix -> a
gix ix
ix))
    PrefIndexLinear Int -> a
gi ->
      RowMajor
-> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f e) -> f ()
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ RowMajor
defRowMajor Int
0 Sz ix
sz ix
forall ix. Index ix => ix
zeroIndex Stride ix
forall ix. Index ix => Stride ix
oneStride ((Int -> ix -> f e) -> f ()) -> (Int -> ix -> f e) -> f ()
forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix -> ix -> a -> f e
f ix
ix (Int -> a
gi Int
i)
  where
    sz :: Sz ix
sz = Array r ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix a
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 :: (a -> m b) -> Array r' ix a -> m (Array r ix b)
traversePrim a -> m b
f Array r' ix a
arr = do
  let sz :: Sz ix
sz = Array r' ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr
  MArray (PrimState m) r ix b
marr <- Sz ix -> m (MArray (PrimState m) r ix b)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  case Array r' ix a -> PrefIndex ix a
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r' ix a
arr of
    PrefIndex ix -> a
gix ->
      RowMajor
-> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> m ()) -> m ()
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ RowMajor
defRowMajor Int
0 Sz ix
sz ix
forall ix. Index ix => ix
zeroIndex Stride ix
forall ix. Index ix => Stride ix
oneStride ((Int -> ix -> m ()) -> m ()) -> (Int -> ix -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix ->
        a -> m b
f (ix -> a
gix ix
ix) m b -> (b -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState m) r ix b -> Int -> b -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix b
marr Int
i
    PrefIndexLinear Int -> a
gi ->
      Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
        a -> m b
f (Int -> a
gi Int
i) m b -> (b -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState m) r ix b -> Int -> b -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix b
marr Int
i
  Comp -> MArray (PrimState m) r ix b -> m (Array r ix b)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) MArray (PrimState m) r ix b
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 :: (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
itraversePrim ix -> a -> m b
f Array r' ix a
arr = do
  let sz :: Sz ix
sz = Array r' ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr
  MArray (PrimState m) r ix b
marr <- Sz ix -> m (MArray (PrimState m) r ix b)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  case Array r' ix a -> PrefIndex ix a
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r' ix a
arr of
    PrefIndex ix -> a
gix ->
      RowMajor
-> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> m ()) -> m ()
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ RowMajor
defRowMajor Int
0 Sz ix
sz ix
forall ix. Index ix => ix
zeroIndex Stride ix
forall ix. Index ix => Stride ix
oneStride ((Int -> ix -> m ()) -> m ()) -> (Int -> ix -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix ->
        ix -> a -> m b
f ix
ix (ix -> a
gix ix
ix) m b -> (b -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState m) r ix b -> Int -> b -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix b
marr Int
i
    PrefIndexLinear Int -> a
gi ->
      RowMajor
-> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> m ()) -> m ()
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ RowMajor
defRowMajor Int
0 Sz ix
sz ix
forall ix. Index ix => ix
zeroIndex Stride ix
forall ix. Index ix => Stride ix
oneStride ((Int -> ix -> m ()) -> m ()) -> (Int -> ix -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix ->
        ix -> a -> m b
f ix
ix (Int -> a
gi Int
i) m b -> (b -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState m) r ix b -> Int -> b -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix b
marr Int
i
  Comp -> MArray (PrimState m) r ix b -> m (Array r ix b)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) MArray (PrimState m) r ix b
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 :: (a -> m b) -> Array r' ix a -> m (Array r ix b)
mapM = (a -> m b) -> Array r' ix a -> m (Array r ix b)
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
{-# 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 :: Array r' ix a -> (a -> m b) -> m (Array r ix b)
forM = ((a -> m b) -> Array r' ix a -> m (Array r ix b))
-> Array r' ix a -> (a -> m b) -> m (Array r ix b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> m b) -> Array r' ix a -> m (Array r ix b)
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
{-# 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 :: (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
imapM = (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
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
{-# 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 :: Array r' ix a -> (ix -> a -> m b) -> m (Array r ix b)
iforM = ((ix -> a -> m b) -> Array r' ix a -> m (Array r ix b))
-> Array r' ix a -> (ix -> a -> m b) -> m (Array r ix b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
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
{-# 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_ :: (a -> m b) -> Array r ix a -> m ()
mapM_ = (a -> m b) -> Array r ix a -> m ()
forall r ix e a (f :: * -> *).
(Index ix, Source r e, Applicative f) =>
(e -> f a) -> Array r ix e -> f ()
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_ :: Array r ix a -> (a -> m b) -> m ()
forM_ = ((a -> m b) -> Array r ix a -> m ())
-> Array r ix a -> (a -> m b) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> m b) -> Array r ix a -> m ()
forall r ix e a (f :: * -> *).
(Index ix, Source r e, Applicative f) =>
(e -> f a) -> Array r ix e -> f ()
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_ :: (ix -> a -> m b) -> Array r ix a -> m ()
imapM_ = (ix -> a -> m b) -> Array r ix a -> m ()
forall r ix e a (f :: * -> *).
(Source r a, Index ix, Applicative f) =>
(ix -> a -> f e) -> Array r ix a -> f ()
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_ :: Array r ix a -> (ix -> a -> m b) -> m ()
iforM_ = ((ix -> a -> m b) -> Array r ix a -> m ())
-> Array r ix a -> (ix -> a -> m b) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (ix -> a -> m b) -> Array r ix a -> m ()
forall r ix e a (f :: * -> *).
(Source r a, Index ix, Applicative f) =>
(ix -> a -> f e) -> Array r ix a -> f ()
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 :: (a -> m b) -> Array r' ix a -> m (Array r ix b)
mapIO a -> m b
action = (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
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 ((a -> m b) -> ix -> a -> m b
forall a b. a -> b -> a
const a -> m b
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_ :: (e -> m a) -> Array r ix e -> m ()
mapIO_ e -> m a
action Array r ix e
arr =
  ((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 ->
    Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
      Scheduler RealWorld ()
-> Array r ix e -> (Int -> e -> IO ()) -> IO ()
forall r ix e (m :: * -> *) s.
(Load r ix e, MonadPrimBase s m) =>
Scheduler s () -> Array r ix e -> (Int -> e -> m ()) -> m ()
iterArrayLinearM_ Scheduler RealWorld ()
scheduler Array r ix e
arr (\Int
_ -> IO a -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO a -> IO ()) -> (e -> IO a) -> e -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> IO a
forall a. m a -> IO a
run (m a -> IO a) -> (e -> m a) -> e -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> m a
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_ :: (ix -> e -> m a) -> Array r ix e -> m ()
imapIO_ ix -> e -> m a
action Array r ix e
arr =
  ((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 ->
    Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
      let sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr
          -- 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.
       in Scheduler RealWorld ()
-> Array r ix e -> (Int -> e -> IO ()) -> IO ()
forall r ix e (m :: * -> *) s.
(Load r ix e, MonadPrimBase s m) =>
Scheduler s () -> Array r ix e -> (Int -> e -> m ()) -> m ()
iterArrayLinearM_ Scheduler RealWorld ()
scheduler Array r ix e
arr (\Int
i -> IO a -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO a -> IO ()) -> (e -> IO a) -> e -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> IO a
forall a. m a -> IO a
run (m a -> IO a) -> (e -> m a) -> e -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e -> m a
action (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
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 :: (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
imapIO ix -> a -> m b
action Array r' ix a
arr = do
  let sz :: Sz ix
sz = Array r' ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr
  ((forall a. m a -> IO a) -> IO (Array r ix b)) -> m (Array r ix b)
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO (Array r ix b))
 -> m (Array r ix b))
-> ((forall a. m a -> IO a) -> IO (Array r ix b))
-> m (Array r ix b)
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> do
    MArray RealWorld r ix b
marr <- Sz ix -> IO (MArray (PrimState IO) r ix b)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
    Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (Array r' ix a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
      Scheduler RealWorld ()
-> Array r' ix a -> (Int -> a -> IO ()) -> IO ()
forall r ix e (m :: * -> *) s.
(Load r ix e, MonadPrimBase s m) =>
Scheduler s () -> Array r ix e -> (Int -> e -> m ()) -> m ()
iterArrayLinearM_ Scheduler RealWorld ()
scheduler Array r' ix a
arr ((Int -> a -> IO ()) -> IO ()) -> (Int -> a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i a
e ->
        m b -> IO b
forall a. m a -> IO a
run (ix -> a -> m b
action (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i) a
e) IO b -> (b -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState IO) r ix b -> Int -> b -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r ix b
MArray (PrimState IO) r ix b
marr Int
i
    Comp -> MArray (PrimState IO) r ix b -> IO (Array r ix b)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) MArray RealWorld r ix b
MArray (PrimState IO) r ix b
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 :: Array r' ix a -> (a -> m b) -> m (Array r ix b)
forIO = ((a -> m b) -> Array r' ix a -> m (Array r ix b))
-> Array r' ix a -> (a -> m b) -> m (Array r ix b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> m b) -> Array r' ix a -> m (Array r ix b)
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
{-# 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 :: WorkerStates s
-> (ix -> a -> s -> m b) -> Array r' ix a -> m (Array r ix b)
imapWS WorkerStates s
states ix -> a -> s -> m b
f Array r' ix a
arr = WorkerStates s -> Sz ix -> (ix -> s -> m b) -> m (Array r ix b)
forall r ix e s (m :: * -> *).
(Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) =>
WorkerStates s -> Sz ix -> (ix -> s -> m e) -> m (Array r ix e)
generateArrayWS WorkerStates s
states (Array r' ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr) (\ix
ix s
s -> ix -> a -> s -> m b
f ix
ix (Array r' ix a -> ix -> a
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r' ix a
arr ix
ix) s
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 :: WorkerStates s
-> (a -> s -> m b) -> Array r' ix a -> m (Array r ix b)
mapWS WorkerStates s
states a -> s -> m b
f = WorkerStates s
-> (ix -> a -> s -> m b) -> Array r' ix a -> m (Array r ix b)
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 WorkerStates s
states (\ ix
_ -> a -> s -> m b
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 :: WorkerStates s
-> Array r' ix a -> (ix -> a -> s -> m b) -> m (Array r ix b)
iforWS WorkerStates s
states Array r' ix a
f ix -> a -> s -> m b
arr = WorkerStates s
-> (ix -> a -> s -> m b) -> Array r' ix a -> m (Array r ix b)
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 WorkerStates s
states ix -> a -> s -> m b
arr Array r' ix a
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 :: WorkerStates s
-> Array r' ix a -> (a -> s -> m b) -> m (Array r ix b)
forWS WorkerStates s
states Array r' ix a
arr a -> s -> m b
f = WorkerStates s
-> (ix -> a -> s -> m b) -> Array r' ix a -> m (Array r ix b)
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 WorkerStates s
states (\ ix
_ -> a -> s -> m b
f) Array r' ix a
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_ :: Array r ix e -> (e -> m a) -> m ()
forIO_ = ((e -> m a) -> Array r ix e -> m ())
-> Array r ix e -> (e -> m a) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (e -> m a) -> Array r ix e -> m ()
forall r ix e a (m :: * -> *).
(Load r ix e, MonadUnliftIO m) =>
(e -> m a) -> Array r ix e -> m ()
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 :: Array r' ix a -> (ix -> a -> m b) -> m (Array r ix b)
iforIO = ((ix -> a -> m b) -> Array r' ix a -> m (Array r ix b))
-> Array r' ix a -> (ix -> a -> m b) -> m (Array r ix b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
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
{-# 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_ :: Array r ix e -> (ix -> e -> m a) -> m ()
iforIO_ = ((ix -> e -> m a) -> Array r ix e -> m ())
-> Array r ix e -> (ix -> e -> m a) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (ix -> e -> m a) -> Array r ix e -> m ()
forall r ix e a (m :: * -> *).
(Load r ix e, MonadUnliftIO m) =>
(ix -> e -> m a) -> Array r ix e -> m ()
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 s () -> Array r ix e -> (Int -> e -> m ()) -> m ()
iterArrayLinearM_ Scheduler s ()
scheduler Array r ix e
arr Int -> e -> m ()
f =
  ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array r ix e
arr (\Int
i -> m () -> ST s ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim (m () -> ST s ()) -> (e -> m ()) -> e -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e -> m ()
f Int
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 s ()
-> Array r ix e
-> (Int -> e -> m ())
-> (Int -> Sz1 -> e -> m ())
-> m ()
iterArrayLinearWithSetM_ Scheduler s ()
scheduler Array r ix e
arr Int -> e -> m ()
f Int -> Sz1 -> e -> m ()
set =
  ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$
  Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler s ()
scheduler Array r ix e
arr (\Int
i -> m () -> ST s ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim (m () -> ST s ()) -> (e -> m ()) -> e -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e -> m ()
f Int
i) (\Int
i Sz1
n -> m () -> ST s ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim (m () -> ST s ()) -> (e -> m ()) -> e -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Sz1 -> e -> m ()
set Int
i Sz1
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 s ()
-> Stride ix -> Sz ix -> Array r ix e -> (Int -> e -> m ()) -> m ()
iterArrayLinearWithStrideM_ Scheduler s ()
scheduler Stride ix
stride Sz ix
sz Array r ix e
arr Int -> e -> m ()
f =
  ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ Scheduler s ()
-> Stride ix
-> Sz ix
-> Array r ix e
-> (Int -> e -> ST s ())
-> ST s ()
forall r ix e s.
StrideLoad r ix e =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array r ix e
-> (Int -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride ix
stride Sz ix
sz Array r ix e
arr (\Int
i -> m () -> ST s ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim (m () -> ST s ()) -> (e -> m ()) -> e -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e -> m ()
f Int
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 s () -> (ix -> e -> m a) -> Array r ix e -> m ()
imapSchedulerM_ Scheduler s ()
scheduler ix -> e -> m a
action Array r ix e
arr = do
  let sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
  Scheduler s () -> Int -> (Int -> e) -> (Int -> e -> m ()) -> m ()
forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_
    Scheduler s ()
scheduler
    (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
    (Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr)
    (\Int
i -> m a -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m a -> m ()) -> (e -> m a) -> e -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e -> m a
action (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
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 s () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler s ()
scheduler Array r ix e
arr ix -> e -> m a
action = Scheduler s () -> (ix -> e -> m a) -> Array r ix e -> m ()
forall ix r e s (m :: * -> *) a.
(Index ix, Source r e, MonadPrimBase s m) =>
Scheduler s () -> (ix -> e -> m a) -> Array r ix e -> m ()
imapSchedulerM_ Scheduler s ()
scheduler ix -> e -> m a
action Array r ix e
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 ()