git.iko.soy/streamly
on ee691b6 help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Generate.hs
398 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. concatUnfold). 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 , yield
42 , yieldM
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
179------------------------------------------------------------------------------
181-- This would not work properly for floats, therefore we put an Integral
182-- constraint.
183-- | Can be used to enumerate unbounded integrals. This does not check for
184-- overflow or underflow for bounded integrals.
185{-# INLINE_NORMAL enumerateFromStepIntegral #-}
186enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
187enumerateFromStepIntegral from stride =
188 from `seq` stride `seq` Stream step from
189 where
190 {-# INLINE_LATE step #-}
191 step _ !x = return $ Yield x $! (x + stride)
193data EnumState a = EnumInit | EnumYield a a a | EnumStop
195{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
196enumerateFromThenToIntegralUp
197 :: (Monad m, Integral a)
198 => a -> a -> a -> Stream m a
199enumerateFromThenToIntegralUp from next to = Stream step EnumInit
200 where
201 {-# INLINE_LATE step #-}
202 step _ EnumInit =
203 return $
204 if to < next
205 then if to < from
206 then Stop
207 else Yield from EnumStop
208 else -- from <= next <= to
209 let stride = next - from
210 in Skip $ EnumYield from stride (to - stride)
212 step _ (EnumYield x stride toMinus) =
213 return $
214 if x > toMinus
215 then Yield x EnumStop
216 else Yield x $ EnumYield (x + stride) stride toMinus
218 step _ EnumStop = return Stop
220{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
221enumerateFromThenToIntegralDn
222 :: (Monad m, Integral a)
223 => a -> a -> a -> Stream m a
224enumerateFromThenToIntegralDn from next to = Stream step EnumInit
225 where
226 {-# INLINE_LATE step #-}
227 step _ EnumInit =
228 return $ if to > next
229 then if to > from
230 then Stop
231 else Yield from EnumStop
232 else -- from >= next >= to
233 let stride = next - from
234 in Skip $ EnumYield from stride (to - stride)
236 step _ (EnumYield x stride toMinus) =
237 return $
238 if x < toMinus
239 then Yield x EnumStop
240 else Yield x $ EnumYield (x + stride) stride toMinus
242 step _ EnumStop = return Stop
244{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
245enumerateFromThenToIntegral
246 :: (Monad m, Integral a)
247 => a -> a -> a -> Stream m a
248enumerateFromThenToIntegral from next to
249 | next >= from = enumerateFromThenToIntegralUp from next to
250 | otherwise = enumerateFromThenToIntegralDn from next to
252{-# INLINE_NORMAL enumerateFromThenIntegral #-}
253enumerateFromThenIntegral
254 :: (Monad m, Integral a, Bounded a)
255 => a -> a -> Stream m a
256enumerateFromThenIntegral from next =
257 if next > from
258 then enumerateFromThenToIntegralUp from next maxBound
259 else enumerateFromThenToIntegralDn from next minBound
261-- We are assuming that "to" is constrained by the type to be within
262-- max/min bounds.
263{-# INLINE enumerateFromToIntegral #-}
264enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
265enumerateFromToIntegral from to =
266 takeWhile (<= to) $ enumerateFromStepIntegral from 1
268{-# INLINE enumerateFromIntegral #-}
269enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
270enumerateFromIntegral from = enumerateFromToIntegral from maxBound
272-- We cannot write a general function for Num. The only way to write code
273-- portable between the two is to use a 'Real' constraint and convert between
274-- Fractional and Integral using fromRational which is horribly slow.
275{-# INLINE_NORMAL enumerateFromToFractional #-}
276enumerateFromToFractional
277 :: (Monad m, Fractional a, Ord a)
278 => a -> a -> Stream m a
279enumerateFromToFractional from to =
280 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
282{-# INLINE_NORMAL enumerateFromThenToFractional #-}
283enumerateFromThenToFractional
284 :: (Monad m, Fractional a, Ord a)
285 => a -> a -> a -> Stream m a
286enumerateFromThenToFractional from next to =
287 takeWhile predicate $ numFromThen from next
288 where
289 mid = (next - from) / 2
290 predicate | next >= from = (<= to + mid)
291 | otherwise = (>= to + mid)
293-- For floating point numbers if the increment is less than the precision then
294-- it just gets lost. Therefore we cannot always increment it correctly by just
295-- repeated addition.
296-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
297-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
299-- Instead we accumulate the increment counter and compute the increment
300-- every time before adding it to the starting number.
301--
302-- This works for Integrals as well as floating point numbers, but
303-- enumerateFromStepIntegral is faster for integrals.
304{-# INLINE_NORMAL enumerateFromStepNum #-}
305enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
306enumerateFromStepNum from stride = Stream step 0
307 where
308 {-# INLINE_LATE step #-}
309 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
311{-# INLINE_NORMAL numFrom #-}
312numFrom :: (Monad m, Num a) => a -> Stream m a
313numFrom from = enumerateFromStepNum from 1
315{-# INLINE_NORMAL numFromThen #-}
316numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
317numFromThen from next = enumerateFromStepNum from (next - from)
319------------------------------------------------------------------------------
320-- Time Enumeration
321------------------------------------------------------------------------------
323{-# INLINE_NORMAL times #-}
324times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
325times g = Stream step Nothing
327 where
329 {-# INLINE_LATE step #-}
330 step _ Nothing = do
331 clock <- liftIO $ asyncClock Monotonic g
332 a <- liftIO $ readClock clock
333 return $ Skip $ Just (clock, a)
335 step _ s@(Just (clock, t0)) = do
336 a <- liftIO $ readClock clock
337 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
338 -- efficiency. or maybe we can use a representation using Double for
339 -- floating precision time
340 return $ Yield (toAbsTime t0, toRelTime64 (a - t0)) s
342-------------------------------------------------------------------------------
343-- From Generators
344-------------------------------------------------------------------------------
346{-# INLINE_NORMAL fromIndicesM #-}
347fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
348fromIndicesM gen = Stream step 0
349 where
350 {-# INLINE_LATE step #-}
351 step _ i = do
352 x <- gen i
353 return $ Yield x (i + 1)
355{-# INLINE fromIndices #-}
356fromIndices :: Monad m => (Int -> a) -> Stream m a
357fromIndices gen = fromIndicesM (return . gen)
359-- Adapted from the vector package
360{-# INLINE_NORMAL generateM #-}
361generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
362generateM n gen = n `seq` Stream step 0
363 where
364 {-# INLINE_LATE step #-}
365 step _ i | i < n = do
366 x <- gen i
367 return $ Yield x (i + 1)
368 | otherwise = return Stop
370{-# INLINE generate #-}
371generate :: Monad m => Int -> (Int -> a) -> Stream m a
372generate n gen = generateM n (return . gen)
374-------------------------------------------------------------------------------
375-- Iteration
376-------------------------------------------------------------------------------
378{-# INLINE_NORMAL iterateM #-}
379iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
380iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
382{-# INLINE_NORMAL iterate #-}
383iterate :: Monad m => (a -> a) -> a -> Stream m a
384iterate step st = iterateM (return . step) (return st)
386-------------------------------------------------------------------------------
387-- From containers
388-------------------------------------------------------------------------------
390-- XXX we need the MonadAsync constraint because of a rewrite rule.
391-- | Convert a list of monadic actions to a 'Stream'
392{-# INLINE_LATE fromListM #-}
393fromListM :: MonadAsync m => [m a] -> Stream m a
394fromListM = Stream step
395 where
396 {-# INLINE_LATE step #-}
397 step _ (m:ms) = m >>= \x -> return $ Yield x ms
398 step _ [] = return Stop