git.iko.soy/streamly
on fd7bd71 help sign in
streamly/src/Streamly/Streams/StreamD/Type.hs
419 lines 13 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--
18-- License : BSD3
19-- Maintainer : harendra.kumar@gmail.com
20-- Stability : experimental
21-- Portability : GHC
23module Streamly.Streams.StreamD.Type
24 (
25 -- * The stream type
26 Step (..)
27 -- XXX UnStream is exported to avoid a performance issue in concatMap if we
28 -- use the pattern synonym "Stream".
29#if __GLASGOW_HASKELL__ >= 800
30 , Stream (Stream, UnStream)
31#else
32 , Stream (UnStream)
33 , pattern Stream
34#endif
35 , fromStreamK
36 , toStreamK
37 , map
38 , mapM
40 , foldrT
41 , foldrM
42 , foldrMx
43 , foldr
45 , foldl'
46 , foldlM'
47 , foldlx'
48 , foldlMx'
50 , toList
51 , fromList
53 , eqBy
54 , cmpBy
55 , take
56 , groupsOf
57 )
58where
60import Control.Applicative (liftA2)
61import Control.Monad.Trans (lift, MonadTrans)
62import GHC.Types (SPEC(..))
63import Prelude hiding (map, mapM, foldr, take)
65import Streamly.SVar (State(..), adaptState, defState)
66import Streamly.Fold.Types (Fold(..))
68import qualified Streamly.Streams.StreamK as K
70------------------------------------------------------------------------------
71-- The direct style stream type
72------------------------------------------------------------------------------
74-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
75-- the next state of the stream. 'Stop' indicates there are no more values in
76-- the stream.
77data Step s a = Yield a s | Skip s | Stop
79instance Functor (Step s) where
80 {-# INLINE fmap #-}
81 fmap f (Yield x s) = Yield (f x) s
82 fmap _ (Skip s) = Skip s
83 fmap _ Stop = Stop
85-- gst = global state
86-- | A stream consists of a step function that generates the next step given a
87-- current state, and the current state.
88data Stream m a =
89 forall s. UnStream (State K.Stream m a -> s -> m (Step s a)) s
91unShare :: Stream m a -> Stream m a
92unShare (UnStream step state) = UnStream step' state
93 where step' gst = step (adaptState gst)
95pattern Stream :: (State K.Stream m a -> s -> m (Step s a)) -> s -> Stream m a
96pattern Stream step state <- (unShare -> UnStream step state)
97 where Stream = UnStream
99#if __GLASGOW_HASKELL__ >= 802
100{-# COMPLETE Stream #-}
101#endif
103{-# INLINE_LATE fromStreamK #-}
104fromStreamK :: Monad m => K.Stream m a -> Stream m a
105fromStreamK = Stream step
106 where
107 step gst m1 =
108 let stop = return Stop
109 single a = return $ Yield a K.nil
110 yieldk a r = return $ Yield a r
111 in K.foldStreamShared gst yieldk single stop m1
113-- Convert a direct stream to and from CPS encoded stream
114{-# INLINE_LATE toStreamK #-}
115toStreamK :: Monad m => Stream m a -> K.Stream m a
116toStreamK (Stream step state) = go state
117 where
118 go st = K.mkStream $ \gst yld sng stp -> do
119 r <- step gst st
120 case r of
121 Yield x s -> yld x (go s)
122 Skip s -> K.foldStreamShared gst yld sng stp $ go s
123 Stop -> stp
125#ifndef DISABLE_FUSION
126{-# RULES "fromStreamK/toStreamK fusion"
127 forall s. toStreamK (fromStreamK s) = s #-}
128{-# RULES "toStreamK/fromStreamK fusion"
129 forall s. fromStreamK (toStreamK s) = s #-}
130#endif
132------------------------------------------------------------------------------
133-- Instances
134------------------------------------------------------------------------------
136-- | Map a monadic function over a 'Stream'
137{-# INLINE_NORMAL mapM #-}
138mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
139mapM f (Stream step state) = Stream step' state
140 where
141 {-# INLINE_LATE step' #-}
142 step' gst st = do
143 r <- step (adaptState gst) st
144 case r of
145 Yield x s -> f x >>= \a -> return $ Yield a s
146 Skip s -> return $ Skip s
147 Stop -> return Stop
149{-# INLINE map #-}
150map :: Monad m => (a -> b) -> Stream m a -> Stream m b
151map f = mapM (return . f)
153instance Monad m => Functor (Stream m) where
154 {-# INLINE fmap #-}
155 fmap = map
157-- XXX Use of SPEC constructor in folds causes 2x performance degradation in
158-- one shot operations, but helps immensely in operations composed of multiple
159-- combinators or the same combinator many times. There seems to be an
160-- opportunity to optimize here, can we get both, better perf for single ops
161-- as well as composed ops? Without SPEC, all single operation benchmarks
162-- become 2x faster.
164-- The way we want a left fold to be strict, dually we want the right fold to
165-- be lazy. The correct signature of the fold function to keep it lazy must be
166-- (a -> m b -> m b) instead of (a -> b -> m b). We were using the latter
167-- earlier, which is incorrect. In the latter signature we have to feed the
168-- value to the fold function after evaluating the monadic action, depending on
169-- the bind behavior of the monad, the action may get evaluated immediately
170-- introducing unnecessary strictness to the fold. If the implementation is
171-- lazy the following example, must work:
172--
173-- S.foldrM (\x t -> if x then return t else return False) (return True)
174-- (S.fromList [False,undefined] :: SerialT IO Bool)
175--
176{-# INLINE_NORMAL foldrM #-}
177foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
178foldrM f z (Stream step state) = go SPEC state
179 where
180 go !_ st = do
181 r <- step defState st
182 case r of
183 Yield x s -> f x (go SPEC s)
184 Skip s -> go SPEC s
185 Stop -> z
187{-# INLINE_NORMAL foldrMx #-}
188foldrMx :: Monad m
189 => (a -> m x -> m x) -> m x -> (m x -> m b) -> Stream m a -> m b
190foldrMx fstep final convert (Stream step state) = convert $ go SPEC state
191 where
192 go !_ st = do
193 r <- step defState st
194 case r of
195 Yield x s -> fstep x (go SPEC s)
196 Skip s -> go SPEC s
197 Stop -> final
199-- Note that foldr works on pure values, therefore it becomes necessarily
200-- strict when the monad m is strict. In that case it cannot terminate early,
201-- it would evaluate all of its input. Though, this should work fine with lazy
202-- monads. For example, if "any" is implemented using "foldr" instead of
203-- "foldrM" it performs the same with Identity monad but performs 1000x slower
204-- with IO monad.
205--
206{-# INLINE_NORMAL foldr #-}
207foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
208foldr f z = foldrM (\a b -> liftA2 f (return a) b) (return z)
210-- Right fold to some transformer (T) monad. This can be useful to implement
211-- stateless combinators like map, filtering, insertions, takeWhile, dropWhile.
212--
213{-# INLINE_NORMAL foldrT #-}
214foldrT :: (Monad m, Monad (t m), MonadTrans t)
215 => (a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
216foldrT f final (Stream step state) = go SPEC state
217 where
218 go !_ st = do
219 r <- lift $ step defState st
220 case r of
221 Yield x s -> f x (go SPEC s)
222 Skip s -> go SPEC s
223 Stop -> final
225{-# INLINE toList #-}
226toList :: Monad m => Stream m a -> m [a]
227toList = foldr (:) []
229-- XXX run begin action only if the stream is not empty.
230{-# INLINE_NORMAL foldlMx' #-}
231foldlMx' :: Monad m => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
232foldlMx' fstep begin done (Stream step state) =
233 begin >>= \x -> go SPEC x state
234 where
235 -- XXX !acc?
236 go !_ acc st = acc `seq` do
237 r <- step defState st
238 case r of
239 Yield x s -> do
240 acc' <- fstep acc x
241 go SPEC acc' s
242 Skip s -> go SPEC acc s
243 Stop -> done acc
245{-# INLINE foldlx' #-}
246foldlx' :: Monad m => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
247foldlx' fstep begin done m =
248 foldlMx' (\b a -> return (fstep b a)) (return begin) (return . done) m
250-- XXX implement in terms of foldlMx'?
251{-# INLINE_NORMAL foldlM' #-}
252foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
253foldlM' fstep begin (Stream step state) = go SPEC begin state
254 where
255 go !_ acc st = acc `seq` do
256 r <- step defState st
257 case r of
258 Yield x s -> do
259 acc' <- fstep acc x
260 go SPEC acc' s
261 Skip s -> go SPEC acc s
262 Stop -> return acc
264{-# INLINE foldl' #-}
265foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
266foldl' fstep = foldlM' (\b a -> return (fstep b a))
268-- | Convert a list of pure values to a 'Stream'
269{-# INLINE_LATE fromList #-}
270fromList :: Monad m => [a] -> Stream m a
271fromList = Stream step
272 where
273 {-# INLINE_LATE step #-}
274 step _ (x:xs) = return $ Yield x xs
275 step _ [] = return Stop
277------------------------------------------------------------------------------
278-- Comparisons
279------------------------------------------------------------------------------
281{-# INLINE_NORMAL eqBy #-}
282eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
283eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
284 where
285 eq_loop0 !_ s1 s2 = do
286 r <- step1 defState s1
287 case r of
288 Yield x s1' -> eq_loop1 SPEC x s1' s2
289 Skip s1' -> eq_loop0 SPEC s1' s2
290 Stop -> eq_null s2
292 eq_loop1 !_ x s1 s2 = do
293 r <- step2 defState s2
294 case r of
295 Yield y s2'
296 | eq x y -> eq_loop0 SPEC s1 s2'
297 | otherwise -> return False
298 Skip s2' -> eq_loop1 SPEC x s1 s2'
299 Stop -> return False
301 eq_null s2 = do
302 r <- step2 defState s2
303 case r of
304 Yield _ _ -> return False
305 Skip s2' -> eq_null s2'
306 Stop -> return True
308-- | Compare two streams lexicographically
309{-# INLINE_NORMAL cmpBy #-}
310cmpBy
311 :: Monad m
312 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
313cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
314 where
315 cmp_loop0 !_ s1 s2 = do
316 r <- step1 defState s1
317 case r of
318 Yield x s1' -> cmp_loop1 SPEC x s1' s2
319 Skip s1' -> cmp_loop0 SPEC s1' s2
320 Stop -> cmp_null s2
322 cmp_loop1 !_ x s1 s2 = do
323 r <- step2 defState s2
324 case r of
325 Yield y s2' -> case x `cmp` y of
326 EQ -> cmp_loop0 SPEC s1 s2'
327 c -> return c
328 Skip s2' -> cmp_loop1 SPEC x s1 s2'
329 Stop -> return GT
331 cmp_null s2 = do
332 r <- step2 defState s2
333 case r of
334 Yield _ _ -> return LT
335 Skip s2' -> cmp_null s2'
336 Stop -> return EQ
338{-# INLINE_NORMAL take #-}
339take :: Monad m => Int -> Stream m a -> Stream m a
340take n (Stream step state) = n `seq` Stream step' (state, 0)
341 where
342 {-# INLINE_LATE step' #-}
343 step' gst (st, i) | i < n = do
344 r <- step gst st
345 return $ case r of
346 Yield x s -> Yield x (s, i + 1)
347 Skip s -> Skip (s, i)
348 Stop -> Stop
349 step' _ (_, _) = return Stop
351------------------------------------------------------------------------------
352-- Grouping/Splitting
353------------------------------------------------------------------------------
355{-# INLINE_LATE foldOneGroup #-}
356foldOneGroup
357 :: Monad m
358 => Int
359 -> Fold m a b
360 -> a
361 -> State K.Stream m a
362 -> (State K.Stream m a -> s -> m (Step s a))
363 -> s
364 -> m (b, Maybe s)
365foldOneGroup n (Fold fstep begin done) x gst step state = do
366 acc0 <- begin
367 acc <- fstep acc0 x
368 go SPEC state acc 1
370 where
372 -- XXX is it strict enough?
373 go !_ st !acc i | i < n = do
374 r <- step gst st
375 case r of
376 Yield y s -> do
377 acc1 <- fstep acc y
378 go SPEC s acc1 (i + 1)
379 Skip s -> go SPEC s acc i
380 Stop -> do
381 res <- done acc
382 return (res, Nothing)
383 go !_ st acc _ = do
384 r <- done acc
385 return (r, Just st)
387-- groupsOf takes 11 sec to write the file from a stream, whereas just
388-- reading in the same file as a stream and folding each element of the stream
389-- using (+) takes just 1.5 sec, so we are still pretty slow (7x slow), there
390-- is scope to make it faster. There is a possibility of better fusion here.
391--
392{-# INLINE_NORMAL groupsOf #-}
393groupsOf
394 :: Monad m
395 => Int
396 -> Fold m a b
397 -> Stream m a
398 -> Stream m b
399groupsOf n f (Stream step state) =
400 n `seq` Stream stepOuter (Just state)
402 where
404 {-# INLINE_LATE stepOuter #-}
405 stepOuter gst (Just st) = do
406 -- We retrieve the first element of the stream before we start to fold
407 -- a chunk so that we do not return an empty chunk in case the stream
408 -- is empty.
409 res <- step (adaptState gst) st
410 case res of
411 Yield x s -> do
412 -- XXX how to make sure that stepInner and f get fused
413 -- This problem seems to be similar to the concatMap problem
414 (r, s1) <- foldOneGroup n f x (adaptState gst) step s
415 return $ Yield r s1
416 Skip s -> return $ Skip $ Just s
417 Stop -> return Stop
419 stepOuter _ Nothing = return Stop