{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE UnboxedTuples       #-}
{-# LANGUAGE ScopedTypeVariables #-}
#include "inline.hs"
module Streamly.Internal.Memory.Array
    (
      Array
    
    
    
    , A.fromListN
    , A.fromList
    
    , fromStreamN
    , fromStream
    
    
    , A.writeN      
    , A.writeNAligned
    , A.write       
    
    
    , A.toList
    , toStream
    , toStreamRev
    , read
    , unsafeRead
    
    
    , length
    , null
    , last
    
    , readIndex
    , A.unsafeIndex
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    , writeIndex
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    , streamTransform
    
    , streamFold
    , fold
    
    , D.lastN
    )
where
import  Control.Monad (when)
import Control.Monad.IO.Class (MonadIO(..))
import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr)
import Foreign.Ptr (plusPtr)
import Foreign.Storable (Storable(..))
import Prelude hiding (length, null, last, map, (!!), read, concat)
import GHC.ForeignPtr (ForeignPtr(..))
import GHC.Ptr (Ptr(..))
import GHC.Prim (touch#)
import GHC.IO (IO(..))
import Streamly.Internal.Data.Fold.Types (Fold(..))
import Streamly.Internal.Data.Unfold.Types (Unfold(..))
import Streamly.Internal.Memory.Array.Types (Array(..), length)
import Streamly.Internal.Data.Stream.Serial (SerialT)
import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
import qualified Streamly.Internal.Memory.Array.Types as A
import qualified Streamly.Internal.Data.Stream.Prelude as P
import qualified Streamly.Internal.Data.Stream.Serial as Serial
import qualified Streamly.Internal.Data.Stream.StreamD as D
import qualified Streamly.Internal.Data.Stream.StreamK as K
{-# INLINE fromStreamN #-}
fromStreamN :: (MonadIO m, Storable a) => Int -> SerialT m a -> m (Array a)
fromStreamN :: Int -> SerialT m a -> m (Array a)
fromStreamN Int
n SerialT m a
m = do
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"writeN: negative write count specified"
    Int -> Stream m a -> m (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> m (Array a)
A.fromStreamDN Int
n (Stream m a -> m (Array a)) -> Stream m a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ SerialT m a -> Stream m a
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, Storable a) => SerialT m a -> m (Array a)
fromStream :: SerialT m a -> m (Array a)
fromStream = Fold m a (Array a) -> SerialT m a -> m (Array a)
forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
(Monad m, IsStream t) =>
Fold m a b -> t m a -> m b
P.runFold Fold m a (Array a)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Fold m a (Array a)
A.write
{-# INLINE_EARLY toStream #-}
toStream :: (Monad m, K.IsStream t, Storable a) => Array a -> t m a
toStream :: Array a -> t m a
toStream = Stream m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
D.fromStreamD (Stream m a -> t m a)
-> (Array a -> Stream m a) -> Array a -> t m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
A.toStreamD
{-# INLINE_EARLY toStreamRev #-}
toStreamRev :: (Monad m, IsStream t, Storable a) => Array a -> t m a
toStreamRev :: Array a -> t m a
toStreamRev = Stream m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
D.fromStreamD (Stream m a -> t m a)
-> (Array a -> Stream m a) -> Array a -> t m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Stream m a
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> Stream m a
A.toStreamDRev
data ReadUState a = ReadUState
    {-# UNPACK #-} !(ForeignPtr a)  
    {-# UNPACK #-} !(Ptr a)         
{-# INLINE_NORMAL read #-}
read :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a
read :: Unfold m (Array a) a
read = (ReadUState a -> m (Step (ReadUState a) a))
-> (Array a -> m (ReadUState a)) -> Unfold m (Array a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ReadUState a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a.
(Monad m, Storable a) =>
ReadUState a -> m (Step (ReadUState a) a)
step Array a -> m (ReadUState a)
forall (m :: * -> *) a a. Monad m => Array a -> m (ReadUState a)
inject
    where
    inject :: Array a -> m (ReadUState a)
inject (Array (ForeignPtr Addr#
start ForeignPtrContents
contents) (Ptr Addr#
end) Ptr a
_) =
        ReadUState a -> m (ReadUState a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ReadUState a -> m (ReadUState a))
-> ReadUState a -> m (ReadUState a)
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> Ptr a -> ReadUState a
forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState (Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
end ForeignPtrContents
contents) (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
start)
    {-# INLINE_LATE step #-}
    step :: ReadUState a -> m (Step (ReadUState a) a)
step (ReadUState fp :: ForeignPtr a
fp@(ForeignPtr Addr#
end ForeignPtrContents
_) Ptr a
p) | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
end =
        let x :: ()
x = IO () -> ()
forall a. IO a -> a
A.unsafeInlineIO (IO () -> ()) -> IO () -> ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
fp
        in ()
x () -> m (Step (ReadUState a) a) -> m (Step (ReadUState a) a)
`seq` Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ReadUState a) a
forall s a. Step s a
D.Stop
    step (ReadUState ForeignPtr a
fp Ptr a
p) = do
            
            
            
            
            
            
            let !x :: a
x = IO a -> a
forall a. IO a -> a
A.unsafeInlineIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
            Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ReadUState a) a -> m (Step (ReadUState a) a))
-> Step (ReadUState a) a -> m (Step (ReadUState a) a)
forall a b. (a -> b) -> a -> b
$ a -> ReadUState a -> Step (ReadUState a) a
forall s a. a -> s -> Step s a
D.Yield a
x
                (ForeignPtr a -> Ptr a -> ReadUState a
forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState ForeignPtr a
fp (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)))
{-# INLINE_NORMAL unsafeRead #-}
unsafeRead :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a
unsafeRead :: Unfold m (Array a) a
unsafeRead = (ForeignPtr a -> m (Step (ForeignPtr a) a))
-> (Array a -> m (ForeignPtr a)) -> Unfold m (Array a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ForeignPtr a -> m (Step (ForeignPtr a) a)
forall (m :: * -> *) a a a.
(Monad m, Storable a) =>
ForeignPtr a -> m (Step (ForeignPtr a) a)
step Array a -> m (ForeignPtr a)
forall (m :: * -> *) a. Monad m => Array a -> m (ForeignPtr a)
inject
    where
    inject :: Array a -> m (ForeignPtr a)
inject (Array ForeignPtr a
fp Ptr a
_ Ptr a
_) = ForeignPtr a -> m (ForeignPtr a)
forall (m :: * -> *) a. Monad m => a -> m a
return ForeignPtr a
fp
    {-# INLINE_LATE step #-}
    step :: ForeignPtr a -> m (Step (ForeignPtr a) a)
step (ForeignPtr Addr#
p ForeignPtrContents
contents) = do
            
            
            
            
            
            
            let !x :: a
x = IO a -> a
forall a. IO a -> a
A.unsafeInlineIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ do
                        a
r <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
p)
                        ForeignPtrContents -> IO ()
forall a. a -> IO ()
touch ForeignPtrContents
contents
                        a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
            let !(Ptr Addr#
p1) = Addr# -> Ptr Any
forall a. Addr# -> Ptr a
Ptr Addr#
p Ptr Any -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
            Step (ForeignPtr a) a -> m (Step (ForeignPtr a) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ForeignPtr a) a -> m (Step (ForeignPtr a) a))
-> Step (ForeignPtr a) a -> m (Step (ForeignPtr a) a)
forall a b. (a -> b) -> a -> b
$ a -> ForeignPtr a -> Step (ForeignPtr a) a
forall s a. a -> s -> Step s a
D.Yield a
x (Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
p1 ForeignPtrContents
contents)
    touch :: a -> IO ()
touch a
r = (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case a -> State# RealWorld -> State# RealWorld
forall a. a -> State# RealWorld -> State# RealWorld
touch# a
r State# RealWorld
s of State# RealWorld
s' -> (# State# RealWorld
s', () #)
{-# INLINE null #-}
null :: Storable a => Array a -> Bool
null :: Array a -> Bool
null Array a
arr = Array a -> Int
forall a. Storable a => Array a -> Int
length Array a
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE last #-}
last :: Storable a => Array a -> Maybe a
last :: Array a -> Maybe a
last Array a
arr = Array a -> Int -> Maybe a
forall a. Storable a => Array a -> Int -> Maybe a
readIndex Array a
arr (Array a -> Int
forall a. Storable a => Array a -> Int
length Array a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE readIndex #-}
readIndex :: Storable a => Array a -> Int -> Maybe a
readIndex :: Array a -> Int -> Maybe a
readIndex Array a
arr Int
i =
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Array a -> Int
forall a. Storable a => Array a -> Int
length Array a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        then Maybe a
forall a. Maybe a
Nothing
        else IO (Maybe a) -> Maybe a
forall a. IO a -> a
A.unsafeInlineIO (IO (Maybe a) -> Maybe a) -> IO (Maybe a) -> Maybe a
forall a b. (a -> b) -> a -> b
$
             ForeignPtr a -> (Ptr a -> IO (Maybe a)) -> IO (Maybe a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
arr) ((Ptr a -> IO (Maybe a)) -> IO (Maybe a))
-> (Ptr a -> IO (Maybe a)) -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> IO a -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> Int -> IO a
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
i
{-# INLINE writeIndex #-}
writeIndex :: (MonadIO m, Storable a) => Array a -> Int -> a -> m ()
writeIndex :: Array a -> Int -> a -> m ()
writeIndex Array a
arr Int
i a
a = do
    let maxIndex :: Int
maxIndex = Array a -> Int
forall a. Storable a => Array a -> Int
length Array a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
    then [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"writeIndex: negative array index"
    else if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxIndex
         then [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"writeIndex: specified array index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i
                    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is beyond the maximum index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
maxIndex
         else
            IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> (Ptr a -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (Array a -> ForeignPtr a
forall a. Array a -> ForeignPtr a
aStart Array a
arr) ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
p ->
                Ptr a -> Int -> a -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr a
p Int
i a
a
{-# INLINE streamTransform #-}
streamTransform :: forall m a b. (MonadIO m, Storable a, Storable b)
    => (SerialT m a -> SerialT m b) -> Array a -> m (Array b)
streamTransform :: (SerialT m a -> SerialT m b) -> Array a -> m (Array b)
streamTransform SerialT m a -> SerialT m b
f Array a
arr =
    Fold m b (Array b) -> SerialT m b -> m (Array b)
forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
(Monad m, IsStream t) =>
Fold m a b -> t m a -> m b
P.runFold (Int -> Int -> Fold m b (Array b)
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
A.toArrayMinChunk (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a)) (Array a -> Int
forall a. Storable a => Array a -> Int
length Array a
arr))
        (SerialT m b -> m (Array b)) -> SerialT m b -> m (Array b)
forall a b. (a -> b) -> a -> b
$ SerialT m a -> SerialT m b
f (Array a -> SerialT m a
forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t, Storable a) =>
Array a -> t m a
toStream Array a
arr)
{-# INLINE fold #-}
fold :: forall m a b. (MonadIO m, Storable a) => Fold m a b -> Array a -> m b
fold :: Fold m a b -> Array a -> m b
fold Fold m a b
f Array a
arr = Fold m a b -> SerialT m a -> m b
forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
(Monad m, IsStream t) =>
Fold m a b -> t m a -> m b
P.runFold Fold m a b
f (Array a -> SerialT m a
forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t, Storable a) =>
Array a -> t m a
toStream Array a
arr :: Serial.SerialT m a)
{-# INLINE streamFold #-}
streamFold :: (MonadIO m, Storable a) => (SerialT m a -> m b) -> Array a -> m b
streamFold :: (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 (Array a -> SerialT m a
forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t, Storable a) =>
Array a -> t m a
toStream Array a
arr)