git.iko.soy/streamly
on 3093e79 help sign in
streamly/src/Streamly/Streams/StreamD/Type.hs
527 lines 17 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 : harendra.kumar@gmail.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 GHC.Types (SPEC(..))
72import Prelude hiding (map, mapM, foldr, take, concatMap)
74import Streamly.Internal.Data.SVar (State(..), adaptState, defState)
75import Streamly.Internal.Data.Fold.Types (Fold(..))
77import qualified Streamly.Streams.StreamK as K
79------------------------------------------------------------------------------
80-- The direct style stream type
81------------------------------------------------------------------------------
83-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
84-- the next state of the stream. 'Stop' indicates there are no more values in
85-- the stream.
86data Step s a = Yield a s | Skip s | Stop
88{-
89instance Functor (Step s) where
90 {-# INLINE fmap #-}
91 fmap f (Yield x s) = Yield (f x) s
92 fmap _ (Skip s) = Skip s
93 fmap _ Stop = Stop
94-}
96-- gst = global state
97-- | A stream consists of a step function that generates the next step given a
98-- current state, and the current state.
99data Stream m a =
100 forall s. UnStream (State K.Stream m a -> s -> m (Step s a)) s
102unShare :: Stream m a -> Stream m a
103unShare (UnStream step state) = UnStream step' state
104 where step' gst = step (adaptState gst)
106pattern Stream :: (State K.Stream m a -> s -> m (Step s a)) -> s -> Stream m a
107pattern Stream step state <- (unShare -> UnStream step state)
108 where Stream = UnStream
110#if __GLASGOW_HASKELL__ >= 802
111{-# COMPLETE Stream #-}
112#endif
114{-# INLINE_LATE fromStreamK #-}
115fromStreamK :: Monad m => K.Stream m a -> Stream m a
116fromStreamK = Stream step
117 where
118 step gst m1 =
119 let stop = return Stop
120 single a = return $ Yield a K.nil
121 yieldk a r = return $ Yield a r
122 in K.foldStreamShared gst yieldk single stop m1
124-- Convert a direct stream to and from CPS encoded stream
125{-# INLINE_LATE toStreamK #-}
126toStreamK :: Monad m => Stream m a -> K.Stream m a
127toStreamK (Stream step state) = go state
128 where
129 go st = K.mkStream $ \gst yld sng stp -> do
130 r <- step gst st
131 case r of
132 Yield x s -> yld x (go s)
133 Skip s -> K.foldStreamShared gst yld sng stp $ go s
134 Stop -> stp
136#ifndef DISABLE_FUSION
137{-# RULES "fromStreamK/toStreamK fusion"
138 forall s. toStreamK (fromStreamK s) = s #-}
139{-# RULES "toStreamK/fromStreamK fusion"
140 forall s. fromStreamK (toStreamK s) = s #-}
141#endif
143------------------------------------------------------------------------------
144-- Converting folds
145------------------------------------------------------------------------------
147{-# INLINE fromStreamD #-}
148fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
149fromStreamD = K.fromStream . toStreamK
151------------------------------------------------------------------------------
152-- Instances
153------------------------------------------------------------------------------
155-- | Map a monadic function over a 'Stream'
156{-# INLINE_NORMAL mapM #-}
157mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
158mapM f (Stream step state) = Stream step' state
159 where
160 {-# INLINE_LATE step' #-}
161 step' gst st = do
162 r <- step (adaptState gst) st
163 case r of
164 Yield x s -> f x >>= \a -> return $ Yield a s
165 Skip s -> return $ Skip s
166 Stop -> return Stop
168{-# INLINE map #-}
169map :: Monad m => (a -> b) -> Stream m a -> Stream m b
170map f = mapM (return . f)
172instance Monad m => Functor (Stream m) where
173 {-# INLINE fmap #-}
174 fmap = map
176------------------------------------------------------------------------------
177-- concatMap
178------------------------------------------------------------------------------
180{-# INLINE_NORMAL concatMapM #-}
181concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
182concatMapM f (Stream step state) = Stream step' (Left state)
183 where
184 {-# INLINE_LATE step' #-}
185 step' gst (Left st) = do
186 r <- step (adaptState gst) st
187 case r of
188 Yield a s -> do
189 b_stream <- f a
190 return $ Skip (Right (b_stream, s))
191 Skip s -> return $ Skip (Left s)
192 Stop -> return Stop
194 -- XXX flattenArrays is 5x faster than "concatMap fromArray". if somehow we
195 -- can get inner_step to inline and fuse here we can perhaps get the same
196 -- performance using "concatMap fromArray".
197 --
198 -- XXX using the pattern synonym "Stream" causes a major performance issue
199 -- here even if the synonym does not include an adaptState call. Need to
200 -- find out why. Is that something to be fixed in GHC?
201 step' gst (Right (UnStream inner_step inner_st, st)) = do
202 r <- inner_step (adaptState gst) inner_st
203 case r of
204 Yield b inner_s ->
205 return $ Yield b (Right (Stream inner_step inner_s, st))
206 Skip inner_s ->
207 return $ Skip (Right (Stream inner_step inner_s, st))
208 Stop -> return $ Skip (Left st)
210{-# INLINE concatMap #-}
211concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
212concatMap f = concatMapM (return . f)
214-- XXX The idea behind this rule is to rewrite any calls to "concatMap
215-- fromArray" automatically to flattenArrays which is much faster. However, we
216-- need an INLINE_EARLY on concatMap for this rule to fire. But if we use
217-- INLINE_EARLY on concatMap or fromArray then direct uses of
218-- "concatMap fromArray" (without the RULE) become much slower, this means
219-- "concatMap f" in general would become slower. Need to find a solution to
220-- this.
221--
222-- {-# RULES "concatMap Array.toStreamD"
223-- concatMap Array.toStreamD = Array.flattenArray #-}
225-- | Create a singleton 'Stream' from a pure value.
226{-# INLINE_NORMAL yield #-}
227yield :: Monad m => a -> Stream m a
228yield x = Stream (\_ s -> return $ step undefined s) True
229 where
230 {-# INLINE_LATE step #-}
231 step _ True = Yield x False
232 step _ False = Stop
234instance Monad m => Applicative (Stream m) where
235 {-# INLINE pure #-}
236 pure = yield
237 {-# INLINE (<*>) #-}
238 (<*>) = ap
240-- NOTE: even though concatMap for StreamD is 4x faster compared to StreamK,
241-- the monad instance does not seem to be significantly faster.
242instance Monad m => Monad (Stream m) where
243 {-# INLINE return #-}
244 return = pure
245 {-# INLINE (>>=) #-}
246 (>>=) = flip concatMap
248-- XXX Use of SPEC constructor in folds causes 2x performance degradation in
249-- one shot operations, but helps immensely in operations composed of multiple
250-- combinators or the same combinator many times. There seems to be an
251-- opportunity to optimize here, can we get both, better perf for single ops
252-- as well as composed ops? Without SPEC, all single operation benchmarks
253-- become 2x faster.
255-- The way we want a left fold to be strict, dually we want the right fold to
256-- be lazy. The correct signature of the fold function to keep it lazy must be
257-- (a -> m b -> m b) instead of (a -> b -> m b). We were using the latter
258-- earlier, which is incorrect. In the latter signature we have to feed the
259-- value to the fold function after evaluating the monadic action, depending on
260-- the bind behavior of the monad, the action may get evaluated immediately
261-- introducing unnecessary strictness to the fold. If the implementation is
262-- lazy the following example, must work:
263--
264-- S.foldrM (\x t -> if x then return t else return False) (return True)
265-- (S.fromList [False,undefined] :: SerialT IO Bool)
266--
267{-# INLINE_NORMAL foldrM #-}
268foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
269foldrM f z (Stream step state) = go SPEC state
270 where
271 {-# INLINE_LATE go #-}
272 go !_ st = do
273 r <- step defState st
274 case r of
275 Yield x s -> f x (go SPEC s)
276 Skip s -> go SPEC s
277 Stop -> z
279{-# INLINE_NORMAL foldrMx #-}
280foldrMx :: Monad m
281 => (a -> m x -> m x) -> m x -> (m x -> m b) -> Stream m a -> m b
282foldrMx fstep final convert (Stream step state) = convert $ go SPEC state
283 where
284 {-# INLINE_LATE go #-}
285 go !_ st = do
286 r <- step defState st
287 case r of
288 Yield x s -> fstep x (go SPEC s)
289 Skip s -> go SPEC s
290 Stop -> final
292-- Note that foldr works on pure values, therefore it becomes necessarily
293-- strict when the monad m is strict. In that case it cannot terminate early,
294-- it would evaluate all of its input. Though, this should work fine with lazy
295-- monads. For example, if "any" is implemented using "foldr" instead of
296-- "foldrM" it performs the same with Identity monad but performs 1000x slower
297-- with IO monad.
298--
299{-# INLINE_NORMAL foldr #-}
300foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
301foldr f z = foldrM (\a b -> liftA2 f (return a) b) (return z)
303-- | Create a singleton 'Stream' from a monadic action.
304{-# INLINE_NORMAL yieldM #-}
305yieldM :: Monad m => m a -> Stream m a
306yieldM m = Stream step True
307 where
308 {-# INLINE_LATE step #-}
309 step _ True = m >>= \x -> return $ Yield x False
310 step _ False = return Stop
312-- this performs horribly, should not be used
313{-# INLINE_NORMAL foldrS #-}
314foldrS
315 :: Monad m
316 => (a -> Stream m b -> Stream m b)
317 -> Stream m b
318 -> Stream m a
319 -> Stream m b
320foldrS f final (Stream step state) = go SPEC state
321 where
322 {-# INLINE_LATE go #-}
323 go !_ st = do
324 -- defState??
325 r <- yieldM $ step defState st
326 case r of
327 Yield x s -> f x (go SPEC s)
328 Skip s -> go SPEC s
329 Stop -> final
331-- Right fold to some transformer (T) monad. This can be useful to implement
332-- stateless combinators like map, filtering, insertions, takeWhile, dropWhile.
333--
334{-# INLINE_NORMAL foldrT #-}
335foldrT :: (Monad m, Monad (t m), MonadTrans t)
336 => (a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
337foldrT f final (Stream step state) = go SPEC state
338 where
339 {-# INLINE_LATE go #-}
340 go !_ st = do
341 r <- lift $ step defState st
342 case r of
343 Yield x s -> f x (go SPEC s)
344 Skip s -> go SPEC s
345 Stop -> final
347{-# INLINE toList #-}
348toList :: Monad m => Stream m a -> m [a]
349toList = foldr (:) []
351-- XXX run begin action only if the stream is not empty.
352{-# INLINE_NORMAL foldlMx' #-}
353foldlMx' :: Monad m => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
354foldlMx' fstep begin done (Stream step state) =
355 begin >>= \x -> go SPEC x state
356 where
357 -- XXX !acc?
358 {-# INLINE_LATE go #-}
359 go !_ acc st = acc `seq` do
360 r <- step defState st
361 case r of
362 Yield x s -> do
363 acc' <- fstep acc x
364 go SPEC acc' s
365 Skip s -> go SPEC acc s
366 Stop -> done acc
368{-# INLINE foldlx' #-}
369foldlx' :: Monad m => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
370foldlx' fstep begin done m =
371 foldlMx' (\b a -> return (fstep b a)) (return begin) (return . done) m
373-- XXX implement in terms of foldlMx'?
374{-# INLINE_NORMAL foldlM' #-}
375foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
376foldlM' fstep begin (Stream step state) = go SPEC begin state
377 where
378 {-# INLINE_LATE go #-}
379 go !_ acc st = acc `seq` do
380 r <- step defState st
381 case r of
382 Yield x s -> do
383 acc' <- fstep acc x
384 go SPEC acc' s
385 Skip s -> go SPEC acc s
386 Stop -> return acc
388{-# INLINE foldl' #-}
389foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
390foldl' fstep = foldlM' (\b a -> return (fstep b a))
392-- | Convert a list of pure values to a 'Stream'
393{-# INLINE_LATE fromList #-}
394fromList :: Monad m => [a] -> Stream m a
395fromList = Stream step
396 where
397 {-# INLINE_LATE step #-}
398 step _ (x:xs) = return $ Yield x xs
399 step _ [] = return Stop
401------------------------------------------------------------------------------
402-- Comparisons
403------------------------------------------------------------------------------
405{-# INLINE_NORMAL eqBy #-}
406eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
407eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
408 where
409 eq_loop0 !_ s1 s2 = do
410 r <- step1 defState s1
411 case r of
412 Yield x s1' -> eq_loop1 SPEC x s1' s2
413 Skip s1' -> eq_loop0 SPEC s1' s2
414 Stop -> eq_null s2
416 eq_loop1 !_ x s1 s2 = do
417 r <- step2 defState s2
418 case r of
419 Yield y s2'
420 | eq x y -> eq_loop0 SPEC s1 s2'
421 | otherwise -> return False
422 Skip s2' -> eq_loop1 SPEC x s1 s2'
423 Stop -> return False
425 eq_null s2 = do
426 r <- step2 defState s2
427 case r of
428 Yield _ _ -> return False
429 Skip s2' -> eq_null s2'
430 Stop -> return True
432-- | Compare two streams lexicographically
433{-# INLINE_NORMAL cmpBy #-}
434cmpBy
435 :: Monad m
436 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
437cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
438 where
439 cmp_loop0 !_ s1 s2 = do
440 r <- step1 defState s1
441 case r of
442 Yield x s1' -> cmp_loop1 SPEC x s1' s2
443 Skip s1' -> cmp_loop0 SPEC s1' s2
444 Stop -> cmp_null s2
446 cmp_loop1 !_ x s1 s2 = do
447 r <- step2 defState s2
448 case r of
449 Yield y s2' -> case x `cmp` y of
450 EQ -> cmp_loop0 SPEC s1 s2'
451 c -> return c
452 Skip s2' -> cmp_loop1 SPEC x s1 s2'
453 Stop -> return GT
455 cmp_null s2 = do
456 r <- step2 defState s2
457 case r of
458 Yield _ _ -> return LT
459 Skip s2' -> cmp_null s2'
460 Stop -> return EQ
462{-# INLINE_NORMAL take #-}
463take :: Monad m => Int -> Stream m a -> Stream m a
464take n (Stream step state) = n `seq` Stream step' (state, 0)
465 where
466 {-# INLINE_LATE step' #-}
467 step' gst (st, i) | i < n = do
468 r <- step gst st
469 return $ case r of
470 Yield x s -> Yield x (s, i + 1)
471 Skip s -> Skip (s, i)
472 Stop -> Stop
473 step' _ (_, _) = return Stop
475------------------------------------------------------------------------------
476-- Grouping/Splitting
477------------------------------------------------------------------------------
479-- s = stream state, fs = fold state
480data GroupState s fs
481 = GroupStart s
482 | GroupBuffer s fs Int
483 | GroupYield fs (GroupState s fs)
484 | GroupFinish
486{-# INLINE_NORMAL groupsOf #-}
487groupsOf
488 :: Monad m
489 => Int
490 -> Fold m a b
491 -> Stream m a
492 -> Stream m b
493groupsOf n (Fold fstep initial extract) (Stream step state) =
494 n `seq` Stream step' (GroupStart state)
496 where
498 {-# INLINE_LATE step' #-}
499 step' _ (GroupStart st) = do
500 -- XXX shall we use the Natural type instead? Need to check performance
501 -- implications.
502 when (n <= 0) $
503 -- XXX we can pass the module string from the higher level API
504 error $ "Streamly.Streams.StreamD.Type.groupsOf: the size of "
505 ++ "groups [" ++ show n ++ "] must be a natural number"
506 -- fs = fold state
507 fs <- initial
508 return $ Skip (GroupBuffer st fs 0)
510 step' gst (GroupBuffer st fs i) = do
511 r <- step (adaptState gst) st
512 case r of
513 Yield x s -> do
514 fs' <- fstep fs x
515 let i' = i + 1
516 return $
517 if i' >= n
518 then Skip (GroupYield fs' (GroupStart s))
519 else Skip (GroupBuffer s fs' i')
520 Skip s -> return $ Skip (GroupBuffer s fs i)
521 Stop -> return $ Skip (GroupYield fs GroupFinish)
523 step' _ (GroupYield fs next) = do
524 r <- extract fs
525 return $ Yield r next
527 step' _ GroupFinish = return Stop