@@ -136,11 +136,12 @@ import qualified Streamly.Streams.StreamK as K
136136-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
137137-- the next state of the stream. 'Stop' indicates there are no more values in
138138-- the stream.
139−data Step s a = Yield a s | Stop
139+data Step s a = Yield a s | Skip s | Stop
140140
141141instance Functor (Step s) where
142142 {-# INLINE fmap #-}
143143 fmap f (Yield x s) = Yield (f x) s
144+ fmap _ (Skip s) = Skip s
144145 fmap _ Stop = Stop
145146
146147-- gst = global state
@@ -166,9 +167,11 @@ cons x (Stream step state) = Stream step1 Nothing
166167 step1 _ Nothing = return $ Yield x (Just state)
167168 step1 gst (Just st) = do
168169 r <- step (rstState gst) st
169− case r of
170− Yield a s -> return $ Yield a (Just s)
171− Stop -> return Stop
170+ return $
171+ case r of
172+ Yield a s -> Yield a (Just s)
173+ Skip s -> Skip (Just s)
174+ Stop -> Stop
172175
173176-------------------------------------------------------------------------------
174177-- Deconstruction
@@ -181,9 +184,10 @@ uncons (Stream step state) = go state
181184 where
182185 go st = do
183186 r <- step defState st
184− return $ case r of
185− Yield x s -> Just (x, Stream step s)
186− Stop -> Nothing
187+ case r of
188+ Yield x s -> return $ Just (x, Stream step s)
189+ Skip s -> go s
190+ Stop -> return Nothing
187191
188192------------------------------------------------------------------------------
189193-- Generation by unfold
@@ -284,6 +288,7 @@ foldrM f z (Stream step state) = go SPEC state
284288 r <- step defState st
285289 case r of
286290 Yield x s -> go SPEC s >>= f x
291+ Skip s -> go SPEC s
287292 Stop -> return z
288293
289294{-# INLINE_NORMAL foldr #-}
@@ -300,7 +305,8 @@ foldlM' fstep begin (Stream step state) = go SPEC begin state
300305 Yield x s -> do
301306 acc' <- fstep acc x
302307 go SPEC acc' s
303− Stop -> return acc
308+ Skip s -> go SPEC acc s
309+ Stop -> return acc
304310
305311{-# INLINE foldl' #-}
306312foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
@@ -319,6 +325,7 @@ runStream (Stream step state) = go SPEC state
319325 r <- step defState st
320326 case r of
321327 Yield _ s -> go SPEC s
328+ Skip s -> go SPEC s
322329 Stop -> return ()
323330
324331{-# INLINE_NORMAL null #-}
@@ -329,7 +336,8 @@ null (Stream step state) = go state
329336 r <- step defState st
330337 case r of
331338 Yield _ _ -> return False
332− Stop -> return True
339+ Skip s -> go s
340+ Stop -> return True
333341
334342-- XXX SPEC?
335343{-# INLINE_NORMAL head #-}
@@ -340,7 +348,8 @@ head (Stream step state) = go state
340348 r <- step defState st
341349 case r of
342350 Yield x _ -> return (Just x)
343− Stop -> return Nothing
351+ Skip s -> go s
352+ Stop -> return Nothing
344353
345354-- Does not fuse, has the same performance as the StreamK version.
346355{-# INLINE_NORMAL tail #-}
@@ -351,7 +360,8 @@ tail (Stream step state) = go state
351360 r <- step defState st
352361 case r of
353362 Yield _ s -> return (Just $ Stream step s)
354− Stop -> return Nothing
363+ Skip s -> go s
364+ Stop -> return Nothing
355365
356366-- XXX will it fuse? need custom impl?
357367{-# INLINE_NORMAL last #-}
@@ -365,24 +375,15 @@ elem e (Stream step state) = go state
365375 go st = do
366376 r <- step defState st
367377 case r of
368− Yield x s ->
369− if x == e
370− then return True
371− else go s
372− Stop -> return False
378+ Yield x s
379+ | x == e -> return True
380+ | otherwise -> go s
381+ Skip s -> go s
382+ Stop -> return False
373383
374384{-# INLINE_NORMAL notElem #-}
375385notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
376−notElem e (Stream step state) = go state
377− where
378− go st = do
379− r <- step defState st
380− case r of
381− Yield x s ->
382− if x == e
383− then return False
384− else go s
385− Stop -> return True
386+notElem e s = fmap not (elem e s)
386387
387388{-# INLINE_NORMAL all #-}
388389all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
@@ -391,11 +392,11 @@ all p (Stream step state) = go state
391392 go st = do
392393 r <- step defState st
393394 case r of
394− Yield x s ->
395− if p x
396− then go s
397− else return False
398− Stop -> return True
395+ Yield x s
396+ | p x -> go s
397+ | otherwise -> return False
398+ Skip s -> go s
399+ Stop -> return True
399400
400401{-# INLINE_NORMAL any #-}
401402any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
@@ -404,11 +405,11 @@ any p (Stream step state) = go state
404405 go st = do
405406 r <- step defState st
406407 case r of
407− Yield x s ->
408− if p x
409− then return True
410− else go s
411− Stop -> return False
408+ Yield x s
409+ | p x -> return True
410+ | otherwise -> go s
411+ Skip s -> go s
412+ Stop -> return False
412413
413414{-# INLINE_NORMAL maximum #-}
414415maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
@@ -418,15 +419,16 @@ maximum (Stream step state) = go Nothing state
418419 r <- step defState st
419420 case r of
420421 Yield x s -> go (Just x) s
421− Stop -> return Nothing
422+ Skip s -> go Nothing s
423+ Stop -> return Nothing
422424 go (Just acc) st = do
423425 r <- step defState st
424426 case r of
425− Yield x s ->
426− if acc <= x
427− then go (Just x) s
428− else go (Just acc) s
429− Stop -> return (Just acc)
427+ Yield x s
428+ | acc <= x -> go (Just x) s
429+ | otherwise -> go (Just acc) s
430+ Skip s -> go (Just acc) s
431+ Stop -> return (Just acc)
430432
431433{-# INLINE_NORMAL minimum #-}
432434minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
@@ -436,15 +438,16 @@ minimum (Stream step state) = go Nothing state
436438 r <- step defState st
437439 case r of
438440 Yield x s -> go (Just x) s
439− Stop -> return Nothing
441+ Skip s -> go Nothing s
442+ Stop -> return Nothing
440443 go (Just acc) st = do
441444 r <- step defState st
442445 case r of
443− Yield x s ->
444− if acc <= x
445− then go (Just acc) s
446− else go (Just x) s
447− Stop -> return (Just acc)
446+ Yield x s
447+ | acc <= x -> go (Just acc) s
448+ | otherwise -> go (Just x) s
449+ Skip s -> go (Just acc) s
450+ Stop -> return (Just acc)
448451
449452------------------------------------------------------------------------------
450453-- Map and Fold
@@ -468,10 +471,11 @@ toList = foldr (:) []
468471toStreamK :: Monad m => Stream m a -> K.Stream m a
469472toStreamK (Stream step state) = go state
470473 where
471− go st = K.Stream $ \gst stp _ yld -> do
474+ go st = K.Stream $ \gst stp sng yld -> do
472475 r <- step gst st
473476 case r of
474477 Yield x s -> yld x (go s)
478+ Skip s -> K.unStream (go s) gst stp sng yld
475479 Stop -> stp
476480
477481#ifndef DISABLE_FUSION
@@ -497,7 +501,8 @@ postscanlM' fstep begin (Stream step state) =
497501 Yield x s -> do
498502 y <- fstep acc x
499503 y `seq` return (Yield y (s, y))
500− Stop -> return Stop
504+ Skip s -> return $ Skip (s, acc)
505+ Stop -> return Stop
501506
502507{-# INLINE_LATE scanlM' #-}
503508scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
@@ -520,6 +525,7 @@ take n (Stream step state) = n `seq` Stream step' (state, 0)
520525 r <- step (rstState gst) st
521526 return $ case r of
522527 Yield x s -> Yield x (s, i + 1)
528+ Skip s -> Skip (s, i)
523529 Stop -> Stop
524530 step' _ (_, _) = return Stop
525531
@@ -534,7 +540,8 @@ takeWhileM f (Stream step state) = Stream step' state
534540 Yield x s -> do
535541 b <- f x
536542 return $ if b then Yield x s else Stop
537− Stop -> return Stop
543+ Skip s -> return $ Skip s
544+ Stop -> return Stop
538545
539546{-# INLINE takeWhile #-}
540547takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
@@ -542,15 +549,26 @@ takeWhile f = takeWhileM (return . f)
542549
543550{-# INLINE_NORMAL drop #-}
544551drop :: Monad m => Int -> Stream m a -> Stream m a
545−drop n (Stream step state) = Stream step' (state, n)
552+drop n (Stream step state) = Stream step' (state, Just n)
546553 where
547554 {-# INLINE_LATE step' #-}
548− step' gst (st, i) = do
549− r <- step (rstState gst) st
555+ step' gst (st, Just i)
556+ | i > 0 = do
557+ r <- step (rstState gst) st
558+ return $
559+ case r of
560+ Yield _ s -> Skip (s, Just (i - 1))
561+ Skip s -> Skip (s, Just i)
562+ Stop -> Stop
563+ | otherwise = return $ Skip (st, Nothing)
564+
565+ step' gst (st, Nothing) = do
566+ r <- step (rstState gst) st
567+ return $
550568 case r of
551− Yield _ s | i > 0 -> step' gst (s, i - 1)
552− Yield x s -> return $ Yield x (s, 0)
553− Stop -> return Stop
569+ Yield x s -> Yield x (s, Nothing)
570+ Skip s -> Skip (s, Nothing)
571+ Stop -> Stop
554572
555573data DropWhileState s a
556574 = DropWhileDrop s
@@ -568,14 +586,16 @@ dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
568586 Yield x s -> do
569587 b <- f x
570588 if b
571− then step' gst (DropWhileDrop s)
572− else step' gst (DropWhileYield x s)
589+ then return $ Skip (DropWhileDrop s)
590+ else return $ Skip (DropWhileYield x s)
591+ Skip s -> return $ Skip (DropWhileDrop s)
573592 Stop -> return Stop
574593
575594 step' gst (DropWhileNext st) = do
576595 r <- step (rstState gst) st
577596 case r of
578− Yield x s -> step' gst (DropWhileYield x s)
597+ Yield x s -> return $ Skip (DropWhileYield x s)
598+ Skip s -> return $ Skip (DropWhileNext s)
579599 Stop -> return Stop
580600
581601 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
@@ -594,10 +614,11 @@ filterM f (Stream step state) = Stream step' state
594614 case r of
595615 Yield x s -> do
596616 b <- f x
597− if b
598− then return $ Yield x s
599− else step' gst s
600− Stop -> return Stop
617+ return $ if b
618+ then Yield x s
619+ else Skip s
620+ Skip s -> return $ Skip s
621+ Stop -> return Stop
601622
602623{-# INLINE filter #-}
603624filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
@@ -617,6 +638,7 @@ mapM f (Stream step state) = Stream step' state
617638 r <- step (rstState gst) st
618639 case r of
619640 Yield x s -> f x >>= \a -> return $ Yield a s
641+ Skip s -> return $ Skip s
620642 Stop -> return Stop
621643
622644{-# INLINE map #-}
@@ -648,9 +670,11 @@ zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
648670 {-# INLINE_LATE step #-}
649671 step gst (sa, sb, Nothing) = do
650672 r <- stepa (rstState gst) sa
651− case r of
652− Yield x sa' -> step gst (sa', sb, Just x)
653− Stop -> return Stop
673+ return $
674+ case r of
675+ Yield x sa' -> Skip (sa', sb, Just x)
676+ Skip sa' -> Skip (sa', sb, Nothing)
677+ Stop -> Stop
654678
655679 step gst (sa, sb, Just x) = do
656680 r <- stepb (rstState gst) sb
@@ -658,7 +682,8 @@ zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
658682 Yield y sb' -> do
659683 z <- f x y
660684 return $ Yield z (sa, sb', Nothing)
661− Stop -> return Stop
685+ Skip sb' -> return $ Skip (sa, sb', Just x)
686+ Stop -> return Stop
662687
663688{-# RULES "zipWithM xs xs"
664689 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}