{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Massiv.Core.List (
L (..),
Array (..),
List (..),
toListArray,
showsArrayPrec,
showArrayList,
ListItem,
) where
import Control.Monad (unless, when)
import Control.Scheduler
import Data.Coerce
import Data.Functor.Identity
import Data.Kind
import qualified Data.List as L
import Data.Massiv.Core.Common
import qualified Data.Massiv.Vector.Stream as S
import Data.Monoid
import Data.Typeable
import GHC.Exts (IsList(..))
import GHC.TypeLits
import System.IO.Unsafe (unsafePerformIO)
type family ListItem ix e :: Type where
ListItem Ix1 e = e
ListItem ix e = [ListItem (Lower ix) e]
type family Elt ix e :: Type where
Elt Ix1 e = e
Elt ix e = List (Lower ix) e
newtype List ix e = List {forall ix e. List ix e -> [Elt ix e]
unList :: [Elt ix e]}
instance Coercible (Elt ix e) (ListItem ix e) => IsList (List ix e) where
type Item (List ix e) = ListItem ix e
fromList :: [Item (List ix e)] -> List ix e
fromList = coerce :: forall a b. Coercible a b => a -> b
coerce
{-# INLINE fromList #-}
toList :: List ix e -> [Item (List ix e)]
toList = coerce :: forall a b. Coercible a b => a -> b
coerce
{-# INLINE toList #-}
data L = L
data instance Array L ix e = LArray
{ forall ix e. Array L ix e -> Comp
lComp :: Comp
, forall ix e. Array L ix e -> List ix e
lData :: !(List ix e)
}
instance Coercible (Elt ix e) (ListItem 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 = forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq forall b c a. (b -> c) -> (a -> b) -> a -> c
. coerce :: forall a b. Coercible a b => a -> b
coerce
{-# INLINE fromList #-}
toList :: Array L ix e -> [Item (Array L ix e)]
toList = coerce :: forall a b. Coercible a b => a -> b
coerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE toList #-}
lengthHintList :: [a] -> LengthHint
lengthHintList :: forall a. [a] -> LengthHint
lengthHintList =
\case
[] -> Sz1 -> LengthHint
LengthExact forall ix. Index ix => Sz ix
zeroSz
[a]
_ -> LengthHint
LengthUnknown
{-# INLINE lengthHintList #-}
instance Shape L Ix1 where
linearSize :: forall e. Array L Ix1 e -> Sz1
linearSize = forall ix e. Array L ix e -> Sz1
outerLength
{-# INLINE linearSize #-}
linearSizeHint :: forall e. Array L Ix1 e -> LengthHint
linearSizeHint = forall a. [a] -> LengthHint
lengthHintList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE linearSizeHint #-}
isNull :: forall e. Array L Ix1 e -> Bool
isNull = forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE isNull #-}
outerSize :: forall e. Array L Ix1 e -> Sz1
outerSize = forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize
{-# INLINE outerSize #-}
instance Shape L Ix2 where
linearSize :: forall e. Array L Ix2 e -> Sz1
linearSize = forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Sum a -> a
getSum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a. a -> Sum a
Sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Ix1
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE linearSize #-}
linearSizeHint :: forall e. Array L Ix2 e -> LengthHint
linearSizeHint = forall a. [a] -> LengthHint
lengthHintList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE linearSizeHint #-}
isNull :: forall e. Array L Ix2 e -> Bool
isNull = All -> Bool
getAll forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Bool -> All
All forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE isNull #-}
outerSize :: forall e. Array L Ix2 e -> Sz Ix2
outerSize Array L Ix2 e
arr =
case forall ix e. List ix e -> [Elt ix e]
unList (forall ix e. Array L ix e -> List ix e
lData Array L Ix2 e
arr) of
[] -> forall ix. Index ix => Sz ix
zeroSz
(Elt Ix2 e
x : [Elt Ix2 e]
xs) -> forall ix. ix -> Sz ix
SafeSz ((Ix1
1 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Ix1
length [Elt Ix2 e]
xs) Ix1 -> Ix1 -> Ix2
:. forall (t :: * -> *) a. Foldable t => t a -> Ix1
length (forall ix e. List ix e -> [Elt ix e]
unList Elt Ix2 e
x))
{-# INLINE outerSize #-}
instance (Shape L (Ix (n - 1)), Index (IxN n)) => Shape L (IxN n) where
linearSize :: forall e. Array L (IxN n) e -> Sz1
linearSize = forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Sum a -> a
getSum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a. a -> Sum a
Sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Sz ix -> ix
unSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE linearSize #-}
linearSizeHint :: forall e. Array L (IxN n) e -> LengthHint
linearSizeHint = forall a. [a] -> LengthHint
lengthHintList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE linearSizeHint #-}
isNull :: forall e. Array L (IxN n) e -> Bool
isNull = All -> Bool
getAll forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Bool -> All
All forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e. Shape r ix => Array r ix e -> Bool
isNull forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE isNull #-}
outerSize :: forall e. Array L (IxN n) e -> Sz (IxN n)
outerSize Array L (IxN n) e
arr =
case forall ix e. List ix e -> [Elt ix e]
unList (forall ix e. Array L ix e -> List ix e
lData Array L (IxN n) e
arr) of
[] -> forall ix. Index ix => Sz ix
zeroSz
(Elt (IxN n) e
x : [Elt (IxN n) e]
xs) -> forall ix. ix -> Sz ix
SafeSz ((Ix1
1 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Ix1
length [Elt (IxN n) e]
xs) forall (n :: Natural). Ix1 -> Ix (n - 1) -> IxN n
:> forall ix. Sz ix -> ix
unSz (forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize (forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq Elt (IxN n) e
x)))
{-# INLINE outerSize #-}
outerLength :: Array L ix e -> Sz Int
outerLength :: forall ix e. Array L ix e -> Sz1
outerLength = forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Ix1
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE outerLength #-}
instance Ragged L Ix1 e where
flattenRagged :: Array L Ix1 e -> Array L Ix1 e
flattenRagged = forall a. a -> a
id
{-# INLINE flattenRagged #-}
generateRaggedM :: forall (m :: * -> *).
Monad m =>
Comp -> Sz1 -> (Ix1 -> m e) -> m (Array L Ix1 e)
generateRaggedM !Comp
comp !Sz1
k Ix1 -> m e
f = do
[e]
xs <-
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopDeepM Ix1
0 (forall a. Ord a => a -> a -> Bool
< coerce :: forall a b. Coercible a b => a -> b
coerce Sz1
k) (forall a. Num a => a -> a -> a
+ Ix1
1) [] forall a b. (a -> b) -> a -> b
$ \Ix1
i [e]
acc -> do
e
e <- Ix1 -> m e
f Ix1
i
forall (m :: * -> *) a. Monad m => a -> m a
return (e
e forall a. a -> [a] -> [a]
: [e]
acc)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
comp forall a b. (a -> b) -> a -> b
$ coerce :: forall a b. Coercible a b => a -> b
coerce [e]
xs
{-# INLINE generateRaggedM #-}
loadRaggedST :: forall s.
Scheduler s ()
-> Array L Ix1 e
-> (Ix1 -> e -> ST s ())
-> Ix1
-> Ix1
-> Sz1
-> ST s ()
loadRaggedST Scheduler s ()
_scheduler Array L Ix1 e
xs Ix1 -> e -> ST s ()
uWrite Ix1
start Ix1
end Sz1
sz = [e] -> Ix1 -> ST s ()
go (forall ix e. List ix e -> [Elt ix e]
unList (forall ix e. Array L ix e -> List ix e
lData Array L Ix1 e
xs)) Ix1
start
where
go :: [e] -> Ix1 -> ST s ()
go (e
y : [e]
ys) Ix1
i
| Ix1
i forall a. Ord a => a -> a -> Bool
< Ix1
end = Ix1 -> e -> ST s ()
uWrite Ix1
i e
y forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [e] -> Ix1 -> ST s ()
go [e]
ys (Ix1
i forall a. Num a => a -> a -> a
+ Ix1
1)
| Bool
otherwise = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException Dim
1 Sz1
sz (forall ix e. Array L ix e -> Sz1
outerLength Array L Ix1 e
xs)
go [] Ix1
i = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
i forall a. Eq a => a -> a -> Bool
/= Ix1
end) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException Dim
1 Sz1
sz (forall ix e. Array L ix e -> Sz1
outerLength Array L Ix1 e
xs)
{-# INLINE loadRaggedST #-}
raggedFormat :: (e -> String) -> String -> Array L Ix1 e -> String
raggedFormat e -> String
f String
_ Array L Ix1 e
arr = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat forall a b. (a -> b) -> a -> b
$ String
"[ " forall a. a -> [a] -> [a]
: forall a. a -> [a] -> [a]
L.intersperse String
", " (forall a b. (a -> b) -> [a] -> [b]
map e -> String
f (coerce :: forall a b. Coercible a b => a -> b
coerce (forall ix e. Array L ix e -> List ix e
lData Array L Ix1 e
arr))) forall a. [a] -> [a] -> [a]
++ [String
" ]"]
instance (Shape L ix, Ragged L ix e) => Load L ix e where
makeArray :: Comp -> Sz ix -> (ix -> e) -> Array L ix e
makeArray Comp
comp Sz ix
sz ix -> e
f = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ 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 ix
sz (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
f)
{-# INLINE makeArray #-}
iterArrayLinearST_ :: forall s.
Scheduler s () -> Array L ix e -> (Ix1 -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array L ix e
arr Ix1 -> e -> ST s ()
uWrite =
forall r ix e s.
Ragged r ix e =>
Scheduler s ()
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> Ix1
-> Ix1
-> Sz ix
-> ST s ()
loadRaggedST Scheduler s ()
scheduler Array L ix e
arr Ix1 -> e -> ST s ()
uWrite Ix1
0 (forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) Sz ix
sz
where
!sz :: Sz ix
sz = forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array L ix e
arr
{-# INLINE iterArrayLinearST_ #-}
instance Ragged L Ix2 e where
generateRaggedM :: forall (m :: * -> *).
Monad m =>
Comp -> Sz Ix2 -> (Ix2 -> m e) -> m (Array L Ix2 e)
generateRaggedM = forall ix e (m :: * -> *).
(Elt ix e ~ List (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 Ix2 e -> Vector L e
flattenRagged Array L Ix2 e
arr = LArray{lComp :: Comp
lComp = forall ix e. Array L ix e -> Comp
lComp Array L Ix2 e
arr, lData :: List Ix1 e
lData = coerce :: forall a b. Coercible a b => a -> b
coerce [e]
xs}
where
xs :: [e]
xs = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Comp -> List ix e -> Array L ix e
LArray (forall ix e. Array L ix e -> Comp
lComp Array L Ix2 e
arr)) (forall ix e. List ix e -> [Elt ix e]
unList (forall ix e. Array L ix e -> List ix e
lData Array L Ix2 e
arr))
{-# INLINE flattenRagged #-}
loadRaggedST :: forall s.
Scheduler s ()
-> Array L Ix2 e
-> (Ix1 -> e -> ST s ())
-> Ix1
-> Ix1
-> Sz Ix2
-> ST s ()
loadRaggedST Scheduler s ()
scheduler Array L Ix2 e
xs Ix1 -> e -> ST s ()
uWrite Ix1
start Ix1
end Sz Ix2
sz
| forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz Ix2
sz = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall r ix e. Shape r ix => Array r ix e -> Bool
isNotNull (forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged Array L Ix2 e
xs)) (forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM ShapeException
ShapeNonEmpty)
| Bool
otherwise = do
let (Sz1
k, Sz (Lower Ix2)
szL) = forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz Ix2
sz
step :: Ix1
step = forall ix. Index ix => Sz ix -> Ix1
totalElem Sz (Lower Ix2)
szL
[[e]]
leftOver <-
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
start (forall a. Ord a => a -> a -> Bool
< Ix1
end) (forall a. Num a => a -> a -> a
+ Ix1
step) (coerce :: forall a b. Coercible a b => a -> b
coerce (forall ix e. Array L ix e -> List ix e
lData Array L Ix2 e
xs)) forall a b. (a -> b) -> a -> b
$ \Ix1
i [[e]]
zs ->
case [[e]]
zs of
[] -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException Dim
2 Sz1
k (forall ix e. Array L ix e -> Sz1
outerLength Array L Ix2 e
xs))
([e]
y : [[e]]
ys) -> do
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
let end' :: Ix1
end' = Ix1
i forall a. Num a => a -> a -> a
+ Ix1
step
go :: [e] -> Ix1 -> ST s ()
go (e
a : [e]
as) Ix1
j
| Ix1
j forall a. Ord a => a -> a -> Bool
< Ix1
end' = Ix1 -> e -> ST s ()
uWrite Ix1
j e
a forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [e] -> Ix1 -> ST s ()
go [e]
as (Ix1
j forall a. Num a => a -> a -> a
+ Ix1
1)
| Bool
otherwise = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException Dim
1 Sz (Lower Ix2)
szL (forall ix. Index ix => ix -> Sz ix
Sz (forall (t :: * -> *) a. Foldable t => t a -> Ix1
length [e]
y))
go [] Ix1
j = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
j forall a. Eq a => a -> a -> Bool
/= Ix1
end') forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException Dim
1 Sz (Lower Ix2)
szL (forall ix. Index ix => ix -> Sz ix
Sz (forall (t :: * -> *) a. Foldable t => t a -> Ix1
length [e]
y)))
in [e] -> Ix1 -> ST s ()
go [e]
y Ix1
i
forall (f :: * -> *) a. Applicative f => a -> f a
pure [[e]]
ys
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[e]]
leftOver) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException Dim
2 Sz1
k (forall ix e. Array L ix e -> Sz1
outerLength Array L Ix2 e
xs)
{-# INLINE loadRaggedST #-}
raggedFormat :: (e -> String) -> String -> Array L Ix2 e -> String
raggedFormat e -> String
f String
sep (LArray Comp
comp List Ix2 e
xs) =
forall a. (String -> a -> String) -> String -> [a] -> String
showN (\String
s List Ix1 e
y -> forall r ix e.
Ragged r ix e =>
(e -> String) -> String -> Array r ix e -> String
raggedFormat e -> String
f String
s (forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
comp List Ix1 e
y :: Array L Ix1 e)) String
sep (coerce :: forall a b. Coercible a b => a -> b
coerce List Ix2 e
xs)
instance
( Shape L (IxN n)
, Ragged L (Ix (n - 1)) e
, Coercible (Elt (Ix (n - 1)) e) (ListItem (Ix (n - 1)) e)
)
=> Ragged L (IxN n) e
where
generateRaggedM :: forall (m :: * -> *).
Monad m =>
Comp -> Sz (IxN n) -> (IxN n -> m e) -> m (Array L (IxN n) e)
generateRaggedM = forall ix e (m :: * -> *).
(Elt ix e ~ List (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 (IxN n) e -> Vector L e
flattenRagged Array L (IxN n) e
arr = LArray{lComp :: Comp
lComp = forall ix e. Array L ix e -> Comp
lComp Array L (IxN n) e
arr, lData :: List Ix1 e
lData = coerce :: forall a b. Coercible a b => a -> b
coerce [e]
xs}
where
xs :: [e]
xs = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Comp -> List ix e -> Array L ix e
LArray (forall ix e. Array L ix e -> Comp
lComp Array L (IxN n) e
arr)) (forall ix e. List ix e -> [Elt ix e]
unList (forall ix e. Array L ix e -> List ix e
lData Array L (IxN n) e
arr))
{-# INLINE flattenRagged #-}
loadRaggedST :: forall s.
Scheduler s ()
-> Array L (IxN n) e
-> (Ix1 -> e -> ST s ())
-> Ix1
-> Ix1
-> Sz (IxN n)
-> ST s ()
loadRaggedST Scheduler s ()
scheduler Array L (IxN n) e
xs Ix1 -> e -> ST s ()
uWrite Ix1
start Ix1
end Sz (IxN n)
sz
| forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz (IxN n)
sz = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall r ix e. Shape r ix => Array r ix e -> Bool
isNotNull (forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged Array L (IxN n) e
xs)) (forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM ShapeException
ShapeNonEmpty)
| Bool
otherwise = do
let (Sz1
k, Sz (Lower (IxN n))
szL) = forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz (IxN n)
sz
step :: Ix1
step = forall ix. Index ix => Sz ix -> Ix1
totalElem Sz (Lower (IxN n))
szL
subScheduler :: Scheduler s ()
subScheduler
| Ix1
end forall a. Num a => a -> a -> a
- Ix1
start forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler forall a. Num a => a -> a -> a
* Ix1
step = Scheduler s ()
scheduler
| Bool
otherwise = forall s. Scheduler s ()
trivialScheduler_
[List (Ix (n - 1)) e]
leftOver <-
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
start (forall a. Ord a => a -> a -> Bool
< Ix1
end) (forall a. Num a => a -> a -> a
+ Ix1
step) (forall ix e. List ix e -> [Elt ix e]
unList (forall ix e. Array L ix e -> List ix e
lData Array L (IxN n) e
xs)) forall a b. (a -> b) -> a -> b
$ \Ix1
i [List (Ix (n - 1)) e]
zs ->
case [List (Ix (n - 1)) e]
zs of
[] -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException (forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions Sz (IxN n)
sz) Sz1
k (forall ix e. Array L ix e -> Sz1
outerLength Array L (IxN n) e
xs))
(List (Ix (n - 1)) e
y : [List (Ix (n - 1)) e]
ys) -> do
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
forall r ix e s.
Ragged r ix e =>
Scheduler s ()
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> Ix1
-> Ix1
-> Sz ix
-> ST s ()
loadRaggedST Scheduler s ()
subScheduler (forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq List (Ix (n - 1)) e
y) Ix1 -> e -> ST s ()
uWrite Ix1
i (Ix1
i forall a. Num a => a -> a -> a
+ Ix1
step) Sz (Lower (IxN n))
szL
forall (f :: * -> *) a. Applicative f => a -> f a
pure [List (Ix (n - 1)) e]
ys
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [List (Ix (n - 1)) e]
leftOver) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException (forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions Sz (IxN n)
sz) Sz1
k (forall ix e. Array L ix e -> Sz1
outerLength Array L (IxN n) e
xs)
{-# INLINE loadRaggedST #-}
raggedFormat :: (e -> String) -> String -> Array L (IxN n) e -> String
raggedFormat e -> String
f String
sep (LArray Comp
comp List (IxN n) e
xs) =
forall a. (String -> a -> String) -> String -> [a] -> String
showN (\String
s List (Ix (n - 1)) e
y -> forall r ix e.
Ragged r ix e =>
(e -> String) -> String -> Array r ix e -> String
raggedFormat e -> String
f String
s (forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
comp List (Ix (n - 1)) e
y :: Array L (Ix (n - 1)) e)) String
sep (coerce :: forall a b. Coercible a b => a -> b
coerce List (IxN n) e
xs)
unsafeGenerateParM
:: (Elt ix e ~ List (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e)
=> Comp
-> Sz ix
-> (ix -> m e)
-> m (Array L ix e)
unsafeGenerateParM :: forall ix e (m :: * -> *).
(Elt ix e ~ List (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
comp !Sz ix
sz ix -> m e
f = do
[[List (Lower ix) e]]
res <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall a b. (a -> b) -> a -> b
$ forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
let !(Sz1
ksz, Sz (Lower ix)
szL) = forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz ix
sz
!k :: Ix1
k = forall ix. Sz ix -> ix
unSz Sz1
ksz
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler Comp
comp forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld (m [List (Lower ix) e])
scheduler ->
forall a. Ix1 -> Ix1 -> (Ix1 -> Ix1 -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Ix1
numWorkers Scheduler RealWorld (m [List (Lower ix) e])
scheduler) Ix1
k forall a b. (a -> b) -> a -> b
$ \Ix1
chunkLength Ix1
slackStart -> do
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
0 (forall a. Ord a => a -> a -> Bool
< Ix1
slackStart) (forall a. Num a => a -> a -> a
+ Ix1
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Ix1
start ->
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld (m [List (Lower ix) e])
scheduler forall a b. (a -> b) -> a -> b
$ do
[m (List (Lower ix) e)]
res <- forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopDeepM Ix1
start (forall a. Ord a => a -> a -> Bool
< (Ix1
start forall a. Num a => a -> a -> a
+ Ix1
chunkLength)) (forall a. Num a => a -> a -> a
+ Ix1
1) [] forall a b. (a -> b) -> a -> b
$ \Ix1
i [m (List (Lower ix) e)]
acc ->
forall (m :: * -> *) a. Monad m => a -> m a
return (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall ix e. Array L ix e -> List ix e
lData (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 (forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i Lower ix
ixL))) forall a. a -> [a] -> [a]
: [m (List (Lower ix) e)]
acc)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [m (List (Lower ix) e)]
res
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
slackStart forall a. Ord a => a -> a -> Bool
< Ix1
k) forall a b. (a -> b) -> a -> b
$
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld (m [List (Lower ix) e])
scheduler forall a b. (a -> b) -> a -> b
$ do
[m (List (Lower ix) e)]
res <- forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopDeepM Ix1
slackStart (forall a. Ord a => a -> a -> Bool
< Ix1
k) (forall a. Num a => a -> a -> a
+ Ix1
1) [] forall a b. (a -> b) -> a -> b
$ \Ix1
i [m (List (Lower ix) e)]
acc ->
forall (m :: * -> *) a. Monad m => a -> m a
return (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall ix e. Array L ix e -> List ix e
lData (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 (forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i Lower ix
ixL))) forall a. a -> [a] -> [a]
: [m (List (Lower ix) e)]
acc)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [m (List (Lower ix) e)]
res
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
comp forall a b. (a -> b) -> a -> b
$ forall ix e. [Elt ix e] -> List ix e
List forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[List (Lower ix) e]]
res
{-# INLINE unsafeGenerateParM #-}
instance Strategy L where
setComp :: forall ix e. Comp -> Array L ix e -> Array L ix e
setComp Comp
c Array L ix e
arr = Array L ix e
arr{lComp :: Comp
lComp = Comp
c}
{-# INLINE setComp #-}
getComp :: forall ix e. Array L ix e -> Comp
getComp = forall ix e. Array L ix e -> Comp
lComp
{-# INLINE getComp #-}
repr :: L
repr = L
L
toListArray :: (Ragged L ix e, Shape r ix, Source r e) => Array r ix e -> Array L ix e
toListArray :: forall ix e r.
(Ragged L ix e, Shape r ix, Source r e) =>
Array r ix e -> Array L ix e
toListArray !Array r ix e
arr = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) (forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr) (forall r e ix. (Source r e, Index ix) => 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 Ix1
n Array L ix e
arr = forall r ix e.
(Ragged L ix e, Typeable r, Show e) =>
Proxy r -> Sz ix -> Ix1 -> Array L ix e -> ShowS
showsArrayLAsPrec (forall {k} (t :: k). Proxy t
Proxy :: Proxy L) (forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array L ix e
arr) Ix1
n Array L ix e
arr
instance (Ragged L ix e, Show e) => Show (List ix e) where
show :: List ix e -> String
show List ix e
xs = String
" " forall a. [a] -> [a] -> [a]
++ forall r ix e.
Ragged r ix e =>
(e -> String) -> String -> Array r ix e -> String
raggedFormat forall a. Show a => a -> String
show String
"\n " Array L ix e
arrL
where
arrL :: Array L ix e
arrL = forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq List ix e
xs :: Array L ix e
showN :: (String -> a -> String) -> String -> [a] -> String
showN :: forall a. (String -> a -> String) -> String -> [a] -> String
showN String -> a -> String
_ String
_ [] = String
"[ ]"
showN String -> a -> String
fShow String
lnPrefix [a]
ls =
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat
( [String
"[ "]
forall a. [a] -> [a] -> [a]
++ forall a. a -> [a] -> [a]
L.intersperse (String
lnPrefix forall a. [a] -> [a] -> [a]
++ String
", ") (forall a b. (a -> b) -> [a] -> [b]
map (String -> a -> String
fShow (String
lnPrefix forall a. [a] -> [a] -> [a]
++ String
" ")) [a]
ls)
forall a. [a] -> [a] -> [a]
++ [String
lnPrefix, String
"]"]
)
showsArrayLAsPrec
:: forall r ix e
. (Ragged L ix e, Typeable r, Show e)
=> Proxy r
-> Sz ix
-> Int
-> Array L ix e
-> ShowS
showsArrayLAsPrec :: forall r ix e.
(Ragged L ix e, Typeable r, Show e) =>
Proxy r -> Sz ix -> Ix1 -> Array L ix e -> ShowS
showsArrayLAsPrec Proxy r
pr Sz ix
sz Ix1
n Array L ix e
arr =
ShowS
opp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"Array " forall a. [a] -> [a] -> [a]
++)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRep -> ShowS
showsTypeRep (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep Proxy r
pr)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
' ' forall a. a -> [a] -> [a]
:)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Ix1 -> a -> ShowS
showsPrec Ix1
1 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array L ix e
arr)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
" (" forall a. [a] -> [a] -> [a]
++)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows Sz ix
sz
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
")\n" forall a. [a] -> [a] -> [a]
++)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows List ix e
lnarr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
clp
where
(ShowS
opp, ShowS
clp) =
if Ix1
n forall a. Eq a => a -> a -> Bool
== Ix1
0
then (forall a. a -> a
id, forall a. a -> a
id)
else ((Char
'(' forall a. a -> [a] -> [a]
:), (String
"\n)" forall a. [a] -> [a] -> [a]
++))
lnarr :: List ix e
lnarr = forall ix e. Array L ix e -> List ix e
lData Array L ix e
arr
showsArrayPrec
:: forall r r' ix e
. (Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e)
=> (Array r ix e -> Array r' ix e)
-> Int
-> Array r ix e
-> ShowS
showsArrayPrec :: forall r r' ix e.
(Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) =>
(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 = forall r ix e.
(Ragged L ix e, Typeable r, Show e) =>
Proxy r -> Sz ix -> Ix1 -> Array L ix e -> ShowS
showsArrayLAsPrec (forall {k} (t :: k). Proxy t
Proxy :: Proxy r) Sz ix
sz Ix1
n Array L ix e
larr
where
sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix e
arr'
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 = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr') Sz ix
sz (forall ix r e.
(HasCallStack, Index ix, Source r e) =>
Array r ix e -> ix -> e
evaluate' Array r' ix e
arr') :: Array L ix e
showArrayList
:: Show arr => [arr] -> String -> String
showArrayList :: forall arr. Show arr => [arr] -> ShowS
showArrayList [arr]
arrs = (Char
'[' forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall arr. Show arr => [arr] -> ShowS
go [arr]
arrs forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
']' forall a. a -> [a] -> [a]
:)
where
go :: [a] -> ShowS
go [] = forall a. a -> a
id
go [a
x] = (Char
' ' forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows a
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
'\n' forall a. a -> [a] -> [a]
:)
go (a
x : [a]
xs) = (Char
' ' forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows a
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"\n," forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> ShowS
go [a]
xs
instance Stream L Ix1 e where
toStream :: Array L Ix1 e -> Steps Id e
toStream = forall (m :: * -> *) e. Monad m => [e] -> Steps m e
S.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE toStream #-}
toStreamIx :: Array L Ix1 e -> Steps Id (Ix1, e)
toStreamIx = forall (m :: * -> *) e. Monad m => Steps m e -> Steps m (Ix1, e)
S.indexed forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) e. Monad m => [e] -> Steps m e
S.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. List ix e -> [Elt ix e]
unList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array L ix e -> List ix e
lData
{-# INLINE toStreamIx #-}