git.iko.soy/streamly
on master help sign in
Add Skip constructor to the core Stream type to enhance stream fusion
81cf15bAbhiroop Sarkar
src/Streamly/Streams/StreamD.hs+93−68
compare with master
unifiedsplit
@@ -136,11 +136,12 @@ import qualified Streamly.Streams.StreamK as K
136-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
136-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
137-- the next state of the stream. 'Stop' indicates there are no more values in
137-- the next state of the stream. 'Stop' indicates there are no more values in
138-- the stream.
138-- the stream.
139−data Step s a = Yield a s | Stop
139+data Step s a = Yield a s | Skip s | Stop
140
140
141instance Functor (Step s) where
141instance Functor (Step s) where
142 {-# INLINE fmap #-}
142 {-# INLINE fmap #-}
143 fmap f (Yield x s) = Yield (f x) s
143 fmap f (Yield x s) = Yield (f x) s
144+ fmap _ (Skip s) = Skip s
144 fmap _ Stop = Stop
145 fmap _ Stop = Stop
145
146
146-- gst = global state
147-- gst = global state
@@ -166,9 +167,11 @@ cons x (Stream step state) = Stream step1 Nothing
166 step1 _ Nothing = return $ Yield x (Just state)
167 step1 _ Nothing = return $ Yield x (Just state)
167 step1 gst (Just st) = do
168 step1 gst (Just st) = do
168 r <- step (rstState gst) st
169 r <- step (rstState gst) st
170+ return $
169− case r of
171+ case r of
170− Yield a s -> return $ Yield a (Just s)
172+ Yield a s -> Yield a (Just s)
171− Stop -> return Stop
173+ Skip s -> Skip (Just s)
174+ Stop -> Stop
172
175
173-------------------------------------------------------------------------------
176-------------------------------------------------------------------------------
174-- Deconstruction
177-- Deconstruction
@@ -181,9 +184,10 @@ uncons (Stream step state) = go state
181 where
184 where
182 go st = do
185 go st = do
183 r <- step defState st
186 r <- step defState st
184− return $ case r of
187+ case r of
185− Yield x s -> Just (x, Stream step s)
188+ Yield x s -> return $ Just (x, Stream step s)
189+ Skip s -> go s
186− Stop -> Nothing
190+ Stop -> return Nothing
187
191
188------------------------------------------------------------------------------
192------------------------------------------------------------------------------
189-- Generation by unfold
193-- Generation by unfold
@@ -284,6 +288,7 @@ foldrM f z (Stream step state) = go SPEC state
284 r <- step defState st
288 r <- step defState st
285 case r of
289 case r of
286 Yield x s -> go SPEC s >>= f x
290 Yield x s -> go SPEC s >>= f x
291+ Skip s -> go SPEC s
287 Stop -> return z
292 Stop -> return z
288
293
289{-# INLINE_NORMAL foldr #-}
294{-# INLINE_NORMAL foldr #-}
@@ -300,7 +305,8 @@ foldlM' fstep begin (Stream step state) = go SPEC begin state
300 Yield x s -> do
305 Yield x s -> do
301 acc' <- fstep acc x
306 acc' <- fstep acc x
302 go SPEC acc' s
307 go SPEC acc' s
308+ Skip s -> go SPEC acc s
303− Stop -> return acc
309+ Stop -> return acc
304
310
305{-# INLINE foldl' #-}
311{-# INLINE foldl' #-}
306foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
312foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
@@ -319,6 +325,7 @@ runStream (Stream step state) = go SPEC state
319 r <- step defState st
325 r <- step defState st
320 case r of
326 case r of
321 Yield _ s -> go SPEC s
327 Yield _ s -> go SPEC s
328+ Skip s -> go SPEC s
322 Stop -> return ()
329 Stop -> return ()
323
330
324{-# INLINE_NORMAL null #-}
331{-# INLINE_NORMAL null #-}
@@ -329,7 +336,8 @@ null (Stream step state) = go state
329 r <- step defState st
336 r <- step defState st
330 case r of
337 case r of
331 Yield _ _ -> return False
338 Yield _ _ -> return False
339+ Skip s -> go s
332− Stop -> return True
340+ Stop -> return True
333
341
334-- XXX SPEC?
342-- XXX SPEC?
335{-# INLINE_NORMAL head #-}
343{-# INLINE_NORMAL head #-}
@@ -340,7 +348,8 @@ head (Stream step state) = go state
340 r <- step defState st
348 r <- step defState st
341 case r of
349 case r of
342 Yield x _ -> return (Just x)
350 Yield x _ -> return (Just x)
351+ Skip s -> go s
343− Stop -> return Nothing
352+ Stop -> return Nothing
344
353
345-- Does not fuse, has the same performance as the StreamK version.
354-- Does not fuse, has the same performance as the StreamK version.
346{-# INLINE_NORMAL tail #-}
355{-# INLINE_NORMAL tail #-}
@@ -351,7 +360,8 @@ tail (Stream step state) = go state
351 r <- step defState st
360 r <- step defState st
352 case r of
361 case r of
353 Yield _ s -> return (Just $ Stream step s)
362 Yield _ s -> return (Just $ Stream step s)
363+ Skip s -> go s
354− Stop -> return Nothing
364+ Stop -> return Nothing
355
365
356-- XXX will it fuse? need custom impl?
366-- XXX will it fuse? need custom impl?
357{-# INLINE_NORMAL last #-}
367{-# INLINE_NORMAL last #-}
@@ -365,24 +375,15 @@ elem e (Stream step state) = go state
365 go st = do
375 go st = do
366 r <- step defState st
376 r <- step defState st
367 case r of
377 case r of
368− Yield x s ->
378+ Yield x s
369− if x == e
370− then return True
379+ | x == e -> return True
380+ | otherwise -> go s
371− else go s
381+ Skip s -> go s
372− Stop -> return False
382+ Stop -> return False
373
383
374{-# INLINE_NORMAL notElem #-}
384{-# INLINE_NORMAL notElem #-}
375notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
385notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
376−notElem e (Stream step state) = go state
386+notElem e s = fmap not (elem e s)
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
387
387{-# INLINE_NORMAL all #-}
388{-# INLINE_NORMAL all #-}
388all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
389all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
@@ -391,11 +392,11 @@ all p (Stream step state) = go state
391 go st = do
392 go st = do
392 r <- step defState st
393 r <- step defState st
393 case r of
394 case r of
394− Yield x s ->
395+ Yield x s
395− if p x
396− then go s
396+ | p x -> go s
397− else return False
397+ | otherwise -> return False
398+ Skip s -> go s
398− Stop -> return True
399+ Stop -> return True
399
400
400{-# INLINE_NORMAL any #-}
401{-# INLINE_NORMAL any #-}
401any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
402any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
@@ -404,11 +405,11 @@ any p (Stream step state) = go state
404 go st = do
405 go st = do
405 r <- step defState st
406 r <- step defState st
406 case r of
407 case r of
407− Yield x s ->
408+ Yield x s
408− if p x
409− then return True
409+ | p x -> return True
410+ | otherwise -> go s
410− else go s
411+ Skip s -> go s
411− Stop -> return False
412+ Stop -> return False
412
413
413{-# INLINE_NORMAL maximum #-}
414{-# INLINE_NORMAL maximum #-}
414maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
415maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
@@ -418,15 +419,16 @@ maximum (Stream step state) = go Nothing state
418 r <- step defState st
419 r <- step defState st
419 case r of
420 case r of
420 Yield x s -> go (Just x) s
421 Yield x s -> go (Just x) s
422+ Skip s -> go Nothing s
421− Stop -> return Nothing
423+ Stop -> return Nothing
422 go (Just acc) st = do
424 go (Just acc) st = do
423 r <- step defState st
425 r <- step defState st
424 case r of
426 case r of
425− Yield x s ->
427+ Yield x s
426− if acc <= x
428+ | acc <= x -> go (Just x) s
427− then go (Just x) s
429+ | otherwise -> go (Just acc) s
428− else go (Just acc) s
430+ Skip s -> go (Just acc) s
429− Stop -> return (Just acc)
431+ Stop -> return (Just acc)
430
432
431{-# INLINE_NORMAL minimum #-}
433{-# INLINE_NORMAL minimum #-}
432minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
434minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
@@ -436,15 +438,16 @@ minimum (Stream step state) = go Nothing state
436 r <- step defState st
438 r <- step defState st
437 case r of
439 case r of
438 Yield x s -> go (Just x) s
440 Yield x s -> go (Just x) s
441+ Skip s -> go Nothing s
439− Stop -> return Nothing
442+ Stop -> return Nothing
440 go (Just acc) st = do
443 go (Just acc) st = do
441 r <- step defState st
444 r <- step defState st
442 case r of
445 case r of
443− Yield x s ->
446+ Yield x s
444− if acc <= x
445− then go (Just acc) s
447+ | acc <= x -> go (Just acc) s
446− else go (Just x) s
448+ | otherwise -> go (Just x) s
449+ Skip s -> go (Just acc) s
447− Stop -> return (Just acc)
450+ Stop -> return (Just acc)
448
451
449------------------------------------------------------------------------------
452------------------------------------------------------------------------------
450-- Map and Fold
453-- Map and Fold
@@ -468,10 +471,11 @@ toList = foldr (:) []
468toStreamK :: Monad m => Stream m a -> K.Stream m a
471toStreamK :: Monad m => Stream m a -> K.Stream m a
469toStreamK (Stream step state) = go state
472toStreamK (Stream step state) = go state
470 where
473 where
471− go st = K.Stream $ \gst stp _ yld -> do
474+ go st = K.Stream $ \gst stp sng yld -> do
472 r <- step gst st
475 r <- step gst st
473 case r of
476 case r of
474 Yield x s -> yld x (go s)
477 Yield x s -> yld x (go s)
478+ Skip s -> K.unStream (go s) gst stp sng yld
475 Stop -> stp
479 Stop -> stp
476
480
477#ifndef DISABLE_FUSION
481#ifndef DISABLE_FUSION
@@ -497,7 +501,8 @@ postscanlM' fstep begin (Stream step state) =
497 Yield x s -> do
501 Yield x s -> do
498 y <- fstep acc x
502 y <- fstep acc x
499 y `seq` return (Yield y (s, y))
503 y `seq` return (Yield y (s, y))
504+ Skip s -> return $ Skip (s, acc)
500− Stop -> return Stop
505+ Stop -> return Stop
501
506
502{-# INLINE_LATE scanlM' #-}
507{-# INLINE_LATE scanlM' #-}
503scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
508scanlM' :: 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)
520 r <- step (rstState gst) st
525 r <- step (rstState gst) st
521 return $ case r of
526 return $ case r of
522 Yield x s -> Yield x (s, i + 1)
527 Yield x s -> Yield x (s, i + 1)
528+ Skip s -> Skip (s, i)
523 Stop -> Stop
529 Stop -> Stop
524 step' _ (_, _) = return Stop
530 step' _ (_, _) = return Stop
525
531
@@ -534,7 +540,8 @@ takeWhileM f (Stream step state) = Stream step' state
534 Yield x s -> do
540 Yield x s -> do
535 b <- f x
541 b <- f x
536 return $ if b then Yield x s else Stop
542 return $ if b then Yield x s else Stop
543+ Skip s -> return $ Skip s
537− Stop -> return Stop
544+ Stop -> return Stop
538
545
539{-# INLINE takeWhile #-}
546{-# INLINE takeWhile #-}
540takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
547takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
@@ -542,15 +549,26 @@ takeWhile f = takeWhileM (return . f)
542
549
543{-# INLINE_NORMAL drop #-}
550{-# INLINE_NORMAL drop #-}
544drop :: Monad m => Int -> Stream m a -> Stream m a
551drop :: 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)
546 where
553 where
547 {-# INLINE_LATE step' #-}
554 {-# INLINE_LATE step' #-}
548− step' gst (st, i) = do
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
549− r <- step (rstState gst) st
566+ r <- step (rstState gst) st
567+ return $
550 case r of
568 case r of
551− Yield _ s | i > 0 -> step' gst (s, i - 1)
569+ Yield x s -> Yield x (s, Nothing)
552− Yield x s -> return $ Yield x (s, 0)
570+ Skip s -> Skip (s, Nothing)
553− Stop -> return Stop
571+ Stop -> Stop
554
572
555data DropWhileState s a
573data DropWhileState s a
556 = DropWhileDrop s
574 = DropWhileDrop s
@@ -568,14 +586,16 @@ dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
568 Yield x s -> do
586 Yield x s -> do
569 b <- f x
587 b <- f x
570 if b
588 if b
571− then step' gst (DropWhileDrop s)
589+ then return $ Skip (DropWhileDrop s)
572− else step' gst (DropWhileYield x s)
590+ else return $ Skip (DropWhileYield x s)
591+ Skip s -> return $ Skip (DropWhileDrop s)
573 Stop -> return Stop
592 Stop -> return Stop
574
593
575 step' gst (DropWhileNext st) = do
594 step' gst (DropWhileNext st) = do
576 r <- step (rstState gst) st
595 r <- step (rstState gst) st
577 case r of
596 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)
579 Stop -> return Stop
599 Stop -> return Stop
580
600
581 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
601 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
@@ -594,10 +614,11 @@ filterM f (Stream step state) = Stream step' state
594 case r of
614 case r of
595 Yield x s -> do
615 Yield x s -> do
596 b <- f x
616 b <- f x
597− if b
617+ return $ if b
598− then return $ Yield x s
618+ then Yield x s
599− else step' gst s
619+ else Skip s
620+ Skip s -> return $ Skip s
600− Stop -> return Stop
621+ Stop -> return Stop
601
622
602{-# INLINE filter #-}
623{-# INLINE filter #-}
603filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
624filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
@@ -617,6 +638,7 @@ mapM f (Stream step state) = Stream step' state
617 r <- step (rstState gst) st
638 r <- step (rstState gst) st
618 case r of
639 case r of
619 Yield x s -> f x >>= \a -> return $ Yield a s
640 Yield x s -> f x >>= \a -> return $ Yield a s
641+ Skip s -> return $ Skip s
620 Stop -> return Stop
642 Stop -> return Stop
621
643
622{-# INLINE map #-}
644{-# INLINE map #-}
@@ -648,9 +670,11 @@ zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
648 {-# INLINE_LATE step #-}
670 {-# INLINE_LATE step #-}
649 step gst (sa, sb, Nothing) = do
671 step gst (sa, sb, Nothing) = do
650 r <- stepa (rstState gst) sa
672 r <- stepa (rstState gst) sa
673+ return $
651− case r of
674+ case r of
652− Yield x sa' -> step gst (sa', sb, Just x)
675+ Yield x sa' -> Skip (sa', sb, Just x)
676+ Skip sa' -> Skip (sa', sb, Nothing)
653− Stop -> return Stop
677+ Stop -> Stop
654
678
655 step gst (sa, sb, Just x) = do
679 step gst (sa, sb, Just x) = do
656 r <- stepb (rstState gst) sb
680 r <- stepb (rstState gst) sb
@@ -658,7 +682,8 @@ zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
658 Yield y sb' -> do
682 Yield y sb' -> do
659 z <- f x y
683 z <- f x y
660 return $ Yield z (sa, sb', Nothing)
684 return $ Yield z (sa, sb', Nothing)
685+ Skip sb' -> return $ Skip (sa, sb', Just x)
661− Stop -> return Stop
686+ Stop -> return Stop
662
687
663{-# RULES "zipWithM xs xs"
688{-# RULES "zipWithM xs xs"
664 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
689 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}