{-# 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 :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
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]
"writeN: negative write count specified"
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> m (Array a)
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, Storable a) => SerialT m a -> m (Array a)
fromStream :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
SerialT m a -> m (Array a)
fromStream = forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
(Monad m, IsStream t) =>
Fold m a b -> t m a -> m b
P.runFold 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 :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t, Storable a) =>
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, 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 :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, IsStream t, Storable a) =>
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, 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 :: forall (m :: * -> *) a.
(Monad m, Storable a) =>
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, Storable a) =>
ReadUState a -> m (Step (ReadUState a) a)
step 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
_) =
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState (forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
end ForeignPtrContents
contents) (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 forall a. Eq a => a -> a -> Bool
== forall a. Addr# -> Ptr a
Ptr Addr#
end =
let x :: ()
x = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
fp
in ()
x seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
step (ReadUState ForeignPtr a
fp Ptr a
p) = do
let !x :: a
x = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
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 a
x
(forall a. ForeignPtr a -> Ptr a -> ReadUState a
ReadUState ForeignPtr a
fp (Ptr a
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a)))
{-# INLINE_NORMAL unsafeRead #-}
unsafeRead :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a
unsafeRead :: forall (m :: * -> *) a.
(Monad m, Storable a) =>
Unfold m (Array a) a
unsafeRead = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {m :: * -> *} {a} {a} {a}.
(Monad m, Storable a) =>
ForeignPtr a -> m (Step (ForeignPtr a) a)
step 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
_) = 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 = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ do
a
r <- forall a. Storable a => Ptr a -> IO a
peek (forall a. Addr# -> Ptr a
Ptr Addr#
p)
forall {a}. a -> IO ()
touch ForeignPtrContents
contents
forall (m :: * -> *) a. Monad m => a -> m a
return a
r
let !(Ptr Addr#
p1) = forall a. Addr# -> Ptr a
Ptr Addr#
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a)
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 a
x (forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
p1 ForeignPtrContents
contents)
touch :: a -> IO ()
touch a
r = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case touch# :: 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 :: forall a. Storable a => Array a -> Bool
null Array a
arr = forall a. Storable a => Array a -> Int
length Array a
arr forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE last #-}
last :: Storable a => Array a -> Maybe a
last :: forall a. Storable a => Array a -> Maybe a
last Array a
arr = forall a. Storable a => Array a -> Int -> Maybe a
readIndex Array a
arr (forall a. Storable a => Array a -> Int
length Array a
arr forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE readIndex #-}
readIndex :: Storable a => Array a -> Int -> Maybe a
readIndex :: forall a. Storable a => Array a -> Int -> Maybe a
readIndex Array a
arr Int
i =
if Int
i forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i forall a. Ord a => a -> a -> Bool
> forall a. Storable a => Array a -> Int
length Array a
arr forall a. Num a => a -> a -> a
- Int
1
then forall a. Maybe a
Nothing
else forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (forall a. Array a -> ForeignPtr a
aStart Array a
arr) forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Int -> a -> m ()
writeIndex Array a
arr Int
i a
a = do
let maxIndex :: Int
maxIndex = forall a. Storable a => Array a -> Int
length Array a
arr forall a. Num a => a -> a -> a
- Int
1
if Int
i forall a. Ord a => a -> a -> Bool
< Int
0
then forall a. HasCallStack => [Char] -> a
error [Char]
"writeIndex: negative array index"
else if Int
i forall a. Ord a => a -> a -> Bool
> Int
maxIndex
then forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"writeIndex: specified array index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
i
forall a. [a] -> [a] -> [a]
++ [Char]
" is beyond the maximum index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
maxIndex
else
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (forall a. Array a -> ForeignPtr a
aStart Array a
arr) forall a b. (a -> b) -> a -> b
$ \Ptr a
p ->
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 :: 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
f Array a
arr =
forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
(Monad m, IsStream t) =>
Fold m a b -> t m a -> m b
P.runFold (forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Int -> Fold m a (Array a)
A.toArrayMinChunk (forall a. Storable a => a -> Int
alignment (forall a. HasCallStack => a
undefined :: a)) (forall a. Storable a => Array a -> Int
length Array a
arr))
forall a b. (a -> b) -> a -> b
$ SerialT m a -> SerialT m b
f (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 :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Fold m a b -> Array a -> m b
fold Fold m a b
f Array a
arr = 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 (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 :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
(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, Storable a) =>
Array a -> t m a
toStream Array a
arr)