{-# 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
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
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 Seq !sz f = do
  --   let !(k, szL) = unconsSz sz
  --   loopDeepM 0 (< coerce k) (+ 1) (emptyR Seq) $ \i acc -> do
  --     e <- generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))
  --     return (cons e acc)
  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 #-}

 -- 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 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 #-}


-- | Construct an array backed by linked lists from any source array
--
-- @since 0.4.0
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 -- Array to show
  -> 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

-- | Helper function for declaring `Show` instances for arrays
--
-- @since 0.4.0
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) -- ^ Modifier
  -> Int
  -> Array r ix e -- Array to show
  -> 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


-- | 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 {-# 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 #-}