{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Massiv.Core.List
( LN
, L(..)
, Array(..)
, toListArray
, showsArrayPrec
, showArrayList
, ListItem
) where
import Control.Exception
import Control.Monad (unless, when)
import Control.Scheduler
import Data.Coerce
import Data.Foldable (foldr')
import qualified Data.List as L
import qualified Data.Massiv.Vector.Stream as S
import Data.Massiv.Core.Common
import Data.Typeable
import GHC.Exts
import System.IO.Unsafe (unsafePerformIO)
data LN
type family ListItem ix e :: * where
ListItem Ix1 e = e
ListItem ix e = [ListItem (Lower ix) e]
type instance NestedStruct LN ix e = [ListItem ix e]
newtype instance Array LN ix e = List { Array LN ix e -> [Elt LN ix e]
unList :: [Elt LN ix e] }
instance Construct LN Ix1 e where
setComp :: Comp -> Array LN Ix1 e -> Array LN Ix1 e
setComp Comp
_ = Array LN Ix1 e -> Array LN Ix1 e
forall a. a -> a
id
{-# INLINE setComp #-}
makeArray :: Comp -> Sz Ix1 -> (Ix1 -> e) -> Array LN Ix1 e
makeArray Comp
_ (Sz Ix1
n) Ix1 -> e
f = [e] -> Array LN Ix1 e
coerce ((Ix1 -> e) -> [Ix1] -> [e]
forall a b. (a -> b) -> [a] -> [b]
L.map Ix1 -> e
f [Ix1
0 .. Ix1
n Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
- Ix1
1])
{-# INLINE makeArray #-}
makeArrayLinear :: Comp -> Sz Ix1 -> (Ix1 -> e) -> Array LN Ix1 e
makeArrayLinear Comp
_ (Sz Ix1
n) Ix1 -> e
f = [e] -> Array LN Ix1 e
coerce ((Ix1 -> e) -> [Ix1] -> [e]
forall a b. (a -> b) -> [a] -> [b]
L.map Ix1 -> e
f [Ix1
0 .. Ix1
n Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
- Ix1
1])
{-# INLINE makeArrayLinear #-}
instance {-# OVERLAPPING #-} Nested LN Ix1 e where
fromNested :: NestedStruct LN Ix1 e -> Array LN Ix1 e
fromNested = NestedStruct LN Ix1 e -> Array LN Ix1 e
coerce
{-# INLINE fromNested #-}
toNested :: Array LN Ix1 e -> NestedStruct LN Ix1 e
toNested = Array LN Ix1 e -> NestedStruct LN Ix1 e
coerce
{-# INLINE toNested #-}
instance ( Elt LN ix e ~ Array LN (Lower ix) e
, ListItem ix e ~ [ListItem (Lower ix) e]
, Coercible (Elt LN ix e) (ListItem ix e)
) =>
Nested LN ix e where
fromNested :: NestedStruct LN ix e -> Array LN ix e
fromNested = NestedStruct LN ix e -> Array LN ix e
coerce
{-# INLINE fromNested #-}
toNested :: Array LN ix e -> NestedStruct LN ix e
toNested = Array LN ix e -> NestedStruct LN ix e
coerce
{-# INLINE toNested #-}
instance Nested LN ix e => IsList (Array LN ix e) where
type Item (Array LN ix e) = ListItem ix e
fromList :: [Item (Array LN ix e)] -> Array LN ix e
fromList = [Item (Array LN ix e)] -> Array LN ix e
forall r ix e. Nested r ix e => NestedStruct r ix e -> Array r ix e
fromNested
{-# INLINE fromList #-}
toList :: Array LN ix e -> [Item (Array LN ix e)]
toList = Array LN ix e -> [Item (Array LN ix e)]
forall r ix e. Nested r ix e => Array r ix e -> NestedStruct r ix e
toNested
{-# INLINE toList #-}
data L = L
type instance NestedStruct L ix e = Array LN ix e
data instance Array L ix e = LArray { Array L ix e -> Comp
lComp :: Comp
, Array L ix e -> Array LN ix e
lData :: !(Array LN ix e) }
instance Nested L ix e where
fromNested :: NestedStruct L ix e -> Array L ix e
fromNested = Comp -> Array LN ix e -> Array L ix e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
Seq
{-# INLINE fromNested #-}
toNested :: Array L ix e -> NestedStruct L ix e
toNested = Array L ix e -> NestedStruct L ix e
forall ix e. Array L ix e -> Array LN ix e
lData
{-# INLINE toNested #-}
instance Nested LN ix e => IsList (Array L ix e) where
type Item (Array L ix e) = ListItem ix e
fromList :: [Item (Array L ix e)] -> Array L ix e
fromList = Comp -> Array LN ix e -> Array L ix e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
Seq (Array LN ix e -> Array L ix e)
-> ([ListItem ix e] -> Array LN ix e)
-> [ListItem ix e]
-> Array L ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ListItem ix e] -> Array LN ix e
forall r ix e. Nested r ix e => NestedStruct r ix e -> Array r ix e
fromNested
{-# INLINE fromList #-}
toList :: Array L ix e -> [Item (Array L ix e)]
toList = Array LN ix e -> [ListItem ix e]
forall r ix e. Nested r ix e => Array r ix e -> NestedStruct r ix e
toNested (Array LN ix e -> [ListItem ix e])
-> (Array L ix e -> Array LN ix e)
-> Array L ix e
-> [ListItem ix e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L ix e -> Array LN ix e
forall ix e. Array L ix e -> Array LN ix e
lData
{-# INLINE toList #-}
instance {-# OVERLAPPING #-} Ragged L Ix1 e where
isNull :: Array L Ix1 e -> Bool
isNull = [e] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([e] -> Bool) -> (Array L Ix1 e -> [e]) -> Array L Ix1 e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array LN Ix1 e -> [e]
forall ix e. Array LN ix e -> [Elt LN ix e]
unList (Array LN Ix1 e -> [e])
-> (Array L Ix1 e -> Array LN Ix1 e) -> Array L Ix1 e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix1 e -> Array LN Ix1 e
forall ix e. Array L ix e -> Array LN ix e
lData
{-# INLINE isNull #-}
emptyR :: Comp -> Array L Ix1 e
emptyR Comp
comp = Comp -> Array LN Ix1 e -> Array L Ix1 e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
comp ([Elt LN Ix1 e] -> Array LN Ix1 e
forall ix e. [Elt LN ix e] -> Array LN ix e
List [])
{-# INLINE emptyR #-}
edgeSize :: Array L Ix1 e -> Sz Ix1
edgeSize = Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz (Ix1 -> Sz Ix1)
-> (Array L Ix1 e -> Ix1) -> Array L Ix1 e -> Sz Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [e] -> Ix1
forall (t :: * -> *) a. Foldable t => t a -> Ix1
length ([e] -> Ix1) -> (Array L Ix1 e -> [e]) -> Array L Ix1 e -> Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array LN Ix1 e -> [e]
forall ix e. Array LN ix e -> [Elt LN ix e]
unList (Array LN Ix1 e -> [e])
-> (Array L Ix1 e -> Array LN Ix1 e) -> Array L Ix1 e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix1 e -> Array LN Ix1 e
forall ix e. Array L ix e -> Array LN ix e
lData
{-# INLINE edgeSize #-}
consR :: Elt L Ix1 e -> Array L Ix1 e -> Array L Ix1 e
consR Elt L Ix1 e
x Array L Ix1 e
arr = Array L Ix1 e
R:ArrayLixe Ix1 e
arr { lData :: Array LN Ix1 e
lData = [e] -> Array LN Ix1 e
coerce (e
Elt L Ix1 e
x e -> [e] -> [e]
forall a. a -> [a] -> [a]
: Array LN Ix1 e -> [e]
coerce (Array L Ix1 e -> Array LN Ix1 e
forall ix e. Array L ix e -> Array LN ix e
lData Array L Ix1 e
arr)) }
{-# INLINE consR #-}
unconsR :: Array L Ix1 e -> Maybe (Elt L Ix1 e, Array L Ix1 e)
unconsR LArray {..} =
case [e] -> Maybe (e, [e])
forall a. [a] -> Maybe (a, [a])
L.uncons ([e] -> Maybe (e, [e])) -> [e] -> Maybe (e, [e])
forall a b. (a -> b) -> a -> b
$ Array LN Ix1 e -> [e]
coerce Array LN Ix1 e
lData of
Maybe (e, [e])
Nothing -> Maybe (Elt L Ix1 e, Array L Ix1 e)
forall a. Maybe a
Nothing
Just (e
x, [e]
xs) -> (e, Array L Ix1 e) -> Maybe (e, Array L Ix1 e)
forall a. a -> Maybe a
Just (e
x, Comp -> Array LN Ix1 e -> Array L Ix1 e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
lComp ([e] -> Array LN Ix1 e
coerce [e]
xs))
{-# INLINE unconsR #-}
flattenRagged :: Array L Ix1 e -> Array L Ix1 e
flattenRagged = Array L Ix1 e -> Array L Ix1 e
forall a. a -> a
id
{-# INLINE flattenRagged #-}
generateRaggedM :: Comp -> Sz Ix1 -> (Ix1 -> m e) -> m (Array L Ix1 e)
generateRaggedM !Comp
comp !Sz Ix1
k Ix1 -> m e
f = do
[e]
xs <- Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> [e]
-> (Ix1 -> [e] -> m [e])
-> m [e]
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopDeepM Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz Ix1 -> Ix1
coerce Sz Ix1
k) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) [] ((Ix1 -> [e] -> m [e]) -> m [e]) -> (Ix1 -> [e] -> m [e]) -> m [e]
forall a b. (a -> b) -> a -> b
$ \Ix1
i [e]
acc -> do
e
e <- Ix1 -> m e
f Ix1
i
[e] -> m [e]
forall (m :: * -> *) a. Monad m => a -> m a
return (e
ee -> [e] -> [e]
forall a. a -> [a] -> [a]
:[e]
acc)
Array L Ix1 e -> m (Array L Ix1 e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array L Ix1 e -> m (Array L Ix1 e))
-> Array L Ix1 e -> m (Array L Ix1 e)
forall a b. (a -> b) -> a -> b
$ Comp -> Array LN Ix1 e -> Array L Ix1 e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
comp (Array LN Ix1 e -> Array L Ix1 e)
-> Array LN Ix1 e -> Array L Ix1 e
forall a b. (a -> b) -> a -> b
$ [e] -> Array LN Ix1 e
coerce [e]
xs
{-# INLINE generateRaggedM #-}
loadRagged :: (m () -> m ())
-> (Ix1 -> e -> m a)
-> Ix1
-> Ix1
-> Sz Ix1
-> Array L Ix1 e
-> m ()
loadRagged m () -> m ()
using Ix1 -> e -> m a
uWrite Ix1
start Ix1
end Sz Ix1
sz Array L Ix1 e
xs =
m () -> m ()
using (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Array L Ix1 e
leftOver <-
Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> Array L Ix1 e
-> (Ix1 -> Array L Ix1 e -> m (Array L Ix1 e))
-> m (Array L Ix1 e)
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
start (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
end) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) Array L Ix1 e
xs ((Ix1 -> Array L Ix1 e -> m (Array L Ix1 e)) -> m (Array L Ix1 e))
-> (Ix1 -> Array L Ix1 e -> m (Array L Ix1 e)) -> m (Array L Ix1 e)
forall a b. (a -> b) -> a -> b
$ \Ix1
i Array L Ix1 e
xs' ->
case Array L Ix1 e -> Maybe (Elt L Ix1 e, Array L Ix1 e)
forall r ix e.
Ragged r ix e =>
Array r ix e -> Maybe (Elt r ix e, Array r ix e)
unconsR Array L Ix1 e
xs' of
Maybe (Elt L Ix1 e, Array L Ix1 e)
Nothing -> Array L Ix1 e -> m (Array L Ix1 e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array L Ix1 e -> m (Array L Ix1 e))
-> Array L Ix1 e -> m (Array L Ix1 e)
forall a b. (a -> b) -> a -> b
$! ShapeException -> Array L Ix1 e
forall a e. Exception e => e -> a
throw (Sz Ix1 -> Sz Ix1 -> ShapeException
DimTooShortException Sz Ix1
sz (Array L Ix1 e -> Sz Ix1
forall ix e. Array L ix e -> Sz Ix1
outerLength Array L Ix1 e
xs))
Just (Elt L Ix1 e
y, Array L Ix1 e
ys) -> Ix1 -> e -> m a
uWrite Ix1
i e
Elt L Ix1 e
y m a -> m (Array L Ix1 e) -> m (Array L Ix1 e)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Array L Ix1 e -> m (Array L Ix1 e)
forall (m :: * -> *) a. Monad m => a -> m a
return Array L Ix1 e
ys
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Array L Ix1 e -> Bool
forall r ix e. Ragged r ix e => Array r ix e -> Bool
isNull Array L Ix1 e
leftOver) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> m ()) -> () -> m ()
forall a b. (a -> b) -> a -> b
$! ShapeException -> ()
forall a e. Exception e => e -> a
throw ShapeException
DimTooLongException)
{-# INLINE loadRagged #-}
raggedFormat :: (e -> String) -> String -> Array L Ix1 e -> String
raggedFormat e -> String
f String
_ Array L Ix1 e
arr = [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ String
"[ " String -> [String] -> [String]
forall a. a -> [a] -> [a]
: String -> [String] -> [String]
forall a. a -> [a] -> [a]
L.intersperse String
", " ((e -> String) -> [e] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map e -> String
f (Array LN Ix1 e -> [e]
coerce (Array L Ix1 e -> Array LN Ix1 e
forall ix e. Array L ix e -> Array LN ix e
lData Array L Ix1 e
arr))) [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String
" ]"]
instance (Index ix, Ragged L ix e) => Load L ix e where
size :: Array L ix e -> Sz ix
size = Sz ix -> Sz ix
coerce (Sz ix -> Sz ix)
-> (Array L ix e -> Sz ix) -> Array L ix e -> Sz ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L ix e -> Sz ix
forall r ix e. Ragged r ix e => Array r ix e -> Sz ix
edgeSize
{-# INLINE size #-}
getComp :: Array L ix e -> Comp
getComp = Array L ix e -> Comp
forall ix e. Array L ix e -> Comp
lComp
{-# INLINE getComp #-}
loadArrayM :: Scheduler m () -> Array L ix e -> (Ix1 -> e -> m ()) -> m ()
loadArrayM Scheduler m ()
scheduler Array L ix e
arr Ix1 -> e -> m ()
uWrite =
(m () -> m ())
-> (Ix1 -> e -> m ())
-> Ix1
-> Ix1
-> Sz ix
-> Array L ix e
-> m ()
forall r ix e (m :: * -> *) a.
(Ragged r ix e, Monad m) =>
(m () -> m ())
-> (Ix1 -> e -> m a) -> Ix1 -> Ix1 -> Sz ix -> Array r ix e -> m ()
loadRagged (Scheduler m () -> m () -> m ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler m ()
scheduler) Ix1 -> e -> m ()
uWrite Ix1
0 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) Sz ix
sz Array L ix e
arr
where !sz :: Sz ix
sz = Array L ix e -> Sz ix
forall r ix e. Ragged r ix e => Array r ix e -> Sz ix
edgeSize Array L ix e
arr
{-# INLINE loadArrayM #-}
instance (Index ix, Load L ix e, Ragged L ix e) => Load LN ix e where
size :: Array LN ix e -> Sz ix
size = Array L ix e -> Sz ix
forall r ix e. Ragged r ix e => Array r ix e -> Sz ix
edgeSize (Array L ix e -> Sz ix)
-> (Array LN ix e -> Array L ix e) -> Array LN ix e -> Sz ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> Array LN ix e -> Array L ix e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
Seq
{-# INLINE size #-}
getComp :: Array LN ix e -> Comp
getComp Array LN ix e
_ = Comp
Seq
{-# INLINE getComp #-}
loadArrayM :: Scheduler m () -> Array LN ix e -> (Ix1 -> e -> m ()) -> m ()
loadArrayM Scheduler m ()
scheduler Array LN ix e
arr Ix1 -> e -> m ()
uWrite =
(m () -> m ())
-> (Ix1 -> e -> m ())
-> Ix1
-> Ix1
-> Sz ix
-> Array L ix e
-> m ()
forall r ix e (m :: * -> *) a.
(Ragged r ix e, Monad m) =>
(m () -> m ())
-> (Ix1 -> e -> m a) -> Ix1 -> Ix1 -> Sz ix -> Array r ix e -> m ()
loadRagged (Scheduler m () -> m () -> m ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler m ()
scheduler) Ix1 -> e -> m ()
uWrite Ix1
0 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) Sz ix
sz Array L ix e
arrL
where
!arrL :: Array L ix e
arrL = Comp -> Array LN ix e -> Array L ix e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
Seq Array LN ix e
arr
!sz :: Sz ix
sz = Array L ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array L ix e
arrL
{-# INLINE loadArrayM #-}
outerLength :: Array L ix e -> Sz Int
outerLength :: Array L ix e -> Sz Ix1
outerLength = Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz (Ix1 -> Sz Ix1) -> (Array L ix e -> Ix1) -> Array L ix e -> Sz Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Elt LN ix e] -> Ix1
forall (t :: * -> *) a. Foldable t => t a -> Ix1
length ([Elt LN ix e] -> Ix1)
-> (Array L ix e -> [Elt LN ix e]) -> Array L ix e -> Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array LN ix e -> [Elt LN ix e]
forall ix e. Array LN ix e -> [Elt LN ix e]
unList (Array LN ix e -> [Elt LN ix e])
-> (Array L ix e -> Array LN ix e) -> Array L ix e -> [Elt LN ix e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L ix e -> Array LN ix e
forall ix e. Array L ix e -> Array LN ix e
lData
instance ( Index ix
, Index (Lower ix)
, Ragged L (Lower ix) e
, Elt L ix e ~ Array L (Lower ix) e
, Elt LN ix e ~ Array LN (Lower ix) e
, Coercible (Elt LN ix e) [Elt LN (Lower ix) e]
) =>
Ragged L ix e where
isNull :: Array L ix e -> Bool
isNull = [Array LN (Lower ix) e] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Array LN (Lower ix) e] -> Bool)
-> (Array L ix e -> [Array LN (Lower ix) e])
-> Array L ix e
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array LN ix e -> [Array LN (Lower ix) e]
forall ix e. Array LN ix e -> [Elt LN ix e]
unList (Array LN ix e -> [Array LN (Lower ix) e])
-> (Array L ix e -> Array LN ix e)
-> Array L ix e
-> [Array LN (Lower ix) e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L ix e -> Array LN ix e
forall ix e. Array L ix e -> Array LN ix e
lData
{-# INLINE isNull #-}
emptyR :: Comp -> Array L ix e
emptyR Comp
comp = Comp -> Array LN ix e -> Array L ix e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
comp ([Elt LN ix e] -> Array LN ix e
forall ix e. [Elt LN ix e] -> Array LN ix e
List [])
{-# INLINE emptyR #-}
edgeSize :: Array L ix e -> Sz ix
edgeSize Array L ix e
arr =
ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz
(Ix1 -> Lower ix -> ix
forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim ([Array LN (Lower ix) e] -> Ix1
forall (t :: * -> *) a. Foldable t => t a -> Ix1
length (Array LN ix e -> [Elt LN ix e]
forall ix e. Array LN ix e -> [Elt LN ix e]
unList (Array L ix e -> Array LN ix e
forall ix e. Array L ix e -> Array LN ix e
lData Array L ix e
arr))) (Lower ix -> ix) -> Lower ix -> ix
forall a b. (a -> b) -> a -> b
$
case Array L ix e -> Maybe (Elt L ix e, Array L ix e)
forall r ix e.
Ragged r ix e =>
Array r ix e -> Maybe (Elt r ix e, Array r ix e)
unconsR Array L ix e
arr of
Maybe (Elt L ix e, Array L ix e)
Nothing -> Lower ix
forall ix. Index ix => ix
zeroIndex
Just (Elt L ix e
x, Array L ix e
_) -> Sz (Lower ix) -> Lower ix
coerce (Array L (Lower ix) e -> Sz (Lower ix)
forall r ix e. Ragged r ix e => Array r ix e -> Sz ix
edgeSize Elt L ix e
Array L (Lower ix) e
x))
{-# INLINE edgeSize #-}
consR :: Elt L ix e -> Array L ix e -> Array L ix e
consR (LArray _ x) Array L ix e
arr = Array L ix e
newArr
where
newArr :: Array L ix e
newArr = Array L ix e
R:ArrayLixe ix e
arr {lData :: Array LN ix e
lData = [Array LN (Lower ix) e] -> Array LN ix e
coerce (Array LN (Lower ix) e
x Array LN (Lower ix) e
-> [Array LN (Lower ix) e] -> [Array LN (Lower ix) e]
forall a. a -> [a] -> [a]
: Array LN ix e -> [Array LN (Lower ix) e]
coerce (Array L ix e -> Array LN ix e
forall ix e. Array L ix e -> Array LN ix e
lData Array L ix e
arr))}
{-# INLINE consR #-}
unconsR :: Array L ix e -> Maybe (Elt L ix e, Array L ix e)
unconsR LArray {..} =
case [Array LN (Lower ix) e]
-> Maybe (Array LN (Lower ix) e, [Array LN (Lower ix) e])
forall a. [a] -> Maybe (a, [a])
L.uncons (Array LN ix e -> [Array LN (Lower ix) e]
coerce Array LN ix e
lData) of
Maybe (Array LN (Lower ix) e, [Array LN (Lower ix) e])
Nothing -> Maybe (Elt L ix e, Array L ix e)
forall a. Maybe a
Nothing
Just (Array LN (Lower ix) e
x, [Array LN (Lower ix) e]
xs) ->
let newArr :: Array L ix e
newArr = Comp -> Array LN ix e -> Array L ix e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
lComp ([Array LN (Lower ix) e] -> Array LN ix e
coerce [Array LN (Lower ix) e]
xs)
newX :: Array L (Lower ix) e
newX = Comp -> Array LN (Lower ix) e -> Array L (Lower ix) e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
lComp Array LN (Lower ix) e
x
in (Array L (Lower ix) e, Array L ix e)
-> Maybe (Array L (Lower ix) e, Array L ix e)
forall a. a -> Maybe a
Just (Array L (Lower ix) e
newX, Array L ix e
newArr)
{-# INLINE unconsR #-}
generateRaggedM :: Comp -> Sz ix -> (ix -> m e) -> m (Array L ix e)
generateRaggedM = Comp -> Sz ix -> (ix -> m e) -> m (Array L ix e)
forall ix e (m :: * -> *).
(Elt LN ix e ~ Array LN (Lower ix) e, Index ix, Monad m,
Ragged L (Lower ix) e) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array L ix e)
unsafeGenerateParM
{-# INLINE generateRaggedM #-}
flattenRagged :: Array L ix e -> Array L Ix1 e
flattenRagged Array L ix e
arr = LArray :: forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray {lComp :: Comp
lComp = Array L ix e -> Comp
forall ix e. Array L ix e -> Comp
lComp Array L ix e
arr, lData :: Array LN Ix1 e
lData = [e] -> Array LN Ix1 e
coerce [e]
xs}
where
xs :: [e]
xs = (Array LN (Lower ix) e -> [e]) -> [Array LN (Lower ix) e] -> [e]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Array LN Ix1 e -> [e]
forall ix e. Array LN ix e -> [Elt LN ix e]
unList (Array LN Ix1 e -> [e])
-> (Array LN (Lower ix) e -> Array LN Ix1 e)
-> Array LN (Lower ix) e
-> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix1 e -> Array LN Ix1 e
forall ix e. Array L ix e -> Array LN ix e
lData (Array L Ix1 e -> Array LN Ix1 e)
-> (Array LN (Lower ix) e -> Array L Ix1 e)
-> Array LN (Lower ix) e
-> Array LN Ix1 e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L (Lower ix) e -> Array L Ix1 e
forall r ix e. Ragged r ix e => Array r ix e -> Array r Ix1 e
flattenRagged (Array L (Lower ix) e -> Array L Ix1 e)
-> (Array LN (Lower ix) e -> Array L (Lower ix) e)
-> Array LN (Lower ix) e
-> Array L Ix1 e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> Array LN (Lower ix) e -> Array L (Lower ix) e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray (Array L ix e -> Comp
forall ix e. Array L ix e -> Comp
lComp Array L ix e
arr)) (Array LN ix e -> [Elt LN ix e]
forall ix e. Array LN ix e -> [Elt LN ix e]
unList (Array L ix e -> Array LN ix e
forall ix e. Array L ix e -> Array LN ix e
lData Array L ix e
arr))
{-# INLINE flattenRagged #-}
loadRagged :: (m () -> m ())
-> (Ix1 -> e -> m a) -> Ix1 -> Ix1 -> Sz ix -> Array L ix e -> m ()
loadRagged m () -> m ()
using Ix1 -> e -> m a
uWrite Ix1
start Ix1
end Sz ix
sz Array L ix e
xs = do
let (Sz Ix1
k, Sz (Lower ix)
szL) = Sz ix -> (Sz Ix1, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz Sz ix
sz
step :: Ix1
step = Sz (Lower ix) -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz (Lower ix)
szL
isZero :: Bool
isZero = Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz Ix1 -> Ix1 -> Bool
forall a. Eq a => a -> a -> Bool
== Ix1
0
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool
isZero Bool -> Bool -> Bool
&& Bool -> Bool
not (Array L Ix1 e -> Bool
forall r ix e. Ragged r ix e => Array r ix e -> Bool
isNull (Array L ix e -> Array L Ix1 e
forall r ix e. Ragged r ix e => Array r ix e -> Array r Ix1 e
flattenRagged Array L ix e
xs))) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> m ()) -> () -> m ()
forall a b. (a -> b) -> a -> b
$! ShapeException -> ()
forall a e. Exception e => e -> a
throw ShapeException
DimTooLongException)
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
isZero (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Array L ix e
leftOver <-
Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> Array L ix e
-> (Ix1 -> Array L ix e -> m (Array L ix e))
-> m (Array L ix e)
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
start (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
end) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
step) Array L ix e
xs ((Ix1 -> Array L ix e -> m (Array L ix e)) -> m (Array L ix e))
-> (Ix1 -> Array L ix e -> m (Array L ix e)) -> m (Array L ix e)
forall a b. (a -> b) -> a -> b
$ \Ix1
i Array L ix e
zs ->
case Array L ix e -> Maybe (Elt L ix e, Array L ix e)
forall r ix e.
Ragged r ix e =>
Array r ix e -> Maybe (Elt r ix e, Array r ix e)
unconsR Array L ix e
zs of
Maybe (Elt L ix e, Array L ix e)
Nothing -> Array L ix e -> m (Array L ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array L ix e -> m (Array L ix e))
-> Array L ix e -> m (Array L ix e)
forall a b. (a -> b) -> a -> b
$! ShapeException -> Array L ix e
forall a e. Exception e => e -> a
throw (Sz Ix1 -> Sz Ix1 -> ShapeException
DimTooShortException Sz Ix1
k (Array L ix e -> Sz Ix1
forall ix e. Array L ix e -> Sz Ix1
outerLength Array L ix e
xs))
Just (Elt L ix e
y, Array L ix e
ys) -> do
()
_ <- (m () -> m ())
-> (Ix1 -> e -> m a)
-> Ix1
-> Ix1
-> Sz (Lower ix)
-> Array L (Lower ix) e
-> m ()
forall r ix e (m :: * -> *) a.
(Ragged r ix e, Monad m) =>
(m () -> m ())
-> (Ix1 -> e -> m a) -> Ix1 -> Ix1 -> Sz ix -> Array r ix e -> m ()
loadRagged m () -> m ()
using Ix1 -> e -> m a
uWrite Ix1
i (Ix1
i Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
step) Sz (Lower ix)
szL Elt L ix e
Array L (Lower ix) e
y
Array L ix e -> m (Array L ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return Array L ix e
ys
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Array L ix e -> Bool
forall r ix e. Ragged r ix e => Array r ix e -> Bool
isNull Array L ix e
leftOver) (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> m ()) -> () -> m ()
forall a b. (a -> b) -> a -> b
$! ShapeException -> ()
forall a e. Exception e => e -> a
throw ShapeException
DimTooLongException)
{-# INLINE loadRagged #-}
raggedFormat :: (e -> String) -> String -> Array L ix e -> String
raggedFormat e -> String
f String
sep (LArray comp xs) =
(String -> Array LN (Lower ix) e -> String)
-> String -> [Array LN (Lower ix) e] -> String
forall a. (String -> a -> String) -> String -> [a] -> String
showN (\String
s Array LN (Lower ix) e
y -> (e -> String) -> String -> Array L (Lower ix) e -> String
forall r ix e.
Ragged r ix e =>
(e -> String) -> String -> Array r ix e -> String
raggedFormat e -> String
f String
s (Comp -> Array LN (Lower ix) e -> Array L (Lower ix) e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
comp Array LN (Lower ix) e
y :: Array L (Lower ix) e)) String
sep (Array LN ix e -> [Array LN (Lower ix) e]
coerce Array LN ix e
xs)
unsafeGenerateParM ::
(Elt LN ix e ~ Array LN (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e)
=> Comp
-> Sz ix
-> (ix -> m e)
-> m (Array L ix e)
unsafeGenerateParM :: Comp -> Sz ix -> (ix -> m e) -> m (Array L ix e)
unsafeGenerateParM Comp
comp !Sz ix
sz ix -> m e
f = do
[[Array LN (Lower ix) e]]
res <- [m [Array LN (Lower ix) e]] -> m [[Array LN (Lower ix) e]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([m [Array LN (Lower ix) e]] -> m [[Array LN (Lower ix) e]])
-> [m [Array LN (Lower ix) e]] -> m [[Array LN (Lower ix) e]]
forall a b. (a -> b) -> a -> b
$ IO [m [Array LN (Lower ix) e]] -> [m [Array LN (Lower ix) e]]
forall a. IO a -> a
unsafePerformIO (IO [m [Array LN (Lower ix) e]] -> [m [Array LN (Lower ix) e]])
-> IO [m [Array LN (Lower ix) e]] -> [m [Array LN (Lower ix) e]]
forall a b. (a -> b) -> a -> b
$ do
let !(Sz Ix1
ksz, Sz (Lower ix)
szL) = Sz ix -> (Sz Ix1, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz Sz ix
sz
!k :: Ix1
k = Sz Ix1 -> Ix1
forall ix. Sz ix -> ix
unSz Sz Ix1
ksz
Comp
-> (Scheduler IO (m [Array LN (Lower ix) e]) -> IO ())
-> IO [m [Array LN (Lower ix) e]]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler m a -> m b) -> m [a]
withScheduler Comp
comp ((Scheduler IO (m [Array LN (Lower ix) e]) -> IO ())
-> IO [m [Array LN (Lower ix) e]])
-> (Scheduler IO (m [Array LN (Lower ix) e]) -> IO ())
-> IO [m [Array LN (Lower ix) e]]
forall a b. (a -> b) -> a -> b
$ \ Scheduler IO (m [Array LN (Lower ix) e])
scheduler ->
Ix1 -> Ix1 -> (Ix1 -> Ix1 -> IO ()) -> IO ()
forall a. Ix1 -> Ix1 -> (Ix1 -> Ix1 -> a) -> a
splitLinearly (Scheduler IO (m [Array LN (Lower ix) e]) -> Ix1
forall (m :: * -> *) a. Scheduler m a -> Ix1
numWorkers Scheduler IO (m [Array LN (Lower ix) e])
scheduler) Ix1
k ((Ix1 -> Ix1 -> IO ()) -> IO ()) -> (Ix1 -> Ix1 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ix1
chunkLength Ix1
slackStart -> do
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
slackStart) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
chunkLength) ((Ix1 -> IO ()) -> IO ()) -> (Ix1 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Ix1
start ->
Scheduler IO (m [Array LN (Lower ix) e])
-> IO (m [Array LN (Lower ix) e]) -> IO ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler IO (m [Array LN (Lower ix) e])
scheduler (IO (m [Array LN (Lower ix) e]) -> IO ())
-> IO (m [Array LN (Lower ix) e]) -> IO ()
forall a b. (a -> b) -> a -> b
$ do
[m (Array LN (Lower ix) e)]
res <- Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> [m (Array LN (Lower ix) e)]
-> (Ix1
-> [m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)])
-> IO [m (Array LN (Lower ix) e)]
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopDeepM Ix1
start (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< (Ix1
start Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
chunkLength)) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) [] ((Ix1
-> [m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)])
-> IO [m (Array LN (Lower ix) e)])
-> (Ix1
-> [m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)])
-> IO [m (Array LN (Lower ix) e)]
forall a b. (a -> b) -> a -> b
$ \Ix1
i [m (Array LN (Lower ix) e)]
acc ->
[m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((Array L (Lower ix) e -> Array LN (Lower ix) e)
-> m (Array L (Lower ix) e) -> m (Array LN (Lower ix) e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array L (Lower ix) e -> Array LN (Lower ix) e
forall ix e. Array L ix e -> Array LN ix e
lData (Comp
-> Sz (Lower ix) -> (Lower ix -> m e) -> m (Array L (Lower ix) e)
forall r ix e (m :: * -> *).
(Ragged r ix e, Monad m) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateRaggedM Comp
Seq Sz (Lower ix)
szL (\ !Lower ix
ixL -> ix -> m e
f (Ix1 -> Lower ix -> ix
forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i Lower ix
ixL)))m (Array LN (Lower ix) e)
-> [m (Array LN (Lower ix) e)] -> [m (Array LN (Lower ix) e)]
forall a. a -> [a] -> [a]
:[m (Array LN (Lower ix) e)]
acc)
m [Array LN (Lower ix) e] -> IO (m [Array LN (Lower ix) e])
forall (m :: * -> *) a. Monad m => a -> m a
return (m [Array LN (Lower ix) e] -> IO (m [Array LN (Lower ix) e]))
-> m [Array LN (Lower ix) e] -> IO (m [Array LN (Lower ix) e])
forall a b. (a -> b) -> a -> b
$! [m (Array LN (Lower ix) e)] -> m [Array LN (Lower ix) e]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [m (Array LN (Lower ix) e)]
res
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
slackStart Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
k) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
Scheduler IO (m [Array LN (Lower ix) e])
-> IO (m [Array LN (Lower ix) e]) -> IO ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler IO (m [Array LN (Lower ix) e])
scheduler (IO (m [Array LN (Lower ix) e]) -> IO ())
-> IO (m [Array LN (Lower ix) e]) -> IO ()
forall a b. (a -> b) -> a -> b
$ do
[m (Array LN (Lower ix) e)]
res <- Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> [m (Array LN (Lower ix) e)]
-> (Ix1
-> [m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)])
-> IO [m (Array LN (Lower ix) e)]
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopDeepM Ix1
slackStart (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
k) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) [] ((Ix1
-> [m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)])
-> IO [m (Array LN (Lower ix) e)])
-> (Ix1
-> [m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)])
-> IO [m (Array LN (Lower ix) e)]
forall a b. (a -> b) -> a -> b
$ \Ix1
i [m (Array LN (Lower ix) e)]
acc ->
[m (Array LN (Lower ix) e)] -> IO [m (Array LN (Lower ix) e)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((Array L (Lower ix) e -> Array LN (Lower ix) e)
-> m (Array L (Lower ix) e) -> m (Array LN (Lower ix) e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array L (Lower ix) e -> Array LN (Lower ix) e
forall ix e. Array L ix e -> Array LN ix e
lData (Comp
-> Sz (Lower ix) -> (Lower ix -> m e) -> m (Array L (Lower ix) e)
forall r ix e (m :: * -> *).
(Ragged r ix e, Monad m) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateRaggedM Comp
Seq Sz (Lower ix)
szL (\ !Lower ix
ixL -> ix -> m e
f (Ix1 -> Lower ix -> ix
forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i Lower ix
ixL)))m (Array LN (Lower ix) e)
-> [m (Array LN (Lower ix) e)] -> [m (Array LN (Lower ix) e)]
forall a. a -> [a] -> [a]
:[m (Array LN (Lower ix) e)]
acc)
m [Array LN (Lower ix) e] -> IO (m [Array LN (Lower ix) e])
forall (m :: * -> *) a. Monad m => a -> m a
return (m [Array LN (Lower ix) e] -> IO (m [Array LN (Lower ix) e]))
-> m [Array LN (Lower ix) e] -> IO (m [Array LN (Lower ix) e])
forall a b. (a -> b) -> a -> b
$! [m (Array LN (Lower ix) e)] -> m [Array LN (Lower ix) e]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [m (Array LN (Lower ix) e)]
res
Array L ix e -> m (Array L ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array L ix e -> m (Array L ix e))
-> Array L ix e -> m (Array L ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> Array LN ix e -> Array L ix e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
comp (Array LN ix e -> Array L ix e) -> Array LN ix e -> Array L ix e
forall a b. (a -> b) -> a -> b
$ [Elt LN ix e] -> Array LN ix e
forall ix e. [Elt LN ix e] -> Array LN ix e
List ([Elt LN ix e] -> Array LN ix e) -> [Elt LN ix e] -> Array LN ix e
forall a b. (a -> b) -> a -> b
$ [[Array LN (Lower ix) e]] -> [Array LN (Lower ix) e]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Array LN (Lower ix) e]]
res
{-# INLINE unsafeGenerateParM #-}
instance {-# OVERLAPPING #-} Construct L Ix1 e where
setComp :: Comp -> Array L Ix1 e -> Array L Ix1 e
setComp Comp
c Array L Ix1 e
arr = Array L Ix1 e
R:ArrayLixe Ix1 e
arr { lComp :: Comp
lComp = Comp
c }
{-# INLINE setComp #-}
makeArray :: Comp -> Sz Ix1 -> (Ix1 -> e) -> Array L Ix1 e
makeArray Comp
comp Sz Ix1
sz Ix1 -> e
f = Comp -> Array LN Ix1 e -> Array L Ix1 e
forall ix e. Comp -> Array LN ix e -> Array L ix e
LArray Comp
comp (Array LN Ix1 e -> Array L Ix1 e)
-> Array LN Ix1 e -> Array L Ix1 e
forall a b. (a -> b) -> a -> b
$ [Elt LN Ix1 e] -> Array LN Ix1 e
forall ix e. [Elt LN ix e] -> Array LN ix e
List ([Elt LN Ix1 e] -> Array LN Ix1 e)
-> [Elt LN Ix1 e] -> Array LN Ix1 e
forall a b. (a -> b) -> a -> b
$ IO [e] -> [e]
forall a. IO a -> a
unsafePerformIO (IO [e] -> [e]) -> IO [e] -> [e]
forall a b. (a -> b) -> a -> b
$
Comp -> (Scheduler IO e -> IO ()) -> IO [e]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler m a -> m b) -> m [a]
withScheduler Comp
comp ((Scheduler IO e -> IO ()) -> IO [e])
-> (Scheduler IO e -> IO ()) -> IO [e]
forall a b. (a -> b) -> a -> b
$ \Scheduler IO e
scheduler ->
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz Ix1 -> Ix1
coerce Sz Ix1
sz) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) (Scheduler IO e -> IO e -> IO ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler IO e
scheduler (IO e -> IO ()) -> (Ix1 -> IO e) -> Ix1 -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> IO e
forall (m :: * -> *) a. Monad m => a -> m a
return (e -> IO e) -> (Ix1 -> e) -> Ix1 -> IO e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> e
f)
{-# INLINE makeArray #-}
instance ( Index ix
, Ragged L ix e
, Ragged L (Lower ix) e
, Elt L ix e ~ Array L (Lower ix) e
) =>
Construct L ix e where
setComp :: Comp -> Array L ix e -> Array L ix e
setComp Comp
c Array L ix e
arr = Array L ix e
R:ArrayLixe ix e
arr {lComp :: Comp
lComp = Comp
c}
{-# INLINE setComp #-}
makeArray :: Comp -> Sz ix -> (ix -> e) -> Array L ix e
makeArray = Comp -> Sz ix -> (ix -> e) -> Array L ix e
forall r ix e.
(Ragged r ix e, Ragged r (Lower ix) e,
Elt r ix e ~ Array r (Lower ix) e) =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
unsafeGenerateN
{-# INLINE makeArray #-}
unsafeGenerateN ::
( Ragged r ix e
, Ragged r (Lower ix) e
, Elt r ix e ~ Array r (Lower ix) e )
=> Comp
-> Sz ix
-> (ix -> e)
-> Array r ix e
unsafeGenerateN :: Comp -> Sz ix -> (ix -> e) -> Array r ix e
unsafeGenerateN Comp
comp Sz ix
sz ix -> e
f = IO (Array r ix e) -> Array r ix e
forall a. IO a -> a
unsafePerformIO (IO (Array r ix e) -> Array r ix e)
-> IO (Array r ix e) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ do
let !(Sz Ix1
m, Sz (Lower ix)
szL) = Sz ix -> (Sz Ix1, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz Sz ix
sz
[Array r (Lower ix) e]
xs <- Comp
-> (Scheduler IO (Array r (Lower ix) e) -> IO ())
-> IO [Array r (Lower ix) e]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler m a -> m b) -> m [a]
withScheduler Comp
comp ((Scheduler IO (Array r (Lower ix) e) -> IO ())
-> IO [Array r (Lower ix) e])
-> (Scheduler IO (Array r (Lower ix) e) -> IO ())
-> IO [Array r (Lower ix) e]
forall a b. (a -> b) -> a -> b
$ \Scheduler IO (Array r (Lower ix) e)
scheduler ->
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz Ix1 -> Ix1
coerce Sz Ix1
m) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) ((Ix1 -> IO ()) -> IO ()) -> (Ix1 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ix1
i -> Scheduler IO (Array r (Lower ix) e)
-> IO (Array r (Lower ix) e) -> IO ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler IO (Array r (Lower ix) e)
scheduler (IO (Array r (Lower ix) e) -> IO ())
-> IO (Array r (Lower ix) e) -> IO ()
forall a b. (a -> b) -> a -> b
$
Comp
-> Sz (Lower ix) -> (Lower ix -> IO e) -> IO (Array r (Lower ix) e)
forall r ix e (m :: * -> *).
(Ragged r ix e, Monad m) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateRaggedM Comp
comp Sz (Lower ix)
szL ((Lower ix -> IO e) -> IO (Array r (Lower ix) e))
-> (Lower ix -> IO e) -> IO (Array r (Lower ix) e)
forall a b. (a -> b) -> a -> b
$ \Lower ix
ix -> e -> IO e
forall (m :: * -> *) a. Monad m => a -> m a
return (e -> IO e) -> e -> IO e
forall a b. (a -> b) -> a -> b
$ ix -> e
f (Ix1 -> Lower ix -> ix
forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i Lower ix
ix)
Array r ix e -> IO (Array r ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array r ix e -> IO (Array r ix e))
-> Array r ix e -> IO (Array r ix e)
forall a b. (a -> b) -> a -> b
$! (Array r (Lower ix) e -> Array r ix e -> Array r ix e)
-> Array r ix e -> [Array r (Lower ix) e] -> Array r ix e
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' Array r (Lower ix) e -> Array r ix e -> Array r ix e
forall r ix e.
Ragged r ix e =>
Elt r ix e -> Array r ix e -> Array r ix e
consR (Comp -> Array r ix e
forall r ix e. Ragged r ix e => Comp -> Array r ix e
emptyR Comp
comp) [Array r (Lower ix) e]
xs
{-# INLINE unsafeGenerateN #-}
toListArray :: (Construct L ix e, Source r ix e)
=> Array r ix e
-> Array L ix e
toListArray :: Array r ix e -> Array L ix e
toListArray !Array r ix e
arr = Comp -> Sz ix -> (ix -> e) -> Array L 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) (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 -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr)
{-# INLINE toListArray #-}
instance (Ragged L ix e, Show e) => Show (Array L ix e) where
showsPrec :: Ix1 -> Array L ix e -> ShowS
showsPrec = Proxy L -> Ix1 -> Array L ix e -> ShowS
forall r ix e.
(Ragged L ix e, Typeable r, Show e) =>
Proxy r -> Ix1 -> Array L ix e -> ShowS
showsArrayLAsPrec (Proxy L
forall k (t :: k). Proxy t
Proxy :: Proxy L)
instance (Ragged L ix e, Show e) => Show (Array LN ix e) where
show :: Array LN ix e -> String
show Array LN ix e
arr = String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (e -> String) -> String -> Array L ix e -> String
forall r ix e.
Ragged r ix e =>
(e -> String) -> String -> Array r ix e -> String
raggedFormat e -> String
forall a. Show a => a -> String
show String
"\n " Array L ix e
arrL
where arrL :: Array L ix e
arrL = NestedStruct L ix e -> Array L ix e
forall r ix e. Nested r ix e => NestedStruct r ix e -> Array r ix e
fromNested NestedStruct L ix e
Array LN ix e
arr :: Array L ix e
showN :: (String -> a -> String) -> String -> [a] -> String
showN :: (String -> a -> String) -> String -> [a] -> String
showN String -> a -> String
_ String
_ [] = String
"[ ]"
showN String -> a -> String
fShow String
lnPrefix [a]
ls =
[String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat
([String
"[ "] [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++
String -> [String] -> [String]
forall a. a -> [a] -> [a]
L.intersperse (String
lnPrefix String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", ") ((a -> String) -> [a] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String -> a -> String
fShow (String
lnPrefix String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" ")) [a]
ls) [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String
lnPrefix, String
"]"])
showsArrayLAsPrec ::
forall r ix e. (Ragged L ix e, Typeable r, Show e)
=> Proxy r
-> Int
-> Array L ix e
-> ShowS
showsArrayLAsPrec :: Proxy r -> Ix1 -> Array L ix e -> ShowS
showsArrayLAsPrec Proxy r
pr Ix1
n Array L ix e
arr =
ShowS
opp ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
(String
"Array " String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
TypeRep -> ShowS
showsTypeRep (Proxy r -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep Proxy r
pr) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
(Char
' 'Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
Ix1 -> Comp -> ShowS
forall a. Show a => Ix1 -> a -> ShowS
showsPrec Ix1
1 (Array L ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array L ix e
arr) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
" (" String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ShowS
forall a. Show a => a -> ShowS
shows (Array L ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array L ix e
arr) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
")\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array LN ix e -> ShowS
forall a. Show a => a -> ShowS
shows NestedStruct L ix e
Array LN ix e
lnarr ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
clp
where
(ShowS
opp, ShowS
clp) =
if Ix1
n Ix1 -> Ix1 -> Bool
forall a. Eq a => a -> a -> Bool
== Ix1
0
then (ShowS
forall a. a -> a
id, ShowS
forall a. a -> a
id)
else ((Char
'('Char -> ShowS
forall a. a -> [a] -> [a]
:), (String
"\n)" String -> ShowS
forall a. [a] -> [a] -> [a]
++))
lnarr :: NestedStruct L ix e
lnarr = Array L ix e -> NestedStruct L ix e
forall r ix e. Nested r ix e => Array r ix e -> NestedStruct r ix e
toNested Array L ix e
arr
showsArrayPrec ::
forall r r' ix ix' e. (Ragged L ix' e, Load r ix e, Source r' ix' e, Show e)
=> (Array r ix e -> Array r' ix' e)
-> Int
-> Array r ix e
-> ShowS
showsArrayPrec :: (Array r ix e -> Array r' ix' e) -> Ix1 -> Array r ix e -> ShowS
showsArrayPrec Array r ix e -> Array r' ix' e
f Ix1
n Array r ix e
arr = Proxy r -> Ix1 -> Array L ix' e -> ShowS
forall r ix e.
(Ragged L ix e, Typeable r, Show e) =>
Proxy r -> Ix1 -> Array L ix e -> ShowS
showsArrayLAsPrec (Proxy r
forall k (t :: k). Proxy t
Proxy :: Proxy r) Ix1
n Array L ix' e
larr
where
arr' :: Array r' ix' e
arr' = Array r ix e -> Array r' ix' e
f Array r ix e
arr
larr :: Array L ix' e
larr = Comp -> Sz ix' -> (ix' -> e) -> Array L 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') (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 -> ix' -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r' ix' e
arr') :: Array L ix' e
showArrayList
:: Show arr => [arr] -> String -> String
showArrayList :: [arr] -> ShowS
showArrayList [arr]
arrs = (Char
'['Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [arr] -> ShowS
forall a. Show a => [a] -> ShowS
go [arr]
arrs ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
']'Char -> ShowS
forall a. a -> [a] -> [a]
:)
where
go :: [a] -> ShowS
go [] = ShowS
forall a. a -> a
id
go [a
x] = (Char
' 'Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ShowS
forall a. Show a => a -> ShowS
shows a
x ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
'\n'Char -> ShowS
forall a. a -> [a] -> [a]
:)
go (a
x:[a]
xs) = (Char
' 'Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ShowS
forall a. Show a => a -> ShowS
shows a
x ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"\n," String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> ShowS
go [a]
xs
instance {-# OVERLAPPING #-} OuterSlice L Ix1 e where
unsafeOuterSlice :: Array L Ix1 e -> Ix1 -> Elt L Ix1 e
unsafeOuterSlice (LArray _ xs) = (Array LN Ix1 e -> [e]
coerce Array LN Ix1 e
xs [e] -> Ix1 -> e
forall a. [a] -> Ix1 -> a
!!)
{-# INLINE unsafeOuterSlice #-}
instance Ragged L ix e => OuterSlice L ix e where
unsafeOuterSlice :: Array L ix e -> Ix1 -> Elt L ix e
unsafeOuterSlice Array L ix e
arr' Ix1
i = Ix1 -> Array L ix e -> Elt L ix e
go Ix1
0 Array L ix e
arr'
where
go :: Ix1 -> Array L ix e -> Elt L ix e
go Ix1
n Array L ix e
arr =
case Array L ix e -> Maybe (Elt L ix e, Array L ix e)
forall r ix e.
Ragged r ix e =>
Array r ix e -> Maybe (Elt r ix e, Array r ix e)
unconsR Array L ix e
arr of
Maybe (Elt L ix e, Array L ix e)
Nothing -> IndexException -> Elt L ix e
forall a e. Exception e => e -> a
throw (IndexException -> Elt L ix e) -> IndexException -> Elt L ix e
forall a b. (a -> b) -> a -> b
$ Sz Ix1 -> Ix1 -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (Ix1 -> Sz Ix1
forall ix. Index ix => ix -> Sz ix
Sz (ix -> Ix1
forall ix. Index ix => ix -> Ix1
headDim (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array L ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array L ix e
arr')))) Ix1
i
Just (Elt L ix e
x, Array L ix e
_) | Ix1
n Ix1 -> Ix1 -> Bool
forall a. Eq a => a -> a -> Bool
== Ix1
i -> Elt L ix e
x
Just (Elt L ix e
_, Array L ix e
xs) -> Ix1 -> Array L ix e -> Elt L ix e
go (Ix1
n Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) Array L ix e
xs
{-# INLINE unsafeOuterSlice #-}
instance Stream LN Ix1 e where
toStream :: Array LN Ix1 e -> Steps Id e
toStream = [e] -> Steps Id e
forall (m :: * -> *) e. Monad m => [e] -> Steps m e
S.fromList ([e] -> Steps Id e)
-> (Array LN Ix1 e -> [e]) -> Array LN Ix1 e -> Steps Id e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array LN Ix1 e -> [e]
coerce
{-# INLINE toStream #-}
toStreamIx :: Array LN Ix1 e -> Steps Id (Ix1, e)
toStreamIx = Steps Id e -> Steps Id (Ix1, e)
forall (m :: * -> *) e. Monad m => Steps m e -> Steps m (Ix1, e)
S.indexed (Steps Id e -> Steps Id (Ix1, e))
-> (Array LN Ix1 e -> Steps Id e)
-> Array LN Ix1 e
-> Steps Id (Ix1, e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [e] -> Steps Id e
forall (m :: * -> *) e. Monad m => [e] -> Steps m e
S.fromList ([e] -> Steps Id e)
-> (Array LN Ix1 e -> [e]) -> Array LN Ix1 e -> Steps Id e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array LN Ix1 e -> [e]
coerce
{-# INLINE toStreamIx #-}
instance Stream L Ix1 e where
toStream :: Array L Ix1 e -> Steps Id e
toStream = Array LN Ix1 e -> Steps Id e
forall r ix e. Stream r ix e => Array r ix e -> Steps Id e
toStream (Array LN Ix1 e -> Steps Id e)
-> (Array L Ix1 e -> Array LN Ix1 e) -> Array L Ix1 e -> Steps Id e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix1 e -> Array LN Ix1 e
forall ix e. Array L ix e -> Array LN ix e
lData
{-# INLINE toStream #-}
toStreamIx :: Array L Ix1 e -> Steps Id (Ix1, e)
toStreamIx = Array LN Ix1 e -> Steps Id (Ix1, e)
forall r ix e. Stream r ix e => Array r ix e -> Steps Id (ix, e)
toStreamIx (Array LN Ix1 e -> Steps Id (Ix1, e))
-> (Array L Ix1 e -> Array LN Ix1 e)
-> Array L Ix1 e
-> Steps Id (Ix1, e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix1 e -> Array LN Ix1 e
forall ix e. Array L ix e -> Array LN ix e
lData
{-# INLINE toStreamIx #-}