@@ -6,6 +6,7 @@
6{-# LANGUAGE InstanceSigs #-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
8{-# LANGUAGE PatternSynonyms #-}
9+{-# LANGUAGE KindSignatures #-}
9{-# LANGUAGE ViewPatterns #-}
10{-# LANGUAGE ViewPatterns #-}
10{-# LANGUAGE RankNTypes #-}
11{-# LANGUAGE RankNTypes #-}
11{-# LANGUAGE UndecidableInstances #-} -- XXX
12{-# LANGUAGE UndecidableInstances #-} -- XXX
@@ -45,23 +46,45 @@ module Streamly.Streams.StreamK.Type
45 , foldStreamShared
46 , foldStreamShared
46 , foldStreamSVar
47 , foldStreamSVar
48 -- instances
60 -- instances
49 , consMStream
63 , consMStream
56 , Streaming -- deprecated
79 , Streaming -- deprecated
60−import Control.Monad (void)
83+import Control.Monad (void, ap, (>=>))
61import Control.Monad.IO.Class (MonadIO(liftIO))
84import Control.Monad.IO.Class (MonadIO(liftIO))
62import Control.Monad.Trans.Class (MonadTrans(lift))
85import Control.Monad.Trans.Class (MonadTrans(lift))
63import Data.Semigroup (Semigroup(..))
86import Data.Semigroup (Semigroup(..))
64−import Prelude hiding (map)
87+import Prelude hiding (map, mapM, concatMap, foldr)
66import Streamly.SVar
89import Streamly.SVar
@@ -247,6 +270,68 @@ consK k r = mkStream $ \_ yld _ _ -> k (\x -> yld x r)
248-- XXX Build a stream from a repeating callback function.
271-- XXX Build a stream from a repeating callback function.
273+------------------------------------------------------------------------------
275+------------------------------------------------------------------------------
279+-- faster than consM because there is no bind.
280+-- | Construct a stream by adding a pure value at the head of an existing
281+-- stream. For serial streams this is the same as @(return a) \`consM` r@ but
282+-- more efficient. For concurrent streams this is not concurrent whereas
283+-- 'consM' is concurrent. For example:
286+-- > toList $ 1 \`cons` 2 \`cons` 3 \`cons` nil
291+{-# INLINE_NORMAL cons #-}
292+cons :: IsStream t => a -> t m a -> t m a
293+cons a r = mkStream $ \_ yld _ _ -> yld a r
297+-- | Operator equivalent of 'cons'.
300+-- > toList $ 1 .: 2 .: 3 .: nil
306+(.:) :: IsStream t => a -> t m a -> t m a
309+-- | An empty stream.
317+{-# INLINE_NORMAL nil #-}
318+nil :: IsStream t => t m a
319+nil = mkStream $ \_ _ _ stp -> stp
321+{-# INLINE_NORMAL yield #-}
322+yield :: IsStream t => a -> t m a
323+yield a = mkStream $ \_ _ single _ -> single a
325+{-# INLINE_NORMAL yieldM #-}
326+yieldM :: (Monad m, IsStream t) => m a -> t m a
327+yieldM m = fromStream $ mkStream $ \_ _ single _ -> m >>= single
329+-- XXX specialize to IO?
330+{-# INLINE consMBy #-}
331+consMBy :: (IsStream t, MonadAsync m) => (t m a -> t m a -> t m a)
332+ -> m a -> t m a -> t m a
333+consMBy f m r = (fromStream $ yieldM m) `f` r
250------------------------------------------------------------------------------
335------------------------------------------------------------------------------
251-- Folding a stream
336-- Folding a stream
252------------------------------------------------------------------------------
337------------------------------------------------------------------------------
@@ -345,6 +430,299 @@ instance IsStream Stream where
345 (|:) :: Monad m => m a -> Stream m a -> Stream m a
430 (|:) :: Monad m => m a -> Stream m a -> Stream m a
346 (|:) = consMStream
431 (|:) = consMStream
433+-------------------------------------------------------------------------------
434+-- foldr/build fusion
435+-------------------------------------------------------------------------------
437+-- XXX perhaps we can just use foldrSM/buildM everywhere as they are more
438+-- general and cover foldrS/buildS as well.
440+-- | The function 'f' decides how to reconstruct the stream. We could
441+-- reconstruct using a shared state (SVar) or without sharing the state.
443+{-# INLINE foldrSWith #-}
444+foldrSWith :: IsStream t
445+ => (forall r. State Stream m b
446+ -> (b -> t m b -> m r)
451+ -> (a -> t m b -> t m b) -> t m b -> t m a -> t m b
452+foldrSWith f step final m = go m
454+ go m1 = mkStream $ \st yld sng stp ->
455+ let run x = f st yld sng stp x
457+ single a = run $ step a final
458+ yieldk a r = run $ step a (go r)
459+ -- XXX if type a and b are the same we do not need adaptState, can we
460+ -- save some perf with that?
461+ -- XXX since we are using adaptState anyway here we can use
462+ -- foldStreamShared instead, will that save some perf?
463+ in foldStream (adaptState st) yieldk single stop m1
465+-- XXX we can use rewrite rules just for foldrSWith, if the function f is the
466+-- same we can rewrite it.
468+-- | Fold sharing the SVar state within the reconstructed stream
469+{-# INLINE_NORMAL foldrSShared #-}
470+foldrSShared :: IsStream t => (a -> t m b -> t m b) -> t m b -> t m a -> t m b
471+foldrSShared = foldrSWith foldStreamShared
473+-- XXX consM is a typeclass method, therefore rewritten already. Instead maybe
474+-- we can make consM polymorphic using rewrite rules.
475+-- {-# RULES "foldrSShared/id" foldrSShared consM nil = \x -> x #-}
476+{-# RULES "foldrSShared/nil"
477+ forall k z. foldrSShared k z nil = z #-}
478+{-# RULES "foldrSShared/single"
479+ forall k z x. foldrSShared k z (yield x) = k x z #-}
480+-- {-# RULES "foldrSShared/app" [1]
481+-- forall ys. foldrSShared consM ys = \xs -> xs `conjoin` ys #-}
483+-- | Lazy right associative fold to a stream.
484+{-# INLINE_NORMAL foldrS #-}
485+foldrS :: IsStream t => (a -> t m b -> t m b) -> t m b -> t m a -> t m b
486+foldrS = foldrSWith foldStream
488+{-# RULES "foldrS/id" foldrS cons nil = \x -> x #-}
489+{-# RULES "foldrS/nil" forall k z. foldrS k z nil = z #-}
490+-- See notes in GHC.Base about this rule
491+-- {-# RULES "foldr/cons"
492+-- forall k z x xs. foldrS k z (x `cons` xs) = k x (foldrS k z xs) #-}
493+{-# RULES "foldrS/single" forall k z x. foldrS k z (yield x) = k x z #-}
494+-- {-# RULES "foldrS/app" [1]
495+-- forall ys. foldrS cons ys = \xs -> xs `conjoin` ys #-}
497+-------------------------------------------------------------------------------
498+-- foldrS with monadic cons i.e. consM
499+-------------------------------------------------------------------------------
501+{-# INLINE foldrSMWith #-}
502+foldrSMWith :: (IsStream t, Monad m)
503+ => (forall r. State Stream m b
504+ -> (b -> t m b -> m r)
509+ -> (m a -> t m b -> t m b) -> t m b -> t m a -> t m b
510+foldrSMWith f step final m = go m
512+ go m1 = mkStream $ \st yld sng stp ->
513+ let run x = f st yld sng stp x
515+ single a = run $ step (return a) final
516+ yieldk a r = run $ step (return a) (go r)
517+ in foldStream (adaptState st) yieldk single stop m1
519+{-# INLINE_NORMAL foldrSM #-}
520+foldrSM :: (IsStream t, Monad m)
521+ => (m a -> t m b -> t m b) -> t m b -> t m a -> t m b
522+foldrSM = foldrSMWith foldStream
524+-- {-# RULES "foldrSM/id" foldrSM consM nil = \x -> x #-}
525+{-# RULES "foldrSM/nil" forall k z. foldrSM k z nil = z #-}
526+{-# RULES "foldrSM/single" forall k z x. foldrSM k z (yieldM x) = k x z #-}
527+-- {-# RULES "foldrSM/app" [1]
528+-- forall ys. foldrSM consM ys = \xs -> xs `conjoin` ys #-}
530+-- Like foldrSM but sharing the SVar state within the recostructed stream.
531+{-# INLINE_NORMAL foldrSMShared #-}
532+foldrSMShared :: (IsStream t, Monad m)
533+ => (m a -> t m b -> t m b) -> t m b -> t m a -> t m b
534+foldrSMShared = foldrSMWith foldStreamShared
536+-- {-# RULES "foldrSM/id" foldrSM consM nil = \x -> x #-}
537+{-# RULES "foldrSMShared/nil"
538+ forall k z. foldrSMShared k z nil = z #-}
539+{-# RULES "foldrSMShared/single"
540+ forall k z x. foldrSMShared k z (yieldM x) = k x z #-}
541+-- {-# RULES "foldrSM/app" [1]
542+-- forall ys. foldrSM consM ys = \xs -> xs `conjoin` ys #-}
544+-------------------------------------------------------------------------------
546+-------------------------------------------------------------------------------
548+-- build a stream by applying cons and nil to a build function
549+{-# INLINE_NORMAL buildS #-}
550+buildS :: IsStream t => ((a -> t m a -> t m a) -> t m a -> t m a) -> t m a
551+buildS g = g cons nil
553+{-# RULES "foldrS/buildS"
554+ forall k z (g :: (a -> t m a -> t m a) -> t m a -> t m a).
555+ foldrS k z (buildS g) = g k z #-}
557+{-# RULES "foldrS/cons/buildS"
558+ forall k z x (g :: (a -> t m a -> t m a) -> t m a -> t m a).
559+ foldrS k z (x `cons` buildS g) = k x (g k z) #-}
561+{-# RULES "foldrSShared/buildS"
562+ forall k z (g :: (a -> t m a -> t m a) -> t m a -> t m a).
563+ foldrSShared k z (buildS g) = g k z #-}
565+{-# RULES "foldrSShared/cons/buildS"
566+ forall k z x (g :: (a -> t m a -> t m a) -> t m a -> t m a).
567+ foldrSShared k z (x `cons` buildS g) = k x (g k z) #-}
569+-- build a stream by applying consM and nil to a build function
570+{-# INLINE_NORMAL buildSM #-}
571+buildSM :: (IsStream t, MonadAsync m)
572+ => ((m a -> t m a -> t m a) -> t m a -> t m a) -> t m a
573+buildSM g = g consM nil
575+{-# RULES "foldrSM/buildSM"
576+ forall k z (g :: (m a -> t m a -> t m a) -> t m a -> t m a).
577+ foldrSM k z (buildSM g) = g k z #-}
579+{-# RULES "foldrSMShared/buildSM"
580+ forall k z (g :: (m a -> t m a -> t m a) -> t m a -> t m a).
581+ foldrSMShared k z (buildSM g) = g k z #-}
583+-- Disabled because this may not fire as consM is a class Op
585+{-# RULES "foldrS/consM/buildSM"
586+ forall k z x (g :: (m a -> t m a -> t m a) -> t m a -> t m a)
587+ . foldrSM k z (x `consM` buildSM g)
592+-- Build using monadic build functions (continuations) instead of
593+-- reconstructing a stream.
594+{-# INLINE_NORMAL buildM #-}
595+buildM :: (IsStream t, MonadAsync m)
596+ => (forall r. (a -> t m a -> m r)
602+buildM g = mkStream $ \st yld sng stp ->
603+ g (\a r -> foldStream st yld sng stp (return a `consM` r)) sng stp
605+-- | Like 'buildM' but shares the SVar state across computations.
606+{-# INLINE_NORMAL sharedM #-}
607+sharedM :: (IsStream t, MonadAsync m)
608+ => (forall r. (a -> t m a -> m r)
614+sharedM g = mkStream $ \st yld sng stp ->
615+ g (\a r -> foldStreamShared st yld sng stp (return a `consM` r)) sng stp
617+-------------------------------------------------------------------------------
619+-------------------------------------------------------------------------------
621+{-# INLINE_NORMAL augmentS #-}
622+augmentS :: IsStream t
623+ => ((a -> t m a -> t m a) -> t m a -> t m a) -> t m a -> t m a
624+augmentS g xs = g cons xs
626+{-# RULES "augmentS/nil"
627+ forall (g :: (a -> t m a -> t m a) -> t m a -> t m a)
628+ . augmentS g nil = buildS g
631+{-# RULES "foldrS/augmentS"
632+ forall k z xs (g :: (a -> t m a -> t m a) -> t m a -> t m a)
633+ . foldrS k z (augmentS g xs) = g k (foldrS k z xs)
636+{-# RULES "augmentS/buildS"
637+ forall (g :: (a -> t m a -> t m a) -> t m a -> t m a)
638+ (h :: (a -> t m a -> t m a) -> t m a -> t m a)
639+ . augmentS g (buildS h) = buildS (\c n -> g c (h c n))
642+{-# INLINE_NORMAL augmentSM #-}
643+augmentSM :: (IsStream t, MonadAsync m)
644+ => ((m a -> t m a -> t m a) -> t m a -> t m a) -> t m a -> t m a
645+augmentSM g xs = g consM xs
647+{-# RULES "augmentSM/nil"
648+ forall (g :: (m a -> t m a -> t m a) -> t m a -> t m a)
649+ . augmentSM g nil = buildSM g
652+{-# RULES "foldrSM/augmentSM"
653+ forall k z xs (g :: (m a -> t m a -> t m a) -> t m a -> t m a)
654+ . foldrSM k z (augmentSM g xs) = g k (foldrSM k z xs)
657+{-# RULES "augmentSM/buildSM"
658+ forall (g :: (m a -> t m a -> t m a) -> t m a -> t m a)
659+ (h :: (m a -> t m a -> t m a) -> t m a -> t m a)
660+ . augmentSM g (buildSM h) = buildSM (\c n -> g c (h c n))
663+-------------------------------------------------------------------------------
664+-- Experimental foldrM/buildM
665+-------------------------------------------------------------------------------
667+-- | Lazy right fold with a monadic step function.
668+{-# INLINE_NORMAL foldrM #-}
669+foldrM :: IsStream t => (a -> m b -> m b) -> m b -> t m a -> m b
670+foldrM step acc m = go m
674+ single a = step a acc
675+ yieldk a r = step a (go r)
676+ in foldStream defState yieldk single stop m1
678+{-# INLINE_NORMAL foldrMKWith #-}
680+ :: (State Stream m a
681+ -> (a -> t m a -> m b)
686+ -> (a -> m b -> m b)
688+ -> ((a -> t m a -> m b) -> (a -> m b) -> m b -> m b)
690+foldrMKWith f step acc g = go g
694+ single a = step a acc
695+ yieldk a r = step a (go (\yld sng stp -> f defState yld sng stp r))
696+ in k yieldk single stop
699+{-# RULES "foldrM/buildS"
700+ forall k z (g :: (a -> t m a -> t m a) -> t m a -> t m a)
701+ . foldrM k z (buildS g)
705+-- XXX in which case will foldrM/buildM fusion be useful?
706+{-# RULES "foldrM/buildM"
707+ forall step acc (g :: (forall r.
713+ foldrM step acc (buildM g) = foldrMKWith foldStream step acc g
716+{-# RULES "foldrM/sharedM"
717+ forall step acc (g :: (forall r.
723+ foldrM step acc (sharedM g) = foldrMKWith foldStreamShared step acc g
348------------------------------------------------------------------------------
726------------------------------------------------------------------------------
349-- Semigroup
727-- Semigroup
350------------------------------------------------------------------------------
728------------------------------------------------------------------------------
@@ -356,6 +734,8 @@ instance IsStream Stream where
356-- @since 0.2.0
734-- @since 0.2.0
357{-# INLINE serial #-}
735{-# INLINE serial #-}
358serial :: IsStream t => t m a -> t m a -> t m a
736serial :: IsStream t => t m a -> t m a -> t m a
737+serial xs ys = augmentS (\c n -> foldrS c n xs) ys
359serial m1 m2 = go m1
739serial m1 m2 = go m1
361 go m = mkStream $ \st yld sng stp ->
741 go m = mkStream $ \st yld sng stp ->
@@ -363,6 +743,12 @@ serial m1 m2 = go m1
363 single a = yld a m2
743 single a = yld a m2
364 yieldk a r = yld a (go r)
744 yieldk a r = yld a (go r)
365 in foldStream st yieldk single stop m
745 in foldStream st yieldk single stop m
748+-- join/merge/append streams depending on consM
749+{-# INLINE conjoin #-}
750+conjoin :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
751+conjoin xs ys = augmentSM (\c n -> foldrSM c n xs) ys
367instance Semigroup (Stream m a) where
753instance Semigroup (Stream m a) where
368 (<>) = serial
754 (<>) = serial
@@ -371,18 +757,6 @@ instance Semigroup (Stream m a) where
372------------------------------------------------------------------------------
758------------------------------------------------------------------------------
374−-- | An empty stream.
383−nil :: IsStream t => t m a
384−nil = mkStream $ \_ _ _ stp -> stp
386instance Monoid (Stream m a) where
760instance Monoid (Stream m a) where
387 mempty = nil
761 mempty = nil
388 mappend = (<>)
762 mappend = (<>)
@@ -391,8 +765,26 @@ instance Monoid (Stream m a) where
391-- Functor
765-- Functor
392-------------------------------------------------------------------------------
766-------------------------------------------------------------------------------
769+{-# INLINE_LATE mapFB #-}
770+mapFB :: forall (t :: (* -> *) -> * -> *) b m a.
771+ (b -> t m b -> t m b) -> (a -> b) -> a -> t m b -> t m b
772+mapFB c f = \x ys -> c (f x) ys
775+"mapFB/mapFB" forall c f g. mapFB (mapFB c f) g = mapFB c (f . g)
776+"mapFB/id" forall c. mapFB c (\x -> x) = c
394{-# INLINE map #-}
779{-# INLINE map #-}
395map :: IsStream t => (a -> b) -> t m a -> t m b
780map :: IsStream t => (a -> b) -> t m a -> t m b
781+map f xs = buildS (\c n -> foldrS (mapFB c f) n xs)
783+-- XXX This definition might potentially be more efficient, but the cost in the
784+-- benchmark is dominated by unfoldrM cost so we cannot correctly determine
785+-- differences in the mapping cost. We should perhaps deduct the cost of
786+-- unfoldrM from the benchmarks and then compare.
396map f m = go m
788map f m = go m
@@ -400,6 +792,47 @@ map f m = go m
400 let single = sng . f
792 let single = sng . f
401 yieldk a r = yld (f a) (go r)
793 yieldk a r = yld (f a) (go r)
402 in foldStream (adaptState st) yieldk single stp m1
794 in foldStream (adaptState st) yieldk single stp m1
797+{-# INLINE_LATE mapMFB #-}
798+mapMFB :: Monad m => (m b -> t m b -> t m b) -> (a -> m b) -> m a -> t m b -> t m b
799+mapMFB c f = \x ys -> c (x >>= f) ys
802+"mapMFB/mapMFB" forall c f g. mapMFB (mapMFB c f) g = mapMFB c (f >=> g)
804+-- XXX These rules may never fire because pure/return type class rules will
807+"mapMFB/pure" forall c. mapMFB c (\x -> pure x) = c
808+"mapMFB/return" forall c. mapMFB c (\x -> return x) = c
811+-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
812+-- other stream types to overload it.
814+mapM :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m b
815+mapM f = foldrSShared (\x xs -> f x `consM` xs) nil
816+-- See note under map definition above.
820+ go m1 = mkStream $ \st yld sng stp ->
821+ let single a = f a >>= sng
822+ yieldk a r = foldStreamShared st yld sng stp $ f a |: go r
823+ in foldStream (adaptState st) yieldk single stp m1
826+-- This is experimental serial version supporting fusion.
828+-- XXX what if we do not want to fuse two concurrent mapMs?
829+-- XXX we can combine two concurrent mapM only if the SVar is of the same type
830+-- So for now we use it only for serial streams.
831+-- XXX fusion would be easier for monomoprhic stream types.
832+-- {-# RULES "mapM serial" mapM = mapMSerial #-}
833+{-# INLINE mapMSerial #-}
834+mapMSerial :: MonadAsync m => (a -> m b) -> Stream m a -> Stream m b
835+mapMSerial f xs = buildSM (\c n -> foldrSMShared (mapMFB c f) n xs)
404-- XXX in fact use the Stream type everywhere and only use polymorphism in the
837-- XXX in fact use the Stream type everywhere and only use polymorphism in the
405-- high level modules/prelude.
838-- high level modules/prelude.
@@ -410,9 +843,91 @@ instance Monad m => Functor (Stream m) where
410-- Transformers
843-- Transformers
411-------------------------------------------------------------------------------
844-------------------------------------------------------------------------------
413−{-# INLINE yieldM #-}
414−yieldM :: (Monad m, IsStream t) => m a -> t m a
415−yieldM m = fromStream $ mkStream $ \_ _ single _ -> m >>= single
417instance MonadTrans Stream where
846instance MonadTrans Stream where
418 lift = yieldM
847 lift = yieldM
849+-------------------------------------------------------------------------------
851+-------------------------------------------------------------------------------
853+-- | Detach a stream from an SVar
854+{-# INLINE unShare #-}
855+unShare :: IsStream t => t m a -> t m a
856+unShare x = mkStream $ \st yld sng stp ->
857+ foldStream st yld sng stp x
859+-- XXX This is just concatMapBy with arguments flipped. We need to keep this
860+-- instead of using a concatMap style definition because the bind
861+-- implementation in Async and WAsync streams show significant perf degradation
862+-- if the argument order is changed.
863+{-# INLINE bindWith #-}
866+ => (forall c. t m c -> t m c -> t m c)
870+bindWith par m1 f = go m1
873+ mkStream $ \st yld sng stp ->
874+ let foldShared = foldStreamShared st yld sng stp
875+ single a = foldShared $ unShare (f a)
876+ yieldk a r = foldShared $ unShare (f a) `par` go r
877+ in foldStream (adaptState st) yieldk single stp m
879+-- XXX express in terms of foldrS?
880+-- XXX can we use a different stream type for the generated stream being
881+-- falttened so that we can combine them differently and keep the resulting
883+-- XXX do we need specialize to IO?
884+-- XXX can we optimize when c and a are same, by removing the forall using
885+-- rewrite rules with type applications?
886+{-# INLINE concatMapBy #-}
889+ => (forall c. t m c -> t m c -> t m c)
893+concatMapBy par f xs = bindWith par xs f
895+{-# INLINE concatMap #-}
896+concatMap :: IsStream t => (a -> t m b) -> t m a -> t m b
897+concatMap f m = fromStream $
899+ (\a -> adapt $ toStream $ f a)
900+ (adapt $ toStream m)
904+-- XXX This fuses but when the stream is nil this performs poorly.
905+-- The filterAllOut benchmark degrades. Need to investigate and fix that.
906+{-# INLINE concatMap #-}
907+concatMap :: IsStream t => (a -> t m b) -> t m a -> t m b
908+concatMap f xs = buildS
909+ (\c n -> foldrS (\x b -> foldrS c b (f x)) n xs)
911+-- Stream polymorphic concatMap implementation
912+-- XXX need to use buildSM/foldrSMShared for parallel behavior
913+-- XXX unShare seems to degrade the fused performance
914+{-# INLINE_EARLY concatMap_ #-}
915+concatMap_ :: IsStream t => (a -> t m b) -> t m a -> t m b
916+concatMap_ f xs = buildS
917+ (\c n -> foldrSShared (\x b -> foldrSShared c b (unShare $ f x)) n xs)
920+instance Monad m => Applicative (Stream m) where
923+ {-# INLINE (<*>) #-}
926+-- NOTE: even though concatMap for StreamD is 3x faster compared to StreamK,
927+-- the monad instance of StreamD is slower than StreamK after foldr/build
929+instance Monad m => Monad (Stream m) where
930+ {-# INLINE return #-}
932+ {-# INLINE (>>=) #-}
933+ (>>=) = flip concatMap