@@ -98,15 +98,18 @@ module Streamly.Internal.Data.Stream.Concurrent
9898
9999 -- ** Reactive
100100 , fromCallback
101+ , pollCountsD
102+ , pollCounts
101103 )
102104where
103105
104106#include "inline.hs"
105107
106−import Control.Concurrent (myThreadId)
107−import Control.Monad (void)
108+import Control.Concurrent (myThreadId, killThread)
109+import Control.Monad (void, when)
108110import Control.Monad.IO.Class (MonadIO(liftIO))
109111import Streamly.Internal.Control.Concurrent (MonadAsync, askRunInIO)
112+import Streamly.Internal.Control.ForkLifted (forkManaged)
110113import Streamly.Internal.Data.Stream.Channel.Dispatcher (modifyThread)
111114import Streamly.Internal.Data.Stream.Channel.Types
112115 ( ChildEvent(..)
@@ -114,8 +117,11 @@ import Streamly.Internal.Data.Stream.Channel.Types
114117 )
115118import Streamly.Internal.Data.Stream.Channel.Worker (sendWithDoorBell)
116119import Streamly.Internal.Data.Stream.Type (Stream)
120+import Streamly.Internal.Data.Stream.StreamD (Step(..))
117121
122+import qualified Streamly.Internal.Data.IORef.Unboxed as Unboxed
118123import qualified Streamly.Internal.Data.Stream as Stream
124+import qualified Streamly.Internal.Data.Stream.StreamD as D
119125import qualified Streamly.Internal.Data.Stream.StreamK as K
120126import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
121127
@@ -695,3 +701,65 @@ fromCallback setCallback = Stream.concatM $ do
695701 (callback, stream) <- newCallbackStream
696702 setCallback callback
697703 return stream
704+
705+{-# INLINE_NORMAL pollCountsD #-}
706+pollCountsD
707+ :: MonadAsync m
708+ => (a -> Bool)
709+ -> (D.Stream m Int -> m b)
710+ -> D.Stream m a
711+ -> D.Stream m a
712+pollCountsD predicate fld (D.Stream step state) = D.Stream step' Nothing
713+ where
714+
715+ {-# INLINE_LATE step' #-}
716+ step' _ Nothing = do
717+ -- As long as we are using an "Int" for counts lockfree reads from
718+ -- Var should work correctly on both 32-bit and 64-bit machines.
719+ -- However, an Int on a 32-bit machine may overflow quickly.
720+ countVar <- liftIO $ Unboxed.newIORef (0 :: Int)
721+ tid <- forkManaged
722+ $ void $ fld
723+ $ Unboxed.toStreamD countVar
724+ return $ Skip (Just (countVar, tid, state))
725+
726+ step' gst (Just (countVar, tid, st)) = do
727+ r <- step gst st
728+ case r of
729+ Yield x s -> do
730+ when (predicate x)
731+ $ liftIO $ Unboxed.modifyIORef' countVar (+ 1)
732+ return $ Yield x (Just (countVar, tid, s))
733+ Skip s -> return $ Skip (Just (countVar, tid, s))
734+ Stop -> do
735+ liftIO $ killThread tid
736+ return Stop
737+
738+-- | @pollCounts predicate transform fold stream@ counts those elements in the
739+-- stream that pass the @predicate@. The resulting count stream is sent to
740+-- another thread which transforms it using @transform@ and then folds it using
741+-- @fold@. The thread is automatically cleaned up if the stream stops or
742+-- aborts due to exception.
743+--
744+-- For example, to print the count of elements processed every second:
745+--
746+-- @
747+-- > Stream.drain $ Concur.pollCounts (const True) (Stream.rollingMap (-) . Stream.delayPost 1) (Fold.drainBy print)
748+-- $ Stream.enumerateFrom 0
749+-- @
750+--
751+-- Note: This may not work correctly on 32-bit machines.
752+--
753+-- /Pre-release/
754+--
755+{-# INLINE pollCounts #-}
756+pollCounts ::
757+ (MonadAsync m)
758+ => (a -> Bool)
759+ -> (Stream m Int -> m b)
760+ -> Stream m a
761+ -> Stream m a
762+pollCounts predicate f xs =
763+ Stream.fromStreamD
764+ $ pollCountsD predicate (f . Stream.fromStreamD)
765+ $ Stream.toStreamD xs