{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
module Data.Massiv.Array.Ops.Transform
(
transpose
, transposeInner
, transposeOuter
, reverse
, reverse'
, reverseM
, backpermuteM
, backpermute'
, resizeM
, resize'
, flatten
, extractM
, extract'
, extractFromToM
, extractFromTo'
, deleteRowsM
, deleteColumnsM
, deleteRegionM
, appendOuterM
, appendM
, append'
, concatOuterM
, concatM
, concat'
, stackSlicesM
, stackOuterSlicesM
, stackInnerSlicesM
, splitAtM
, splitAt'
, splitExtractM
, upsample
, downsample
, zoom
, zoomWithGrid
, transformM
, transform'
, transform2M
, transform2'
) where
import Control.Scheduler (traverse_)
import Control.Monad as M (foldM_, unless, forM_)
import Data.Bifunctor (bimap)
import Data.Foldable as F (foldl', foldrM, toList, length)
import qualified Data.List as L (uncons)
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Delayed.Push
import Data.Massiv.Array.Mutable
import Data.Massiv.Array.Ops.Construct
import Data.Massiv.Array.Ops.Map
import Data.Massiv.Core.Common
import Prelude as P hiding (concat, splitAt, traverse, mapM_, reverse, take, drop)
extractM :: (MonadThrow m, Extract r ix e)
=> ix
-> Sz ix
-> Array r ix e
-> m (Array (R r) ix e)
!ix
sIx !Sz ix
newSz !Array r ix e
arr
| Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz1 ix
sIx Bool -> Bool -> Bool
&& Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
eIx1 ix
sIx Bool -> Bool -> Bool
&& Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz1 ix
eIx =
Array (R r) ix e -> m (Array (R r) ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array (R r) ix e -> m (Array (R r) ix e))
-> Array (R r) ix e -> m (Array (R r) ix e)
forall a b. (a -> b) -> a -> b
$ ix -> Sz ix -> Array r ix e -> Array (R r) ix e
forall r ix e.
Extract r ix e =>
ix -> Sz ix -> Array r ix e -> Array (R r) ix e
unsafeExtract ix
sIx Sz ix
newSz Array r ix e
arr
| Bool
otherwise = SizeException -> m (Array (R r) ix e)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> m (Array (R r) ix e))
-> SizeException -> m (Array (R r) ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> ix -> Sz ix -> SizeException
SizeSubregionException (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) ix
sIx Sz ix
newSz
where
sz1 :: Sz ix
sz1 = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (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)))
eIx1 :: Sz ix
eIx1 = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ix
eIx)
eIx :: ix
eIx = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
sIx (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
newSz
{-# INLINE extractM #-}
extract' :: Extract r ix e
=> ix
-> Sz ix
-> Array r ix e
-> Array (R r) ix e
ix
sIx Sz ix
newSz = (SomeException -> Array (R r) ix e)
-> (Array (R r) ix e -> Array (R r) ix e)
-> Either SomeException (Array (R r) ix e)
-> Array (R r) ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array (R r) ix e
forall a e. Exception e => e -> a
throw Array (R r) ix e -> Array (R r) ix e
forall a. a -> a
id (Either SomeException (Array (R r) ix e) -> Array (R r) ix e)
-> (Array r ix e -> Either SomeException (Array (R r) ix e))
-> Array r ix e
-> Array (R r) ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix
-> Sz ix -> Array r ix e -> Either SomeException (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
sIx Sz ix
newSz
{-# INLINE extract' #-}
extractFromToM :: (MonadThrow m, Extract r ix e) =>
ix
-> ix
-> Array r ix e
-> m (Array (R r) ix e)
ix
sIx ix
eIx = ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
sIx (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
eIx ix
sIx))
{-# INLINE extractFromToM #-}
extractFromTo' :: Extract r ix e =>
ix
-> ix
-> Array r ix e
-> Array (R r) ix e
ix
sIx ix
eIx = ix -> Sz ix -> Array r ix e -> Array (R r) ix e
forall r ix e.
Extract r ix e =>
ix -> Sz ix -> Array r ix e -> Array (R r) ix e
extract' ix
sIx (Sz ix -> Array r ix e -> Array (R r) ix e)
-> Sz ix -> Array r ix e -> Array (R r) ix e
forall a b. (a -> b) -> a -> b
$ ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
eIx ix
sIx)
{-# INLINE extractFromTo' #-}
resizeM ::
(MonadThrow m, Index ix', Load r ix e, Resize r ix)
=> Sz ix'
-> Array r ix e
-> m (Array r ix' e)
resizeM :: Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM Sz ix'
sz Array r ix e
arr = Sz ix -> Sz ix' -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements (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) Sz ix'
sz m () -> m (Array r ix' e) -> m (Array r ix' e)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Array r ix' e -> m (Array r ix' e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Sz ix' -> Array r ix e -> Array r ix' e
forall r ix ix' e.
(Resize r ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix'
sz Array r ix e
arr)
{-# INLINE resizeM #-}
resize' :: (Index ix', Load r ix e, Resize r ix) => Sz ix' -> Array r ix e -> Array r ix' e
resize' :: Sz ix' -> Array r ix e -> Array r ix' e
resize' Sz ix'
sz = (SomeException -> Array r ix' e)
-> (Array r ix' e -> Array r ix' e)
-> Either SomeException (Array r ix' e)
-> Array r ix' e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array r ix' e
forall a e. Exception e => e -> a
throw Array r ix' e -> Array r ix' e
forall a. a -> a
id (Either SomeException (Array r ix' e) -> Array r ix' e)
-> (Array r ix e -> Either SomeException (Array r ix' e))
-> Array r ix e
-> Array r ix' e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix' -> Array r ix e -> Either SomeException (Array r ix' e)
forall (m :: * -> *) ix' r ix e.
(MonadThrow m, Index ix', Load r ix e, Resize r ix) =>
Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM Sz ix'
sz
{-# INLINE resize' #-}
flatten :: (Load r ix e, Resize r ix) => Array r ix e -> Array r Ix1 e
flatten :: Array r ix e -> Array r Int e
flatten Array r ix e
arr = Sz Int -> Array r ix e -> Array r Int e
forall r ix ix' e.
(Resize r ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (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))) Array r ix e
arr
{-# INLINE flatten #-}
transpose :: Source r Ix2 e => Array r Ix2 e -> Array D Ix2 e
transpose :: Array r Ix2 e -> Array D Ix2 e
transpose = Array r Ix2 e -> Array D Ix2 e
forall ix r' e.
(Index (Lower ix), Source r' ix e) =>
Array r' ix e -> Array D ix e
transposeInner
{-# INLINE [1] transpose #-}
{-# RULES
"transpose . transpose" [~1] forall arr . transpose (transpose arr) = delay arr
"transposeInner . transposeInner" [~1] forall arr . transposeInner (transposeInner arr) = delay arr
"transposeOuter . transposeOuter" [~1] forall arr . transposeOuter (transposeOuter arr) = delay arr
#-}
transposeInner :: (Index (Lower ix), Source r' ix e)
=> Array r' ix e -> Array D ix e
transposeInner :: Array r' ix e -> Array D ix e
transposeInner !Array r' ix e
arr = Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr) Sz ix
newsz ix -> e
newVal
where
transInner :: ix -> ix
transInner !ix
ix =
(SomeException -> ix)
-> (ix -> ix) -> Either SomeException ix -> ix
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> ix
forall e a. Exception e => e -> a
throwImpossible ix -> ix
forall a. a -> a
id (Either SomeException ix -> ix) -> Either SomeException ix -> ix
forall a b. (a -> b) -> a -> b
$ do
Int
n <- ix -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM ix
ix Dim
dix
Int
m <- ix -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM ix
ix (Dim
dix Dim -> Dim -> Dim
forall a. Num a => a -> a -> a
- Dim
1)
ix
ix' <- ix -> Dim -> Int -> Either SomeException ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
ix Dim
dix Int
m
ix -> Dim -> Int -> Either SomeException ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
ix' (Dim
dix Dim -> Dim -> Dim
forall a. Num a => a -> a -> a
- Dim
1) Int
n
{-# INLINE transInner #-}
newVal :: ix -> e
newVal = Array r' ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r' ix e
arr (ix -> e) -> (ix -> ix) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix
transInner
{-# INLINE newVal #-}
!newsz :: Sz ix
newsz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (ix -> ix
transInner (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (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)))
!dix :: Dim
dix = Sz ix -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions Sz ix
newsz
{-# INLINE [1] transposeInner #-}
transposeOuter :: (Index (Lower ix), Source r' ix e)
=> Array r' ix e -> Array D ix e
transposeOuter :: Array r' ix e -> Array D ix e
transposeOuter !Array r' ix e
arr = Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr) Sz ix
newsz ix -> e
newVal
where
transOuter :: c -> c
transOuter !c
ix =
(SomeException -> c) -> (c -> c) -> Either SomeException c -> c
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> c
forall e a. Exception e => e -> a
throwImpossible c -> c
forall a. a -> a
id (Either SomeException c -> c) -> Either SomeException c -> c
forall a b. (a -> b) -> a -> b
$ do
Int
n <- c -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM c
ix Dim
1
Int
m <- c -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM c
ix Dim
2
c
ix' <- c -> Dim -> Int -> Either SomeException c
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM c
ix Dim
1 Int
m
c -> Dim -> Int -> Either SomeException c
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM c
ix' Dim
2 Int
n
{-# INLINE transOuter #-}
newVal :: ix -> e
newVal = Array r' ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r' ix e
arr (ix -> e) -> (ix -> ix) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix
forall c. Index c => c -> c
transOuter
{-# INLINE newVal #-}
!newsz :: Sz ix
newsz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (ix -> ix
forall c. Index c => c -> c
transOuter (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (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)))
{-# INLINE [1] transposeOuter #-}
reverse :: (IsIndexDimension ix n, Source r ix e) => Dimension n -> Array r ix e -> Array D ix e
reverse :: Dimension n -> Array r ix e -> Array D ix e
reverse Dimension n
dim = Dim -> Array r ix e -> Array D ix e
forall r ix e. Source r ix e => Dim -> Array r ix e -> Array D ix e
reverse' (Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE reverse #-}
reverseM :: (MonadThrow m, Source r ix e) => Dim -> Array r ix e -> m (Array D ix e)
reverseM :: Dim -> Array r ix e -> m (Array D ix e)
reverseM Dim
dim 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
Int
k <- ix -> Dim -> m Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) Dim
dim
Array D ix e -> m (Array D ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D ix e -> m (Array D ix e))
-> Array D ix e -> m (Array D ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix
sz ((ix -> e) -> Array D ix e) -> (ix -> e) -> Array D ix e
forall a b. (a -> b) -> a -> b
$ \ ix
ix ->
Array r ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ((Int, ix) -> ix
forall a b. (a, b) -> b
snd ((Int, ix) -> ix) -> (Int, ix) -> ix
forall a b. (a -> b) -> a -> b
$ ix -> Dim -> (Int -> Int) -> (Int, ix)
forall ix. Index ix => ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' ix
ix Dim
dim (\Int
i -> Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
{-# INLINE reverseM #-}
reverse' :: Source r ix e => Dim -> Array r ix e -> Array D ix e
reverse' :: Dim -> Array r ix e -> Array D ix e
reverse' Dim
dim = (SomeException -> Array D ix e)
-> (Array D ix e -> Array D ix e)
-> Either SomeException (Array D ix e)
-> Array D ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D ix e
forall a e. Exception e => e -> a
throw Array D ix e -> Array D ix e
forall a. a -> a
id (Either SomeException (Array D ix e) -> Array D ix e)
-> (Array r ix e -> Either SomeException (Array D ix e))
-> Array r ix e
-> Array D ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> Array r ix e -> Either SomeException (Array D ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Source r ix e) =>
Dim -> Array r ix e -> m (Array D ix e)
reverseM Dim
dim
{-# INLINE reverse' #-}
backpermuteM ::
forall r ix e r' ix' m.
(Mutable r ix e, Source r' ix' e, MonadUnliftIO m, PrimMonad m, MonadThrow m)
=> Sz ix
-> (ix -> ix')
-> Array r' ix' e
-> m (Array r ix e)
backpermuteM :: Sz ix -> (ix -> ix') -> Array r' ix' e -> m (Array r ix e)
backpermuteM Sz ix
sz ix -> ix'
ixF !Array r' ix' e
arr = Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
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' e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e
arr) Sz ix
sz (Array r' ix' e -> ix' -> m e
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r' ix' e
arr (ix' -> m e) -> (ix -> ix') -> ix -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix'
ixF)
{-# INLINE backpermuteM #-}
backpermute' :: (Source r' ix' e, Index ix) =>
Sz ix
-> (ix -> ix')
-> Array r' ix' e
-> Array D ix e
backpermute' :: Sz ix -> (ix -> ix') -> Array r' ix' e -> Array D ix e
backpermute' Sz ix
sz ix -> ix'
ixF !Array r' ix' e
arr = Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r' ix' e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e
arr) Sz ix
sz (Array r' ix' e -> ix' -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r' ix' e
arr (ix' -> e) -> (ix -> ix') -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix'
ixF)
{-# INLINE backpermute' #-}
appendM ::
forall r1 r2 ix e m. (MonadThrow m, Source r1 ix e, Source r2 ix e)
=> Dim
-> Array r1 ix e
-> Array r2 ix e
-> m (Array DL ix e)
appendM :: Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e)
appendM Dim
n !Array r1 ix e
arr1 !Array r2 ix e
arr2 = do
let !sz1 :: Sz ix
sz1 = Array r1 ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix e
arr1
!sz2 :: Sz ix
sz2 = Array r2 ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix e
arr2
(Sz Int
k1, Sz (Lower ix)
szl1) <- Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM Sz ix
sz1 Dim
n
(Sz Int
k2, Sz (Lower ix)
szl2) <- Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM Sz ix
sz2 Dim
n
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz (Lower ix)
szl1 Sz (Lower ix) -> Sz (Lower ix) -> Bool
forall a. Eq a => a -> a -> Bool
== Sz (Lower ix)
szl2) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ SizeException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> m ()) -> SizeException -> m ()
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz ix
sz1 Sz ix
sz2
let !k1' :: Int
k1' = Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
k1
Sz ix
newSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM Sz (Lower ix)
szl1 Dim
n (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int
k1' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
k2))
let load :: Monad n =>
Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()
load :: Scheduler n ()
-> Int
-> (Int -> e -> n ())
-> (Int -> Sz Int -> e -> n ())
-> n ()
load Scheduler n ()
scheduler !Int
startAt Int -> e -> n ()
dlWrite Int -> Sz Int -> e -> n ()
_dlSet = do
Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> n ()) -> n ()
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 Sz ix
sz1) (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((ix -> n ()) -> n ()) -> (ix -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
Int -> e -> n ()
dlWrite (Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix) (Array r1 ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e
arr1 ix
ix)
Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> n ()) -> n ()
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 Sz ix
sz2) (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((ix -> n ()) -> n ()) -> (ix -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
let i :: Int
i = ix -> Dim -> Int
forall ix. Index ix => ix -> Dim -> Int
getDim' ix
ix Dim
n
ix' :: ix
ix' = ix -> Dim -> Int -> ix
forall ix. Index ix => ix -> Dim -> Int -> ix
setDim' ix
ix Dim
n (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k1')
in Int -> e -> n ()
dlWrite (Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') (Array r2 ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e
arr2 ix
ix)
{-# INLINE load #-}
Array DL ix e -> m (Array DL ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array DL ix e -> m (Array DL ix e))
-> Array DL ix e -> m (Array DL ix e)
forall a b. (a -> b) -> a -> b
$
DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ())
-> Array DL ix e
DLArray
{dlComp :: Comp
dlComp = Array r1 ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r1 ix e
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r2 ix e
arr2, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
{-# INLINE appendM #-}
append' :: (Source r1 ix e, Source r2 ix e) =>
Dim -> Array r1 ix e -> Array r2 ix e -> Array DL ix e
append' :: Dim -> Array r1 ix e -> Array r2 ix e -> Array DL ix e
append' Dim
dim Array r1 ix e
arr1 Array r2 ix e
arr2 = (SomeException -> Array DL ix e)
-> (Array DL ix e -> Array DL ix e)
-> Either SomeException (Array DL ix e)
-> Array DL ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array DL ix e
forall a e. Exception e => e -> a
throw Array DL ix e -> Array DL ix e
forall a. a -> a
id (Either SomeException (Array DL ix e) -> Array DL ix e)
-> Either SomeException (Array DL ix e) -> Array DL ix e
forall a b. (a -> b) -> a -> b
$ Dim
-> Array r1 ix e
-> Array r2 ix e
-> Either SomeException (Array DL ix e)
forall r1 r2 ix e (m :: * -> *).
(MonadThrow m, Source r1 ix e, Source r2 ix e) =>
Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e)
appendM Dim
dim Array r1 ix e
arr1 Array r2 ix e
arr2
{-# INLINE append' #-}
concat' :: (Foldable f, Source r ix e) => Dim -> f (Array r ix e) -> Array DL ix e
concat' :: Dim -> f (Array r ix e) -> Array DL ix e
concat' Dim
n f (Array r ix e)
arrs = (SomeException -> Array DL ix e)
-> (Array DL ix e -> Array DL ix e)
-> Either SomeException (Array DL ix e)
-> Array DL ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array DL ix e
forall a e. Exception e => e -> a
throw Array DL ix e -> Array DL ix e
forall a. a -> a
id (Either SomeException (Array DL ix e) -> Array DL ix e)
-> Either SomeException (Array DL ix e) -> Array DL ix e
forall a b. (a -> b) -> a -> b
$ Dim -> f (Array r ix e) -> Either SomeException (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(MonadThrow m, Foldable f, Source r ix e) =>
Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM Dim
n f (Array r ix e)
arrs
{-# INLINE concat' #-}
concatM ::
forall r ix e f m. (MonadThrow m, Foldable f, Source r ix e)
=> Dim
-> f (Array r ix e)
-> m (Array DL ix e)
concatM :: Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM Dim
n !f (Array r ix e)
arrsF =
case [Array r ix e] -> Maybe (Array r ix e, [Array r ix e])
forall a. [a] -> Maybe (a, [a])
L.uncons (f (Array r ix e) -> [Array r ix e]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList f (Array r ix e)
arrsF) of
Maybe (Array r ix e, [Array r ix e])
Nothing -> Array DL ix e -> m (Array DL ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array DL ix e
forall r ix e. Construct r ix e => Array r ix e
empty
Just (Array r ix e
a, [Array r ix e]
arrs) -> do
let sz :: ix
sz = Sz ix -> ix
forall ix. Sz ix -> ix
unSz (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
a)
szs :: [ix]
szs = Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Sz ix -> ix) -> (Array r ix e -> Sz ix) -> Array r ix e -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 -> ix) -> [Array r ix e] -> [ix]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Array r ix e]
arrs
(Int
k, Lower ix
szl) <- ix -> Dim -> m (Int, Lower ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
sz Dim
n
([Int]
ks, [Lower ix]
szls) <-
(ix -> ([Int], [Lower ix]) -> m ([Int], [Lower ix]))
-> ([Int], [Lower ix]) -> [ix] -> m ([Int], [Lower ix])
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
F.foldrM (\ !ix
csz ([Int]
ks, [Lower ix]
szls) -> (Int -> [Int])
-> (Lower ix -> [Lower ix])
-> (Int, Lower ix)
-> ([Int], [Lower ix])
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ks) (Lower ix -> [Lower ix] -> [Lower ix]
forall a. a -> [a] -> [a]
: [Lower ix]
szls) ((Int, Lower ix) -> ([Int], [Lower ix]))
-> m (Int, Lower ix) -> m ([Int], [Lower ix])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ix -> Dim -> m (Int, Lower ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
csz Dim
n) ([], []) [ix]
szs
((ix, Lower ix) -> m ()) -> [(ix, Lower ix)] -> m ()
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
traverse_
(\(ix
sz', Lower ix
_) -> SizeException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ix
sz) (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ix
sz')))
(((ix, Lower ix) -> Bool) -> [(ix, Lower ix)] -> [(ix, Lower ix)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile ((Lower ix -> Lower ix -> Bool
forall a. Eq a => a -> a -> Bool
== Lower ix
szl) (Lower ix -> Bool)
-> ((ix, Lower ix) -> Lower ix) -> (ix, Lower ix) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix, Lower ix) -> Lower ix
forall a b. (a, b) -> b
snd) ([(ix, Lower ix)] -> [(ix, Lower ix)])
-> [(ix, Lower ix)] -> [(ix, Lower ix)]
forall a b. (a -> b) -> a -> b
$ [ix] -> [Lower ix] -> [(ix, Lower ix)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip [ix]
szs [Lower ix]
szls)
let kTotal :: Sz Int
kTotal = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int -> Sz Int) -> Int -> Sz Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Int -> [Int] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
k [Int]
ks
Sz ix
newSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
szl) Dim
n Sz Int
kTotal
let load :: Monad n =>
Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()
load :: Scheduler n ()
-> Int
-> (Int -> e -> n ())
-> (Int -> Sz Int -> e -> n ())
-> n ()
load Scheduler n ()
scheduler Int
startAt Int -> e -> n ()
dlWrite Int -> Sz Int -> e -> n ()
_dlSet =
let arrayLoader :: Int -> (Int, Array r ix e) -> n Int
arrayLoader !Int
kAcc (Int
kCur, Array r ix e
arr) = do
Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$
Array r ix e -> (ix -> e -> n ()) -> n ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
Array r ix a -> (ix -> a -> m b) -> m ()
iforM_ Array r ix e
arr ((ix -> e -> n ()) -> n ()) -> (ix -> e -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \ix
ix e
e ->
let i :: Int
i = ix -> Dim -> Int
forall ix. Index ix => ix -> Dim -> Int
getDim' ix
ix Dim
n
ix' :: ix
ix' = ix -> Dim -> Int -> ix
forall ix. Index ix => ix -> Dim -> Int -> ix
setDim' ix
ix Dim
n (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kAcc)
in Int -> e -> n ()
dlWrite (Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') e
e
Int -> n Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
kAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kCur)
in (Int -> (Int, Array r ix e) -> n Int)
-> Int -> [(Int, Array r ix e)] -> n ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
M.foldM_ Int -> (Int, Array r ix e) -> n Int
arrayLoader Int
0 ([(Int, Array r ix e)] -> n ()) -> [(Int, Array r ix e)] -> n ()
forall a b. (a -> b) -> a -> b
$ (Int
k, Array r ix e
a) (Int, Array r ix e)
-> [(Int, Array r ix e)] -> [(Int, Array r ix e)]
forall a. a -> [a] -> [a]
: [Int] -> [Array r ix e] -> [(Int, Array r ix e)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip [Int]
ks [Array r ix e]
arrs
{-# INLINE load #-}
Array DL ix e -> m (Array DL ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array DL ix e -> m (Array DL ix e))
-> Array DL ix e -> m (Array DL ix e)
forall a b. (a -> b) -> a -> b
$
DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = (Array r ix e -> Comp) -> f (Array r ix e) -> Comp
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp f (Array r ix e)
arrsF, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
{-# INLINE concatM #-}
stackSlicesM ::
forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix)
=> Dim
-> f (Array r (Lower ix) e)
-> m (Array DL ix e)
stackSlicesM :: Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
stackSlicesM Dim
dim !f (Array r (Lower ix) e)
arrsF = do
case [Array r (Lower ix) e]
-> Maybe (Array r (Lower ix) e, [Array r (Lower ix) e])
forall a. [a] -> Maybe (a, [a])
L.uncons (f (Array r (Lower ix) e) -> [Array r (Lower ix) e]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList f (Array r (Lower ix) e)
arrsF) of
Maybe (Array r (Lower ix) e, [Array r (Lower ix) e])
Nothing -> Array DL ix e -> m (Array DL ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array DL ix e
forall r ix e. Construct r ix e => Array r ix e
empty
Just (Array r (Lower ix) e
a, [Array r (Lower ix) e]
arrs) -> do
let sz :: Sz (Lower ix)
sz = Array r (Lower ix) e -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) e
a
len :: Sz Int
len = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (f (Array r (Lower ix) e) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
F.length f (Array r (Lower ix) e)
arrsF)
f (Array r (Lower ix) e) -> (Array r (Lower ix) e -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
M.forM_ f (Array r (Lower ix) e)
arrsF ((Array r (Lower ix) e -> m ()) -> m ())
-> (Array r (Lower ix) e -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Array r (Lower ix) e
arr ->
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz (Lower ix)
sz Sz (Lower ix) -> Sz (Lower ix) -> Bool
forall a. Eq a => a -> a -> Bool
== Array r (Lower ix) e -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) e
arr) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ SizeException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz (Lower ix) -> Sz (Lower ix) -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz (Lower ix)
sz (Array r (Lower ix) e -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) e
arr))
Sz ix
newSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM Sz (Lower ix)
sz Dim
dim Sz Int
len
let load :: Monad n =>
Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()
load :: Scheduler n ()
-> Int
-> (Int -> e -> n ())
-> (Int -> Sz Int -> e -> n ())
-> n ()
load Scheduler n ()
scheduler Int
startAt Int -> e -> n ()
dlWrite Int -> Sz Int -> e -> n ()
_dlSet =
let loadIndex :: Int -> Lower ix -> e -> n ()
loadIndex Int
k Lower ix
ix = Int -> e -> n ()
dlWrite (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz (Lower ix -> Dim -> Int -> ix
forall ix. Index ix => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ix Dim
dim Int
k) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt)
arrayLoader :: Int -> Array r (Lower ix) e -> n Int
arrayLoader !Int
k Array r (Lower ix) e
arr = (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> n () -> n Int
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler ((Lower ix -> e -> n ()) -> Array r (Lower ix) e -> n ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(ix -> a -> m b) -> Array r ix a -> m ()
imapM_ (Int -> Lower ix -> e -> n ()
loadIndex Int
k) Array r (Lower ix) e
arr)
{-# INLINE arrayLoader #-}
in (Int -> Array r (Lower ix) e -> n Int)
-> Int -> f (Array r (Lower ix) e) -> n ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
M.foldM_ Int -> Array r (Lower ix) e -> n Int
arrayLoader Int
0 f (Array r (Lower ix) e)
arrsF
{-# INLINE load #-}
Array DL ix e -> m (Array DL ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array DL ix e -> m (Array DL ix e))
-> Array DL ix e -> m (Array DL ix e)
forall a b. (a -> b) -> a -> b
$
DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = (Array r (Lower ix) e -> Comp) -> [Array r (Lower ix) e] -> Comp
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Array r (Lower ix) e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp [Array r (Lower ix) e]
arrs, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
{-# INLINE stackSlicesM #-}
stackOuterSlicesM ::
forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix)
=> f (Array r (Lower ix) e)
-> m (Array DL ix e)
stackOuterSlicesM :: f (Array r (Lower ix) e) -> m (Array DL ix e)
stackOuterSlicesM = Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix) =>
Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
stackSlicesM (Proxy ix -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions (Proxy ix
forall k (t :: k). Proxy t
Proxy :: Proxy ix))
{-# INLINE stackOuterSlicesM #-}
stackInnerSlicesM ::
forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix)
=> f (Array r (Lower ix) e)
-> m (Array DL ix e)
stackInnerSlicesM :: f (Array r (Lower ix) e) -> m (Array DL ix e)
stackInnerSlicesM = Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix) =>
Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
stackSlicesM Dim
1
{-# INLINE stackInnerSlicesM #-}
splitAtM ::
(MonadThrow m, Extract r ix e)
=> Dim
-> Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e)
splitAtM :: Dim
-> Int -> Array r ix e -> m (Array (R r) ix e, Array (R r) ix e)
splitAtM Dim
dim Int
i Array r ix e
arr = do
let 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
ix
eIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
sz Dim
dim Int
i
ix
sIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim Int
i
Array (R r) ix e
arr1 <- ix -> ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> ix -> Array r ix e -> m (Array (R r) ix e)
extractFromToM ix
forall ix. Index ix => ix
zeroIndex ix
eIx Array r ix e
arr
Array (R r) ix e
arr2 <- ix -> ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> ix -> Array r ix e -> m (Array (R r) ix e)
extractFromToM ix
sIx ix
sz Array r ix e
arr
(Array (R r) ix e, Array (R r) ix e)
-> m (Array (R r) ix e, Array (R r) ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array (R r) ix e
arr1, Array (R r) ix e
arr2)
{-# INLINE splitAtM #-}
splitAt' :: Extract r ix e =>
Dim -> Int -> Array r ix e -> (Array (R r) ix e, Array (R r) ix e)
splitAt' :: Dim -> Int -> Array r ix e -> (Array (R r) ix e, Array (R r) ix e)
splitAt' Dim
dim Int
i Array r ix e
arr = (SomeException -> (Array (R r) ix e, Array (R r) ix e))
-> ((Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e))
-> Either SomeException (Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> (Array (R r) ix e, Array (R r) ix e)
forall a e. Exception e => e -> a
throw (Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e)
forall a. a -> a
id (Either SomeException (Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e))
-> Either SomeException (Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e)
forall a b. (a -> b) -> a -> b
$ Dim
-> Int
-> Array r ix e
-> Either SomeException (Array (R r) ix e, Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
Dim
-> Int -> Array r ix e -> m (Array (R r) ix e, Array (R r) ix e)
splitAtM Dim
dim Int
i Array r ix e
arr
{-# INLINE splitAt' #-}
splitExtractM ::
(MonadThrow m, Extract r ix e, Source (R r) ix e)
=> Dim
-> Ix1
-> Sz Ix1
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
Dim
dim Int
startIx1 (Sz Int
extractSzIx1) Array r ix e
arr = do
let Sz ix
szIx = 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
ix
midStartIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim Int
startIx1
ix
midExtractSzIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
szIx Dim
dim Int
extractSzIx1
Array (R r) ix e
midArr <- ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
midStartIx (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
midExtractSzIx) Array r ix e
arr
ix
leftArrSzIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
szIx Dim
dim Int
startIx1
Array (R r) ix e
leftArr <- ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
forall ix. Index ix => ix
zeroIndex (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
leftArrSzIx) Array r ix e
arr
ix
rightArrStartIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim (Int
startIx1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
extractSzIx1)
Array (R r) ix e
rightArr <- ix -> ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> ix -> Array r ix e -> m (Array (R r) ix e)
extractFromToM ix
rightArrStartIx ix
szIx Array r ix e
arr
(Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array (R r) ix e
leftArr, Array (R r) ix e
midArr, Array (R r) ix e
rightArr)
{-# INLINE splitExtractM #-}
deleteRegionM ::
(MonadThrow m, Extract r ix e, Source (R r) ix e)
=> Dim
-> Ix1
-> Sz Ix1
-> Array r ix e
-> m (Array DL ix e)
deleteRegionM :: Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRegionM Dim
dim Int
ix Sz Int
sz Array r ix e
arr = do
(Array (R r) ix e
leftArr, Array (R r) ix e
_, Array (R r) ix e
rightArr) <- Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e) =>
Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
splitExtractM Dim
dim Int
ix Sz Int
sz Array r ix e
arr
Dim -> Array (R r) ix e -> Array (R r) ix e -> m (Array DL ix e)
forall r1 r2 ix e (m :: * -> *).
(MonadThrow m, Source r1 ix e, Source r2 ix e) =>
Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e)
appendM Dim
dim Array (R r) ix e
leftArr Array (R r) ix e
rightArr
{-# INLINE deleteRegionM #-}
deleteRowsM ::
(MonadThrow m, Extract r ix e, Source (R r) ix e, Index (Lower ix))
=> Ix1
-> Sz Ix1
-> Array r ix e
-> m (Array DL ix e)
deleteRowsM :: Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRowsM = Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e) =>
Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRegionM Dim
2
{-# INLINE deleteRowsM #-}
deleteColumnsM ::
(MonadThrow m, Extract r ix e, Source (R r) ix e)
=> Ix1
-> Sz Ix1
-> Array r ix e
-> m (Array DL ix e)
deleteColumnsM :: Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteColumnsM = Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e) =>
Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRegionM Dim
1
{-# INLINE deleteColumnsM #-}
downsample ::
forall r ix e. Source r ix e
=> Stride ix
-> Array r ix e
-> Array DL ix e
downsample :: Stride ix -> Array r ix e -> Array DL ix e
downsample Stride ix
stride Array r ix e
arr =
DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr, dlSize :: Sz ix
dlSize = Sz ix
resultSize, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
where
resultSize :: Sz ix
resultSize = Stride ix -> Sz ix -> Sz ix
forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride ix
stride (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)
strideIx :: ix
strideIx = Stride ix -> ix
forall ix. Stride ix -> ix
unStride Stride ix
stride
unsafeLinearWriteWithStride :: Int -> e
unsafeLinearWriteWithStride =
Array r ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr (ix -> e) -> (Int -> ix) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ix
strideIx (ix -> ix) -> (Int -> ix) -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
resultSize
{-# INLINE unsafeLinearWriteWithStride #-}
load :: Monad m =>
Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()
load :: Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
dlWrite Int -> Sz Int -> e -> m ()
_ =
Scheduler m ()
-> Int -> Int -> (Int -> m e) -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *) b c.
Monad m =>
Scheduler m ()
-> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_
Scheduler m ()
scheduler
Int
startAt
(Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
resultSize)
(e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> (Int -> e) -> Int -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
unsafeLinearWriteWithStride)
Int -> e -> m ()
dlWrite
{-# INLINE load #-}
{-# INLINE downsample #-}
upsample ::
forall r ix e. Load r ix e
=> e
-> Stride ix
-> Array r ix e
-> Array DL ix e
upsample :: e -> Stride ix -> Array r ix e -> Array DL ix e
upsample !e
fillWith Stride ix
safeStride Array r ix e
arr =
DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ())
-> Array DL ix e
DLArray
{ dlComp :: Comp
dlComp = Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr
, dlSize :: Sz ix
dlSize = Sz ix
newsz
, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load
}
where
load :: Monad m =>
Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()
load :: Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
uWrite Int -> Sz Int -> e -> m ()
uSet = do
Int -> Sz Int -> e -> m ()
uSet Int
startAt (Sz ix -> Sz Int
forall ix. Index ix => Sz ix -> Sz Int
toLinearSz Sz ix
newsz) e
fillWith
Scheduler m () -> Array r ix e -> (Int -> e -> m ()) -> m ()
forall r ix e (m :: * -> *).
(Load r ix e, Monad m) =>
Scheduler m () -> Array r ix e -> (Int -> e -> m ()) -> m ()
loadArrayM Scheduler m ()
scheduler Array r ix e
arr (\Int
i -> Int -> e -> m ()
uWrite (Int -> Int
adjustLinearStride (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt)))
{-# INLINE load #-}
adjustLinearStride :: Int -> Int
adjustLinearStride = Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newsz (ix -> Int) -> (Int -> ix) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix
timesStride (ix -> ix) -> (Int -> ix) -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz
{-# INLINE adjustLinearStride #-}
timesStride :: ix -> ix
timesStride !ix
ix = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ix
stride ix
ix
{-# INLINE timesStride #-}
!stride :: ix
stride = Stride ix -> ix
forall ix. Stride ix -> ix
unStride Stride ix
safeStride
!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
!newsz :: Sz ix
newsz = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> ix
timesStride (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz)
{-# INLINE upsample #-}
transformM ::
forall r ix e r' ix' e' a m.
(Mutable r ix e, Source r' ix' e', MonadUnliftIO m, PrimMonad m, MonadThrow m)
=> (Sz ix' -> m (Sz ix, a))
-> (a -> (ix' -> m e') -> ix -> m e)
-> Array r' ix' e'
-> m (Array r ix e)
transformM :: (Sz ix' -> m (Sz ix, a))
-> (a -> (ix' -> m e') -> ix -> m e)
-> Array r' ix' e'
-> m (Array r ix e)
transformM Sz ix' -> m (Sz ix, a)
getSzM a -> (ix' -> m e') -> ix -> m e
getM Array r' ix' e'
arr = do
(Sz ix
sz, a
a) <- Sz ix' -> m (Sz ix, a)
getSzM (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)
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
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' e' -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e'
arr) Sz ix
sz (a -> (ix' -> m e') -> ix -> m e
getM a
a (Array r' ix' e' -> ix' -> m e'
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r' ix' e'
arr))
{-# INLINE transformM #-}
transform' ::
(Source r' ix' e', Index ix)
=> (Sz ix' -> (Sz ix, a))
-> (a -> (ix' -> e') -> ix -> e)
-> Array r' ix' e'
-> Array D ix e
transform' :: (Sz ix' -> (Sz ix, a))
-> (a -> (ix' -> e') -> ix -> e) -> Array r' ix' e' -> Array D ix e
transform' Sz ix' -> (Sz ix, a)
getSz a -> (ix' -> e') -> ix -> e
get Array r' ix' e'
arr = Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r' ix' e' -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e'
arr) Sz ix
sz (a -> (ix' -> e') -> ix -> e
get a
a (Array r' ix' e' -> ix' -> e'
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r' ix' e'
arr))
where
(Sz ix
sz, a
a) = Sz ix' -> (Sz ix, a)
getSz (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)
{-# INLINE transform' #-}
transform2M ::
(Mutable r ix e, Source r1 ix1 e1, Source r2 ix2 e2, MonadUnliftIO m, PrimMonad m, MonadThrow m)
=> (Sz ix1 -> Sz ix2 -> m (Sz ix, a))
-> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e)
-> Array r1 ix1 e1
-> Array r2 ix2 e2
-> m (Array r ix e)
transform2M :: (Sz ix1 -> Sz ix2 -> m (Sz ix, a))
-> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e)
-> Array r1 ix1 e1
-> Array r2 ix2 e2
-> m (Array r ix e)
transform2M Sz ix1 -> Sz ix2 -> m (Sz ix, a)
getSzM a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e
getM Array r1 ix1 e1
arr1 Array r2 ix2 e2
arr2 = do
(Sz ix
sz, a
a) <- Sz ix1 -> Sz ix2 -> m (Sz ix, a)
getSzM (Array r1 ix1 e1 -> Sz ix1
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> Sz ix2
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix2 e2
arr2)
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
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 r1 ix1 e1 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r1 ix1 e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix2 e2 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r2 ix2 e2
arr2) Sz ix
sz (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e
getM a
a (Array r1 ix1 e1 -> ix1 -> m e1
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> ix2 -> m e2
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r2 ix2 e2
arr2))
{-# INLINE transform2M #-}
transform2' ::
(Source r1 ix1 e1, Source r2 ix2 e2, Index ix)
=> (Sz ix1 -> Sz ix2 -> (Sz ix, a))
-> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e)
-> Array r1 ix1 e1
-> Array r2 ix2 e2
-> Array D ix e
transform2' :: (Sz ix1 -> Sz ix2 -> (Sz ix, a))
-> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e)
-> Array r1 ix1 e1
-> Array r2 ix2 e2
-> Array D ix e
transform2' Sz ix1 -> Sz ix2 -> (Sz ix, a)
getSz a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e
get Array r1 ix1 e1
arr1 Array r2 ix2 e2
arr2 =
Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r1 ix1 e1 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r1 ix1 e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix2 e2 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r2 ix2 e2
arr2) Sz ix
sz (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e
get a
a (Array r1 ix1 e1 -> ix1 -> e1
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> ix2 -> e2
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r2 ix2 e2
arr2))
where
(Sz ix
sz, a
a) = Sz ix1 -> Sz ix2 -> (Sz ix, a)
getSz (Array r1 ix1 e1 -> Sz ix1
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> Sz ix2
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix2 e2
arr2)
{-# INLINE transform2' #-}
zoomWithGrid ::
forall r ix e. Source r ix e
=> e
-> Stride ix
-> Array r ix e
-> Array DL ix e
zoomWithGrid :: e -> Stride ix -> Array r ix e -> Array DL ix e
zoomWithGrid e
gridVal (Stride ix
zoomFactor) Array r ix e
arr = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
unsafeMakeLoadArray Comp
Seq Sz ix
newSz (e -> Maybe e
forall a. a -> Maybe a
Just e
gridVal) forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
!kx :: ix
kx = (Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ix
zoomFactor
!lastNewIx :: ix
lastNewIx = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ix
kx (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz (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)
!newSz :: Sz ix
newSz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ix
lastNewIx)
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
_ Int -> e -> m ()
writeElement =
Scheduler m () -> Array r ix e -> (ix -> e -> m ()) -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
Scheduler m () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler m ()
scheduler Array r ix e
arr ((ix -> e -> m ()) -> m ()) -> (ix -> e -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !ix
ix !e
e ->
let !kix :: ix
kix = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ix
ix ix
kx
in (ix -> m ()) -> Array D ix ix -> m ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_ (\ !ix
ix' -> Int -> e -> m ()
writeElement (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') e
e) (Array D ix ix -> m ()) -> Array D ix ix -> m ()
forall a b. (a -> b) -> a -> b
$
Comp -> ix -> ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range Comp
Seq ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ix
kix) ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
kix ix
kx)
{-# INLINE load #-}
{-# INLINE zoomWithGrid #-}
zoom ::
forall r ix e. Source r ix e
=> Stride ix
-> Array r ix e
-> Array DL ix e
zoom :: Stride ix -> Array r ix e -> Array DL ix e
zoom (Stride ix
zoomFactor) Array r ix e
arr = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
unsafeMakeLoadArray Comp
Seq Sz ix
newSz Maybe e
forall a. Maybe a
Nothing forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
!lastNewIx :: ix
lastNewIx = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ix
zoomFactor (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz (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)
!newSz :: Sz ix
newSz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
lastNewIx
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
_ Int -> e -> m ()
writeElement =
Scheduler m () -> Array r ix e -> (ix -> e -> m ()) -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
Scheduler m () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler m ()
scheduler Array r ix e
arr ((ix -> e -> m ()) -> m ()) -> (ix -> e -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !ix
ix !e
e ->
let !kix :: ix
kix = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ix
ix ix
zoomFactor
in (ix -> m ()) -> Array D ix ix -> m ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_ (\ !ix
ix' -> Int -> e -> m ()
writeElement (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') e
e) (Array D ix ix -> m ()) -> Array D ix ix -> m ()
forall a b. (a -> b) -> a -> b
$
Comp -> ix -> ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range Comp
Seq ix
kix ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
kix ix
zoomFactor)
{-# INLINE load #-}
{-# INLINE zoom #-}