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