{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Data.Massiv.Array.Ops.Map
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- 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_
  -- ** Zipping
  , zip
  , zip3
  , zip4
  , unzip
  , unzip3
  , unzip4
  , zipWith
  , zipWith3
  , zipWith4
  , izipWith
  , izipWith3
  , izipWith4
  , liftArray2
  -- *** Applicative
  , zipWithA
  , izipWithA
  , zipWith3A
  , izipWith3A
  ) where

import Control.Monad (void)
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
map :: Source r ix 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 = (ix -> e' -> e) -> Array r ix e' -> Array D ix e
forall r ix e' e.
Source r ix e' =>
(ix -> e' -> e) -> Array r ix e' -> Array D ix e
imap ((e' -> e) -> ix -> e' -> e
forall a b. a -> b -> a
const e' -> e
f)
{-# INLINE map #-}


--------------------------------------------------------------------------------
-- zip -------------------------------------------------------------------------
--------------------------------------------------------------------------------

-- | Zip two arrays
zip :: (Source r1 ix e1, Source r2 ix 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 r1 ix e1 r2 e2 e.
(Source r1 ix e1, Source r2 ix e2) =>
(e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e
zipWith (,)
{-# INLINE zip #-}

-- | Zip three arrays
zip3 :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix 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 r1 ix e1 r2 e2 r3 e3 e.
(Source r1 ix e1, Source r2 ix e2, Source r3 ix 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 ::
     (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Source r4 ix 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 r1 ix e1 r2 e2 r3 e3 r4 e4 e.
(Source r1 ix e1, Source r2 ix e2, Source r3 ix e3,
 Source r4 ix 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
unzip :: Source r ix (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 r ix e' e.
Source r ix 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 r ix e' e.
Source r ix 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
unzip3 :: Source r ix (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 r ix e' e.
Source r ix 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 r ix e' e.
Source r ix 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 r ix e' e.
Source r ix 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 :: Source r ix (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 r ix e' e.
Source r ix 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 r ix e' e.
Source r ix 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 r ix e' e.
Source r ix 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 r ix e' e.
Source r ix 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 :: (Source r1 ix e1, Source r2 ix 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 = (ix -> e1 -> e2 -> e)
-> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e
forall r1 ix e1 r2 e2 e.
(Source r1 ix e1, Source r2 ix e2) =>
(ix -> e1 -> e2 -> e)
-> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e
izipWith (\ ix
_ e1
e1 e2
e2 -> e1 -> e2 -> e
f e1
e1 e2
e2)
{-# INLINE zipWith #-}


-- | Just like `zipWith`, except with an index aware function.
izipWith :: (Source r1 ix e1, Source r2 ix 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 -> (ix -> e) -> Array D ix e
forall ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray
    (Array r1 ix e1 -> Comp
forall r ix e. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix e2
arr2)))) ((ix -> e) -> Array D ix e) -> (ix -> e) -> Array D ix e
forall a b. (a -> b) -> a -> b
$ \ !ix
ix ->
    ix -> e1 -> e2 -> e
f ix
ix (Array r1 ix e1 -> ix -> e1
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r ix e. Source r ix e => 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 :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix 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 = (ix -> e1 -> e2 -> e3 -> e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array r3 ix e3
-> Array D ix e
forall r1 ix e1 r2 e2 r3 e3 e.
(Source r1 ix e1, Source r2 ix e2, Source r3 ix 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)
{-# INLINE zipWith3 #-}


-- | Just like `zipWith3`, except with an index aware function.
izipWith3
  :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix 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 -> (ix -> e) -> Array D ix e
forall ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray
    (Array r1 ix e1 -> Comp
forall r ix e. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => Array r ix e -> Sz ix
size Array r3 ix e3
arr3)))) ((ix -> e) -> Array D ix e) -> (ix -> e) -> Array D 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 ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (Array r3 ix e3 -> ix -> e3
forall r ix e. Source r ix e => 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 ::
     (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Source r4 ix 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 = (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 r1 ix e1 r2 e2 r3 e3 r4 e4 e.
(Source r1 ix e1, Source r2 ix e2, Source r3 ix e3,
 Source r4 ix 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)
{-# INLINE zipWith4 #-}


-- | Just like `zipWith4`, except with an index aware function.
--
-- @since 0.5.4
izipWith4
  :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Source r4 ix 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 ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray
    (Array r1 ix e1 -> Comp
forall r ix e. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => Array r ix e -> Sz ix
size Array r4 ix e4
arr4)))) ((ix -> e) -> Array D ix e) -> (ix -> e) -> Array D ix e
forall a b. (a -> b) -> a -> b
$ \ !ix
ix ->
    ix -> e1 -> e2 -> e3 -> e4 -> e
f ix
ix (Array r1 ix e1 -> ix -> e1
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (Array r3 ix e3 -> ix -> e3
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r3 ix e3
arr3 ix
ix) (Array r4 ix e4 -> ix -> e4
forall r ix e. Source r ix e => 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 Mutable representation.
--
-- @since 0.3.0
zipWithA ::
     (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e)
  => (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 = (ix -> e1 -> e2 -> f e)
-> Array r1 ix e1 -> Array r2 ix e2 -> f (Array r ix e)
forall r1 ix e1 r2 e2 (f :: * -> *) r e.
(Source r1 ix e1, Source r2 ix e2, Applicative f,
 Mutable r ix e) =>
(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)
{-# INLINE zipWithA #-}

-- | Similar to `zipWith`, except does it sequentiall and using the `Applicative`. Note that
-- resulting array has Mutable representation.
--
-- @since 0.3.0
izipWithA ::
     (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e)
  => (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.
Construct r ix e =>
Comp -> Array r ix e -> Array r ix e
setComp (Array r1 ix e1 -> Comp
forall r ix e. Load r ix e => 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. Load r ix e => 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 :: * -> *).
(Mutable r ix e, 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. Load r ix e => 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. Load r ix e => 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 ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r ix e. Source r ix e => 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 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e)
  => (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 ix e1 r2 e2 r3 e3 (f :: * -> *) r e.
(Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f,
 Mutable r ix e) =>
(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 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e)
  => (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.
Construct r ix e =>
Comp -> Array r ix e -> Array r ix e
setComp (Array r1 ix e1 -> Comp
forall r ix e. Load r ix e => 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. Load r ix e => 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. Load r ix e => 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 :: * -> *).
(Mutable r ix e, 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 ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (Array r3 ix e3 -> ix -> e3
forall r ix e. Source r ix e => 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. Load r ix e => 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. Load r ix e => 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. Load r ix e => Array r ix e -> Sz ix
size Array r3 ix e3
arr3))
{-# INLINE izipWith3A #-}



-- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays either have to be the
-- same, or at least one of the two array must be a singleton array, in which case it will behave as
-- a `Data.Massiv.Array.map`.
--
-- @since 0.1.4
liftArray2
  :: (Source r1 ix a, Source r2 ix b)
  => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2 :: (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2 a -> b -> e
f !Array r1 ix a
arr1 !Array r2 ix b
arr2
  | Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
forall ix. Index ix => Sz ix
oneSz = (b -> e) -> Array r2 ix b -> Array D ix e
forall r ix e' e.
Source r ix e' =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (a -> b -> e
f (Array r1 ix a -> ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix a
arr1 ix
forall ix. Index ix => ix
zeroIndex)) Array r2 ix b
arr2
  | Sz ix
sz2 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
forall ix. Index ix => Sz ix
oneSz = (a -> e) -> Array r1 ix a -> Array D ix e
forall r ix e' e.
Source r ix e' =>
(e' -> e) -> Array r ix e' -> Array D ix e
map (a -> b -> e
`f` Array r2 ix b -> ix -> b
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r2 ix b
arr2 ix
forall ix. Index ix => ix
zeroIndex) Array r1 ix a
arr1
  | Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 =
    Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray (Array r1 ix a -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r1 ix a
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix b -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r2 ix b
arr2) Sz ix
sz1 (\ !ix
ix -> a -> b -> e
f (Array r1 ix a -> ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix a
arr1 ix
ix) (Array r2 ix b -> ix -> b
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r2 ix b
arr2 ix
ix))
  | Bool
otherwise = SizeException -> Array D ix e
forall a e. Exception e => e -> a
throw (SizeException -> Array D ix e) -> SizeException -> Array D ix e
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (Array r1 ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix a
arr1) (Array r2 ix b -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix b
arr2)
  where
    sz1 :: Sz ix
sz1 = Array r1 ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix a
arr1
    sz2 :: Sz ix
sz2 = Array r2 ix b -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix b
arr2
{-# INLINE liftArray2 #-}


--------------------------------------------------------------------------------
-- 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' ix a, Mutable r ix e, 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 -> (Int -> f e) -> f (Array r ix e)
forall r ix e (f :: * -> *).
(Mutable r ix e, Applicative f) =>
Sz ix -> (Int -> f e) -> f (Array r ix e)
makeArrayLinearA (Array r' ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix a
arr) (a -> f e
f (a -> f e) -> (Int -> a) -> Int -> f e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r' ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex 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 . (Source r ix 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 = 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 (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)) (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
. Array r ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex 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' ix (f e), Mutable r ix e, 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' ix a, Mutable r ix e, 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 . (Source r ix (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 :: * -> *).
(Source r ix 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' ix a, Mutable r ix e, 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.
Construct r ix e =>
Comp -> Array r ix e -> Array r ix e
setComp (Array r' ix a -> Comp
forall r ix e. Load r ix e => 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 :: * -> *).
(Mutable r ix e, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA (Array r' ix a -> Sz ix
forall r ix e. Load r ix e => 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 ix e. Source r ix e => 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 ix a, 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 =
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f e) -> 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) (\ !Int
i -> ix -> a -> f e
f (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i) (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr Int
i))
  where
    sz :: Sz ix
sz = Array r ix a -> Sz ix
forall r ix e. Load r ix e => 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' ix a, Mutable r ix b, 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 = (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
forall r ix b r' a (m :: * -> *).
(Source r' ix a, Mutable r ix b, PrimMonad m) =>
(ix -> a -> m b) -> Array r' ix a -> m (Array r ix b)
itraversePrim ((a -> m b) -> ix -> a -> m b
forall a b. a -> b -> a
const a -> m b
f)
{-# 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' ix a, Mutable r ix b, 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 =
  Comp -> Array r ix b -> Array r ix b
forall r ix e.
Construct r ix e =>
Comp -> Array r ix e -> Array r ix e
setComp (Array r' ix a -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix a
arr) (Array r ix b -> Array r ix b)
-> m (Array r ix b) -> m (Array r ix b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  Sz ix -> (Int -> m b) -> m (Array r ix b)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Sz ix -> (Int -> m e) -> m (Array r ix e)
generateArrayLinearS
    (Array r' ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix a
arr)
    (\ !Int
i ->
       let ix :: ix
ix = Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex (Array r' ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix a
arr) Int
i
        in ix -> a -> m b
f ix
ix (Array r' ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r' ix a
arr Int
i))
{-# 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix e, 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix e, 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix e, 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix e, 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 ix a, Monad m) => (a -> m b) -> Array r ix a -> m ()
mapM_ :: (a -> m b) -> Array r ix a -> m ()
mapM_ a -> m b
f !Array r ix a
arr = ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m b) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ ix
forall ix. Index ix => ix
zeroIndex (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix a
arr)) (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) (a -> m b
f (a -> m b) -> (ix -> a) -> ix -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix a -> ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix a
arr)
{-# 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 ix a, 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 a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_
{-# INLINE forM_ #-}


-- | Just like `imapM_`, except with flipped arguments.
iforM_ :: (Source r ix a, 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 a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(ix -> a -> m b) -> Array r ix a -> m ()
imapM_
{-# 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. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad 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 :: * -> *).
(Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad 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.
--
-- @since 0.2.6
mapIO_ :: (Source r b e, MonadUnliftIO m) => (e -> m a) -> Array r b e -> m ()
mapIO_ :: (e -> m a) -> Array r b e -> m ()
mapIO_ e -> m a
action = (b -> e -> m a) -> Array r b e -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix e, MonadUnliftIO m) =>
(ix -> e -> m a) -> Array r ix e -> m ()
imapIO_ ((e -> m a) -> b -> e -> m a
forall a b. a -> b -> a
const e -> m a
action)
{-# INLINE mapIO_ #-}

-- | Same as `mapIO_`, but map an index aware action instead.
--
-- @since 0.2.6
imapIO_ :: (Source 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 =
  Comp -> (Scheduler m () -> m ()) -> m ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler m a -> m b) -> m ()
withScheduler_ (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler m () -> m ()) -> m ())
-> (Scheduler m () -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Scheduler m ()
scheduler -> Scheduler m () -> (ix -> e -> m a) -> Array r ix e -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
Scheduler m () -> (ix -> e -> m a) -> Array r ix e -> m ()
imapSchedulerM_ Scheduler m ()
scheduler ix -> e -> m a
action Array r ix e
arr
{-# INLINE imapIO_ #-}

-- | Same as `imapM_`, but will use the supplied scheduler.
--
-- @since 0.3.1
imapSchedulerM_ ::
     (Source r ix e, Monad m) => Scheduler m () -> (ix -> e -> m a) -> Array r ix e -> m ()
imapSchedulerM_ :: Scheduler m () -> (ix -> e -> m a) -> Array r ix e -> m ()
imapSchedulerM_ Scheduler m ()
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. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
  Scheduler m () -> Int -> (Int -> e) -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *) b.
Monad m =>
Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_
    Scheduler m ()
scheduler
    (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
    (Array r ix e -> Int -> e
forall r ix e. Source r ix e => 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_ ::
     (Source r ix e, Monad m) => Scheduler m () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ :: Scheduler m () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler m ()
scheduler Array r ix e
arr ix -> e -> m a
action = Scheduler m () -> (ix -> e -> m a) -> Array r ix e -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
Scheduler m () -> (ix -> e -> m a) -> Array r ix e -> m ()
imapSchedulerM_ Scheduler m ()
scheduler ix -> e -> m a
action Array r ix e
arr
{-# INLINE iforSchedulerM_ #-}


-- | 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. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad 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 = Comp -> Sz ix -> (ix -> m b) -> m (Array r ix b)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, PrimMonad m, Mutable r ix e) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray (Array r' ix a -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix a
arr) (Array r' ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix a
arr) ((ix -> m b) -> m (Array r ix b))
-> (ix -> m b) -> m (Array r ix b)
forall a b. (a -> b) -> a -> b
$ \ix
ix -> ix -> a -> m b
action ix
ix (Array r' ix a -> ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r' ix a
arr ix
ix)
{-# INLINE imapIO #-}

-- | Same as `mapIO` but with arguments flipped.
--
-- @since 0.2.6
forIO ::
     forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad 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 :: * -> *).
(Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad 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 `initWorkerStates` for the `WorkerStates`
-- initialization.
--
-- @since 0.3.4
imapWS ::
     forall r ix b r' a s m. (Source r' ix a, Mutable r ix b, 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 :: * -> *).
(Mutable r ix e, 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. Load r ix e => 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 ix e. Source r ix e => 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix b, 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' ix a, Mutable r ix b, 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_ :: (Source 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 b e (m :: * -> *) a.
(Source r b e, MonadUnliftIO m) =>
(e -> m a) -> Array r b e -> m ()
mapIO_
{-# INLINE forIO_ #-}

-- | Same as `imapIO` but with arguments flipped.
--
-- @since 0.2.6
iforIO ::
     forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad 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 :: * -> *).
(Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad 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_ :: (Source r ix a, MonadUnliftIO m) => Array r ix a -> (ix -> a -> m b) -> m ()
iforIO_ :: Array r ix a -> (ix -> a -> m b) -> m ()
iforIO_ = ((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 (m :: * -> *) a.
(Source r ix e, MonadUnliftIO m) =>
(ix -> e -> m a) -> Array r ix e -> m ()
imapIO_
{-# INLINE iforIO_ #-}