3e22f1b · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

flush sessions when classifySessionsOf ends

3e22f1bHarendra Kumar6 years ago

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

compare with master

unifiedsplit

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

226

227 -- ** By folding (scans)

228 , scanlM'

229+ , scanlMAfter'

229 , scanl'

230 , scanl'

230 , scanlM

231 , scanlM

231 , scanl

232 , scanl

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

241 , postscanlM

242 , postscanlM

242 , postscanl'

243 , postscanl'

243 , postscanlM'

244 , postscanlM'

245+ , postscanlMAfter'

244

246

245 , postscanlx'

247 , postscanlx'

246 , postscanlMx'

248 , postscanlMx'

@@ -3609,6 +3611,31 @@ postscanlM' fstep begin (Stream step state) =

3609postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

3611postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

3610postscanl' f = postscanlM' (\a b -> return (f a b))

3612postscanl' f = postscanlM' (\a b -> return (f a b))

3611

3613

3614+-- We can possibly have the "done" function as a Maybe to provide an option to

3615+-- emit or not emit the accumulator when the stream stops.

3616+--

3617+-- TBD: use a single Yield point

3618+--

3619+{-# INLINE_NORMAL postscanlMAfter' #-}

3620+postscanlMAfter' :: Monad m

3621+ => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

3622+postscanlMAfter' fstep initial done (Stream step1 state1) = do

3623+ Stream step (Just (state1, initial))

3624+

3625+ where

3626+

3627+ {-# INLINE_LATE step #-}

3628+ step gst (Just (st, acc)) = do

3629+ r <- step1 (adaptState gst) st

3630+ case r of

3631+ Yield x s -> do

3632+ old <- acc

3633+ y <- fstep old x

3634+ y `seq` return (Yield y (Just (s, return y)))

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

3636+ Stop -> acc >>= done >>= \res -> return (Yield res Nothing)

3637+ step _ Nothing = return Stop

3638+

3612{-# INLINE_NORMAL postscanlM #-}

3639{-# INLINE_NORMAL postscanlM #-}

3613postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

3640postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

3614postscanlM fstep begin (Stream step state) = Stream step' (state, begin)

3641postscanlM fstep begin (Stream step state) = Stream step' (state, begin)

@@ -3631,6 +3658,13 @@ postscanl f = postscanlM (\a b -> return (f a b))

3631scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

3658scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

3632scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)

3659scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)

3633

3660

3661+{-# INLINE scanlMAfter' #-}

3662+scanlMAfter' :: Monad m

3663+ => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

3664+scanlMAfter' fstep initial done s =

3665+ (initial >>= \x -> x `seq` return x) `consM`

3666+ postscanlMAfter' fstep initial done s

3667+

3634{-# INLINE scanl' #-}

3668{-# INLINE scanl' #-}

3635scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

3669scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

3636scanl' f = scanlM' (\a b -> return (f a b))

3670scanl' f = scanlM' (\a b -> return (f a b))

src/Streamly/Internal/Prelude.hs+47−2

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

202 -- ** Left scans

203 , scanl'

204 , scanlM'

205+ , scanlMAfter'

205 , postscanl'

206 , postscanl'

206 , postscanlM'

207 , postscanlM'

207 , prescanl'

208 , prescanl'

@@ -2125,6 +2126,23 @@ scanx = P.scanlx'

2125scanlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> t m b

2126scanlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> t m b

2126scanlM' step begin m = fromStreamD $ D.scanlM' step begin $ toStreamD m

2127scanlM' step begin m = fromStreamD $ D.scanlM' step begin $ toStreamD m

2127

2128

2129+-- | @scanlMAfter' accumulate initial done stream@ is like 'scanlM'' except

2130+-- that it provides an additional @done@ function to be applied on the

2131+-- accumulator when the stream stops. The result of @done@ is also emitted in

2132+-- the stream.

2133+--

2134+-- This function can be used to allocate a resource in the beginning of the

2135+-- scan and release it when the stream ends or to flush the internal state of

2136+-- the scan at the end.

2137+--

2138+-- /Internal/

2139+--

2140+{-# INLINE scanlMAfter' #-}

2141+scanlMAfter' :: (IsStream t, Monad m)

2142+ => (b -> a -> m b) -> m b -> (b -> m b) -> t m a -> t m b

2143+scanlMAfter' step initial done stream =

2144+ fromStreamD $ D.scanlMAfter' step initial done $ toStreamD stream

2145+

2128-- | Strict left scan. Like 'map', 'scanl'' too is a one to one transformation,

2146-- | Strict left scan. Like 'map', 'scanl'' too is a one to one transformation,

2129-- however it adds an extra element.

2147-- however it adds an extra element.

2130--

2148--

@@ -4507,8 +4525,8 @@ classifySessionsBy

4507 -> t m (k, b) -- ^ session key, fold result

4525 -> t m (k, b) -- ^ session key, fold result

4508classifySessionsBy tick tmout reset ejectPred

4526classifySessionsBy tick tmout reset ejectPred

4509 (Fold step initial extract) str =

4527 (Fold step initial extract) str =

4510− concatMap sessionOutputStream

4528+ concatMap sessionOutputStream $

4511− $ scanlM' sstep szero stream

4529+ scanlMAfter' sstep (return szero) flush stream

4512

4530

4513 where

4531 where

4514

4532

@@ -4613,6 +4631,21 @@ classifySessionsBy tick tmout reset ejectPred

4613 let curTime = addToAbsTime sessionCurTime tickMs

4631 let curTime = addToAbsTime sessionCurTime tickMs

4614 in ejectExpired sessionState curTime

4632 in ejectExpired sessionState curTime

4615

4633

4634+ flush session@SessionState{..} = do

4635+ (hp', mp', out, count) <-

4636+ ejectAll

4637+ ( sessionTimerHeap

4638+ , sessionKeyValueMap

4639+ , K.nil

4640+ , sessionCount

4641+ )

4642+ return $ session

4643+ { sessionCount = count

4644+ , sessionTimerHeap = hp'

4645+ , sessionKeyValueMap = mp'

4646+ , sessionOutputStream = out

4647+ }

4648+

4616 fromEither e =

4649 fromEither e =

4617 case e of

4650 case e of

4618 Left x -> x

4651 Left x -> x

@@ -4625,6 +4658,18 @@ classifySessionsBy tick tmout reset ejectPred

4625 let mp1 = Map.delete key mp

4658 let mp1 = Map.delete key mp

4626 return (hp, mp1, out1, cnt - 1)

4659 return (hp, mp1, out1, cnt - 1)

4627

4660

4661+ ejectAll (hp, mp, out, !cnt) = do

4662+ let hres = H.uncons hp

4663+ case hres of

4664+ Just (Entry _ key, hp1) -> do

4665+ r <- case Map.lookup key mp of

4666+ Nothing -> return (hp1, mp, out, cnt)

4667+ Just (Tuple' _ acc) -> ejectEntry hp1 mp out cnt acc key

4668+ ejectAll r

4669+ Nothing -> do

4670+ assert (Map.null mp) (return ())

4671+ return (hp, mp, out, cnt)

4672+

4628 ejectOne (hp, mp, out, !cnt) = do

4673 ejectOne (hp, mp, out, !cnt) = do

4629 let hres = H.uncons hp

4674 let hres = H.uncons hp

4630 case hres of

4675 case hres of