git.iko.soy/streamly
on 1a46cd5 help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamD/Generate.hs
482 lines 16 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 -- * From Pointers
83 , fromPtr
85 -- * Conversions
86 , fromStreamK
87 , toStreamK
88 )
89where
91#include "inline.hs"
92#include "ArrayMacros.h"
94import Control.Monad.IO.Class (MonadIO(..))
95import Foreign.Ptr (Ptr, plusPtr)
96import Foreign.Storable (Storable (peek), sizeOf)
97import Streamly.Internal.Data.Time.Clock
98 (Clock(Monotonic), asyncClock, readClock)
99import Streamly.Internal.Data.Time.Units
100 (toAbsTime, AbsTime, toRelTime64, RelTime64)
102#ifdef USE_UNFOLDS_EVERYWHERE
103import qualified Streamly.Internal.Data.Unfold as Unfold
104import qualified Streamly.Internal.Data.Unfold.Enumeration as Unfold
105#endif
107import Prelude hiding (iterate, repeat, replicate, takeWhile)
108import Streamly.Internal.Data.Stream.StreamD.Type
110------------------------------------------------------------------------------
111-- Primitives
112------------------------------------------------------------------------------
114-- | An empty 'Stream'.
115{-# INLINE_NORMAL nil #-}
116nil :: Monad m => Stream m a
117nil = Stream (\_ _ -> return Stop) ()
119-- XXX implement in terms of consM?
120-- cons x = consM (return x)
121--
122-- | Can fuse but has O(n^2) complexity.
123{-# INLINE_NORMAL cons #-}
124cons :: Monad m => a -> Stream m a -> Stream m a
125cons x (Stream step state) = Stream step1 Nothing
126 where
127 {-# INLINE_LATE step1 #-}
128 step1 _ Nothing = return $ Yield x (Just state)
129 step1 gst (Just st) = do
130 r <- step gst st
131 return $
132 case r of
133 Yield a s -> Yield a (Just s)
134 Skip s -> Skip (Just s)
135 Stop -> Stop
137------------------------------------------------------------------------------
138-- Unfolding
139------------------------------------------------------------------------------
141-- Adapted from vector package
142{-# INLINE_NORMAL unfoldrM #-}
143unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
144#ifdef USE_UNFOLDS_EVERYWHERE
145unfoldrM next = unfold (Unfold.unfoldrM next)
146#else
147unfoldrM next = Stream step
148 where
149 {-# INLINE_LATE step #-}
150 step _ st = do
151 r <- next st
152 return $ case r of
153 Just (x, s) -> Yield x s
154 Nothing -> Stop
155#endif
157{-# INLINE_LATE unfoldr #-}
158unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
159unfoldr f = unfoldrM (return . f)
161------------------------------------------------------------------------------
162-- From values
163------------------------------------------------------------------------------
165{-# INLINE_NORMAL repeatM #-}
166repeatM :: Monad m => m a -> Stream m a
167#ifdef USE_UNFOLDS_EVERYWHERE
168repeatM = unfold Unfold.repeatM
169#else
170repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
171#endif
173{-# INLINE_NORMAL repeat #-}
174repeat :: Monad m => a -> Stream m a
175#ifdef USE_UNFOLDS_EVERYWHERE
176repeat x = repeatM (pure x)
177#else
178repeat x = Stream (\_ _ -> return $ Yield x ()) ()
179#endif
181-- Adapted from the vector package
182{-# INLINE_NORMAL replicateM #-}
183replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
184#ifdef USE_UNFOLDS_EVERYWHERE
185replicateM n p = unfold Unfold.replicateM (n, p)
186#else
187replicateM n p = Stream step n
188 where
189 {-# INLINE_LATE step #-}
190 step _ (i :: Int)
191 | i <= 0 = return Stop
192 | otherwise = do
193 x <- p
194 return $ Yield x (i - 1)
195#endif
197{-# INLINE_NORMAL replicate #-}
198replicate :: Monad m => Int -> a -> Stream m a
199replicate n x = replicateM n (return x)
201------------------------------------------------------------------------------
202-- Enumeration of Num
203------------------------------------------------------------------------------
205-- | For floating point numbers if the increment is less than the precision then
206-- it just gets lost. Therefore we cannot always increment it correctly by just
207-- repeated addition.
208-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
209-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
210--
211-- Instead we accumulate the increment counter and compute the increment
212-- every time before adding it to the starting number.
213--
214-- This works for Integrals as well as floating point numbers, but
215-- enumerateFromStepIntegral is faster for integrals.
216{-# INLINE_NORMAL enumerateFromStepNum #-}
217enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
218#ifdef USE_UNFOLDS_EVERYWHERE
219enumerateFromStepNum from stride =
220 unfold Unfold.enumerateFromStepNum (from, stride)
221#else
222enumerateFromStepNum from stride = Stream step 0
223 where
224 {-# INLINE_LATE step #-}
225 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
226#endif
228{-# INLINE_NORMAL enumerateFromNum #-}
229enumerateFromNum :: (Monad m, Num a) => a -> Stream m a
230enumerateFromNum from = enumerateFromStepNum from 1
232{-# INLINE_NORMAL enumerateFromThenNum #-}
233enumerateFromThenNum :: (Monad m, Num a) => a -> a -> Stream m a
234enumerateFromThenNum from next = enumerateFromStepNum from (next - from)
236------------------------------------------------------------------------------
237-- Enumeration of Integrals
238------------------------------------------------------------------------------
240#ifndef USE_UNFOLDS_EVERYWHERE
241data EnumState a = EnumInit | EnumYield a a a | EnumStop
243{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
244enumerateFromThenToIntegralUp
245 :: (Monad m, Integral a)
246 => a -> a -> a -> Stream m a
247enumerateFromThenToIntegralUp from next to = Stream step EnumInit
248 where
249 {-# INLINE_LATE step #-}
250 step _ EnumInit =
251 return $
252 if to < next
253 then if to < from
254 then Stop
255 else Yield from EnumStop
256 else -- from <= next <= to
257 let stride = next - from
258 in Skip $ EnumYield from stride (to - stride)
260 step _ (EnumYield x stride toMinus) =
261 return $
262 if x > toMinus
263 then Yield x EnumStop
264 else Yield x $ EnumYield (x + stride) stride toMinus
266 step _ EnumStop = return Stop
268{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
269enumerateFromThenToIntegralDn
270 :: (Monad m, Integral a)
271 => a -> a -> a -> Stream m a
272enumerateFromThenToIntegralDn from next to = Stream step EnumInit
273 where
274 {-# INLINE_LATE step #-}
275 step _ EnumInit =
276 return $ if to > next
277 then if to > from
278 then Stop
279 else Yield from EnumStop
280 else -- from >= next >= to
281 let stride = next - from
282 in Skip $ EnumYield from stride (to - stride)
284 step _ (EnumYield x stride toMinus) =
285 return $
286 if x < toMinus
287 then Yield x EnumStop
288 else Yield x $ EnumYield (x + stride) stride toMinus
290 step _ EnumStop = return Stop
291#endif
293-- XXX This can perhaps be simplified and written in terms of
294-- enumeratFromStepIntegral as we have done in unfolds. But anyway we should be
295-- replacing the stream generation module with unfolds.
296{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
297enumerateFromThenToIntegral
298 :: (Monad m, Integral a)
299 => a -> a -> a -> Stream m a
300#ifdef USE_UNFOLDS_EVERYWHERE
301enumerateFromThenToIntegral from next to =
302 unfold Unfold.enumerateFromThenToIntegral (from, next, to)
303#else
304enumerateFromThenToIntegral from next to
305 | next >= from = enumerateFromThenToIntegralUp from next to
306 | otherwise = enumerateFromThenToIntegralDn from next to
307#endif
309{-# INLINE_NORMAL enumerateFromThenIntegral #-}
310enumerateFromThenIntegral
311 :: (Monad m, Integral a, Bounded a)
312 => a -> a -> Stream m a
313#ifdef USE_UNFOLDS_EVERYWHERE
314enumerateFromThenIntegral from next =
315 unfold Unfold.enumerateFromThenIntegralBounded (from, next)
316#else
317enumerateFromThenIntegral from next =
318 if next > from
319 then enumerateFromThenToIntegralUp from next maxBound
320 else enumerateFromThenToIntegralDn from next minBound
321#endif
323-- | Can be used to enumerate unbounded integrals. This does not check for
324-- overflow or underflow for bounded integrals.
325--
326{-# INLINE_NORMAL enumerateFromStepIntegral #-}
327enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
328#ifdef USE_UNFOLDS_EVERYWHERE
329enumerateFromStepIntegral from stride =
330 unfold Unfold.enumerateFromStepIntegral (from, stride)
331#else
332enumerateFromStepIntegral from stride =
333 from `seq` stride `seq` Stream step from
334 where
335 {-# INLINE_LATE step #-}
336 step _ !x = return $ Yield x $! (x + stride)
337#endif
339-- | Enumerate upwards from @from@ to @to@. We are assuming that "to" is
340-- constrained by the type to be within max/min bounds.
341{-# INLINE enumerateFromToIntegral #-}
342enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
343enumerateFromToIntegral from to =
344 takeWhile (<= to) $ enumerateFromStepIntegral from 1
346{-# INLINE enumerateFromIntegral #-}
347enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
348enumerateFromIntegral from = enumerateFromToIntegral from maxBound
350------------------------------------------------------------------------------
351-- Enumeration of Fractionals
352------------------------------------------------------------------------------
354-- | We cannot write a general function for Num. The only way to write code
355-- portable between the two is to use a 'Real' constraint and convert between
356-- Fractional and Integral using fromRational which is horribly slow.
357{-# INLINE_NORMAL enumerateFromToFractional #-}
358enumerateFromToFractional
359 :: (Monad m, Fractional a, Ord a)
360 => a -> a -> Stream m a
361enumerateFromToFractional from to =
362 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
364{-# INLINE_NORMAL enumerateFromThenToFractional #-}
365enumerateFromThenToFractional
366 :: (Monad m, Fractional a, Ord a)
367 => a -> a -> a -> Stream m a
368enumerateFromThenToFractional from next to =
369 takeWhile predicate $ enumerateFromThenNum from next
370 where
371 mid = (next - from) / 2
372 predicate | next >= from = (<= to + mid)
373 | otherwise = (>= to + mid)
375------------------------------------------------------------------------------
376-- Time Enumeration
377------------------------------------------------------------------------------
379{-# INLINE_NORMAL times #-}
380times :: MonadIO m => Double -> Stream m (AbsTime, RelTime64)
381times g = Stream step Nothing
383 where
385 {-# INLINE_LATE step #-}
386 step _ Nothing = do
387 clock <- liftIO $ asyncClock Monotonic g
388 a <- liftIO $ readClock clock
389 return $ Skip $ Just (clock, a)
391 step _ s@(Just (clock, t0)) = do
392 a <- liftIO $ readClock clock
393 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
394 -- efficiency. or maybe we can use a representation using Double for
395 -- floating precision time
396 return $ Yield (toAbsTime t0, toRelTime64 (a - t0)) s
398-------------------------------------------------------------------------------
399-- From Generators
400-------------------------------------------------------------------------------
402{-# INLINE_NORMAL fromIndicesM #-}
403fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
404#ifdef USE_UNFOLDS_EVERYWHERE
405fromIndicesM gen = unfold (Unfold.fromIndicesM gen) 0
406#else
407fromIndicesM gen = Stream step 0
408 where
409 {-# INLINE_LATE step #-}
410 step _ i = do
411 x <- gen i
412 return $ Yield x (i + 1)
413#endif
415{-# INLINE fromIndices #-}
416fromIndices :: Monad m => (Int -> a) -> Stream m a
417fromIndices gen = fromIndicesM (return . gen)
419-- Adapted from the vector package
420{-# INLINE_NORMAL generateM #-}
421generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
422generateM n gen = n `seq` Stream step 0
423 where
424 {-# INLINE_LATE step #-}
425 step _ i | i < n = do
426 x <- gen i
427 return $ Yield x (i + 1)
428 | otherwise = return Stop
430{-# INLINE generate #-}
431generate :: Monad m => Int -> (Int -> a) -> Stream m a
432generate n gen = generateM n (return . gen)
434-------------------------------------------------------------------------------
435-- Iteration
436-------------------------------------------------------------------------------
438{-# INLINE_NORMAL iterateM #-}
439iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
440#ifdef USE_UNFOLDS_EVERYWHERE
441iterateM step = unfold (Unfold.iterateM step)
442#else
443iterateM step = Stream (\_ st -> st >>= \(!x) -> return $ Yield x (step x))
444#endif
446{-# INLINE_NORMAL iterate #-}
447iterate :: Monad m => (a -> a) -> a -> Stream m a
448iterate step st = iterateM (return . step) (return st)
450-------------------------------------------------------------------------------
451-- From containers
452-------------------------------------------------------------------------------
454-- | Convert a list of monadic actions to a 'Stream'
455{-# INLINE_LATE fromListM #-}
456fromListM :: Monad m => [m a] -> Stream m a
457#ifdef USE_UNFOLDS_EVERYWHERE
458fromListM = unfold Unfold.fromListM
459#else
460fromListM = Stream step
461 where
462 {-# INLINE_LATE step #-}
463 step _ (m:ms) = m >>= \x -> return $ Yield x ms
464 step _ [] = return Stop
465#endif
467-------------------------------------------------------------------------------
468-- From pointers
469-------------------------------------------------------------------------------
471-- | Read an infinite stream from a pointer, advancing the pointer as needed.
472-- The caller is responsible to end the stream safely.
473{-# INLINE fromPtr #-}
474fromPtr :: forall m a. (MonadIO m, Storable a) => Ptr a -> Stream m a
475fromPtr = Stream step
477 where
479 {-# INLINE_LATE step #-}
480 step _ p = do
481 x <- liftIO $ peek p
482 return $ Yield x (PTR_NEXT(p, a))