@@ -77,6 +77,11 @@ module Streamly.Internal.Data.Array
7777 , compactOnByte
7878 , compactOnByteSuffix
7979
80+ , foldBreakChunks
81+ , foldChunks
82+ , foldBreakChunksK
83+ , parseBreakChunksK
84+
8085 -- * Serialization
8186 , encodeAs
8287 , serialize
@@ -95,32 +100,42 @@ where
95100#include "ArrayMacros.h"
96101
97102import Control.Monad.IO.Class (MonadIO(..))
103+-- import Data.Bifunctor (first)
104+-- import Data.Either (fromRight)
98105import Data.Functor.Identity (Identity)
99106import Data.Proxy (Proxy(..))
100107import Data.Word (Word8)
101108import Foreign.C.String (CString)
102109import Foreign.Ptr (castPtr)
103110import Foreign.Storable (Storable)
111+import GHC.Types (SPEC(..))
104112import Streamly.Internal.Data.Unbox (Unbox(..))
105113import Prelude hiding (length, null, last, map, (!!), read, concat)
106114
107115import Streamly.Internal.Data.MutByteArray.Type (PinnedState(..), MutByteArray)
108116import Streamly.Internal.Data.Serialize.Type (Serialize)
109117import Streamly.Internal.Data.Fold.Type (Fold(..))
110−import Streamly.Internal.Data.Stream (Stream)
111−import Streamly.Internal.Data.SVar.Type (adaptState)
112−import Streamly.Internal.Data.Tuple.Strict (Tuple3Fused'(..))
118+import Streamly.Internal.Data.Parser (Parser(..), Initial(..), ParseError(..))
119+import Streamly.Internal.Data.Stream (Stream(..))
120+import Streamly.Internal.Data.StreamK (StreamK)
121+import Streamly.Internal.Data.SVar.Type (adaptState, defState)
122+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3Fused'(..))
113123import Streamly.Internal.Data.Unfold.Type (Unfold(..))
114124import Streamly.Internal.System.IO (unsafeInlineIO)
115125
126+import qualified Streamly.Internal.Data.Fold.Type as Fold
116127import qualified Streamly.Internal.Data.Serialize.Type as Serialize
117128import qualified Streamly.Internal.Data.MutByteArray.Type as MBA
118129import qualified Streamly.Internal.Data.MutArray as MA
119130import qualified Streamly.Internal.Data.Fold as FL
120131import qualified Streamly.Internal.Data.Ring as RB
132+import qualified Streamly.Internal.Data.Parser as Parser
133+-- import qualified Streamly.Internal.Data.ParserK as ParserK
121134import qualified Streamly.Internal.Data.Stream as D
122135import qualified Streamly.Internal.Data.Stream as Stream
136+import qualified Streamly.Internal.Data.StreamK as StreamK
123137import qualified Streamly.Internal.Data.Unfold as Unfold
138+import qualified Prelude
124139
125140import Streamly.Internal.Data.Array.Type
126141
@@ -631,3 +646,348 @@ compactOnByteSuffix
631646 -> Stream m (Array Word8)
632647compactOnByteSuffix byte =
633648 fmap unsafeFreeze . MA.compactOnByteSuffix byte . fmap unsafeThaw
649+
650+-------------------------------------------------------------------------------
651+-- Folding Streams of Arrays
652+-------------------------------------------------------------------------------
653+
654+-- XXX This should not be used for breaking a stream as the D.cons used in
655+-- reconstructing the stream could be very bad for performance. This can only
656+-- be useful in folding without breaking.
657+{-# INLINE_NORMAL foldBreakChunks #-}
658+foldBreakChunks :: forall m a b. (MonadIO m, Unbox a) =>
659+ Fold m a b -> Stream m (Array a) -> m (b, Stream m (Array a))
660+foldBreakChunks (Fold fstep initial _ final) stream@(Stream step state) = do
661+ res <- initial
662+ case res of
663+ Fold.Partial fs -> go SPEC state fs
664+ Fold.Done fb -> return $! (fb, stream)
665+
666+ where
667+
668+ {-# INLINE go #-}
669+ go !_ st !fs = do
670+ r <- step defState st
671+ case r of
672+ Stream.Yield (Array contents start end) s ->
673+ let fp = Tuple' end contents
674+ in goArray SPEC s fp start fs
675+ Stream.Skip s -> go SPEC s fs
676+ Stream.Stop -> do
677+ b <- final fs
678+ return (b, D.nil)
679+
680+ goArray !_ s (Tuple' end _) !cur !fs
681+ | cur == end = do
682+ go SPEC s fs
683+ goArray !_ st fp@(Tuple' end contents) !cur !fs = do
684+ x <- liftIO $ peekAt cur contents
685+ res <- fstep fs x
686+ let next = INDEX_NEXT(cur,a)
687+ case res of
688+ Fold.Done b -> do
689+ let arr = Array contents next end
690+ return $! (b, D.cons arr (D.Stream step st))
691+ Fold.Partial fs1 -> goArray SPEC st fp next fs1
692+
693+-- This may be more robust wrt fusion compared to unfoldMany?
694+
695+-- | Fold a stream of arrays using a 'Fold'. This is equivalent to the
696+-- following:
697+--
698+-- >>> foldChunks f = Stream.fold f . Stream.unfoldMany Array.reader
699+--
700+foldChunks :: (MonadIO m, Unbox a) => Fold m a b -> Stream m (Array a) -> m b
701+foldChunks f s = fmap fst (foldBreakChunks f s)
702+-- foldStream f = Stream.fold f . Stream.unfoldMany reader
703+
704+-- | Fold a stream of arrays using a 'Fold' and return the remaining stream.
705+--
706+-- The following alternative to this function allows composing the fold using
707+-- the parser Monad:
708+--
709+-- @
710+-- foldBreakStreamK f s =
711+-- fmap (first (fromRight undefined))
712+-- $ StreamK.parseBreakChunks (ParserK.adaptC (Parser.fromFold f)) s
713+-- @
714+--
715+-- We can compare perf and remove this one or define it in terms of that.
716+--
717+foldBreakChunksK :: forall m a b. (MonadIO m, Unbox a) =>
718+ Fold m a b -> StreamK m (Array a) -> m (b, StreamK m (Array a))
719+{-
720+foldBreakChunksK f s =
721+ fmap (first (fromRight undefined))
722+ $ StreamK.parseBreakChunks (ParserK.adaptC (Parser.fromFold f)) s
723+-}
724+foldBreakChunksK (Fold fstep initial _ final) stream = do
725+ res <- initial
726+ case res of
727+ Fold.Partial fs -> go fs stream
728+ Fold.Done fb -> return (fb, stream)
729+
730+ where
731+
732+ {-# INLINE go #-}
733+ go !fs st = do
734+ let stop = (, StreamK.nil) <$> final fs
735+ single a = yieldk a StreamK.nil
736+ yieldk (Array contents start end) r =
737+ let fp = Tuple' end contents
738+ in goArray fs r fp start
739+ in StreamK.foldStream defState yieldk single stop st
740+
741+ goArray !fs st (Tuple' end _) !cur
742+ | cur == end = do
743+ go fs st
744+ goArray !fs st fp@(Tuple' end contents) !cur = do
745+ x <- liftIO $ peekAt cur contents
746+ res <- fstep fs x
747+ let next = INDEX_NEXT(cur,a)
748+ case res of
749+ Fold.Done b -> do
750+ let arr = Array contents next end
751+ return $! (b, StreamK.cons arr st)
752+ Fold.Partial fs1 -> goArray fs1 st fp next
753+
754+{-
755+-- This can be generalized to any type provided it can be unfolded to a stream
756+-- and it can be combined using a semigroup operation.
757+--
758+{-# INLINE_NORMAL parseBreakD #-}
759+parseBreakD ::
760+ forall m a b. (MonadIO m, MonadThrow m, Unbox a)
761+ => PRD.Parser a m b
762+ -> D.Stream m (Array.Array a)
763+ -> m (b, D.Stream m (Array.Array a))
764+parseBreakD
765+ (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
766+
767+ res <- initial
768+ case res of
769+ PRD.IPartial s -> go SPEC state (List []) s
770+ PRD.IDone b -> return (b, stream)
771+ PRD.IError err -> throwM $ ParseError err
772+
773+ where
774+
775+ -- "backBuf" contains last few items in the stream that we may have to
776+ -- backtrack to.
777+ --
778+ -- XXX currently we are using a dumb list based approach for backtracking
779+ -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
780+ -- That will allow us more efficient random back and forth movement.
781+ go !_ st backBuf !pst = do
782+ r <- step defState st
783+ case r of
784+ D.Yield (Array contents start end) s ->
785+ gobuf SPEC s backBuf
786+ (Tuple' end contents) start pst
787+ D.Skip s -> go SPEC s backBuf pst
788+ D.Stop -> do
789+ b <- extract pst
790+ return (b, D.nil)
791+
792+ -- Use strictness on "cur" to keep it unboxed
793+ gobuf !_ s backBuf (Tuple' end _) !cur !pst
794+ | cur == end = do
795+ go SPEC s backBuf pst
796+ gobuf !_ s backBuf fp@(Tuple' end contents) !cur !pst = do
797+ x <- liftIO $ peekByteIndex contents cur
798+ pRes <- pstep pst x
799+ let next = INDEX_NEXT(cur,a)
800+ case pRes of
801+ PR.Partial 0 pst1 ->
802+ gobuf SPEC s (List []) fp next pst1
803+ PR.Partial n pst1 -> do
804+ assert (n <= Prelude.length (x:getList backBuf)) (return ())
805+ let src0 = Prelude.take n (x:getList backBuf)
806+ arr0 = A.fromListN n (Prelude.reverse src0)
807+ arr1 = Array contents next end
808+ src = arr0 <> arr1
809+ let !(Array cont1 start end1) = src
810+ fp1 = Tuple' end1 cont1
811+ gobuf SPEC s (List []) fp1 start pst1
812+ PR.Continue 0 pst1 ->
813+ gobuf SPEC s (List (x:getList backBuf)) fp next pst1
814+ PR.Continue n pst1 -> do
815+ assert (n <= Prelude.length (x:getList backBuf)) (return ())
816+ let (src0, buf1) = splitAt n (x:getList backBuf)
817+ arr0 = A.fromListN n (Prelude.reverse src0)
818+ arr1 = Array contents next end
819+ src = arr0 <> arr1
820+ let !(Array cont1 start end1) = src
821+ fp1 = Tuple' end1 cont1
822+ gobuf SPEC s (List buf1) fp1 start pst1
823+ PR.Done 0 b -> do
824+ let arr = Array contents next end
825+ return (b, D.cons arr (D.Stream step s))
826+ PR.Done n b -> do
827+ assert (n <= Prelude.length (x:getList backBuf)) (return ())
828+ let src0 = Prelude.take n (x:getList backBuf)
829+ -- XXX create the array in reverse instead
830+ arr0 = A.fromListN n (Prelude.reverse src0)
831+ arr1 = Array contents next end
832+ -- XXX Use StreamK to avoid adding arbitrary layers of
833+ -- constructors every time.
834+ str = D.cons arr0 (D.cons arr1 (D.Stream step s))
835+ return (b, str)
836+ PR.Error err -> throwM $ ParseError err
837+-}
838+
839+-- | Parse an array stream using the supplied 'Parser'. Returns the parse
840+-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
841+--
842+-- The following alternative to this function allows composing the parser using
843+-- the parser Monad:
844+--
845+-- >>> parseBreakStreamK p = StreamK.parseBreakChunks (ParserK.adaptC p)
846+--
847+-- We can compare perf and remove this one or define it in terms of that.
848+--
849+-- /Internal/
850+--
851+{-# INLINE_NORMAL parseBreakChunksK #-}
852+parseBreakChunksK ::
853+ forall m a b. (MonadIO m, Unbox a)
854+ => Parser a m b
855+ -> StreamK m (Array a)
856+ -> m (Either ParseError b, StreamK m (Array a))
857+-- parseBreakStreamK p = StreamK.parseBreakChunks (ParserK.adaptC p)
858+parseBreakChunksK (Parser pstep initial extract) stream = do
859+ res <- initial
860+ case res of
861+ IPartial s -> go s stream []
862+ IDone b -> return (Right b, stream)
863+ IError err -> return (Left (ParseError err), stream)
864+
865+ where
866+
867+ -- "backBuf" contains last few items in the stream that we may have to
868+ -- backtrack to.
869+ --
870+ -- XXX currently we are using a dumb list based approach for backtracking
871+ -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
872+ -- That will allow us more efficient random back and forth movement.
873+ go !pst st backBuf = do
874+ let stop = goStop pst backBuf -- (, K.nil) <$> extract pst
875+ single a = yieldk a StreamK.nil
876+ yieldk arr r = goArray pst backBuf r arr
877+ in StreamK.foldStream defState yieldk single stop st
878+
879+ -- Use strictness on "cur" to keep it unboxed
880+ goArray !pst backBuf st (Array _ cur end) | cur == end = go pst st backBuf
881+ goArray !pst backBuf st (Array contents cur end) = do
882+ x <- liftIO $ peekAt cur contents
883+ pRes <- pstep pst x
884+ let next = INDEX_NEXT(cur,a)
885+ case pRes of
886+ Parser.Partial 0 s ->
887+ goArray s [] st (Array contents next end)
888+ Parser.Partial n s -> do
889+ assert (n <= Prelude.length (x:backBuf)) (return ())
890+ let src0 = Prelude.take n (x:backBuf)
891+ arr0 = fromListN n (Prelude.reverse src0)
892+ arr1 = Array contents next end
893+ src = arr0 <> arr1
894+ goArray s [] st src
895+ Parser.Continue 0 s ->
896+ goArray s (x:backBuf) st (Array contents next end)
897+ Parser.Continue n s -> do
898+ assert (n <= Prelude.length (x:backBuf)) (return ())
899+ let (src0, buf1) = Prelude.splitAt n (x:backBuf)
900+ arr0 = fromListN n (Prelude.reverse src0)
901+ arr1 = Array contents next end
902+ src = arr0 <> arr1
903+ goArray s buf1 st src
904+ Parser.Done 0 b -> do
905+ let arr = Array contents next end
906+ return (Right b, StreamK.cons arr st)
907+ Parser.Done n b -> do
908+ assert (n <= Prelude.length (x:backBuf)) (return ())
909+ let src0 = Prelude.take n (x:backBuf)
910+ -- XXX Use fromListRevN once implemented
911+ -- arr0 = A.fromListRevN n src0
912+ arr0 = fromListN n (Prelude.reverse src0)
913+ arr1 = Array contents next end
914+ str = StreamK.cons arr0 (StreamK.cons arr1 st)
915+ return (Right b, str)
916+ Parser.Error err -> do
917+ let n = Prelude.length backBuf
918+ arr0 = fromListN n (Prelude.reverse backBuf)
919+ arr1 = Array contents cur end
920+ str = StreamK.cons arr0 (StreamK.cons arr1 stream)
921+ return (Left (ParseError err), str)
922+
923+ -- This is a simplified goArray
924+ goExtract !pst backBuf (Array _ cur end)
925+ | cur == end = goStop pst backBuf
926+ goExtract !pst backBuf (Array contents cur end) = do
927+ x <- liftIO $ peekAt cur contents
928+ pRes <- pstep pst x
929+ let next = INDEX_NEXT(cur,a)
930+ case pRes of
931+ Parser.Partial 0 s ->
932+ goExtract s [] (Array contents next end)
933+ Parser.Partial n s -> do
934+ assert (n <= Prelude.length (x:backBuf)) (return ())
935+ let src0 = Prelude.take n (x:backBuf)
936+ arr0 = fromListN n (Prelude.reverse src0)
937+ arr1 = Array contents next end
938+ src = arr0 <> arr1
939+ goExtract s [] src
940+ Parser.Continue 0 s ->
941+ goExtract s backBuf (Array contents next end)
942+ Parser.Continue n s -> do
943+ assert (n <= Prelude.length (x:backBuf)) (return ())
944+ let (src0, buf1) = Prelude.splitAt n (x:backBuf)
945+ arr0 = fromListN n (Prelude.reverse src0)
946+ arr1 = Array contents next end
947+ src = arr0 <> arr1
948+ goExtract s buf1 src
949+ Parser.Done 0 b -> do
950+ let arr = Array contents next end
951+ return (Right b, StreamK.fromPure arr)
952+ Parser.Done n b -> do
953+ assert (n <= Prelude.length backBuf) (return ())
954+ let src0 = Prelude.take n backBuf
955+ -- XXX Use fromListRevN once implemented
956+ -- arr0 = A.fromListRevN n src0
957+ arr0 = fromListN n (Prelude.reverse src0)
958+ arr1 = Array contents next end
959+ str = StreamK.cons arr0 (StreamK.fromPure arr1)
960+ return (Right b, str)
961+ Parser.Error err -> do
962+ let n = Prelude.length backBuf
963+ arr0 = fromListN n (Prelude.reverse backBuf)
964+ arr1 = Array contents cur end
965+ str = StreamK.cons arr0 (StreamK.cons arr1 stream)
966+ return (Left (ParseError err), str)
967+
968+ -- This is a simplified goExtract
969+ {-# INLINE goStop #-}
970+ goStop !pst backBuf = do
971+ pRes <- extract pst
972+ case pRes of
973+ Parser.Partial _ _ -> error "Bug: parseBreak: Partial in extract"
974+ Parser.Continue 0 s ->
975+ goStop s backBuf
976+ Parser.Continue n s -> do
977+ assert (n <= Prelude.length backBuf) (return ())
978+ let (src0, buf1) = Prelude.splitAt n backBuf
979+ arr = fromListN n (Prelude.reverse src0)
980+ goExtract s buf1 arr
981+ Parser.Done 0 b ->
982+ return (Right b, StreamK.nil)
983+ Parser.Done n b -> do
984+ assert (n <= Prelude.length backBuf) (return ())
985+ let src0 = Prelude.take n backBuf
986+ -- XXX Use fromListRevN once implemented
987+ -- arr0 = A.fromListRevN n src0
988+ arr0 = fromListN n (Prelude.reverse src0)
989+ return (Right b, StreamK.fromPure arr0)
990+ Parser.Error err -> do
991+ let n = Prelude.length backBuf
992+ arr0 = fromListN n (Prelude.reverse backBuf)
993+ return (Left (ParseError err), StreamK.fromPure arr0)