{-# OPTIONS_GHC -fno-warn-orphans #-}

{-# LANGUAGE CPP           #-}
{-# LANGUAGE MagicHash     #-}
{-# LANGUAGE UnboxedTuples #-}

#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Array
-- Copyright   : (c) 2019 Composewell Technologies
--
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.Array
    ( Array(..)

    , foldl'
    , foldr

    , length

    , writeN
    , write

    , toStreamD
    , toStreamDRev

    , toStream
    , toStreamRev
    , read

    , fromListN
    , fromList
    , fromStreamDN
    , fromStreamD

    , fromStreamN
    , fromStream

    , streamFold
    , fold
    )
where

import Prelude hiding (foldr, length, read)
#if !MIN_VERSION_primitive(0,7,1)
import Control.DeepSeq (NFData(..))
#endif
import Control.Monad (when)
import Control.Monad.IO.Class (liftIO, MonadIO)
import GHC.IO (unsafePerformIO)
import GHC.Base (Int(..))
import Data.Functor.Identity (runIdentity)
import Data.Primitive.Array hiding (fromList, fromListN)
import qualified GHC.Exts as Exts

import Streamly.Internal.Data.Unfold.Types (Unfold(..))
import Streamly.Internal.Data.Fold.Types (Fold(..))
import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
import Streamly.Internal.Data.Stream.Serial (SerialT)

import qualified Streamly.Internal.Data.Stream.StreamD as D

{-# NOINLINE bottomElement #-}
bottomElement :: a
bottomElement :: forall a. a
bottomElement = forall a. HasCallStack => a
undefined

{-# INLINE_NORMAL toStreamD #-}
toStreamD :: Monad m => Array a -> D.Stream m a
toStreamD :: forall (m :: * -> *) a. Monad m => Array a -> Stream m a
toStreamD Array a
arr = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {p}. Monad m => p -> Int -> m (Step Int a)
step Int
0
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i
        | Int
i forall a. Eq a => a -> a -> Bool
== forall a. Array a -> Int
length Array a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step p
_ (I# Int#
i) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
        case forall a. Array# a -> Int# -> (# a #)
Exts.indexArray# (forall a. Array a -> Array# a
array# Array a
arr) Int#
i of
            (# a
x #) -> forall s a. a -> s -> Step s a
D.Yield a
x ((Int# -> Int
I# Int#
i) forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE length #-}
length :: Array a -> Int
length :: forall a. Array a -> Int
length Array a
arr = forall a. Array a -> Int
sizeofArray Array a
arr

{-# INLINE_NORMAL toStreamDRev #-}
toStreamDRev :: Monad m => Array a -> D.Stream m a
toStreamDRev :: forall (m :: * -> *) a. Monad m => Array a -> Stream m a
toStreamDRev Array a
arr = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {p}. Monad m => p -> Int -> m (Step Int a)
step (forall a. Array a -> Int
length Array a
arr forall a. Num a => a -> a -> a
- Int
1)
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step p
_ (I# Int#
i) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
        case forall a. Array# a -> Int# -> (# a #)
Exts.indexArray# (forall a. Array a -> Array# a
array# Array a
arr) Int#
i of
            (# a
x #) -> forall s a. a -> s -> Step s a
D.Yield a
x ((Int# -> Int
I# Int#
i) forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE_NORMAL foldl' #-}
foldl' :: (b -> a -> b) -> b -> Array a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Array a -> b
foldl' b -> a -> b
f b
z Array a
arr = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
D.foldl' b -> a -> b
f b
z forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Array a -> Stream m a
toStreamD Array a
arr

{-# INLINE_NORMAL foldr #-}
foldr :: (a -> b -> b) -> b -> Array a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
foldr a -> b -> b
f b
z Array a
arr = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr a -> b -> b
f b
z forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Array a -> Stream m a
toStreamD Array a
arr

-- writeN n = S.evertM (fromStreamDN n)
{-# INLINE_NORMAL writeN #-}
writeN :: MonadIO m => Int -> Fold m a (Array a)
writeN :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (Array a)
writeN Int
limit = forall (m :: * -> *) a b s.
(s -> a -> m s) -> m s -> (s -> m b) -> Fold m a b
Fold forall {m :: * -> *} {a}.
MonadIO m =>
(MutableArray RealWorld a, Int)
-> a -> m (MutableArray RealWorld a, Int)
step forall {a}. m (MutableArray RealWorld a, Int)
initial forall {m :: * -> *} {a}.
MonadIO m =>
(MutableArray RealWorld a, Int) -> m (Array a)
extract
  where
    initial :: m (MutableArray RealWorld a, Int)
initial = do
        MutableArray RealWorld a
marr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
limit forall a. a
bottomElement
        forall (m :: * -> *) a. Monad m => a -> m a
return (MutableArray RealWorld a
marr, Int
0)
    step :: (MutableArray RealWorld a, Int)
-> a -> m (MutableArray RealWorld a, Int)
step (MutableArray RealWorld a
marr, Int
i) a
x
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
limit = forall (m :: * -> *) a. Monad m => a -> m a
return (MutableArray RealWorld a
marr, Int
i)
        | Bool
otherwise = do
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld a
marr Int
i a
x
            forall (m :: * -> *) a. Monad m => a -> m a
return (MutableArray RealWorld a
marr, Int
i forall a. Num a => a -> a -> a
+ Int
1)
    extract :: (MutableArray RealWorld a, Int) -> m (Array a)
extract (MutableArray RealWorld a
marr, Int
len) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
freezeArray MutableArray RealWorld a
marr Int
0 Int
len

{-# INLINE_NORMAL write #-}
write :: MonadIO m => Fold m a (Array a)
write :: forall (m :: * -> *) a. MonadIO m => Fold m a (Array a)
write = forall (m :: * -> *) a b s.
(s -> a -> m s) -> m s -> (s -> m b) -> Fold m a b
Fold forall {m :: * -> *} {a}.
MonadIO m =>
(MutableArray RealWorld a, Int, Int)
-> a -> m (MutableArray RealWorld a, Int, Int)
step forall {a}. m (MutableArray RealWorld a, Int, Int)
initial forall {m :: * -> *} {a} {c}.
MonadIO m =>
(MutableArray RealWorld a, Int, c) -> m (Array a)
extract
  where
    initial :: m (MutableArray RealWorld a, Int, Int)
initial = do
        MutableArray RealWorld a
marr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
0 forall a. a
bottomElement
        forall (m :: * -> *) a. Monad m => a -> m a
return (MutableArray RealWorld a
marr, Int
0, Int
0)
    step :: (MutableArray RealWorld a, Int, Int)
-> a -> m (MutableArray RealWorld a, Int, Int)
step (MutableArray RealWorld a
marr, Int
i, Int
capacity) a
x
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
capacity =
            let newCapacity :: Int
newCapacity = forall a. Ord a => a -> a -> a
max (Int
capacity forall a. Num a => a -> a -> a
* Int
2) Int
1
             in do MutableArray RealWorld a
newMarr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
newCapacity forall a. a
bottomElement
                   forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
copyMutableArray MutableArray RealWorld a
newMarr Int
0 MutableArray RealWorld a
marr Int
0 Int
i
                   forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld a
newMarr Int
i a
x
                   forall (m :: * -> *) a. Monad m => a -> m a
return (MutableArray RealWorld a
newMarr, Int
i forall a. Num a => a -> a -> a
+ Int
1, Int
newCapacity)
        | Bool
otherwise = do
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld a
marr Int
i a
x
            forall (m :: * -> *) a. Monad m => a -> m a
return (MutableArray RealWorld a
marr, Int
i forall a. Num a => a -> a -> a
+ Int
1, Int
capacity)
    extract :: (MutableArray RealWorld a, Int, c) -> m (Array a)
extract (MutableArray RealWorld a
marr, Int
len, c
_) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
freezeArray MutableArray RealWorld a
marr Int
0 Int
len

{-# INLINE_NORMAL fromStreamDN #-}
fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
fromStreamDN :: forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (Array a)
fromStreamDN Int
limit Stream m a
str = do
    MutableArray RealWorld a
marr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray (forall a. Ord a => a -> a -> a
max Int
limit Int
0) forall a. a
bottomElement
    Int
i <-
        forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> m b
D.foldlM'
            (\Int
i a
x -> Int
i seq :: forall a b. a -> b -> b
`seq` (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld a
marr Int
i a
x) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i forall a. Num a => a -> a -> a
+ Int
1))
            Int
0 forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *) a. Monad m => Int -> Stream m a -> Stream m a
D.take Int
limit Stream m a
str
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
freezeArray MutableArray RealWorld a
marr Int
0 Int
i

{-# INLINE fromStreamD #-}
fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
fromStreamD :: forall (m :: * -> *) a. MonadIO m => Stream m a -> m (Array a)
fromStreamD Stream m a
str = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.runFold forall (m :: * -> *) a. MonadIO m => Fold m a (Array a)
write Stream m a
str

{-# INLINABLE fromListN #-}
fromListN :: Int -> [a] -> Array a
fromListN :: forall a. Int -> [a] -> Array a
fromListN Int
n [a]
xs = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (Array a)
fromStreamDN Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

{-# INLINABLE fromList #-}
fromList :: [a] -> Array a
fromList :: forall a. [a] -> Array a
fromList [a]
xs = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => Stream m a -> m (Array a)
fromStreamD forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

#if !MIN_VERSION_primitive(0,7,1)
instance NFData a => NFData (Array a) where
    {-# INLINE rnf #-}
    rnf = foldl' (\_ x -> rnf x) ()
#endif

{-# INLINE fromStreamN #-}
fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
fromStreamN :: forall (m :: * -> *) a.
MonadIO m =>
Int -> SerialT m a -> m (Array a)
fromStreamN Int
n SerialT m a
m = do
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n forall a. Ord a => a -> a -> Bool
< Int
0) forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => [Char] -> a
error [Char]
"fromStreamN: negative write count specified"
    forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (Array a)
fromStreamDN Int
n forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
t m a -> Stream m a
D.toStreamD SerialT m a
m

{-# INLINE fromStream #-}
fromStream :: MonadIO m => SerialT m a -> m (Array a)
fromStream :: forall (m :: * -> *) a. MonadIO m => SerialT m a -> m (Array a)
fromStream SerialT m a
m = forall (m :: * -> *) a. MonadIO m => Stream m a -> m (Array a)
fromStreamD forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
t m a -> Stream m a
D.toStreamD SerialT m a
m

{-# INLINE_EARLY toStream #-}
toStream :: (Monad m, IsStream t) => Array a -> t m a
toStream :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t) =>
Array a -> t m a
toStream = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
D.fromStreamD forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Array a -> Stream m a
toStreamD

{-# INLINE_EARLY toStreamRev #-}
toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
toStreamRev :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t) =>
Array a -> t m a
toStreamRev = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
D.fromStreamD forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Array a -> Stream m a
toStreamDRev

{-# INLINE fold #-}
fold :: Monad m => Fold m a b -> Array a -> m b
fold :: forall (m :: * -> *) a b. Monad m => Fold m a b -> Array a -> m b
fold Fold m a b
f Array a
arr = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.runFold Fold m a b
f (forall (m :: * -> *) a. Monad m => Array a -> Stream m a
toStreamD Array a
arr)

{-# INLINE streamFold #-}
streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
streamFold :: forall (m :: * -> *) a b.
Monad m =>
(SerialT m a -> m b) -> Array a -> m b
streamFold SerialT m a -> m b
f Array a
arr = SerialT m a -> m b
f (forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t) =>
Array a -> t m a
toStream Array a
arr)

{-# INLINE_NORMAL read #-}
read :: Monad m => Unfold m (Array a) a
read :: forall (m :: * -> *) a. Monad m => Unfold m (Array a) a
read = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {m :: * -> *} {a}.
Monad m =>
(Array a, Int) -> m (Step (Array a, Int) a)
step forall {m :: * -> *} {b} {a}. (Monad m, Num b) => a -> m (a, b)
inject
  where
    inject :: a -> m (a, b)
inject a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return (a
arr, b
0)
    step :: (Array a, Int) -> m (Step (Array a, Int) a)
step (Array a
arr, Int
i)
        | Int
i forall a. Eq a => a -> a -> Bool
== forall a. Array a -> Int
length Array a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (Array a
arr, (I# Int#
i)) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
        case forall a. Array# a -> Int# -> (# a #)
Exts.indexArray# (forall a. Array a -> Array# a
array# Array a
arr) Int#
i of
            (# a
x #) -> forall s a. a -> s -> Step s a
D.Yield a
x (Array a
arr, Int# -> Int
I# Int#
i forall a. Num a => a -> a -> a
+ Int
1)