{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module      : Data.Massiv.Array.Delayed.Stream
-- Copyright   : (c) Alexey Kuleshevich 2019
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Delayed.Stream
  ( DS(..)
  , Array (..)
  , toStreamArray
  , toStreamM
  , toStreamIxM
  , toSteps
  , fromSteps
  , fromStepsM
  ) where

import Control.Applicative
import Control.Monad (void)
import Data.Coerce
import Data.Foldable
import Data.Massiv.Array.Delayed.Pull
import qualified Data.Massiv.Vector.Stream as S
import Data.Massiv.Core.Common
import GHC.Exts
import Prelude hiding (take, drop)
import Data.Vector.Fusion.Bundle.Size (upperBound)

-- | Delayed stream array that represents a sequence of values that can be loaded
-- sequentially. Important distinction from other arrays is that its size might no be
-- known until it is computed.
data DS = DS

newtype instance Array DS Ix1 e = DSArray
  { Array DS Ix1 e -> Steps Id e
dsArray :: S.Steps S.Id e
  }

-- | /O(1)/ - Convert delayed stream array into `Steps`.
--
-- @since 0.4.1
toSteps :: Array DS Ix1 e -> Steps Id e
toSteps :: Array DS Ix1 e -> Steps Id e
toSteps = Array DS Ix1 e -> Steps Id e
coerce
{-# INLINE toSteps #-}

-- | /O(1)/ - Convert `Steps` into delayed stream array
--
-- @since 0.4.1
fromSteps :: Steps Id e -> Array DS Ix1 e
fromSteps :: Steps Id e -> Array DS Ix1 e
fromSteps = Steps Id e -> Array DS Ix1 e
coerce
{-# INLINE fromSteps #-}

-- | /O(1)/ - Convert monadic `Steps` into delayed stream array
--
-- @since 0.5.0
fromStepsM :: Monad m => Steps m e -> m (Array DS Ix1 e)
fromStepsM :: Steps m e -> m (Array DS Ix1 e)
fromStepsM = (Steps Id e -> Array DS Ix1 e)
-> m (Steps Id e) -> m (Array DS Ix1 e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
DSArray (m (Steps Id e) -> m (Array DS Ix1 e))
-> (Steps m e -> m (Steps Id e)) -> Steps m e -> m (Array DS Ix1 e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Steps m e -> m (Steps Id e)
forall (m :: * -> *) (n :: * -> *) e.
(Monad m, Monad n) =>
Steps m e -> m (Steps n e)
S.transSteps
{-# INLINE fromStepsM #-}


instance Functor (Array DS Ix1) where
  fmap :: (a -> b) -> Array DS Ix1 a -> Array DS Ix1 b
fmap a -> b
f = Steps Id b -> Array DS Ix1 b
coerce (Steps Id b -> Array DS Ix1 b)
-> (Array DS Ix1 a -> Steps Id b)
-> Array DS Ix1 a
-> Array DS Ix1 b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Steps Id a -> Steps Id b
forall (m :: * -> *) e a.
Monad m =>
(e -> a) -> Steps m e -> Steps m a
S.map a -> b
f (Steps Id a -> Steps Id b)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Steps Id b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
dsArray
  {-# INLINE fmap #-}
  <$ :: a -> Array DS Ix1 b -> Array DS Ix1 a
(<$) a
e = Steps Id a -> Array DS Ix1 a
coerce (Steps Id a -> Array DS Ix1 a)
-> (Array DS Ix1 b -> Steps Id a)
-> Array DS Ix1 b
-> Array DS Ix1 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
e a -> Steps Id b -> Steps Id a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$) (Steps Id b -> Steps Id a)
-> (Array DS Ix1 b -> Steps Id b) -> Array DS Ix1 b -> Steps Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 b -> Steps Id b
forall e. Array DS Ix1 e -> Steps Id e
dsArray
  {-# INLINE (<$) #-}

instance Applicative (Array DS Ix1) where
  pure :: a -> Array DS Ix1 a
pure = Steps Id a -> Array DS Ix1 a
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (Steps Id a -> Array DS Ix1 a)
-> (a -> Steps Id a) -> a -> Array DS Ix1 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Steps Id a
forall (m :: * -> *) e. Monad m => e -> Steps m e
S.singleton
  {-# INLINE pure #-}
  <*> :: Array DS Ix1 (a -> b) -> Array DS Ix1 a -> Array DS Ix1 b
(<*>) Array DS Ix1 (a -> b)
a1 Array DS Ix1 a
a2 = Steps Id b -> Array DS Ix1 b
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (((a -> b) -> a -> b)
-> Steps Id (a -> b) -> Steps Id a -> Steps Id b
forall (m :: * -> *) a b e.
Monad m =>
(a -> b -> e) -> Steps m a -> Steps m b -> Steps m e
S.zipWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($) (Array DS Ix1 (a -> b) -> Steps Id (a -> b)
coerce Array DS Ix1 (a -> b)
a1) (Array DS Ix1 a -> Steps Id a
coerce Array DS Ix1 a
a2))
  {-# INLINE (<*>) #-}

#if MIN_VERSION_base(4,10,0)
  liftA2 :: (a -> b -> c) -> Array DS Ix1 a -> Array DS Ix1 b -> Array DS Ix1 c
liftA2 a -> b -> c
f Array DS Ix1 a
a1 Array DS Ix1 b
a2 = Steps Id c -> Array DS Ix1 c
forall e. Steps Id e -> Array DS Ix1 e
fromSteps ((a -> b -> c) -> Steps Id a -> Steps Id b -> Steps Id c
forall (m :: * -> *) a b e.
Monad m =>
(a -> b -> e) -> Steps m a -> Steps m b -> Steps m e
S.zipWith a -> b -> c
f (Array DS Ix1 a -> Steps Id a
coerce Array DS Ix1 a
a1) (Array DS Ix1 b -> Steps Id b
coerce Array DS Ix1 b
a2))
  {-# INLINE liftA2 #-}
#endif

instance Monad (Array DS Ix1) where
  return :: a -> Array DS Ix1 a
return = Steps Id a -> Array DS Ix1 a
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (Steps Id a -> Array DS Ix1 a)
-> (a -> Steps Id a) -> a -> Array DS Ix1 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Steps Id a
forall (m :: * -> *) e. Monad m => e -> Steps m e
S.singleton
  {-# INLINE return #-}
  >>= :: Array DS Ix1 a -> (a -> Array DS Ix1 b) -> Array DS Ix1 b
(>>=) Array DS Ix1 a
arr a -> Array DS Ix1 b
f = Steps Id b -> Array DS Ix1 b
coerce ((a -> Steps Id b) -> Steps Id a -> Steps Id b
forall (m :: * -> *) a e.
Monad m =>
(a -> Steps m e) -> Steps m a -> Steps m e
S.concatMap (Array DS Ix1 b -> Steps Id b
coerce (Array DS Ix1 b -> Steps Id b)
-> (a -> Array DS Ix1 b) -> a -> Steps Id b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Array DS Ix1 b
f) (Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
dsArray Array DS Ix1 a
arr))
  {-# INLINE (>>=) #-}


instance Foldable (Array DS Ix1) where
  foldr :: (a -> b -> b) -> b -> Array DS Ix1 a -> b
foldr a -> b -> b
f b
acc = Id b -> b
forall a. Id a -> a
S.unId (Id b -> b) -> (Array DS Ix1 a -> Id b) -> Array DS Ix1 a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b -> b) -> b -> Steps Id a -> Id b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Steps m a -> m b
S.foldrLazy a -> b -> b
f b
acc (Steps Id a -> Id b)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE foldr #-}
  foldl :: (b -> a -> b) -> b -> Array DS Ix1 a -> b
foldl b -> a -> b
f b
acc = Id b -> b
forall a. Id a -> a
S.unId (Id b -> b) -> (Array DS Ix1 a -> Id b) -> Array DS Ix1 a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> a -> b) -> b -> Steps Id a -> Id b
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Steps m a -> m b
S.foldlLazy b -> a -> b
f b
acc (Steps Id a -> Id b)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE foldl #-}
  foldl' :: (b -> a -> b) -> b -> Array DS Ix1 a -> b
foldl' b -> a -> b
f b
acc = Id b -> b
forall a. Id a -> a
S.unId (Id b -> b) -> (Array DS Ix1 a -> Id b) -> Array DS Ix1 a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> a -> b) -> b -> Steps Id a -> Id b
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Steps m a -> m b
S.foldl b -> a -> b
f b
acc (Steps Id a -> Id b)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE foldl' #-}
  foldr1 :: (a -> a -> a) -> Array DS Ix1 a -> a
foldr1 a -> a -> a
f = Id a -> a
forall a. Id a -> a
S.unId (Id a -> a) -> (Array DS Ix1 a -> Id a) -> Array DS Ix1 a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> Steps Id a -> Id a
forall (m :: * -> *) a.
Monad m =>
(a -> a -> a) -> Steps m a -> m a
S.foldr1Lazy a -> a -> a
f (Steps Id a -> Id a)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE foldr1 #-}
  foldl1 :: (a -> a -> a) -> Array DS Ix1 a -> a
foldl1 a -> a -> a
f = Id a -> a
forall a. Id a -> a
S.unId (Id a -> a) -> (Array DS Ix1 a -> Id a) -> Array DS Ix1 a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> Steps Id a -> Id a
forall (m :: * -> *) a.
Monad m =>
(a -> a -> a) -> Steps m a -> m a
S.foldl1Lazy a -> a -> a
f (Steps Id a -> Id a)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE foldl1 #-}
  toList :: Array DS Ix1 a -> [a]
toList = Steps Id a -> [a]
forall e. Steps Id e -> [e]
S.toList (Steps Id a -> [a])
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
coerce
  {-# INLINE toList #-}
  length :: Array DS Ix1 a -> Ix1
length = Id Ix1 -> Ix1
forall a. Id a -> a
S.unId (Id Ix1 -> Ix1)
-> (Array DS Ix1 a -> Id Ix1) -> Array DS Ix1 a -> Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Steps Id a -> Id Ix1
forall (m :: * -> *) a. Monad m => Steps m a -> m Ix1
S.length (Steps Id a -> Id Ix1)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
coerce
  {-# INLINE length #-}
  null :: Array DS Ix1 a -> Bool
null = Id Bool -> Bool
forall a. Id a -> a
S.unId (Id Bool -> Bool)
-> (Array DS Ix1 a -> Id Bool) -> Array DS Ix1 a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Steps Id a -> Id Bool
forall (m :: * -> *) a. Monad m => Steps m a -> m Bool
S.null (Steps Id a -> Id Bool)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
coerce
  {-# INLINE null #-}
  sum :: Array DS Ix1 a -> a
sum = Id a -> a
forall a. Id a -> a
S.unId (Id a -> a) -> (Array DS Ix1 a -> Id a) -> Array DS Ix1 a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> a -> Steps Id a -> Id a
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Steps m a -> m b
S.foldl a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0 (Steps Id a -> Id a)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE sum #-}
  product :: Array DS Ix1 a -> a
product = Id a -> a
forall a. Id a -> a
S.unId (Id a -> a) -> (Array DS Ix1 a -> Id a) -> Array DS Ix1 a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> a -> Steps Id a -> Id a
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Steps m a -> m b
S.foldl a -> a -> a
forall a. Num a => a -> a -> a
(*) a
1 (Steps Id a -> Id a)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE product #-}
  maximum :: Array DS Ix1 a -> a
maximum = Id a -> a
forall a. Id a -> a
S.unId (Id a -> a) -> (Array DS Ix1 a -> Id a) -> Array DS Ix1 a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> Steps Id a -> Id a
forall (m :: * -> *) a.
Monad m =>
(a -> a -> a) -> Steps m a -> m a
S.foldl1 a -> a -> a
forall a. Ord a => a -> a -> a
max (Steps Id a -> Id a)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE maximum #-}
  minimum :: Array DS Ix1 a -> a
minimum = Id a -> a
forall a. Id a -> a
S.unId (Id a -> a) -> (Array DS Ix1 a -> Id a) -> Array DS Ix1 a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> Steps Id a -> Id a
forall (m :: * -> *) a.
Monad m =>
(a -> a -> a) -> Steps m a -> m a
S.foldl1 a -> a -> a
forall a. Ord a => a -> a -> a
min (Steps Id a -> Id a)
-> (Array DS Ix1 a -> Steps Id a) -> Array DS Ix1 a -> Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 a -> Steps Id a
forall e. Array DS Ix1 e -> Steps Id e
toSteps
  {-# INLINE minimum #-}



instance Semigroup (Array DS Ix1 e) where
  <> :: Array DS Ix1 e -> Array DS Ix1 e -> Array DS Ix1 e
(<>) Array DS Ix1 e
a1 Array DS Ix1 e
a2 = Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (Array DS Ix1 e -> Steps Id e
coerce Array DS Ix1 e
a1 Steps Id e -> Steps Id e -> Steps Id e
forall (m :: * -> *) e.
Monad m =>
Steps m e -> Steps m e -> Steps m e
`S.append` Array DS Ix1 e -> Steps Id e
coerce Array DS Ix1 e
a2)
  {-# INLINE (<>) #-}


instance Monoid (Array DS Ix1 e) where
  mempty :: Array DS Ix1 e
mempty = Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
DSArray Steps Id e
forall (m :: * -> *) e. Monad m => Steps m e
S.empty
  {-# INLINE mempty #-}
  mappend :: Array DS Ix1 e -> Array DS Ix1 e -> Array DS Ix1 e
mappend = Array DS Ix1 e -> Array DS Ix1 e -> Array DS Ix1 e
forall a. Semigroup a => a -> a -> a
(<>)
  {-# INLINE mappend #-}

instance IsList (Array DS Ix1 e) where
  type Item (Array DS Ix1 e) = e
  fromList :: [Item (Array DS Ix1 e)] -> Array DS Ix1 e
fromList = Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (Steps Id e -> Array DS Ix1 e)
-> ([e] -> Steps Id e) -> [e] -> Array DS 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
  {-# INLINE fromList #-}
  fromListN :: Ix1 -> [Item (Array DS Ix1 e)] -> Array DS Ix1 e
fromListN Ix1
n = Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (Steps Id e -> Array DS Ix1 e)
-> ([e] -> Steps Id e) -> [e] -> Array DS Ix1 e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> [e] -> Steps Id e
forall (m :: * -> *) e. Monad m => Ix1 -> [e] -> Steps m e
S.fromListN Ix1
n
  {-# INLINE fromListN #-}
  toList :: Array DS Ix1 e -> [Item (Array DS Ix1 e)]
toList = Steps Id e -> [e]
forall e. Steps Id e -> [e]
S.toList (Steps Id e -> [e])
-> (Array DS Ix1 e -> Steps Id e) -> Array DS Ix1 e -> [e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 e -> Steps Id e
coerce
  {-# INLINE toList #-}


instance S.Stream DS Ix1 e where
  toStream :: Array DS Ix1 e -> Steps Id e
toStream = Array DS Ix1 e -> Steps Id e
coerce
  {-# INLINE toStream #-}
  toStreamIx :: Array DS 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 DS Ix1 e -> Steps Id e)
-> Array DS Ix1 e
-> Steps Id (Ix1, e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 e -> Steps Id e
coerce
  {-# INLINE toStreamIx #-}


-- | Flatten an array into a stream of values.
--
-- @since 0.4.1
toStreamArray :: Source r ix e => Array r ix e -> Array DS Ix1 e
toStreamArray :: Array r ix e -> Array DS Ix1 e
toStreamArray = Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
DSArray (Steps Id e -> Array DS Ix1 e)
-> (Array r ix e -> Steps Id e) -> Array r ix e -> Array DS Ix1 e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Steps Id e
forall r ix e (m :: * -> *).
(Monad m, Source r ix e) =>
Array r ix e -> Steps m e
S.steps
{-# INLINE[1] toStreamArray #-}
{-# RULES "toStreamArray/id" toStreamArray = id #-}

-- | /O(1)/ - Convert an array into monadic `Steps`
--
-- @since 0.5.0
toStreamM :: (Stream r ix e, Monad m) => Array r ix e -> Steps m e
toStreamM :: Array r ix e -> Steps m e
toStreamM = Steps Id e -> Steps m e
forall (m :: * -> *) e. Monad m => Steps Id e -> Steps m e
S.transStepsId (Steps Id e -> Steps m e)
-> (Array r ix e -> Steps Id e) -> Array r ix e -> Steps m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Steps Id e
forall r ix e. Stream r ix e => Array r ix e -> Steps Id e
toStream
{-# INLINE toStreamM #-}

-- | /O(1)/ - Convert an array into monadic `Steps`
--
-- @since 0.5.0
toStreamIxM :: (Stream r ix e, Monad m) => Array r ix e -> Steps m (ix, e)
toStreamIxM :: Array r ix e -> Steps m (ix, e)
toStreamIxM = Steps Id (ix, e) -> Steps m (ix, e)
forall (m :: * -> *) e. Monad m => Steps Id e -> Steps m e
S.transStepsId (Steps Id (ix, e) -> Steps m (ix, e))
-> (Array r ix e -> Steps Id (ix, e))
-> Array r ix e
-> Steps m (ix, e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Steps Id (ix, e)
forall r ix e. Stream r ix e => Array r ix e -> Steps Id (ix, e)
toStreamIx
{-# INLINE toStreamIxM #-}


instance Construct DS Ix1 e where
  setComp :: Comp -> Array DS Ix1 e -> Array DS Ix1 e
setComp Comp
_ Array DS Ix1 e
arr = Array DS Ix1 e
arr
  {-# INLINE setComp #-}

  makeArrayLinear :: Comp -> Sz Ix1 -> (Ix1 -> e) -> Array DS Ix1 e
makeArrayLinear Comp
_ (Sz Ix1
k) = Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (Steps Id e -> Array DS Ix1 e)
-> ((Ix1 -> e) -> Steps Id e) -> (Ix1 -> e) -> Array DS Ix1 e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> (Ix1 -> e) -> Steps Id e
forall (m :: * -> *) e. Monad m => Ix1 -> (Ix1 -> e) -> Steps m e
S.generate Ix1
k
  {-# INLINE makeArrayLinear #-}


instance Extract DS Ix1 e where
  unsafeExtract :: Ix1 -> Sz Ix1 -> Array DS Ix1 e -> Array (R DS) Ix1 e
unsafeExtract Ix1
sIx Sz Ix1
newSz = Steps Id e -> Array DS Ix1 e
forall e. Steps Id e -> Array DS Ix1 e
fromSteps (Steps Id e -> Array DS Ix1 e)
-> (Array DS Ix1 e -> Steps Id e)
-> Array DS Ix1 e
-> Array DS Ix1 e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> Ix1 -> Steps Id e -> Steps Id e
forall (m :: * -> *) a.
Monad m =>
Ix1 -> Ix1 -> Steps m a -> Steps m a
S.slice Ix1
sIx (Sz Ix1 -> Ix1
forall ix. Sz ix -> ix
unSz Sz Ix1
newSz) (Steps Id e -> Steps Id e)
-> (Array DS Ix1 e -> Steps Id e) -> Array DS Ix1 e -> Steps Id e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 e -> Steps Id e
forall e. Array DS Ix1 e -> Steps Id e
dsArray
  {-# INLINE unsafeExtract #-}

-- | /O(n)/ - `size` implementation.
instance Load DS Ix1 e where
  size :: Array DS Ix1 e -> Sz Ix1
size = Ix1 -> Sz Ix1
coerce (Ix1 -> Sz Ix1)
-> (Array DS Ix1 e -> Ix1) -> Array DS Ix1 e -> Sz Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id Ix1 -> Ix1
forall a. Id a -> a
S.unId (Id Ix1 -> Ix1)
-> (Array DS Ix1 e -> Id Ix1) -> Array DS Ix1 e -> Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Steps Id e -> Id Ix1
forall (m :: * -> *) a. Monad m => Steps m a -> m Ix1
S.length (Steps Id e -> Id Ix1)
-> (Array DS Ix1 e -> Steps Id e) -> Array DS Ix1 e -> Id Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 e -> Steps Id e
coerce
  {-# INLINE size #-}

  maxSize :: Array DS Ix1 e -> Maybe (Sz Ix1)
maxSize = Maybe Ix1 -> Maybe (Sz Ix1)
coerce (Maybe Ix1 -> Maybe (Sz Ix1))
-> (Array DS Ix1 e -> Maybe Ix1)
-> Array DS Ix1 e
-> Maybe (Sz Ix1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Size -> Maybe Ix1
upperBound (Size -> Maybe Ix1)
-> (Array DS Ix1 e -> Size) -> Array DS Ix1 e -> Maybe Ix1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Steps Id e -> Size
forall (m :: * -> *) e. Steps m e -> Size
stepsSize (Steps Id e -> Size)
-> (Array DS Ix1 e -> Steps Id e) -> Array DS Ix1 e -> Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 e -> Steps Id e
forall e. Array DS Ix1 e -> Steps Id e
dsArray
  {-# INLINE maxSize #-}

  isEmpty :: Array DS Ix1 e -> Bool
isEmpty = Id Bool -> Bool
forall a. Id a -> a
S.unId (Id Bool -> Bool)
-> (Array DS Ix1 e -> Id Bool) -> Array DS Ix1 e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Steps Id e -> Id Bool
forall (m :: * -> *) a. Monad m => Steps m a -> m Bool
S.null (Steps Id e -> Id Bool)
-> (Array DS Ix1 e -> Steps Id e) -> Array DS Ix1 e -> Id Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DS Ix1 e -> Steps Id e
coerce
  {-# INLINE isEmpty #-}

  getComp :: Array DS Ix1 e -> Comp
getComp Array DS Ix1 e
_ = Comp
Seq
  {-# INLINE getComp #-}

  loadArrayM :: Scheduler m () -> Array DS Ix1 e -> (Ix1 -> e -> m ()) -> m ()
loadArrayM Scheduler m ()
_scheduler Array DS Ix1 e
arr Ix1 -> e -> m ()
uWrite =
    case Steps Id e -> Size
forall (m :: * -> *) e. Steps m e -> Size
stepsSize (Array DS Ix1 e -> Steps Id e
forall e. Array DS Ix1 e -> Steps Id e
dsArray Array DS Ix1 e
arr) of
      S.Exact Ix1
_ ->
        m Ix1 -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Ix1 -> m ()) -> m Ix1 -> m ()
forall a b. (a -> b) -> a -> b
$ (Ix1 -> e -> m Ix1) -> Ix1 -> Steps m e -> m Ix1
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Steps m b -> m a
S.foldlM (\Ix1
i e
e -> Ix1 -> e -> m ()
uWrite Ix1
i e
e m () -> m Ix1 -> m Ix1
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix1 -> m Ix1
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ix1
i Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1)) Ix1
0 (Steps Id e -> Steps m e
forall (m :: * -> *) e. Monad m => Steps Id e -> Steps m e
S.transStepsId (Array DS Ix1 e -> Steps Id e
coerce Array DS Ix1 e
arr))
      Size
_ -> [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Loading Stream array is not supported with loadArrayM"
  {-# INLINE loadArrayM #-}

  unsafeLoadIntoS :: MArray (PrimState m) r' Ix1 e
-> Array DS Ix1 e -> m (MArray (PrimState m) r' Ix1 e)
unsafeLoadIntoS MArray (PrimState m) r' Ix1 e
marr (DSArray sts) =
    MArray (PrimState m) r' Ix1 e
-> Size -> Stream Id e -> m (MArray (PrimState m) r' Ix1 e)
forall r a (m :: * -> *).
(Mutable r Ix1 a, PrimMonad m) =>
MArray (PrimState m) r Ix1 a
-> Size -> Stream Id a -> m (MArray (PrimState m) r Ix1 a)
S.unstreamIntoM MArray (PrimState m) r' Ix1 e
marr (Steps Id e -> Size
forall (m :: * -> *) e. Steps m e -> Size
stepsSize Steps Id e
sts) (Steps Id e -> Stream Id e
forall (m :: * -> *) e. Steps m e -> Stream m e
stepsStream Steps Id e
sts)
  {-# INLINE unsafeLoadIntoS #-}

  unsafeLoadIntoM :: MArray RealWorld r' Ix1 e
-> Array DS Ix1 e -> m (MArray RealWorld r' Ix1 e)
unsafeLoadIntoM MArray RealWorld r' Ix1 e
marr Array DS Ix1 e
arr = IO (MArray RealWorld r' Ix1 e) -> m (MArray RealWorld r' Ix1 e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MArray RealWorld r' Ix1 e) -> m (MArray RealWorld r' Ix1 e))
-> IO (MArray RealWorld r' Ix1 e) -> m (MArray RealWorld r' Ix1 e)
forall a b. (a -> b) -> a -> b
$ MArray (PrimState IO) r' Ix1 e
-> Array DS Ix1 e -> IO (MArray (PrimState IO) r' Ix1 e)
forall r ix e r' (m :: * -> *).
(Load r ix e, Mutable r' ix e, PrimMonad m) =>
MArray (PrimState m) r' ix e
-> Array r ix e -> m (MArray (PrimState m) r' ix e)
unsafeLoadIntoS MArray RealWorld r' Ix1 e
MArray (PrimState IO) r' Ix1 e
marr Array DS Ix1 e
arr
  {-# INLINE unsafeLoadIntoM #-}


-- cons :: e -> Array DS Ix1 e -> Array DS Ix1 e
-- cons e = coerce . S.cons e . dsArray
-- {-# INLINE cons #-}

-- uncons :: Array DS Ix1 e -> Maybe (e, Array DS Ix1 e)
-- uncons = coerce . S.uncons . dsArray
-- {-# INLINE uncons #-}

-- snoc :: Array DS Ix1 e -> e -> Array DS Ix1 e
-- snoc (DSArray sts) e = DSArray (S.snoc sts e)
-- {-# INLINE snoc #-}


-- TODO: skip the stride while loading
-- instance StrideLoad DS Ix1 e where
--   loadArrayWithStrideM scheduler stride resultSize arr uWrite =
--     let strideIx = unStride stride
--         DIArray (DArray _ _ f) = arr
--     in loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !start ->
--           scheduleWork scheduler $
--           iterLinearM_ resultSize start (totalElem resultSize) (numWorkers scheduler) (<) $
--             \ !i ix -> uWrite i (f (liftIndex2 (*) strideIx ix))
--   {-# INLINE loadArrayWithStrideM #-}