git.iko.soy/streamly
on a296fac help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Generate.hs
405 lines 13 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 , numFrom
57 , numFromThen
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 , fromStreamD
86 , toStreamD
87 )
88where
90#include "inline.hs"
92import Control.Monad.IO.Class (MonadIO(..))
93import Streamly.Internal.Control.Concurrent (MonadAsync)
94import Streamly.Internal.Data.Time.Clock
95 (Clock(Monotonic), asyncClock, readClock)
96import Streamly.Internal.Data.Time.Units
97 (toAbsTime, AbsTime, toRelTime64, RelTime64)
99import Prelude hiding (iterate, repeat, replicate, takeWhile)
100import Streamly.Internal.Data.Stream.StreamD.Type
102------------------------------------------------------------------------------
103-- Primitives
104------------------------------------------------------------------------------
106-- | An empty 'Stream'.
107{-# INLINE_NORMAL nil #-}
108nil :: Monad m => Stream m a
109nil = Stream (\_ _ -> return Stop) ()
111-- XXX implement in terms of consM?
112-- cons x = consM (return x)
113--
114-- | Can fuse but has O(n^2) complexity.
115{-# INLINE_NORMAL cons #-}
116cons :: Monad m => a -> Stream m a -> Stream m a
117cons x (Stream step state) = Stream step1 Nothing
118 where
119 {-# INLINE_LATE step1 #-}
120 step1 _ Nothing = return $ Yield x (Just state)
121 step1 gst (Just st) = do
122 r <- step gst st
123 return $
124 case r of
125 Yield a s -> Yield a (Just s)
126 Skip s -> Skip (Just s)
127 Stop -> Stop
129------------------------------------------------------------------------------
130-- Unfolding
131------------------------------------------------------------------------------
133-- Adapted from vector package
134{-# INLINE_NORMAL unfoldrM #-}
135unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
136unfoldrM next state = Stream step state
137 where
138 {-# INLINE_LATE step #-}
139 step _ st = do
140 r <- next st
141 return $ case r of
142 Just (x, s) -> Yield x s
143 Nothing -> Stop
145{-# INLINE_LATE unfoldr #-}
146unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
147unfoldr f = unfoldrM (return . f)
149------------------------------------------------------------------------------
150-- From values
151------------------------------------------------------------------------------
153{-# INLINE_NORMAL repeatM #-}
154repeatM :: Monad m => m a -> Stream m a
155repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
157{-# INLINE_NORMAL repeat #-}
158repeat :: Monad m => a -> Stream m a
159repeat x = Stream (\_ _ -> return $ Yield x ()) ()
161-- Adapted from the vector package
162{-# INLINE_NORMAL replicateM #-}
163replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
164replicateM n p = Stream step n
165 where
166 {-# INLINE_LATE step #-}
167 step _ (i :: Int)
168 | i <= 0 = return Stop
169 | otherwise = do
170 x <- p
171 return $ Yield x (i - 1)
173{-# INLINE_NORMAL replicate #-}
174replicate :: Monad m => Int -> a -> Stream m a
175replicate n x = replicateM n (return x)
177------------------------------------------------------------------------------
178-- Enumeration of Num
179------------------------------------------------------------------------------
181-- | For floating point numbers if the increment is less than the precision then
182-- it just gets lost. Therefore we cannot always increment it correctly by just
183-- repeated addition.
184-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
185-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
186--
187-- Instead we accumulate the increment counter and compute the increment
188-- every time before adding it to the starting number.
189--
190-- This works for Integrals as well as floating point numbers, but
191-- enumerateFromStepIntegral is faster for integrals.
192{-# INLINE_NORMAL enumerateFromStepNum #-}
193enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
194enumerateFromStepNum from stride = Stream step 0
195 where
196 {-# INLINE_LATE step #-}
197 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
199{-# INLINE_NORMAL numFrom #-}
200numFrom :: (Monad m, Num a) => a -> Stream m a
201numFrom from = enumerateFromStepNum from 1
203{-# INLINE_NORMAL numFromThen #-}
204numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
205numFromThen from next = enumerateFromStepNum from (next - from)
207------------------------------------------------------------------------------
208-- Enumeration of Integrals
209------------------------------------------------------------------------------
211data EnumState a = EnumInit | EnumYield a a a | EnumStop
213{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
214enumerateFromThenToIntegralUp
215 :: (Monad m, Integral a)
216 => a -> a -> a -> Stream m a
217enumerateFromThenToIntegralUp from next to = Stream step EnumInit
218 where
219 {-# INLINE_LATE step #-}
220 step _ EnumInit =
221 return $
222 if to < next
223 then if to < from
224 then Stop
225 else Yield from EnumStop
226 else -- from <= next <= to
227 let stride = next - from
228 in Skip $ EnumYield from stride (to - stride)
230 step _ (EnumYield x stride toMinus) =
231 return $
232 if x > toMinus
233 then Yield x EnumStop
234 else Yield x $ EnumYield (x + stride) stride toMinus
236 step _ EnumStop = return Stop
238{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
239enumerateFromThenToIntegralDn
240 :: (Monad m, Integral a)
241 => a -> a -> a -> Stream m a
242enumerateFromThenToIntegralDn from next to = Stream step EnumInit
243 where
244 {-# INLINE_LATE step #-}
245 step _ EnumInit =
246 return $ if to > next
247 then if to > from
248 then Stop
249 else Yield from EnumStop
250 else -- from >= next >= to
251 let stride = next - from
252 in Skip $ EnumYield from stride (to - stride)
254 step _ (EnumYield x stride toMinus) =
255 return $
256 if x < toMinus
257 then Yield x EnumStop
258 else Yield x $ EnumYield (x + stride) stride toMinus
260 step _ EnumStop = return Stop
262{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
263enumerateFromThenToIntegral
264 :: (Monad m, Integral a)
265 => a -> a -> a -> Stream m a
266enumerateFromThenToIntegral from next to
267 | next >= from = enumerateFromThenToIntegralUp from next to
268 | otherwise = enumerateFromThenToIntegralDn from next to
270{-# INLINE_NORMAL enumerateFromThenIntegral #-}
271enumerateFromThenIntegral
272 :: (Monad m, Integral a, Bounded a)
273 => a -> a -> Stream m a
274enumerateFromThenIntegral from next =
275 if next > from
276 then enumerateFromThenToIntegralUp from next maxBound
277 else enumerateFromThenToIntegralDn from next minBound
279-- | Can be used to enumerate unbounded integrals. This does not check for
280-- overflow or underflow for bounded integrals.
281--
282{-# INLINE_NORMAL enumerateFromStepIntegral #-}
283enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
284enumerateFromStepIntegral from stride =
285 from `seq` stride `seq` Stream step from
286 where
287 {-# INLINE_LATE step #-}
288 step _ !x = return $ Yield x $! (x + stride)
290-- | Enumerate upwards from @from@ to @to@. We are assuming that "to" is
291-- constrained by the type to be within max/min bounds.
292{-# INLINE enumerateFromToIntegral #-}
293enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
294enumerateFromToIntegral from to =
295 takeWhile (<= to) $ enumerateFromStepIntegral from 1
297{-# INLINE enumerateFromIntegral #-}
298enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
299enumerateFromIntegral from = enumerateFromToIntegral from maxBound
301------------------------------------------------------------------------------
302-- Enumeration of Fractionals
303------------------------------------------------------------------------------
305-- | We cannot write a general function for Num. The only way to write code
306-- portable between the two is to use a 'Real' constraint and convert between
307-- Fractional and Integral using fromRational which is horribly slow.
308{-# INLINE_NORMAL enumerateFromToFractional #-}
309enumerateFromToFractional
310 :: (Monad m, Fractional a, Ord a)
311 => a -> a -> Stream m a
312enumerateFromToFractional from to =
313 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
315{-# INLINE_NORMAL enumerateFromThenToFractional #-}
316enumerateFromThenToFractional
317 :: (Monad m, Fractional a, Ord a)
318 => a -> a -> a -> Stream m a
319enumerateFromThenToFractional from next to =
320 takeWhile predicate $ numFromThen from next
321 where
322 mid = (next - from) / 2
323 predicate | next >= from = (<= to + mid)
324 | otherwise = (>= to + mid)
326------------------------------------------------------------------------------
327-- Time Enumeration
328------------------------------------------------------------------------------
330{-# INLINE_NORMAL times #-}
331times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
332times g = Stream step Nothing
334 where
336 {-# INLINE_LATE step #-}
337 step _ Nothing = do
338 clock <- liftIO $ asyncClock Monotonic g
339 a <- liftIO $ readClock clock
340 return $ Skip $ Just (clock, a)
342 step _ s@(Just (clock, t0)) = do
343 a <- liftIO $ readClock clock
344 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
345 -- efficiency. or maybe we can use a representation using Double for
346 -- floating precision time
347 return $ Yield (toAbsTime t0, toRelTime64 (a - t0)) s
349-------------------------------------------------------------------------------
350-- From Generators
351-------------------------------------------------------------------------------
353{-# INLINE_NORMAL fromIndicesM #-}
354fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
355fromIndicesM gen = Stream step 0
356 where
357 {-# INLINE_LATE step #-}
358 step _ i = do
359 x <- gen i
360 return $ Yield x (i + 1)
362{-# INLINE fromIndices #-}
363fromIndices :: Monad m => (Int -> a) -> Stream m a
364fromIndices gen = fromIndicesM (return . gen)
366-- Adapted from the vector package
367{-# INLINE_NORMAL generateM #-}
368generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
369generateM n gen = n `seq` Stream step 0
370 where
371 {-# INLINE_LATE step #-}
372 step _ i | i < n = do
373 x <- gen i
374 return $ Yield x (i + 1)
375 | otherwise = return Stop
377{-# INLINE generate #-}
378generate :: Monad m => Int -> (Int -> a) -> Stream m a
379generate n gen = generateM n (return . gen)
381-------------------------------------------------------------------------------
382-- Iteration
383-------------------------------------------------------------------------------
385{-# INLINE_NORMAL iterateM #-}
386iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
387iterateM step = Stream (\_ st -> st >>= \(!x) -> return $ Yield x (step x))
389{-# INLINE_NORMAL iterate #-}
390iterate :: Monad m => (a -> a) -> a -> Stream m a
391iterate step st = iterateM (return . step) (return st)
393-------------------------------------------------------------------------------
394-- From containers
395-------------------------------------------------------------------------------
397-- XXX we need the MonadAsync constraint because of a rewrite rule.
398-- | Convert a list of monadic actions to a 'Stream'
399{-# INLINE_LATE fromListM #-}
400fromListM :: MonadAsync m => [m a] -> Stream m a
401fromListM = Stream step
402 where
403 {-# INLINE_LATE step #-}
404 step _ (m:ms) = m >>= \x -> return $ Yield x ms
405 step _ [] = return Stop