{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Massiv.Array.Ops.Map (
map,
imap,
traverseA,
traverseA_,
itraverseA,
itraverseA_,
sequenceA,
sequenceA_,
traversePrim,
itraversePrim,
mapM,
forM,
imapM,
iforM,
mapM_,
forM_,
imapM_,
iforM_,
mapIO,
mapWS,
mapIO_,
imapIO,
imapWS,
imapIO_,
forIO,
forWS,
forIO_,
iforIO,
iforWS,
iforIO_,
imapSchedulerM_,
iforSchedulerM_,
iterArrayLinearM_,
iterArrayLinearWithSetM_,
iterArrayLinearWithStrideM_,
zip,
zip3,
zip4,
unzip,
unzip3,
unzip4,
zipWith,
zipWith3,
zipWith4,
izipWith,
izipWith3,
izipWith4,
zipWithA,
izipWithA,
zipWith3A,
izipWith3A,
) where
import Control.Monad (void)
import Control.Monad.Primitive
import Control.Scheduler
import Data.Coerce
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Manifest.List
import Data.Massiv.Array.Mutable
import Data.Massiv.Array.Ops.Construct (makeArrayA, makeArrayLinearA)
import Data.Massiv.Core.Common
import Data.Traversable (traverse)
import Prelude hiding (
map,
mapM,
mapM_,
sequenceA,
traverse,
unzip,
unzip3,
zip,
zip3,
zipWith,
zipWith3,
)
map :: (Index ix, Source r e') => (e' -> e) -> Array r ix e' -> Array D ix e
map :: forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map e' -> e
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap e' -> e
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay
{-# INLINE map #-}
zip
:: (Index ix, Source r1 e1, Source r2 e2)
=> Array r1 ix e1
-> Array r2 ix e2
-> Array D ix (e1, e2)
zip :: forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
Array r1 ix e1 -> Array r2 ix e2 -> Array D ix (e1, e2)
zip = 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 #-}
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 :: forall ix r1 e1 r2 e2 r3 e3.
(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 = 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 #-}
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 :: forall ix r1 e1 r2 e2 r3 e3 r4 e4.
(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 = 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 :: (Index ix, Source r (e1, e2)) => Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2)
unzip :: forall ix r e1 e2.
(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)
arr = (forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map forall a b. (a, b) -> a
fst Array r ix (e1, e2)
arr, forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
map forall a b. (a, b) -> b
snd Array r ix (e1, e2)
arr)
{-# INLINE unzip #-}
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 :: forall ix r e1 e2 e3.
(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)
arr = (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, 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, 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 #-}
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 :: forall ix r e1 e2 e3 e4.
(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)
arr =
( 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
, 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
, 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
, 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
:: (Index ix, Source r1 e1, Source r2 e2)
=> (e1 -> e2 -> e)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array D ix e
zipWith :: 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 e1 -> e2 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 = 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 = forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2)))
comp :: Comp
comp = forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2
prefIndex :: PrefIndex ix e
prefIndex = forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (\ix
ix -> e1 -> e2 -> e
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix))
{-# INLINE zipWith #-}
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 :: forall ix r1 e1 r2 e2 e.
(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
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 =
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray
(forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2)
(forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2))))
(forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (\ix
ix -> ix -> e1 -> e2 -> e
f ix
ix (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (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 #-}
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 :: 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 e1 -> e2 -> e3 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 = 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
{-# INLINE zipWith3 #-}
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 :: 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 -> e2 -> e3 -> e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 =
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray
(forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r3 ix e3
arr3)
( forall ix. ix -> Sz ix
SafeSz
( forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2
forall a. Ord a => a -> a -> a
min
(forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2)))
(coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r3 ix e3
arr3))
)
)
(forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex forall a b. (a -> b) -> a -> b
$ \ !ix
ix -> ix -> e1 -> e2 -> e3 -> e
f ix
ix (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (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 #-}
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 :: 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 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 =
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
{-# INLINE zipWith4 #-}
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 :: 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 -> 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 =
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray
(forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r3 ix e3
arr3 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r4 ix e4
arr4)
( forall ix. ix -> Sz ix
SafeSz
( forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2
forall a. Ord a => a -> a -> a
min
( forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2
forall a. Ord a => a -> a -> a
min
(forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2)))
(coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r3 ix e3
arr3))
)
(coerce :: forall a b. Coercible a b => a -> b
coerce (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 (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r3 ix e3
arr3 ix
ix) (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 #-}
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 :: forall r1 e1 r2 e2 (f :: * -> *) r e ix.
(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
f Array r1 ix e1
arr1 Array r2 ix e2
arr2
| Sz ix
sz1 forall a. Eq a => a -> a -> Bool
== forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2
, PrefIndexLinear Int -> e1
gi1 <- 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 <- forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r2 ix e2
arr2 =
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 = 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 (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 = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1
{-# INLINE zipWithA #-}
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 :: 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 ix -> e1 -> e2 -> f e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 =
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA
(forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (coerce :: forall a b. Coercible a b => a -> b
coerce (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 (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (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 #-}
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 :: 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) =>
(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 = 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 (forall a b. a -> b -> a
const e1 -> e2 -> e3 -> f e
f)
{-# INLINE zipWith3A #-}
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 :: 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 ix -> e1 -> e2 -> e3 -> f e
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 Array r3 ix e3
arr3 =
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r1 ix e1
arr1 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r2 ix e2
arr2 forall a. Semigroup a => a -> a -> a
<> forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r3 ix e3
arr3)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix) (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 =
forall ix. ix -> Sz ix
SafeSz forall a b. (a -> b) -> a -> b
$
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r1 ix e1
arr1)) (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r2 ix e2
arr2))) (coerce :: forall a b. Coercible a b => a -> b
coerce (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r3 ix e3
arr3))
{-# INLINE izipWith3A #-}
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 :: 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
f Array r' ix a
arr =
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r e. Manifest r e => Comp -> [e] -> Vector r e
fromList (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f e
f (forall ix r e. (Index ix, Source r e) => Array r ix e -> [e]
toList Array r' ix a
arr)
{-# INLINE traverseA #-}
traverseA_
:: forall r ix e a f
. (Index ix, Source r e, Applicative f)
=> (e -> f a)
-> Array r ix e
-> f ()
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
f Array r ix e
arr =
case 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 -> forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ forall ix. Index ix => ix
zeroIndex (forall ix. Sz ix -> ix
unSz Sz ix
sz) forall ix. Index ix => ix
oneIndex forall a. Ord a => a -> a -> Bool
(<) (e -> f a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
gix)
PrefIndexLinear Int -> e
gi -> forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (forall a. Num a => a -> a -> a
+ Int
1) (e -> f a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
gi)
where
sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
{-# INLINE traverseA_ #-}
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 :: 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 = 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 forall a. a -> a
id
{-# INLINE sequenceA #-}
sequenceA_
:: forall r ix e f
. (Index ix, Source r (f e), Applicative f)
=> Array r ix (f e)
-> f ()
sequenceA_ :: forall r ix e (f :: * -> *).
(Index ix, Source r (f e), Applicative f) =>
Array r ix (f e) -> f ()
sequenceA_ = forall r ix e a (f :: * -> *).
(Index ix, Source r e, Applicative f) =>
(e -> f a) -> Array r ix e -> f ()
traverseA_ forall a. a -> a
id
{-# INLINE sequenceA_ #-}
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 :: 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
f Array r' ix a
arr =
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA (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 (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 #-}
itraverseA_
:: forall r ix e a f
. (Source r a, Index ix, Applicative f)
=> (ix -> a -> f e)
-> Array r ix a
-> f ()
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
f Array r ix a
arr =
case 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 ->
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ forall ix. Index ix => ix
zeroIndex (forall ix. Sz ix -> ix
unSz Sz ix
sz) forall ix. Index ix => ix
oneIndex 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 ->
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 forall ix. Index ix => ix
zeroIndex forall ix. Index ix => Stride ix
oneStride 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 = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix a
arr
{-# INLINE itraverseA_ #-}
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 :: 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
f Array r' ix a
arr = do
let sz :: Sz ix
sz = 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 <- 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 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 ->
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 forall ix. Index ix => ix
zeroIndex forall ix. Index ix => Stride ix
oneStride forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix ->
a -> m b
f (ix -> a
gix ix
ix) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 ->
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \Int
i ->
a -> m b
f (Int -> a
gi Int
i) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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
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 (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 #-}
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 :: 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
f Array r' ix a
arr = do
let sz :: Sz ix
sz = 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 <- 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 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 ->
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 forall ix. Index ix => ix
zeroIndex forall ix. Index ix => Stride ix
oneStride forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix ->
ix -> a -> m b
f ix
ix (ix -> a
gix ix
ix) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 ->
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 forall ix. Index ix => ix
zeroIndex forall ix. Index ix => Stride ix
oneStride forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix ->
ix -> a -> m b
f ix
ix (Int -> a
gi Int
i) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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
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 (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
:: forall r ix b r' a m
. (Source r' a, Manifest r b, Index ix, Monad m)
=> (a -> m b)
-> Array r' ix a
-> m (Array r ix b)
mapM :: forall r ix b r' a (m :: * -> *).
(Source r' a, Manifest r b, Index ix, Monad m) =>
(a -> m b) -> Array r' ix a -> m (Array r ix b)
mapM = 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 #-}
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 :: 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 = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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 #-}
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 :: 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 = 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 #-}
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 :: 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 = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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 #-}
mapM_ :: (Source r a, Index ix, Monad m) => (a -> m b) -> Array r ix a -> m ()
mapM_ :: forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_ = forall r ix e a (f :: * -> *).
(Index ix, Source r e, Applicative f) =>
(e -> f a) -> Array r ix e -> f ()
traverseA_
{-# INLINE mapM_ #-}
forM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (a -> m b) -> m ()
forM_ :: forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
Array r ix a -> (a -> m b) -> m ()
forM_ = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall r ix e a (f :: * -> *).
(Index ix, Source r e, Applicative f) =>
(e -> f a) -> Array r ix e -> f ()
traverseA_
{-# INLINE forM_ #-}
imapM_ :: (Index ix, Source r a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m ()
imapM_ :: forall ix r a (m :: * -> *) b.
(Index ix, Source r a, Monad m) =>
(ix -> a -> m b) -> Array r ix a -> m ()
imapM_ = 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_ #-}
iforM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (ix -> a -> m b) -> m ()
iforM_ :: forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
Array r ix a -> (ix -> a -> m b) -> m ()
iforM_ = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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_ #-}
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 :: 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
action = 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 (forall a b. a -> b -> a
const a -> m b
action)
{-# INLINE mapIO #-}
mapIO_
:: forall r ix e a m
. (Load r ix e, MonadUnliftIO m)
=> (e -> m a)
-> Array r ix e
-> m ()
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
action Array r ix e
arr =
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
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
_ -> forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. m a -> IO a
run forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> m a
action)
{-# INLINE mapIO_ #-}
imapIO_
:: forall r ix e a m
. (Load r ix e, MonadUnliftIO m)
=> (ix -> e -> m a)
-> Array r ix e
-> m ()
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
action Array r ix e
arr =
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
let sz :: Sz ix
sz = forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr
in
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 -> forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. m a -> IO a
run forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e -> m a
action (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i))
{-# INLINE imapIO_ #-}
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 :: 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
action Array r' ix a
arr = do
let sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix a
arr
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> do
MArray RealWorld r ix b
marr <- 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_ (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
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 forall a b. (a -> b) -> a -> b
$ \ !Int
i a
e ->
forall a. m a -> IO a
run (ix -> a -> m b
action (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i) a
e) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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
marr Int
i
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 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix a
arr) MArray RealWorld r ix b
marr
{-# INLINE imapIO #-}
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 :: 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 = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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 #-}
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 :: 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 = 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 (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 (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 #-}
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 :: 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
states a -> s -> m b
f = 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 #-}
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 :: 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
states Array r' ix a
f ix -> a -> s -> m b
arr = 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 #-}
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 :: 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
states Array r' ix a
arr a -> s -> m b
f = 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 #-}
forIO_ :: (Load r ix e, MonadUnliftIO m) => Array r ix e -> (e -> m a) -> m ()
forIO_ :: forall r ix e (m :: * -> *) a.
(Load r ix e, MonadUnliftIO m) =>
Array r ix e -> (e -> m a) -> m ()
forIO_ = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall r ix e a (m :: * -> *).
(Load r ix e, MonadUnliftIO m) =>
(e -> m a) -> Array r ix e -> m ()
mapIO_
{-# INLINE forIO_ #-}
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 :: 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 = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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 #-}
iforIO_
:: forall r ix e a m
. (Load r ix e, MonadUnliftIO m)
=> Array r ix e
-> (ix -> e -> m a)
-> m ()
iforIO_ :: forall r ix e a (m :: * -> *).
(Load r ix e, MonadUnliftIO m) =>
Array r ix e -> (ix -> e -> m a) -> m ()
iforIO_ = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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
-> (Int -> e -> m ())
-> m ()
iterArrayLinearM_ :: 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 s ()
scheduler Array r ix e
arr Int -> e -> m ()
f =
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim forall a b. (a -> b) -> a -> b
$ 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 -> forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim 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
-> (Int -> e -> m ())
-> (Ix1 -> Sz1 -> e -> m ())
-> m ()
iterArrayLinearWithSetM_ :: forall r ix e (m :: * -> *) s.
(Load r ix e, MonadPrimBase s m) =>
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 =
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim forall a b. (a -> b) -> a -> b
$
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 -> forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e -> m ()
f Int
i) (\Int
i Sz1
n -> forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim 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
-> Sz ix
-> Array r ix e
-> (Int -> e -> m ())
-> m ()
iterArrayLinearWithStrideM_ :: forall r ix e (m :: * -> *) s.
(StrideLoad r ix e, MonadPrimBase s m) =>
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 =
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim forall a b. (a -> b) -> a -> b
$ 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 -> forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e -> m ()
f Int
i)
{-# INLINE iterArrayLinearWithStrideM_ #-}
imapSchedulerM_
:: (Index ix, Source r e, MonadPrimBase s m)
=> Scheduler s ()
-> (ix -> e -> m a)
-> Array r ix e
-> m ()
imapSchedulerM_ :: 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 = do
let sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_
Scheduler s ()
scheduler
(forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
(forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr)
(\Int
i -> forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e -> m a
action (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i))
{-# INLINE imapSchedulerM_ #-}
iforSchedulerM_
:: (Index ix, Source r e, MonadPrimBase s m)
=> Scheduler s ()
-> Array r ix e
-> (ix -> e -> m a)
-> m ()
iforSchedulerM_ :: forall ix r e s (m :: * -> *) a.
(Index ix, Source r e, MonadPrimBase s m) =>
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 = 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_ #-}