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