git.iko.soy/streamly
on ffcfcac help sign in
streamly/src/Streamly/Streams/StreamD/Type.hs
545 lines 18 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE ViewPatterns #-}
10{-# LANGUAGE RankNTypes #-}
12#include "../inline.hs"
14-- |
15-- Module : Streamly.Streams.StreamD.Type
16-- Copyright : (c) 2018 Harendra Kumar
17-- Copyright : (c) Roman Leshchinskiy 2008-2010
18--
19-- License : BSD3
20-- Maintainer : streamly@composewell.com
21-- Stability : experimental
22-- Portability : GHC
24module Streamly.Streams.StreamD.Type
25 (
26 -- * The stream type
27 Step (..)
28 -- XXX UnStream is exported to avoid a performance issue in concatMap if we
29 -- use the pattern synonym "Stream".
30#if __GLASGOW_HASKELL__ >= 800
31 , Stream (Stream, UnStream)
32#else
33 , Stream (UnStream)
34 , pattern Stream
35#endif
36 , fromStreamK
37 , toStreamK
38 , fromStreamD
39 , map
40 , mapM
41 , yield
42 , yieldM
43 , concatMap
44 , concatMapM
46 , foldrT
47 , foldrM
48 , foldrMx
49 , foldr
50 , foldrS
52 , foldl'
53 , foldlM'
54 , foldlx'
55 , foldlMx'
57 , toList
58 , fromList
60 , eqBy
61 , cmpBy
62 , take
63 , GroupState (..) -- for inspection testing
64 , groupsOf
65 )
66where
68import Control.Applicative (liftA2)
69import Control.Monad (ap, when)
70import Control.Monad.Trans (lift, MonadTrans)
71import Data.Functor.Identity (Identity(..))
72import GHC.Base (build)
73import GHC.Types (SPEC(..))
74import Prelude hiding (map, mapM, foldr, take, concatMap)
76import Streamly.Internal.Data.SVar (State(..), adaptState, defState)
77import Streamly.Internal.Data.Fold.Types (Fold(..))
79import qualified Streamly.Streams.StreamK as K
81------------------------------------------------------------------------------
82-- The direct style stream type
83------------------------------------------------------------------------------
85-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
86-- the next state of the stream. 'Stop' indicates there are no more values in
87-- the stream.
88data Step s a = Yield a s | Skip s | Stop
90{-
91instance Functor (Step s) where
92 {-# INLINE fmap #-}
93 fmap f (Yield x s) = Yield (f x) s
94 fmap _ (Skip s) = Skip s
95 fmap _ Stop = Stop
96-}
98-- gst = global state
99-- | A stream consists of a step function that generates the next step given a
100-- current state, and the current state.
101data Stream m a =
102 forall s. UnStream (State K.Stream m a -> s -> m (Step s a)) s
104unShare :: Stream m a -> Stream m a
105unShare (UnStream step state) = UnStream step' state
106 where step' gst = step (adaptState gst)
108pattern Stream :: (State K.Stream m a -> s -> m (Step s a)) -> s -> Stream m a
109pattern Stream step state <- (unShare -> UnStream step state)
110 where Stream = UnStream
112#if __GLASGOW_HASKELL__ >= 802
113{-# COMPLETE Stream #-}
114#endif
116{-# INLINE_LATE fromStreamK #-}
117fromStreamK :: Monad m => K.Stream m a -> Stream m a
118fromStreamK = Stream step
119 where
120 step gst m1 =
121 let stop = return Stop
122 single a = return $ Yield a K.nil
123 yieldk a r = return $ Yield a r
124 in K.foldStreamShared gst yieldk single stop m1
126-- Convert a direct stream to and from CPS encoded stream
127{-# INLINE_LATE toStreamK #-}
128toStreamK :: Monad m => Stream m a -> K.Stream m a
129toStreamK (Stream step state) = go state
130 where
131 go st = K.mkStream $ \gst yld sng stp -> do
132 r <- step gst st
133 case r of
134 Yield x s -> yld x (go s)
135 Skip s -> K.foldStreamShared gst yld sng stp $ go s
136 Stop -> stp
138#ifndef DISABLE_FUSION
139{-# RULES "fromStreamK/toStreamK fusion"
140 forall s. toStreamK (fromStreamK s) = s #-}
141{-# RULES "toStreamK/fromStreamK fusion"
142 forall s. fromStreamK (toStreamK s) = s #-}
143#endif
145------------------------------------------------------------------------------
146-- Converting folds
147------------------------------------------------------------------------------
149{-# INLINE fromStreamD #-}
150fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
151fromStreamD = K.fromStream . toStreamK
153------------------------------------------------------------------------------
154-- Instances
155------------------------------------------------------------------------------
157-- | Map a monadic function over a 'Stream'
158{-# INLINE_NORMAL mapM #-}
159mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
160mapM f (Stream step state) = Stream step' state
161 where
162 {-# INLINE_LATE step' #-}
163 step' gst st = do
164 r <- step (adaptState gst) st
165 case r of
166 Yield x s -> f x >>= \a -> return $ Yield a s
167 Skip s -> return $ Skip s
168 Stop -> return Stop
170{-# INLINE map #-}
171map :: Monad m => (a -> b) -> Stream m a -> Stream m b
172map f = mapM (return . f)
174instance Monad m => Functor (Stream m) where
175 {-# INLINE fmap #-}
176 fmap = map
178------------------------------------------------------------------------------
179-- concatMap
180------------------------------------------------------------------------------
182{-# INLINE_NORMAL concatMapM #-}
183concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
184concatMapM f (Stream step state) = Stream step' (Left state)
185 where
186 {-# INLINE_LATE step' #-}
187 step' gst (Left st) = do
188 r <- step (adaptState gst) st
189 case r of
190 Yield a s -> do
191 b_stream <- f a
192 return $ Skip (Right (b_stream, s))
193 Skip s -> return $ Skip (Left s)
194 Stop -> return Stop
196 -- XXX flattenArrays is 5x faster than "concatMap fromArray". if somehow we
197 -- can get inner_step to inline and fuse here we can perhaps get the same
198 -- performance using "concatMap fromArray".
199 --
200 -- XXX using the pattern synonym "Stream" causes a major performance issue
201 -- here even if the synonym does not include an adaptState call. Need to
202 -- find out why. Is that something to be fixed in GHC?
203 step' gst (Right (UnStream inner_step inner_st, st)) = do
204 r <- inner_step (adaptState gst) inner_st
205 case r of
206 Yield b inner_s ->
207 return $ Yield b (Right (Stream inner_step inner_s, st))
208 Skip inner_s ->
209 return $ Skip (Right (Stream inner_step inner_s, st))
210 Stop -> return $ Skip (Left st)
212{-# INLINE concatMap #-}
213concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
214concatMap f = concatMapM (return . f)
216-- XXX The idea behind this rule is to rewrite any calls to "concatMap
217-- fromArray" automatically to flattenArrays which is much faster. However, we
218-- need an INLINE_EARLY on concatMap for this rule to fire. But if we use
219-- INLINE_EARLY on concatMap or fromArray then direct uses of
220-- "concatMap fromArray" (without the RULE) become much slower, this means
221-- "concatMap f" in general would become slower. Need to find a solution to
222-- this.
223--
224-- {-# RULES "concatMap Array.toStreamD"
225-- concatMap Array.toStreamD = Array.flattenArray #-}
227-- | Create a singleton 'Stream' from a pure value.
228{-# INLINE_NORMAL yield #-}
229yield :: Monad m => a -> Stream m a
230yield x = Stream (\_ s -> return $ step undefined s) True
231 where
232 {-# INLINE_LATE step #-}
233 step _ True = Yield x False
234 step _ False = Stop
236instance Monad m => Applicative (Stream m) where
237 {-# INLINE pure #-}
238 pure = yield
239 {-# INLINE (<*>) #-}
240 (<*>) = ap
242-- NOTE: even though concatMap for StreamD is 4x faster compared to StreamK,
243-- the monad instance does not seem to be significantly faster.
244instance Monad m => Monad (Stream m) where
245 {-# INLINE return #-}
246 return = pure
247 {-# INLINE (>>=) #-}
248 (>>=) = flip concatMap
250-- XXX Use of SPEC constructor in folds causes 2x performance degradation in
251-- one shot operations, but helps immensely in operations composed of multiple
252-- combinators or the same combinator many times. There seems to be an
253-- opportunity to optimize here, can we get both, better perf for single ops
254-- as well as composed ops? Without SPEC, all single operation benchmarks
255-- become 2x faster.
257-- The way we want a left fold to be strict, dually we want the right fold to
258-- be lazy. The correct signature of the fold function to keep it lazy must be
259-- (a -> m b -> m b) instead of (a -> b -> m b). We were using the latter
260-- earlier, which is incorrect. In the latter signature we have to feed the
261-- value to the fold function after evaluating the monadic action, depending on
262-- the bind behavior of the monad, the action may get evaluated immediately
263-- introducing unnecessary strictness to the fold. If the implementation is
264-- lazy the following example, must work:
265--
266-- S.foldrM (\x t -> if x then return t else return False) (return True)
267-- (S.fromList [False,undefined] :: SerialT IO Bool)
268--
269{-# INLINE_NORMAL foldrM #-}
270foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
271foldrM f z (Stream step state) = go SPEC state
272 where
273 {-# INLINE_LATE go #-}
274 go !_ st = do
275 r <- step defState st
276 case r of
277 Yield x s -> f x (go SPEC s)
278 Skip s -> go SPEC s
279 Stop -> z
281{-# INLINE_NORMAL foldrMx #-}
282foldrMx :: Monad m
283 => (a -> m x -> m x) -> m x -> (m x -> m b) -> Stream m a -> m b
284foldrMx fstep final convert (Stream step state) = convert $ go SPEC state
285 where
286 {-# INLINE_LATE go #-}
287 go !_ st = do
288 r <- step defState st
289 case r of
290 Yield x s -> fstep x (go SPEC s)
291 Skip s -> go SPEC s
292 Stop -> final
294-- Note that foldr works on pure values, therefore it becomes necessarily
295-- strict when the monad m is strict. In that case it cannot terminate early,
296-- it would evaluate all of its input. Though, this should work fine with lazy
297-- monads. For example, if "any" is implemented using "foldr" instead of
298-- "foldrM" it performs the same with Identity monad but performs 1000x slower
299-- with IO monad.
300--
301{-# INLINE_NORMAL foldr #-}
302foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
303foldr f z = foldrM (\a b -> liftA2 f (return a) b) (return z)
305-- | Create a singleton 'Stream' from a monadic action.
306{-# INLINE_NORMAL yieldM #-}
307yieldM :: Monad m => m a -> Stream m a
308yieldM m = Stream step True
309 where
310 {-# INLINE_LATE step #-}
311 step _ True = m >>= \x -> return $ Yield x False
312 step _ False = return Stop
314-- this performs horribly, should not be used
315{-# INLINE_NORMAL foldrS #-}
316foldrS
317 :: Monad m
318 => (a -> Stream m b -> Stream m b)
319 -> Stream m b
320 -> Stream m a
321 -> Stream m b
322foldrS f final (Stream step state) = go SPEC state
323 where
324 {-# INLINE_LATE go #-}
325 go !_ st = do
326 -- defState??
327 r <- yieldM $ step defState st
328 case r of
329 Yield x s -> f x (go SPEC s)
330 Skip s -> go SPEC s
331 Stop -> final
333-- Right fold to some transformer (T) monad. This can be useful to implement
334-- stateless combinators like map, filtering, insertions, takeWhile, dropWhile.
335--
336{-# INLINE_NORMAL foldrT #-}
337foldrT :: (Monad m, Monad (t m), MonadTrans t)
338 => (a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
339foldrT f final (Stream step state) = go SPEC state
340 where
341 {-# INLINE_LATE go #-}
342 go !_ st = do
343 r <- lift $ step defState st
344 case r of
345 Yield x s -> f x (go SPEC s)
346 Skip s -> go SPEC s
347 Stop -> final
349{-# INLINE_NORMAL toList #-}
350toList :: Monad m => Stream m a -> m [a]
351toList = foldr (:) []
353-- Use foldr/build fusion to fuse with list consumers
354-- This can be useful when using the IsList instance
355{-# INLINE_LATE toListFB #-}
356toListFB :: (a -> b -> b) -> b -> Stream Identity a -> b
357toListFB c n (Stream step state) = go state
358 where
359 go st = case runIdentity (step defState st) of
360 Yield x s -> x `c` go s
361 Skip s -> go s
362 Stop -> n
364{-# RULES "toList Identity" toList = toListId #-}
365{-# INLINE_EARLY toListId #-}
366toListId :: Stream Identity a -> Identity [a]
367toListId s = Identity $ build (\c n -> toListFB c n s)
369-- XXX run begin action only if the stream is not empty.
370{-# INLINE_NORMAL foldlMx' #-}
371foldlMx' :: Monad m => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
372foldlMx' fstep begin done (Stream step state) =
373 begin >>= \x -> go SPEC x state
374 where
375 -- XXX !acc?
376 {-# INLINE_LATE go #-}
377 go !_ acc st = acc `seq` do
378 r <- step defState st
379 case r of
380 Yield x s -> do
381 acc' <- fstep acc x
382 go SPEC acc' s
383 Skip s -> go SPEC acc s
384 Stop -> done acc
386{-# INLINE foldlx' #-}
387foldlx' :: Monad m => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
388foldlx' fstep begin done m =
389 foldlMx' (\b a -> return (fstep b a)) (return begin) (return . done) m
391-- XXX implement in terms of foldlMx'?
392{-# INLINE_NORMAL foldlM' #-}
393foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
394foldlM' fstep begin (Stream step state) = go SPEC begin state
395 where
396 {-# INLINE_LATE go #-}
397 go !_ acc st = acc `seq` do
398 r <- step defState st
399 case r of
400 Yield x s -> do
401 acc' <- fstep acc x
402 go SPEC acc' s
403 Skip s -> go SPEC acc s
404 Stop -> return acc
406{-# INLINE foldl' #-}
407foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
408foldl' fstep = foldlM' (\b a -> return (fstep b a))
410-- | Convert a list of pure values to a 'Stream'
411{-# INLINE_LATE fromList #-}
412fromList :: Monad m => [a] -> Stream m a
413fromList = Stream step
414 where
415 {-# INLINE_LATE step #-}
416 step _ (x:xs) = return $ Yield x xs
417 step _ [] = return Stop
419------------------------------------------------------------------------------
420-- Comparisons
421------------------------------------------------------------------------------
423{-# INLINE_NORMAL eqBy #-}
424eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
425eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
426 where
427 eq_loop0 !_ s1 s2 = do
428 r <- step1 defState s1
429 case r of
430 Yield x s1' -> eq_loop1 SPEC x s1' s2
431 Skip s1' -> eq_loop0 SPEC s1' s2
432 Stop -> eq_null s2
434 eq_loop1 !_ x s1 s2 = do
435 r <- step2 defState s2
436 case r of
437 Yield y s2'
438 | eq x y -> eq_loop0 SPEC s1 s2'
439 | otherwise -> return False
440 Skip s2' -> eq_loop1 SPEC x s1 s2'
441 Stop -> return False
443 eq_null s2 = do
444 r <- step2 defState s2
445 case r of
446 Yield _ _ -> return False
447 Skip s2' -> eq_null s2'
448 Stop -> return True
450-- | Compare two streams lexicographically
451{-# INLINE_NORMAL cmpBy #-}
452cmpBy
453 :: Monad m
454 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
455cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
456 where
457 cmp_loop0 !_ s1 s2 = do
458 r <- step1 defState s1
459 case r of
460 Yield x s1' -> cmp_loop1 SPEC x s1' s2
461 Skip s1' -> cmp_loop0 SPEC s1' s2
462 Stop -> cmp_null s2
464 cmp_loop1 !_ x s1 s2 = do
465 r <- step2 defState s2
466 case r of
467 Yield y s2' -> case x `cmp` y of
468 EQ -> cmp_loop0 SPEC s1 s2'
469 c -> return c
470 Skip s2' -> cmp_loop1 SPEC x s1 s2'
471 Stop -> return GT
473 cmp_null s2 = do
474 r <- step2 defState s2
475 case r of
476 Yield _ _ -> return LT
477 Skip s2' -> cmp_null s2'
478 Stop -> return EQ
480{-# INLINE_NORMAL take #-}
481take :: Monad m => Int -> Stream m a -> Stream m a
482take n (Stream step state) = n `seq` Stream step' (state, 0)
483 where
484 {-# INLINE_LATE step' #-}
485 step' gst (st, i) | i < n = do
486 r <- step gst st
487 return $ case r of
488 Yield x s -> Yield x (s, i + 1)
489 Skip s -> Skip (s, i)
490 Stop -> Stop
491 step' _ (_, _) = return Stop
493------------------------------------------------------------------------------
494-- Grouping/Splitting
495------------------------------------------------------------------------------
497-- s = stream state, fs = fold state
498data GroupState s fs
499 = GroupStart s
500 | GroupBuffer s fs Int
501 | GroupYield fs (GroupState s fs)
502 | GroupFinish
504{-# INLINE_NORMAL groupsOf #-}
505groupsOf
506 :: Monad m
507 => Int
508 -> Fold m a b
509 -> Stream m a
510 -> Stream m b
511groupsOf n (Fold fstep initial extract) (Stream step state) =
512 n `seq` Stream step' (GroupStart state)
514 where
516 {-# INLINE_LATE step' #-}
517 step' _ (GroupStart st) = do
518 -- XXX shall we use the Natural type instead? Need to check performance
519 -- implications.
520 when (n <= 0) $
521 -- XXX we can pass the module string from the higher level API
522 error $ "Streamly.Streams.StreamD.Type.groupsOf: the size of "
523 ++ "groups [" ++ show n ++ "] must be a natural number"
524 -- fs = fold state
525 fs <- initial
526 return $ Skip (GroupBuffer st fs 0)
528 step' gst (GroupBuffer st fs i) = do
529 r <- step (adaptState gst) st
530 case r of
531 Yield x s -> do
532 fs' <- fstep fs x
533 let i' = i + 1
534 return $
535 if i' >= n
536 then Skip (GroupYield fs' (GroupStart s))
537 else Skip (GroupBuffer s fs' i')
538 Skip s -> return $ Skip (GroupBuffer s fs i)
539 Stop -> return $ Skip (GroupYield fs GroupFinish)
541 step' _ (GroupYield fs next) = do
542 r <- extract fs
543 return $ Yield r next
545 step' _ GroupFinish = return Stop