{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Massiv.Core.List
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
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
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 { 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 = [Item (List ix e)] -> List ix e
coerce
  {-# INLINE fromList #-}
  toList :: List ix e -> [Item (List ix e)]
toList = List ix e -> [Item (List ix e)]
coerce
  {-# INLINE toList #-}


data L = L

data instance Array L ix e = LArray { Array L ix e -> Comp
lComp :: Comp
                                    , 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 = Comp -> List ix e -> Array L ix e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq (List ix e -> Array L ix e)
-> ([ListItem ix e] -> List ix e)
-> [ListItem ix e]
-> Array L ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ListItem ix e] -> List ix e
coerce
  {-# INLINE fromList #-}
  toList :: Array L ix e -> [Item (Array L ix e)]
toList = List ix e -> [ListItem ix e]
coerce (List ix e -> [ListItem ix e])
-> (Array L ix e -> List ix e) -> Array L ix e -> [ListItem ix e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L ix e -> List ix e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE toList #-}

lengthHintList :: [a] -> LengthHint
lengthHintList :: [a] -> LengthHint
lengthHintList =
  \case
    [] -> Sz1 -> LengthHint
LengthExact Sz1
forall ix. Index ix => Sz ix
zeroSz
    [a]
_  -> LengthHint
LengthUnknown
{-# INLINE lengthHintList #-}

instance Shape L Ix1 where
  linearSize :: Array L Int e -> Sz1
linearSize = Array L Int e -> Sz1
forall ix e. Array L ix e -> Sz1
outerLength
  {-# INLINE linearSize #-}
  linearSizeHint :: Array L Int e -> LengthHint
linearSizeHint = [e] -> LengthHint
forall a. [a] -> LengthHint
lengthHintList ([e] -> LengthHint)
-> (Array L Int e -> [e]) -> Array L Int e -> LengthHint
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList (List Int e -> [e])
-> (Array L Int e -> List Int e) -> Array L Int e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Int e -> List Int e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE linearSizeHint #-}
  isNull :: Array L Int e -> Bool
isNull = [e] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([e] -> Bool) -> (Array L Int e -> [e]) -> Array L Int e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList (List Int e -> [e])
-> (Array L Int e -> List Int e) -> Array L Int e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Int e -> List Int e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE isNull #-}
  outerSize :: Array L Int e -> Sz1
outerSize = Array L Int e -> Sz1
forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize
  {-# INLINE outerSize #-}

instance Shape L Ix2 where
  linearSize :: Array L Ix2 e -> Sz1
linearSize = Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Int -> Sz1) -> (Array L Ix2 e -> Int) -> Array L Ix2 e -> Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sum Int -> Int
forall a. Sum a -> a
getSum (Sum Int -> Int)
-> (Array L Ix2 e -> Sum Int) -> Array L Ix2 e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (List Int e -> Sum Int) -> [List Int e] -> Sum Int
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Int -> Sum Int
forall a. a -> Sum a
Sum (Int -> Sum Int) -> (List Int e -> Int) -> List Int e -> Sum Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([e] -> Int) -> (List Int e -> [e]) -> List Int e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList) ([List Int e] -> Sum Int)
-> (Array L Ix2 e -> [List Int e]) -> Array L Ix2 e -> Sum Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Ix2 e -> [List Int e]
forall ix e. List ix e -> [Elt ix e]
unList (List Ix2 e -> [List Int e])
-> (Array L Ix2 e -> List Ix2 e) -> Array L Ix2 e -> [List Int e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix2 e -> List Ix2 e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE linearSize #-}
  linearSizeHint :: Array L Ix2 e -> LengthHint
linearSizeHint = [List Int e] -> LengthHint
forall a. [a] -> LengthHint
lengthHintList ([List Int e] -> LengthHint)
-> (Array L Ix2 e -> [List Int e]) -> Array L Ix2 e -> LengthHint
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Ix2 e -> [List Int e]
forall ix e. List ix e -> [Elt ix e]
unList (List Ix2 e -> [List Int e])
-> (Array L Ix2 e -> List Ix2 e) -> Array L Ix2 e -> [List Int e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix2 e -> List Ix2 e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE linearSizeHint #-}
  isNull :: Array L Ix2 e -> Bool
isNull = All -> Bool
getAll (All -> Bool) -> (Array L Ix2 e -> All) -> Array L Ix2 e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (List Int e -> All) -> [List Int e] -> All
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Bool -> All
All (Bool -> All) -> (List Int e -> Bool) -> List Int e -> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [e] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([e] -> Bool) -> (List Int e -> [e]) -> List Int e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList) ([List Int e] -> All)
-> (Array L Ix2 e -> [List Int e]) -> Array L Ix2 e -> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Ix2 e -> [List Int e]
forall ix e. List ix e -> [Elt ix e]
unList (List Ix2 e -> [List Int e])
-> (Array L Ix2 e -> List Ix2 e) -> Array L Ix2 e -> [List Int e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Ix2 e -> List Ix2 e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE isNull #-}
  outerSize :: Array L Ix2 e -> Sz Ix2
outerSize Array L Ix2 e
arr =
    case List Ix2 e -> [Elt Ix2 e]
forall ix e. List ix e -> [Elt ix e]
unList (Array L Ix2 e -> List Ix2 e
forall ix e. Array L ix e -> List ix e
lData Array L Ix2 e
arr) of
      []     -> Sz Ix2
forall ix. Index ix => Sz ix
zeroSz
      (Elt Ix2 e
x:[Elt Ix2 e]
xs) -> Ix2 -> Sz Ix2
forall ix. ix -> Sz ix
SafeSz ((Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [List Int e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [List Int e]
[Elt Ix2 e]
xs) Int -> Int -> Ix2
:. [e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (List Int e -> [Elt Int e]
forall ix e. List ix e -> [Elt ix e]
unList List Int e
Elt Ix2 e
x))
  {-# INLINE outerSize #-}

instance (Shape L (Ix (n - 1)), Index (IxN n)) => Shape L (IxN n) where
  linearSize :: Array L (IxN n) e -> Sz1
linearSize = Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Int -> Sz1)
-> (Array L (IxN n) e -> Int) -> Array L (IxN n) e -> Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sum Int -> Int
forall a. Sum a -> a
getSum (Sum Int -> Int)
-> (Array L (IxN n) e -> Sum Int) -> Array L (IxN n) e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (List (Ix (n - 1)) e -> Sum Int)
-> [List (Ix (n - 1)) e] -> Sum Int
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Int -> Sum Int
forall a. a -> Sum a
Sum (Int -> Sum Int)
-> (List (Ix (n - 1)) e -> Int) -> List (Ix (n - 1)) e -> Sum Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz1 -> Int
forall ix. Sz ix -> ix
unSz (Sz1 -> Int)
-> (List (Ix (n - 1)) e -> Sz1) -> List (Ix (n - 1)) e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L (Ix (n - 1)) e -> Sz1
forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize (Array L (Ix (n - 1)) e -> Sz1)
-> (List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e)
-> List (Ix (n - 1)) e
-> Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq) ([List (Ix (n - 1)) e] -> Sum Int)
-> (Array L (IxN n) e -> [List (Ix (n - 1)) e])
-> Array L (IxN n) e
-> Sum Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List (IxN n) e -> [List (Ix (n - 1)) e]
forall ix e. List ix e -> [Elt ix e]
unList (List (IxN n) e -> [List (Ix (n - 1)) e])
-> (Array L (IxN n) e -> List (IxN n) e)
-> Array L (IxN n) e
-> [List (Ix (n - 1)) e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L (IxN n) e -> List (IxN n) e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE linearSize #-}
  linearSizeHint :: Array L (IxN n) e -> LengthHint
linearSizeHint = [List (Ix (n - 1)) e] -> LengthHint
forall a. [a] -> LengthHint
lengthHintList ([List (Ix (n - 1)) e] -> LengthHint)
-> (Array L (IxN n) e -> [List (Ix (n - 1)) e])
-> Array L (IxN n) e
-> LengthHint
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List (IxN n) e -> [List (Ix (n - 1)) e]
forall ix e. List ix e -> [Elt ix e]
unList (List (IxN n) e -> [List (Ix (n - 1)) e])
-> (Array L (IxN n) e -> List (IxN n) e)
-> Array L (IxN n) e
-> [List (Ix (n - 1)) e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L (IxN n) e -> List (IxN n) e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE linearSizeHint #-}
  isNull :: Array L (IxN n) e -> Bool
isNull = All -> Bool
getAll (All -> Bool)
-> (Array L (IxN n) e -> All) -> Array L (IxN n) e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (List (Ix (n - 1)) e -> All) -> [List (Ix (n - 1)) e] -> All
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Bool -> All
All (Bool -> All)
-> (List (Ix (n - 1)) e -> Bool) -> List (Ix (n - 1)) e -> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L (Ix (n - 1)) e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNull (Array L (Ix (n - 1)) e -> Bool)
-> (List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e)
-> List (Ix (n - 1)) e
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq) ([List (Ix (n - 1)) e] -> All)
-> (Array L (IxN n) e -> [List (Ix (n - 1)) e])
-> Array L (IxN n) e
-> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List (IxN n) e -> [List (Ix (n - 1)) e]
forall ix e. List ix e -> [Elt ix e]
unList (List (IxN n) e -> [List (Ix (n - 1)) e])
-> (Array L (IxN n) e -> List (IxN n) e)
-> Array L (IxN n) e
-> [List (Ix (n - 1)) e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L (IxN n) e -> List (IxN n) e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE isNull #-}
  outerSize :: Array L (IxN n) e -> Sz (IxN n)
outerSize Array L (IxN n) e
arr =
    case List (IxN n) e -> [Elt (IxN n) e]
forall ix e. List ix e -> [Elt ix e]
unList (Array L (IxN n) e -> List (IxN n) e
forall ix e. Array L ix e -> List ix e
lData Array L (IxN n) e
arr) of
      []     -> Sz (IxN n)
forall ix. Index ix => Sz ix
zeroSz
      (Elt (IxN n) e
x:[Elt (IxN n) e]
xs) -> IxN n -> Sz (IxN n)
forall ix. ix -> Sz ix
SafeSz ((Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [List (Ix (n - 1)) e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [List (Ix (n - 1)) e]
[Elt (IxN n) e]
xs) Int -> Ix (n - 1) -> IxN n
forall (n :: Nat). Int -> Ix (n - 1) -> IxN n
:> Sz (Ix (n - 1)) -> Ix (n - 1)
forall ix. Sz ix -> ix
unSz (Array L (Ix (n - 1)) e -> Sz (Ix (n - 1))
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize (Comp -> List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq List (Ix (n - 1)) e
Elt (IxN n) e
x)))
  {-# INLINE outerSize #-}


outerLength :: Array L ix e -> Sz Int
outerLength :: Array L ix e -> Sz1
outerLength = Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Int -> Sz1) -> (Array L ix e -> Int) -> Array L ix e -> Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Elt ix e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Elt ix e] -> Int)
-> (Array L ix e -> [Elt ix e]) -> Array L ix e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List ix e -> [Elt ix e]
forall ix e. List ix e -> [Elt ix e]
unList (List ix e -> [Elt ix e])
-> (Array L ix e -> List ix e) -> Array L ix e -> [Elt ix e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L ix e -> List ix e
forall ix e. Array L ix e -> List ix e
lData


instance Ragged L Ix1 e where
  flattenRagged :: Array L Int e -> Array L Int e
flattenRagged = Array L Int e -> Array L Int e
forall a. a -> a
id
  {-# INLINE flattenRagged #-}
  generateRaggedM :: Comp -> Sz1 -> (Int -> m e) -> m (Array L Int e)
generateRaggedM !Comp
comp !Sz1
k Int -> m e
f = do
    [e]
xs <-
      Int
-> (Int -> Bool)
-> (Int -> Int)
-> [e]
-> (Int -> [e] -> m [e])
-> m [e]
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Sz1 -> Int
coerce Sz1
k) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [] ((Int -> [e] -> m [e]) -> m [e]) -> (Int -> [e] -> m [e]) -> m [e]
forall a b. (a -> b) -> a -> b
$ \Int
i [e]
acc -> do
        e
e <- Int -> m e
f Int
i
        [e] -> m [e]
forall (m :: * -> *) a. Monad m => a -> m a
return (e
e e -> [e] -> [e]
forall a. a -> [a] -> [a]
: [e]
acc)
    Array L Int e -> m (Array L Int e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array L Int e -> m (Array L Int e))
-> Array L Int e -> m (Array L Int e)
forall a b. (a -> b) -> a -> b
$ Comp -> List Int e -> Array L Int e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
comp (List Int e -> Array L Int e) -> List Int e -> Array L Int e
forall a b. (a -> b) -> a -> b
$ [e] -> List Int e
coerce [e]
xs
  {-# INLINE generateRaggedM #-}
  loadRaggedST :: Scheduler s ()
-> Array L Int e
-> (Int -> e -> ST s ())
-> Int
-> Int
-> Sz1
-> ST s ()
loadRaggedST Scheduler s ()
_scheduler Array L Int e
xs Int -> e -> ST s ()
uWrite Int
start Int
end Sz1
sz = [e] -> Int -> ST s ()
go (List Int e -> [Elt Int e]
forall ix e. List ix e -> [Elt ix e]
unList (Array L Int e -> List Int e
forall ix e. Array L ix e -> List ix e
lData Array L Int e
xs)) Int
start
    where
      go :: [e] -> Int -> ST s ()
go (e
y:[e]
ys) Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end = Int -> e -> ST s ()
uWrite Int
i e
y ST s () -> ST s () -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [e] -> Int -> ST s ()
go [e]
ys (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ShapeException -> ST s ()) -> ShapeException -> ST s ()
forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException Dim
1 Sz1
sz (Array L Int e -> Sz1
forall ix e. Array L ix e -> Sz1
outerLength Array L Int e
xs)
      go [] Int
i = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
end) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ShapeException -> ST s ()) -> ShapeException -> ST s ()
forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException Dim
1 Sz1
sz (Array L Int e -> Sz1
forall ix e. Array L ix e -> Sz1
outerLength Array L Int e
xs)
  {-# INLINE loadRaggedST #-}
  raggedFormat :: (e -> String) -> String -> Array L Int e -> String
raggedFormat e -> String
f String
_ Array L Int 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 (List Int e -> [e]
coerce (Array L Int e -> List Int e
forall ix e. Array L ix e -> List ix e
lData Array L Int e
arr))) [String] -> [String] -> [String]
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 = Identity (Array L ix e) -> Array L ix e
forall a. Identity a -> a
runIdentity (Identity (Array L ix e) -> Array L ix e)
-> Identity (Array L ix e) -> Array L ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Sz ix -> (ix -> Identity e) -> Identity (Array L 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 ix
sz (e -> Identity e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> Identity e) -> (ix -> e) -> ix -> Identity e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
f)
  {-# INLINE makeArray #-}
  iterArrayLinearST_ :: Scheduler s () -> Array L ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array L ix e
arr Int -> e -> ST s ()
uWrite =
    Scheduler s ()
-> Array L ix e
-> (Int -> e -> ST s ())
-> Int
-> Int
-> Sz ix
-> ST s ()
forall r ix e s.
Ragged r ix e =>
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> Int
-> Int
-> Sz ix
-> ST s ()
loadRaggedST Scheduler s ()
scheduler Array L ix e
arr Int -> e -> ST s ()
uWrite Int
0 (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) Sz ix
sz
    where !sz :: Sz ix
sz = Array L ix e -> Sz ix
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 :: Comp -> Sz Ix2 -> (Ix2 -> m e) -> m (Array L Ix2 e)
generateRaggedM = Comp -> Sz Ix2 -> (Ix2 -> m e) -> m (Array L Ix2 e)
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 :: forall ix e. Comp -> List ix e -> Array L ix e
LArray {lComp :: Comp
lComp = Array L Ix2 e -> Comp
forall ix e. Array L ix e -> Comp
lComp Array L Ix2 e
arr, lData :: List Int e
lData = [e] -> List Int e
coerce [e]
xs}
    where
      xs :: [e]
xs = (List Int e -> [e]) -> [List Int e] -> [e]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList (List Int e -> [e])
-> (List Int e -> List Int e) -> List Int e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector L e -> List Int e
forall ix e. Array L ix e -> List ix e
lData (Vector L e -> List Int e)
-> (List Int e -> Vector L e) -> List Int e -> List Int e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector L e -> Vector L e
forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged (Vector L e -> Vector L e)
-> (List Int e -> Vector L e) -> List Int e -> Vector L e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> List Int e -> Vector L e
forall ix e. Comp -> List ix e -> Array L ix e
LArray (Array L Ix2 e -> Comp
forall ix e. Array L ix e -> Comp
lComp Array L Ix2 e
arr)) (List Ix2 e -> [Elt Ix2 e]
forall ix e. List ix e -> [Elt ix e]
unList (Array L Ix2 e -> List Ix2 e
forall ix e. Array L ix e -> List ix e
lData Array L Ix2 e
arr))
  {-# INLINE flattenRagged #-}
  loadRaggedST :: Scheduler s ()
-> Array L Ix2 e
-> (Int -> e -> ST s ())
-> Int
-> Int
-> Sz Ix2
-> ST s ()
loadRaggedST Scheduler s ()
scheduler Array L Ix2 e
xs Int -> e -> ST s ()
uWrite Int
start Int
end Sz Ix2
sz
    | Sz Ix2 -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz Ix2
sz = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Vector L e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNotNull (Array L Ix2 e -> Vector L e
forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged Array L Ix2 e
xs)) (ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM ShapeException
ShapeNonEmpty)
    | Bool
otherwise = do
      let (Sz1
k, Sz1
szL) = Sz Ix2 -> (Sz1, Sz (Lower Ix2))
forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz Ix2
sz
          step :: Int
step = Sz1 -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz1
szL
      [[e]]
leftOver <-
        Int
-> (Int -> Bool)
-> (Int -> Int)
-> [[e]]
-> (Int -> [[e]] -> ST s [[e]])
-> ST s [[e]]
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step) (List Ix2 e -> [[e]]
coerce (Array L Ix2 e -> List Ix2 e
forall ix e. Array L ix e -> List ix e
lData Array L Ix2 e
xs)) ((Int -> [[e]] -> ST s [[e]]) -> ST s [[e]])
-> (Int -> [[e]] -> ST s [[e]]) -> ST s [[e]]
forall a b. (a -> b) -> a -> b
$ \Int
i [[e]]
zs ->
          case [[e]]
zs of
            [] -> ShapeException -> ST s [[e]]
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException Dim
2 Sz1
k (Array L Ix2 e -> Sz1
forall ix e. Array L ix e -> Sz1
outerLength Array L Ix2 e
xs))
            ([e]
y:[[e]]
ys) -> do
              Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
                let end' :: Int
end' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step
                    go :: [e] -> Int -> ST s ()
go (e
a:[e]
as) Int
j
                      | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end' = Int -> e -> ST s ()
uWrite Int
j e
a ST s () -> ST s () -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [e] -> Int -> ST s ()
go [e]
as (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                      | Bool
otherwise = ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ShapeException -> ST s ()) -> ShapeException -> ST s ()
forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException Dim
1 Sz1
szL (Int -> Sz1
forall ix. Index ix => ix -> Sz ix
Sz ([e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [e]
y))
                    go [] Int
j = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
end') (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException Dim
1 Sz1
szL (Int -> Sz1
forall ix. Index ix => ix -> Sz ix
Sz ([e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [e]
y)))
                 in [e] -> Int -> ST s ()
go [e]
y Int
i
              [[e]] -> ST s [[e]]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [[e]]
ys
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([[e]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[e]]
leftOver) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ShapeException -> ST s ()) -> ShapeException -> ST s ()
forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException Dim
2 Sz1
k (Array L Ix2 e -> Sz1
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 xs) =
    (String -> List Int e -> String)
-> String -> [List Int e] -> String
forall a. (String -> a -> String) -> String -> [a] -> String
showN (\String
s List Int e
y -> (e -> String) -> String -> Vector L 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 -> List Int e -> Vector L e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
comp List Int e
y :: Array L Ix1 e)) String
sep (List Ix2 e -> [List Int e]
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 :: Comp -> Sz (IxN n) -> (IxN n -> m e) -> m (Array L (IxN n) e)
generateRaggedM = Comp -> Sz (IxN n) -> (IxN n -> m e) -> m (Array L (IxN n) e)
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 :: forall ix e. Comp -> List ix e -> Array L ix e
LArray {lComp :: Comp
lComp = Array L (IxN n) e -> Comp
forall ix e. Array L ix e -> Comp
lComp Array L (IxN n) e
arr, lData :: List Int e
lData = [e] -> List Int e
coerce [e]
xs}
    where
      xs :: [e]
xs = (List (Ix (n - 1)) e -> [e]) -> [List (Ix (n - 1)) e] -> [e]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList (List Int e -> [e])
-> (List (Ix (n - 1)) e -> List Int e)
-> List (Ix (n - 1)) e
-> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector L e -> List Int e
forall ix e. Array L ix e -> List ix e
lData (Vector L e -> List Int e)
-> (List (Ix (n - 1)) e -> Vector L e)
-> List (Ix (n - 1)) e
-> List Int e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L (Ix (n - 1)) e -> Vector L e
forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged (Array L (Ix (n - 1)) e -> Vector L e)
-> (List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e)
-> List (Ix (n - 1)) e
-> Vector L e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e
forall ix e. Comp -> List ix e -> Array L ix e
LArray (Array L (IxN n) e -> Comp
forall ix e. Array L ix e -> Comp
lComp Array L (IxN n) e
arr)) (List (IxN n) e -> [Elt (IxN n) e]
forall ix e. List ix e -> [Elt ix e]
unList (Array L (IxN n) e -> List (IxN n) e
forall ix e. Array L ix e -> List ix e
lData Array L (IxN n) e
arr))
  {-# INLINE flattenRagged #-}
  loadRaggedST :: Scheduler s ()
-> Array L (IxN n) e
-> (Int -> e -> ST s ())
-> Int
-> Int
-> Sz (IxN n)
-> ST s ()
loadRaggedST Scheduler s ()
scheduler Array L (IxN n) e
xs Int -> e -> ST s ()
uWrite Int
start Int
end Sz (IxN n)
sz
    | Sz (IxN n) -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz (IxN n)
sz = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Vector L e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNotNull (Array L (IxN n) e -> Vector L e
forall r ix e. Ragged r ix e => Array r ix e -> Vector r e
flattenRagged Array L (IxN n) e
xs)) (ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM ShapeException
ShapeNonEmpty)
    | Bool
otherwise = do
      let (Sz1
k, Sz (Ix (n - 1))
szL) = Sz (IxN n) -> (Sz1, Sz (Lower (IxN n)))
forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz (IxN n)
sz
          step :: Int
step = Sz (Ix (n - 1)) -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz (Ix (n - 1))
szL
          subScheduler :: Scheduler s ()
subScheduler
            | Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
step = Scheduler s ()
scheduler
            | Bool
otherwise = Scheduler s ()
forall s. Scheduler s ()
trivialScheduler_
      [List (Ix (n - 1)) e]
leftOver <-
        Int
-> (Int -> Bool)
-> (Int -> Int)
-> [List (Ix (n - 1)) e]
-> (Int -> [List (Ix (n - 1)) e] -> ST s [List (Ix (n - 1)) e])
-> ST s [List (Ix (n - 1)) e]
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step) (List (IxN n) e -> [Elt (IxN n) e]
forall ix e. List ix e -> [Elt ix e]
unList (Array L (IxN n) e -> List (IxN n) e
forall ix e. Array L ix e -> List ix e
lData Array L (IxN n) e
xs)) ((Int -> [List (Ix (n - 1)) e] -> ST s [List (Ix (n - 1)) e])
 -> ST s [List (Ix (n - 1)) e])
-> (Int -> [List (Ix (n - 1)) e] -> ST s [List (Ix (n - 1)) e])
-> ST s [List (Ix (n - 1)) e]
forall a b. (a -> b) -> a -> b
$ \Int
i [List (Ix (n - 1)) e]
zs ->
          case [List (Ix (n - 1)) e]
zs of
            [] -> ShapeException -> ST s [List (Ix (n - 1)) e]
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Dim -> Sz1 -> Sz1 -> ShapeException
DimTooShortException (Sz (IxN n) -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions Sz (IxN n)
sz) Sz1
k (Array L (IxN n) e -> Sz1
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
              Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
                Scheduler s ()
-> Array L (Ix (n - 1)) e
-> (Int -> e -> ST s ())
-> Int
-> Int
-> Sz (Ix (n - 1))
-> ST s ()
forall r ix e s.
Ragged r ix e =>
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> Int
-> Int
-> Sz ix
-> ST s ()
loadRaggedST Scheduler s ()
subScheduler (Comp -> List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
Seq List (Ix (n - 1)) e
y) Int -> e -> ST s ()
uWrite Int
i (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step) Sz (Ix (n - 1))
szL
              [List (Ix (n - 1)) e] -> ST s [List (Ix (n - 1)) e]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [List (Ix (n - 1)) e]
ys
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([List (Ix (n - 1)) e] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [List (Ix (n - 1)) e]
leftOver) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ ShapeException -> ST s ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ShapeException -> ST s ()) -> ShapeException -> ST s ()
forall a b. (a -> b) -> a -> b
$ Dim -> Sz1 -> Sz1 -> ShapeException
DimTooLongException (Sz (IxN n) -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions Sz (IxN n)
sz) Sz1
k (Array L (IxN n) e -> Sz1
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 xs) =
    (String -> List (Ix (n - 1)) e -> String)
-> String -> [List (Ix (n - 1)) e] -> String
forall a. (String -> a -> String) -> String -> [a] -> String
showN (\String
s List (Ix (n - 1)) e
y -> (e -> String) -> String -> Array L (Ix (n - 1)) 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 -> List (Ix (n - 1)) e -> Array L (Ix (n - 1)) e
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 (List (IxN n) e -> [List (Ix (n - 1)) e]
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 :: 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 <- [m [List (Lower ix) e]] -> m [[List (Lower ix) e]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([m [List (Lower ix) e]] -> m [[List (Lower ix) e]])
-> [m [List (Lower ix) e]] -> m [[List (Lower ix) e]]
forall a b. (a -> b) -> a -> b
$ IO [m [List (Lower ix) e]] -> [m [List (Lower ix) e]]
forall a. IO a -> a
unsafePerformIO (IO [m [List (Lower ix) e]] -> [m [List (Lower ix) e]])
-> IO [m [List (Lower ix) e]] -> [m [List (Lower ix) e]]
forall a b. (a -> b) -> a -> b
$ do
    let !(Sz1
ksz, Sz (Lower ix)
szL) = Sz ix -> (Sz1, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz ix
sz
        !k :: Int
k = Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
ksz
    Comp
-> (Scheduler RealWorld (m [List (Lower ix) e]) -> IO ())
-> IO [m [List (Lower ix) e]]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler Comp
comp ((Scheduler RealWorld (m [List (Lower ix) e]) -> IO ())
 -> IO [m [List (Lower ix) e]])
-> (Scheduler RealWorld (m [List (Lower ix) e]) -> IO ())
-> IO [m [List (Lower ix) e]]
forall a b. (a -> b) -> a -> b
$ \ Scheduler RealWorld (m [List (Lower ix) e])
scheduler ->
      Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler RealWorld (m [List (Lower ix) e]) -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler RealWorld (m [List (Lower ix) e])
scheduler) Int
k ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Int
chunkLength Int
slackStart -> do
        Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
          Scheduler RealWorld (m [List (Lower ix) e])
-> IO (m [List (Lower ix) e]) -> IO ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld (m [List (Lower ix) e])
scheduler (IO (m [List (Lower ix) e]) -> IO ())
-> IO (m [List (Lower ix) e]) -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            [m (List (Lower ix) e)]
res <- Int
-> (Int -> Bool)
-> (Int -> Int)
-> [m (List (Lower ix) e)]
-> (Int -> [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)])
-> IO [m (List (Lower ix) e)]
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [] ((Int -> [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)])
 -> IO [m (List (Lower ix) e)])
-> (Int -> [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)])
-> IO [m (List (Lower ix) e)]
forall a b. (a -> b) -> a -> b
$ \Int
i [m (List (Lower ix) e)]
acc ->
              [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((Array L (Lower ix) e -> List (Lower ix) e)
-> m (Array L (Lower ix) e) -> m (List (Lower ix) e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array L (Lower ix) e -> List (Lower ix) e
forall ix e. Array L ix e -> List 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 (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ixL)))m (List (Lower ix) e)
-> [m (List (Lower ix) e)] -> [m (List (Lower ix) e)]
forall a. a -> [a] -> [a]
:[m (List (Lower ix) e)]
acc)
            m [List (Lower ix) e] -> IO (m [List (Lower ix) e])
forall (m :: * -> *) a. Monad m => a -> m a
return (m [List (Lower ix) e] -> IO (m [List (Lower ix) e]))
-> m [List (Lower ix) e] -> IO (m [List (Lower ix) e])
forall a b. (a -> b) -> a -> b
$! [m (List (Lower ix) e)] -> m [List (Lower ix) e]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [m (List (Lower ix) e)]
res
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          Scheduler RealWorld (m [List (Lower ix) e])
-> IO (m [List (Lower ix) e]) -> IO ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld (m [List (Lower ix) e])
scheduler (IO (m [List (Lower ix) e]) -> IO ())
-> IO (m [List (Lower ix) e]) -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            [m (List (Lower ix) e)]
res <- Int
-> (Int -> Bool)
-> (Int -> Int)
-> [m (List (Lower ix) e)]
-> (Int -> [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)])
-> IO [m (List (Lower ix) e)]
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM Int
slackStart (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [] ((Int -> [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)])
 -> IO [m (List (Lower ix) e)])
-> (Int -> [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)])
-> IO [m (List (Lower ix) e)]
forall a b. (a -> b) -> a -> b
$ \Int
i [m (List (Lower ix) e)]
acc ->
              [m (List (Lower ix) e)] -> IO [m (List (Lower ix) e)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((Array L (Lower ix) e -> List (Lower ix) e)
-> m (Array L (Lower ix) e) -> m (List (Lower ix) e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array L (Lower ix) e -> List (Lower ix) e
forall ix e. Array L ix e -> List 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 (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ixL)))m (List (Lower ix) e)
-> [m (List (Lower ix) e)] -> [m (List (Lower ix) e)]
forall a. a -> [a] -> [a]
:[m (List (Lower ix) e)]
acc)
            m [List (Lower ix) e] -> IO (m [List (Lower ix) e])
forall (m :: * -> *) a. Monad m => a -> m a
return (m [List (Lower ix) e] -> IO (m [List (Lower ix) e]))
-> m [List (Lower ix) e] -> IO (m [List (Lower ix) e])
forall a b. (a -> b) -> a -> b
$! [m (List (Lower ix) e)] -> m [List (Lower ix) e]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [m (List (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 -> List ix e -> Array L ix e
forall ix e. Comp -> List ix e -> Array L ix e
LArray Comp
comp (List ix e -> Array L ix e) -> List ix e -> Array L ix e
forall a b. (a -> b) -> a -> b
$ [Elt ix e] -> List ix e
forall ix e. [Elt ix e] -> List ix e
List ([Elt ix e] -> List ix e) -> [Elt ix e] -> List ix e
forall a b. (a -> b) -> a -> b
$ [[List (Lower ix) e]] -> [List (Lower ix) e]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[List (Lower ix) e]]
res
{-# INLINE unsafeGenerateParM #-}

instance Strategy L 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 #-}
  getComp :: Array L ix e -> Comp
getComp = Array L ix e -> Comp
forall ix e. Array L ix e -> Comp
lComp
  {-# INLINE getComp #-}

-- -- TODO: benchmark against using unsafeGenerateM directly
-- 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 f = unsafePerformIO $ do
--   let !(m, szL) = unconsSz sz
--   xs <- withScheduler comp $ \scheduler ->
--     loopM_ 0 (< coerce m) (+ 1) $ \i -> scheduleWork scheduler $
--       generateRaggedM comp szL $ \ix -> return $ f (consDim i ix)
--   return $! foldr' consR (emptyR comp) xs
-- {-# INLINE unsafeGenerateN #-}


-- | Construct an array backed by linked lists from any source array
--
-- @since 0.4.0
toListArray :: (Ragged L ix e, Shape r ix, Source r 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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) (Array r ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr) (Array r ix e -> ix -> e
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 :: Int -> Array L ix e -> ShowS
showsPrec Int
n Array L ix e
arr  = Proxy L -> Sz ix -> Int -> Array L ix e -> ShowS
forall r ix e.
(Ragged L ix e, Typeable r, Show e) =>
Proxy r -> Sz ix -> Int -> Array L ix e -> ShowS
showsArrayLAsPrec (Proxy L
forall k (t :: k). Proxy t
Proxy :: Proxy L) (Array L ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array L ix e
arr) Int
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
"  " 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 = Comp -> List ix e -> Array L ix e
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 :: (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
  -> Sz ix
  -> Int
  -> Array L ix e -- Array to show
  -> ShowS
showsArrayLAsPrec :: Proxy r -> Sz ix -> Int -> Array L ix e -> ShowS
showsArrayLAsPrec Proxy r
pr Sz ix
sz Int
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
.
  Int -> Comp -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 (Array L ix e -> Comp
forall r ix e. Strategy r => 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 Sz ix
sz 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
. List ix e -> ShowS
forall a. Show a => a -> ShowS
shows List ix e
lnarr ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
clp
  where
    (ShowS
opp, ShowS
clp) =
      if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
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 :: List ix e
lnarr = Array L ix e -> List ix e
forall ix e. Array L ix e -> List ix e
lData Array L ix e
arr

-- | Helper function for declaring `Show` instances for arrays
--
-- @since 0.4.0
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) -- ^ Modifier
  -> Int
  -> Array r ix e -- Array to show
  -> ShowS
showsArrayPrec :: (Array r ix e -> Array r' ix e) -> Int -> Array r ix e -> ShowS
showsArrayPrec Array r ix e -> Array r' ix e
f Int
n Array r ix e
arr = Proxy r -> Sz ix -> Int -> Array L ix e -> ShowS
forall r ix e.
(Ragged L ix e, Typeable r, Show e) =>
Proxy r -> Sz ix -> Int -> Array L ix e -> ShowS
showsArrayLAsPrec (Proxy r
forall k (t :: k). Proxy t
Proxy :: Proxy r) Sz ix
sz Int
n Array L ix e
larr
  where
    sz :: Sz ix
sz = Array r' ix e -> Sz ix
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 = Comp -> Sz ix -> (ix -> e) -> Array L ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr') Sz ix
sz (Array r' ix e -> ix -> e
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


-- | Helper function for declaring `Show` instances for arrays
--
-- @since 0.4.0
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 Stream L Ix1 e where
  toStream :: Array L Int 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 L Int e -> [e]) -> Array L Int e -> Steps Id e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList (List Int e -> [e])
-> (Array L Int e -> List Int e) -> Array L Int e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Int e -> List Int e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE toStream #-}
  toStreamIx :: Array L Int e -> Steps Id (Int, e)
toStreamIx = Steps Id e -> Steps Id (Int, e)
forall (m :: * -> *) e. Monad m => Steps m e -> Steps m (Int, e)
S.indexed (Steps Id e -> Steps Id (Int, e))
-> (Array L Int e -> Steps Id e)
-> Array L Int e
-> Steps Id (Int, 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 L Int e -> [e]) -> Array L Int e -> Steps Id e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List Int e -> [e]
forall ix e. List ix e -> [Elt ix e]
unList (List Int e -> [e])
-> (Array L Int e -> List Int e) -> Array L Int e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array L Int e -> List Int e
forall ix e. Array L ix e -> List ix e
lData
  {-# INLINE toStreamIx #-}