git.iko.soy/streamly
on c8fc696 help sign in
streamly/src/Streamly/Streams/StreamD/Type.hs
181 lines 6.3 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 , map
36 , mapM
37 , foldrT
38 , foldrM
39 , foldrMx
40 , foldr
41 , toList
42 )
43where
45import Control.Applicative (liftA2)
46import Control.Monad.Trans (lift, MonadTrans)
47import GHC.Types (SPEC(..))
48import Prelude hiding (map, mapM, foldr)
49import Streamly.SVar (State(..), adaptState, defState)
51import qualified Streamly.Streams.StreamK as K
53------------------------------------------------------------------------------
54-- The direct style stream type
55------------------------------------------------------------------------------
57-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
58-- the next state of the stream. 'Stop' indicates there are no more values in
59-- the stream.
60data Step s a = Yield a s | Skip s | Stop
62instance Functor (Step s) where
63 {-# INLINE fmap #-}
64 fmap f (Yield x s) = Yield (f x) s
65 fmap _ (Skip s) = Skip s
66 fmap _ Stop = Stop
68-- gst = global state
69-- | A stream consists of a step function that generates the next step given a
70-- current state, and the current state.
71data Stream m a =
72 forall s. UnStream (State K.Stream m a -> s -> m (Step s a)) s
74unShare :: Stream m a -> Stream m a
75unShare (UnStream step state) = UnStream step' state
76 where step' gst = step (adaptState gst)
78pattern Stream :: (State K.Stream m a -> s -> m (Step s a)) -> s -> Stream m a
79pattern Stream step state <- (unShare -> UnStream step state)
80 where Stream = UnStream
82#if __GLASGOW_HASKELL__ >= 802
83{-# COMPLETE Stream #-}
84#endif
86------------------------------------------------------------------------------
87-- Instances
88------------------------------------------------------------------------------
90-- | Map a monadic function over a 'Stream'
91{-# INLINE_NORMAL mapM #-}
92mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
93mapM f (Stream step state) = Stream step' state
94 where
95 {-# INLINE_LATE step' #-}
96 step' gst st = do
97 r <- step (adaptState gst) st
98 case r of
99 Yield x s -> f x >>= \a -> return $ Yield a s
100 Skip s -> return $ Skip s
101 Stop -> return Stop
103{-# INLINE map #-}
104map :: Monad m => (a -> b) -> Stream m a -> Stream m b
105map f = mapM (return . f)
107instance Monad m => Functor (Stream m) where
108 {-# INLINE fmap #-}
109 fmap = map
111-- XXX Use of SPEC constructor in folds causes 2x performance degradation in
112-- one shot operations, but helps immensely in operations composed of multiple
113-- combinators or the same combinator many times. There seems to be an
114-- opportunity to optimize here, can we get both, better perf for single ops
115-- as well as composed ops? Without SPEC, all single operation benchmarks
116-- become 2x faster.
118-- The way we want a left fold to be strict, dually we want the right fold to
119-- be lazy. The correct signature of the fold function to keep it lazy must be
120-- (a -> m b -> m b) instead of (a -> b -> m b). We were using the latter
121-- earlier, which is incorrect. In the latter signature we have to feed the
122-- value to the fold function after evaluating the monadic action, depending on
123-- the bind behavior of the monad, the action may get evaluated immediately
124-- introducing unnecessary strictness to the fold. If the implementation is
125-- lazy the following example, must work:
126--
127-- S.foldrM (\x t -> if x then return t else return False) (return True)
128-- (S.fromList [False,undefined] :: SerialT IO Bool)
129--
130{-# INLINE_NORMAL foldrM #-}
131foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
132foldrM f z (Stream step state) = go SPEC state
133 where
134 go !_ st = do
135 r <- step defState st
136 case r of
137 Yield x s -> f x (go SPEC s)
138 Skip s -> go SPEC s
139 Stop -> z
141{-# INLINE_NORMAL foldrMx #-}
142foldrMx :: Monad m
143 => (a -> m x -> m x) -> m x -> (m x -> m b) -> Stream m a -> m b
144foldrMx fstep final convert (Stream step state) = convert $ go SPEC state
145 where
146 go !_ st = do
147 r <- step defState st
148 case r of
149 Yield x s -> fstep x (go SPEC s)
150 Skip s -> go SPEC s
151 Stop -> final
153-- Note that foldr works on pure values, therefore it becomes necessarily
154-- strict when the monad m is strict. In that case it cannot terminate early,
155-- it would evaluate all of its input. Though, this should work fine with lazy
156-- monads. For example, if "any" is implemented using "foldr" instead of
157-- "foldrM" it performs the same with Identity monad but performs 1000x slower
158-- with IO monad.
159--
160{-# INLINE_NORMAL foldr #-}
161foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
162foldr f z = foldrM (\a b -> liftA2 f (return a) b) (return z)
164-- Right fold to some transformer (T) monad. This can be useful to implement
165-- stateless combinators like map, filtering, insertions, takeWhile, dropWhile.
166--
167{-# INLINE_NORMAL foldrT #-}
168foldrT :: (Monad m, Monad (t m), MonadTrans t)
169 => (a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
170foldrT f final (Stream step state) = go SPEC state
171 where
172 go !_ st = do
173 r <- lift $ step defState st
174 case r of
175 Yield x s -> f x (go SPEC s)
176 Skip s -> go SPEC s
177 Stop -> final
179{-# INLINE toList #-}
180toList :: Monad m => Stream m a -> m [a]
181toList = foldr (:) []