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