@@ -22,6 +22,10 @@ module Streamly.Internal.Data.Array.Stream.Foreign
2222 , unlines
2323
2424 -- * Elimination
25+ , fold
26+ -- , parse
27+ , parseD
28+ , foldMany
2529 , toArray
2630
2731 -- * Compaction
@@ -39,30 +43,40 @@ where
3943
4044#include "inline.hs"
4145
46+import Control.Exception (assert)
47+import Control.Monad.Catch (MonadThrow, throwM)
4248import Control.Monad.IO.Class (MonadIO(..))
4349import Data.Word (Word8)
4450import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr)
4551import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
4652import Foreign.Ptr (minusPtr, plusPtr, castPtr)
4753import Foreign.Storable (Storable(..))
54+import Fusion.Plugin.Types (Fuse(..))
55+import GHC.Exts (SpecConstrAnnotation(..))
4856import GHC.ForeignPtr (ForeignPtr(..))
4957import GHC.Ptr (Ptr(..))
50−import Prelude hiding (length, null, last, map, (!!), read, concat, unlines)
58+import GHC.Types (SPEC(..))
59+import Prelude hiding (null, last, (!!), read, concat, unlines)
5160
5261#if !defined(mingw32_HOST_OS)
5362import Streamly.Internal.FileSystem.FDIO (IOVec(..))
5463#endif
55−import Streamly.Internal.Data.Array.Foreign.Type (Array(..), length)
64+import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
5665import Streamly.Internal.Data.Fold.Type (Fold(..))
66+import Streamly.Internal.Data.Parser (ParseError(..))
5767import Streamly.Internal.Data.Stream.Serial (SerialT)
5868import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
59−import Streamly.Internal.Data.SVar (adaptState)
69+import Streamly.Internal.Data.SVar (adaptState, defState)
6070
6171import qualified Streamly.Internal.Data.Array.Foreign as A
72+import qualified Streamly.Internal.Data.Array.Foreign as Array
6273import qualified Streamly.Internal.Data.Array.Foreign.Type as A
6374import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
6475import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
76+import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign as ASF
6577import qualified Streamly.Internal.Data.Fold.Type as FL
78+import qualified Streamly.Internal.Data.Parser as PR
79+import qualified Streamly.Internal.Data.Parser.ParserD as PRD
6680import qualified Streamly.Internal.Data.Stream.IsStream as S
6781import qualified Streamly.Internal.Data.Stream.StreamD as D
6882
@@ -300,6 +314,56 @@ splitOnSuffix byte s =
300314-- Elimination
301315-------------------------------------------------------------------------------
302316
317+-- When we have to take an array partially, take the last part of the array.
318+{-# INLINE takeArrayListRev #-}
319+takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
320+takeArrayListRev = go
321+
322+ where
323+
324+ go _ [] = []
325+ go n _ | n <= 0 = []
326+ go n (x:xs) =
327+ let len = Array.length x
328+ in if n > len
329+ then x : go (n - len) xs
330+ else if n == len
331+ then [x]
332+ else let !(Array (ForeignPtr _ contents) end) = x
333+ sz = sizeOf (undefined :: a)
334+ !(Ptr start) = end `plusPtr` negate (n * sz)
335+ in [Array (ForeignPtr start contents) end]
336+
337+-- When we have to take an array partially, take the last part of the array in
338+-- the first split.
339+{-# INLINE splitAtArrayListRev #-}
340+splitAtArrayListRev :: forall a. Storable a =>
341+ Int -> [Array a] -> ([Array a],[Array a])
342+splitAtArrayListRev n ls
343+ | n <= 0 = ([], ls)
344+ | otherwise = go n ls
345+ where
346+ go :: Int -> [Array a] -> ([Array a], [Array a])
347+ go _ [] = ([], [])
348+ go m (x:xs) =
349+ let len = Array.length x
350+ (xs', xs'') = go (m - len) xs
351+ in if m > len
352+ then (x:xs', xs'')
353+ else if m == len
354+ then ([x],xs)
355+ else let !(Array (ForeignPtr start contents) end) = x
356+ sz = sizeOf (undefined :: a)
357+ end1 = end `plusPtr` negate (m * sz)
358+ arr2 = Array (ForeignPtr start contents) end1
359+ !(Ptr addrEnd1) = end1
360+ arr1 = Array (ForeignPtr addrEnd1 contents) end
361+ in ([arr1], arr2:xs)
362+
363+-------------------------------------------------------------------------------
364+-- Fold to a single Array
365+-------------------------------------------------------------------------------
366+
303367-- XXX Both of these implementations of splicing seem to perform equally well.
304368-- We need to perform benchmarks over a range of sizes though.
305369
@@ -326,7 +390,7 @@ _spliceArrays :: (MonadIO m, Storable a)
326390 => SerialT m (Array a) -> m (Array a)
327391_spliceArrays s = do
328392 buffered <- S.foldr S.cons S.nil s
329− len <- S.sum (S.map length buffered)
393+ len <- S.sum (S.map Array.length buffered)
330394 arr <- liftIO $ MA.newArray len
331395 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
332396 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
@@ -344,7 +408,7 @@ _spliceArraysBuffered :: (MonadIO m, Storable a)
344408 => SerialT m (Array a) -> m (Array a)
345409_spliceArraysBuffered s = do
346410 buffered <- S.foldr S.cons S.nil s
347− len <- S.sum (S.map length buffered)
411+ len <- S.sum (S.map Array.length buffered)
348412 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
349413
350414{-# INLINE spliceArraysRealloced #-}
@@ -357,6 +421,8 @@ spliceArraysRealloced s = do
357421 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
358422 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
359423
424+-- XXX This should just be "fold A.write"
425+--
360426-- | Given a stream of arrays, splice them all together to generate a single
361427-- array. The stream must be /finite/.
362428--
@@ -385,3 +451,242 @@ _toArraysOf :: (MonadIO m, Storable a)
385451 => Int -> Fold m a (SerialT Identity (Array a))
386452_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
387453-}
454+
455+-------------------------------------------------------------------------------
456+-- Parsing
457+-------------------------------------------------------------------------------
458+
459+-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
460+-- to make a newtype.
461+{-# ANN type List NoSpecConstr #-}
462+newtype List a = List {getList :: [a]}
463+
464+{-# INLINE_NORMAL parseD #-}
465+parseD ::
466+ forall m a b. (MonadIO m, MonadThrow m, Storable a)
467+ => PRD.Parser m (Array a) b
468+ -> D.Stream m (Array.Array a)
469+ -> m (b, D.Stream m (Array.Array a))
470+parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
471+ res <- initial
472+ case res of
473+ PRD.IPartial s -> go SPEC state (List []) s
474+ PRD.IDone b -> return (b, stream)
475+ PRD.IError err -> throwM $ ParseError err
476+
477+ where
478+
479+ -- "backBuf" contains last few items in the stream that we may have to
480+ -- backtrack to.
481+ --
482+ -- XXX currently we are using a dumb list based approach for backtracking
483+ -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
484+ -- That will allow us more efficient random back and forth movement.
485+ {-# INLINE go #-}
486+ go !_ st backBuf !pst = do
487+ r <- step defState st
488+ case r of
489+ D.Yield x s -> gobuf SPEC [x] s backBuf pst
490+ D.Skip s -> go SPEC s backBuf pst
491+ D.Stop -> do
492+ b <- extract pst
493+ return (b, D.nil)
494+
495+ gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
496+ gobuf !_ (x:xs) s backBuf !pst = do
497+ pRes <- pstep pst x
498+ case pRes of
499+ PR.Partial 0 pst1 ->
500+ gobuf SPEC xs s (List []) pst1
501+ PR.Partial n pst1 -> do
502+ assert
503+ (n <= sum (map Array.length (x:getList backBuf)))
504+ (return ())
505+ let src0 = takeArrayListRev n (x:getList backBuf)
506+ src = Prelude.reverse src0 ++ xs
507+ gobuf SPEC src s (List []) pst1
508+ PR.Continue 0 pst1 ->
509+ gobuf SPEC xs s (List (x:getList backBuf)) pst1
510+ PR.Continue n pst1 -> do
511+ assert
512+ (n <= sum (map Array.length (x:getList backBuf)))
513+ (return ())
514+ let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
515+ src = Prelude.reverse src0 ++ xs
516+ gobuf SPEC src s (List buf1) pst1
517+ PR.Done 0 b ->
518+ return (b, D.Stream step s)
519+ PR.Done n b -> do
520+ assert
521+ (n <= sum (map Array.length (x:getList backBuf)))
522+ (return ())
523+ let src0 = takeArrayListRev n (x:getList backBuf)
524+ src = Prelude.reverse src0 ++ xs
525+ return (b, D.append (D.fromList src) (D.Stream step s))
526+ PR.Error err -> throwM $ ParseError err
527+
528+{-
529+-- | Parse an array stream using the supplied 'Parser'. Returns the parse
530+-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
531+--
532+-- /Internal/
533+--
534+{-# INLINE parse #-}
535+parse ::
536+ (MonadIO m, MonadThrow m, Storable a)
537+ => PRD.Parser m a b
538+ -> SerialT m (A.Array a)
539+ -> m (b, SerialT m (A.Array a))
540+parse p s = fmap D.fromStreamD <$> parseD p (D.toStreamD s)
541+-}
542+
543+-- | Fold an array stream using the supplied array stream 'Fold'.
544+--
545+-- /Pre-release/
546+--
547+{-# INLINE fold #-}
548+fold :: (MonadIO m, MonadThrow m, Storable a) =>
549+ ASF.Fold m a b -> SerialT m (A.Array a) -> m b
550+fold (ASF.Fold p) s = fst <$> parseD p (D.toStreamD s)
551+
552+{-# ANN type ParseChunksState Fuse #-}
553+data ParseChunksState x inpBuf st pst =
554+ ParseChunksInit inpBuf st
555+ | ParseChunksInitLeftOver inpBuf
556+ | ParseChunksStream st inpBuf !pst
557+ | ParseChunksBuf inpBuf st inpBuf !pst
558+ | ParseChunksYield x (ParseChunksState x inpBuf st pst)
559+
560+{-# INLINE_NORMAL foldManyD #-}
561+foldManyD
562+ :: (MonadThrow m, Storable a)
563+ => ASF.Fold m a b
564+ -> D.Stream m (Array a)
565+ -> D.Stream m b
566+foldManyD (ASF.Fold (PRD.Parser pstep initial extract)) (D.Stream step state) =
567+ D.Stream stepOuter (ParseChunksInit [] state)
568+
569+ where
570+
571+ {-# INLINE_LATE stepOuter #-}
572+ -- Buffer is empty, get the first element from the stream, initialize the
573+ -- fold and then go to stream processing loop.
574+ stepOuter gst (ParseChunksInit [] st) = do
575+ r <- step (adaptState gst) st
576+ case r of
577+ D.Yield x s -> do
578+ res <- initial
579+ case res of
580+ PRD.IPartial ps ->
581+ return $ D.Skip $ ParseChunksBuf [x] s [] ps
582+ PRD.IDone pb ->
583+ let next = ParseChunksInit [x] s
584+ in return $ D.Skip $ ParseChunksYield pb next
585+ PRD.IError err -> throwM $ ParseError err
586+ D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
587+ D.Stop -> return D.Stop
588+
589+ -- Buffer is not empty, go to buffered processing loop
590+ stepOuter _ (ParseChunksInit src st) = do
591+ res <- initial
592+ case res of
593+ PRD.IPartial ps ->
594+ return $ D.Skip $ ParseChunksBuf src st [] ps
595+ PRD.IDone pb ->
596+ let next = ParseChunksInit src st
597+ in return $ D.Skip $ ParseChunksYield pb next
598+ PRD.IError err -> throwM $ ParseError err
599+
600+ -- XXX we just discard any leftover input at the end
601+ stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
602+
603+ -- Buffer is empty, process elements from the stream
604+ stepOuter gst (ParseChunksStream st backBuf pst) = do
605+ r <- step (adaptState gst) st
606+ case r of
607+ D.Yield x s -> do
608+ pRes <- pstep pst x
609+ case pRes of
610+ PR.Partial 0 pst1 ->
611+ return $ D.Skip $ ParseChunksStream s [] pst1
612+ PR.Partial n pst1 -> do
613+ assert
614+ (n <= sum (map Array.length (x:backBuf)))
615+ (return ())
616+ let src0 = takeArrayListRev n (x:backBuf)
617+ src = Prelude.reverse src0
618+ return $ D.Skip $ ParseChunksBuf src s [] pst1
619+ PR.Continue 0 pst1 ->
620+ return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
621+ PR.Continue n pst1 -> do
622+ assert
623+ (n <= sum (map Array.length (x:backBuf)))
624+ (return ())
625+ let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
626+ src = Prelude.reverse src0
627+ return $ D.Skip $ ParseChunksBuf src s buf1 pst1
628+ PR.Done 0 b -> do
629+ return $ D.Skip $
630+ ParseChunksYield b (ParseChunksInit [] s)
631+ PR.Done n b -> do
632+ assert
633+ (n <= sum (map Array.length (x:backBuf)))
634+ (return ())
635+ let src0 = takeArrayListRev n (x:backBuf)
636+ let src = Prelude.reverse src0
637+ return $ D.Skip $
638+ ParseChunksYield b (ParseChunksInit src s)
639+ PR.Error err -> throwM $ ParseError err
640+ D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
641+ D.Stop -> do
642+ b <- extract pst
643+ let src = Prelude.reverse backBuf
644+ return $ D.Skip $
645+ ParseChunksYield b (ParseChunksInitLeftOver src)
646+
647+ -- go back to stream processing mode
648+ stepOuter _ (ParseChunksBuf [] s buf pst) =
649+ return $ D.Skip $ ParseChunksStream s buf pst
650+
651+ -- buffered processing loop
652+ stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
653+ pRes <- pstep pst x
654+ case pRes of
655+ PR.Partial 0 pst1 ->
656+ return $ D.Skip $ ParseChunksBuf xs s [] pst1
657+ PR.Partial n pst1 -> do
658+ assert (n <= sum (map Array.length (x:backBuf))) (return ())
659+ let src0 = takeArrayListRev n (x:backBuf)
660+ src = Prelude.reverse src0 ++ xs
661+ return $ D.Skip $ ParseChunksBuf src s [] pst1
662+ PR.Continue 0 pst1 ->
663+ return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
664+ PR.Continue n pst1 -> do
665+ assert (n <= sum (map Array.length (x:backBuf))) (return ())
666+ let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
667+ src = Prelude.reverse src0 ++ xs
668+ return $ D.Skip $ ParseChunksBuf src s buf1 pst1
669+ PR.Done 0 b ->
670+ return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
671+ PR.Done n b -> do
672+ assert (n <= sum (map Array.length (x:backBuf))) (return ())
673+ let src0 = takeArrayListRev n (x:backBuf)
674+ src = Prelude.reverse src0 ++ xs
675+ return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
676+ PR.Error err -> throwM $ ParseError err
677+
678+ stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
679+
680+-- | Apply an array stream 'Fold' repeatedly on an array stream and emit the
681+-- fold outputs in the output stream.
682+--
683+-- See "Streamly.Prelude.foldMany" for more details.
684+--
685+-- /Pre-release/
686+{-# INLINE foldMany #-}
687+foldMany
688+ :: (IsStream t, MonadThrow m, Storable a)
689+ => ASF.Fold m a b
690+ -> t m (Array a)
691+ -> t m b
692+foldMany p m = D.fromStreamD $ foldManyD p (D.toStreamD m)