{-# 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)