on ea342a4 help sign in
984 lines 33 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE FlexibleInstances #-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE KindSignatures #-}
10{-# LANGUAGE ViewPatterns #-}
11{-# LANGUAGE RankNTypes #-}
12{-# LANGUAGE UndecidableInstances #-} -- XXX
14#include "../inline.hs"
16-- |
17-- Module : Streamly.Streams.StreamK.Type
18-- Copyright : (c) 2017 Harendra Kumar
19--
20-- License : BSD3
21-- Maintainer : harendra.kumar@gmail.com
22-- Stability : experimental
23-- Portability : GHC
24--
25--
26-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
27-- denotes a function as well as a Kontinuation.
28--
29module Streamly.Streams.StreamK.Type
30 (
31 -- * A class for streams
32 IsStream (..)
33 , adapt
35 -- * The stream type
36 , Stream ()
38 -- * Construction
39 , mkStream
40 , fromStopK
41 , fromYieldK
42 , consK
44 -- * Elimination
45 , foldStream
46 , foldStreamShared
47 , foldStreamSVar
49 -- * foldr/build
50 , foldrM
51 , foldrS
52 , foldrSM
53 , buildS
54 , buildM
55 , buildSM
56 , sharedM
57 , augmentS
58 , augmentSM
60 -- instances
61 , cons
62 , (.:)
63 , consMStream
64 , consMBy
65 , yieldM
66 , yield
68 , nil
69 , conjoin
70 , serial
71 , map
72 , mapM
73 , mapMSerial
74 , unShare
75 , concatMapBy
76 , concatMap
77 , bindWith
79 , Streaming -- deprecated
80 )
81where
83import Control.Monad (void, ap, (>=>))
84import Control.Monad.IO.Class (MonadIO(liftIO))
85import Control.Monad.Trans.Class (MonadTrans(lift))
86import Data.Semigroup (Semigroup(..))
87import Prelude hiding (map, mapM, concatMap, foldr)
89import Streamly.SVar
91------------------------------------------------------------------------------
92-- Basic stream type
93------------------------------------------------------------------------------
95-- | The type @Stream m a@ represents a monadic stream of values of type 'a'
96-- constructed using actions in monad 'm'. It uses stop, singleton and yield
97-- continuations equivalent to the following direct style type:
98--
99-- @
100-- data Stream m a = Stop | Singleton a | Yield a (Stream m a)
101-- @
102--
103-- To facilitate parallel composition we maintain a local state in an 'SVar'
104-- that is shared across and is used for synchronization of the streams being
105-- composed.
106--
107-- The singleton case can be expressed in terms of stop and yield but we have
108-- it as a separate case to optimize composition operations for streams with
109-- single element. We build singleton streams in the implementation of 'pure'
110-- for Applicative and Monad, and in 'lift' for MonadTrans.
111--
112-- XXX remove the Stream type parameter from State as it is always constant.
113-- We can remove it from SVar as well
114--
115newtype Stream m a =
116 MkStream (forall r.
117 State Stream m a -- state
118 -> (a -> Stream m a -> m r) -- yield
119 -> (a -> m r) -- singleton
120 -> m r -- stop
121 -> m r
122 )
124------------------------------------------------------------------------------
125-- Types that can behave as a Stream
126------------------------------------------------------------------------------
128infixr 5 `consM`
129infixr 5 |:
131-- XXX Use a different SVar based on the stream type. But we need to make sure
132-- that we do not lose performance due to polymorphism.
133--
134-- | Class of types that can represent a stream of elements of some type 'a' in
135-- some monad 'm'.
136--
137-- @since 0.2.0
138class IsStream t where
139 toStream :: t m a -> Stream m a
140 fromStream :: Stream m a -> t m a
141 -- | Constructs a stream by adding a monadic action at the head of an
142 -- existing stream. For example:
143 --
144 -- @
145 -- > toList $ getLine \`consM` getLine \`consM` nil
146 -- hello
147 -- world
148 -- ["hello","world"]
149 -- @
150 --
151 -- /Concurrent (do not use 'parallely' to construct infinite streams)/
152 --
153 -- @since 0.2.0
154 consM :: MonadAsync m => m a -> t m a -> t m a
155 -- | Operator equivalent of 'consM'. We can read it as "@parallel colon@"
156 -- to remember that @|@ comes before ':'.
157 --
158 -- @
159 -- > toList $ getLine |: getLine |: nil
160 -- hello
161 -- world
162 -- ["hello","world"]
163 -- @
164 --
165 -- @
166 -- let delay = threadDelay 1000000 >> print 1
167 -- drain $ serially $ delay |: delay |: delay |: nil
168 -- drain $ parallely $ delay |: delay |: delay |: nil
169 -- @
170 --
171 -- /Concurrent (do not use 'parallely' to construct infinite streams)/
172 --
173 -- @since 0.2.0
174 (|:) :: MonadAsync m => m a -> t m a -> t m a
175 -- We can define (|:) just as 'consM' but it is defined explicitly for each
176 -- type because we want to use SPECIALIZE pragma on the definition.
178-- | Same as 'IsStream'.
179--
180-- @since 0.1.0
181{-# DEPRECATED Streaming "Please use IsStream instead." #-}
182type Streaming = IsStream
184-------------------------------------------------------------------------------
185-- Type adapting combinators
186-------------------------------------------------------------------------------
188-- XXX Move/reset the State here by reconstructing the stream with cleared
189-- state. Can we make sure we do not do that when t1 = t2? If we do this then
190-- we do not need to do that explicitly using svarStyle. It would act as
191-- unShare when the stream type is the same.
192--
193-- | Adapt any specific stream type to any other specific stream type.
194--
195-- @since 0.1.0
196adapt :: (IsStream t1, IsStream t2) => t1 m a -> t2 m a
197adapt = fromStream . toStream
199------------------------------------------------------------------------------
200-- Building a stream
201------------------------------------------------------------------------------
203-- XXX The State is always parameterized by "Stream" which means State is not
204-- different for different stream types. So we have to manually make sure that
205-- when converting from one stream to another we migrate the state correctly.
206-- This can be fixed if we use a different SVar type for different streams.
207-- Currently we always use "SVar Stream" and therefore a different State type
208-- parameterized by that stream.
209--
210-- XXX Since t is coercible we should be able to coerce k
211-- mkStream k = fromStream $ MkStream $ coerce k
212--
213-- | Build a stream from an 'SVar', a stop continuation, a singleton stream
214-- continuation and a yield continuation.
215{-# INLINE_EARLY mkStream #-}
216mkStream :: IsStream t
217 => (forall r. State Stream m a
218 -> (a -> t m a -> m r)
219 -> (a -> m r)
220 -> m r
221 -> m r)
222 -> t m a
223mkStream k = fromStream $ MkStream $ \st yld sng stp ->
224 let yieldk a r = yld a (toStream r)
225 in k st yieldk sng stp
227{-# RULES "mkStream from stream" mkStream = mkStreamFromStream #-}
228mkStreamFromStream :: IsStream t
229 => (forall r. State Stream m a
230 -> (a -> Stream m a -> m r)
231 -> (a -> m r)
232 -> m r
233 -> m r)
234 -> t m a
235mkStreamFromStream k = fromStream $ MkStream k
237{-# RULES "mkStream stream" mkStream = mkStreamStream #-}
238mkStreamStream
239 :: (forall r. State Stream m a
240 -> (a -> Stream m a -> m r)
241 -> (a -> m r)
242 -> m r
243 -> m r)
244 -> Stream m a
245mkStreamStream = MkStream
247-- | A terminal function that has no continuation to follow.
248type StopK m = forall r. m r -> m r
250-- | A monadic continuation, it is a function that yields a value of type "a"
251-- and calls the argument (a -> m r) as a continuation with that value. We can
252-- also think of it as a callback with a handler (a -> m r). Category
253-- theorists call it a codensity type, a special type of right kan extension.
254type YieldK m a = forall r. (a -> m r) -> m r
256_wrapM :: Monad m => m a -> YieldK m a
257_wrapM m = \k -> m >>= k
259-- | Make an empty stream from a stop function.
260fromStopK :: IsStream t => StopK m -> t m a
261fromStopK k = mkStream $ \_ _ _ stp -> k stp
263-- | Make a singleton stream from a yield function.
264fromYieldK :: IsStream t => YieldK m a -> t m a
265fromYieldK k = mkStream $ \_ _ sng _ -> k sng
267-- | Add a yield function at the head of the stream.
268consK :: IsStream t => YieldK m a -> t m a -> t m a
269consK k r = mkStream $ \_ yld _ _ -> k (\x -> yld x r)
271-- XXX Build a stream from a repeating callback function.
273------------------------------------------------------------------------------
274-- Construction
275------------------------------------------------------------------------------
277infixr 5 `cons`
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:
284--
285-- @
286-- > toList $ 1 \`cons` 2 \`cons` 3 \`cons` nil
287-- [1,2,3]
288-- @
289--
290-- @since 0.1.0
291{-# INLINE_NORMAL cons #-}
292cons :: IsStream t => a -> t m a -> t m a
293cons a r = mkStream $ \_ yld _ _ -> yld a r
295infixr 5 .:
297-- | Operator equivalent of 'cons'.
298--
299-- @
300-- > toList $ 1 .: 2 .: 3 .: nil
301-- [1,2,3]
302-- @
303--
304-- @since 0.1.1
305{-# INLINE (.:) #-}
306(.:) :: IsStream t => a -> t m a -> t m a
307(.:) = cons
309-- | An empty stream.
310--
311-- @
312-- > toList nil
313-- []
314-- @
315--
316-- @since 0.1.0
317{-# INLINE_NORMAL nil #-}
318nil :: IsStream t => t m a
319nil = mkStream $ \_ _ _ stp -> stp
321{-# INLINE_NORMAL yield #-}
322yield :: IsStream t => a -> t m a
323yield a = mkStream $ \_ _ single _ -> single a
325{-# INLINE_NORMAL yieldM #-}
326yieldM :: (Monad m, IsStream t) => m a -> t m a
327yieldM m = fromStream $ mkStream $ \_ _ single _ -> m >>= single
329-- XXX specialize to IO?
330{-# INLINE consMBy #-}
331consMBy :: (IsStream t, MonadAsync m) => (t m a -> t m a -> t m a)
332 -> m a -> t m a -> t m a
333consMBy f m r = (fromStream $ yieldM m) `f` r
335------------------------------------------------------------------------------
336-- Folding a stream
337------------------------------------------------------------------------------
339-- | Fold a stream by providing an SVar, a stop continuation, a singleton
340-- continuation and a yield continuation. The stream would share the current
341-- SVar passed via the State.
342{-# INLINE_EARLY foldStreamShared #-}
343foldStreamShared
344 :: IsStream t
345 => State Stream m a
346 -> (a -> t m a -> m r)
347 -> (a -> m r)
348 -> m r
349 -> t m a
350 -> m r
351foldStreamShared st yld sng stp m =
352 let yieldk a x = yld a (fromStream x)
353 MkStream k = toStream m
354 in k st yieldk sng stp
356-- XXX write a similar rule for foldStream as well?
357{-# RULES "foldStreamShared from stream"
358 foldStreamShared = foldStreamSharedStream #-}
359foldStreamSharedStream
360 :: State Stream m a
361 -> (a -> Stream m a -> m r)
362 -> (a -> m r)
363 -> m r
364 -> Stream m a
365 -> m r
366foldStreamSharedStream st yld sng stp m =
367 let MkStream k = toStream m
368 in k st yld sng stp
370-- | Fold a stream by providing a State, stop continuation, a singleton
371-- continuation and a yield continuation. The stream will not use the SVar
372-- passed via State.
373{-# INLINE foldStream #-}
374foldStream
375 :: IsStream t
376 => State Stream m a
377 -> (a -> t m a -> m r)
378 -> (a -> m r)
379 -> m r
380 -> t m a
381 -> m r
382foldStream st yld sng stp m =
383 let yieldk a x = yld a (fromStream x)
384 MkStream k = toStream m
385 in k (adaptState st) yieldk sng stp
387-- Run the stream using a run function associated with the SVar that runs the
388-- streams with a captured snapshot of the monadic state.
389{-# INLINE foldStreamSVar #-}
390foldStreamSVar
391 :: (IsStream t, MonadIO m)
392 => SVar Stream m a
393 -> State Stream m a -- state
394 -> (a -> t m a -> m r) -- yield
395 -> (a -> m r) -- singleton
396 -> m r -- stop
397 -> t m a
398 -> m ()
399foldStreamSVar sv st yld sng stp m =
400 let mrun = runInIO $ svarMrun sv
401 in void $ liftIO $ mrun $ foldStreamShared st yld sng stp m
403-------------------------------------------------------------------------------
404-- Instances
405-------------------------------------------------------------------------------
407-- NOTE: specializing the function outside the instance definition seems to
408-- improve performance quite a bit at times, even if we have the same
409-- SPECIALIZE in the instance definition.
410{-# INLINE consMStream #-}
411{-# SPECIALIZE consMStream :: IO a -> Stream IO a -> Stream IO a #-}
412consMStream :: (Monad m) => m a -> Stream m a -> Stream m a
413consMStream m r = MkStream $ \_ yld _ _ -> m >>= \a -> yld a r
415-------------------------------------------------------------------------------
416-- IsStream Stream
417-------------------------------------------------------------------------------
419instance IsStream Stream where
420 toStream = id
421 fromStream = id
423 {-# INLINE consM #-}
424 {-# SPECIALIZE consM :: IO a -> Stream IO a -> Stream IO a #-}
425 consM :: Monad m => m a -> Stream m a -> Stream m a
426 consM = consMStream
428 {-# INLINE (|:) #-}
429 {-# SPECIALIZE (|:) :: IO a -> Stream IO a -> Stream IO a #-}
430 (|:) :: Monad m => m a -> Stream m a -> Stream m a
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.
442--
443{-# INLINE foldrSWith #-}
444foldrSWith :: IsStream t
445 => (forall r. State Stream m b
446 -> (b -> t m b -> m r)
447 -> (b -> m r)
448 -> m r
449 -> t m b
450 -> m r)
451 -> (a -> t m b -> t m b) -> t m b -> t m a -> t m b
452foldrSWith f step final m = go m
453 where
454 go m1 = mkStream $ \st yld sng stp ->
455 let run x = f st yld sng stp x
456 stop = run final
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 #-}
470foldrSShared :: IsStream t => (a -> t m b -> t m b) -> t m b -> t m a -> t m b
471foldrSShared = 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 #-}
485foldrS :: IsStream t => (a -> t m b -> t m b) -> t m b -> t m a -> t m b
486foldrS = 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 #-}
502foldrSMWith :: (IsStream t, Monad m)
503 => (forall r. State Stream m b
504 -> (b -> t m b -> m r)
505 -> (b -> m r)
506 -> m r
507 -> t m b
508 -> m r)
509 -> (m a -> t m b -> t m b) -> t m b -> t m a -> t m b
510foldrSMWith f step final m = go m
511 where
512 go m1 = mkStream $ \st yld sng stp ->
513 let run x = f st yld sng stp x
514 stop = run final
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 #-}
520foldrSM :: (IsStream t, Monad m)
521 => (m a -> t m b -> t m b) -> t m b -> t m a -> t m b
522foldrSM = 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 #-}
532foldrSMShared :: (IsStream t, Monad m)
533 => (m a -> t m b -> t m b) -> t m b -> t m a -> t m b
534foldrSMShared = 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-------------------------------------------------------------------------------
545-- build
546-------------------------------------------------------------------------------
548-- build a stream by applying cons and nil to a build function
549{-# INLINE_NORMAL buildS #-}
550buildS :: IsStream t => ((a -> t m a -> t m a) -> t m a -> t m a) -> t m a
551buildS 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 #-}
571buildSM :: (IsStream t, MonadAsync m)
572 => ((m a -> t m a -> t m a) -> t m a -> t m a) -> t m a
573buildSM 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
584{-
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)
588 = k x (g k z)
589#-}
590-}
592-- Build using monadic build functions (continuations) instead of
593-- reconstructing a stream.
594{-# INLINE_NORMAL buildM #-}
595buildM :: (IsStream t, MonadAsync m)
596 => (forall r. (a -> t m a -> m r)
597 -> (a -> m r)
598 -> m r
599 -> m r
600 )
601 -> t m a
602buildM 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 #-}
607sharedM :: (IsStream t, MonadAsync m)
608 => (forall r. (a -> t m a -> m r)
609 -> (a -> m r)
610 -> m r
611 -> m r
612 )
613 -> t m a
614sharedM g = mkStream $ \st yld sng stp ->
615 g (\a r -> foldStreamShared st yld sng stp (return a `consM` r)) sng stp
617-------------------------------------------------------------------------------
618-- augment
619-------------------------------------------------------------------------------
621{-# INLINE_NORMAL augmentS #-}
622augmentS :: IsStream t
623 => ((a -> t m a -> t m a) -> t m a -> t m a) -> t m a -> t m a
624augmentS 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
629 #-}
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)
634 #-}
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))
640 #-}
642{-# INLINE_NORMAL augmentSM #-}
643augmentSM :: (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
645augmentSM 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
650 #-}
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)
655 #-}
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))
661 #-}
663-------------------------------------------------------------------------------
664-- Experimental foldrM/buildM
665-------------------------------------------------------------------------------
667-- | Lazy right fold with a monadic step function.
668{-# INLINE_NORMAL foldrM #-}
669foldrM :: IsStream t => (a -> m b -> m b) -> m b -> t m a -> m b
670foldrM step acc m = go m
671 where
672 go m1 =
673 let stop = acc
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 #-}
679foldrMKWith
680 :: (State Stream m a
681 -> (a -> t m a -> m b)
682 -> (a -> m b)
683 -> m b
684 -> t m a
685 -> m b)
686 -> (a -> m b -> m b)
687 -> m b
688 -> ((a -> t m a -> m b) -> (a -> m b) -> m b -> m b)
689 -> m b
690foldrMKWith f step acc g = go g
691 where
692 go k =
693 let stop = acc
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
698{-
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)
702 = g k z
703#-}
704-}
705-- XXX in which case will foldrM/buildM fusion be useful?
706{-# RULES "foldrM/buildM"
707 forall step acc (g :: (forall r.
708 (a -> t m a -> m r)
709 -> (a -> m r)
710 -> m r
711 -> m r
712 )).
713 foldrM step acc (buildM g) = foldrMKWith foldStream step acc g
714 #-}
716{-# RULES "foldrM/sharedM"
717 forall step acc (g :: (forall r.
718 (a -> t m a -> m r)
719 -> (a -> m r)
720 -> m r
721 -> m r
722 )).
723 foldrM step acc (sharedM g) = foldrMKWith foldStreamShared step acc g
724 #-}
726------------------------------------------------------------------------------
727-- Semigroup
728------------------------------------------------------------------------------
730-- | Polymorphic version of the 'Semigroup' operation '<>' of 'SerialT'.
731-- Appends two streams sequentially, yielding all elements from the first
732-- stream, and then all elements from the second stream.
733--
734-- @since 0.2.0
735{-# INLINE serial #-}
736serial :: IsStream t => t m a -> t m a -> t m a
737serial xs ys = augmentS (\c n -> foldrS c n xs) ys
738{-
739serial m1 m2 = go m1
740 where
741 go m = mkStream $ \st yld sng stp ->
742 let stop = foldStream st yld sng stp m2
743 single a = yld a m2
744 yieldk a r = yld a (go r)
745 in foldStream st yieldk single stop m
746-}
748-- join/merge/append streams depending on consM
749{-# INLINE conjoin #-}
750conjoin :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
751conjoin xs ys = augmentSM (\c n -> foldrSM c n xs) ys
753instance Semigroup (Stream m a) where
754 (<>) = serial
756------------------------------------------------------------------------------
757-- Monoid
758------------------------------------------------------------------------------
760instance Monoid (Stream m a) where
761 mempty = nil
762 mappend = (<>)
764-------------------------------------------------------------------------------
765-- Functor
766-------------------------------------------------------------------------------
768-- Note eta expanded
769{-# INLINE_LATE mapFB #-}
770mapFB :: forall (t :: (* -> *) -> * -> *) b m a.
771 (b -> t m b -> t m b) -> (a -> b) -> a -> t m b -> t m b
772mapFB c f = \x ys -> c (f x) ys
774{-# RULES
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
777 #-}
779{-# INLINE map #-}
780map :: IsStream t => (a -> b) -> t m a -> t m b
781map 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.
787{-
788map f m = go m
789 where
790 go m1 =
791 mkStream $ \st yld sng stp ->
792 let single = sng . f
793 yieldk a r = yld (f a) (go r)
794 in foldStream (adaptState st) yieldk single stp m1
795-}
797{-# INLINE_LATE mapMFB #-}
798mapMFB :: Monad m => (m b -> t m b -> t m b) -> (a -> m b) -> m a -> t m b -> t m b
799mapMFB c f = \x ys -> c (x >>= f) ys
801{-# RULES
802 "mapMFB/mapMFB" forall c f g. mapMFB (mapMFB c f) g = mapMFB c (f >=> g)
803 #-}
804-- XXX These rules may never fire because pure/return type class rules will
805-- fire first.
806{-
807"mapMFB/pure" forall c. mapMFB c (\x -> pure x) = c
808"mapMFB/return" forall c. mapMFB c (\x -> return x) = c
809-}
811-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
812-- other stream types to overload it.
813{-# INLINE mapM #-}
814mapM :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m b
815mapM f = foldrSShared (\x xs -> f x `consM` xs) nil
816-- See note under map definition above.
817{-
818mapM f m = go m
819 where
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
824 -}
826-- This is experimental serial version supporting fusion.
827--
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 #-}
834mapMSerial :: MonadAsync m => (a -> m b) -> Stream m a -> Stream m b
835mapMSerial f xs = buildSM (\c n -> foldrSMShared (mapMFB c f) n xs)
837-- XXX in fact use the Stream type everywhere and only use polymorphism in the
838-- high level modules/prelude.
839instance Monad m => Functor (Stream m) where
840 fmap = map
842-------------------------------------------------------------------------------
843-- Transformers
844-------------------------------------------------------------------------------
846instance MonadTrans Stream where
847 lift = yieldM
849-------------------------------------------------------------------------------
850-- Nesting
851-------------------------------------------------------------------------------
853-- | Detach a stream from an SVar
854{-# INLINE unShare #-}
855unShare :: IsStream t => t m a -> t m a
856unShare 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 #-}
864bindWith
865 :: IsStream t
866 => (forall c. t m c -> t m c -> t m c)
867 -> t m a
868 -> (a -> t m b)
869 -> t m b
870bindWith par m1 f = go m1
871 where
872 go m =
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
882-- stream different?
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?
887-- | Perform a 'concatMap' using a specified concat strategy. The first
888-- argument specifies a merge or concat function that is used to merge the
889-- streams generated by the map function. For example, the concat function
890-- could be 'serial', 'parallel', 'async', 'ahead' or any other zip or merge
891-- function.
892--
893-- @since 0.7.0
894{-# INLINE concatMapBy #-}
895concatMapBy
896 :: IsStream t
897 => (forall c. t m c -> t m c -> t m c)
898 -> (a -> t m b)
899 -> t m a
900 -> t m b
901concatMapBy par f xs = bindWith par xs f
903{-# INLINE concatMap #-}
904concatMap :: IsStream t => (a -> t m b) -> t m a -> t m b
905concatMap f m = fromStream $
906 concatMapBy serial
907 (\a -> adapt $ toStream $ f a)
908 (adapt $ toStream m)
910{-
911-- Fused version.
912-- XXX This fuses but when the stream is nil this performs poorly.
913-- The filterAllOut benchmark degrades. Need to investigate and fix that.
914{-# INLINE concatMap #-}
915concatMap :: IsStream t => (a -> t m b) -> t m a -> t m b
916concatMap f xs = buildS
917 (\c n -> foldrS (\x b -> foldrS c b (f x)) n xs)
919-- Stream polymorphic concatMap implementation
920-- XXX need to use buildSM/foldrSMShared for parallel behavior
921-- XXX unShare seems to degrade the fused performance
922{-# INLINE_EARLY concatMap_ #-}
923concatMap_ :: IsStream t => (a -> t m b) -> t m a -> t m b
924concatMap_ f xs = buildS
925 (\c n -> foldrSShared (\x b -> foldrSShared c b (unShare $ f x)) n xs)
926-}
928instance Monad m => Applicative (Stream m) where
929 {-# INLINE pure #-}
930 pure = yield
931 {-# INLINE (<*>) #-}
932 (<*>) = ap
934-- NOTE: even though concatMap for StreamD is 3x faster compared to StreamK,
935-- the monad instance of StreamD is slower than StreamK after foldr/build
936-- fusion.
937instance Monad m => Monad (Stream m) where
938 {-# INLINE return #-}
939 return = pure
940 {-# INLINE (>>=) #-}
941 (>>=) = flip concatMap
943{-
944-- ConcatMap recursively on itself using a merge strategy.
945concatLoopBy :: IsStream t
946 => (forall c. t m c -> t m c -> t m c)
947 -> (a -> t m a) -> t m a
948concatLoopBy = undefined
950-- This is mfix. Put another way, concatMap recursively on the output of a
951-- stream. Compare with iterate.
952concatLoop :: IsStream t => (a -> t m a) -> t m a
953concatLoop = concatLoopBy serial
955instance MonadFix (Stream m) where
956 mfix = concatLoop
958-- The SVar implementation is something similar to concatFeedBack, so maybe
959-- there is an opportunity to share the implementation here. In an SVar we run
960-- a part of an action (a stream), it yields an output and the rest of the
961-- stream, we yield the output and queue back the rest of the stream for
962-- further evaluation. Also see unfoldrA.
963--
964-- There could be multiple variants of this combinator, for example use a
965-- specific way of concating i.e. concatLoopBy. Which also includes combinators
966-- with different stop behaviors. For example if the Left values are errors we
967-- can stop the whole composition on errors or on specific errors.
968--
969-- We can also flip the serial/ahead append behavior e.g. instead of processing
970-- the Left output after the original stream we can reverse the order.
971--
972-- | Concat map on the 'Left' output of a stream and merge it back into the
973-- stream. The right output is yielded in the output stream.
974concatFeedBack :: IsStream t
975 => (b -> t m (Either a b)) -> t m (Either a b) -> t m a
976concatFeedBack = undefined
978-- Compare this with unfoldr. Start with a seed stream and generate a stream
979-- with a value and new seeds. The new seeds are fed back to generate a seed
980-- stream and so on. This is a stream level unfoldr.
981concatUnfoldr :: IsStream t
982 => (b -> t m (Maybe (a, b))) -> t m (Maybe (a, b)) -> t m a
983concatUnfoldr = undefined
984-}