@@ -1,706 +0,0 @@
1−-- Copyright : (c) 2020 Composewell Technologies
2−-- License : BSD-3-Clause
3−-- Maintainer : streamly@composewell.com
4−-- Stability : experimental
7−import Control.DeepSeq (NFData(..))
8−import Control.Monad (when)
9−import Control.Monad.IO.Class (MonadIO(..))
10−import Control.Monad.Primitive (PrimMonad(..), primitive_)
11−import Data.Bifunctor (first)
12−import Data.Primitive.Types (Prim(..), sizeOf)
13−#if __GLASGOW_HASKELL__ < 808
14−import Data.Semigroup (Semigroup(..))
16−import Data.Word (Word8)
17−import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
18−import Streamly.Internal.Data.Maybe.Strict (Maybe'(..))
19−import Streamly.Internal.Data.Fold.Type (Fold(..))
20−import Streamly.Internal.Data.SVar (adaptState)
21−import System.IO.Unsafe (unsafePerformIO)
22−import Text.Read (readPrec, readListPrec, readListPrecDefault)
24−import qualified GHC.Exts as Exts
25−import qualified Prelude as P
26−import qualified Streamly.Internal.Data.Fold as FL
27−import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
28−import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
30−import GHC.Exts hiding (fromListN, fromList, toList)
31−import Prelude hiding (length, unlines, foldr)
33−-------------------------------------------------------------------------------
35−-------------------------------------------------------------------------------
37−data Array a = Array ByteArray# Int Int
39−-------------------------------------------------------------------------------
41−-------------------------------------------------------------------------------
43−-- | Both arrays must fully contain the specified ranges, but this is not
44−-- checked. The two arrays must not be the same array in different states, but
45−-- this is not checked either.
46−{-# INLINE unsafeCopy #-}
48− forall m a. (MonadIO m, Prim a)
49− => MA.Array a -- ^ destination array
50− -> Int -- ^ offset into destination array
51− -> Array a -- ^ source array
52− -> Int -- ^ offset into source array
53− -> Int -- ^ number of elements to copy
55−unsafeCopy (MA.Array dst#) (I# doff#) (Array src# (I# off#) _) (I# soff#) (I# n#) =
57− let toBytes cnt# = cnt# *# (sizeOf# (undefined :: a))
61− (toBytes (off# +# soff#))
66−-------------------------------------------------------------------------------
67−-- Basic Byte Array Operations
68−-------------------------------------------------------------------------------
70−{-# INLINE unsafeFreeze #-}
71−unsafeFreeze :: (Prim a, MonadIO m) => MA.Array a -> m (Array a)
72−unsafeFreeze marr@(MA.Array arr#) =
74− len <- MA.length marr
76− case unsafeFreezeByteArray# arr# s# of
77− (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
79−{-# INLINE unsafeFreezeWithShrink #-}
80−unsafeFreezeWithShrink ::
81− (Prim a, MonadIO m) => MA.Array a -> Int -> m (Array a)
82−unsafeFreezeWithShrink marr@(MA.Array arr#) n =
84− MA.shrinkArray marr n
85− len <- MA.length marr
87− case unsafeFreezeByteArray# arr# s# of
88− (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
91−-- Should never be used in general
92−{-# INLINE unsafeThaw #-}
93−unsafeThaw :: MonadIO m => Array a -> m (MA.Array a)
94−unsafeThaw (Array arr#) =
95− primitive $ \s# -> (# s#, MA.Array (unsafeCoerce# arr#) #)
98−-- Unsafe because the index bounds are not checked
99−{-# INLINE unsafeIndex #-}
100−unsafeIndex :: Prim a => Array a -> Int -> a
101−unsafeIndex (Array arr# (I# off#) _) (I# i#) = indexByteArray# arr# (off# +# i#)
104−sameByteArray :: ByteArray# -> ByteArray# -> Bool
105−sameByteArray ba1 ba2 =
106− case reallyUnsafePtrEquality#
107− (unsafeCoerce# ba1 :: ())
108− (unsafeCoerce# ba2 :: ()) of
111−-------------------------------------------------------------------------------
113−-------------------------------------------------------------------------------
115−-- XXX move this section to mutable array module?
117−mkChunkSizeKB :: Int -> Int
118−mkChunkSizeKB n = n * k
121−-- | Default maximum buffer size in bytes, for reading from and writing to IO
122−-- devices, the value is 32KB minus GHC allocation overhead, which is a few
123−-- bytes, so that the actual allocation is 32KB.
124−defaultChunkSize :: Int
125−defaultChunkSize = mkChunkSizeKB 32
127−-------------------------------------------------------------------------------
129−-------------------------------------------------------------------------------
131−sizeOfPrimElem :: forall a . (Prim a) => a -> Int
132−sizeOfPrimElem _ = max 1 $ sizeOf (undefined :: a)
134−-- XXX rename to byteCount?
135−{-# INLINE byteLength #-}
136−byteLength :: forall a. Prim a => Array a -> Int
137−byteLength (Array _ _ len) = len * sizeOfPrimElem (undefined :: a)
139−-- XXX Also, rename to elemCount
140−-- XXX I would prefer length to keep the API consistent
141−-- XXX Also, re-export sizeOf from Primitive
142−{-# INLINE length #-}
143−length :: Array a -> Int
144−length (Array _ _ len) = len
146−-------------------------------------------------------------------------------
148−-------------------------------------------------------------------------------
150−-- | Use a slice of an array as another array. Note that this is unsafe and does
151−-- not check the bounds
152−slice :: Array a -> Int -> Int -> Array a
153−slice (Array arr# off _) off1 len1 = Array arr# (off + off1) len1
155−nil :: Prim a => Array a
157− unsafePerformIO $ do
158− arr <- MA.newArray 0
161−-- | Fold the whole input to a single array.
163−-- /Caution! Do not use this on infinite streams./
166−{-# INLINE_NORMAL write #-}
167−write :: (MonadIO m, Prim a) => Fold m a (Array a)
168−write = FL.rmapM unsafeFreeze MA.write
170−-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
174−{-# INLINE_NORMAL writeN #-}
175−writeN :: (MonadIO m, Prim a) => Int -> Fold m a (Array a)
176−writeN limit = FL.rmapM unsafeFreeze (MA.writeN limit)
178−-- | Like 'writeN' but does not check the array bounds when writing. The fold
179−-- driver must not call the step function more than 'n' times otherwise it will
180−-- corrupt the memory and crash. This function exists mainly because any
181−-- conditional in the step function blocks fusion causing 10x performance
185−{-# INLINE_NORMAL writeNUnsafe #-}
186−writeNUnsafe :: (MonadIO m, Prim a) => Int -> Fold m a (Array a)
187−writeNUnsafe limit = FL.rmapM unsafeFreeze (MA.writeNUnsafe limit)
189−{-# INLINE_NORMAL fromStreamDN #-}
190−fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (Array a)
191−fromStreamDN limit str = MA.fromStreamDN limit str >>= unsafeFreeze
193−{-# INLINE fromStreamD #-}
194−fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (Array a)
195−fromStreamD str = MA.fromStreamD str >>= unsafeFreeze
197−-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
199−{-# INLINE_NORMAL fromStreamDArraysOf #-}
200−fromStreamDArraysOf ::
201− (MonadIO m, Prim a) => Int -> D.Stream m a -> D.Stream m (Array a)
202−fromStreamDArraysOf n str = D.mapM unsafeFreeze (MA.fromStreamDArraysOf n str)
204−-- XXX derive from MA.fromListN?
205−{-# INLINE fromListN #-}
206−fromListN :: Prim a => Int -> [a] -> Array a
207−fromListN len xs = unsafePerformIO $ MA.fromListNM len xs >>= unsafeFreeze
209−-- XXX derive from MA.fromList?
210−{-# INLINE fromList #-}
211−fromList :: Prim a => [a] -> Array a
212−fromList xs = fromListN (P.length xs) xs
214−-------------------------------------------------------------------------------
216−-------------------------------------------------------------------------------
218−-- | Splice two immutable arrays creating a new immutable array.
219−{-# INLINE spliceTwo #-}
220−spliceTwo :: (MonadIO m, Prim a) => Array a -> Array a -> m (Array a)
224− a3 <- MA.newArray (l1 + l2)
225− unsafeCopy a3 0 a1 0 l1
226− unsafeCopy a3 l1 a2 0 l2
227− unsafeFreeze a3 -- Use `unsafeFreezeWith off len`?
229−-------------------------------------------------------------------------------
231−-------------------------------------------------------------------------------
233−{-# INLINE_LATE toListFB #-}
234−toListFB :: Prim a => (a -> b -> b) -> b -> Array a -> b
235−toListFB c n arr = go 0
240− let !x = unsafeIndex arr p
243−-- | Convert an 'Array' into a list.
246−{-# INLINE toList #-}
247−toList :: Prim a => Array a -> [a]
248−toList s = build (\c n -> toListFB c n s)
250−-------------------------------------------------------------------------------
252−-------------------------------------------------------------------------------
254−instance (Eq a, Prim a) => Eq (Array a) where
256− a1@(Array ba1# _ len1) == a2@(Array ba2# _ len2)
257− | sameByteArray ba1# ba2# = True
258− | len1 /= len2 = False
259− | otherwise = loop (len1 - 1)
265− | otherwise = unsafeIndex a1 i == unsafeIndex a2 i && loop (i - 1)
267−-- | Lexicographic ordering. Subject to change between major versions.
268−instance (Ord a, Prim a) => Ord (Array a) where
269− {-# INLINE compare #-}
270− compare a1@(Array ba1# _ len1) a2@(Array ba2# _ len2)
271− | sameByteArray ba1# ba2# = EQ
272− | otherwise = loop 0
280− compare (unsafeIndex a1 i) (unsafeIndex a2 i) <> loop (i + 1)
281− | otherwise = compare len1 len2
283−instance Prim a => Semigroup (Array a) where
284− -- XXX can't we use runST instead of inlineIO?
285− -- XXX I plan to remove MonadIO and replace it with IO
286− a <> b = unsafePerformIO (spliceTwo a b :: IO (Array a))
288−instance Prim a => Monoid (Array a) where
292−instance NFData (Array a) where
297−-- XXX check if this is compatible with Memory.Array?
298−-- XXX It isn't. I might prefer this Show instance though
299−-- XXX Memory.Array: showsPrec _ = shows . toList
300−instance (Show a, Prim a) => Show (Array a) where
302− showParen (p > 10) $
303− showString "fromListN "
308−instance (a ~ Char) => IsString (Array a) where
309− {-# INLINE fromString #-}
310− fromString = fromList
312−-- GHC versions 8.0 and below cannot derive IsList
313−instance Prim a => IsList (Array a) where
314− type (Item (Array a)) = a
316− {-# INLINE fromList #-}
319− {-# INLINE fromListN #-}
320− fromListN = fromListN
322− {-# INLINE toList #-}
325−instance (Prim a, Read a, Show a) => Read (Array a) where
326− {-# INLINE readPrec #-}
327− readPrec = fromList <$> readPrec
328− readListPrec = readListPrecDefault
330−-- XXX these folds can be made common with mutable arrays by defining a
331−-- unsafeIndex in the specific module?
333−-------------------------------------------------------------------------------
335−-------------------------------------------------------------------------------
338−foldr :: Prim a => (a -> b -> b) -> b -> Array a -> b
346− | len > i = f (unsafeIndex arr i) (go (i + 1))
349−-- | Strict right-associated fold over the elements of an 'Array'.
350−{-# INLINE foldr' #-}
351−foldr' :: Prim a => (a -> b -> b) -> b -> Array a -> b
352−foldr' f z0 arr = go (length arr - 1) z0
358− | otherwise = go (i - 1) (f (unsafeIndex arr i) acc)
360−-- | Strict left-associated fold over the elements of an 'Array'.
361−{-# INLINE foldl' #-}
362−foldl' :: Prim a => (b -> a -> b) -> b -> Array a -> b
363−foldl' f z0 arr = go 0 z0
370− | i < len = go (i + 1) (f acc (unsafeIndex arr i))
373−-- | Strict left-associated fold over the elements of an 'Array'.
374−{-# INLINE foldlM' #-}
375−foldlM' :: (Prim a, Monad m) => (b -> a -> m b) -> b -> Array a -> m b
376−foldlM' f z0 arr = go 0 z0
384− acc2 <- f acc1 (unsafeIndex arr i)
386− | otherwise = return acc1
388−-------------------------------------------------------------------------------
389−-- Converting to streams
390−-------------------------------------------------------------------------------
392−{-# INLINE_NORMAL toStreamD #-}
393−toStreamD :: (Prim a, Monad m) => Array a -> D.Stream m a
394−toStreamD arr = D.Stream step 0
398− {-# INLINE_LATE step #-}
400− | i == length arr = return D.Stop
401− step _ i = return $ D.Yield (unsafeIndex arr i) (i + 1)
403−{-# INLINE toStreamK #-}
404−toStreamK :: Prim a => Array a -> K.Stream m a
414− let !x = unsafeIndex arr p
415− in x `K.cons` go (p + 1)
417−{-# INLINE_NORMAL toStreamDRev #-}
418−toStreamDRev :: (Prim a, Monad m) => Array a -> D.Stream m a
419−toStreamDRev arr = D.Stream step (length arr - 1)
423− {-# INLINE_LATE step #-}
425− | i < 0 = return D.Stop
426− step _ i = return $ D.Yield (unsafeIndex arr i) (i - 1)
428−{-# INLINE toStreamKRev #-}
429−toStreamKRev :: Prim a => Array a -> K.Stream m a
430−toStreamKRev arr = go (length arr - 1)
434− go p | p == -1 = K.nil
436− let !x = unsafeIndex arr p
437− in x `K.cons` go (p - 1)
439−-------------------------------------------------------------------------------
441−-------------------------------------------------------------------------------
443−data FlattenState s a =
445− | InnerLoop s !(Array a) !Int !Int
447−{-# INLINE_NORMAL flattenArrays #-}
448−flattenArrays :: (MonadIO m, Prim a) => D.Stream m (Array a) -> D.Stream m a
449−flattenArrays (D.Stream step state) = D.Stream step' (OuterLoop state)
453− {-# INLINE_LATE step' #-}
454− step' gst (OuterLoop st) = do
455− r <- step (adaptState gst) st
458− let len = length arr
460− then D.Skip (OuterLoop s)
461− else D.Skip (InnerLoop s arr len 0)
462− D.Skip s -> D.Skip (OuterLoop s)
465− step' _ (InnerLoop st _ len i) | i == len =
466− return $ D.Skip $ OuterLoop st
468− step' _ (InnerLoop st arr len i) = do
469− let x = unsafeIndex arr i
470− return $ D.Yield x (InnerLoop st arr len (i + 1))
472−{-# INLINE_NORMAL flattenArraysRev #-}
473−flattenArraysRev :: (MonadIO m, Prim a) => D.Stream m (Array a) -> D.Stream m a
474−flattenArraysRev (D.Stream step state) = D.Stream step' (OuterLoop state)
478− {-# INLINE_LATE step' #-}
479− step' gst (OuterLoop st) = do
480− r <- step (adaptState gst) st
483− let len = length arr
485− then D.Skip (OuterLoop s)
486− else D.Skip (InnerLoop s arr len (len - 1))
487− D.Skip s -> D.Skip (OuterLoop s)
490− step' _ (InnerLoop st _ _ i) | i == -1 =
491− return $ D.Skip $ OuterLoop st
493− step' _ (InnerLoop st arr len i) = do
494− let x = unsafeIndex arr i
495− return $ D.Yield x (InnerLoop st arr len (i - 1))
497−{-# INLINE_NORMAL unlines #-}
498−unlines :: (MonadIO m, Prim a) => a -> D.Stream m (Array a) -> D.Stream m a
499−unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
503− {-# INLINE_LATE step' #-}
504− step' gst (OuterLoop st) = do
505− r <- step (adaptState gst) st
508− let len = length arr
509− in D.Skip (InnerLoop s arr len 0)
510− D.Skip s -> D.Skip (OuterLoop s)
513− step' _ (InnerLoop st _ len i) | i == len =
514− return $ D.Yield sep $ OuterLoop st
516− step' _ (InnerLoop st arr len i) = do
517− let x = unsafeIndex arr i
518− return $ D.Yield x (InnerLoop st arr len (i + 1))
520−-- Splice an array into a pre-reserved mutable array. The user must ensure
521−-- that there is enough space in the mutable array.
522−{-# INLINE spliceInto #-}
523−spliceInto :: (MonadIO m, Prim a) => MA.Array a -> Int -> Array a -> m Int
524−spliceInto dst doff src@(Array _ _ len) = do
525− unsafeCopy dst doff src 0 len
528−data SpliceState s arr1 arr2
530− | SpliceBuffering s arr2
531− | SpliceYielding arr1 (SpliceState s arr1 arr2)
534−-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
535−-- maximum specified size in bytes. Note that if a single array is bigger than
536−-- the specified size we do not split it to fit. When we coalesce multiple
537−-- arrays if the size would exceed the specified size we do not coalesce
538−-- therefore the actual array size may be less than the specified chunk size.
541−{-# INLINE_NORMAL packArraysChunksOf #-}
542−packArraysChunksOf ::
543− forall m a. (MonadIO m, Prim a)
545− -> D.Stream m (Array a)
546− -> D.Stream m (Array a)
547−packArraysChunksOf n (D.Stream step state) =
548− D.Stream step' (SpliceInitial state)
552− nElem = n `quot` sizeOfPrimElem (undefined :: a)
554− {-# INLINE_LATE step' #-}
555− step' gst (SpliceInitial st) = do
557− -- XXX we can pass the module string from the higher level API
558− error $ "Streamly.Internal.Data.Array.Foreign.Type.packArraysChunksOf: the size of "
559− ++ "arrays [" ++ show n ++ "] must be a natural number"
563− if length arr >= nElem
564− then return $ D.Skip (SpliceYielding arr (SpliceInitial s))
566− buf <- MA.newArray nElem
567− noff <- spliceInto buf 0 arr
568− return $ D.Skip (SpliceBuffering s (buf, noff))
569− D.Skip s -> return $ D.Skip (SpliceInitial s)
570− D.Stop -> return $ D.Stop
572− step' gst (SpliceBuffering st arr2@(buf, boff)) = do
576− if boff + length arr > nElem
578− nArr <- unsafeFreeze buf
579− return $ D.Skip (SpliceYielding (slice nArr 0 boff) (SpliceBuffering s arr2))
581− noff <- spliceInto buf boff arr
582− return $ D.Skip (SpliceBuffering s (buf, noff))
583− D.Skip s -> return $ D.Skip (SpliceBuffering s arr2)
585− nArr <- unsafeFreeze buf
586− return $ D.Skip (SpliceYielding (slice nArr 0 boff) SpliceFinish)
588− step' _ SpliceFinish = return D.Stop
590− step' _ (SpliceYielding arr next) = return $ D.Yield arr next
592−{-# INLINE_NORMAL lpackArraysChunksOf #-}
593−lpackArraysChunksOf ::
594− forall m a. (MonadIO m, Prim a)
596− -> Fold m (Array a) ()
597− -> Fold m (Array a) ()
598−lpackArraysChunksOf n (Fold step1 initial1 extract1) =
599− Fold step initial extract
603− nElem = n `quot` sizeOfPrimElem (undefined :: a)
607− -- XXX we can pass the module string from the higher level API
608− error $ "Streamly.Internal.Data.Array.Foreign.Type.packArraysChunksOf: the size of "
609− ++ "arrays [" ++ show n ++ "] must be a natural number"
611− return $ first (Tuple3' Nothing' 0) res
613− extract (Tuple3' Nothing' _ r1) = extract1 r1
614− extract (Tuple3' (Just' buf) boff r1) = do
615− nArr <- unsafeFreeze buf
616− r <- step1 r1 (slice nArr 0 boff)
618− FL.Partial rr -> extract1 rr
619− FL.Done _ -> return ()
621− step (Tuple3' Nothing' _ r1) arr =
622− if length arr >= nElem
626− FL.Done _ -> return $ FL.Done ()
630− return $ first (Tuple3' Nothing' 0) res
632− buf <- MA.newArray nElem
633− noff <- spliceInto buf 0 arr
634− return $ FL.Partial $ Tuple3' (Just' buf) noff r1
636− step (Tuple3' (Just' buf) boff r1) arr = do
637− noff <- spliceInto buf boff arr
641− nArr <- unsafeFreeze buf
642− r <- step1 r1 (slice nArr 0 noff)
644− FL.Done _ -> return $ FL.Done ()
648− return $ first (Tuple3' Nothing' 0) res
649− else return $ FL.Partial $ Tuple3' (Just' buf) noff r1
651−data SplitState s arr
655− | Yielding arr (SplitState s arr)
658−-- | Split a stream of arrays on a given separator byte, dropping the separator
659−-- and coalescing all the arrays between two separators into a single array.
662−{-# INLINE_NORMAL splitOn #-}
666− -> D.Stream m (Array Word8)
667− -> D.Stream m (Array Word8)
668−splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
672− {-# INLINE_LATE step' #-}
673− step' gst (Initial st) = do
677− (arr1, marr2) <- breakOn byte arr
678− return $ case marr2 of
679− Nothing -> D.Skip (Buffering s arr1)
680− Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
681− D.Skip s -> return $ D.Skip (Initial s)
682− D.Stop -> return $ D.Stop
684− step' gst (Buffering st buf) = do
688− (arr1, marr2) <- breakOn byte arr
689− buf' <- spliceTwo buf arr1
690− return $ case marr2 of
691− Nothing -> D.Skip (Buffering s buf')
692− Just x -> D.Skip (Yielding buf' (Splitting s x))
693− D.Skip s -> return $ D.Skip (Buffering s buf)
695− if byteLength buf == 0
697− else D.Skip (Yielding buf Finishing)
699− step' _ (Splitting st buf) = do
700− (arr1, marr2) <- breakOn byte buf
701− return $ case marr2 of
702− Nothing -> D.Skip $ Buffering st arr1
703− Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
705− step' _ (Yielding arr next) = return $ D.Yield arr next
706− step' _ Finishing = return $ D.Stop