{-# LANGUAGE CPP #-}
module Streamly.Internal.Data.Array
(
Array
, A.fromListN
, A.fromList
, fromStreamN
, fromStream
, A.writeN
, A.writeNAligned
, A.write
, writeLastN
, A.toList
, A.read
, A.readRev
, reader
, readerUnsafe
, A.readerRev
, producer
, getIndex
, A.unsafeIndex
, getIndexRev
, last
, getIndices
, getIndicesFromThenTo
, length
, null
, binarySearch
, findIndicesOf
, cast
, asBytes
, castUnsafe
, asPtrUnsafe
, asCStringUnsafe
, A.unsafeFreeze
, A.unsafeThaw
, getSliceUnsafe
, genSlicesFromLen
, getSlicesFromLen
, splitOn
, streamTransform
, streamFold
, fold
, A.toStream
, A.toStreamRev
)
where
#include "inline.hs"
#include "ArrayMacros.h"
import Control.Exception (assert)
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Functor.Identity (Identity)
import Data.Proxy (Proxy(..))
import Data.Word (Word8)
import Foreign.C.String (CString)
import Foreign.Ptr (castPtr)
import Foreign.Storable (Storable)
import Streamly.Internal.Data.Unboxed
( Unbox
, peekWith
, sizeOf
)
import Prelude hiding (length, null, last, map, (!!), read, concat)
import Streamly.Internal.Data.Array.Mut.Type (ArrayUnsafe(..))
import Streamly.Internal.Data.Array.Type
(Array(..), length, asPtrUnsafe)
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Producer.Type (Producer(..))
import Streamly.Internal.Data.Stream.StreamD (Stream)
import Streamly.Internal.Data.Tuple.Strict (Tuple3Fused'(..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.System.IO (unsafeInlineIO)
import qualified Streamly.Internal.Data.Array.Mut.Type as MA
import qualified Streamly.Internal.Data.Array.Mut as MA
import qualified Streamly.Internal.Data.Array.Type as A
import qualified Streamly.Internal.Data.Fold as FL
import qualified Streamly.Internal.Data.Producer.Type as Producer
import qualified Streamly.Internal.Data.Producer as Producer
import qualified Streamly.Internal.Data.Ring.Unboxed as RB
import qualified Streamly.Internal.Data.Stream.StreamD as D
import qualified Streamly.Internal.Data.Stream.StreamD as Stream
import qualified Streamly.Internal.Data.Unfold as Unfold
#include "DocTestDataArray.hs"
{-# INLINE fromStreamN #-}
fromStreamN :: (MonadIO m, Unbox a) => Int -> Stream m a -> m (Array a)
fromStreamN :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Stream m a -> m (Array a)
fromStreamN Int
n Stream 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, Unbox a) =>
Int -> Stream m a -> m (Array a)
A.fromStreamDN Int
n Stream m a
m
{-# INLINE fromStream #-}
fromStream :: (MonadIO m, Unbox a) => Stream m a -> m (Array a)
fromStream :: forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Stream m a -> m (Array a)
fromStream = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
Stream.fold forall (m :: * -> *) a. (MonadIO m, Unbox a) => Fold m a (Array a)
A.write
{-# INLINE_NORMAL producer #-}
producer :: forall m a. (Monad m, Unbox a) => Producer m (Array a) a
producer :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Producer m (Array a) a
producer =
forall (m :: * -> *) a c b.
Functor m =>
(a -> c) -> (c -> a) -> Producer m c b -> Producer m a b
Producer.translate forall a. Array a -> MutArray a
A.unsafeThaw forall a. MutArray a -> Array a
A.unsafeFreeze
forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Producer m (MutArray a) a
MA.producerWith (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IO a -> a
unsafeInlineIO)
{-# INLINE_NORMAL reader #-}
reader :: forall m a. (Monad m, Unbox a) => Unfold m (Array a) a
reader :: forall (m :: * -> *) a. (Monad m, Unbox a) => Unfold m (Array a) a
reader = forall (m :: * -> *) a b. Producer m a b -> Unfold m a b
Producer.simplify forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Producer m (Array a) a
producer
{-# INLINE_NORMAL readerUnsafe #-}
readerUnsafe :: forall m a. (Monad m, Unbox a) => Unfold m (Array a) a
readerUnsafe :: forall (m :: * -> *) a. (Monad m, Unbox a) => Unfold m (Array a) a
readerUnsafe = 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, Unbox a) =>
ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step forall {m :: * -> *} {a} {a}.
Monad m =>
Array a -> m (ArrayUnsafe a)
inject
where
inject :: Array a -> m (ArrayUnsafe a)
inject (Array MutableByteArray
contents Int
start Int
end) =
forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. MutableByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutableByteArray
contents Int
end Int
start)
{-# INLINE_LATE step #-}
step :: ArrayUnsafe a -> m (Step (ArrayUnsafe a) a)
step (ArrayUnsafe MutableByteArray
contents Int
end Int
p) = do
let !x :: a
x = forall a. IO a -> a
unsafeInlineIO forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith MutableByteArray
contents Int
p
let !p1 :: Int
p1 = INDEX_NEXT(p,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. MutableByteArray -> Int -> Int -> ArrayUnsafe a
ArrayUnsafe MutableByteArray
contents Int
end Int
p1)
{-# INLINE null #-}
null :: Array a -> Bool
null :: forall a. Array a -> Bool
null Array a
arr = forall a. Array a -> Int
A.byteLength Array a
arr forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE getIndexRev #-}
getIndexRev :: forall a. Unbox a => Int -> Array a -> Maybe a
getIndexRev :: forall a. Unbox a => Int -> Array a -> Maybe a
getIndexRev Int
i Array a
arr =
forall a. IO a -> a
unsafeInlineIO
forall a b. (a -> b) -> a -> b
$ do
let elemPtr :: Int
elemPtr = RINDEX_OF(arrEnd arr, i, a)
if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
elemPtr forall a. Ord a => a -> a -> Bool
>= forall a. Array a -> Int
arrStart Array a
arr
then forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith (forall a. Array a -> MutableByteArray
arrContents Array a
arr) Int
elemPtr
else forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
{-# INLINE last #-}
last :: Unbox a => Array a -> Maybe a
last :: forall a. Unbox a => Array a -> Maybe a
last = forall a. Unbox a => Int -> Array a -> Maybe a
getIndexRev Int
0
{-# INLINE writeLastN #-}
writeLastN ::
(Storable a, Unbox a, MonadIO m) => Int -> Fold m a (Array a)
writeLastN :: forall a (m :: * -> *).
(Storable a, Unbox a, MonadIO m) =>
Int -> Fold m a (Array a)
writeLastN Int
n
| Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. a -> b -> a
const forall a. Monoid a => a
mempty) forall (m :: * -> *) a. Monad m => Fold m a ()
FL.drain
| Bool
otherwise = forall a. MutArray a -> Array a
A.unsafeFreeze forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold forall {m :: * -> *} {a} {c} {b}.
(MonadIO m, Storable a, Num c) =>
Tuple3Fused' (Ring a) (Ptr a) c
-> a -> m (Step (Tuple3Fused' (Ring a) (Ptr a) c) b)
step forall {b}. m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
initial forall {m :: * -> *} {a}.
(MonadIO m, Unbox a, Storable a) =>
Tuple3Fused' (Ring a) (Ptr a) Int -> m (MutArray a)
done
where
step :: Tuple3Fused' (Ring a) (Ptr a) c
-> a -> m (Step (Tuple3Fused' (Ring a) (Ptr a) c) b)
step (Tuple3Fused' Ring a
rb Ptr a
rh c
i) a
a = do
Ptr a
rh1 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
a
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
FL.Partial forall a b. (a -> b) -> a -> b
$ forall a b c. a -> b -> c -> Tuple3Fused' a b c
Tuple3Fused' Ring a
rb Ptr a
rh1 (c
i forall a. Num a => a -> a -> a
+ c
1)
initial :: m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
initial =
let f :: (a, b) -> Step (Tuple3Fused' a b Int) b
f (a
a, b
b) = forall s b. s -> Step s b
FL.Partial forall a b. (a -> b) -> a -> b
$ forall a b c. a -> b -> c -> Tuple3Fused' a b c
Tuple3Fused' a
a b
b (Int
0 :: Int)
in forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {a} {b} {b}. (a, b) -> Step (Tuple3Fused' a b Int) b
f forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Int -> IO (Ring a, Ptr a)
RB.new Int
n
done :: Tuple3Fused' (Ring a) (Ptr a) Int -> m (MutArray a)
done (Tuple3Fused' Ring a
rb Ptr a
rh Int
i) = do
MutArray a
arr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> m (MutArray a)
MA.newPinned Int
n
forall {m :: * -> *} {a} {b}.
(MonadIO m, Storable a) =>
Int -> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
foldFunc Int
i Ptr a
rh forall {m :: * -> *} {a}.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
snoc' MutArray a
arr Ring a
rb
snoc' :: MutArray a -> a -> m (MutArray a)
snoc' MutArray a
b a
a = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall {m :: * -> *} {a}.
(MonadIO m, Unbox a) =>
MutArray a -> a -> m (MutArray a)
MA.snocUnsafe MutArray a
b a
a
foldFunc :: Int -> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
foldFunc Int
i
| Int
i forall a. Ord a => a -> a -> Bool
< Int
n = forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingM
| Bool
otherwise = forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingFullM
{-# INLINE binarySearch #-}
binarySearch :: a -> Array a -> Maybe Int
binarySearch :: forall a. a -> Array a -> Maybe Int
binarySearch = forall a. HasCallStack => a
undefined
findIndicesOf :: (a -> Bool) -> Unfold Identity (Array a) Int
findIndicesOf :: forall a. (a -> Bool) -> Unfold Identity (Array a) Int
findIndicesOf = forall a. HasCallStack => a
undefined
{-# INLINE getSliceUnsafe #-}
getSliceUnsafe ::
forall a. Unbox a
=> Int
-> Int
-> Array a
-> Array a
getSliceUnsafe :: forall a. Unbox a => Int -> Int -> Array a -> Array a
getSliceUnsafe Int
index Int
len (Array MutableByteArray
contents Int
start Int
e) =
let size :: Int
size = SIZE_OF(a)
start1 :: Int
start1 = Int
start forall a. Num a => a -> a -> a
+ (Int
index forall a. Num a => a -> a -> a
* Int
size)
end1 :: Int
end1 = Int
start1 forall a. Num a => a -> a -> a
+ (Int
len forall a. Num a => a -> a -> a
* Int
size)
in forall a. HasCallStack => Bool -> a -> a
assert (Int
end1 forall a. Ord a => a -> a -> Bool
<= Int
e) (forall a. MutableByteArray -> Int -> Int -> Array a
Array MutableByteArray
contents Int
start1 Int
end1)
{-# INLINE splitOn #-}
splitOn :: (Monad m, Unbox a) =>
(a -> Bool) -> Array a -> Stream m (Array a)
splitOn :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(a -> Bool) -> Array a -> Stream m (Array a)
splitOn a -> Bool
predicate Array a
arr =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Int
i, Int
len) -> forall a. Unbox a => Int -> Int -> Array a -> Array a
getSliceUnsafe Int
i Int
len Array a
arr)
forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m (Int, Int)
D.sliceOnSuffix a -> Bool
predicate (forall (m :: * -> *) a. (Monad m, Unbox a) => Array a -> Stream m a
A.toStreamD Array a
arr)
{-# INLINE genSlicesFromLen #-}
genSlicesFromLen :: forall m a. (Monad m, Unbox a)
=> Int
-> Int
-> Unfold m (Array a) (Int, Int)
genSlicesFromLen :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> Int -> Unfold m (Array a) (Int, Int)
genSlicesFromLen Int
from Int
len =
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
Unfold.lmap forall a. Array a -> MutArray a
A.unsafeThaw (forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> Int -> Unfold m (MutArray a) (Int, Int)
MA.genSlicesFromLen Int
from Int
len)
{-# INLINE getSlicesFromLen #-}
getSlicesFromLen :: forall m a. (Monad m, Unbox a)
=> Int
-> Int
-> Unfold m (Array a) (Array a)
getSlicesFromLen :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> Int -> Unfold m (Array a) (Array a)
getSlicesFromLen Int
from Int
len =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. MutArray a -> Array a
A.unsafeFreeze
forall a b. (a -> b) -> a -> b
$ forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
Unfold.lmap forall a. Array a -> MutArray a
A.unsafeThaw (forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> Int -> Unfold m (MutArray a) (MutArray a)
MA.getSlicesFromLen Int
from Int
len)
{-# INLINE getIndex #-}
getIndex :: forall a. Unbox a => Int -> Array a -> Maybe a
getIndex :: forall a. Unbox a => Int -> Array a -> Maybe a
getIndex Int
i Array a
arr =
forall a. IO a -> a
unsafeInlineIO
forall a b. (a -> b) -> a -> b
$ do
let elemPtr :: Int
elemPtr = forall a. Array a -> Int
INDEX_OF(arrStart arr, i, a)
if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& INDEX_VALID(elemPtr, arrEnd arr, a)
then forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Unbox a => MutableByteArray -> Int -> IO a
peekWith (forall a. Array a -> MutableByteArray
arrContents Array a
arr) Int
elemPtr
else forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
{-# INLINE getIndices #-}
getIndices :: (Monad m, Unbox a) => Stream m Int -> Unfold m (Array a) a
getIndices :: forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Stream m Int -> Unfold m (Array a) a
getIndices Stream m Int
m =
let unf :: Unfold m (MutArray a) a
unf = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
(forall b. IO b -> m b) -> Stream m Int -> Unfold m (MutArray a) a
MA.getIndicesD (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IO a -> a
unsafeInlineIO) Stream m Int
m
in forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
Unfold.lmap forall a. Array a -> MutArray a
A.unsafeThaw Unfold m (MutArray a) a
unf
{-# INLINE getIndicesFromThenTo #-}
getIndicesFromThenTo :: Unfold m (Int, Int, Int, Array a) a
getIndicesFromThenTo :: forall (m :: * -> *) a. Unfold m (Int, Int, Int, Array a) a
getIndicesFromThenTo = forall a. HasCallStack => a
undefined
{-# INLINE streamTransform #-}
streamTransform :: forall m a b. (MonadIO m, Unbox a, Unbox b)
=> (Stream m a -> Stream m b) -> Array a -> m (Array b)
streamTransform :: forall (m :: * -> *) a b.
(MonadIO m, Unbox a, Unbox b) =>
(Stream m a -> Stream m b) -> Array a -> m (Array b)
streamTransform Stream m a -> Stream m b
f Array a
arr =
forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
Stream.fold (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (Array a)
A.writeWith (forall a. Unbox a => Array a -> Int
length Array a
arr)) forall a b. (a -> b) -> a -> b
$ Stream m a -> Stream m b
f (forall (m :: * -> *) a. (Monad m, Unbox a) => Array a -> Stream m a
A.read Array a
arr)
castUnsafe ::
#ifdef DEVBUILD
Unbox b =>
#endif
Array a -> Array b
castUnsafe :: forall a b. Array a -> Array b
castUnsafe (Array MutableByteArray
contents Int
start Int
end) =
forall a. MutableByteArray -> Int -> Int -> Array a
Array MutableByteArray
contents Int
start Int
end
asBytes :: Array a -> Array Word8
asBytes :: forall a. Array a -> Array Word8
asBytes = forall a b. Array a -> Array b
castUnsafe
cast :: forall a b. (Unbox b) => Array a -> Maybe (Array b)
cast :: forall a b. Unbox b => Array a -> Maybe (Array b)
cast Array a
arr =
let len :: Int
len = forall a. Array a -> Int
A.byteLength Array a
arr
r :: Int
r = Int
len forall a. Integral a => a -> a -> a
`mod` SIZE_OF(b)
in if Int
r forall a. Eq a => a -> a -> Bool
/= Int
0
then forall a. Maybe a
Nothing
else forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. Array a -> Array b
castUnsafe Array a
arr
asCStringUnsafe :: Array a -> (CString -> IO b) -> IO b
asCStringUnsafe :: forall a b. Array a -> (CString -> IO b) -> IO b
asCStringUnsafe Array a
arr CString -> IO b
act = do
let arr1 :: Array Word8
arr1 = forall a. Array a -> Array Word8
asBytes Array a
arr forall a. Semigroup a => a -> a -> a
<> forall a. Unbox a => [a] -> Array a
A.fromList [Word8
0]
forall (m :: * -> *) a b.
MonadIO m =>
Array a -> (Ptr a -> m b) -> m b
asPtrUnsafe Array Word8
arr1 forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> CString -> IO b
act (forall a b. Ptr a -> Ptr b
castPtr Ptr Word8
ptr)
{-# INLINE fold #-}
fold :: forall m a b. (Monad m, Unbox a) => Fold m a b -> Array a -> m b
fold :: forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
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
Stream.fold Fold m a b
f (forall (m :: * -> *) a. (Monad m, Unbox a) => Array a -> Stream m a
A.read Array a
arr)
{-# INLINE streamFold #-}
streamFold :: (Monad m, Unbox a) => (Stream m a -> m b) -> Array a -> m b
streamFold :: forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Stream m a -> m b) -> Array a -> m b
streamFold Stream m a -> m b
f Array a
arr = Stream m a -> m b
f (forall (m :: * -> *) a. (Monad m, Unbox a) => Array a -> Stream m a
A.read Array a
arr)