@@ -0,0 +1,694 @@
1+-- Copyright : (c) 2020 Composewell Technologies
2+-- License : BSD-3-Clause
3+-- Maintainer : streamly@composewell.com
4+-- Stability : experimental
5+-- Portability : GHC
6+
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(..))
14+#endif
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)
21+
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
27+
28+import GHC.Exts hiding (fromListN, fromList, toList)
29+import Prelude hiding (length, unlines, foldr)
30+
31+-------------------------------------------------------------------------------
32+-- Array Data Type
33+-------------------------------------------------------------------------------
34+
35+data Array a = Array ByteArray# Int Int
36+
37+-------------------------------------------------------------------------------
38+-- Utilities
39+-------------------------------------------------------------------------------
40+
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 #-}
45+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
52+ -> m ()
53+unsafeCopy (MA.Array dst#) (I# doff#) (Array src# (I# off#) _) (I# soff#) (I# n#) =
54+ liftIO $ do
55+ let toBytes cnt# = cnt# *# (sizeOf# (undefined :: a))
56+ primitive_ $
57+ copyByteArray#
58+ src#
59+ (toBytes (off# +# soff#))
60+ dst#
61+ (toBytes doff#)
62+ (toBytes n#)
63+
64+-------------------------------------------------------------------------------
65+-- Basic Byte Array Operations
66+-------------------------------------------------------------------------------
67+
68+{-# INLINE unsafeFreeze #-}
69+unsafeFreeze :: (Prim a, MonadIO m) => MA.Array a -> m (Array a)
70+unsafeFreeze marr@(MA.Array arr#) =
71+ liftIO $ do
72+ len <- MA.length marr
73+ primitive $ \s# ->
74+ case unsafeFreezeByteArray# arr# s# of
75+ (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
76+
77+{-# INLINE unsafeFreezeWithShrink #-}
78+unsafeFreezeWithShrink ::
79+ (Prim a, MonadIO m) => MA.Array a -> Int -> m (Array a)
80+unsafeFreezeWithShrink marr@(MA.Array arr#) n =
81+ liftIO $ do
82+ MA.shrinkArray marr n
83+ len <- MA.length marr
84+ primitive $ \s# ->
85+ case unsafeFreezeByteArray# arr# s# of
86+ (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
87+
88+{-
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#) #)
94+-}
95+
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#)
100+
101+-- unsafe
102+sameByteArray :: ByteArray# -> ByteArray# -> Bool
103+sameByteArray ba1 ba2 =
104+ case reallyUnsafePtrEquality#
105+ (unsafeCoerce# ba1 :: ())
106+ (unsafeCoerce# ba2 :: ()) of
107+ r -> isTrue# r
108+
109+-------------------------------------------------------------------------------
110+-- Chunk Size
111+-------------------------------------------------------------------------------
112+
113+-- XXX move this section to mutable array module?
114+
115+mkChunkSizeKB :: Int -> Int
116+mkChunkSizeKB n = n * k
117+ where k = 1024
118+
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
124+
125+-------------------------------------------------------------------------------
126+-- Length
127+-------------------------------------------------------------------------------
128+
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)
133+
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
140+
141+-------------------------------------------------------------------------------
142+-- Construction
143+-------------------------------------------------------------------------------
144+
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
149+
150+nil :: Prim a => Array a
151+nil =
152+ unsafePerformIO $ do
153+ arr <- MA.newArray 0
154+ unsafeFreeze arr
155+
156+-- | Fold the whole input to a single array.
157+--
158+-- /Caution! Do not use this on infinite streams./
159+--
160+-- /Internal/
161+{-# INLINE_NORMAL write #-}
162+write :: (MonadIO m, Prim a) => Fold m a (Array a)
163+write = FL.mapM unsafeFreeze MA.write
164+
165+-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
166+-- 'Array'.
167+--
168+-- /Internal/
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)
172+
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
177+-- slowdown.
178+--
179+-- /Internal/
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)
183+
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
187+
188+{-# INLINE fromStreamD #-}
189+fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (Array a)
190+fromStreamD str = MA.fromStreamD str >>= unsafeFreeze
191+
192+-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
193+-- arrays of size n.
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)
198+
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
203+
204+-- XXX derive from MA.fromList?
205+{-# INLINE fromList #-}
206+fromList :: Prim a => [a] -> Array a
207+fromList xs = fromListN (P.length xs) xs
208+
209+-------------------------------------------------------------------------------
210+-- Combining
211+-------------------------------------------------------------------------------
212+
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)
216+spliceTwo a1 a2 = do
217+ let l1 = length a1
218+ l2 = length a2
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`?
223+
224+-------------------------------------------------------------------------------
225+-- Elimination
226+-------------------------------------------------------------------------------
227+
228+{-# INLINE_LATE toListFB #-}
229+toListFB :: Prim a => (a -> b -> b) -> b -> Array a -> b
230+toListFB c n arr = go 0
231+ where
232+ len = length arr
233+ go p | p == len = n
234+ go p =
235+ let !x = unsafeIndex arr p
236+ in c x (go (p + 1))
237+
238+-- | Convert an 'Array' into a list.
239+--
240+-- /Internal/
241+{-# INLINE toList #-}
242+toList :: Prim a => Array a -> [a]
243+toList s = build (\c n -> toListFB c n s)
244+
245+-------------------------------------------------------------------------------
246+-- Instances
247+-------------------------------------------------------------------------------
248+
249+instance (Eq a, Prim a) => Eq (Array a) where
250+ {-# INLINE (==) #-}
251+ a1@(Array ba1# _ len1) == a2@(Array ba2# _ len2)
252+ | sameByteArray ba1# ba2# = True
253+ | len1 /= len2 = False
254+ | otherwise = loop (len1 - 1)
255+
256+ where
257+
258+ loop !i
259+ | i < 0 = True
260+ | otherwise = unsafeIndex a1 i == unsafeIndex a2 i && loop (i - 1)
261+
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
268+
269+ where
270+
271+ sz = min len1 len2
272+
273+ loop !i
274+ | i < sz =
275+ compare (unsafeIndex a1 i) (unsafeIndex a2 i) <> loop (i + 1)
276+ | otherwise = compare len1 len2
277+
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))
282+
283+instance Prim a => Monoid (Array a) where
284+ mempty = nil
285+ mappend = (<>)
286+
287+instance NFData (Array a) where
288+ {-# INLINE rnf #-}
289+ rnf _ = ()
290+
291+
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
296+ showsPrec p a =
297+ showParen (p > 10) $
298+ showString "fromListN "
299+ . shows (length a)
300+ . showString " "
301+ . shows (toList a)
302+
303+instance (a ~ Char) => IsString (Array a) where
304+ {-# INLINE fromString #-}
305+ fromString = fromList
306+
307+-- GHC versions 8.0 and below cannot derive IsList
308+instance Prim a => IsList (Array a) where
309+ type (Item (Array a)) = a
310+
311+ {-# INLINE fromList #-}
312+ fromList = fromList
313+
314+ {-# INLINE fromListN #-}
315+ fromListN = fromListN
316+
317+ {-# INLINE toList #-}
318+ toList = toList
319+
320+instance (Prim a, Read a, Show a) => Read (Array a) where
321+ {-# INLINE readPrec #-}
322+ readPrec = fromList <$> readPrec
323+ readListPrec = readListPrecDefault
324+
325+-- XXX these folds can be made common with mutable arrays by defining a
326+-- unsafeIndex in the specific module?
327+
328+-------------------------------------------------------------------------------
329+-- Folds
330+-------------------------------------------------------------------------------
331+
332+{-# INLINE foldr #-}
333+foldr :: Prim a => (a -> b -> b) -> b -> Array a -> b
334+foldr f z arr = go 0
335+
336+ where
337+
338+ !len = length arr
339+
340+ go !i
341+ | len > i = f (unsafeIndex arr i) (go (i + 1))
342+ | otherwise = z
343+
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
348+
349+ where
350+
351+ go !i !acc
352+ | i < 0 = acc
353+ | otherwise = go (i - 1) (f (unsafeIndex arr i) acc)
354+
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
359+
360+ where
361+
362+ !len = length arr
363+
364+ go !i !acc
365+ | i < len = go (i + 1) (f acc (unsafeIndex arr i))
366+ | otherwise = acc
367+
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
372+
373+ where
374+
375+ !len = length arr
376+
377+ go !i !acc1
378+ | i < len = do
379+ acc2 <- f acc1 (unsafeIndex arr i)
380+ go (i + 1) acc2
381+ | otherwise = return acc1
382+
383+-------------------------------------------------------------------------------
384+-- Converting to streams
385+-------------------------------------------------------------------------------
386+
387+{-# INLINE_NORMAL toStreamD #-}
388+toStreamD :: (Prim a, Monad m) => Array a -> D.Stream m a
389+toStreamD arr = D.Stream step 0
390+
391+ where
392+
393+ {-# INLINE_LATE step #-}
394+ step _ i
395+ | i == length arr = return D.Stop
396+ step _ i = return $ D.Yield (unsafeIndex arr i) (i + 1)
397+
398+{-# INLINE toStreamK #-}
399+toStreamK :: (K.IsStream t, Prim a) => Array a -> t m a
400+toStreamK arr = go 0
401+
402+ where
403+
404+ len = length arr
405+
406+ go p
407+ | p == len = K.nil
408+ | otherwise =
409+ let !x = unsafeIndex arr p
410+ in x `K.cons` go (p + 1)
411+
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)
415+
416+ where
417+
418+ {-# INLINE_LATE step #-}
419+ step _ i
420+ | i < 0 = return D.Stop
421+ step _ i = return $ D.Yield (unsafeIndex arr i) (i - 1)
422+
423+{-# INLINE toStreamKRev #-}
424+toStreamKRev :: (K.IsStream t, Prim a) => Array a -> t m a
425+toStreamKRev arr = go (length arr - 1)
426+
427+ where
428+
429+ go p | p == -1 = K.nil
430+ | otherwise =
431+ let !x = unsafeIndex arr p
432+ in x `K.cons` go (p - 1)
433+
434+-------------------------------------------------------------------------------
435+-- Stream of Arrays
436+-------------------------------------------------------------------------------
437+
438+data FlattenState s a =
439+ OuterLoop s
440+ | InnerLoop s !(Array a) !Int !Int
441+
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)
445+
446+ where
447+
448+ {-# INLINE_LATE step' #-}
449+ step' gst (OuterLoop st) = do
450+ r <- step (adaptState gst) st
451+ return $ case r of
452+ D.Yield arr s ->
453+ let len = length arr
454+ in if len == 0
455+ then D.Skip (OuterLoop s)
456+ else D.Skip (InnerLoop s arr len 0)
457+ D.Skip s -> D.Skip (OuterLoop s)
458+ D.Stop -> D.Stop
459+
460+ step' _ (InnerLoop st _ len i) | i == len =
461+ return $ D.Skip $ OuterLoop st
462+
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))
466+
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)
470+
471+ where
472+
473+ {-# INLINE_LATE step' #-}
474+ step' gst (OuterLoop st) = do
475+ r <- step (adaptState gst) st
476+ return $ case r of
477+ D.Yield arr s ->
478+ let len = length arr
479+ in if len == 0
480+ then D.Skip (OuterLoop s)
481+ else D.Skip (InnerLoop s arr len (len - 1))
482+ D.Skip s -> D.Skip (OuterLoop s)
483+ D.Stop -> D.Stop
484+
485+ step' _ (InnerLoop st _ _ i) | i == -1 =
486+ return $ D.Skip $ OuterLoop st
487+
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))
491+
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)
495+
496+ where
497+
498+ {-# INLINE_LATE step' #-}
499+ step' gst (OuterLoop st) = do
500+ r <- step (adaptState gst) st
501+ return $ case r of
502+ D.Yield arr s ->
503+ let len = length arr
504+ in D.Skip (InnerLoop s arr len 0)
505+ D.Skip s -> D.Skip (OuterLoop s)
506+ D.Stop -> D.Stop
507+
508+ step' _ (InnerLoop st _ len i) | i == len =
509+ return $ D.Yield sep $ OuterLoop st
510+
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))
514+
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
521+ return $ doff + len
522+
523+data SpliceState s arr1 arr2
524+ = SpliceInitial s
525+ | SpliceBuffering s arr2
526+ | SpliceYielding arr1 (SpliceState s arr1 arr2)
527+ | SpliceFinish
528+
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.
534+--
535+-- /Internal/
536+{-# INLINE_NORMAL packArraysChunksOf #-}
537+packArraysChunksOf ::
538+ forall m a. (MonadIO m, Prim a)
539+ => Int
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)
544+
545+ where
546+
547+ nElem = n `quot` sizeOf (undefined :: a)
548+
549+ {-# INLINE_LATE step' #-}
550+ step' gst (SpliceInitial st) = do
551+ when (n <= 0) $
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"
555+ r <- step gst st
556+ case r of
557+ D.Yield arr s ->
558+ if length arr >= nElem
559+ then return $ D.Skip (SpliceYielding arr (SpliceInitial s))
560+ else do
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
566+
567+ step' gst (SpliceBuffering st arr2@(buf, boff)) = do
568+ r <- step gst st
569+ case r of
570+ D.Yield arr s -> do
571+ if boff + length arr > nElem
572+ then do
573+ nArr <- unsafeFreeze buf
574+ return $ D.Skip (SpliceYielding (slice nArr 0 boff) (SpliceBuffering s arr2))
575+ else do
576+ noff <- spliceInto buf boff arr
577+ return $ D.Skip (SpliceBuffering s (buf, noff))
578+ D.Skip s -> return $ D.Skip (SpliceBuffering s arr2)
579+ D.Stop -> do
580+ nArr <- unsafeFreeze buf
581+ return $ D.Skip (SpliceYielding (slice nArr 0 boff) SpliceFinish)
582+
583+ step' _ SpliceFinish = return D.Stop
584+
585+ step' _ (SpliceYielding arr next) = return $ D.Yield arr next
586+
587+{-# INLINE_NORMAL lpackArraysChunksOf #-}
588+lpackArraysChunksOf ::
589+ forall m a. (MonadIO m, Prim a)
590+ => Int
591+ -> Fold m (Array a) ()
592+ -> Fold m (Array a) ()
593+lpackArraysChunksOf n (Fold step1 initial1 extract1) =
594+ Fold step initial extract
595+
596+ where
597+
598+ nElem = n `quot` sizeOf (undefined :: a)
599+
600+ initial = do
601+ when (n <= 0) $
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"
605+ r1 <- initial1
606+ return (Tuple3' Nothing' 0 r1)
607+
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)
612+ extract1 r
613+
614+ step (Tuple3' Nothing' _ r1) arr =
615+
616+ if length arr >= nElem
617+ then do
618+ r <- step1 r1 arr
619+ extract1 r
620+ r1' <- initial1
621+ return (Tuple3' Nothing' 0 r1')
622+ else do
623+ buf <- MA.newArray nElem
624+ noff <- spliceInto buf 0 arr
625+ return (Tuple3' (Just' buf) noff r1)
626+
627+ step (Tuple3' (Just' buf) boff r1) arr = do
628+ noff <- spliceInto buf boff arr
629+
630+ if noff >= nElem
631+ then do
632+ nArr <- unsafeFreeze buf
633+ r <- step1 r1 (slice nArr 0 noff)
634+ extract1 r
635+ r1' <- initial1
636+ return (Tuple3' Nothing' 0 r1')
637+ else return (Tuple3' (Just' buf) noff r1)
638+
639+data SplitState s arr
640+ = Initial s
641+ | Buffering s arr
642+ | Splitting s arr
643+ | Yielding arr (SplitState s arr)
644+ | Finishing
645+
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.
648+--
649+-- /Internal/
650+{-# INLINE_NORMAL splitOn #-}
651+splitOn ::
652+ MonadIO m
653+ => Word8
654+ -> D.Stream m (Array Word8)
655+ -> D.Stream m (Array Word8)
656+splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
657+
658+ where
659+
660+ {-# INLINE_LATE step' #-}
661+ step' gst (Initial st) = do
662+ r <- step gst st
663+ case r of
664+ D.Yield arr s -> 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
671+
672+ step' gst (Buffering st buf) = do
673+ r <- step gst st
674+ case r of
675+ D.Yield arr s -> 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)
682+ D.Stop -> return $
683+ if byteLength buf == 0
684+ then D.Stop
685+ else D.Skip (Yielding buf Finishing)
686+
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)
692+
693+ step' _ (Yielding arr next) = return $ D.Yield arr next
694+ step' _ Finishing = return $ D.Stop