{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Massiv.Array.Delayed.Push
( DL(..)
, Array(..)
, toLoadArray
, makeLoadArrayS
, makeLoadArray
, unsafeMakeLoadArray
, unsafeMakeLoadArrayAdjusted
, fromStrideLoad
, appendOuterM
, concatOuterM
) where
import Control.Monad
import Data.Massiv.Core.Common
import Prelude hiding (map, zipWith)
import Control.Scheduler as S (traverse_)
import Data.Foldable as F
#include "massiv.h"
data DL = DL deriving Int -> DL -> ShowS
[DL] -> ShowS
DL -> String
(Int -> DL -> ShowS)
-> (DL -> String) -> ([DL] -> ShowS) -> Show DL
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DL] -> ShowS
$cshowList :: [DL] -> ShowS
show :: DL -> String
$cshow :: DL -> String
showsPrec :: Int -> DL -> ShowS
$cshowsPrec :: Int -> DL -> ShowS
Show
data instance Array DL ix e = DLArray
{ Array DL ix e -> Comp
dlComp :: !Comp
, Array DL ix e -> Sz ix
dlSize :: !(Sz ix)
, Array DL ix e -> Maybe e
dlDefault :: !(Maybe e)
, Array DL ix e
-> forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad :: forall m . Monad m
=> Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> m ()
}
instance Index ix => Construct DL ix e where
setComp :: Comp -> Array DL ix e -> Array DL ix e
setComp Comp
c Array DL ix e
arr = Array DL ix e
R:ArrayDLixe ix e
arr {dlComp :: Comp
dlComp = Comp
c}
{-# INLINE setComp #-}
makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array DL ix e
makeArrayLinear Comp
comp Sz ix
sz Int -> e
f = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray Comp
comp Sz ix
sz Maybe e
forall a. Maybe a
Nothing forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
dlWrite =
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
sz) (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
f) Int -> e -> m ()
dlWrite
{-# INLINE load #-}
{-# INLINE makeArrayLinear #-}
instance Index ix => Resize DL ix where
unsafeResize :: Sz ix' -> Array DL ix e -> Array DL ix' e
unsafeResize !Sz ix'
sz Array DL ix e
arr = Array DL ix e
R:ArrayDLixe ix e
arr { dlSize :: Sz ix'
dlSize = Sz ix'
sz }
{-# INLINE unsafeResize #-}
instance Semigroup (Array DL Ix1 e) where
<> :: Array DL Int e -> Array DL Int e -> Array DL Int e
(<>) = Array DL Int e -> Array DL Int e -> Array DL Int e
forall e. Array DL Int e -> Array DL Int e -> Array DL Int e
mappendDL
{-# INLINE (<>) #-}
instance Monoid (Array DL Ix1 e) where
mempty :: Array DL Int e
mempty =
DLArray :: forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray
{dlComp :: Comp
dlComp = Comp
forall a. Monoid a => a
mempty, dlSize :: Sz Int
dlSize = Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz Int
forall ix. Index ix => ix
zeroIndex, dlDefault :: Maybe e
dlDefault = Maybe e
forall a. Maybe a
Nothing, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad = \Scheduler m ()
_ Int
_ Int -> e -> m ()
_ -> () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()}
{-# INLINE mempty #-}
mappend :: Array DL Int e -> Array DL Int e -> Array DL Int e
mappend = Array DL Int e -> Array DL Int e -> Array DL Int e
forall e. Array DL Int e -> Array DL Int e -> Array DL Int e
mappendDL
{-# INLINE mappend #-}
mconcat :: [Array DL Int e] -> Array DL Int e
mconcat [] = Array DL Int e
forall a. Monoid a => a
mempty
mconcat [Array DL Int e
x] = Array DL Int e
x
mconcat [Array DL Int e
x, Array DL Int e
y] = Array DL Int e
x Array DL Int e -> Array DL Int e -> Array DL Int e
forall a. Semigroup a => a -> a -> a
<> Array DL Int e
y
mconcat [Array DL Int e]
xs = [Array DL Int e] -> Array DL Int e
forall e. [Array DL Int e] -> Array DL Int e
mconcatDL [Array DL Int e]
xs
{-# INLINE mconcat #-}
mconcatDL :: forall e . [Array DL Ix1 e] -> Array DL Ix1 e
mconcatDL :: [Array DL Int e] -> Array DL Int e
mconcatDL ![Array DL Int e]
arrs =
DLArray :: forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = (Array DL Int e -> Comp) -> [Array DL Int e] -> Comp
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Array DL Int e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp [Array DL Int e]
arrs, dlSize :: Sz Int
dlSize = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz Int
k, dlDefault :: Maybe e
dlDefault = Maybe e
forall a. Maybe a
Nothing, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load}
where
!k :: Int
k = (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
0 (Sz Int -> Int
forall ix. Sz ix -> ix
unSz (Sz Int -> Int)
-> (Array DL Int e -> Sz Int) -> Array DL Int e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DL Int e -> Sz Int
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size (Array DL Int e -> Int) -> [Array DL Int e] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Array DL Int e]
arrs)
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
dlWrite =
let loadArr :: Int -> Array DL Int e -> m Int
loadArr !Int
startAtCur DLArray {dlSize = SafeSz kCur, dlDefault, dlLoad} = do
let !endAtCur :: Int
endAtCur = Int
startAtCur Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kCur
Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
(e -> m ()) -> Maybe e -> m ()
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_
(\e
def -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAtCur (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
endAtCur) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> e -> m ()
`dlWrite` e
def))
Maybe e
dlDefault
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad Scheduler m ()
scheduler Int
startAtCur Int -> e -> m ()
dlWrite
Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
endAtCur
{-# INLINE loadArr #-}
in (Int -> Array DL Int e -> m Int) -> Int -> [Array DL Int e] -> m ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ Int -> Array DL Int e -> m Int
loadArr Int
startAt [Array DL Int e]
arrs
{-# INLINE load #-}
{-# INLINE mconcatDL #-}
mappendDL :: forall e . Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
mappendDL :: Array DL Int e -> Array DL Int e -> Array DL Int e
mappendDL (DLArray c1 sz1 mDef1 load1) (DLArray c2 sz2 mDef2 load2) =
DLArray :: forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = Comp
c1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Comp
c2, dlSize :: Sz Int
dlSize = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int
k1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k2), dlDefault :: Maybe e
dlDefault = Maybe e
forall a. Maybe a
Nothing, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load}
where
!k1 :: Int
k1 = Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
sz1
!k2 :: Int
k2 = Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
sz2
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
dlWrite = do
Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
(e -> m ()) -> Maybe e -> m ()
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (\e
def1 -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> e -> m ()
`dlWrite` e
def1)) Maybe e
mDef1
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load1 Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
dlWrite
Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
let startAt2 :: Int
startAt2 = Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k1
(e -> m ()) -> Maybe e -> m ()
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (\e
def2 -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAt2 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
startAt2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k2) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> e -> m ()
`dlWrite` e
def2)) Maybe e
mDef2
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load2 Scheduler m ()
scheduler Int
startAt2 Int -> e -> m ()
dlWrite
{-# INLINE load #-}
{-# INLINE mappendDL #-}
appendOuterM ::
forall ix e m. (Index ix, MonadThrow m)
=> Array DL ix e
-> Array DL ix e
-> m (Array DL ix e)
appendOuterM :: Array DL ix e -> Array DL ix e -> m (Array DL ix e)
appendOuterM (DLArray c1 sz1 mDef1 load1) (DLArray c2 sz2 mDef2 load2) = do
let (!Sz Int
i1, !Sz (Lower ix)
szl1) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz1
(!Sz Int
i2, !Sz (Lower ix)
szl2) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz2
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
Array DL ix e -> m (Array DL ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = Comp
c1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Comp
c2, dlSize :: Sz ix
dlSize = Sz Int -> Sz (Lower ix) -> Sz ix
forall ix. Index ix => Sz Int -> Sz (Lower ix) -> Sz ix
consSz (Sz Int
i1 Sz Int -> Sz Int -> Sz Int
forall a. Num a => a -> a -> a
+ Sz Int
i2) Sz (Lower ix)
szl1, dlDefault :: Maybe e
dlDefault = Maybe e
forall a. Maybe a
Nothing, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load}
where
!k1 :: Int
k1 = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz1
!k2 :: Int
k2 = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz2
load :: Monad n => Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()
load :: Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()
load Scheduler n ()
scheduler !Int
startAt Int -> e -> n ()
dlWrite = do
Scheduler n () -> n () -> n ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$ do
(e -> n ()) -> Maybe e -> n ()
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (\e
def1 -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> n ()) -> n ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> e -> n ()
`dlWrite` e
def1)) Maybe e
mDef1
Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load1 Scheduler n ()
scheduler Int
startAt Int -> e -> n ()
dlWrite
Scheduler n () -> n () -> n ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$ do
let !startAt2 :: Int
startAt2 = Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k1
(e -> n ()) -> Maybe e -> n ()
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (\e
def2 -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> n ()) -> n ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAt2 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
startAt2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k2) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> e -> n ()
`dlWrite` e
def2)) Maybe e
mDef2
Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load2 Scheduler n ()
scheduler Int
startAt2 Int -> e -> n ()
dlWrite
{-# INLINE load #-}
{-# INLINE appendOuterM #-}
concatOuterM ::
forall ix e m. (Index ix, MonadThrow m)
=> [Array DL ix e]
-> m (Array DL ix e)
concatOuterM :: [Array DL ix e] -> m (Array DL ix e)
concatOuterM =
\case
[] -> 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
(Array DL ix e
x:[Array DL ix e]
xs) -> (Array DL ix e -> Array DL ix e -> m (Array DL ix e))
-> Array DL ix e -> [Array DL ix e] -> m (Array DL ix e)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM Array DL ix e -> Array DL ix e -> m (Array DL ix e)
forall ix e (m :: * -> *).
(Index ix, MonadThrow m) =>
Array DL ix e -> Array DL ix e -> m (Array DL ix e)
appendOuterM Array DL ix e
x [Array DL ix e]
xs
{-# INLINE concatOuterM #-}
makeLoadArrayS ::
forall ix e. Index ix
=> Sz ix
-> e
-> (forall m. Monad m =>
(ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArrayS :: Sz ix
-> e
-> (forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArrayS Sz ix
sz e
defVal forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ()
writer = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray Comp
Seq Sz ix
sz (e -> Maybe e
forall a. a -> Maybe a
Just e
defVal) forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
_scheduler !Int
startAt Int -> e -> m ()
uWrite =
let safeWrite :: ix -> e -> m Bool
safeWrite !ix
ix !e
e
| Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix = Int -> e -> m ()
uWrite (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
sz ix
ix) e
e m () -> m Bool -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
| Bool
otherwise = Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
{-# INLINE safeWrite #-}
in (ix -> e -> m Bool) -> m ()
forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ()
writer ix -> e -> m Bool
safeWrite
{-# INLINE load #-}
{-# INLINE makeLoadArrayS #-}
makeLoadArray ::
forall ix e. Index ix
=> Comp
-> Sz ix
-> e
-> (forall m. Monad m =>
Scheduler m () -> (ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArray :: Comp
-> Sz ix
-> e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> (ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArray Comp
comp Sz ix
sz e
defVal forall (m :: * -> *).
Monad m =>
Scheduler m () -> (ix -> e -> m Bool) -> m ()
writer = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray Comp
comp Sz ix
sz (e -> Maybe e
forall a. a -> Maybe a
Just e
defVal) forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler !Int
startAt Int -> e -> m ()
uWrite =
let safeWrite :: ix -> e -> m Bool
safeWrite !ix
ix !e
e
| Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix = Int -> e -> m ()
uWrite (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
sz ix
ix) e
e m () -> m Bool -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
| Bool
otherwise = Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
{-# INLINE safeWrite #-}
in Scheduler m () -> (ix -> e -> m Bool) -> m ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> (ix -> e -> m Bool) -> m ()
writer Scheduler m ()
scheduler ix -> e -> m Bool
safeWrite
{-# INLINE load #-}
{-# INLINE makeLoadArray #-}
unsafeMakeLoadArray ::
Comp
-> Sz ix
-> Maybe e
-> (forall m. Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
unsafeMakeLoadArray :: Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
unsafeMakeLoadArray = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray
{-# INLINE unsafeMakeLoadArray #-}
unsafeMakeLoadArrayAdjusted ::
forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall m. Monad m =>
Scheduler m () -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
unsafeMakeLoadArrayAdjusted :: Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
unsafeMakeLoadArrayAdjusted Comp
comp Sz ix
sz Maybe e
mDefVal forall (m :: * -> *).
Monad m =>
Scheduler m () -> (Int -> e -> m ()) -> m ()
writer = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray Comp
comp Sz ix
sz Maybe e
mDefVal forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler !Int
startAt Int -> e -> m ()
uWrite = Scheduler m () -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> (Int -> e -> m ()) -> m ()
writer Scheduler m ()
scheduler (\Int
i -> Int -> e -> m ()
uWrite (Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i))
{-# INLINE load #-}
{-# INLINE unsafeMakeLoadArrayAdjusted #-}
toLoadArray ::
forall r ix e. Load r ix e
=> Array r ix e
-> Array DL ix e
toLoadArray :: Array r ix e -> Array DL ix e
toLoadArray 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.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) (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 -> Maybe e
forall r ix e. Load r ix e => Array r ix e -> Maybe e
defaultElement Array r ix e
arr) forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler !Int
startAt Int -> e -> m ()
dlWrite = 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 -> e -> m ()
dlWrite (Int -> e -> m ()) -> (Int -> Int) -> Int -> e -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt))
{-# INLINE load #-}
{-# INLINE[1] toLoadArray #-}
{-# RULES "toLoadArray/id" toLoadArray = id #-}
fromStrideLoad ::
forall r ix e. StrideLoad r ix e
=> Stride ix
-> Array r ix e
-> Array DL ix e
fromStrideLoad :: Stride ix -> Array r ix e -> Array DL ix e
fromStrideLoad Stride ix
stride 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.
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
DLArray (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 Maybe e
forall a. Maybe a
Nothing forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
where
newsz :: Sz ix
newsz = 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)
load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler !Int
startAt Int -> e -> m ()
dlWrite =
Scheduler m ()
-> Stride ix -> Sz ix -> Array r ix e -> (Int -> e -> m ()) -> m ()
forall r ix e (m :: * -> *).
(StrideLoad r ix e, Monad m) =>
Scheduler m ()
-> Stride ix -> Sz ix -> Array r ix e -> (Int -> e -> m ()) -> m ()
loadArrayWithStrideM Scheduler m ()
scheduler Stride ix
stride Sz ix
newsz Array r ix e
arr (\ !Int
i -> Int -> e -> m ()
dlWrite (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt))
{-# INLINE load #-}
{-# INLINE fromStrideLoad #-}
instance Index ix => Load DL ix e where
size :: Array DL ix e -> Sz ix
size = Array DL ix e -> Sz ix
forall ix e. Array DL ix e -> Sz ix
dlSize
{-# INLINE size #-}
getComp :: Array DL ix e -> Comp
getComp = Array DL ix e -> Comp
forall ix e. Array DL ix e -> Comp
dlComp
{-# INLINE getComp #-}
loadArrayM :: Scheduler m () -> Array DL ix e -> (Int -> e -> m ()) -> m ()
loadArrayM Scheduler m ()
scheduler DLArray {dlLoad} = Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad Scheduler m ()
scheduler Int
0
{-# INLINE loadArrayM #-}
defaultElement :: Array DL ix e -> Maybe e
defaultElement = Array DL ix e -> Maybe e
forall ix e. Array DL ix e -> Maybe e
dlDefault
{-# INLINE defaultElement #-}
instance Functor (Array DL ix) where
fmap :: (a -> b) -> Array DL ix a -> Array DL ix b
fmap a -> b
f Array DL ix a
arr = Array DL ix a
R:ArrayDLixe ix a
arr {dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> b -> m ()) -> m ()
dlLoad = Array DL ix a
-> (a -> b) -> Scheduler m () -> Int -> (Int -> b -> m ()) -> m ()
forall (m :: * -> *) ix t1 t2.
Monad m =>
Array DL ix t1
-> (t1 -> t2)
-> Scheduler m ()
-> Int
-> (Int -> t2 -> m ())
-> m ()
loadFunctor Array DL ix a
arr a -> b
f, dlDefault :: Maybe b
dlDefault = a -> b
f (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array DL ix a -> Maybe a
forall ix e. Array DL ix e -> Maybe e
dlDefault Array DL ix a
arr}
{-# INLINE fmap #-}
<$ :: a -> Array DL ix b -> Array DL ix a
(<$) a
e Array DL ix b
arr = Array DL ix b
R:ArrayDLixe ix b
arr {dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> a -> m ()) -> m ()
dlLoad = \Scheduler m ()
_ Int
_ Int -> a -> m ()
_ -> () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure (), dlDefault :: Maybe a
dlDefault = a -> Maybe a
forall a. a -> Maybe a
Just a
e}
{-# INLINE (<$) #-}
loadFunctor ::
Monad m => Array DL ix t1 -> (t1 -> t2) -> Scheduler m () -> Int -> (Int -> t2 -> m ()) -> m ()
loadFunctor :: Array DL ix t1
-> (t1 -> t2)
-> Scheduler m ()
-> Int
-> (Int -> t2 -> m ())
-> m ()
loadFunctor Array DL ix t1
arr t1 -> t2
f Scheduler m ()
scheduler Int
startAt Int -> t2 -> m ()
uWrite = Array DL ix t1
-> Scheduler m () -> Int -> (Int -> t1 -> m ()) -> m ()
forall ix e.
Array DL ix e
-> forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
dlLoad Array DL ix t1
arr Scheduler m ()
scheduler Int
startAt (\ !Int
i t1
e -> Int -> t2 -> m ()
uWrite Int
i (t1 -> t2
f t1
e))
{-# INLINE loadFunctor #-}