{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
#include "inline.hs"
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
{-# 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)