{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module      : Data.Massiv.Array.Delayed.Stream
-- Copyright   : (c) Alexey Kuleshevich 2019-2021
-- 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.ST
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)

-- | 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 :: Vector DS e -> Steps Id e
toSteps :: Vector DS e -> Steps Id e
toSteps = Vector DS e -> Steps Id e
coerce
{-# INLINE toSteps #-}

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

-- | /O(1)/ - Convert monadic `Steps` into delayed stream array
--
-- @since 0.5.0
fromStepsM :: Monad m => Steps m e -> m (Vector DS e)
fromStepsM :: Steps m e -> m (Vector DS e)
fromStepsM = (Steps Id e -> Vector DS e) -> m (Steps Id e) -> m (Vector DS e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Steps Id e -> Vector DS e
forall e. Steps Id e -> Array DS Ix1 e
DSArray (m (Steps Id e) -> m (Vector DS e))
-> (Steps m e -> m (Steps Id e)) -> Steps m e -> m (Vector DS 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 Shape DS Ix1 where
  linearSizeHint :: Array DS Ix1 e -> LengthHint
linearSizeHint = Steps Id e -> LengthHint
forall (m :: * -> *) e. Steps m e -> LengthHint
stepsSize (Steps Id e -> LengthHint)
-> (Array DS Ix1 e -> Steps Id e) -> Array DS Ix1 e -> LengthHint
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 linearSizeHint #-}

  linearSize :: Array DS Ix1 e -> Sz1
linearSize = Ix1 -> Sz1
forall ix. ix -> Sz ix
SafeSz (Ix1 -> Sz1) -> (Array DS Ix1 e -> Ix1) -> Array DS Ix1 e -> Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id Ix1 -> Ix1
forall a. Id a -> a
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
forall e. Array DS Ix1 e -> Steps Id e
dsArray
  {-# INLINE linearSize #-}

  outerSize :: Array DS Ix1 e -> Sz1
outerSize = Array DS Ix1 e -> Sz1
forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize
  {-# INLINE outerSize #-}

  isNull :: Array DS Ix1 e -> Bool
isNull = 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 isNull #-}


--TODO remove
instance Strategy DS where
  getComp :: Array DS ix e -> Comp
getComp Array DS ix e
_ = Comp
Seq
  setComp :: Comp -> Array DS ix e -> Array DS ix e
setComp Comp
_ = Array DS ix e -> Array DS ix e
forall a. a -> a
id


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 l. IsList l => [Item l] -> l
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 -> [Item (Steps Id e)] -> Steps Id e
forall l. IsList l => Ix1 -> [Item l] -> l
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 :: (Index ix, Source r e) => Array r ix e -> Vector DS e
toStreamArray :: Array r ix e -> Vector DS e
toStreamArray = Steps Id e -> Vector DS e
forall e. Steps Id e -> Array DS Ix1 e
DSArray (Steps Id e -> Vector DS e)
-> (Array r ix e -> Steps Id e) -> Array r ix e -> Vector DS 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, Index ix, Source r 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 #-}


-- | /O(n)/ - `size` implementation.
instance Load DS Ix1 e where

  makeArrayLinear :: Comp -> Sz1 -> (Ix1 -> e) -> Array DS Ix1 e
makeArrayLinear Comp
_ Sz1
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
. Sz1 -> (Ix1 -> e) -> Steps Id e
forall (m :: * -> *) e. Monad m => Sz1 -> (Ix1 -> e) -> Steps m e
S.generate Sz1
k
  {-# INLINE makeArrayLinear #-}
  replicate :: Comp -> Sz1 -> e -> Array DS Ix1 e
replicate Comp
_ Sz1
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)
-> (e -> Steps Id e) -> e -> Array DS Ix1 e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz1 -> e -> Steps Id e
forall (m :: * -> *) a. Monad m => Sz1 -> a -> Steps m a
S.replicate Sz1
k
  {-# INLINE replicate #-}

  iterArrayLinearST_ :: Scheduler s ()
-> Array DS Ix1 e -> (Ix1 -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
_scheduler Array DS Ix1 e
arr Ix1 -> e -> ST s ()
uWrite =
    ((Ix1, e) -> ST s ()) -> Steps (ST s) (Ix1, e) -> ST s ()
forall (m :: * -> *) e a.
Monad m =>
(e -> m a) -> Steps m e -> m ()
S.mapM_ ((Ix1 -> e -> ST s ()) -> (Ix1, e) -> ST s ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Ix1 -> e -> ST s ()
uWrite) (Steps (ST s) (Ix1, e) -> ST s ())
-> Steps (ST s) (Ix1, e) -> ST s ()
forall a b. (a -> b) -> a -> b
$ Steps (ST s) e -> Steps (ST s) (Ix1, e)
forall (m :: * -> *) e. Monad m => Steps m e -> Steps m (Ix1, e)
S.indexed (Steps (ST s) e -> Steps (ST s) (Ix1, e))
-> Steps (ST s) e -> Steps (ST s) (Ix1, e)
forall a b. (a -> b) -> a -> b
$ Steps Id e -> Steps (ST s) 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)
  {-# INLINE iterArrayLinearST_ #-}

  unsafeLoadIntoST :: MVector s r' e -> Array DS Ix1 e -> ST s (MVector s r' e)
unsafeLoadIntoST MVector s r' e
marr (DSArray sts) =
    MVector (PrimState (ST s)) r' e
-> LengthHint
-> Stream Id e
-> ST s (MVector (PrimState (ST s)) r' e)
forall r a (m :: * -> *).
(Manifest r a, PrimMonad m) =>
MVector (PrimState m) r a
-> LengthHint -> Stream Id a -> m (MVector (PrimState m) r a)
S.unstreamIntoM MVector s r' e
MVector (PrimState (ST s)) r' e
marr (Steps Id e -> LengthHint
forall (m :: * -> *) e. Steps m e -> LengthHint
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 unsafeLoadIntoST #-}

  unsafeLoadIntoIO :: MVector RealWorld r' e
-> Array DS Ix1 e -> IO (MVector RealWorld r' e)
unsafeLoadIntoIO MVector RealWorld r' e
marr Array DS Ix1 e
arr = ST RealWorld (MVector RealWorld r' e)
-> IO (MVector RealWorld r' e)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (MVector RealWorld r' e)
 -> IO (MVector RealWorld r' e))
-> ST RealWorld (MVector RealWorld r' e)
-> IO (MVector RealWorld r' e)
forall a b. (a -> b) -> a -> b
$ MVector RealWorld r' e
-> Array DS Ix1 e -> ST RealWorld (MVector RealWorld r' e)
forall r ix e r' s.
(Load r ix e, Manifest r' e) =>
MVector s r' e -> Array r ix e -> ST s (MArray s r' ix e)
unsafeLoadIntoST MVector RealWorld r' e
marr Array DS Ix1 e
arr
  {-# INLINE unsafeLoadIntoIO #-}


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