@@ -22,6 +22,11 @@ module Streamly.Internal.Data.Array.Stream.Foreign
2222 , unlines
2323
2424 -- * Elimination
25+ -- ** Element Folds
26+ , fold
27+ , parse
28+ , parseD
29+
2530 -- ** Array Folds
2631 , foldArr
2732 , foldArr_
@@ -50,6 +55,9 @@ import Data.Bifunctor (second)
5055import Control.Exception (assert)
5156import Control.Monad.Catch (MonadThrow, throwM)
5257import Control.Monad.IO.Class (MonadIO(..))
58+#if __GLASGOW_HASKELL__ < 808
59+import Data.Semigroup (Semigroup(..))
60+#endif
5361import Data.Word (Word8)
5462import Foreign.ForeignPtr (touchForeignPtr)
5563import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
@@ -317,7 +325,65 @@ splitOnSuffix byte s =
317325 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.toStreamD s
318326
319327-------------------------------------------------------------------------------
320−-- Elimination
328+-- Elimination - Running folds
329+-------------------------------------------------------------------------------
330+
331+-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
332+-- to the number of times it can be called on the same stream.
333+--
334+{-# INLINE_NORMAL foldD #-}
335+foldD :: forall m a b. (MonadIO m, Storable a) =>
336+ Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
337+foldD (Fold fstep initial extract) stream@(D.Stream step state) = do
338+ res <- initial
339+ case res of
340+ FL.Partial fs -> go SPEC state fs
341+ FL.Done fb -> return $! (fb, stream)
342+
343+ where
344+
345+ {-# INLINE go #-}
346+ go !_ st !fs = do
347+ r <- step defState st
348+ case r of
349+ D.Yield (Array (ForeignPtr start contents) (Ptr end)) s ->
350+ goArray SPEC s (ForeignPtr end contents) (Ptr start) fs
351+ D.Skip s -> go SPEC s fs
352+ D.Stop -> do
353+ b <- extract fs
354+ return (b, D.nil)
355+
356+ goArray !_ s fp@(ForeignPtr end _) !cur !fs
357+ | cur == Ptr end = do
358+ liftIO $ touchForeignPtr fp
359+ go SPEC s fs
360+ goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
361+ x <- liftIO $ peek cur
362+ res <- fstep fs x
363+ let elemSize = sizeOf (undefined :: a)
364+ next = cur `plusPtr` elemSize
365+ case res of
366+ FL.Done b -> do
367+ let !(Ptr curAddr) = next
368+ arr = Array (ForeignPtr curAddr contents) (Ptr end)
369+ return $! (b, D.cons arr (D.Stream step st))
370+ FL.Partial fs1 -> goArray SPEC st fp next fs1
371+
372+-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
373+-- and the unconsumed stream.
374+--
375+-- /Internal/
376+--
377+{-# INLINE_NORMAL fold #-}
378+fold ::
379+ (MonadIO m, Storable a)
380+ => FL.Fold m a b
381+ -> SerialT m (A.Array a)
382+ -> m (b, SerialT m (A.Array a))
383+fold f s = fmap D.fromStreamD <$> foldD f (D.toStreamD s)
384+
385+-------------------------------------------------------------------------------
386+-- Fold to a single Array
321387-------------------------------------------------------------------------------
322388
323389-- When we have to take an array partially, take the last part of the array.
@@ -459,7 +525,7 @@ _toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
459525-}
460526
461527-------------------------------------------------------------------------------
462−-- Parsing
528+-- Elimination - running element parsers
463529-------------------------------------------------------------------------------
464530
465531-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
@@ -467,6 +533,110 @@ _toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
467533{-# ANN type List NoSpecConstr #-}
468534newtype List a = List {getList :: [a]}
469535
536+-- This can be generalized to any type provided it can be unfolded to a stream
537+-- and it can be combined using a semigroup operation.
538+--
539+-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
540+-- to the number of times it can be called on the same stream.
541+{-# INLINE_NORMAL parseD #-}
542+parseD ::
543+ forall m a b. (MonadIO m, MonadThrow m, Storable a)
544+ => PRD.Parser m a b
545+ -> D.Stream m (Array.Array a)
546+ -> m (b, D.Stream m (Array.Array a))
547+parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
548+ res <- initial
549+ case res of
550+ PRD.IPartial s -> go SPEC state (List []) s
551+ PRD.IDone b -> return (b, stream)
552+ PRD.IError err -> throwM $ ParseError err
553+
554+ where
555+
556+ -- "backBuf" contains last few items in the stream that we may have to
557+ -- backtrack to.
558+ --
559+ -- XXX currently we are using a dumb list based approach for backtracking
560+ -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
561+ -- That will allow us more efficient random back and forth movement.
562+ {-# INLINE go #-}
563+ go !_ st backBuf !pst = do
564+ r <- step defState st
565+ case r of
566+ D.Yield (Array (ForeignPtr start contents) (Ptr end)) s ->
567+ gobuf SPEC s backBuf (ForeignPtr end contents) (Ptr start) pst
568+ D.Skip s -> go SPEC s backBuf pst
569+ D.Stop -> do
570+ b <- extract pst
571+ return (b, D.nil)
572+
573+ -- Use strictness on "cur" to keep it unboxed
574+ gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
575+ | cur == Ptr end = do
576+ liftIO $ touchForeignPtr fp
577+ go SPEC s backBuf pst
578+ gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
579+ x <- liftIO $ peek cur
580+ pRes <- pstep pst x
581+ let elemSize = sizeOf (undefined :: a)
582+ next = cur `plusPtr` elemSize
583+ case pRes of
584+ PR.Partial 0 pst1 ->
585+ gobuf SPEC s (List []) fp next pst1
586+ PR.Partial n pst1 -> do
587+ assert (n <= Prelude.length (x:getList backBuf)) (return ())
588+ let src0 = Prelude.take n (x:getList backBuf)
589+ arr0 = A.fromListN n (Prelude.reverse src0)
590+ !(Ptr curAddr) = next
591+ arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
592+ src = arr0 <> arr1
593+ let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
594+ gobuf SPEC s (List []) (ForeignPtr end1 cont1) (Ptr start) pst1
595+ PR.Continue 0 pst1 ->
596+ gobuf SPEC s (List (x:getList backBuf)) fp next pst1
597+ PR.Continue n pst1 -> do
598+ assert (n <= Prelude.length (x:getList backBuf)) (return ())
599+ let (src0, buf1) = splitAt n (x:getList backBuf)
600+ arr0 = A.fromListN n (Prelude.reverse src0)
601+ !(Ptr curAddr) = next
602+ arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
603+ src = arr0 <> arr1
604+ let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
605+ gobuf SPEC s (List buf1) (ForeignPtr end1 cont1) (Ptr start) pst1
606+ PR.Done 0 b -> do
607+ let !(Ptr curAddr) = next
608+ arr = Array (ForeignPtr curAddr contents) (Ptr end)
609+ return (b, D.cons arr (D.Stream step s))
610+ PR.Done n b -> do
611+ assert (n <= Prelude.length (x:getList backBuf)) (return ())
612+ let src0 = Prelude.take n (x:getList backBuf)
613+ -- XXX create the array in reverse instead
614+ arr0 = A.fromListN n (Prelude.reverse src0)
615+ !(Ptr curAddr) = next
616+ arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
617+ -- XXX Use StreamK to avoid adding arbitrary layers of
618+ -- constructors every time.
619+ str = D.cons arr0 (D.cons arr1 (D.Stream step s))
620+ return (b, str)
621+ PR.Error err -> throwM $ ParseError err
622+
623+-- | Parse an array stream using the supplied 'Parser'. Returns the parse
624+-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
625+--
626+-- /Internal/
627+--
628+{-# INLINE_NORMAL parse #-}
629+parse ::
630+ (MonadIO m, MonadThrow m, Storable a)
631+ => PRD.Parser m a b
632+ -> SerialT m (A.Array a)
633+ -> m (b, SerialT m (A.Array a))
634+parse p s = fmap D.fromStreamD <$> parseD p (D.toStreamD s)
635+
636+-------------------------------------------------------------------------------
637+-- Elimination - Running Array Folds and parsers
638+-------------------------------------------------------------------------------
639+
470640{-# INLINE_NORMAL parseArrD #-}
471641parseArrD ::
472642 forall m a b. (MonadIO m, MonadThrow m, Storable a)
@@ -540,7 +710,7 @@ parseArrD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
540710{-# INLINE parseArr #-}
541711parseArr ::
542712 (MonadIO m, MonadThrow m, Storable a)
543− => PRD.Parser m a b
713+ => ASF.Parser m a b
544714 -> SerialT m (A.Array a)
545715 -> m (b, SerialT m (A.Array a))
546716parseArr p s = fmap D.fromStreamD <$> parseD p (D.toStreamD s)