@@ -27,20 +27,15 @@ module Streamly.Internal.Data.Unfold.Enumeration
2727 , enumerateFromStepNum
2828 , enumerateFromNum
2929 , enumerateFromThenNum
30− , enumerateFromToNum
31− , enumerateFromThenToNum
3230
3331 -- ** Enumerating unbounded 'Integral' Types
3432 , enumerateFromStepIntegral
3533 , enumerateFromIntegral
3634 , enumerateFromThenIntegral
37−
38− , enumerateFromToStepIntegral
3935 , enumerateFromToIntegral
4036 , enumerateFromThenToIntegral
4137
4238 -- ** Enumerating 'Bounded' 'Integral' Types
43− , enumerateFromStepIntegralBounded
4439 , enumerateFromIntegralBounded
4540 , enumerateFromThenIntegralBounded
4641 , enumerateFromToIntegralBounded
@@ -160,61 +155,6 @@ enumerateFromThenNum =
160155enumerateFromNum :: (Monad m, Num a) => Unfold m a a
161156enumerateFromNum = lmap (\from -> (from, 1)) enumerateFromStepNum
162157
163−-- | Generate a finite stream starting from a from and next values with
164−-- stride of (next-from) up to value. The implementation is numerically
165−-- stable for floating point values.
166−--
167−-- Note that 'enumerateFromThenToIntegral' is faster for integrals.
168−--
169−-- /Internal/
170−--
171−{-# INLINE enumerateFromThenToNum #-}
172−enumerateFromThenToNum :: (Monad m, Ord a, Fractional a) =>
173− Unfold m (a, a, a) a
174−enumerateFromThenToNum = Unfold step inject
175−
176− where
177−
178− inject (!from, !next, !to) = return (from, next - from, to, 0)
179−
180− {-# INLINE_LATE step #-}
181− step (from, stride, to, !i) = return $
182− if stride == 0
183− then if from <= to then Yield from $! (from, stride, to, i + 1) else Stop
184− else
185− if stride > 0
186− then
187− if (from + i * stride) <= to + stride / 2
188− then (Yield $! (from + i * stride)) $! (from, stride, to, i + 1)
189− else Stop
190− else
191− if (from + i * stride) >= to + stride / 2
192− then (Yield $! (from + i * stride)) $! (from, stride, to, i + 1)
193− else Stop
194−
195−-- | Generate a finite stream starting from a from till to value.
196−-- stride of value 1 will be used. The implementation is numerically
197−-- stable for floating point values.
198−--
199−-- Note that 'enumerateFromToIntegral' is faster for integrals.
200−--
201−-- /Internal/
202−--
203−{-# INLINE enumerateFromToNum #-}
204−enumerateFromToNum :: (Monad m, Ord a, Fractional a) =>
205− Unfold m (a, a) a
206−enumerateFromToNum = Unfold step inject
207−
208− where
209−
210− inject (!from, !to) = return (from, 1, to, 0)
211−
212− {-# INLINE_LATE step #-}
213− step (from, !stride, to, !i) = return $
214− if (from + i * stride) <= to + stride / 2
215− then (Yield $! (from + i * stride)) $! (from, stride, to, i + 1)
216− else Stop
217−
218158------------------------------------------------------------------------------
219159-- Enumeration of Integrals
220160------------------------------------------------------------------------------
@@ -224,147 +164,91 @@ enumerateFromToNum = Unfold step inject
224164--
225165-- /Internal/
226166{-# INLINE_NORMAL enumerateFromStepIntegral #-}
227−enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
167+enumerateFromStepIntegral :: (Monad m, Integral a) => Unfold m (a, a) a
228168enumerateFromStepIntegral = Unfold step inject
229− where
230− inject (from, next) = from `seq` next `seq` return (from, next - from)
231− {-# INLINE_LATE step #-}
232− step (x, stride) = return $ Yield x $! (x + stride, stride)
233169
234−-- | Can be used to enumerate unbounded integrals. Starting with value from
235−-- with stride of next-from till value of to.
236−--
237−-- /Internal/
238−--
239−{-# INLINE_NORMAL enumerateFromToStepIntegral #-}
240−enumerateFromToStepIntegral :: (Integral a, Monad m) =>
241− Unfold m (a, a, a) a
242−enumerateFromToStepIntegral = Unfold step inject
243170 where
244− inject (from, next, to) =
245− from `seq` next `seq` to `seq` return (from, next - from, to)
246− {-# INLINE_LATE step #-}
247− step (x, stride, to) = return $
248− if stride == 0
249− then
250− if x <= to then Yield x $! (x, stride, to) else Stop
251−
252− else
253− if stride > 0
254− then
255− if x <= to
256− then Yield x $! (x + stride, stride, to)
257− else Stop
258− else
259− if x >= to
260− then Yield x $! (x + stride, stride, to)
261− else Stop
262171
263−------------------------------------------------------------------------------
264−-- Enumeration of Bounded Integrals
265−------------------------------------------------------------------------------
172+ inject (from, stride) = from `seq` stride `seq` return (from, stride)
266173
267−-- | Can be used to enumerate bounded integrals from starting from with
268−-- stride of (next-from) . This check for overflow or underflow for
269−-- bounded integrals.
270−--
271−-- /Internal/
272−--
273−{-# INLINE_NORMAL enumerateFromStepIntegralBounded #-}
274−enumerateFromStepIntegralBounded ::
275− forall a m. (Integral a, Bounded a, Monad m) =>
276− Unfold m (a, a) a
277−enumerateFromStepIntegralBounded = Unfold step inject
278− where
279− inject (from, next) =
280− from `seq` next `seq` return (from, next - from)
281174 {-# INLINE_LATE step #-}
282− step (x, stride) = return $
283− if stride == 0
284− then
285− Yield x $! (x, stride)
286−
287− else
288− if stride > 0
289− then
290− if x <= (maxBound :: a)
291− then Yield x $! (x + stride, stride)
292− else Stop
293− else
294− if x >= (minBound :: a)
295− then Yield x $! (x + stride, stride)
296− else Stop
297−
298−-- | Can be used to enumerate bounded integrals from starting from with
299−-- stride of (next-from) till to value. This check for overflow or underflow
300−-- for bounded integrals.
301−--
302−-- /Internal/
303−--
304−{-# INLINE_NORMAL enumerateFromToStepIntegralBounded #-}
305−enumerateFromToStepIntegralBounded ::
306− forall a m. (Integral a, Bounded a, Monad m) =>
307− Unfold m (a, a, a) a
308−enumerateFromToStepIntegralBounded = Unfold step inject
309− where
310− inject (from, next, to) =
311− from `seq` next `seq` to `seq` return (from, next - from, to)
312− {-# INLINE_LATE step #-}
313− step (x, stride, to) = return $
314− if stride == 0
315− then
316− if x <= to then Yield x $! (x, stride, to) else Stop
317−
318− else
319− if stride > 0
320− then
321− if x <= to && x <= (maxBound :: a)
322− then Yield x $! (x + stride, stride, to)
323− else Stop
324− else
325− if x >= to && x >= (minBound :: a)
326− then Yield x $! (x + stride, stride, to)
327− else Stop
175+ step (x, stride) = return $ Yield x $! (x + stride, stride)
328176
329177-- Enumerate Unbounded Integrals ----------------------------------------------
330178{-# INLINE enumerateFromIntegral #-}
331179enumerateFromIntegral :: (Monad m, Integral a) => Unfold m a a
332−enumerateFromIntegral = lmap (\a -> (a, a + 1)) enumerateFromStepIntegral
180+enumerateFromIntegral = lmap (\from -> (from, 1)) enumerateFromStepIntegral
333181
334182{-# INLINE enumerateFromThenIntegral #-}
335183enumerateFromThenIntegral :: (Monad m, Integral a ) => Unfold m (a, a) a
336−enumerateFromThenIntegral = enumerateFromStepIntegral
184+enumerateFromThenIntegral =
185+ lmap (\(from, next) -> (from, next - from)) enumerateFromStepIntegral
337186
338187{-# INLINE enumerateFromToIntegral #-}
339188enumerateFromToIntegral :: (Monad m, Integral a) => Unfold m (a, a) a
340−enumerateFromToIntegral = lmap (\(a, c) -> (a, a + 1, c)) enumerateFromToStepIntegral
189+enumerateFromToIntegral =
190+ takeWhileMWithInput (\(_, to) b -> return $ b <= to)
191+ $ lmap (\(from, _) -> (from, 1)) enumerateFromStepIntegral
341192
342193{-# INLINE enumerateFromThenToIntegral #-}
343194enumerateFromThenToIntegral :: (Monad m, Integral a) => Unfold m (a, a, a) a
344−enumerateFromThenToIntegral = enumerateFromToStepIntegral
195+enumerateFromThenToIntegral =
196+ takeWhileMWithInput cond $ lmap toFromStep enumerateFromStepIntegral
197+
198+ where
199+
200+ toFromStep (from, next, _) = (from, next - from)
201+
202+ cond (from, next, to) b =
203+ return
204+ $ if next >= from
205+ then b <= to
206+ else b >= to
345207
346208-- Enumerate Bounded Integrals ------------------------------------------------
347209{-# INLINE enumerateFromIntegralBounded #-}
348−enumerateFromIntegralBounded :: (Monad m, Bounded a, Integral a) =>
210+enumerateFromIntegralBounded :: (Monad m, Integral a, Bounded a) =>
349211 Unfold m a a
350−enumerateFromIntegralBounded =
351− lmap (\a -> (a, a + 1, maxBound)) enumerateFromToStepIntegralBounded
212+enumerateFromIntegralBounded = supplySecond maxBound enumerateFromToIntegral
352213
353214{-# INLINE enumerateFromThenIntegralBounded #-}
354−enumerateFromThenIntegralBounded :: (Monad m, Bounded a, Integral a ) =>
215+enumerateFromThenIntegralBounded :: (Monad m, Integral a, Bounded a ) =>
355216 Unfold m (a, a) a
356−enumerateFromThenIntegralBounded = enumerateFromStepIntegralBounded
217+enumerateFromThenIntegralBounded =
218+ takeWhileMWithInput cond $ lmap toFromStep enumerateFromStepIntegral
219+
220+ where
221+
222+ toFromStep (from, next) = (from, next - from)
223+
224+ cond (from, next) b =
225+ return
226+ $ if next >= from
227+ then b <= maxBound
228+ else b >= minBound
357229
358230{-# INLINE enumerateFromToIntegralBounded #-}
359231enumerateFromToIntegralBounded :: (Monad m, Integral a, Bounded a) =>
360232 Unfold m (a, a) a
361233enumerateFromToIntegralBounded =
362− lmap (\(a, c) -> (a, a + 1, c)) enumerateFromToStepIntegralBounded
234+ takeWhileMWithInput (\(_, to) b -> return $ b <= to)
235+ $ lmap fst enumerateFromIntegralBounded
363236
364237{-# INLINE enumerateFromThenToIntegralBounded #-}
365238enumerateFromThenToIntegralBounded :: (Monad m, Integral a, Bounded a) =>
366239 Unfold m (a, a, a) a
367−enumerateFromThenToIntegralBounded = enumerateFromToStepIntegralBounded
240+enumerateFromThenToIntegralBounded =
241+ takeWhileMWithInput cond $ lmap toFromThen enumerateFromThenIntegralBounded
242+
243+ where
244+
245+ toFromThen (from, next, _) = (from, next)
246+
247+ cond (from, next, to) b =
248+ return
249+ $ if next >= from
250+ then b <= to
251+ else b >= to
368252
369253------------------------------------------------------------------------------
370254-- Enumeration of Fractionals
@@ -378,23 +262,41 @@ enumerateFromFractional = enumerateFromNum
378262enumerateFromThenFractional :: (Monad m, Fractional a) => Unfold m (a, a) a
379263enumerateFromThenFractional = enumerateFromThenNum
380264
265+-- | Same as 'enumerateFromStepNum' with a step of 1 and enumerating up to the
266+-- specified upper limit rounded to the nearest integral value:
267+--
268+-- >>> enumerateFromToFractional =
269+--
270+-- /Internal/
271+--
381272{-# INLINE_NORMAL enumerateFromToFractional #-}
382273enumerateFromToFractional :: (Monad m, Fractional a, Ord a) =>
383274 Unfold m (a, a) a
384−enumerateFromToFractional = enumerateFromToNum
275+enumerateFromToFractional =
276+ takeWhileMWithInput (\(_, to) b -> return $ b <= to + 1 / 2)
277+ $ lmap (\(from, _) -> (from, 1)) enumerateFromStepNum
385278
386−{-# INLINE_NORMAL enumerateFromThenToFractional #-}
279+{-# INLINE enumerateFromThenToFractional #-}
387280enumerateFromThenToFractional :: (Monad m, Fractional a, Ord a) =>
388281 Unfold m (a, a, a) a
389−enumerateFromThenToFractional = enumerateFromThenToNum
282+enumerateFromThenToFractional =
283+ takeWhileMWithInput cond $ lmap toFromStep enumerateFromStepNum
284+
285+ where
286+
287+ toFromStep (from, next, _) = (from, next - from)
288+
289+ cond (from, next, to) b =
290+ let stride = next - from
291+ in return
292+ $ if next >= from
293+ then b <= to + stride / 2
294+ else b >= to + stride / 2
390295
391296-------------------------------------------------------------------------------
392297-- Enumeration of Enum types not larger than Int
393298-------------------------------------------------------------------------------
394299
395−-- XXX should we check for bounds here? Otherwise it will cycle back to "to" if
396−-- to is smaller?
397−--
398300-- | Enumerate from given starting Enum value 'from' and to Enum value 'to'
399301-- with stride of 1 till to value.
400302--
@@ -402,16 +304,8 @@ enumerateFromThenToFractional = enumerateFromThenToNum
402304--
403305{-# INLINE enumerateFromToSmall #-}
404306enumerateFromToSmall :: (Monad m, Enum a) => Unfold m (a, a) a
405−enumerateFromToSmall = map toEnum $
406− Unfold step inject
407− where
408− inject (from, to) =
409− fromEnum from
410− `seq` fromEnum to
411− `seq` return (fromEnum from, 1, fromEnum to)
412− {-# INLINE_LATE step #-}
413− step (x, stride, to) = return $
414− if x <= to then Yield x $! (x + stride, stride, to) else Stop
307+enumerateFromToSmall =
308+ fmap toEnum (lmap (bimap fromEnum fromEnum) enumerateFromToIntegral)
415309
416310-- | Enumerate from given starting Enum value 'from' and then Enum value 'next'
417311-- and to Enum value 'to' with stride of (fromEnum next - fromEnum from)
@@ -421,31 +315,9 @@ enumerateFromToSmall = map toEnum $
421315--
422316{-# INLINE enumerateFromThenToSmall #-}
423317enumerateFromThenToSmall :: (Monad m, Enum a) => Unfold m (a, a, a) a
424−enumerateFromThenToSmall = map toEnum $
425− Unfold step inject
426− where
427− inject (from, next, to) =
428− fromEnum from
429− `seq` fromEnum next
430− `seq` fromEnum to
431− `seq` return
432− (fromEnum from, fromEnum next - fromEnum from, fromEnum to)
433− {-# INLINE_LATE step #-}
434− step (x, stride, to) = return $
435− if stride == 0
436− then
437− if x <= to then Yield x $! (x, stride, to) else Stop
438−
439− else
440− if stride > 0
441− then
442− if x <= to
443− then Yield x $! (x + stride, stride, to)
444− else Stop
445− else
446− if x >= 0 && x >= to
447− then Yield x $! (x + stride, stride, to)
448− else Stop
318+enumerateFromThenToSmall =
319+ let toInts (x, y, z) = (fromEnum x, fromEnum y, fromEnum z)
320+ in fmap toEnum (lmap toInts enumerateFromThenToIntegral)
449321
450322-------------------------------------------------------------------------------
451323-- Bounded Enumeration of Enum types not larger than Int
@@ -457,16 +329,8 @@ enumerateFromThenToSmall = map toEnum $
457329-- /Internal/
458330--
459331{-# INLINE enumerateFromSmallBounded #-}
460−enumerateFromSmallBounded :: forall a m. (Monad m, Bounded a, Enum a) => Unfold m a a
461−enumerateFromSmallBounded = map toEnum $
462− Unfold step inject
463− where
464− inject from =
465− fromEnum from
466− `seq` return (fromEnum from, 1, fromEnum (maxBound :: a))
467− {-# INLINE_LATE step #-}
468− step (x, stride, to) = return $
469− if x <= to then Yield x $! (x + stride, stride, to) else Stop
332+enumerateFromSmallBounded :: (Monad m, Enum a, Bounded a) => Unfold m a a
333+enumerateFromSmallBounded = supplySecond maxBound enumerateFromToSmall
470334
471335-- | Enumerate from given starting Enum value 'from' and next Enum value 'next'
472336-- with stride of (fromEnum next - fromEnum from) till maxBound.
@@ -474,35 +338,18 @@ enumerateFromSmallBounded = map toEnum $
474338-- /Internal/
475339--
476340{-# INLINE enumerateFromThenSmallBounded #-}
477−enumerateFromThenSmallBounded :: forall a m. (Monad m, Bounded a, Enum a) =>
341+enumerateFromThenSmallBounded :: forall m a. (Monad m, Enum a, Bounded a) =>
478342 Unfold m (a, a) a
479−enumerateFromThenSmallBounded = map toEnum $
480− Unfold step inject
481− where
482− inject (from, next) =
483− fromEnum from
484− `seq` fromEnum next
485− `seq` return
486− ( fromEnum from
487− , fromEnum next - fromEnum from
488− , fromEnum (maxBound :: a)
489− )
490− {-# INLINE_LATE step #-}
491− step (x, stride, to) = return $
492− if stride == 0
493− then
494− if x <= to then Yield x $! (x, stride, to) else Stop
495−
496− else
497− if stride > 0
498− then
499− if x <= to
500− then Yield x $! (x + stride, stride, to)
501− else Stop
502− else
503− if x >= 0
504− then Yield x $! (x + stride, stride, to)
505− else Stop
343+enumerateFromThenSmallBounded =
344+ let adapt (from, next) =
345+ let frm = fromEnum from
346+ nxt = fromEnum next
347+ stride = nxt - frm
348+ to = if stride >= 0
349+ then fromEnum (maxBound :: a)
350+ else fromEnum (minBound :: a)
351+ in (frm, nxt, to)
352+ in fmap toEnum (lmap adapt enumerateFromThenToIntegral)
506353
507354-------------------------------------------------------------------------------
508355-- Enumerable type class