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

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

#include "inline.hs"

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

    -- XXX should it be just Array instead? We should be able to replace one
    -- array type with another easily.
      PrimArray(..)

    -- XXX Prim should be exported from Data.Prim module?
    , Prim(..)

    , foldl'
    , foldr

    , length

    , writeN
    , write

    , toStreamD
    , toStreamDRev

    , toStream
    , toStreamRev
    , read
    , readSlice

    , fromListN
    , fromList
    , fromStreamDN
    , fromStreamD

    , fromStreamN
    , fromStream

    , streamFold
    , fold
    )
where

import Prelude hiding (foldr, length, read)
import Control.DeepSeq (NFData(..))
import Control.Monad (when)
import Control.Monad.IO.Class (liftIO, MonadIO)
import GHC.IO (unsafePerformIO)
import Data.Primitive.Types (Prim(..))

import Streamly.Internal.Data.Prim.Array.Types
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

{-# INLINE_NORMAL toStreamD #-}
toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
toStreamD :: forall a (m :: * -> *).
(Prim a, Monad m) =>
PrimArray a -> Stream m a
toStreamD PrimArray 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. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step p
_ Int
i = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i) (Int
i forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE length #-}
length :: Prim a => PrimArray a -> Int
length :: forall a. Prim a => PrimArray a -> Int
length PrimArray a
arr = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray a
arr

{-# INLINE_NORMAL toStreamDRev #-}
toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
toStreamDRev :: forall a (m :: * -> *).
(Prim a, Monad m) =>
PrimArray a -> Stream m a
toStreamDRev PrimArray 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. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray 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
_ Int
i = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i) (Int
i forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE_NORMAL foldl' #-}
foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldl' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldl' = forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldlPrimArray'

{-# INLINE_NORMAL foldr #-}
foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldr :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldr = forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
foldrPrimArray

-- writeN n = S.evertM (fromStreamDN n)
{-# INLINE_NORMAL writeN #-}
writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
writeN :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Int -> Fold m a (PrimArray 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, Prim a) =>
(MutablePrimArray RealWorld a, Int)
-> a -> m (MutablePrimArray RealWorld a, Int)
step m (MutablePrimArray RealWorld a, Int)
initial forall {m :: * -> *} {a} {b}.
MonadIO m =>
(MutablePrimArray RealWorld a, b) -> m (PrimArray a)
extract
  where
    initial :: m (MutablePrimArray RealWorld a, Int)
initial = do
        MutablePrimArray 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, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
limit
        forall (m :: * -> *) a. Monad m => a -> m a
return (MutablePrimArray RealWorld a
marr, Int
0)
    step :: (MutablePrimArray RealWorld a, Int)
-> a -> m (MutablePrimArray RealWorld a, Int)
step (MutablePrimArray 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 (MutablePrimArray 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 a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld a
marr Int
i a
x
            forall (m :: * -> *) a. Monad m => a -> m a
return (MutablePrimArray RealWorld a
marr, Int
i forall a. Num a => a -> a -> a
+ Int
1)
    extract :: (MutablePrimArray RealWorld a, b) -> m (PrimArray a)
extract (MutablePrimArray RealWorld a
marr, b
_) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld a
marr

{-# INLINE_NORMAL write #-}
write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
write :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Fold m a (PrimArray 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, Prim a) =>
(MutablePrimArray RealWorld a, Int, Int)
-> a -> m (MutablePrimArray RealWorld a, Int, Int)
step m (MutablePrimArray RealWorld a, Int, Int)
initial forall {m :: * -> *} {a} {c}.
(MonadIO m, Prim a) =>
(MutablePrimArray RealWorld a, Int, c) -> m (PrimArray a)
extract
  where
    initial :: m (MutablePrimArray RealWorld a, Int, Int)
initial = do
        MutablePrimArray 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, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
0
        forall (m :: * -> *) a. Monad m => a -> m a
return (MutablePrimArray RealWorld a
marr, Int
0, Int
0)
    step :: (MutablePrimArray RealWorld a, Int, Int)
-> a -> m (MutablePrimArray RealWorld a, Int, Int)
step (MutablePrimArray 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 MutablePrimArray 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, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray RealWorld a
marr Int
newCapacity
                   forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld a
newMarr Int
i a
x
                   forall (m :: * -> *) a. Monad m => a -> m a
return (MutablePrimArray 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 a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld a
marr Int
i a
x
            forall (m :: * -> *) a. Monad m => a -> m a
return (MutablePrimArray RealWorld a
marr, Int
i forall a. Num a => a -> a -> a
+ Int
1, Int
capacity)
    extract :: (MutablePrimArray RealWorld a, Int, c) -> m (PrimArray a)
extract (MutablePrimArray RealWorld a
marr, Int
len, c
_) = do forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> m ()
shrinkMutablePrimArray MutablePrimArray 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 =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld a
marr

{-# INLINE_NORMAL fromStreamDN #-}
fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
fromStreamDN :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Int -> Stream m a -> m (PrimArray a)
fromStreamDN Int
limit Stream m a
str = do
    MutablePrimArray 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, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (forall a. Ord a => a -> a -> a
max Int
limit Int
0)
    Int
_ <-
        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 a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray 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 =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld a
marr

{-# INLINE fromStreamD #-}
fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
fromStreamD :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Stream m a -> m (PrimArray 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, Prim a) =>
Fold m a (PrimArray a)
write Stream m a
str

{-# INLINABLE fromListN #-}
fromListN :: Prim a => Int -> [a] -> PrimArray a
fromListN :: forall a. Prim a => Int -> [a] -> PrimArray a
fromListN Int
n [a]
xs = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Int -> Stream m a -> m (PrimArray 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 :: Prim a => [a] -> PrimArray a
fromList :: forall a. Prim a => [a] -> PrimArray a
fromList [a]
xs = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Stream m a -> m (PrimArray a)
fromStreamD forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

instance Prim a => NFData (PrimArray a) where
    {-# INLINE rnf #-}
    rnf :: PrimArray a -> ()
rnf = forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldl' (\()
_ a
_ -> ()) ()

{-# INLINE fromStreamN #-}
fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
fromStreamN :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Int -> SerialT m a -> m (PrimArray 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, Prim a) =>
Int -> Stream m a -> m (PrimArray 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, Prim a) => SerialT m a -> m (PrimArray a)
fromStream :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
SerialT m a -> m (PrimArray a)
fromStream SerialT m a
m = forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Stream m a -> m (PrimArray 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 :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
toStream :: forall a (m :: * -> *) (t :: (* -> *) -> * -> *).
(Prim a, Monad m, IsStream t) =>
PrimArray 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 a (m :: * -> *).
(Prim a, Monad m) =>
PrimArray a -> Stream m a
toStreamD

{-# INLINE_EARLY toStreamRev #-}
toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
toStreamRev :: forall a (m :: * -> *) (t :: (* -> *) -> * -> *).
(Prim a, Monad m, IsStream t) =>
PrimArray 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 a (m :: * -> *).
(Prim a, Monad m) =>
PrimArray a -> Stream m a
toStreamDRev

{-# INLINE fold #-}
fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
fold :: forall a (m :: * -> *) b.
(Prim a, Monad m) =>
Fold m a b -> PrimArray a -> m b
fold Fold m a b
f PrimArray 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 a (m :: * -> *).
(Prim a, Monad m) =>
PrimArray a -> Stream m a
toStreamD PrimArray a
arr)

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

{-# INLINE_NORMAL read #-}
read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
read :: forall a (m :: * -> *).
(Prim a, Monad m) =>
Unfold m (PrimArray a) a
read = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {a} {m :: * -> *}.
(Prim a, Monad m) =>
(PrimArray a, Int) -> m (Step (PrimArray 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 :: (PrimArray a, Int) -> m (Step (PrimArray a, Int) a)
step (PrimArray a
arr, Int
i)
        | Int
i forall a. Eq a => a -> a -> Bool
== forall a. Prim a => PrimArray a -> Int
length PrimArray a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (PrimArray a
arr, Int
i) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i) (PrimArray a
arr, Int
i forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE_NORMAL readSlice #-}
readSlice :: (Prim a, Monad m) => Int -> Int -> Unfold m (PrimArray a) a
readSlice :: forall a (m :: * -> *).
(Prim a, Monad m) =>
Int -> Int -> Unfold m (PrimArray a) a
readSlice Int
off Int
len = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {a} {m :: * -> *}.
(Prim a, Monad m) =>
(PrimArray a, Int) -> m (Step (PrimArray a, Int) a)
step forall {m :: * -> *} {a}. Monad m => a -> m (a, Int)
inject
  where
    inject :: a -> m (a, Int)
inject a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return (a
arr, Int
off)
    step :: (PrimArray a, Int) -> m (Step (PrimArray a, Int) a)
step (PrimArray a
arr, Int
i)
        | Int
i forall a. Eq a => a -> a -> Bool
== forall a. Ord a => a -> a -> a
min (Int
off forall a. Num a => a -> a -> a
+ Int
len) (forall a. Prim a => PrimArray a -> Int
length PrimArray a
arr) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
    step (PrimArray a
arr, Int
i) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
i) (PrimArray a
arr, Int
i forall a. Num a => a -> a -> a
+ Int
1)