@@ -1,706 +0,0 @@
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.Bifunctor (first)
12−import Data.Primitive.Types (Prim(..), sizeOf)
13−#if __GLASGOW_HASKELL__ < 808
14−import Data.Semigroup (Semigroup(..))
15−#endif
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)
23−
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
29−
30−import GHC.Exts hiding (fromListN, fromList, toList)
31−import Prelude hiding (length, unlines, foldr)
32−
33−-------------------------------------------------------------------------------
34−-- Array Data Type
35−-------------------------------------------------------------------------------
36−
37−data Array a = Array ByteArray# Int Int
38−
39−-------------------------------------------------------------------------------
40−-- Utilities
41−-------------------------------------------------------------------------------
42−
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 #-}
47−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
54− -> m ()
55−unsafeCopy (MA.Array dst#) (I# doff#) (Array src# (I# off#) _) (I# soff#) (I# n#) =
56− liftIO $ do
57− let toBytes cnt# = cnt# *# (sizeOf# (undefined :: a))
58− primitive_ $
59− copyByteArray#
60− src#
61− (toBytes (off# +# soff#))
62− dst#
63− (toBytes doff#)
64− (toBytes n#)
65−
66−-------------------------------------------------------------------------------
67−-- Basic Byte Array Operations
68−-------------------------------------------------------------------------------
69−
70−{-# INLINE unsafeFreeze #-}
71−unsafeFreeze :: (Prim a, MonadIO m) => MA.Array a -> m (Array a)
72−unsafeFreeze marr@(MA.Array arr#) =
73− liftIO $ do
74− len <- MA.length marr
75− primitive $ \s# ->
76− case unsafeFreezeByteArray# arr# s# of
77− (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
78−
79−{-# INLINE unsafeFreezeWithShrink #-}
80−unsafeFreezeWithShrink ::
81− (Prim a, MonadIO m) => MA.Array a -> Int -> m (Array a)
82−unsafeFreezeWithShrink marr@(MA.Array arr#) n =
83− liftIO $ do
84− MA.shrinkArray marr n
85− len <- MA.length marr
86− primitive $ \s# ->
87− case unsafeFreezeByteArray# arr# s# of
88− (# s1#, arr1# #) -> (# s1#, Array arr1# 0 len #)
89−
90−{-
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#) #)
96−-}
97−
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#)
102−
103−-- unsafe
104−sameByteArray :: ByteArray# -> ByteArray# -> Bool
105−sameByteArray ba1 ba2 =
106− case reallyUnsafePtrEquality#
107− (unsafeCoerce# ba1 :: ())
108− (unsafeCoerce# ba2 :: ()) of
109− r -> isTrue# r
110−
111−-------------------------------------------------------------------------------
112−-- Chunk Size
113−-------------------------------------------------------------------------------
114−
115−-- XXX move this section to mutable array module?
116−
117−mkChunkSizeKB :: Int -> Int
118−mkChunkSizeKB n = n * k
119− where k = 1024
120−
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
126−
127−-------------------------------------------------------------------------------
128−-- Length
129−-------------------------------------------------------------------------------
130−
131−sizeOfPrimElem :: forall a . (Prim a) => a -> Int
132−sizeOfPrimElem _ = max 1 $ sizeOf (undefined :: a)
133−
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)
138−
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
145−
146−-------------------------------------------------------------------------------
147−-- Construction
148−-------------------------------------------------------------------------------
149−
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
154−
155−nil :: Prim a => Array a
156−nil =
157− unsafePerformIO $ do
158− arr <- MA.newArray 0
159− unsafeFreeze arr
160−
161−-- | Fold the whole input to a single array.
162−--
163−-- /Caution! Do not use this on infinite streams./
164−--
165−-- /Pre-release/
166−{-# INLINE_NORMAL write #-}
167−write :: (MonadIO m, Prim a) => Fold m a (Array a)
168−write = FL.rmapM unsafeFreeze MA.write
169−
170−-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
171−-- 'Array'.
172−--
173−-- /Pre-release/
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)
177−
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
182−-- slowdown.
183−--
184−-- /Pre-release/
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)
188−
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
192−
193−{-# INLINE fromStreamD #-}
194−fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (Array a)
195−fromStreamD str = MA.fromStreamD str >>= unsafeFreeze
196−
197−-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
198−-- arrays of size n.
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)
203−
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
208−
209−-- XXX derive from MA.fromList?
210−{-# INLINE fromList #-}
211−fromList :: Prim a => [a] -> Array a
212−fromList xs = fromListN (P.length xs) xs
213−
214−-------------------------------------------------------------------------------
215−-- Combining
216−-------------------------------------------------------------------------------
217−
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)
221−spliceTwo a1 a2 = do
222− let l1 = length a1
223− l2 = length a2
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`?
228−
229−-------------------------------------------------------------------------------
230−-- Elimination
231−-------------------------------------------------------------------------------
232−
233−{-# INLINE_LATE toListFB #-}
234−toListFB :: Prim a => (a -> b -> b) -> b -> Array a -> b
235−toListFB c n arr = go 0
236− where
237− len = length arr
238− go p | p == len = n
239− go p =
240− let !x = unsafeIndex arr p
241− in c x (go (p + 1))
242−
243−-- | Convert an 'Array' into a list.
244−--
245−-- /Pre-release/
246−{-# INLINE toList #-}
247−toList :: Prim a => Array a -> [a]
248−toList s = build (\c n -> toListFB c n s)
249−
250−-------------------------------------------------------------------------------
251−-- Instances
252−-------------------------------------------------------------------------------
253−
254−instance (Eq a, Prim a) => Eq (Array a) where
255− {-# INLINE (==) #-}
256− a1@(Array ba1# _ len1) == a2@(Array ba2# _ len2)
257− | sameByteArray ba1# ba2# = True
258− | len1 /= len2 = False
259− | otherwise = loop (len1 - 1)
260−
261− where
262−
263− loop !i
264− | i < 0 = True
265− | otherwise = unsafeIndex a1 i == unsafeIndex a2 i && loop (i - 1)
266−
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
273−
274− where
275−
276− sz = min len1 len2
277−
278− loop !i
279− | i < sz =
280− compare (unsafeIndex a1 i) (unsafeIndex a2 i) <> loop (i + 1)
281− | otherwise = compare len1 len2
282−
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))
287−
288−instance Prim a => Monoid (Array a) where
289− mempty = nil
290− mappend = (<>)
291−
292−instance NFData (Array a) where
293− {-# INLINE rnf #-}
294− rnf _ = ()
295−
296−
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
301− showsPrec p a =
302− showParen (p > 10) $
303− showString "fromListN "
304− . shows (length a)
305− . showString " "
306− . shows (toList a)
307−
308−instance (a ~ Char) => IsString (Array a) where
309− {-# INLINE fromString #-}
310− fromString = fromList
311−
312−-- GHC versions 8.0 and below cannot derive IsList
313−instance Prim a => IsList (Array a) where
314− type (Item (Array a)) = a
315−
316− {-# INLINE fromList #-}
317− fromList = fromList
318−
319− {-# INLINE fromListN #-}
320− fromListN = fromListN
321−
322− {-# INLINE toList #-}
323− toList = toList
324−
325−instance (Prim a, Read a, Show a) => Read (Array a) where
326− {-# INLINE readPrec #-}
327− readPrec = fromList <$> readPrec
328− readListPrec = readListPrecDefault
329−
330−-- XXX these folds can be made common with mutable arrays by defining a
331−-- unsafeIndex in the specific module?
332−
333−-------------------------------------------------------------------------------
334−-- Folds
335−-------------------------------------------------------------------------------
336−
337−{-# INLINE foldr #-}
338−foldr :: Prim a => (a -> b -> b) -> b -> Array a -> b
339−foldr f z arr = go 0
340−
341− where
342−
343− !len = length arr
344−
345− go !i
346− | len > i = f (unsafeIndex arr i) (go (i + 1))
347− | otherwise = z
348−
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
353−
354− where
355−
356− go !i !acc
357− | i < 0 = acc
358− | otherwise = go (i - 1) (f (unsafeIndex arr i) acc)
359−
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
364−
365− where
366−
367− !len = length arr
368−
369− go !i !acc
370− | i < len = go (i + 1) (f acc (unsafeIndex arr i))
371− | otherwise = acc
372−
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
377−
378− where
379−
380− !len = length arr
381−
382− go !i !acc1
383− | i < len = do
384− acc2 <- f acc1 (unsafeIndex arr i)
385− go (i + 1) acc2
386− | otherwise = return acc1
387−
388−-------------------------------------------------------------------------------
389−-- Converting to streams
390−-------------------------------------------------------------------------------
391−
392−{-# INLINE_NORMAL toStreamD #-}
393−toStreamD :: (Prim a, Monad m) => Array a -> D.Stream m a
394−toStreamD arr = D.Stream step 0
395−
396− where
397−
398− {-# INLINE_LATE step #-}
399− step _ i
400− | i == length arr = return D.Stop
401− step _ i = return $ D.Yield (unsafeIndex arr i) (i + 1)
402−
403−{-# INLINE toStreamK #-}
404−toStreamK :: Prim a => Array a -> K.Stream m a
405−toStreamK arr = go 0
406−
407− where
408−
409− len = length arr
410−
411− go p
412− | p == len = K.nil
413− | otherwise =
414− let !x = unsafeIndex arr p
415− in x `K.cons` go (p + 1)
416−
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)
420−
421− where
422−
423− {-# INLINE_LATE step #-}
424− step _ i
425− | i < 0 = return D.Stop
426− step _ i = return $ D.Yield (unsafeIndex arr i) (i - 1)
427−
428−{-# INLINE toStreamKRev #-}
429−toStreamKRev :: Prim a => Array a -> K.Stream m a
430−toStreamKRev arr = go (length arr - 1)
431−
432− where
433−
434− go p | p == -1 = K.nil
435− | otherwise =
436− let !x = unsafeIndex arr p
437− in x `K.cons` go (p - 1)
438−
439−-------------------------------------------------------------------------------
440−-- Stream of Arrays
441−-------------------------------------------------------------------------------
442−
443−data FlattenState s a =
444− OuterLoop s
445− | InnerLoop s !(Array a) !Int !Int
446−
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)
450−
451− where
452−
453− {-# INLINE_LATE step' #-}
454− step' gst (OuterLoop st) = do
455− r <- step (adaptState gst) st
456− return $ case r of
457− D.Yield arr s ->
458− let len = length arr
459− in if len == 0
460− then D.Skip (OuterLoop s)
461− else D.Skip (InnerLoop s arr len 0)
462− D.Skip s -> D.Skip (OuterLoop s)
463− D.Stop -> D.Stop
464−
465− step' _ (InnerLoop st _ len i) | i == len =
466− return $ D.Skip $ OuterLoop st
467−
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))
471−
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)
475−
476− where
477−
478− {-# INLINE_LATE step' #-}
479− step' gst (OuterLoop st) = do
480− r <- step (adaptState gst) st
481− return $ case r of
482− D.Yield arr s ->
483− let len = length arr
484− in if len == 0
485− then D.Skip (OuterLoop s)
486− else D.Skip (InnerLoop s arr len (len - 1))
487− D.Skip s -> D.Skip (OuterLoop s)
488− D.Stop -> D.Stop
489−
490− step' _ (InnerLoop st _ _ i) | i == -1 =
491− return $ D.Skip $ OuterLoop st
492−
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))
496−
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)
500−
501− where
502−
503− {-# INLINE_LATE step' #-}
504− step' gst (OuterLoop st) = do
505− r <- step (adaptState gst) st
506− return $ case r of
507− D.Yield arr s ->
508− let len = length arr
509− in D.Skip (InnerLoop s arr len 0)
510− D.Skip s -> D.Skip (OuterLoop s)
511− D.Stop -> D.Stop
512−
513− step' _ (InnerLoop st _ len i) | i == len =
514− return $ D.Yield sep $ OuterLoop st
515−
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))
519−
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
526− return $ doff + len
527−
528−data SpliceState s arr1 arr2
529− = SpliceInitial s
530− | SpliceBuffering s arr2
531− | SpliceYielding arr1 (SpliceState s arr1 arr2)
532− | SpliceFinish
533−
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.
539−--
540−-- /Pre-release/
541−{-# INLINE_NORMAL packArraysChunksOf #-}
542−packArraysChunksOf ::
543− forall m a. (MonadIO m, Prim a)
544− => Int
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)
549−
550− where
551−
552− nElem = n `quot` sizeOfPrimElem (undefined :: a)
553−
554− {-# INLINE_LATE step' #-}
555− step' gst (SpliceInitial st) = do
556− when (n <= 0) $
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"
560− r <- step gst st
561− case r of
562− D.Yield arr s ->
563− if length arr >= nElem
564− then return $ D.Skip (SpliceYielding arr (SpliceInitial s))
565− else do
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
571−
572− step' gst (SpliceBuffering st arr2@(buf, boff)) = do
573− r <- step gst st
574− case r of
575− D.Yield arr s -> do
576− if boff + length arr > nElem
577− then do
578− nArr <- unsafeFreeze buf
579− return $ D.Skip (SpliceYielding (slice nArr 0 boff) (SpliceBuffering s arr2))
580− else do
581− noff <- spliceInto buf boff arr
582− return $ D.Skip (SpliceBuffering s (buf, noff))
583− D.Skip s -> return $ D.Skip (SpliceBuffering s arr2)
584− D.Stop -> do
585− nArr <- unsafeFreeze buf
586− return $ D.Skip (SpliceYielding (slice nArr 0 boff) SpliceFinish)
587−
588− step' _ SpliceFinish = return D.Stop
589−
590− step' _ (SpliceYielding arr next) = return $ D.Yield arr next
591−
592−{-# INLINE_NORMAL lpackArraysChunksOf #-}
593−lpackArraysChunksOf ::
594− forall m a. (MonadIO m, Prim a)
595− => Int
596− -> Fold m (Array a) ()
597− -> Fold m (Array a) ()
598−lpackArraysChunksOf n (Fold step1 initial1 extract1) =
599− Fold step initial extract
600−
601− where
602−
603− nElem = n `quot` sizeOfPrimElem (undefined :: a)
604−
605− initial = do
606− when (n <= 0) $
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"
610− res <- initial1
611− return $ first (Tuple3' Nothing' 0) res
612−
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)
617− case r of
618− FL.Partial rr -> extract1 rr
619− FL.Done _ -> return ()
620−
621− step (Tuple3' Nothing' _ r1) arr =
622− if length arr >= nElem
623− then do
624− r <- step1 r1 arr
625− case r of
626− FL.Done _ -> return $ FL.Done ()
627− FL.Partial s -> do
628− extract1 s
629− res <- initial1
630− return $ first (Tuple3' Nothing' 0) res
631− else do
632− buf <- MA.newArray nElem
633− noff <- spliceInto buf 0 arr
634− return $ FL.Partial $ Tuple3' (Just' buf) noff r1
635−
636− step (Tuple3' (Just' buf) boff r1) arr = do
637− noff <- spliceInto buf boff arr
638−
639− if noff >= nElem
640− then do
641− nArr <- unsafeFreeze buf
642− r <- step1 r1 (slice nArr 0 noff)
643− case r of
644− FL.Done _ -> return $ FL.Done ()
645− FL.Partial s -> do
646− extract1 s
647− res <- initial1
648− return $ first (Tuple3' Nothing' 0) res
649− else return $ FL.Partial $ Tuple3' (Just' buf) noff r1
650−
651−data SplitState s arr
652− = Initial s
653− | Buffering s arr
654− | Splitting s arr
655− | Yielding arr (SplitState s arr)
656− | Finishing
657−
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.
660−--
661−-- /Pre-release/
662−{-# INLINE_NORMAL splitOn #-}
663−splitOn ::
664− MonadIO m
665− => Word8
666− -> D.Stream m (Array Word8)
667− -> D.Stream m (Array Word8)
668−splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
669−
670− where
671−
672− {-# INLINE_LATE step' #-}
673− step' gst (Initial st) = do
674− r <- step gst st
675− case r of
676− D.Yield arr s -> 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
683−
684− step' gst (Buffering st buf) = do
685− r <- step gst st
686− case r of
687− D.Yield arr s -> 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)
694− D.Stop -> return $
695− if byteLength buf == 0
696− then D.Stop
697− else D.Skip (Yielding buf Finishing)
698−
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)
704−
705− step' _ (Yielding arr next) = return $ D.Yield arr next
706− step' _ Finishing = return $ D.Stop