git.iko.soy/streamly
on bc5617f help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Generate.hs
406 lines 14 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.StreamD.Generate
3-- Copyright : (c) 2020 Composewell Technologies and Contributors
4-- (c) Roman Leshchinskiy 2008-2010
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- Prefer unfolds ("Streamly.Internal.Data.Unfold") over the combinators in
11-- this module. They are more powerful and efficient as they can be transformed
12-- and composed on the input side efficiently and they can fuse in nested
13-- operations (e.g. unfoldMany). All the combinators in this module can be
14-- expressed using unfolds with the same efficiency.
15--
16-- Operations in this module that are not in "Streamly.Internal.Data.Unfold":
17-- generate, times, fromPrimIORef.
18--
19-- We should plan to replace this module with "Streamly.Internal.Data.Unfold"
20-- in future.
22-- A few combinators in this module have been adapted from the vector package
23-- (c) Roman Leshchinskiy. See the notes in specific combinators.
24--
25module Streamly.Internal.Data.Stream.StreamD.Generate
26 (
27 -- * Primitives
28 nil
29 , nilM
30 , cons
31 , consM
33 -- * From 'Unfold'
34 , unfold
36 -- * Unfolding
37 , unfoldr
38 , unfoldrM
40 -- * From Values
41 , fromPure
42 , fromEffect
43 , repeat
44 , repeatM
45 , replicate
46 , replicateM
48 -- * Enumeration
49 , enumerateFromStepIntegral
50 , enumerateFromIntegral
51 , enumerateFromThenIntegral
52 , enumerateFromToIntegral
53 , enumerateFromThenToIntegral
55 , enumerateFromStepNum
56 , enumerateFromNum
57 , enumerateFromThenNum
58 , enumerateFromToFractional
59 , enumerateFromThenToFractional
61 -- * Time Enumeration
62 , times
64 -- * From Generators
65 -- | Generate a monadic stream from a seed.
66 , fromIndices
67 , fromIndicesM
68 , generate
69 , generateM
71 -- * Iteration
72 , iterate
73 , iterateM
75 -- * From Containers
76 -- | Transform an input structure into a stream.
78 -- Note: Direct style stream does not support @fromFoldable@.
79 , fromList
80 , fromListM
82 -- * Conversions
83 , fromStreamK
84 , toStreamK
85 )
86where
88#include "inline.hs"
90import Control.Monad.IO.Class (MonadIO(..))
91import Streamly.Internal.Control.Concurrent (MonadAsync)
92import Streamly.Internal.Data.Time.Clock
93 (Clock(Monotonic), asyncClock, readClock)
94import Streamly.Internal.Data.Time.Units
95 (toAbsTime, AbsTime, toRelTime64, RelTime64)
97import Prelude hiding (iterate, repeat, replicate, takeWhile)
98import Streamly.Internal.Data.Stream.StreamD.Type
100------------------------------------------------------------------------------
101-- Primitives
102------------------------------------------------------------------------------
104-- | An empty 'Stream'.
105{-# INLINE_NORMAL nil #-}
106nil :: Monad m => Stream m a
107nil = Stream (\_ _ -> return Stop) ()
109-- XXX implement in terms of consM?
110-- cons x = consM (return x)
111--
112-- | Can fuse but has O(n^2) complexity.
113{-# INLINE_NORMAL cons #-}
114cons :: Monad m => a -> Stream m a -> Stream m a
115cons x (Stream step state) = Stream step1 Nothing
116 where
117 {-# INLINE_LATE step1 #-}
118 step1 _ Nothing = return $ Yield x (Just state)
119 step1 gst (Just st) = do
120 r <- step gst st
121 return $
122 case r of
123 Yield a s -> Yield a (Just s)
124 Skip s -> Skip (Just s)
125 Stop -> Stop
127------------------------------------------------------------------------------
128-- Unfolding
129------------------------------------------------------------------------------
131-- Adapted from vector package
132{-# INLINE_NORMAL unfoldrM #-}
133unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
134unfoldrM next = Stream step
135 where
136 {-# INLINE_LATE step #-}
137 step _ st = do
138 r <- next st
139 return $ case r of
140 Just (x, s) -> Yield x s
141 Nothing -> Stop
143{-# INLINE_LATE unfoldr #-}
144unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
145unfoldr f = unfoldrM (return . f)
147------------------------------------------------------------------------------
148-- From values
149------------------------------------------------------------------------------
151{-# INLINE_NORMAL repeatM #-}
152repeatM :: Monad m => m a -> Stream m a
153repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
155{-# INLINE_NORMAL repeat #-}
156repeat :: Monad m => a -> Stream m a
157repeat x = Stream (\_ _ -> return $ Yield x ()) ()
159-- Adapted from the vector package
160{-# INLINE_NORMAL replicateM #-}
161replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
162replicateM n p = Stream step n
163 where
164 {-# INLINE_LATE step #-}
165 step _ (i :: Int)
166 | i <= 0 = return Stop
167 | otherwise = do
168 x <- p
169 return $ Yield x (i - 1)
171{-# INLINE_NORMAL replicate #-}
172replicate :: Monad m => Int -> a -> Stream m a
173replicate n x = replicateM n (return x)
175------------------------------------------------------------------------------
176-- Enumeration of Num
177------------------------------------------------------------------------------
179-- | For floating point numbers if the increment is less than the precision then
180-- it just gets lost. Therefore we cannot always increment it correctly by just
181-- repeated addition.
182-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
183-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
184--
185-- Instead we accumulate the increment counter and compute the increment
186-- every time before adding it to the starting number.
187--
188-- This works for Integrals as well as floating point numbers, but
189-- enumerateFromStepIntegral is faster for integrals.
190{-# INLINE_NORMAL enumerateFromStepNum #-}
191enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
192enumerateFromStepNum from stride = Stream step 0
193 where
194 {-# INLINE_LATE step #-}
195 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
197{-# INLINE_NORMAL enumerateFromNum #-}
198enumerateFromNum :: (Monad m, Num a) => a -> Stream m a
199enumerateFromNum from = enumerateFromStepNum from 1
201{-# INLINE_NORMAL enumerateFromThenNum #-}
202enumerateFromThenNum :: (Monad m, Num a) => a -> a -> Stream m a
203enumerateFromThenNum from next = enumerateFromStepNum from (next - from)
205------------------------------------------------------------------------------
206-- Enumeration of Integrals
207------------------------------------------------------------------------------
209data EnumState a = EnumInit | EnumYield a a a | EnumStop
211{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
212enumerateFromThenToIntegralUp
213 :: (Monad m, Integral a)
214 => a -> a -> a -> Stream m a
215enumerateFromThenToIntegralUp from next to = Stream step EnumInit
216 where
217 {-# INLINE_LATE step #-}
218 step _ EnumInit =
219 return $
220 if to < next
221 then if to < from
222 then Stop
223 else Yield from EnumStop
224 else -- from <= next <= to
225 let stride = next - from
226 in Skip $ EnumYield from stride (to - stride)
228 step _ (EnumYield x stride toMinus) =
229 return $
230 if x > toMinus
231 then Yield x EnumStop
232 else Yield x $ EnumYield (x + stride) stride toMinus
234 step _ EnumStop = return Stop
236{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
237enumerateFromThenToIntegralDn
238 :: (Monad m, Integral a)
239 => a -> a -> a -> Stream m a
240enumerateFromThenToIntegralDn from next to = Stream step EnumInit
241 where
242 {-# INLINE_LATE step #-}
243 step _ EnumInit =
244 return $ if to > next
245 then if to > from
246 then Stop
247 else Yield from EnumStop
248 else -- from >= next >= to
249 let stride = next - from
250 in Skip $ EnumYield from stride (to - stride)
252 step _ (EnumYield x stride toMinus) =
253 return $
254 if x < toMinus
255 then Yield x EnumStop
256 else Yield x $ EnumYield (x + stride) stride toMinus
258 step _ EnumStop = return Stop
260-- XXX This can perhaps be simplified and written in terms of
261-- enumeratFromStepIntegral as we have done in unfolds. But anyway we should be
262-- replacing the stream generation module with unfolds.
263{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
264enumerateFromThenToIntegral
265 :: (Monad m, Integral a)
266 => a -> a -> a -> Stream m a
267enumerateFromThenToIntegral from next to
268 | next >= from = enumerateFromThenToIntegralUp from next to
269 | otherwise = enumerateFromThenToIntegralDn from next to
271{-# INLINE_NORMAL enumerateFromThenIntegral #-}
272enumerateFromThenIntegral
273 :: (Monad m, Integral a, Bounded a)
274 => a -> a -> Stream m a
275enumerateFromThenIntegral from next =
276 if next > from
277 then enumerateFromThenToIntegralUp from next maxBound
278 else enumerateFromThenToIntegralDn from next minBound
280-- | Can be used to enumerate unbounded integrals. This does not check for
281-- overflow or underflow for bounded integrals.
282--
283{-# INLINE_NORMAL enumerateFromStepIntegral #-}
284enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
285enumerateFromStepIntegral from stride =
286 from `seq` stride `seq` Stream step from
287 where
288 {-# INLINE_LATE step #-}
289 step _ !x = return $ Yield x $! (x + stride)
291-- | Enumerate upwards from @from@ to @to@. We are assuming that "to" is
292-- constrained by the type to be within max/min bounds.
293{-# INLINE enumerateFromToIntegral #-}
294enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
295enumerateFromToIntegral from to =
296 takeWhile (<= to) $ enumerateFromStepIntegral from 1
298{-# INLINE enumerateFromIntegral #-}
299enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
300enumerateFromIntegral from = enumerateFromToIntegral from maxBound
302------------------------------------------------------------------------------
303-- Enumeration of Fractionals
304------------------------------------------------------------------------------
306-- | We cannot write a general function for Num. The only way to write code
307-- portable between the two is to use a 'Real' constraint and convert between
308-- Fractional and Integral using fromRational which is horribly slow.
309{-# INLINE_NORMAL enumerateFromToFractional #-}
310enumerateFromToFractional
311 :: (Monad m, Fractional a, Ord a)
312 => a -> a -> Stream m a
313enumerateFromToFractional from to =
314 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
316{-# INLINE_NORMAL enumerateFromThenToFractional #-}
317enumerateFromThenToFractional
318 :: (Monad m, Fractional a, Ord a)
319 => a -> a -> a -> Stream m a
320enumerateFromThenToFractional from next to =
321 takeWhile predicate $ enumerateFromThenNum from next
322 where
323 mid = (next - from) / 2
324 predicate | next >= from = (<= to + mid)
325 | otherwise = (>= to + mid)
327------------------------------------------------------------------------------
328-- Time Enumeration
329------------------------------------------------------------------------------
331{-# INLINE_NORMAL times #-}
332times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
333times g = Stream step Nothing
335 where
337 {-# INLINE_LATE step #-}
338 step _ Nothing = do
339 clock <- liftIO $ asyncClock Monotonic g
340 a <- liftIO $ readClock clock
341 return $ Skip $ Just (clock, a)
343 step _ s@(Just (clock, t0)) = do
344 a <- liftIO $ readClock clock
345 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
346 -- efficiency. or maybe we can use a representation using Double for
347 -- floating precision time
348 return $ Yield (toAbsTime t0, toRelTime64 (a - t0)) s
350-------------------------------------------------------------------------------
351-- From Generators
352-------------------------------------------------------------------------------
354{-# INLINE_NORMAL fromIndicesM #-}
355fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
356fromIndicesM gen = Stream step 0
357 where
358 {-# INLINE_LATE step #-}
359 step _ i = do
360 x <- gen i
361 return $ Yield x (i + 1)
363{-# INLINE fromIndices #-}
364fromIndices :: Monad m => (Int -> a) -> Stream m a
365fromIndices gen = fromIndicesM (return . gen)
367-- Adapted from the vector package
368{-# INLINE_NORMAL generateM #-}
369generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
370generateM n gen = n `seq` Stream step 0
371 where
372 {-# INLINE_LATE step #-}
373 step _ i | i < n = do
374 x <- gen i
375 return $ Yield x (i + 1)
376 | otherwise = return Stop
378{-# INLINE generate #-}
379generate :: Monad m => Int -> (Int -> a) -> Stream m a
380generate n gen = generateM n (return . gen)
382-------------------------------------------------------------------------------
383-- Iteration
384-------------------------------------------------------------------------------
386{-# INLINE_NORMAL iterateM #-}
387iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
388iterateM step = Stream (\_ st -> st >>= \(!x) -> return $ Yield x (step x))
390{-# INLINE_NORMAL iterate #-}
391iterate :: Monad m => (a -> a) -> a -> Stream m a
392iterate step st = iterateM (return . step) (return st)
394-------------------------------------------------------------------------------
395-- From containers
396-------------------------------------------------------------------------------
398-- XXX we need the MonadAsync constraint because of a rewrite rule.
399-- | Convert a list of monadic actions to a 'Stream'
400{-# INLINE_LATE fromListM #-}
401fromListM :: MonadAsync m => [m a] -> Stream m a
402fromListM = Stream step
403 where
404 {-# INLINE_LATE step #-}
405 step _ (m:ms) = m >>= \x -> return $ Yield x ms
406 step _ [] = return Stop