git.iko.soy/streamly
on edbb7b8 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.Data.Time.Clock
94 (Clock(Monotonic), asyncClock, readClock)
95import Streamly.Internal.Data.Time.Units
96 (toAbsTime, AbsTime, toRelTime64, RelTime64)
98import Prelude hiding (iterate, repeat, replicate, takeWhile)
99import Streamly.Internal.Data.Stream.StreamD.Type
100import Streamly.Internal.Data.SVar
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