@@ -5,7 +5,6 @@
55-- Copyright : (c) 2018 Composewell Technologies
66-- (c) Roman Leshchinskiy 2008-2010
77-- (c) The University of Glasgow, 2009
8−--
98-- License : BSD3
109-- Maintainer : streamly@composewell.com
1110-- Stability : experimental
@@ -302,11 +301,11 @@ module Streamly.Internal.Data.Stream.StreamD
302301 )
303302where
304303
305−import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
304+import Control.Concurrent (myThreadId)
306305import Control.Exception
307− (assert, Exception, SomeException, AsyncException, mask_)
308−import Control.Monad (void, when)
309−import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
306+ (Exception, SomeException, mask_)
307+import Control.Monad (void)
308+import Control.Monad.Catch (MonadCatch)
310309import Control.Monad.IO.Class (MonadIO(..))
311310import Control.Monad.Trans.Reader (ReaderT)
312311import Control.Monad.Trans.State.Strict (StateT)
@@ -315,31 +314,19 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
315314import Data.Functor.Identity (Identity(..))
316315import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
317316import Data.Maybe (fromJust, isJust)
318−import Foreign.Storable (Storable(..))
319−import GHC.Exts (SpecConstrAnnotation(..))
320−import GHC.Types (SPEC(..))
321317import Fusion.Plugin.Types (Fuse(..))
322318import Streamly.Internal.Data.Time.Units
323319 (TimeUnit64, toRelTime64, diffAbsTime64)
324−import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
325−import Streamly.Internal.Data.Fold.Types (Fold(..))
326−import Streamly.Internal.Data.Parser (ParseError(..))
327320import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
328321import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
329322import Streamly.Internal.Data.Unfold.Types (Unfold(..))
330−import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
331323
332−import qualified Streamly.Internal.Data.IORef.Prim as Prim
333324import qualified Streamly.Internal.Data.Pipe.Types as Pipe
334−import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
335325import qualified Streamly.Internal.Data.Fold as FL
336326import qualified Streamly.Internal.Data.Stream.StreamK as K
337−import qualified Streamly.Internal.Data.Parser as PR
338−import qualified Streamly.Internal.Data.Parser.ParserD as PRD
339327import qualified Control.Monad.Catch as MC
340328import qualified Control.Monad.Trans.Reader as Reader
341329import qualified Control.Monad.Trans.State.Strict as State
342−import qualified Prelude
343330
344331import Prelude hiding
345332 ( map, mapM, mapM_, repeat, foldr, last, take, filter
@@ -349,6 +336,7 @@ import Prelude hiding
349336 , reverse, iterate, splitAt)
350337import Streamly.Internal.Data.Stream.StreamD.Type
351338import Streamly.Internal.Data.Stream.StreamD.Generate
339+import Streamly.Internal.Data.Stream.StreamD.Eliminate
352340import Streamly.Internal.Data.Stream.StreamD.SplitGroup
353341import Streamly.Internal.Data.SVar
354342
@@ -534,621 +522,6 @@ foldr1 f m = do
534522 Nothing -> return Nothing
535523 Just (h, t) -> fmap Just (foldr f h t)
536524
537−------------------------------------------------------------------------------
538−-- Left Folds
539−------------------------------------------------------------------------------
540−
541−{-# INLINE_NORMAL foldlT #-}
542−foldlT :: (Monad m, Monad (s m), MonadTrans s)
543− => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
544−foldlT fstep begin (Stream step state) = go SPEC begin state
545− where
546− go !_ acc st = do
547− r <- lift $ step defState st
548− case r of
549− Yield x s -> go SPEC (fstep acc x) s
550− Skip s -> go SPEC acc s
551− Stop -> acc
552−
553−-- Note, this is going to have horrible performance, because of the nature of
554−-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
555−-- likely be practically unusable.
556−{-# INLINE_NORMAL foldlS #-}
557−foldlS :: Monad m
558− => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
559−foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
560− where
561− step' gst (Left (st, acc)) = do
562− r <- step (adaptState gst) st
563− return $ case r of
564− Yield x s -> Skip (Left (s, fstep acc x))
565− Skip s -> Skip (Left (s, acc))
566− Stop -> Skip (Right acc)
567−
568− step' gst (Right (Stream stp stt)) = do
569− r <- stp (adaptState gst) stt
570− return $ case r of
571− Yield x s -> Yield x (Right (Stream stp s))
572− Skip s -> Skip (Right (Stream stp s))
573− Stop -> Stop
574−
575−------------------------------------------------------------------------------
576−-- Parses
577−------------------------------------------------------------------------------
578−
579−-- Inlined definition. Without the inline "serially/parser/take" benchmark
580−-- degrades and parseMany does not fuse. Even using "inline" at the callsite
581−-- does not help.
582−{-# INLINE splitAt #-}
583−splitAt :: Int -> [a] -> ([a],[a])
584−splitAt n ls
585− | n <= 0 = ([], ls)
586− | otherwise = splitAt' n ls
587− where
588− splitAt' :: Int -> [a] -> ([a], [a])
589− splitAt' _ [] = ([], [])
590− splitAt' 1 (x:xs) = ([x], xs)
591− splitAt' m (x:xs) = (x:xs', xs'')
592− where
593− (xs', xs'') = splitAt' (m - 1) xs
594−
595−-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
596−-- to make a newtype.
597−{-# ANN type List NoSpecConstr #-}
598−newtype List a = List {getList :: [a]}
599−
600−-- | Run a 'Parse' over a stream.
601−{-# INLINE_NORMAL parselMx' #-}
602−parselMx'
603− :: MonadThrow m
604− => (s -> a -> m (PR.Step s b))
605− -> m s
606− -> (s -> m b)
607− -> Stream m a
608− -> m b
609−parselMx' pstep initial extract (Stream step state) = do
610− initial >>= go SPEC state (List [])
611−
612− where
613−
614− -- XXX currently we are using a dumb list based approach for backtracking
615− -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
616− -- That will allow us more efficient random back and forth movement.
617− {-# INLINE go #-}
618− go !_ st buf !pst = do
619− r <- step defState st
620− case r of
621− Yield x s -> do
622− pRes <- pstep pst x
623− case pRes of
624− PR.Partial 0 pst1 -> go SPEC s (List []) pst1
625− PR.Partial n pst1 -> do
626− assert (n <= length (x:getList buf)) (return ())
627− let src0 = Prelude.take n (x:getList buf)
628− src = Prelude.reverse src0
629− gobuf SPEC s (List []) (List src) pst1
630− PR.Continue 0 pst1 -> go SPEC s (List (x:getList buf)) pst1
631− PR.Continue n pst1 -> do
632− assert (n <= length (x:getList buf)) (return ())
633− let (src0, buf1) = splitAt n (x:getList buf)
634− src = Prelude.reverse src0
635− gobuf SPEC s (List buf1) (List src) pst1
636− PR.Done _ b -> return b
637− PR.Error err -> throwM $ ParseError err
638− Skip s -> go SPEC s buf pst
639− Stop -> extract pst
640−
641− gobuf !_ s buf (List []) !pst = go SPEC s buf pst
642− gobuf !_ s buf (List (x:xs)) !pst = do
643− pRes <- pstep pst x
644− case pRes of
645− PR.Partial 0 pst1 ->
646− gobuf SPEC s (List []) (List xs) pst1
647− PR.Partial n pst1 -> do
648− assert (n <= length (x:getList buf)) (return ())
649− let src0 = Prelude.take n (x:getList buf)
650− src = Prelude.reverse src0 ++ xs
651− gobuf SPEC s (List []) (List src) pst1
652− PR.Continue 0 pst1 ->
653− gobuf SPEC s (List (x:getList buf)) (List xs) pst1
654− PR.Continue n pst1 -> do
655− assert (n <= length (x:getList buf)) (return ())
656− let (src0, buf1) = splitAt n (x:getList buf)
657− src = Prelude.reverse src0 ++ xs
658− gobuf SPEC s (List buf1) (List src) pst1
659− PR.Done _ b -> return b
660− PR.Error err -> throwM $ ParseError err
661−
662−------------------------------------------------------------------------------
663−-- Repeated parsing
664−------------------------------------------------------------------------------
665−
666−{-# ANN type ParseChunksState Fuse #-}
667−data ParseChunksState x inpBuf st pst =
668− ParseChunksInit inpBuf st
669− | ParseChunksInitLeftOver inpBuf
670− | ParseChunksStream st inpBuf !pst
671− | ParseChunksBuf inpBuf st inpBuf !pst
672− | ParseChunksYield x (ParseChunksState x inpBuf st pst)
673−
674−{-# INLINE_NORMAL parseMany #-}
675−parseMany
676− :: MonadThrow m
677− => PRD.Parser m a b
678− -> Stream m a
679− -> Stream m b
680−parseMany (PRD.Parser pstep initial extract) (Stream step state) =
681− Stream stepOuter (ParseChunksInit [] state)
682−
683− where
684−
685− {-# INLINE_LATE stepOuter #-}
686− -- Buffer is empty, get the first element from the stream, initialize the
687− -- fold and then go to stream processing loop.
688− stepOuter gst (ParseChunksInit [] st) = do
689− r <- step (adaptState gst) st
690− case r of
691− Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
692− Skip s -> return $ Skip $ ParseChunksInit [] s
693− Stop -> return Stop
694−
695− -- Buffer is not empty, go to buffered processing loop
696− stepOuter _ (ParseChunksInit src st) = do
697− initial >>= return . Skip . ParseChunksBuf src st []
698−
699− -- XXX we just discard any leftover input at the end
700− stepOuter _ (ParseChunksInitLeftOver _) = return Stop
701−
702− -- Buffer is empty, process elements from the stream
703− stepOuter gst (ParseChunksStream st buf pst) = do
704− r <- step (adaptState gst) st
705− case r of
706− Yield x s -> do
707− pRes <- pstep pst x
708− case pRes of
709− PR.Partial 0 pst1 ->
710− return $ Skip $ ParseChunksStream s [] pst1
711− PR.Partial n pst1 -> do
712− assert (n <= length (x:buf)) (return ())
713− let src0 = Prelude.take n (x:buf)
714− src = Prelude.reverse src0
715− return $ Skip $ ParseChunksBuf src s [] pst1
716− PR.Continue 0 pst1 ->
717− return $ Skip $ ParseChunksStream s (x:buf) pst1
718− PR.Continue n pst1 -> do
719− assert (n <= length (x:buf)) (return ())
720− let (src0, buf1) = splitAt n (x:buf)
721− src = Prelude.reverse src0
722− return $ Skip $ ParseChunksBuf src s buf1 pst1
723− PR.Done 0 b -> do
724− return $ Skip $
725− ParseChunksYield b (ParseChunksInit [] s)
726− PR.Done n b -> do
727− assert (n <= length (x:buf)) (return ())
728− let src = Prelude.reverse (Prelude.take n (x:buf))
729− return $ Skip $
730− ParseChunksYield b (ParseChunksInit src s)
731− PR.Error err -> throwM $ ParseError err
732− Skip s -> return $ Skip $ ParseChunksStream s buf pst
733− Stop -> do
734− b <- extract pst
735− let src = Prelude.reverse buf
736− return $ Skip $
737− ParseChunksYield b (ParseChunksInitLeftOver src)
738−
739− -- go back to stream processing mode
740− stepOuter _ (ParseChunksBuf [] s buf pst) =
741− return $ Skip $ ParseChunksStream s buf pst
742−
743− -- buffered processing loop
744− stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
745− pRes <- pstep pst x
746− case pRes of
747− PR.Partial 0 pst1 ->
748− return $ Skip $ ParseChunksBuf xs s [] pst1
749− PR.Partial n pst1 -> do
750− assert (n <= length (x:buf)) (return ())
751− let src0 = Prelude.take n (x:buf)
752− src = Prelude.reverse src0 ++ xs
753− return $ Skip $ ParseChunksBuf src s [] pst1
754− PR.Continue 0 pst1 ->
755− return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
756− PR.Continue n pst1 -> do
757− assert (n <= length (x:buf)) (return ())
758− let (src0, buf1) = splitAt n (x:buf)
759− src = Prelude.reverse src0 ++ xs
760− return $ Skip $ ParseChunksBuf src s buf1 pst1
761− PR.Done 0 b ->
762− return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
763− PR.Done n b -> do
764− assert (n <= length (x:buf)) (return ())
765− let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
766− return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
767− PR.Error err -> throwM $ ParseError err
768−
769− stepOuter _ (ParseChunksYield a next) = return $ Yield a next
770−
771−{-# ANN type ConcatParseState Fuse #-}
772−data ConcatParseState x inpBuf st p =
773− ConcatParseInit inpBuf st p
774− | ConcatParseInitLeftOver inpBuf
775− | ConcatParseStream st inpBuf p
776− | ConcatParseBuf inpBuf st inpBuf p
777− | ConcatParseYield x (ConcatParseState x inpBuf st p)
778−
779−{-# INLINE_NORMAL parseIterate #-}
780−parseIterate
781− :: MonadThrow m
782− => (b -> PRD.Parser m a b)
783− -> b
784− -> Stream m a
785− -> Stream m b
786−parseIterate func seed (Stream step state) =
787− Stream stepOuter (ConcatParseInit [] state (func seed))
788−
789− where
790−
791− {-# INLINE_LATE stepOuter #-}
792− -- Buffer is empty, go to stream processing loop
793− stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
794− initial >>= \r -> return $ Skip $ ConcatParseStream st []
795− (PRD.Parser pstep (return r) extract)
796−
797− -- Buffer is not empty, go to buffered processing loop
798− stepOuter _ (ConcatParseInit src st
799− (PRD.Parser pstep initial extract)) = do
800− initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
801− (PRD.Parser pstep (return r) extract)
802−
803− -- XXX we just discard any leftover input at the end
804− stepOuter _ (ConcatParseInitLeftOver _) = return Stop
805−
806− -- Buffer is empty process elements from the stream
807− stepOuter gst (ConcatParseStream st buf
808− p@(PRD.Parser pstep initial extract)) = do
809− r <- step (adaptState gst) st
810− case r of
811− Yield x s -> do
812− pst <- initial
813− pRes <- pstep pst x
814− case pRes of
815− PR.Partial 0 pst1 ->
816− return $ Skip $ ConcatParseStream s []
817− (PRD.Parser pstep (return pst1) extract)
818− PR.Partial n pst1 -> do
819− assert (n <= length (x:buf)) (return ())
820− let src0 = Prelude.take n (x:buf)
821− src = Prelude.reverse src0
822− return $ Skip $ ConcatParseBuf src s []
823− (PRD.Parser pstep (return pst1) extract)
824− -- PR.Continue 0 pst1 ->
825− -- return $ Skip $ ConcatParseStream s (x:buf) pst1
826− PR.Continue n pst1 -> do
827− assert (n <= length (x:buf)) (return ())
828− let (src0, buf1) = splitAt n (x:buf)
829− src = Prelude.reverse src0
830− return $ Skip $ ConcatParseBuf src s buf1
831− (PRD.Parser pstep (return pst1) extract)
832− -- XXX Specialize for Stop 0 common case?
833− PR.Done n b -> do
834− assert (n <= length (x:buf)) (return ())
835− let src = Prelude.reverse (Prelude.take n (x:buf))
836− return $ Skip $
837− ConcatParseYield b (ConcatParseInit src s (func b))
838− PR.Error err -> throwM $ ParseError err
839− Skip s -> return $ Skip $ ConcatParseStream s buf p
840− Stop -> do
841− pst <- initial
842− b <- extract pst
843− let src = Prelude.reverse buf
844− return $ Skip $
845− ConcatParseYield b (ConcatParseInitLeftOver src)
846−
847− -- go back to stream processing mode
848− stepOuter _ (ConcatParseBuf [] s buf p) =
849− return $ Skip $ ConcatParseStream s buf p
850−
851− -- buffered processing loop
852− stepOuter _ (ConcatParseBuf (x:xs) s buf
853− (PRD.Parser pstep initial extract)) = do
854− pst <- initial
855− pRes <- pstep pst x
856− case pRes of
857− PR.Partial 0 pst1 ->
858− return $ Skip $ ConcatParseBuf xs s []
859− (PRD.Parser pstep (return pst1) extract)
860− PR.Partial n pst1 -> do
861− assert (n <= length (x:buf)) (return ())
862− let src0 = Prelude.take n (x:buf)
863− src = Prelude.reverse src0 ++ xs
864− return $ Skip $ ConcatParseBuf src s []
865− (PRD.Parser pstep (return pst1) extract)
866− -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
867− PR.Continue n pst1 -> do
868− assert (n <= length (x:buf)) (return ())
869− let (src0, buf1) = splitAt n (x:buf)
870− src = Prelude.reverse src0 ++ xs
871− return $ Skip $ ConcatParseBuf src s buf1
872− (PRD.Parser pstep (return pst1) extract)
873− -- XXX Specialize for Stop 0 common case?
874− PR.Done n b -> do
875− assert (n <= length (x:buf)) (return ())
876− let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
877− return $ Skip $ ConcatParseYield b
878− (ConcatParseInit src s (func b))
879− PR.Error err -> throwM $ ParseError err
880−
881− stepOuter _ (ConcatParseYield a next) = return $ Yield a next
882−
883−------------------------------------------------------------------------------
884−-- Specialized Folds
885−------------------------------------------------------------------------------
886−
887−-- | Run a streaming composition, discard the results.
888−{-# INLINE_LATE drain #-}
889−drain :: Monad m => Stream m a -> m ()
890−-- drain = foldrM (\_ xs -> xs) (return ())
891−drain (Stream step state) = go SPEC state
892− where
893− go !_ st = do
894− r <- step defState st
895− case r of
896− Yield _ s -> go SPEC s
897− Skip s -> go SPEC s
898− Stop -> return ()
899−
900−{-# INLINE_NORMAL null #-}
901−null :: Monad m => Stream m a -> m Bool
902−null m = foldrM (\_ _ -> return False) (return True) m
903−
904−{-# INLINE_NORMAL head #-}
905−head :: Monad m => Stream m a -> m (Maybe a)
906−head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
907−
908−{-# INLINE_NORMAL headElse #-}
909−headElse :: Monad m => a -> Stream m a -> m a
910−headElse a m = foldrM (\x _ -> return x) (return a) m
911−
912−-- Does not fuse, has the same performance as the StreamK version.
913−{-# INLINE_NORMAL tail #-}
914−tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
915−tail (UnStream step state) = go state
916− where
917− go st = do
918− r <- step defState st
919− case r of
920− Yield _ s -> return (Just $ Stream step s)
921− Skip s -> go s
922− Stop -> return Nothing
923−
924−-- XXX will it fuse? need custom impl?
925−{-# INLINE_NORMAL last #-}
926−last :: Monad m => Stream m a -> m (Maybe a)
927−last = foldl' (\_ y -> Just y) Nothing
928−
929−{-# INLINE_NORMAL elem #-}
930−elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
931−-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
932−elem e (Stream step state) = go state
933− where
934− go st = do
935− r <- step defState st
936− case r of
937− Yield x s
938− | x == e -> return True
939− | otherwise -> go s
940− Skip s -> go s
941− Stop -> return False
942−
943−{-# INLINE_NORMAL notElem #-}
944−notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
945−notElem e s = fmap not (elem e s)
946−
947−{-# INLINE_NORMAL all #-}
948−all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
949−-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
950−all p (Stream step state) = go state
951− where
952− go st = do
953− r <- step defState st
954− case r of
955− Yield x s
956− | p x -> go s
957− | otherwise -> return False
958− Skip s -> go s
959− Stop -> return True
960−
961−{-# INLINE_NORMAL any #-}
962−any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
963−-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
964−any p (Stream step state) = go state
965− where
966− go st = do
967− r <- step defState st
968− case r of
969− Yield x s
970− | p x -> return True
971− | otherwise -> go s
972− Skip s -> go s
973− Stop -> return False
974−
975−{-# INLINE_NORMAL maximum #-}
976−maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
977−maximum (Stream step state) = go Nothing state
978− where
979− go Nothing st = do
980− r <- step defState st
981− case r of
982− Yield x s -> go (Just x) s
983− Skip s -> go Nothing s
984− Stop -> return Nothing
985− go (Just acc) st = do
986− r <- step defState st
987− case r of
988− Yield x s
989− | acc <= x -> go (Just x) s
990− | otherwise -> go (Just acc) s
991− Skip s -> go (Just acc) s
992− Stop -> return (Just acc)
993−
994−{-# INLINE_NORMAL maximumBy #-}
995−maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
996−maximumBy cmp (Stream step state) = go Nothing state
997− where
998− go Nothing st = do
999− r <- step defState st
1000− case r of
1001− Yield x s -> go (Just x) s
1002− Skip s -> go Nothing s
1003− Stop -> return Nothing
1004− go (Just acc) st = do
1005− r <- step defState st
1006− case r of
1007− Yield x s -> case cmp acc x of
1008− GT -> go (Just acc) s
1009− _ -> go (Just x) s
1010− Skip s -> go (Just acc) s
1011− Stop -> return (Just acc)
1012−
1013−{-# INLINE_NORMAL minimum #-}
1014−minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1015−minimum (Stream step state) = go Nothing state
1016− where
1017− go Nothing st = do
1018− r <- step defState st
1019− case r of
1020− Yield x s -> go (Just x) s
1021− Skip s -> go Nothing s
1022− Stop -> return Nothing
1023− go (Just acc) st = do
1024− r <- step defState st
1025− case r of
1026− Yield x s
1027− | acc <= x -> go (Just acc) s
1028− | otherwise -> go (Just x) s
1029− Skip s -> go (Just acc) s
1030− Stop -> return (Just acc)
1031−
1032−{-# INLINE_NORMAL minimumBy #-}
1033−minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1034−minimumBy cmp (Stream step state) = go Nothing state
1035− where
1036− go Nothing st = do
1037− r <- step defState st
1038− case r of
1039− Yield x s -> go (Just x) s
1040− Skip s -> go Nothing s
1041− Stop -> return Nothing
1042− go (Just acc) st = do
1043− r <- step defState st
1044− case r of
1045− Yield x s -> case cmp acc x of
1046− GT -> go (Just x) s
1047− _ -> go (Just acc) s
1048− Skip s -> go (Just acc) s
1049− Stop -> return (Just acc)
1050−
1051−{-# INLINE_NORMAL (!!) #-}
1052−(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1053−(Stream step state) !! i = go i state
1054− where
1055− go n st = do
1056− r <- step defState st
1057− case r of
1058− Yield x s | n < 0 -> return Nothing
1059− | n == 0 -> return $ Just x
1060− | otherwise -> go (n - 1) s
1061− Skip s -> go n s
1062− Stop -> return Nothing
1063−
1064−{-# INLINE_NORMAL lookup #-}
1065−lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1066−lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1067− (return Nothing) m
1068−
1069−{-# INLINE_NORMAL findM #-}
1070−findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1071−findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1072− (return Nothing) m
1073−
1074−{-# INLINE find #-}
1075−find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1076−find p = findM (return . p)
1077−
1078−{-# INLINE_NORMAL findIndices #-}
1079−findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1080−findIndices p (Stream step state) = Stream step' (state, 0)
1081− where
1082− {-# INLINE_LATE step' #-}
1083− step' gst (st, i) = i `seq` do
1084− r <- step (adaptState gst) st
1085− return $ case r of
1086− Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1087− Skip s -> Skip (s, i)
1088− Stop -> Stop
1089−
1090−{-# INLINE toListRev #-}
1091−toListRev :: Monad m => Stream m a -> m [a]
1092−toListRev = foldl' (flip (:)) []
1093−
1094−-- We can implement reverse as:
1095−--
1096−-- > reverse = foldlS (flip cons) nil
1097−--
1098−-- However, this implementation is unusable because of the horrible performance
1099−-- of cons. So we just convert it to a list first and then stream from the
1100−-- list.
1101−--
1102−-- XXX Maybe we can use an Array instead of a list here?
1103−{-# INLINE_NORMAL reverse #-}
1104−reverse :: Monad m => Stream m a -> Stream m a
1105−reverse m = Stream step Nothing
1106− where
1107− {-# INLINE_LATE step #-}
1108− step _ Nothing = do
1109− xs <- toListRev m
1110− return $ Skip (Just xs)
1111− step _ (Just (x:xs)) = return $ Yield x (Just xs)
1112− step _ (Just []) = return Stop
1113−
1114−-- Much faster reverse for Storables
1115−{-# INLINE_NORMAL reverse' #-}
1116−reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1117−{-
1118−-- This commented implementation copies the whole stream into one single array
1119−-- and then streams from that array, this is 3-4x faster than the chunked code
1120−-- that follows. Though this could be problematic due to unbounded large
1121−-- allocations. We need to figure out why the chunked code is slower and if we
1122−-- can optimize the chunked code to work as fast as this one. It may be a
1123−-- fusion issue?
1124−import Foreign.ForeignPtr (touchForeignPtr)
1125−import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1126−import Foreign.Ptr (Ptr, plusPtr)
1127−reverse' m = Stream step Nothing
1128− where
1129− {-# INLINE_LATE step #-}
1130− step _ Nothing = do
1131− arr <- A.fromStreamD m
1132− let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1133− return $ Skip $ Just (aStart arr, p)
1134−
1135− step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1136−
1137− step _ (Just (start, p)) = do
1138− let !x = A.unsafeInlineIO $ do
1139− r <- peek p
1140− touchForeignPtr start
1141− return r
1142− next = p `plusPtr` negate (sizeOf (undefined :: a))
1143− return $ Yield x (Just (start, next))
1144−-}
1145−reverse' m =
1146− A.flattenArraysRev
1147− $ fromStreamK
1148− $ K.reverse
1149− $ toStreamK
1150− $ A.fromStreamDArraysOf A.defaultChunkSize m
1151−
1152525------------------------------------------------------------------------------
1153526-- Substreams
1154527------------------------------------------------------------------------------
@@ -2671,168 +2044,6 @@ rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
26712044rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
26722045rollingMap f = rollingMapM (\x y -> return $ f x y)
26732046
2674−------------------------------------------------------------------------------
2675−-- Tapping/Distributing
2676−------------------------------------------------------------------------------
2677−
2678−data TapState fs st a
2679− = TapInit | Tapping !fs st | TapDone st
2680−
2681−-- XXX Multiple yield points
2682−{-# INLINE tap #-}
2683−tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
2684−tap (Fold fstep initial extract) (Stream step state) = Stream step' TapInit
2685−
2686− where
2687−
2688− step' _ TapInit = do
2689− r <- initial
2690− return $ Skip (Tapping r state)
2691− step' gst (Tapping acc st) = do
2692− r <- step gst st
2693− case r of
2694− Yield x s -> do
2695− res <- fstep acc x
2696− return
2697− $ Yield x
2698− $ case res of
2699− FL.Partial fs -> Tapping fs s
2700− FL.Done _ -> TapDone s
2701− Skip s -> return $ Skip (Tapping acc s)
2702− Stop -> do
2703− void $ extract acc
2704− return Stop
2705− step' gst (TapDone st) = do
2706− r <- step gst st
2707− return
2708− $ case r of
2709− Yield x s -> Yield x (TapDone s)
2710− Skip s -> Skip (TapDone s)
2711− Stop -> Stop
2712−
2713−data TapOffState fs s a
2714− = TapOffInit
2715− | TapOffTapping !fs s Int
2716− | TapOffDone s
2717−
2718−-- XXX Multiple yield points
2719−{-# INLINE_NORMAL tapOffsetEvery #-}
2720−tapOffsetEvery :: Monad m
2721− => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
2722−tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
2723− Stream step' TapOffInit
2724−
2725− where
2726−
2727− {-# INLINE_LATE step' #-}
2728− step' _ TapOffInit = do
2729− r <- initial
2730− return $ Skip $ TapOffTapping r state (offset `mod` n)
2731− step' gst (TapOffTapping acc st count) = do
2732− r <- step gst st
2733− case r of
2734− Yield x s -> do
2735− next <-
2736− if count <= 0
2737− then do
2738− res <- fstep acc x
2739− return
2740− $ case res of
2741− FL.Partial sres ->
2742− TapOffTapping sres s (n - 1)
2743− FL.Done _ -> TapOffDone s
2744− else return $ TapOffTapping acc s (count - 1)
2745− return $ Yield x next
2746− Skip s -> return $ Skip (TapOffTapping acc s count)
2747− Stop -> do
2748− void $ extract acc
2749− return Stop
2750− step' gst (TapOffDone st) = do
2751− r <- step gst st
2752− return
2753− $ case r of
2754− Yield x s -> Yield x (TapOffDone s)
2755− Skip s -> Skip (TapOffDone s)
2756− Stop -> Stop
2757−
2758−{-# INLINE_NORMAL pollCounts #-}
2759−pollCounts
2760− :: MonadAsync m
2761− => (a -> Bool)
2762− -> (Stream m Int -> Stream m Int)
2763− -> Fold m Int b
2764− -> Stream m a
2765− -> Stream m a
2766−pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
2767− where
2768−
2769− {-# INLINE_LATE step' #-}
2770− step' _ Nothing = do
2771− -- As long as we are using an "Int" for counts lockfree reads from
2772− -- Var should work correctly on both 32-bit and 64-bit machines.
2773− -- However, an Int on a 32-bit machine may overflow quickly.
2774− countVar <- liftIO $ Prim.newIORef (0 :: Int)
2775− tid <- forkManaged
2776− $ void $ foldOnce fld
2777− $ transf $ fromPrimIORef countVar
2778− return $ Skip (Just (countVar, tid, state))
2779−
2780− step' gst (Just (countVar, tid, st)) = do
2781− r <- step gst st
2782− case r of
2783− Yield x s -> do
2784− when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
2785− return $ Yield x (Just (countVar, tid, s))
2786− Skip s -> return $ Skip (Just (countVar, tid, s))
2787− Stop -> do
2788− liftIO $ killThread tid
2789− return Stop
2790−
2791−{-# INLINE_NORMAL tapRate #-}
2792−tapRate ::
2793− (MonadAsync m, MonadCatch m)
2794− => Double
2795− -> (Int -> m b)
2796− -> Stream m a
2797− -> Stream m a
2798−tapRate samplingRate action (Stream step state) = Stream step' Nothing
2799− where
2800− {-# NOINLINE loop #-}
2801− loop countVar prev = do
2802− i <-
2803− MC.catch
2804− (do liftIO $ threadDelay (round $ samplingRate * 1000000)
2805− i <- liftIO $ Prim.readIORef countVar
2806− let !diff = i - prev
2807− void $ action diff
2808− return i)
2809− (\(e :: AsyncException) -> do
2810− i <- liftIO $ Prim.readIORef countVar
2811− let !diff = i - prev
2812− void $ action diff
2813− throwM (MC.toException e))
2814− loop countVar i
2815−
2816− {-# INLINE_LATE step' #-}
2817− step' _ Nothing = do
2818− countVar <- liftIO $ Prim.newIORef 0
2819− tid <- fork $ loop countVar 0
2820− ref <- liftIO $ newIORef ()
2821− _ <- liftIO $ mkWeakIORef ref (killThread tid)
2822− return $ Skip (Just (countVar, tid, state, ref))
2823−
2824− step' gst (Just (countVar, tid, st, ref)) = do
2825− r <- step gst st
2826− case r of
2827− Yield x s -> do
2828− liftIO $ Prim.modifyIORef' countVar (+ 1)
2829− return $ Yield x (Just (countVar, tid, s, ref))
2830− Skip s -> return $ Skip (Just (countVar, tid, s, ref))
2831− Stop -> do
2832− liftIO $ killThread tid
2833− return Stop
2834−
2835−
28362047-------------------------------------------------------------------------------
28372048-- Filtering
28382049-------------------------------------------------------------------------------
@@ -3279,97 +2490,10 @@ the (Stream step state) = go state
32792490 Skip s -> go' n s
32802491 Stop -> return (Just n)
32812492
3282−{-# INLINE_NORMAL foldOnce #-}
3283−foldOnce :: (Monad m) => Fold m a b -> Stream m a -> m b
3284−foldOnce (Fold fstep begin done) (Stream step state) =
3285− begin >>= \x -> go SPEC x state
3286−
3287− where
3288−
3289− {-# INLINE go #-}
3290− go !_ !fs st = do
3291− r <- step defState st
3292− case r of
3293− Yield x s -> do
3294− res <- fstep fs x
3295− case res of
3296− FL.Done b -> return b
3297− FL.Partial fs1 -> go SPEC fs1 s
3298− Skip s -> go SPEC fs s
3299− Stop -> done fs
3300−
33012493-------------------------------------------------------------------------------
33022494-- Concurrent application and fold
33032495-------------------------------------------------------------------------------
33042496
3305−-- XXX These functions should be moved to Stream/Parallel.hs
3306−--
3307−-- Using StreamD the worker stream producing code can fuse with the code to
3308−-- queue output to the SVar giving some perf boost.
3309−--
3310−-- Note that StreamD can only be used in limited situations, specifically, we
3311−-- cannot implement joinStreamVarPar using this.
3312−--
3313−-- XXX make sure that the SVar passed is a Parallel style SVar.
3314−
3315−-- | Fold the supplied stream to the SVar asynchronously using Parallel
3316−-- concurrency style.
3317−-- {-# INLINE_NORMAL toSVarParallel #-}
3318−{-# INLINE toSVarParallel #-}
3319−toSVarParallel :: MonadAsync m
3320− => State t m a -> SVar t m a -> Stream m a -> m ()
3321−toSVarParallel st sv xs =
3322− if svarInspectMode sv
3323− then forkWithDiag
3324− else do
3325− tid <-
3326− case getYieldLimit st of
3327− Nothing -> doFork (work Nothing)
3328− (svarMrun sv)
3329− (handleChildException sv)
3330− Just _ -> doFork (workLim Nothing)
3331− (svarMrun sv)
3332− (handleChildException sv)
3333− modifyThread sv tid
3334−
3335− where
3336−
3337− {-# NOINLINE work #-}
3338− work info = (foldOnce (FL.toParallelSVar sv info) xs)
3339−
3340− {-# NOINLINE workLim #-}
3341− workLim info = foldOnce (FL.toParallelSVarLimited sv info) xs
3342−
3343− {-# NOINLINE forkWithDiag #-}
3344− forkWithDiag = do
3345− -- We do not use workerCount in case of ParallelVar but still there is
3346− -- no harm in maintaining it correctly.
3347− liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
3348− recordMaxWorkers sv
3349− -- This allocation matters when significant number of workers are being
3350− -- sent. We allocate it only when needed. The overhead increases by 4x.
3351− winfo <-
3352− case yieldRateInfo sv of
3353− Nothing -> return Nothing
3354− Just _ -> liftIO $ do
3355− cntRef <- newIORef 0
3356− t <- getTime Monotonic
3357− lat <- newIORef (0, t)
3358− return $ Just WorkerInfo
3359− { workerYieldMax = 0
3360− , workerYieldCount = cntRef
3361− , workerLatencyStart = lat
3362− }
3363− tid <-
3364− case getYieldLimit st of
3365− Nothing -> doFork (work winfo)
3366− (svarMrun sv)
3367− (handleChildException sv)
3368− Just _ -> doFork (workLim winfo)
3369− (svarMrun sv)
3370− (handleChildException sv)
3371− modifyThread sv tid
3372−
33732497{-# INLINE_NORMAL mkParallelD #-}
33742498mkParallelD :: MonadAsync m => Stream m a -> Stream m a
33752499mkParallelD m = Stream step Nothing
@@ -3429,75 +2553,6 @@ newCallbackStream = do
34292553 -- SVar?
34302554 return (callback, fromStreamD (fromSVar sv))
34312555
3432−-------------------------------------------------------------------------------
3433−-- Concurrent tap
3434−-------------------------------------------------------------------------------
3435−
3436−-- | Create an SVar with a fold consumer that will fold any elements sent to it
3437−-- using the supplied fold function.
3438−{-# INLINE newFoldSVar #-}
3439−newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
3440−newFoldSVar stt f = do
3441− -- Buffer size for the SVar is derived from the current state
3442− sv <- newParallelVar StopAny (adaptState stt)
3443− -- Add the producer thread-id to the SVar.
3444− liftIO myThreadId >>= modifyThread sv
3445− void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
3446− return sv
3447−
3448− where
3449−
3450− {-# NOINLINE work #-}
3451− work sv = void $ foldOnce f $ fromProducer sv
3452−
3453−{-# INLINE_NORMAL tapAsync #-}
3454−tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
3455−tapAsync f (Stream step1 state1) = Stream step TapInit
3456− where
3457−
3458− drainFold svr = do
3459− -- In general, a Stop event would come equipped with the result
3460− -- of the fold. It is not used here but it would be useful in
3461− -- applicative and distribute.
3462− done <- fromConsumer svr
3463− when (not done) $ do
3464− liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
3465− $ takeMVar (outputDoorBellFromConsumer svr)
3466− drainFold svr
3467−
3468− stopFold svr = do
3469− liftIO $ sendStop svr Nothing
3470− -- drain/wait until a stop event arrives from the fold.
3471− drainFold svr
3472−
3473− {-# INLINE_LATE step #-}
3474− step gst TapInit = do
3475− sv <- newFoldSVar gst f
3476− return $ Skip (Tapping sv state1)
3477−
3478− step gst (Tapping sv st) = do
3479− r <- step1 gst st
3480− case r of
3481− Yield a s -> do
3482− done <- pushToFold sv a
3483− if done
3484− then do
3485− -- XXX we do not need to wait synchronously here
3486− stopFold sv
3487− return $ Yield a (TapDone s)
3488− else return $ Yield a (Tapping sv s)
3489− Skip s -> return $ Skip (Tapping sv s)
3490− Stop -> do
3491− stopFold sv
3492− return $ Stop
3493−
3494− step gst (TapDone st) = do
3495− r <- step1 gst st
3496− return $ case r of
3497− Yield a s -> Yield a (TapDone s)
3498− Skip s -> Skip (TapDone s)
3499− Stop -> Stop
3500−
35012556------------------------------------------------------------------------------
35022557-- Time related
35032558------------------------------------------------------------------------------