b83c2c2 · streamly · git.iko.soy

git.iko.soy/streamly

Add tapOffsetEvery to tap one in n elements

commitb83c2c28fc93a7f53b0a594ad4d8180b8bc06822 authorHarendra Kumar <harendra.kumar@gmail.com> dateSat 28 Dec 2019, 23:34 (7 years ago) parent5073377

2 files changed+53

compare with master

unifiedsplit

src/Streamly/Internal/Data/Stream/StreamD.hs+34 src/Streamly/Internal/Prelude.hs+19

src/Streamly/Internal/Data/Stream/StreamD.hs+34

@@ -128,6 +128,7 @@ module Streamly.Internal.Data.Stream.StreamD

128128

129129 -- ** Specialized Folds

130130 , tap

131+ , tapOffsetEvery

131132 , tapAsync

132133 , tapRate

133134 , pollCounts

@@ -3279,6 +3280,39 @@ tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing

32793280 void $ extract acc

32803281 return $ Stop

32813282

3283+{-# INLINE_NORMAL tapOffsetEvery #-}

3284+tapOffsetEvery :: Monad m

3285+ => Int -> Int -> Fold m a b -> Stream m a -> Stream m a

3286+tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =

3287+ Stream step' Nothing

3288+

3289+ where

3290+

3291+ {-# INLINE_LATE step' #-}

3292+ step' _ Nothing = do

3293+ r <- initial

3294+ return $ Skip (Just (r, state, offset `mod` n))

3295+

3296+ step' gst (Just (acc, st, count)) | count <= 0 = do

3297+ r <- step gst st

3298+ case r of

3299+ Yield x s -> do

3300+ !acc' <- fstep acc x

3301+ return $ Yield x (Just (acc', s, n - 1))

3302+ Skip s -> return $ Skip (Just (acc, s, count))

3303+ Stop -> do

3304+ void $ extract acc

3305+ return $ Stop

3306+

3307+ step' gst (Just (acc, st, count)) = do

3308+ r <- step gst st

3309+ case r of

3310+ Yield x s -> return $ Yield x (Just (acc, s, count - 1))

3311+ Skip s -> return $ Skip (Just (acc, s, count))

3312+ Stop -> do

3313+ void $ extract acc

3314+ return $ Stop

3315+

32823316{-# INLINE_NORMAL pollCounts #-}

32833317pollCounts

32843318 :: MonadAsync m

src/Streamly/Internal/Prelude.hs+19

@@ -345,6 +345,7 @@ module Streamly.Internal.Prelude

345345 -- ** Distributing

346346 , trace

347347 , tap

348+ , tapOffsetEvery

348349 , tapAsync

349350 , tapRate

350351 , pollCounts

@@ -3629,6 +3630,24 @@ reassembleBy = undefined

36293630tap :: (IsStream t, Monad m) => FL.Fold m a b -> t m a -> t m a

36303631tap f xs = D.fromStreamD $ D.tap f (D.toStreamD xs)

36313632

3633+-- | @tapOffsetEvery offset n@ taps every @n@th element in the stream

3634+-- starting at @offset@. @offset@ can be between @0@ and @n - 1@. Offset 0

3635+-- means start at the first element in the stream. If the offset is outside

3636+-- this range then @offset `mod` n@ is used as offset.

3637+--

3638+-- @

3639+-- >>> S.drain $ S.tapOffsetEvery 0 2 (FL.mapM print FL.toList) $ S.enumerateFromTo 0 10

3640+-- > [0,2,4,6,8,10]

3641+-- @

3642+--

3643+-- /Internal/

3644+--

3645+{-# INLINE tapOffsetEvery #-}

3646+tapOffsetEvery :: (IsStream t, Monad m)

3647+ => Int -> Int -> FL.Fold m a b -> t m a -> t m a

3648+tapOffsetEvery offset n f xs =

3649+ D.fromStreamD $ D.tapOffsetEvery offset n f (D.toStreamD xs)

3650+

36323651-- | Redirect a copy of the stream to a supplied fold and run it concurrently

36333652-- in an independent thread. The fold may buffer some elements. The buffer size

36343653-- is determined by the prevailing 'maxBuffer' setting.