@@ -21,10 +21,19 @@ module Streamly.Internal.Data.Stream.StreamD.Type
2121 , fromStreamK
2222 , toStreamK
2323 , fromStreamD
24− , map
25− , mapM
24+ , toStreamD
25+ , fromProducer
26+ , fromPrimIORef
27+ , fromSVar
28+ , foldOnce
29+ , toSVarParallel
30+ , uncons
31+ , nilM
32+ , consM
2633 , yield
2734 , yieldM
35+ , map
36+ , mapM
2837 , concatMap
2938 , concatMapM
3039
@@ -45,6 +54,8 @@ module Streamly.Internal.Data.Stream.StreamD.Type
4554 , eqBy
4655 , cmpBy
4756 , take
57+ , takeWhileM
58+ , takeWhile
4859 , GroupState (..) -- for inspection testing
4960 , foldMany
5061 , foldMany1
@@ -52,21 +63,31 @@ module Streamly.Internal.Data.Stream.StreamD.Type
5263 )
5364where
5465
66+
5567import Control.Applicative (liftA2)
68+import Control.Exception (fromException)
5669import Control.Monad (when)
5770import Control.Monad.Catch (MonadThrow, throwM)
71+import Control.Monad.IO.Class (MonadIO(..))
5872import Control.Monad.Trans.Class (lift, MonadTrans)
73+import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
5974import Data.Functor.Identity (Identity(..))
75+import Data.Maybe (isNothing)
76+import Fusion.Plugin.Types (Fuse(..))
6077import GHC.Base (build)
6178import GHC.Types (SPEC(..))
62−import Prelude hiding (map, mapM, foldr, take, concatMap)
63−import Fusion.Plugin.Types (Fuse(..))
79+import Prelude hiding (map, mapM, foldr, take, concatMap, takeWhile)
80+import System.Mem (performMajorGC)
6481
65−import Streamly.Internal.Data.SVar (State(..), adaptState, defState)
82+import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
83+import Streamly.Internal.Data.IORef.Prim (Prim)
84+import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
6685import Streamly.Internal.Data.Fold.Types (Fold(..), Fold2(..))
86+import Streamly.Internal.Data.SVar
6787
68−import qualified Streamly.Internal.Data.Fold.Types as FL
88+import qualified Streamly.Internal.Data.Fold as FL
6989import qualified Streamly.Internal.Data.Stream.StreamK as K
90+import qualified Streamly.Internal.Data.IORef.Prim as Prim
7091
7192------------------------------------------------------------------------------
7293-- The direct style stream type
@@ -141,6 +162,49 @@ toStreamK (Stream step state) = go state
141162fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
142163fromStreamD = K.fromStream . toStreamK
143164
165+{-# INLINE toStreamD #-}
166+toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
167+toStreamD = fromStreamK . K.toStream
168+
169+-------------------------------------------------------------------------------
170+-- Deconstruction
171+-------------------------------------------------------------------------------
172+
173+-- Does not fuse, has the same performance as the StreamK version.
174+{-# INLINE_NORMAL uncons #-}
175+uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
176+uncons (UnStream step state) = go state
177+ where
178+ go st = do
179+ r <- step defState st
180+ case r of
181+ Yield x s -> return $ Just (x, Stream step s)
182+ Skip s -> go s
183+ Stop -> return Nothing
184+
185+-------------------------------------------------------------------------------
186+-- Construction
187+-------------------------------------------------------------------------------
188+
189+-- | An empty 'Stream' with a side effect.
190+{-# INLINE_NORMAL nilM #-}
191+nilM :: Monad m => m b -> Stream m a
192+nilM m = Stream (\_ _ -> m >> return Stop) ()
193+
194+{-# INLINE_NORMAL consM #-}
195+consM :: Monad m => m a -> Stream m a -> Stream m a
196+consM m (Stream step state) = Stream step1 Nothing
197+ where
198+ {-# INLINE_LATE step1 #-}
199+ step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
200+ step1 gst (Just st) = do
201+ r <- step gst st
202+ return $
203+ case r of
204+ Yield a s -> Yield a (Just s)
205+ Skip s -> Skip (Just s)
206+ Stop -> Stop
207+
144208------------------------------------------------------------------------------
145209-- Instances
146210------------------------------------------------------------------------------
@@ -567,7 +631,9 @@ cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
567631 Skip s2' -> cmp_null s2'
568632 Stop -> return EQ
569633
570−-- XXX The take/drop combinators above should be moved to filtering section.
634+-------------------------------------------------------------------------------
635+-- Filtering
636+-------------------------------------------------------------------------------
571637
572638{-# INLINE_NORMAL take #-}
573639take :: Monad m => Int -> Stream m a -> Stream m a
@@ -582,6 +648,25 @@ take n (Stream step state) = n `seq` Stream step' (state, 0)
582648 Stop -> Stop
583649 step' _ (_, _) = return Stop
584650
651+
652+{-# INLINE_NORMAL takeWhileM #-}
653+takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
654+takeWhileM f (Stream step state) = Stream step' state
655+ where
656+ {-# INLINE_LATE step' #-}
657+ step' gst st = do
658+ r <- step gst st
659+ case r of
660+ Yield x s -> do
661+ b <- f x
662+ return $ if b then Yield x s else Stop
663+ Skip s -> return $ Skip s
664+ Stop -> return Stop
665+
666+{-# INLINE takeWhile #-}
667+takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
668+takeWhile f = takeWhileM (return . f)
669+
585670------------------------------------------------------------------------------
586671-- Grouping/Splitting
587672------------------------------------------------------------------------------
@@ -706,3 +791,228 @@ groupsOf2 n input (Fold2 fstep inject extract) (Stream step state) =
706791 return $ Yield r next
707792
708793 step' _ GroupFinish2 = return Stop
794+
795+data FromSVarState t m a =
796+ FromSVarInit
797+ | FromSVarRead (SVar t m a)
798+ | FromSVarLoop (SVar t m a) [ChildEvent a]
799+ | FromSVarDone (SVar t m a)
800+
801+-------------------------------------------------------------------------------
802+-- Generation from SVar
803+-------------------------------------------------------------------------------
804+
805+{-# INLINE_NORMAL fromSVar #-}
806+fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
807+fromSVar svar = Stream step FromSVarInit
808+ where
809+
810+ {-# INLINE_LATE step #-}
811+ step _ FromSVarInit = do
812+ ref <- liftIO $ newIORef ()
813+ _ <- liftIO $ mkWeakIORef ref hook
814+ -- when this copy of svar gets garbage collected "ref" will get
815+ -- garbage collected and our GC hook will be called.
816+ let sv = svar{svarRef = Just ref}
817+ return $ Skip (FromSVarRead sv)
818+
819+ where
820+
821+ {-# NOINLINE hook #-}
822+ hook = do
823+ when (svarInspectMode svar) $ do
824+ r <- liftIO $ readIORef (svarStopTime (svarStats svar))
825+ when (isNothing r) $
826+ printSVar svar "SVar Garbage Collected"
827+ cleanupSVar svar
828+ -- If there are any SVars referenced by this SVar a GC will prompt
829+ -- them to be cleaned up quickly.
830+ when (svarInspectMode svar) performMajorGC
831+
832+ step _ (FromSVarRead sv) = do
833+ list <- readOutputQ sv
834+ -- Reversing the output is important to guarantee that we process the
835+ -- outputs in the same order as they were generated by the constituent
836+ -- streams.
837+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
838+
839+ step _ (FromSVarLoop sv []) = do
840+ done <- postProcess sv
841+ return $ Skip $ if done
842+ then (FromSVarDone sv)
843+ else (FromSVarRead sv)
844+
845+ step _ (FromSVarLoop sv (ev : es)) = do
846+ case ev of
847+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)
848+ ChildStop tid e -> do
849+ accountThread sv tid
850+ case e of
851+ Nothing -> do
852+ stop <- shouldStop tid
853+ if stop
854+ then do
855+ liftIO (cleanupSVar sv)
856+ return $ Skip (FromSVarDone sv)
857+ else return $ Skip (FromSVarLoop sv es)
858+ Just ex ->
859+ case fromException ex of
860+ Just ThreadAbort ->
861+ return $ Skip (FromSVarLoop sv es)
862+ Nothing -> liftIO (cleanupSVar sv) >> throwM ex
863+ where
864+
865+ shouldStop tid =
866+ case svarStopStyle sv of
867+ StopNone -> return False
868+ StopAny -> return True
869+ StopBy -> do
870+ sid <- liftIO $ readIORef (svarStopBy sv)
871+ return $ if tid == sid then True else False
872+
873+ step _ (FromSVarDone sv) = do
874+ when (svarInspectMode sv) $ do
875+ t <- liftIO $ getTime Monotonic
876+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
877+ liftIO $ printSVar sv "SVar Done"
878+ return Stop
879+
880+{-# INLINE_NORMAL fromPrimIORef #-}
881+fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
882+fromPrimIORef var = Stream step ()
883+ where
884+ {-# INLINE_LATE step #-}
885+ step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
886+
887+-------------------------------------------------------------------------------
888+-- Process events received by a fold consumer from a stream producer
889+-------------------------------------------------------------------------------
890+
891+{-# INLINE_NORMAL fromProducer #-}
892+fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
893+fromProducer svar = Stream step (FromSVarRead svar)
894+ where
895+
896+ {-# INLINE_LATE step #-}
897+ step _ (FromSVarRead sv) = do
898+ list <- readOutputQ sv
899+ -- Reversing the output is important to guarantee that we process the
900+ -- outputs in the same order as they were generated by the constituent
901+ -- streams.
902+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
903+
904+ step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
905+ step _ (FromSVarLoop sv (ev : es)) = do
906+ case ev of
907+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)
908+ ChildStop tid e -> do
909+ accountThread sv tid
910+ case e of
911+ Nothing -> do
912+ sendStopToProducer sv
913+ return $ Skip (FromSVarDone sv)
914+ Just _ -> error "Bug: fromProducer: received exception"
915+
916+ step _ (FromSVarDone sv) = do
917+ when (svarInspectMode sv) $ do
918+ t <- liftIO $ getTime Monotonic
919+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
920+ liftIO $ printSVar sv "SVar Done"
921+ return Stop
922+
923+ step _ FromSVarInit = undefined
924+
925+------------------------------------------------------------------------------
926+-- Left Folds
927+------------------------------------------------------------------------------
928+
929+{-# INLINE_NORMAL foldOnce #-}
930+foldOnce :: (Monad m) => Fold m a b -> Stream m a -> m b
931+foldOnce (Fold fstep begin done) (Stream step state) =
932+ begin >>= \x -> go SPEC x state
933+
934+ where
935+
936+ {-# INLINE go #-}
937+ go !_ !fs st = do
938+ r <- step defState st
939+ case r of
940+ Yield x s -> do
941+ res <- fstep fs x
942+ case res of
943+ FL.Done b -> return b
944+ FL.Partial fs1 -> go SPEC fs1 s
945+ Skip s -> go SPEC fs s
946+ Stop -> done fs
947+
948+-------------------------------------------------------------------------------
949+-- Concurrent application and fold
950+-------------------------------------------------------------------------------
951+
952+-- XXX These functions should be moved to Stream/Parallel.hs
953+--
954+-- Using StreamD the worker stream producing code can fuse with the code to
955+-- queue output to the SVar giving some perf boost.
956+--
957+-- Note that StreamD can only be used in limited situations, specifically, we
958+-- cannot implement joinStreamVarPar using this.
959+--
960+-- XXX make sure that the SVar passed is a Parallel style SVar.
961+
962+-- | Fold the supplied stream to the SVar asynchronously using Parallel
963+-- concurrency style.
964+-- {-# INLINE_NORMAL toSVarParallel #-}
965+{-# INLINE toSVarParallel #-}
966+toSVarParallel :: MonadAsync m
967+ => State t m a -> SVar t m a -> Stream m a -> m ()
968+toSVarParallel st sv xs =
969+ if svarInspectMode sv
970+ then forkWithDiag
971+ else do
972+ tid <-
973+ case getYieldLimit st of
974+ Nothing -> doFork (work Nothing)
975+ (svarMrun sv)
976+ (handleChildException sv)
977+ Just _ -> doFork (workLim Nothing)
978+ (svarMrun sv)
979+ (handleChildException sv)
980+ modifyThread sv tid
981+
982+ where
983+
984+ {-# NOINLINE work #-}
985+ work info = (foldOnce (FL.toParallelSVar sv info) xs)
986+
987+ {-# NOINLINE workLim #-}
988+ workLim info = foldOnce (FL.toParallelSVarLimited sv info) xs
989+
990+ {-# NOINLINE forkWithDiag #-}
991+ forkWithDiag = do
992+ -- We do not use workerCount in case of ParallelVar but still there is
993+ -- no harm in maintaining it correctly.
994+ liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
995+ recordMaxWorkers sv
996+ -- This allocation matters when significant number of workers are being
997+ -- sent. We allocate it only when needed. The overhead increases by 4x.
998+ winfo <-
999+ case yieldRateInfo sv of
1000+ Nothing -> return Nothing
1001+ Just _ -> liftIO $ do
1002+ cntRef <- newIORef 0
1003+ t <- getTime Monotonic
1004+ lat <- newIORef (0, t)
1005+ return $ Just WorkerInfo
1006+ { workerYieldMax = 0
1007+ , workerYieldCount = cntRef
1008+ , workerLatencyStart = lat
1009+ }
1010+ tid <-
1011+ case getYieldLimit st of
1012+ Nothing -> doFork (work winfo)
1013+ (svarMrun sv)
1014+ (handleChildException sv)
1015+ Just _ -> doFork (workLim winfo)
1016+ (svarMrun sv)
1017+ (handleChildException sv)
1018+ modifyThread sv tid