git.iko.soy/streamly
on c8a98dd help sign in
streamly/src/Streamly/Streams/StreamD/Type.hs
576 lines 19 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
45 , flatten
47 , foldrT
48 , foldrM
49 , foldrMx
50 , foldr
51 , foldrS
53 , foldl'
54 , foldlM'
55 , foldlx'
56 , foldlMx'
58 , toList
59 , fromList
61 , eqBy
62 , cmpBy
63 , take
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.SVar (State(..), adaptState, defState)
75import Streamly.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-- | @flatten step inject stream@ uses the unfold @step@ function to generate
226-- streams from the input stream and flatten them into a single output stream.
227-- The @inject@ function is used to wrap the input element into the state of
228-- unfold step.
229--
230-- This is like 'concatMap' but uses an unfold with an explicit state to
231-- generate the stream instead of a 'Stream' type generator. This allows better
232-- optimization via fusion. This can be many times more efficient than
233-- 'concatMap'.
235{-# INLINE_NORMAL flatten #-}
236flatten :: Monad m
237 => (s -> m (Step s b)) -> (a -> m s) -> Stream m a -> Stream m b
238flatten istep inject (Stream ostep u) = Stream step (Left u)
239 where
240 {-# INLINE_LATE step #-}
241 step gst (Left t) = do
242 r <- ostep (adaptState gst) t
243 case r of
244 Yield a t' -> do
245 s <- inject a
246 s `seq` return (Skip (Right (s, t')))
247 Skip t' -> return $ Skip (Left t')
248 Stop -> return $ Stop
250 step _ (Right (s, t)) = do
251 r <- istep s
252 return $ case r of
253 Yield x s' -> Yield x (Right (s', t))
254 Skip s' -> Skip (Right (s', t))
255 Stop -> Skip (Left t)
257-- | Create a singleton 'Stream' from a pure value.
258{-# INLINE_NORMAL yield #-}
259yield :: Monad m => a -> Stream m a
260yield x = Stream (\_ s -> return $ step undefined s) True
261 where
262 {-# INLINE_LATE step #-}
263 step _ True = Yield x False
264 step _ False = Stop
266instance Monad m => Applicative (Stream m) where
267 {-# INLINE pure #-}
268 pure = yield
269 {-# INLINE (<*>) #-}
270 (<*>) = ap
272-- NOTE: even though concatMap for StreamD is 4x faster compared to StreamK,
273-- the monad instance does not seem to be significantly faster.
274instance Monad m => Monad (Stream m) where
275 {-# INLINE return #-}
276 return = pure
277 {-# INLINE (>>=) #-}
278 (>>=) = flip concatMap
280-- XXX Use of SPEC constructor in folds causes 2x performance degradation in
281-- one shot operations, but helps immensely in operations composed of multiple
282-- combinators or the same combinator many times. There seems to be an
283-- opportunity to optimize here, can we get both, better perf for single ops
284-- as well as composed ops? Without SPEC, all single operation benchmarks
285-- become 2x faster.
287-- The way we want a left fold to be strict, dually we want the right fold to
288-- be lazy. The correct signature of the fold function to keep it lazy must be
289-- (a -> m b -> m b) instead of (a -> b -> m b). We were using the latter
290-- earlier, which is incorrect. In the latter signature we have to feed the
291-- value to the fold function after evaluating the monadic action, depending on
292-- the bind behavior of the monad, the action may get evaluated immediately
293-- introducing unnecessary strictness to the fold. If the implementation is
294-- lazy the following example, must work:
295--
296-- S.foldrM (\x t -> if x then return t else return False) (return True)
297-- (S.fromList [False,undefined] :: SerialT IO Bool)
298--
299{-# INLINE_NORMAL foldrM #-}
300foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
301foldrM f z (Stream step state) = go SPEC state
302 where
303 go !_ st = do
304 r <- step defState st
305 case r of
306 Yield x s -> f x (go SPEC s)
307 Skip s -> go SPEC s
308 Stop -> z
310{-# INLINE_NORMAL foldrMx #-}
311foldrMx :: Monad m
312 => (a -> m x -> m x) -> m x -> (m x -> m b) -> Stream m a -> m b
313foldrMx fstep final convert (Stream step state) = convert $ go SPEC state
314 where
315 go !_ st = do
316 r <- step defState st
317 case r of
318 Yield x s -> fstep x (go SPEC s)
319 Skip s -> go SPEC s
320 Stop -> final
322-- Note that foldr works on pure values, therefore it becomes necessarily
323-- strict when the monad m is strict. In that case it cannot terminate early,
324-- it would evaluate all of its input. Though, this should work fine with lazy
325-- monads. For example, if "any" is implemented using "foldr" instead of
326-- "foldrM" it performs the same with Identity monad but performs 1000x slower
327-- with IO monad.
328--
329{-# INLINE_NORMAL foldr #-}
330foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
331foldr f z = foldrM (\a b -> liftA2 f (return a) b) (return z)
333-- | Create a singleton 'Stream' from a monadic action.
334{-# INLINE_NORMAL yieldM #-}
335yieldM :: Monad m => m a -> Stream m a
336yieldM m = Stream step True
337 where
338 {-# INLINE_LATE step #-}
339 step _ True = m >>= \x -> return $ Yield x False
340 step _ False = return Stop
342-- this performs horribly, should not be used
343{-# INLINE_NORMAL foldrS #-}
344foldrS
345 :: Monad m
346 => (a -> Stream m b -> Stream m b)
347 -> Stream m b
348 -> Stream m a
349 -> Stream m b
350foldrS f final (Stream step state) = go SPEC state
351 where
352 go !_ st = do
353 -- defState??
354 r <- yieldM $ step defState st
355 case r of
356 Yield x s -> f x (go SPEC s)
357 Skip s -> go SPEC s
358 Stop -> final
360-- Right fold to some transformer (T) monad. This can be useful to implement
361-- stateless combinators like map, filtering, insertions, takeWhile, dropWhile.
362--
363{-# INLINE_NORMAL foldrT #-}
364foldrT :: (Monad m, Monad (t m), MonadTrans t)
365 => (a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
366foldrT f final (Stream step state) = go SPEC state
367 where
368 go !_ st = do
369 r <- lift $ step defState st
370 case r of
371 Yield x s -> f x (go SPEC s)
372 Skip s -> go SPEC s
373 Stop -> final
375{-# INLINE toList #-}
376toList :: Monad m => Stream m a -> m [a]
377toList = foldr (:) []
379-- XXX run begin action only if the stream is not empty.
380{-# INLINE_NORMAL foldlMx' #-}
381foldlMx' :: Monad m => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
382foldlMx' fstep begin done (Stream step state) =
383 begin >>= \x -> go SPEC x state
384 where
385 -- XXX !acc?
386 go !_ acc st = acc `seq` do
387 r <- step defState st
388 case r of
389 Yield x s -> do
390 acc' <- fstep acc x
391 go SPEC acc' s
392 Skip s -> go SPEC acc s
393 Stop -> done acc
395{-# INLINE foldlx' #-}
396foldlx' :: Monad m => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
397foldlx' fstep begin done m =
398 foldlMx' (\b a -> return (fstep b a)) (return begin) (return . done) m
400-- XXX implement in terms of foldlMx'?
401{-# INLINE_NORMAL foldlM' #-}
402foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
403foldlM' fstep begin (Stream step state) = go SPEC begin state
404 where
405 go !_ acc st = acc `seq` do
406 r <- step defState st
407 case r of
408 Yield x s -> do
409 acc' <- fstep acc x
410 go SPEC acc' s
411 Skip s -> go SPEC acc s
412 Stop -> return acc
414{-# INLINE foldl' #-}
415foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
416foldl' fstep = foldlM' (\b a -> return (fstep b a))
418-- | Convert a list of pure values to a 'Stream'
419{-# INLINE_LATE fromList #-}
420fromList :: Monad m => [a] -> Stream m a
421fromList = Stream step
422 where
423 {-# INLINE_LATE step #-}
424 step _ (x:xs) = return $ Yield x xs
425 step _ [] = return Stop
427------------------------------------------------------------------------------
428-- Comparisons
429------------------------------------------------------------------------------
431{-# INLINE_NORMAL eqBy #-}
432eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
433eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
434 where
435 eq_loop0 !_ s1 s2 = do
436 r <- step1 defState s1
437 case r of
438 Yield x s1' -> eq_loop1 SPEC x s1' s2
439 Skip s1' -> eq_loop0 SPEC s1' s2
440 Stop -> eq_null s2
442 eq_loop1 !_ x s1 s2 = do
443 r <- step2 defState s2
444 case r of
445 Yield y s2'
446 | eq x y -> eq_loop0 SPEC s1 s2'
447 | otherwise -> return False
448 Skip s2' -> eq_loop1 SPEC x s1 s2'
449 Stop -> return False
451 eq_null s2 = do
452 r <- step2 defState s2
453 case r of
454 Yield _ _ -> return False
455 Skip s2' -> eq_null s2'
456 Stop -> return True
458-- | Compare two streams lexicographically
459{-# INLINE_NORMAL cmpBy #-}
460cmpBy
461 :: Monad m
462 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
463cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
464 where
465 cmp_loop0 !_ s1 s2 = do
466 r <- step1 defState s1
467 case r of
468 Yield x s1' -> cmp_loop1 SPEC x s1' s2
469 Skip s1' -> cmp_loop0 SPEC s1' s2
470 Stop -> cmp_null s2
472 cmp_loop1 !_ x s1 s2 = do
473 r <- step2 defState s2
474 case r of
475 Yield y s2' -> case x `cmp` y of
476 EQ -> cmp_loop0 SPEC s1 s2'
477 c -> return c
478 Skip s2' -> cmp_loop1 SPEC x s1 s2'
479 Stop -> return GT
481 cmp_null s2 = do
482 r <- step2 defState s2
483 case r of
484 Yield _ _ -> return LT
485 Skip s2' -> cmp_null s2'
486 Stop -> return EQ
488{-# INLINE_NORMAL take #-}
489take :: Monad m => Int -> Stream m a -> Stream m a
490take n (Stream step state) = n `seq` Stream step' (state, 0)
491 where
492 {-# INLINE_LATE step' #-}
493 step' gst (st, i) | i < n = do
494 r <- step gst st
495 return $ case r of
496 Yield x s -> Yield x (s, i + 1)
497 Skip s -> Skip (s, i)
498 Stop -> Stop
499 step' _ (_, _) = return Stop
501------------------------------------------------------------------------------
502-- Grouping/Splitting
503------------------------------------------------------------------------------
505{-# INLINE_LATE foldOneGroup #-}
506foldOneGroup
507 :: Monad m
508 => Int
509 -> Fold m a b
510 -> a
511 -> State K.Stream m a
512 -> (State K.Stream m a -> s -> m (Step s a))
513 -> s
514 -> m (b, Maybe s)
515foldOneGroup n (Fold fstep begin done) x gst step state = do
516 acc0 <- begin
517 acc <- fstep acc0 x
518 go SPEC state acc 1
520 where
522 -- XXX is it strict enough?
523 go !_ st !acc i | i < n = do
524 r <- step gst st
525 case r of
526 Yield y s -> do
527 acc1 <- fstep acc y
528 go SPEC s acc1 (i + 1)
529 Skip s -> go SPEC s acc i
530 Stop -> do
531 res <- done acc
532 return (res, Nothing)
533 go !_ st acc _ = do
534 r <- done acc
535 return (r, Just st)
537-- groupsOf takes 11 sec to write the file from a stream, whereas just
538-- reading in the same file as a stream and folding each element of the stream
539-- using (+) takes just 1.5 sec, so we are still pretty slow (7x slow), there
540-- is scope to make it faster. There is a possibility of better fusion here.
541--
542{-# INLINE_NORMAL groupsOf #-}
543groupsOf
544 :: Monad m
545 => Int
546 -> Fold m a b
547 -> Stream m a
548 -> Stream m b
549groupsOf n f (Stream step state) =
550 n `seq` Stream stepOuter (Just state)
552 where
554 {-# INLINE_LATE stepOuter #-}
555 stepOuter gst (Just st) = do
556 -- XXX shall we use the Natural type instead? Need to check performance
557 -- implications.
558 when (n <= 0) $
559 -- XXX we can pass the module string from the higher level API
560 error $ "Streamly.Streams.StreamD.Type.groupsOf: the size of "
561 ++ "groups [" ++ show n ++ "] must be a natural number"
563 -- We retrieve the first element of the stream before we start to fold
564 -- a chunk so that we do not return an empty chunk in case the stream
565 -- is empty.
566 res <- step (adaptState gst) st
567 case res of
568 Yield x s -> do
569 -- XXX how to make sure that stepInner and f get fused
570 -- This problem seems to be similar to the concatMap problem
571 (r, s1) <- foldOneGroup n f x (adaptState gst) step s
572 return $ Yield r s1
573 Skip s -> return $ Skip $ Just s
574 Stop -> return Stop
576 stepOuter _ Nothing = return Stop