git.iko.soy/streamly
on master help sign in
Refactor StreamD's elimination ops to StreamD.Eliminate.
a4e3ddbPranay Sashank
src/Streamly/Internal/Data/Stream/StreamD.hs+5−950
compare with master
unifiedsplit
@@ -5,7 +5,6 @@
5-- Copyright : (c) 2018 Composewell Technologies
5-- Copyright : (c) 2018 Composewell Technologies
6-- (c) Roman Leshchinskiy 2008-2010
6-- (c) Roman Leshchinskiy 2008-2010
7-- (c) The University of Glasgow, 2009
7-- (c) The University of Glasgow, 2009
8−--
9-- License : BSD3
8-- License : BSD3
10-- Maintainer : streamly@composewell.com
9-- Maintainer : streamly@composewell.com
11-- Stability : experimental
10-- Stability : experimental
@@ -302,11 +301,11 @@ module Streamly.Internal.Data.Stream.StreamD
302 )
301 )
303where
302where
304
303
305−import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
304+import Control.Concurrent (myThreadId)
306import Control.Exception
305import Control.Exception
307− (assert, Exception, SomeException, AsyncException, mask_)
306+ (Exception, SomeException, mask_)
308−import Control.Monad (void, when)
307+import Control.Monad (void)
309−import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
308+import Control.Monad.Catch (MonadCatch)
310import Control.Monad.IO.Class (MonadIO(..))
309import Control.Monad.IO.Class (MonadIO(..))
311import Control.Monad.Trans.Reader (ReaderT)
310import Control.Monad.Trans.Reader (ReaderT)
312import Control.Monad.Trans.State.Strict (StateT)
311import Control.Monad.Trans.State.Strict (StateT)
@@ -315,31 +314,19 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
315import Data.Functor.Identity (Identity(..))
314import Data.Functor.Identity (Identity(..))
316import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
315import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
317import Data.Maybe (fromJust, isJust)
316import Data.Maybe (fromJust, isJust)
318−import Foreign.Storable (Storable(..))
319−import GHC.Exts (SpecConstrAnnotation(..))
320−import GHC.Types (SPEC(..))
321import Fusion.Plugin.Types (Fuse(..))
317import Fusion.Plugin.Types (Fuse(..))
322import Streamly.Internal.Data.Time.Units
318import Streamly.Internal.Data.Time.Units
323 (TimeUnit64, toRelTime64, diffAbsTime64)
319 (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(..))
327import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
320import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
328import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
321import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
329import Streamly.Internal.Data.Unfold.Types (Unfold(..))
322import Streamly.Internal.Data.Unfold.Types (Unfold(..))
330−import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
331
323
332−import qualified Streamly.Internal.Data.IORef.Prim as Prim
333import qualified Streamly.Internal.Data.Pipe.Types as Pipe
324import qualified Streamly.Internal.Data.Pipe.Types as Pipe
334−import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
335import qualified Streamly.Internal.Data.Fold as FL
325import qualified Streamly.Internal.Data.Fold as FL
336import qualified Streamly.Internal.Data.Stream.StreamK as K
326import 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
339import qualified Control.Monad.Catch as MC
327import qualified Control.Monad.Catch as MC
340import qualified Control.Monad.Trans.Reader as Reader
328import qualified Control.Monad.Trans.Reader as Reader
341import qualified Control.Monad.Trans.State.Strict as State
329import qualified Control.Monad.Trans.State.Strict as State
342−import qualified Prelude
343
330
344import Prelude hiding
331import Prelude hiding
345 ( map, mapM, mapM_, repeat, foldr, last, take, filter
332 ( map, mapM, mapM_, repeat, foldr, last, take, filter
@@ -349,6 +336,7 @@ import Prelude hiding
349 , reverse, iterate, splitAt)
336 , reverse, iterate, splitAt)
350import Streamly.Internal.Data.Stream.StreamD.Type
337import Streamly.Internal.Data.Stream.StreamD.Type
351import Streamly.Internal.Data.Stream.StreamD.Generate
338import Streamly.Internal.Data.Stream.StreamD.Generate
339+import Streamly.Internal.Data.Stream.StreamD.Eliminate
352import Streamly.Internal.Data.Stream.StreamD.SplitGroup
340import Streamly.Internal.Data.Stream.StreamD.SplitGroup
353import Streamly.Internal.Data.SVar
341import Streamly.Internal.Data.SVar
354
342
@@ -534,621 +522,6 @@ foldr1 f m = do
534 Nothing -> return Nothing
522 Nothing -> return Nothing
535 Just (h, t) -> fmap Just (foldr f h t)
523 Just (h, t) -> fmap Just (foldr f h t)
536
524
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−
1152------------------------------------------------------------------------------
525------------------------------------------------------------------------------
1153-- Substreams
526-- Substreams
1154------------------------------------------------------------------------------
527------------------------------------------------------------------------------
@@ -2671,168 +2044,6 @@ rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
2671rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
2044rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
2672rollingMap f = rollingMapM (\x y -> return $ f x y)
2045rollingMap f = rollingMapM (\x y -> return $ f x y)
2673
2046
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−
2836-------------------------------------------------------------------------------
2047-------------------------------------------------------------------------------
2837-- Filtering
2048-- Filtering
2838-------------------------------------------------------------------------------
2049-------------------------------------------------------------------------------
@@ -3279,97 +2490,10 @@ the (Stream step state) = go state
3279 Skip s -> go' n s
2490 Skip s -> go' n s
3280 Stop -> return (Just n)
2491 Stop -> return (Just n)
3281
2492
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−
3301-------------------------------------------------------------------------------
2493-------------------------------------------------------------------------------
3302-- Concurrent application and fold
2494-- Concurrent application and fold
3303-------------------------------------------------------------------------------
2495-------------------------------------------------------------------------------
3304
2496
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−
3373{-# INLINE_NORMAL mkParallelD #-}
2497{-# INLINE_NORMAL mkParallelD #-}
3374mkParallelD :: MonadAsync m => Stream m a -> Stream m a
2498mkParallelD :: MonadAsync m => Stream m a -> Stream m a
3375mkParallelD m = Stream step Nothing
2499mkParallelD m = Stream step Nothing
@@ -3429,75 +2553,6 @@ newCallbackStream = do
3429 -- SVar?
2553 -- SVar?
3430 return (callback, fromStreamD (fromSVar sv))
2554 return (callback, fromStreamD (fromSVar sv))
3431
2555
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−
3501------------------------------------------------------------------------------
2556------------------------------------------------------------------------------
3502-- Time related
2557-- Time related
3503------------------------------------------------------------------------------
2558------------------------------------------------------------------------------
src/Streamly/Internal/Data/Stream/StreamD/Common.hsnew file+88
@@ -0,0 +1,88 @@
1+#include "inline.hs"
2+
3+-- |
4+-- Module : Streamly.Internal.Data.Stream.StreamD.Common
5+-- Copyright : (c) 2018 Composewell Technologies
6+-- (c) Roman Leshchinskiy 2008-2010
7+-- (c) The University of Glasgow, 2009
8+-- License : BSD3
9+-- Maintainer : streamly@composewell.com
10+-- Stability : experimental
11+-- Portability : GHC
12+
13+module Streamly.Internal.Data.Stream.StreamD.Common
14+ (
15+ FromSVarState (..)
16+ , fromProducer
17+ , fromPrimIORef
18+ )
19+where
20+
21+import Control.Monad (when)
22+import Control.Monad.IO.Class (MonadIO(..))
23+import Data.IORef (writeIORef)
24+import Streamly.Internal.Data.IORef.Prim (Prim)
25+import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
26+
27+import qualified Streamly.Internal.Data.IORef.Prim as Prim
28+import qualified Prelude
29+
30+import Prelude hiding
31+ ( map, mapM, mapM_, repeat, foldr, last, take, filter
32+ , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
33+ , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
34+ , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
35+ , reverse, iterate, splitAt)
36+import Streamly.Internal.Data.Stream.StreamD.Type
37+import Streamly.Internal.Data.SVar
38+
39+data FromSVarState t m a =
40+ FromSVarInit
41+ | FromSVarRead (SVar t m a)
42+ | FromSVarLoop (SVar t m a) [ChildEvent a]
43+ | FromSVarDone (SVar t m a)
44+
45+{-# INLINE_NORMAL fromPrimIORef #-}
46+fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
47+fromPrimIORef var = Stream step ()
48+ where
49+ {-# INLINE_LATE step #-}
50+ step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
51+
52+-------------------------------------------------------------------------------
53+-- Process events received by a fold consumer from a stream producer
54+-------------------------------------------------------------------------------
55+
56+{-# INLINE_NORMAL fromProducer #-}
57+fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
58+fromProducer svar = Stream step (FromSVarRead svar)
59+ where
60+
61+ {-# INLINE_LATE step #-}
62+ step _ (FromSVarRead sv) = do
63+ list <- readOutputQ sv
64+ -- Reversing the output is important to guarantee that we process the
65+ -- outputs in the same order as they were generated by the constituent
66+ -- streams.
67+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
68+
69+ step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
70+ step _ (FromSVarLoop sv (ev : es)) = do
71+ case ev of
72+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)
73+ ChildStop tid e -> do
74+ accountThread sv tid
75+ case e of
76+ Nothing -> do
77+ sendStopToProducer sv
78+ return $ Skip (FromSVarDone sv)
79+ Just _ -> error "Bug: fromProducer: received exception"
80+
81+ step _ (FromSVarDone sv) = do
82+ when (svarInspectMode sv) $ do
83+ t <- liftIO $ getTime Monotonic
84+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
85+ liftIO $ printSVar sv "SVar Done"
86+ return Stop
87+
88+ step _ FromSVarInit = undefined
src/Streamly/Internal/Data/Stream/StreamD/Eliminate.hsnew file+1045
@@ -0,0 +1,1045 @@
1+#include "inline.hs"
2+
3+-- |
4+-- Module : Streamly.Internal.Data.Stream.StreamD.Eliminate
5+-- Copyright : (c) 2018 Composewell Technologies
6+-- (c) Roman Leshchinskiy 2008-2010
7+-- (c) The University of Glasgow, 2009
8+-- License : BSD3
9+-- Maintainer : streamly@composewell.com
10+-- Stability : experimental
11+-- Portability : GHC
12+
13+module Streamly.Internal.Data.Stream.StreamD.Eliminate
14+ (
15+ -- * Elimination
16+ -- ** General Folds
17+ foldrS
18+ , foldrT
19+ , foldrM
20+ , foldrMx
21+ , foldr
22+
23+ , foldl'
24+ , foldlM'
25+ , foldlS
26+ , foldlT
27+ , toList
28+ , toListRev
29+ , reverse
30+ , reverse'
31+
32+ , foldlx'
33+ , foldlMx'
34+ , foldOnce
35+
36+ , parselMx'
37+ , parseMany
38+ , parseIterate
39+
40+ -- ** Specialized Folds
41+ , tap
42+ , tapOffsetEvery
43+ , tapAsync
44+ , tapRate
45+ , pollCounts
46+ , drain
47+ , null
48+ , head
49+ , headElse
50+ , tail
51+ , last
52+ , elem
53+ , notElem
54+ , all
55+ , any
56+ , maximum
57+ , maximumBy
58+ , minimum
59+ , minimumBy
60+ , findIndices
61+ , lookup
62+ , findM
63+ , find
64+ , (!!)
65+ , toSVarParallel
66+ )
67+where
68+
69+import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
70+import Control.Exception
71+ (assert, AsyncException)
72+import Control.Monad (void, when)
73+import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
74+import Control.Monad.IO.Class (MonadIO(..))
75+import Control.Monad.Trans.Class (MonadTrans(lift))
76+import Data.IORef (newIORef, mkWeakIORef)
77+import Foreign.Storable (Storable(..))
78+import GHC.Exts (SpecConstrAnnotation(..))
79+import GHC.Types (SPEC(..))
80+import Fusion.Plugin.Types (Fuse(..))
81+import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
82+import Streamly.Internal.Data.Fold.Types (Fold(..))
83+import Streamly.Internal.Data.Parser (ParseError(..))
84+import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
85+import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
86+
87+import qualified Streamly.Internal.Data.IORef.Prim as Prim
88+import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
89+import qualified Streamly.Internal.Data.Fold as FL
90+import qualified Streamly.Internal.Data.Stream.StreamK as K
91+import qualified Streamly.Internal.Data.Parser as PR
92+import qualified Streamly.Internal.Data.Parser.ParserD as PRD
93+import qualified Control.Monad.Catch as MC
94+import qualified Prelude
95+
96+import Prelude hiding
97+ ( map, mapM, mapM_, repeat, foldr, last, take, filter
98+ , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
99+ , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
100+ , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
101+ , reverse, iterate, splitAt)
102+import Streamly.Internal.Data.Stream.StreamD.Type
103+import Streamly.Internal.Data.Stream.StreamD.Common
104+import Streamly.Internal.Data.SVar
105+
106+------------------------------------------------------------------------------
107+-- Elimination by Folds
108+------------------------------------------------------------------------------
109+
110+------------------------------------------------------------------------------
111+-- Left Folds
112+------------------------------------------------------------------------------
113+
114+{-# INLINE_NORMAL foldOnce #-}
115+foldOnce :: (Monad m) => Fold m a b -> Stream m a -> m b
116+foldOnce (Fold fstep begin done) (Stream step state) =
117+ begin >>= \x -> go SPEC x state
118+
119+ where
120+
121+ {-# INLINE go #-}
122+ go !_ !fs st = do
123+ r <- step defState st
124+ case r of
125+ Yield x s -> do
126+ res <- fstep fs x
127+ case res of
128+ FL.Done b -> return b
129+ FL.Partial fs1 -> go SPEC fs1 s
130+ Skip s -> go SPEC fs s
131+ Stop -> done fs
132+
133+{-# INLINE_NORMAL foldlT #-}
134+foldlT :: (Monad m, Monad (s m), MonadTrans s)
135+ => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
136+foldlT fstep begin (Stream step state) = go SPEC begin state
137+ where
138+ go !_ acc st = do
139+ r <- lift $ step defState st
140+ case r of
141+ Yield x s -> go SPEC (fstep acc x) s
142+ Skip s -> go SPEC acc s
143+ Stop -> acc
144+
145+-- Note, this is going to have horrible performance, because of the nature of
146+-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
147+-- likely be practically unusable.
148+{-# INLINE_NORMAL foldlS #-}
149+foldlS :: Monad m
150+ => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
151+foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
152+ where
153+ step' gst (Left (st, acc)) = do
154+ r <- step (adaptState gst) st
155+ return $ case r of
156+ Yield x s -> Skip (Left (s, fstep acc x))
157+ Skip s -> Skip (Left (s, acc))
158+ Stop -> Skip (Right acc)
159+
160+ step' gst (Right (Stream stp stt)) = do
161+ r <- stp (adaptState gst) stt
162+ return $ case r of
163+ Yield x s -> Yield x (Right (Stream stp s))
164+ Skip s -> Skip (Right (Stream stp s))
165+ Stop -> Stop
166+
167+------------------------------------------------------------------------------
168+-- Parses
169+------------------------------------------------------------------------------
170+
171+-- Inlined definition. Without the inline "serially/parser/take" benchmark
172+-- degrades and parseMany does not fuse. Even using "inline" at the callsite
173+-- does not help.
174+{-# INLINE splitAt #-}
175+splitAt :: Int -> [a] -> ([a],[a])
176+splitAt n ls
177+ | n <= 0 = ([], ls)
178+ | otherwise = splitAt' n ls
179+ where
180+ splitAt' :: Int -> [a] -> ([a], [a])
181+ splitAt' _ [] = ([], [])
182+ splitAt' 1 (x:xs) = ([x], xs)
183+ splitAt' m (x:xs) = (x:xs', xs'')
184+ where
185+ (xs', xs'') = splitAt' (m - 1) xs
186+
187+-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
188+-- to make a newtype.
189+{-# ANN type List NoSpecConstr #-}
190+newtype List a = List {getList :: [a]}
191+
192+-- | Run a 'Parse' over a stream.
193+{-# INLINE_NORMAL parselMx' #-}
194+parselMx'
195+ :: MonadThrow m
196+ => (s -> a -> m (PR.Step s b))
197+ -> m s
198+ -> (s -> m b)
199+ -> Stream m a
200+ -> m b
201+parselMx' pstep initial extract (Stream step state) = do
202+ initial >>= go SPEC state (List [])
203+
204+ where
205+
206+ -- XXX currently we are using a dumb list based approach for backtracking
207+ -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
208+ -- That will allow us more efficient random back and forth movement.
209+ {-# INLINE go #-}
210+ go !_ st buf !pst = do
211+ r <- step defState st
212+ case r of
213+ Yield x s -> do
214+ pRes <- pstep pst x
215+ case pRes of
216+ PR.Partial 0 pst1 -> go SPEC s (List []) pst1
217+ PR.Partial n pst1 -> do
218+ assert (n <= length (x:getList buf)) (return ())
219+ let src0 = Prelude.take n (x:getList buf)
220+ src = Prelude.reverse src0
221+ gobuf SPEC s (List []) (List src) pst1
222+ PR.Continue 0 pst1 -> go SPEC s (List (x:getList buf)) pst1
223+ PR.Continue n pst1 -> do
224+ assert (n <= length (x:getList buf)) (return ())
225+ let (src0, buf1) = splitAt n (x:getList buf)
226+ src = Prelude.reverse src0
227+ gobuf SPEC s (List buf1) (List src) pst1
228+ PR.Done _ b -> return b
229+ PR.Error err -> throwM $ ParseError err
230+ Skip s -> go SPEC s buf pst
231+ Stop -> extract pst
232+
233+ gobuf !_ s buf (List []) !pst = go SPEC s buf pst
234+ gobuf !_ s buf (List (x:xs)) !pst = do
235+ pRes <- pstep pst x
236+ case pRes of
237+ PR.Partial 0 pst1 ->
238+ gobuf SPEC s (List []) (List xs) pst1
239+ PR.Partial n pst1 -> do
240+ assert (n <= length (x:getList buf)) (return ())
241+ let src0 = Prelude.take n (x:getList buf)
242+ src = Prelude.reverse src0 ++ xs
243+ gobuf SPEC s (List []) (List src) pst1
244+ PR.Continue 0 pst1 ->
245+ gobuf SPEC s (List (x:getList buf)) (List xs) pst1
246+ PR.Continue n pst1 -> do
247+ assert (n <= length (x:getList buf)) (return ())
248+ let (src0, buf1) = splitAt n (x:getList buf)
249+ src = Prelude.reverse src0 ++ xs
250+ gobuf SPEC s (List buf1) (List src) pst1
251+ PR.Done _ b -> return b
252+ PR.Error err -> throwM $ ParseError err
253+
254+------------------------------------------------------------------------------
255+-- Repeated parsing
256+------------------------------------------------------------------------------
257+
258+{-# ANN type ParseChunksState Fuse #-}
259+data ParseChunksState x inpBuf st pst =
260+ ParseChunksInit inpBuf st
261+ | ParseChunksInitLeftOver inpBuf
262+ | ParseChunksStream st inpBuf !pst
263+ | ParseChunksBuf inpBuf st inpBuf !pst
264+ | ParseChunksYield x (ParseChunksState x inpBuf st pst)
265+
266+{-# INLINE_NORMAL parseMany #-}
267+parseMany
268+ :: MonadThrow m
269+ => PRD.Parser m a b
270+ -> Stream m a
271+ -> Stream m b
272+parseMany (PRD.Parser pstep initial extract) (Stream step state) =
273+ Stream stepOuter (ParseChunksInit [] state)
274+
275+ where
276+
277+ {-# INLINE_LATE stepOuter #-}
278+ -- Buffer is empty, get the first element from the stream, initialize the
279+ -- fold and then go to stream processing loop.
280+ stepOuter gst (ParseChunksInit [] st) = do
281+ r <- step (adaptState gst) st
282+ case r of
283+ Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
284+ Skip s -> return $ Skip $ ParseChunksInit [] s
285+ Stop -> return Stop
286+
287+ -- Buffer is not empty, go to buffered processing loop
288+ stepOuter _ (ParseChunksInit src st) = do
289+ initial >>= return . Skip . ParseChunksBuf src st []
290+
291+ -- XXX we just discard any leftover input at the end
292+ stepOuter _ (ParseChunksInitLeftOver _) = return Stop
293+
294+ -- Buffer is empty, process elements from the stream
295+ stepOuter gst (ParseChunksStream st buf pst) = do
296+ r <- step (adaptState gst) st
297+ case r of
298+ Yield x s -> do
299+ pRes <- pstep pst x
300+ case pRes of
301+ PR.Partial 0 pst1 ->
302+ return $ Skip $ ParseChunksStream s [] pst1
303+ PR.Partial n pst1 -> do
304+ assert (n <= length (x:buf)) (return ())
305+ let src0 = Prelude.take n (x:buf)
306+ src = Prelude.reverse src0
307+ return $ Skip $ ParseChunksBuf src s [] pst1
308+ PR.Continue 0 pst1 ->
309+ return $ Skip $ ParseChunksStream s (x:buf) pst1
310+ PR.Continue n pst1 -> do
311+ assert (n <= length (x:buf)) (return ())
312+ let (src0, buf1) = splitAt n (x:buf)
313+ src = Prelude.reverse src0
314+ return $ Skip $ ParseChunksBuf src s buf1 pst1
315+ PR.Done 0 b -> do
316+ return $ Skip $
317+ ParseChunksYield b (ParseChunksInit [] s)
318+ PR.Done n b -> do
319+ assert (n <= length (x:buf)) (return ())
320+ let src = Prelude.reverse (Prelude.take n (x:buf))
321+ return $ Skip $
322+ ParseChunksYield b (ParseChunksInit src s)
323+ PR.Error err -> throwM $ ParseError err
324+ Skip s -> return $ Skip $ ParseChunksStream s buf pst
325+ Stop -> do
326+ b <- extract pst
327+ let src = Prelude.reverse buf
328+ return $ Skip $
329+ ParseChunksYield b (ParseChunksInitLeftOver src)
330+
331+ -- go back to stream processing mode
332+ stepOuter _ (ParseChunksBuf [] s buf pst) =
333+ return $ Skip $ ParseChunksStream s buf pst
334+
335+ -- buffered processing loop
336+ stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
337+ pRes <- pstep pst x
338+ case pRes of
339+ PR.Partial 0 pst1 ->
340+ return $ Skip $ ParseChunksBuf xs s [] pst1
341+ PR.Partial n pst1 -> do
342+ assert (n <= length (x:buf)) (return ())
343+ let src0 = Prelude.take n (x:buf)
344+ src = Prelude.reverse src0 ++ xs
345+ return $ Skip $ ParseChunksBuf src s [] pst1
346+ PR.Continue 0 pst1 ->
347+ return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
348+ PR.Continue n pst1 -> do
349+ assert (n <= length (x:buf)) (return ())
350+ let (src0, buf1) = splitAt n (x:buf)
351+ src = Prelude.reverse src0 ++ xs
352+ return $ Skip $ ParseChunksBuf src s buf1 pst1
353+ PR.Done 0 b ->
354+ return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
355+ PR.Done n b -> do
356+ assert (n <= length (x:buf)) (return ())
357+ let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
358+ return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
359+ PR.Error err -> throwM $ ParseError err
360+
361+ stepOuter _ (ParseChunksYield a next) = return $ Yield a next
362+
363+{-# ANN type ConcatParseState Fuse #-}
364+data ConcatParseState x inpBuf st p =
365+ ConcatParseInit inpBuf st p
366+ | ConcatParseInitLeftOver inpBuf
367+ | ConcatParseStream st inpBuf p
368+ | ConcatParseBuf inpBuf st inpBuf p
369+ | ConcatParseYield x (ConcatParseState x inpBuf st p)
370+
371+{-# INLINE_NORMAL parseIterate #-}
372+parseIterate
373+ :: MonadThrow m
374+ => (b -> PRD.Parser m a b)
375+ -> b
376+ -> Stream m a
377+ -> Stream m b
378+parseIterate func seed (Stream step state) =
379+ Stream stepOuter (ConcatParseInit [] state (func seed))
380+
381+ where
382+
383+ {-# INLINE_LATE stepOuter #-}
384+ -- Buffer is empty, go to stream processing loop
385+ stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
386+ initial >>= \r -> return $ Skip $ ConcatParseStream st []
387+ (PRD.Parser pstep (return r) extract)
388+
389+ -- Buffer is not empty, go to buffered processing loop
390+ stepOuter _ (ConcatParseInit src st
391+ (PRD.Parser pstep initial extract)) = do
392+ initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
393+ (PRD.Parser pstep (return r) extract)
394+
395+ -- XXX we just discard any leftover input at the end
396+ stepOuter _ (ConcatParseInitLeftOver _) = return Stop
397+
398+ -- Buffer is empty process elements from the stream
399+ stepOuter gst (ConcatParseStream st buf
400+ p@(PRD.Parser pstep initial extract)) = do
401+ r <- step (adaptState gst) st
402+ case r of
403+ Yield x s -> do
404+ pst <- initial
405+ pRes <- pstep pst x
406+ case pRes of
407+ PR.Partial 0 pst1 ->
408+ return $ Skip $ ConcatParseStream s []
409+ (PRD.Parser pstep (return pst1) extract)
410+ PR.Partial n pst1 -> do
411+ assert (n <= length (x:buf)) (return ())
412+ let src0 = Prelude.take n (x:buf)
413+ src = Prelude.reverse src0
414+ return $ Skip $ ConcatParseBuf src s []
415+ (PRD.Parser pstep (return pst1) extract)
416+ -- PR.Continue 0 pst1 ->
417+ -- return $ Skip $ ConcatParseStream s (x:buf) pst1
418+ PR.Continue n pst1 -> do
419+ assert (n <= length (x:buf)) (return ())
420+ let (src0, buf1) = splitAt n (x:buf)
421+ src = Prelude.reverse src0
422+ return $ Skip $ ConcatParseBuf src s buf1
423+ (PRD.Parser pstep (return pst1) extract)
424+ -- XXX Specialize for Stop 0 common case?
425+ PR.Done n b -> do
426+ assert (n <= length (x:buf)) (return ())
427+ let src = Prelude.reverse (Prelude.take n (x:buf))
428+ return $ Skip $
429+ ConcatParseYield b (ConcatParseInit src s (func b))
430+ PR.Error err -> throwM $ ParseError err
431+ Skip s -> return $ Skip $ ConcatParseStream s buf p
432+ Stop -> do
433+ pst <- initial
434+ b <- extract pst
435+ let src = Prelude.reverse buf
436+ return $ Skip $
437+ ConcatParseYield b (ConcatParseInitLeftOver src)
438+
439+ -- go back to stream processing mode
440+ stepOuter _ (ConcatParseBuf [] s buf p) =
441+ return $ Skip $ ConcatParseStream s buf p
442+
443+ -- buffered processing loop
444+ stepOuter _ (ConcatParseBuf (x:xs) s buf
445+ (PRD.Parser pstep initial extract)) = do
446+ pst <- initial
447+ pRes <- pstep pst x
448+ case pRes of
449+ PR.Partial 0 pst1 ->
450+ return $ Skip $ ConcatParseBuf xs s []
451+ (PRD.Parser pstep (return pst1) extract)
452+ PR.Partial n pst1 -> do
453+ assert (n <= length (x:buf)) (return ())
454+ let src0 = Prelude.take n (x:buf)
455+ src = Prelude.reverse src0 ++ xs
456+ return $ Skip $ ConcatParseBuf src s []
457+ (PRD.Parser pstep (return pst1) extract)
458+ -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
459+ PR.Continue n pst1 -> do
460+ assert (n <= length (x:buf)) (return ())
461+ let (src0, buf1) = splitAt n (x:buf)
462+ src = Prelude.reverse src0 ++ xs
463+ return $ Skip $ ConcatParseBuf src s buf1
464+ (PRD.Parser pstep (return pst1) extract)
465+ -- XXX Specialize for Stop 0 common case?
466+ PR.Done n b -> do
467+ assert (n <= length (x:buf)) (return ())
468+ let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
469+ return $ Skip $ ConcatParseYield b
470+ (ConcatParseInit src s (func b))
471+ PR.Error err -> throwM $ ParseError err
472+
473+ stepOuter _ (ConcatParseYield a next) = return $ Yield a next
474+
475+------------------------------------------------------------------------------
476+-- Specialized Folds
477+------------------------------------------------------------------------------
478+
479+-- | Run a streaming composition, discard the results.
480+{-# INLINE_LATE drain #-}
481+drain :: Monad m => Stream m a -> m ()
482+-- drain = foldrM (\_ xs -> xs) (return ())
483+drain (Stream step state) = go SPEC state
484+ where
485+ go !_ st = do
486+ r <- step defState st
487+ case r of
488+ Yield _ s -> go SPEC s
489+ Skip s -> go SPEC s
490+ Stop -> return ()
491+
492+{-# INLINE_NORMAL null #-}
493+null :: Monad m => Stream m a -> m Bool
494+null m = foldrM (\_ _ -> return False) (return True) m
495+
496+{-# INLINE_NORMAL head #-}
497+head :: Monad m => Stream m a -> m (Maybe a)
498+head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
499+
500+{-# INLINE_NORMAL headElse #-}
501+headElse :: Monad m => a -> Stream m a -> m a
502+headElse a m = foldrM (\x _ -> return x) (return a) m
503+
504+-- Does not fuse, has the same performance as the StreamK version.
505+{-# INLINE_NORMAL tail #-}
506+tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
507+tail (UnStream step state) = go state
508+ where
509+ go st = do
510+ r <- step defState st
511+ case r of
512+ Yield _ s -> return (Just $ Stream step s)
513+ Skip s -> go s
514+ Stop -> return Nothing
515+
516+-- XXX will it fuse? need custom impl?
517+{-# INLINE_NORMAL last #-}
518+last :: Monad m => Stream m a -> m (Maybe a)
519+last = foldl' (\_ y -> Just y) Nothing
520+
521+{-# INLINE_NORMAL elem #-}
522+elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
523+-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
524+elem e (Stream step state) = go state
525+ where
526+ go st = do
527+ r <- step defState st
528+ case r of
529+ Yield x s
530+ | x == e -> return True
531+ | otherwise -> go s
532+ Skip s -> go s
533+ Stop -> return False
534+
535+{-# INLINE_NORMAL notElem #-}
536+notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
537+notElem e s = fmap not (elem e s)
538+
539+{-# INLINE_NORMAL all #-}
540+all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
541+-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
542+all p (Stream step state) = go state
543+ where
544+ go st = do
545+ r <- step defState st
546+ case r of
547+ Yield x s
548+ | p x -> go s
549+ | otherwise -> return False
550+ Skip s -> go s
551+ Stop -> return True
552+
553+{-# INLINE_NORMAL any #-}
554+any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
555+-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
556+any p (Stream step state) = go state
557+ where
558+ go st = do
559+ r <- step defState st
560+ case r of
561+ Yield x s
562+ | p x -> return True
563+ | otherwise -> go s
564+ Skip s -> go s
565+ Stop -> return False
566+
567+{-# INLINE_NORMAL maximum #-}
568+maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
569+maximum (Stream step state) = go Nothing state
570+ where
571+ go Nothing st = do
572+ r <- step defState st
573+ case r of
574+ Yield x s -> go (Just x) s
575+ Skip s -> go Nothing s
576+ Stop -> return Nothing
577+ go (Just acc) st = do
578+ r <- step defState st
579+ case r of
580+ Yield x s
581+ | acc <= x -> go (Just x) s
582+ | otherwise -> go (Just acc) s
583+ Skip s -> go (Just acc) s
584+ Stop -> return (Just acc)
585+
586+{-# INLINE_NORMAL maximumBy #-}
587+maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
588+maximumBy cmp (Stream step state) = go Nothing state
589+ where
590+ go Nothing st = do
591+ r <- step defState st
592+ case r of
593+ Yield x s -> go (Just x) s
594+ Skip s -> go Nothing s
595+ Stop -> return Nothing
596+ go (Just acc) st = do
597+ r <- step defState st
598+ case r of
599+ Yield x s -> case cmp acc x of
600+ GT -> go (Just acc) s
601+ _ -> go (Just x) s
602+ Skip s -> go (Just acc) s
603+ Stop -> return (Just acc)
604+
605+{-# INLINE_NORMAL minimum #-}
606+minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
607+minimum (Stream step state) = go Nothing state
608+ where
609+ go Nothing st = do
610+ r <- step defState st
611+ case r of
612+ Yield x s -> go (Just x) s
613+ Skip s -> go Nothing s
614+ Stop -> return Nothing
615+ go (Just acc) st = do
616+ r <- step defState st
617+ case r of
618+ Yield x s
619+ | acc <= x -> go (Just acc) s
620+ | otherwise -> go (Just x) s
621+ Skip s -> go (Just acc) s
622+ Stop -> return (Just acc)
623+
624+{-# INLINE_NORMAL minimumBy #-}
625+minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
626+minimumBy cmp (Stream step state) = go Nothing state
627+ where
628+ go Nothing st = do
629+ r <- step defState st
630+ case r of
631+ Yield x s -> go (Just x) s
632+ Skip s -> go Nothing s
633+ Stop -> return Nothing
634+ go (Just acc) st = do
635+ r <- step defState st
636+ case r of
637+ Yield x s -> case cmp acc x of
638+ GT -> go (Just x) s
639+ _ -> go (Just acc) s
640+ Skip s -> go (Just acc) s
641+ Stop -> return (Just acc)
642+
643+{-# INLINE_NORMAL (!!) #-}
644+(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
645+(Stream step state) !! i = go i state
646+ where
647+ go n st = do
648+ r <- step defState st
649+ case r of
650+ Yield x s | n < 0 -> return Nothing
651+ | n == 0 -> return $ Just x
652+ | otherwise -> go (n - 1) s
653+ Skip s -> go n s
654+ Stop -> return Nothing
655+
656+{-# INLINE_NORMAL lookup #-}
657+lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
658+lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
659+ (return Nothing) m
660+
661+{-# INLINE_NORMAL findM #-}
662+findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
663+findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
664+ (return Nothing) m
665+
666+{-# INLINE find #-}
667+find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
668+find p = findM (return . p)
669+
670+{-# INLINE_NORMAL findIndices #-}
671+findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
672+findIndices p (Stream step state) = Stream step' (state, 0)
673+ where
674+ {-# INLINE_LATE step' #-}
675+ step' gst (st, i) = i `seq` do
676+ r <- step (adaptState gst) st
677+ return $ case r of
678+ Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
679+ Skip s -> Skip (s, i)
680+ Stop -> Stop
681+
682+{-# INLINE toListRev #-}
683+toListRev :: Monad m => Stream m a -> m [a]
684+toListRev = foldl' (flip (:)) []
685+
686+-- We can implement reverse as:
687+--
688+-- > reverse = foldlS (flip cons) nil
689+--
690+-- However, this implementation is unusable because of the horrible performance
691+-- of cons. So we just convert it to a list first and then stream from the
692+-- list.
693+--
694+-- XXX Maybe we can use an Array instead of a list here?
695+{-# INLINE_NORMAL reverse #-}
696+reverse :: Monad m => Stream m a -> Stream m a
697+reverse m = Stream step Nothing
698+ where
699+ {-# INLINE_LATE step #-}
700+ step _ Nothing = do
701+ xs <- toListRev m
702+ return $ Skip (Just xs)
703+ step _ (Just (x:xs)) = return $ Yield x (Just xs)
704+ step _ (Just []) = return Stop
705+
706+-- Much faster reverse for Storables
707+{-# INLINE_NORMAL reverse' #-}
708+reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
709+{-
710+-- This commented implementation copies the whole stream into one single array
711+-- and then streams from that array, this is 3-4x faster than the chunked code
712+-- that follows. Though this could be problematic due to unbounded large
713+-- allocations. We need to figure out why the chunked code is slower and if we
714+-- can optimize the chunked code to work as fast as this one. It may be a
715+-- fusion issue?
716+import Foreign.ForeignPtr (touchForeignPtr)
717+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
718+import Foreign.Ptr (Ptr, plusPtr)
719+reverse' m = Stream step Nothing
720+ where
721+ {-# INLINE_LATE step #-}
722+ step _ Nothing = do
723+ arr <- A.fromStreamD m
724+ let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
725+ return $ Skip $ Just (aStart arr, p)
726+
727+ step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
728+
729+ step _ (Just (start, p)) = do
730+ let !x = A.unsafeInlineIO $ do
731+ r <- peek p
732+ touchForeignPtr start
733+ return r
734+ next = p `plusPtr` negate (sizeOf (undefined :: a))
735+ return $ Yield x (Just (start, next))
736+-}
737+reverse' m =
738+ A.flattenArraysRev
739+ $ fromStreamK
740+ $ K.reverse
741+ $ toStreamK
742+ $ A.fromStreamDArraysOf A.defaultChunkSize m
743+
744+-------------------------------------------------------------------------------
745+-- Concurrent application and fold
746+-------------------------------------------------------------------------------
747+
748+-- XXX These functions should be moved to Stream/Parallel.hs
749+--
750+-- Using StreamD the worker stream producing code can fuse with the code to
751+-- queue output to the SVar giving some perf boost.
752+--
753+-- Note that StreamD can only be used in limited situations, specifically, we
754+-- cannot implement joinStreamVarPar using this.
755+--
756+-- XXX make sure that the SVar passed is a Parallel style SVar.
757+
758+-- | Fold the supplied stream to the SVar asynchronously using Parallel
759+-- concurrency style.
760+-- {-# INLINE_NORMAL toSVarParallel #-}
761+{-# INLINE toSVarParallel #-}
762+toSVarParallel :: MonadAsync m
763+ => State t m a -> SVar t m a -> Stream m a -> m ()
764+toSVarParallel st sv xs =
765+ if svarInspectMode sv
766+ then forkWithDiag
767+ else do
768+ tid <-
769+ case getYieldLimit st of
770+ Nothing -> doFork (work Nothing)
771+ (svarMrun sv)
772+ (handleChildException sv)
773+ Just _ -> doFork (workLim Nothing)
774+ (svarMrun sv)
775+ (handleChildException sv)
776+ modifyThread sv tid
777+
778+ where
779+
780+ {-# NOINLINE work #-}
781+ work info = (foldOnce (FL.toParallelSVar sv info) xs)
782+
783+ {-# NOINLINE workLim #-}
784+ workLim info = foldOnce (FL.toParallelSVarLimited sv info) xs
785+
786+ {-# NOINLINE forkWithDiag #-}
787+ forkWithDiag = do
788+ -- We do not use workerCount in case of ParallelVar but still there is
789+ -- no harm in maintaining it correctly.
790+ liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
791+ recordMaxWorkers sv
792+ -- This allocation matters when significant number of workers are being
793+ -- sent. We allocate it only when needed. The overhead increases by 4x.
794+ winfo <-
795+ case yieldRateInfo sv of
796+ Nothing -> return Nothing
797+ Just _ -> liftIO $ do
798+ cntRef <- newIORef 0
799+ t <- getTime Monotonic
800+ lat <- newIORef (0, t)
801+ return $ Just WorkerInfo
802+ { workerYieldMax = 0
803+ , workerYieldCount = cntRef
804+ , workerLatencyStart = lat
805+ }
806+ tid <-
807+ case getYieldLimit st of
808+ Nothing -> doFork (work winfo)
809+ (svarMrun sv)
810+ (handleChildException sv)
811+ Just _ -> doFork (workLim winfo)
812+ (svarMrun sv)
813+ (handleChildException sv)
814+ modifyThread sv tid
815+
816+------------------------------------------------------------------------------
817+-- Tapping/Distributing
818+------------------------------------------------------------------------------
819+
820+data TapState fs st a
821+ = TapInit | Tapping !fs st | TapDone st
822+
823+-- XXX Multiple yield points
824+{-# INLINE tap #-}
825+tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
826+tap (Fold fstep initial extract) (Stream step state) = Stream step' TapInit
827+
828+ where
829+
830+ step' _ TapInit = do
831+ r <- initial
832+ return $ Skip (Tapping r state)
833+ step' gst (Tapping acc st) = do
834+ r <- step gst st
835+ case r of
836+ -- XXX Abstract Yield?
837+ Yield x s -> do
838+ res <- fstep acc x
839+ return
840+ $ Yield x
841+ $ case res of
842+ FL.Partial fs -> Tapping fs s
843+ FL.Done _ -> TapDone s
844+ Skip s -> return $ Skip (Tapping acc s)
845+ Stop -> do
846+ void $ extract acc
847+ return Stop
848+ step' gst (TapDone st) = do
849+ r <- step gst st
850+ return
851+ $ case r of
852+ Yield x s -> Yield x (TapDone s)
853+ Skip s -> Skip (TapDone s)
854+ Stop -> Stop
855+
856+data TapOffState fs s a
857+ = TapOffInit
858+ | TapOffTapping !fs s Int
859+ | TapOffDone s
860+
861+-- XXX Multiple yield points
862+{-# INLINE_NORMAL tapOffsetEvery #-}
863+tapOffsetEvery :: Monad m
864+ => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
865+tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
866+ Stream step' TapOffInit
867+
868+ where
869+
870+ {-# INLINE_LATE step' #-}
871+ step' _ TapOffInit = do
872+ r <- initial
873+ return $ Skip $ TapOffTapping r state (offset `mod` n)
874+ step' gst (TapOffTapping acc st count) = do
875+ r <- step gst st
876+ case r of
877+ Yield x s -> do
878+ next <-
879+ if count <= 0
880+ then do
881+ res <- fstep acc x
882+ return
883+ $ case res of
884+ FL.Partial sres ->
885+ TapOffTapping sres s (n - 1)
886+ FL.Done _ -> TapOffDone s
887+ else return $ TapOffTapping acc s (count - 1)
888+ return $ Yield x next
889+ Skip s -> return $ Skip (TapOffTapping acc s count)
890+ Stop -> do
891+ void $ extract acc
892+ return Stop
893+ step' gst (TapOffDone st) = do
894+ r <- step gst st
895+ return
896+ $ case r of
897+ Yield x s -> Yield x (TapOffDone s)
898+ Skip s -> Skip (TapOffDone s)
899+ Stop -> Stop
900+
901+{-# INLINE_NORMAL pollCounts #-}
902+pollCounts
903+ :: MonadAsync m
904+ => (a -> Bool)
905+ -> (Stream m Int -> Stream m Int)
906+ -> Fold m Int b
907+ -> Stream m a
908+ -> Stream m a
909+pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
910+ where
911+
912+ {-# INLINE_LATE step' #-}
913+ step' _ Nothing = do
914+ -- As long as we are using an "Int" for counts lockfree reads from
915+ -- Var should work correctly on both 32-bit and 64-bit machines.
916+ -- However, an Int on a 32-bit machine may overflow quickly.
917+ countVar <- liftIO $ Prim.newIORef (0 :: Int)
918+ tid <- forkManaged
919+ $ void $ foldOnce fld
920+ $ transf $ fromPrimIORef countVar
921+ return $ Skip (Just (countVar, tid, state))
922+
923+ step' gst (Just (countVar, tid, st)) = do
924+ r <- step gst st
925+ case r of
926+ Yield x s -> do
927+ when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
928+ return $ Yield x (Just (countVar, tid, s))
929+ Skip s -> return $ Skip (Just (countVar, tid, s))
930+ Stop -> do
931+ liftIO $ killThread tid
932+ return Stop
933+
934+{-# INLINE_NORMAL tapRate #-}
935+tapRate ::
936+ (MonadAsync m, MonadCatch m)
937+ => Double
938+ -> (Int -> m b)
939+ -> Stream m a
940+ -> Stream m a
941+tapRate samplingRate action (Stream step state) = Stream step' Nothing
942+ where
943+ {-# NOINLINE loop #-}
944+ loop countVar prev = do
945+ i <-
946+ MC.catch
947+ (do liftIO $ threadDelay (round $ samplingRate * 1000000)
948+ i <- liftIO $ Prim.readIORef countVar
949+ let !diff = i - prev
950+ void $ action diff
951+ return i)
952+ (\(e :: AsyncException) -> do
953+ i <- liftIO $ Prim.readIORef countVar
954+ let !diff = i - prev
955+ void $ action diff
956+ throwM (MC.toException e))
957+ loop countVar i
958+
959+ {-# INLINE_LATE step' #-}
960+ step' _ Nothing = do
961+ countVar <- liftIO $ Prim.newIORef 0
962+ tid <- fork $ loop countVar 0
963+ ref <- liftIO $ newIORef ()
964+ _ <- liftIO $ mkWeakIORef ref (killThread tid)
965+ return $ Skip (Just (countVar, tid, state, ref))
966+
967+ step' gst (Just (countVar, tid, st, ref)) = do
968+ r <- step gst st
969+ case r of
970+ Yield x s -> do
971+ liftIO $ Prim.modifyIORef' countVar (+ 1)
972+ return $ Yield x (Just (countVar, tid, s, ref))
973+ Skip s -> return $ Skip (Just (countVar, tid, s, ref))
974+ Stop -> do
975+ liftIO $ killThread tid
976+ return Stop
977+
978+-------------------------------------------------------------------------------
979+-- Concurrent tap
980+-------------------------------------------------------------------------------
981+
982+-- | Create an SVar with a fold consumer that will fold any elements sent to it
983+-- using the supplied fold function.
984+{-# INLINE newFoldSVar #-}
985+newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
986+newFoldSVar stt f = do
987+ -- Buffer size for the SVar is derived from the current state
988+ sv <- newParallelVar StopAny (adaptState stt)
989+ -- Add the producer thread-id to the SVar.
990+ liftIO myThreadId >>= modifyThread sv
991+ void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
992+ return sv
993+
994+ where
995+
996+ {-# NOINLINE work #-}
997+ work sv = void $ foldOnce f $ fromProducer sv
998+
999+{-# INLINE_NORMAL tapAsync #-}
1000+tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
1001+tapAsync f (Stream step1 state1) = Stream step TapInit
1002+ where
1003+
1004+ drainFold svr = do
1005+ -- In general, a Stop event would come equipped with the result
1006+ -- of the fold. It is not used here but it would be useful in
1007+ -- applicative and distribute.
1008+ done <- fromConsumer svr
1009+ when (not done) $ do
1010+ liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
1011+ $ takeMVar (outputDoorBellFromConsumer svr)
1012+ drainFold svr
1013+
1014+ stopFold svr = do
1015+ liftIO $ sendStop svr Nothing
1016+ -- drain/wait until a stop event arrives from the fold.
1017+ drainFold svr
1018+
1019+ {-# INLINE_LATE step #-}
1020+ step gst TapInit = do
1021+ sv <- newFoldSVar gst f
1022+ return $ Skip (Tapping sv state1)
1023+
1024+ step gst (Tapping sv st) = do
1025+ r <- step1 gst st
1026+ case r of
1027+ Yield a s -> do
1028+ done <- pushToFold sv a
1029+ if done
1030+ then do
1031+ -- XXX we do not need to wait synchronously here
1032+ stopFold sv
1033+ return $ Yield a (TapDone s)
1034+ else return $ Yield a (Tapping sv s)
1035+ Skip s -> return $ Skip (Tapping sv s)
1036+ Stop -> do
1037+ stopFold sv
1038+ return $ Stop
1039+
1040+ step gst (TapDone st) = do
1041+ r <- step1 gst st
1042+ return $ case r of
1043+ Yield a s -> Yield a (TapDone s)
1044+ Skip s -> Skip (TapDone s)
1045+ Stop -> Stop
src/Streamly/Internal/Data/Stream/StreamD/Generate.hs+1−52
@@ -67,7 +67,6 @@ import Data.Int (Int64)
67import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
67import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
68import Data.Maybe (isNothing)
68import Data.Maybe (isNothing)
69import System.Mem (performMajorGC)
69import System.Mem (performMajorGC)
70−import Streamly.Internal.Data.IORef.Prim (Prim)
71import Streamly.Internal.Data.Time.Units
70import Streamly.Internal.Data.Time.Units
72 (toRelTime64, RelTime64)
71 (toRelTime64, RelTime64)
73import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
72import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
@@ -86,6 +85,7 @@ import Prelude hiding
86 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
85 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
87 , reverse, iterate, splitAt)
86 , reverse, iterate, splitAt)
88import Streamly.Internal.Data.Stream.StreamD.Type
87import Streamly.Internal.Data.Stream.StreamD.Type
88+import Streamly.Internal.Data.Stream.StreamD.Common
89import Streamly.Internal.Data.SVar
89import Streamly.Internal.Data.SVar
90
90
91
91
@@ -363,23 +363,10 @@ fromListM = Stream step
363toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
363toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
364toStreamD = fromStreamK . K.toStream
364toStreamD = fromStreamK . K.toStream
365
365
366−{-# INLINE_NORMAL fromPrimIORef #-}
367−fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
368−fromPrimIORef var = Stream step ()
369− where
370− {-# INLINE_LATE step #-}
371− step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
372−
373-------------------------------------------------------------------------------
366-------------------------------------------------------------------------------
374-- Generation from SVar
367-- Generation from SVar
375-------------------------------------------------------------------------------
368-------------------------------------------------------------------------------
376
369
377−data FromSVarState t m a =
378− FromSVarInit
379− | FromSVarRead (SVar t m a)
380− | FromSVarLoop (SVar t m a) [ChildEvent a]
381− | FromSVarDone (SVar t m a)
382−
383{-# INLINE_NORMAL fromSVar #-}
370{-# INLINE_NORMAL fromSVar #-}
384fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
371fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
385fromSVar svar = Stream step FromSVarInit
372fromSVar svar = Stream step FromSVarInit
@@ -454,41 +441,3 @@ fromSVar svar = Stream step FromSVarInit
454 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
441 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
455 liftIO $ printSVar sv "SVar Done"
442 liftIO $ printSVar sv "SVar Done"
456 return Stop
443 return Stop
457−
458−-------------------------------------------------------------------------------
459−-- Process events received by a fold consumer from a stream producer
460−-------------------------------------------------------------------------------
461−
462−{-# INLINE_NORMAL fromProducer #-}
463−fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
464−fromProducer svar = Stream step (FromSVarRead svar)
465− where
466−
467− {-# INLINE_LATE step #-}
468− step _ (FromSVarRead sv) = do
469− list <- readOutputQ sv
470− -- Reversing the output is important to guarantee that we process the
471− -- outputs in the same order as they were generated by the constituent
472− -- streams.
473− return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
474−
475− step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
476− step _ (FromSVarLoop sv (ev : es)) = do
477− case ev of
478− ChildYield a -> return $ Yield a (FromSVarLoop sv es)
479− ChildStop tid e -> do
480− accountThread sv tid
481− case e of
482− Nothing -> do
483− sendStopToProducer sv
484− return $ Skip (FromSVarDone sv)
485− Just _ -> error "Bug: fromProducer: received exception"
486−
487− step _ (FromSVarDone sv) = do
488− when (svarInspectMode sv) $ do
489− t <- liftIO $ getTime Monotonic
490− liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
491− liftIO $ printSVar sv "SVar Done"
492− return Stop
493−
494− step _ FromSVarInit = undefined
streamly.cabal+2
@@ -452,8 +452,10 @@ library
452 , Streamly.Internal.Data.Stream.StreamK
452 , Streamly.Internal.Data.Stream.StreamK
453 , Streamly.Internal.Data.Stream.StreamD.Type
453 , Streamly.Internal.Data.Stream.StreamD.Type
454 , Streamly.Internal.Data.Stream.StreamD
454 , Streamly.Internal.Data.Stream.StreamD
455+ , Streamly.Internal.Data.Stream.StreamD.Common
455 , Streamly.Internal.Data.Stream.StreamD.Generate
456 , Streamly.Internal.Data.Stream.StreamD.Generate
456 , Streamly.Internal.Data.Stream.StreamD.SplitGroup
457 , Streamly.Internal.Data.Stream.StreamD.SplitGroup
458+ , Streamly.Internal.Data.Stream.StreamD.Eliminate
457 , Streamly.Internal.Data.Stream.StreamDK.Type
459 , Streamly.Internal.Data.Stream.StreamDK.Type
458 , Streamly.Internal.Data.Stream.StreamDK
460 , Streamly.Internal.Data.Stream.StreamDK
459 , Streamly.Internal.Data.Stream.Enumeration
461 , Streamly.Internal.Data.Stream.Enumeration