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