{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module      : Data.Massiv.Array.Delayed.Stream
-- Copyright   : (c) Alexey Kuleshevich 2019-2022
-- 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 Data.Massiv.Core.Common
import qualified Data.Massiv.Vector.Stream as S
import GHC.Exts
import Prelude hiding (drop, take)

-- | 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
  { forall e. 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 :: forall e. Array DS Ix1 e -> Steps Id e
toSteps = coerce :: forall a b. Coercible a b => a -> b
coerce
{-# INLINE toSteps #-}

-- | /O(1)/ - Convert `Steps` into delayed stream array
--
-- @since 0.4.1
fromSteps :: Steps Id e -> Vector DS e
fromSteps :: forall e. Steps Id e -> Vector DS e
fromSteps = coerce :: forall a b. Coercible a b => a -> b
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 :: forall (m :: * -> *) e. Monad m => Steps m e -> m (Vector DS e)
fromStepsM = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall e. Steps Id e -> Vector DS e
DSArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall e. Array DS Ix1 e -> LengthHint
linearSizeHint = forall (m :: * -> *) e. Steps m e -> LengthHint
stepsSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e. Array DS Ix1 e -> Steps Id e
dsArray
  {-# INLINE linearSizeHint #-}

  linearSize :: forall e. Array DS Ix1 e -> Sz1
linearSize = forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Id a -> a
unId forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Steps m a -> m Ix1
S.length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e. Array DS Ix1 e -> Steps Id e
dsArray
  {-# INLINE linearSize #-}

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

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

-- | For now only `Seq` strategy.
instance Strategy DS where
  getComp :: forall ix e. Array DS ix e -> Comp
getComp Array DS ix e
_ = Comp
Seq
  setComp :: forall ix e. Comp -> Array DS ix e -> Array DS ix e
setComp Comp
_ = forall a. a -> a
id
  repr :: DS
repr = DS
DS

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

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

#if MIN_VERSION_base(4,10,0)
  liftA2 :: forall a b c.
(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 = forall e. Steps Id e -> Vector DS e
fromSteps (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 (coerce :: forall a b. Coercible a b => a -> b
coerce Array DS Ix1 a
a1) (coerce :: forall a b. Coercible a b => a -> b
coerce Array DS Ix1 b
a2))
  {-# INLINE liftA2 #-}
#endif

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

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

instance Monoid (Array DS Ix1 e) where
  mempty :: Array DS Ix1 e
mempty = forall e. Steps Id e -> Vector DS e
DSArray forall (m :: * -> *) e. Monad m => Steps m e
S.empty
  {-# INLINE mempty #-}
#if !MIN_VERSION_base(4,11,0)
  mappend = (<>)
  {-# INLINE mappend #-}
#endif

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

instance S.Stream DS Ix1 e where
  toStream :: Array DS Ix1 e -> Steps Id e
toStream = coerce :: forall a b. Coercible a b => a -> b
coerce
  {-# INLINE toStream #-}
  toStreamIx :: Array DS Ix1 e -> Steps Id (Ix1, e)
toStreamIx = forall (m :: * -> *) e. Monad m => Steps m e -> Steps m (Ix1, e)
S.indexed forall b c a. (b -> c) -> (a -> b) -> a -> c
. coerce :: forall a b. Coercible a b => a -> b
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 :: forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Vector DS e
toStreamArray = forall e. Steps Id e -> Vector DS e
DSArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall r ix e (m :: * -> *).
(Stream r ix e, Monad m) =>
Array r ix e -> Steps m e
toStreamM = forall (m :: * -> *) e. Monad m => Steps Id e -> Steps m e
S.transStepsId forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall r ix e (m :: * -> *).
(Stream r ix e, Monad m) =>
Array r ix e -> Steps m (ix, e)
toStreamIxM = forall (m :: * -> *) e. Monad m => Steps Id e -> Steps m e
S.transStepsId forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall e. Steps Id e -> Vector DS e
fromSteps forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall e. Steps Id e -> Vector DS e
fromSteps forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Sz1 -> a -> Steps m a
S.replicate Sz1
k
  {-# INLINE replicate #-}

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

  unsafeLoadIntoST :: forall r' s.
Manifest r' e =>
MVector s r' e -> Array DS Ix1 e -> ST s (MVector s r' e)
unsafeLoadIntoST MVector s r' e
marr (DSArray Steps Id e
sts) =
    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
marr (forall (m :: * -> *) e. Steps m e -> LengthHint
stepsSize Steps Id e
sts) (forall (m :: * -> *) e. Steps m e -> Stream m e
stepsStream Steps Id e
sts)
  {-# INLINE unsafeLoadIntoST #-}

  unsafeLoadIntoIO :: forall r'.
Manifest r' e =>
MVector RealWorld r' e
-> Array DS Ix1 e -> IO (MVector RealWorld r' e)
unsafeLoadIntoIO MVector RealWorld r' e
marr Array DS Ix1 e
arr = forall a. ST RealWorld a -> IO a
stToIO forall a b. (a -> b) -> a -> b
$ 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_ #-}