@@ -0,0 +1,694 @@
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.Primitive.Types (Prim(..), sizeOf)
12+#if __GLASGOW_HASKELL__ < 808
13+import Data.Semigroup (Semigroup(..))
15+import Data.Word (Word8)
16+import Streamly.Internal.Data.Strict (Tuple3'(..), Maybe'(..))
17+import Streamly.Internal.Data.Fold.Types (Fold(..))
18+import Streamly.Internal.Data.SVar (adaptState)
19+import System.IO.Unsafe (unsafePerformIO)
20+import Text.Read (readPrec, readListPrec, readListPrecDefault)
22+import qualified GHC.Exts as Exts
23+import qualified Prelude as P
24+import qualified Streamly.Internal.Data.Fold as FL
25+import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
26+import qualified Streamly.Internal.Data.Stream.StreamK as K
28+import GHC.Exts hiding (fromListN, fromList, toList)
29+import Prelude hiding (length, unlines, foldr)
31+-------------------------------------------------------------------------------
33+-------------------------------------------------------------------------------
35+data Array a = Array ByteArray# Int Int
37+-------------------------------------------------------------------------------
39+-------------------------------------------------------------------------------
41+-- | Both arrays must fully contain the specified ranges, but this is not
42+-- checked. The two arrays must not be the same array in different states, but
43+-- this is not checked either.
44+{-# INLINE unsafeCopy #-}
46+ forall m a. (MonadIO m, Prim a)
47+ => MA.Array a -- ^ destination array
48+ -> Int -- ^ offset into destination array
49+ -> Array a -- ^ source array
50+ -> Int -- ^ offset into source array
51+ -> Int -- ^ number of elements to copy
53+unsafeCopy (MA.Array dst#) (I# doff#) (Array src# (I# off#) _) (I# soff#) (I# n#) =
55+ let toBytes cnt# = cnt# *# (sizeOf# (undefined :: a))
59+ (toBytes (off# +# soff#))
64+-------------------------------------------------------------------------------
65+-- Basic Byte Array Operations
66+-------------------------------------------------------------------------------
68+{-# INLINE unsafeFreeze #-}
69+unsafeFreeze :: (Prim a, MonadIO m) => MA.Array a -> m (Array a)
70+unsafeFreeze marr@(MA.Array arr#) =
72+ len <- MA.length marr
74+ case unsafeFreezeByteArray# arr# s# of
75+ (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
77+{-# INLINE unsafeFreezeWithShrink #-}
78+unsafeFreezeWithShrink ::
79+ (Prim a, MonadIO m) => MA.Array a -> Int -> m (Array a)
80+unsafeFreezeWithShrink marr@(MA.Array arr#) n =
82+ MA.shrinkArray marr n
83+ len <- MA.length marr
85+ case unsafeFreezeByteArray# arr# s# of
86+ (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
89+-- Should never be used in general
90+{-# INLINE unsafeThaw #-}
91+unsafeThaw :: MonadIO m => Array a -> m (MA.Array a)
92+unsafeThaw (Array arr#) =
93+ primitive $ \s# -> (# s#, MA.Array (unsafeCoerce# arr#) #)
96+-- Unsafe because the index bounds are not checked
97+{-# INLINE unsafeIndex #-}
98+unsafeIndex :: Prim a => Array a -> Int -> a
99+unsafeIndex (Array arr# (I# off#) _) (I# i#) = indexByteArray# arr# (off# +# i#)
102+sameByteArray :: ByteArray# -> ByteArray# -> Bool
103+sameByteArray ba1 ba2 =
104+ case reallyUnsafePtrEquality#
105+ (unsafeCoerce# ba1 :: ())
106+ (unsafeCoerce# ba2 :: ()) of
109+-------------------------------------------------------------------------------
111+-------------------------------------------------------------------------------
113+-- XXX move this section to mutable array module?
115+mkChunkSizeKB :: Int -> Int
116+mkChunkSizeKB n = n * k
119+-- | Default maximum buffer size in bytes, for reading from and writing to IO
120+-- devices, the value is 32KB minus GHC allocation overhead, which is a few
121+-- bytes, so that the actual allocation is 32KB.
122+defaultChunkSize :: Int
123+defaultChunkSize = mkChunkSizeKB 32
125+-------------------------------------------------------------------------------
127+-------------------------------------------------------------------------------
129+-- XXX rename to byteCount?
130+{-# INLINE byteLength #-}
131+byteLength :: forall a. Prim a => Array a -> Int
132+byteLength (Array _ _ len) = len * sizeOf (undefined :: a)
134+-- XXX Also, rename to elemCount
135+-- XXX I would prefer length to keep the API consistent
136+-- XXX Also, re-export sizeOf from Primitive
137+{-# INLINE length #-}
138+length :: Array a -> Int
139+length (Array _ _ len) = len
141+-------------------------------------------------------------------------------
143+-------------------------------------------------------------------------------
145+-- | Use a slice of an array as another array. Note that this is unsafe and does
146+-- not check the bounds
147+slice :: Array a -> Int -> Int -> Array a
148+slice (Array arr# off _) off1 len1 = Array arr# (off + off1) len1
150+nil :: Prim a => Array a
152+ unsafePerformIO $ do
153+ arr <- MA.newArray 0
156+-- | Fold the whole input to a single array.
158+-- /Caution! Do not use this on infinite streams./
161+{-# INLINE_NORMAL write #-}
162+write :: (MonadIO m, Prim a) => Fold m a (Array a)
163+write = FL.mapM unsafeFreeze MA.write
165+-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
169+{-# INLINE_NORMAL writeN #-}
170+writeN :: (MonadIO m, Prim a) => Int -> Fold m a (Array a)
171+writeN limit = FL.mapM unsafeFreeze (MA.writeN limit)
173+-- | Like 'writeN' but does not check the array bounds when writing. The fold
174+-- driver must not call the step function more than 'n' times otherwise it will
175+-- corrupt the memory and crash. This function exists mainly because any
176+-- conditional in the step function blocks fusion causing 10x performance
180+{-# INLINE_NORMAL writeNUnsafe #-}
181+writeNUnsafe :: (MonadIO m, Prim a) => Int -> Fold m a (Array a)
182+writeNUnsafe limit = FL.mapM unsafeFreeze (MA.writeNUnsafe limit)
184+{-# INLINE_NORMAL fromStreamDN #-}
185+fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (Array a)
186+fromStreamDN limit str = MA.fromStreamDN limit str >>= unsafeFreeze
188+{-# INLINE fromStreamD #-}
189+fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (Array a)
190+fromStreamD str = MA.fromStreamD str >>= unsafeFreeze
192+-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
194+{-# INLINE_NORMAL fromStreamDArraysOf #-}
195+fromStreamDArraysOf ::
196+ (MonadIO m, Prim a) => Int -> D.Stream m a -> D.Stream m (Array a)
197+fromStreamDArraysOf n str = D.mapM unsafeFreeze (MA.fromStreamDArraysOf n str)
199+-- XXX derive from MA.fromListN?
200+{-# INLINE fromListN #-}
201+fromListN :: Prim a => Int -> [a] -> Array a
202+fromListN len xs = unsafePerformIO $ MA.fromListNM len xs >>= unsafeFreeze
204+-- XXX derive from MA.fromList?
205+{-# INLINE fromList #-}
206+fromList :: Prim a => [a] -> Array a
207+fromList xs = fromListN (P.length xs) xs
209+-------------------------------------------------------------------------------
211+-------------------------------------------------------------------------------
213+-- | Splice two immutable arrays creating a new immutable array.
214+{-# INLINE spliceTwo #-}
215+spliceTwo :: (MonadIO m, Prim a) => Array a -> Array a -> m (Array a)
219+ a3 <- MA.newArray (l1 + l2)
220+ unsafeCopy a3 0 a1 0 l1
221+ unsafeCopy a3 l1 a2 0 l2
222+ unsafeFreeze a3 -- Use `unsafeFreezeWith off len`?
224+-------------------------------------------------------------------------------
226+-------------------------------------------------------------------------------
228+{-# INLINE_LATE toListFB #-}
229+toListFB :: Prim a => (a -> b -> b) -> b -> Array a -> b
230+toListFB c n arr = go 0
235+ let !x = unsafeIndex arr p
238+-- | Convert an 'Array' into a list.
241+{-# INLINE toList #-}
242+toList :: Prim a => Array a -> [a]
243+toList s = build (\c n -> toListFB c n s)
245+-------------------------------------------------------------------------------
247+-------------------------------------------------------------------------------
249+instance (Eq a, Prim a) => Eq (Array a) where
251+ a1@(Array ba1# _ len1) == a2@(Array ba2# _ len2)
252+ | sameByteArray ba1# ba2# = True
253+ | len1 /= len2 = False
254+ | otherwise = loop (len1 - 1)
260+ | otherwise = unsafeIndex a1 i == unsafeIndex a2 i && loop (i - 1)
262+-- | Lexicographic ordering. Subject to change between major versions.
263+instance (Ord a, Prim a) => Ord (Array a) where
264+ {-# INLINE compare #-}
265+ compare a1@(Array ba1# _ len1) a2@(Array ba2# _ len2)
266+ | sameByteArray ba1# ba2# = EQ
267+ | otherwise = loop 0
275+ compare (unsafeIndex a1 i) (unsafeIndex a2 i) <> loop (i + 1)
276+ | otherwise = compare len1 len2
278+instance Prim a => Semigroup (Array a) where
279+ -- XXX can't we use runST instead of inlineIO?
280+ -- XXX I plan to remove MonadIO and replace it with IO
281+ a <> b = unsafePerformIO (spliceTwo a b :: IO (Array a))
283+instance Prim a => Monoid (Array a) where
287+instance NFData (Array a) where
292+-- XXX check if this is compatible with Memory.Array?
293+-- XXX It isn't. I might prefer this Show instance though
294+-- XXX Memory.Array: showsPrec _ = shows . toList
295+instance (Show a, Prim a) => Show (Array a) where
297+ showParen (p > 10) $
298+ showString "fromListN "
303+instance (a ~ Char) => IsString (Array a) where
304+ {-# INLINE fromString #-}
305+ fromString = fromList
307+-- GHC versions 8.0 and below cannot derive IsList
308+instance Prim a => IsList (Array a) where
309+ type (Item (Array a)) = a
311+ {-# INLINE fromList #-}
314+ {-# INLINE fromListN #-}
315+ fromListN = fromListN
317+ {-# INLINE toList #-}
320+instance (Prim a, Read a, Show a) => Read (Array a) where
321+ {-# INLINE readPrec #-}
322+ readPrec = fromList <$> readPrec
323+ readListPrec = readListPrecDefault
325+-- XXX these folds can be made common with mutable arrays by defining a
326+-- unsafeIndex in the specific module?
328+-------------------------------------------------------------------------------
330+-------------------------------------------------------------------------------
333+foldr :: Prim a => (a -> b -> b) -> b -> Array a -> b
341+ | len > i = f (unsafeIndex arr i) (go (i + 1))
344+-- | Strict right-associated fold over the elements of an 'Array'.
345+{-# INLINE foldr' #-}
346+foldr' :: Prim a => (a -> b -> b) -> b -> Array a -> b
347+foldr' f z0 arr = go (length arr - 1) z0
353+ | otherwise = go (i - 1) (f (unsafeIndex arr i) acc)
355+-- | Strict left-associated fold over the elements of an 'Array'.
356+{-# INLINE foldl' #-}
357+foldl' :: Prim a => (b -> a -> b) -> b -> Array a -> b
358+foldl' f z0 arr = go 0 z0
365+ | i < len = go (i + 1) (f acc (unsafeIndex arr i))
368+-- | Strict left-associated fold over the elements of an 'Array'.
369+{-# INLINE foldlM' #-}
370+foldlM' :: (Prim a, Monad m) => (b -> a -> m b) -> b -> Array a -> m b
371+foldlM' f z0 arr = go 0 z0
379+ acc2 <- f acc1 (unsafeIndex arr i)
381+ | otherwise = return acc1
383+-------------------------------------------------------------------------------
384+-- Converting to streams
385+-------------------------------------------------------------------------------
387+{-# INLINE_NORMAL toStreamD #-}
388+toStreamD :: (Prim a, Monad m) => Array a -> D.Stream m a
389+toStreamD arr = D.Stream step 0
393+ {-# INLINE_LATE step #-}
395+ | i == length arr = return D.Stop
396+ step _ i = return $ D.Yield (unsafeIndex arr i) (i + 1)
398+{-# INLINE toStreamK #-}
399+toStreamK :: (K.IsStream t, Prim a) => Array a -> t m a
409+ let !x = unsafeIndex arr p
410+ in x `K.cons` go (p + 1)
412+{-# INLINE_NORMAL toStreamDRev #-}
413+toStreamDRev :: (Prim a, Monad m) => Array a -> D.Stream m a
414+toStreamDRev arr = D.Stream step (length arr - 1)
418+ {-# INLINE_LATE step #-}
420+ | i < 0 = return D.Stop
421+ step _ i = return $ D.Yield (unsafeIndex arr i) (i - 1)
423+{-# INLINE toStreamKRev #-}
424+toStreamKRev :: (K.IsStream t, Prim a) => Array a -> t m a
425+toStreamKRev arr = go (length arr - 1)
429+ go p | p == -1 = K.nil
431+ let !x = unsafeIndex arr p
432+ in x `K.cons` go (p - 1)
434+-------------------------------------------------------------------------------
436+-------------------------------------------------------------------------------
438+data FlattenState s a =
440+ | InnerLoop s !(Array a) !Int !Int
442+{-# INLINE_NORMAL flattenArrays #-}
443+flattenArrays :: (MonadIO m, Prim a) => D.Stream m (Array a) -> D.Stream m a
444+flattenArrays (D.Stream step state) = D.Stream step' (OuterLoop state)
448+ {-# INLINE_LATE step' #-}
449+ step' gst (OuterLoop st) = do
450+ r <- step (adaptState gst) st
453+ let len = length arr
455+ then D.Skip (OuterLoop s)
456+ else D.Skip (InnerLoop s arr len 0)
457+ D.Skip s -> D.Skip (OuterLoop s)
460+ step' _ (InnerLoop st _ len i) | i == len =
461+ return $ D.Skip $ OuterLoop st
463+ step' _ (InnerLoop st arr len i) = do
464+ let x = unsafeIndex arr i
465+ return $ D.Yield x (InnerLoop st arr len (i + 1))
467+{-# INLINE_NORMAL flattenArraysRev #-}
468+flattenArraysRev :: (MonadIO m, Prim a) => D.Stream m (Array a) -> D.Stream m a
469+flattenArraysRev (D.Stream step state) = D.Stream step' (OuterLoop state)
473+ {-# INLINE_LATE step' #-}
474+ step' gst (OuterLoop st) = do
475+ r <- step (adaptState gst) st
478+ let len = length arr
480+ then D.Skip (OuterLoop s)
481+ else D.Skip (InnerLoop s arr len (len - 1))
482+ D.Skip s -> D.Skip (OuterLoop s)
485+ step' _ (InnerLoop st _ _ i) | i == -1 =
486+ return $ D.Skip $ OuterLoop st
488+ step' _ (InnerLoop st arr len i) = do
489+ let x = unsafeIndex arr i
490+ return $ D.Yield x (InnerLoop st arr len (i - 1))
492+{-# INLINE_NORMAL unlines #-}
493+unlines :: (MonadIO m, Prim a) => a -> D.Stream m (Array a) -> D.Stream m a
494+unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
498+ {-# INLINE_LATE step' #-}
499+ step' gst (OuterLoop st) = do
500+ r <- step (adaptState gst) st
503+ let len = length arr
504+ in D.Skip (InnerLoop s arr len 0)
505+ D.Skip s -> D.Skip (OuterLoop s)
508+ step' _ (InnerLoop st _ len i) | i == len =
509+ return $ D.Yield sep $ OuterLoop st
511+ step' _ (InnerLoop st arr len i) = do
512+ let x = unsafeIndex arr i
513+ return $ D.Yield x (InnerLoop st arr len (i + 1))
515+-- Splice an array into a pre-reserved mutable array. The user must ensure
516+-- that there is enough space in the mutable array.
517+{-# INLINE spliceInto #-}
518+spliceInto :: (MonadIO m, Prim a) => MA.Array a -> Int -> Array a -> m Int
519+spliceInto dst doff src@(Array _ _ len) = do
520+ unsafeCopy dst doff src 0 len
523+data SpliceState s arr1 arr2
525+ | SpliceBuffering s arr2
526+ | SpliceYielding arr1 (SpliceState s arr1 arr2)
529+-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
530+-- maximum specified size in bytes. Note that if a single array is bigger than
531+-- the specified size we do not split it to fit. When we coalesce multiple
532+-- arrays if the size would exceed the specified size we do not coalesce
533+-- therefore the actual array size may be less than the specified chunk size.
536+{-# INLINE_NORMAL packArraysChunksOf #-}
537+packArraysChunksOf ::
538+ forall m a. (MonadIO m, Prim a)
540+ -> D.Stream m (Array a)
541+ -> D.Stream m (Array a)
542+packArraysChunksOf n (D.Stream step state) =
543+ D.Stream step' (SpliceInitial state)
547+ nElem = n `quot` sizeOf (undefined :: a)
549+ {-# INLINE_LATE step' #-}
550+ step' gst (SpliceInitial st) = do
552+ -- XXX we can pass the module string from the higher level API
553+ error $ "Streamly.Internal.Memory.Array.Types.packArraysChunksOf: the size of "
554+ ++ "arrays [" ++ show n ++ "] must be a natural number"
558+ if length arr >= nElem
559+ then return $ D.Skip (SpliceYielding arr (SpliceInitial s))
561+ buf <- MA.newArray nElem
562+ noff <- spliceInto buf 0 arr
563+ return $ D.Skip (SpliceBuffering s (buf, noff))
564+ D.Skip s -> return $ D.Skip (SpliceInitial s)
565+ D.Stop -> return $ D.Stop
567+ step' gst (SpliceBuffering st arr2@(buf, boff)) = do
571+ if boff + length arr > nElem
573+ nArr <- unsafeFreeze buf
574+ return $ D.Skip (SpliceYielding (slice nArr 0 boff) (SpliceBuffering s arr2))
576+ noff <- spliceInto buf boff arr
577+ return $ D.Skip (SpliceBuffering s (buf, noff))
578+ D.Skip s -> return $ D.Skip (SpliceBuffering s arr2)
580+ nArr <- unsafeFreeze buf
581+ return $ D.Skip (SpliceYielding (slice nArr 0 boff) SpliceFinish)
583+ step' _ SpliceFinish = return D.Stop
585+ step' _ (SpliceYielding arr next) = return $ D.Yield arr next
587+{-# INLINE_NORMAL lpackArraysChunksOf #-}
588+lpackArraysChunksOf ::
589+ forall m a. (MonadIO m, Prim a)
591+ -> Fold m (Array a) ()
592+ -> Fold m (Array a) ()
593+lpackArraysChunksOf n (Fold step1 initial1 extract1) =
594+ Fold step initial extract
598+ nElem = n `quot` sizeOf (undefined :: a)
602+ -- XXX we can pass the module string from the higher level API
603+ error $ "Streamly.Internal.Memory.Array.Types.packArraysChunksOf: the size of "
604+ ++ "arrays [" ++ show n ++ "] must be a natural number"
606+ return (Tuple3' Nothing' 0 r1)
608+ extract (Tuple3' Nothing' _ r1) = extract1 r1
609+ extract (Tuple3' (Just' buf) boff r1) = do
610+ nArr <- unsafeFreeze buf
611+ r <- step1 r1 (slice nArr 0 boff)
614+ step (Tuple3' Nothing' _ r1) arr =
616+ if length arr >= nElem
621+ return (Tuple3' Nothing' 0 r1')
623+ buf <- MA.newArray nElem
624+ noff <- spliceInto buf 0 arr
625+ return (Tuple3' (Just' buf) noff r1)
627+ step (Tuple3' (Just' buf) boff r1) arr = do
628+ noff <- spliceInto buf boff arr
632+ nArr <- unsafeFreeze buf
633+ r <- step1 r1 (slice nArr 0 noff)
636+ return (Tuple3' Nothing' 0 r1')
637+ else return (Tuple3' (Just' buf) noff r1)
639+data SplitState s arr
643+ | Yielding arr (SplitState s arr)
646+-- | Split a stream of arrays on a given separator byte, dropping the separator
647+-- and coalescing all the arrays between two separators into a single array.
650+{-# INLINE_NORMAL splitOn #-}
654+ -> D.Stream m (Array Word8)
655+ -> D.Stream m (Array Word8)
656+splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
660+ {-# INLINE_LATE step' #-}
661+ step' gst (Initial st) = do
665+ (arr1, marr2) <- breakOn byte arr
666+ return $ case marr2 of
667+ Nothing -> D.Skip (Buffering s arr1)
668+ Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
669+ D.Skip s -> return $ D.Skip (Initial s)
670+ D.Stop -> return $ D.Stop
672+ step' gst (Buffering st buf) = do
676+ (arr1, marr2) <- breakOn byte arr
677+ buf' <- spliceTwo buf arr1
678+ return $ case marr2 of
679+ Nothing -> D.Skip (Buffering s buf')
680+ Just x -> D.Skip (Yielding buf' (Splitting s x))
681+ D.Skip s -> return $ D.Skip (Buffering s buf)
683+ if byteLength buf == 0
685+ else D.Skip (Yielding buf Finishing)
687+ step' _ (Splitting st buf) = do
688+ (arr1, marr2) <- breakOn byte buf
689+ return $ case marr2 of
690+ Nothing -> D.Skip $ Buffering st arr1
691+ Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
693+ step' _ (Yielding arr next) = return $ D.Yield arr next
694+ step' _ Finishing = return $ D.Stop