core/src/Streamly/Internal/Data/Unfold/Enumeration.hs · streamly

git.iko.soy/streamly

on 43dda21 help sign in

streamly/core/src/Streamly/Internal/Data/Unfold/Enumeration.hs

573 lines 21 kB history raw

1-- |

2-- Module : Streamly.Internal.Data.Unfold.Enumeration

4-- License : BSD-3-Clause

5-- Maintainer : streamly@composewell.com

6-- Stability : experimental

7-- Portability : GHC

8--

9-- The functions defined in this module should be rarely needed for direct use,

10-- try to use the operations from the 'Enumerable' type class

11-- instances instead.

12--

13-- This module provides an 'Enumerable' type class to enumerate 'Enum' types

14-- into a stream. The operations in this type class correspond to similar

15-- operations in the 'Enum' type class, the only difference is that they produce

16-- a stream instead of a list. These operations cannot be defined generically

17-- based on the 'Enum' type class. We provide instances for commonly used

18-- types. If instances for other types are needed convenience functions defined

19-- in this module can be used to define them. Alternatively, these functions

20-- can be used directly.

21--

22module Streamly.Internal.Data.Unfold.Enumeration

23 (

24 Enumerable (..)

26 -- ** Enumerating 'Num' Types

27 , enumerateFromStepNum

28 , enumerateFromNum

29 , enumerateFromThenNum

31 -- ** Enumerating unbounded 'Integral' Types

32 , enumerateFromStepIntegral

33 , enumerateFromIntegral

34 , enumerateFromThenIntegral

35 , enumerateFromToIntegral

36 , enumerateFromThenToIntegral

38 -- ** Enumerating 'Bounded' 'Integral' Types

39 , enumerateFromIntegralBounded

40 , enumerateFromThenIntegralBounded

41 , enumerateFromToIntegralBounded

42 , enumerateFromThenToIntegralBounded

44 -- ** Enumerating small 'Integral' Types

45 -- | Small types are always bounded.

46 , enumerateFromSmallBounded

47 , enumerateFromThenSmallBounded

48 , enumerateFromToSmall

49 , enumerateFromThenToSmall

51 -- ** Enumerating 'Fractional' Types

52 -- | Enumeration of 'Num' specialized to 'Fractional' types.

53 , enumerateFromFractional

54 , enumerateFromThenFractional

55 , enumerateFromToFractional

56 , enumerateFromThenToFractional

57 )

58where

60#include "inline.hs"

61import Data.Fixed

62import Data.Bifunctor (bimap)

63import Data.Int

64import Data.Ratio

65import Data.Word

66import Numeric.Natural

67import Data.Functor.Identity (Identity(..))

68import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))

69import Streamly.Internal.Data.Unfold.Type

70import Prelude

71 hiding (map, mapM, takeWhile, take, filter, const, zipWith

72 , drop, dropWhile)

74-- $setup

75-- >>> :m

76-- >>> import qualified Streamly.Prelude as Stream

77-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold

78-- >>> import Streamly.Internal.Data.Unfold.Type

79-- >>> import Data.Word

81------------------------------------------------------------------------------

82-- Enumeration of Num

83------------------------------------------------------------------------------

85-- | Unfolds @(from, stride)@ generating an infinite stream starting from

86-- @from@ and incrementing every time by @stride@. For 'Bounded' types, after

87-- the value overflows it keeps enumerating in a cycle:

88--

89-- @

90-- >>> Stream.toList $ Stream.take 10 $ Stream.unfold Unfold.enumerateFromStepNum (255::Word8,1)

91-- [255,0,1,2,3,4,5,6,7,8]

92--

93-- @

94--

95-- The implementation is numerically stable for floating point values.

96--

97-- Note 'enumerateFromStepIntegral' is faster for integrals.

98--

99-- /Internal/

100--

101{-# INLINE enumerateFromStepNum #-}

102enumerateFromStepNum :: (Monad m, Num a) => Unfold m (a, a) a

103enumerateFromStepNum = Unfold step inject

104

105 where

106

107 inject (!from, !stride) = return (from, stride, 0)

108

109 -- Note that the counter "i" is the same type as the type being enumerated.

110 -- It may overflow, for example, if we are enumerating Word8, after 255 the

111 -- counter will become 0, but the overflow does not affect the enumeration

112 -- behavior.

113 {-# INLINE_LATE step #-}

114 step (from, stride, i) =

115 return $

116 (Yield $! (from + i * stride)) $! (from, stride, i + 1)

117

118-- | Same as 'enumerateFromStepNum (from, next)' using a stride of @next - from@:

119--

120-- @

121-- >>> enumerateFromThenNum = lmap (\(from, next) -> (from, next - from)) Unfold.enumerateFromStepNum

122--

123-- @

124--

125-- Example:

126-- @

127-- >>> Stream.toList $ Stream.take 10 $ Stream.unfold enumerateFromThenNum (255::Word8,0)

128-- [255,0,1,2,3,4,5,6,7,8]

129--

130-- @

131--

132-- The implementation is numerically stable for floating point values.

133--

134-- Note that 'enumerateFromThenIntegral' is faster for integrals.

135--

136-- Note that in the strange world of floating point numbers, using

137-- @enumerateFromThenNum (from, from + 1)@ is almost exactly the same as

138-- @enumerateFromStepNum (from, 1) but not precisely the same. Because @(from +

139-- 1) - from@ is not exactly 1, it may lose some precision, the loss may also

140-- be aggregated in each step, if you want that precision then use

141-- 'enumerateFromStepNum' instead.

142--

143-- /Internal/

144--

145{-# INLINE enumerateFromThenNum #-}

146enumerateFromThenNum :: (Monad m, Num a) => Unfold m (a, a) a

147enumerateFromThenNum =

148 lmap (\(from, next) -> (from, next - from)) enumerateFromStepNum

149

150-- | Same as 'enumerateFromStepNum' using a stride of 1:

151--

152-- @

153-- >>> enumerateFromNum = lmap (\from -> (from, 1)) Unfold.enumerateFromStepNum

154-- >>> Stream.toList $ Stream.take 6 $ Stream.unfold enumerateFromNum (0.9)

155-- [0.9,1.9,2.9,3.9,4.9,5.9]

156--

157-- @

158--

159-- Also, same as 'enumerateFromThenNum' using a stride of 1 but see the note in

160-- 'enumerateFromThenNum' about the loss of precision:

161--

162-- @

163-- >>> enumerateFromNum = lmap (\from -> (from, from + 1)) Unfold.enumerateFromThenNum

164-- >>> Stream.toList $ Stream.take 6 $ Stream.unfold enumerateFromNum (0.9)

165-- [0.9,1.9,2.9,3.8999999999999995,4.8999999999999995,5.8999999999999995]

166--

167-- @

168--

169-- /Internal/

170--

171{-# INLINE enumerateFromNum #-}

172enumerateFromNum :: (Monad m, Num a) => Unfold m a a

173enumerateFromNum = lmap (\from -> (from, 1)) enumerateFromStepNum

174

175------------------------------------------------------------------------------

176-- Enumeration of Integrals

177------------------------------------------------------------------------------

178

179-- | Can be used to enumerate unbounded integrals. This does not check for

180-- overflow or underflow for bounded integrals.

181--

182-- /Internal/

183{-# INLINE_NORMAL enumerateFromStepIntegral #-}

184enumerateFromStepIntegral :: (Monad m, Integral a) => Unfold m (a, a) a

185enumerateFromStepIntegral = Unfold step inject

186

187 where

188

189 inject (from, stride) = from `seq` stride `seq` return (from, stride)

190

191 {-# INLINE_LATE step #-}

192 step (x, stride) = return $ Yield x $! (x + stride, stride)

193

194-- Enumerate Unbounded Integrals ----------------------------------------------

195{-# INLINE enumerateFromIntegral #-}

196enumerateFromIntegral :: (Monad m, Integral a) => Unfold m a a

197enumerateFromIntegral = lmap (\from -> (from, 1)) enumerateFromStepIntegral

198

199{-# INLINE enumerateFromThenIntegral #-}

200enumerateFromThenIntegral :: (Monad m, Integral a ) => Unfold m (a, a) a

201enumerateFromThenIntegral =

202 lmap (\(from, next) -> (from, next - from)) enumerateFromStepIntegral

203

204{-# INLINE enumerateFromToIntegral #-}

205enumerateFromToIntegral :: (Monad m, Integral a) => Unfold m (a, a) a

206enumerateFromToIntegral =

207 takeWhileMWithInput (\(_, to) b -> return $ b <= to)

208 $ lmap (\(from, _) -> (from, 1)) enumerateFromStepIntegral

209

210{-# INLINE enumerateFromThenToIntegral #-}

211enumerateFromThenToIntegral :: (Monad m, Integral a) => Unfold m (a, a, a) a

212enumerateFromThenToIntegral =

213 takeWhileMWithInput cond $ lmap toFromStep enumerateFromStepIntegral

214

215 where

216

217 toFromStep (from, next, _) = (from, next - from)

218

219 cond (from, next, to) b =

220 return

221 $ if next >= from

222 then b <= to

223 else b >= to

224

225-- Enumerate Bounded Integrals ------------------------------------------------

226{-# INLINE enumerateFromIntegralBounded #-}

227enumerateFromIntegralBounded :: (Monad m, Integral a, Bounded a) =>

228 Unfold m a a

229enumerateFromIntegralBounded = second maxBound enumerateFromToIntegral

230

231{-# INLINE enumerateFromThenIntegralBounded #-}

232enumerateFromThenIntegralBounded :: (Monad m, Integral a, Bounded a ) =>

233 Unfold m (a, a) a

234enumerateFromThenIntegralBounded =

235 takeWhileMWithInput cond $ lmap toFromStep enumerateFromStepIntegral

236

237 where

238

239 toFromStep (from, next) = (from, next - from)

240

241 cond (from, next) b =

242 return

243 $ if next >= from

244 then b <= maxBound

245 else b >= minBound

246

247{-# INLINE enumerateFromToIntegralBounded #-}

248enumerateFromToIntegralBounded :: (Monad m, Integral a, Bounded a) =>

249 Unfold m (a, a) a

250enumerateFromToIntegralBounded =

251 takeWhileMWithInput (\(_, to) b -> return $ b <= to)

252 $ lmap fst enumerateFromIntegralBounded

253

254{-# INLINE enumerateFromThenToIntegralBounded #-}

255enumerateFromThenToIntegralBounded :: (Monad m, Integral a, Bounded a) =>

256 Unfold m (a, a, a) a

257enumerateFromThenToIntegralBounded =

258 takeWhileMWithInput cond $ lmap toFromThen enumerateFromThenIntegralBounded

259

260 where

261

262 toFromThen (from, next, _) = (from, next)

263

264 cond (from, next, to) b =

265 return

266 $ if next >= from

267 then b <= to

268 else b >= to

269

270------------------------------------------------------------------------------

271-- Enumeration of Fractionals

272------------------------------------------------------------------------------

273

274{-# INLINE_NORMAL enumerateFromFractional #-}

275enumerateFromFractional :: (Monad m, Fractional a) => Unfold m a a

276enumerateFromFractional = enumerateFromNum

277

278{-# INLINE_NORMAL enumerateFromThenFractional #-}

279enumerateFromThenFractional :: (Monad m, Fractional a) => Unfold m (a, a) a

280enumerateFromThenFractional = enumerateFromThenNum

281

282-- | Same as 'enumerateFromStepNum' with a step of 1 and enumerating up to the

283-- specified upper limit rounded to the nearest integral value:

284--

285-- @

286-- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromToFractional (0.1, 6.3)

287-- [0.1,1.1,2.1,3.1,4.1,5.1,6.1]

288--

289-- @

290--

291-- /Internal/

292--

293{-# INLINE_NORMAL enumerateFromToFractional #-}

294enumerateFromToFractional :: (Monad m, Fractional a, Ord a) =>

295 Unfold m (a, a) a

296enumerateFromToFractional =

297 takeWhileMWithInput (\(_, to) b -> return $ b <= to + 1 / 2)

298 $ lmap (\(from, _) -> (from, 1)) enumerateFromStepNum

299

300{-# INLINE enumerateFromThenToFractional #-}

301enumerateFromThenToFractional :: (Monad m, Fractional a, Ord a) =>

302 Unfold m (a, a, a) a

303enumerateFromThenToFractional =

304 takeWhileMWithInput cond $ lmap toFromStep enumerateFromStepNum

305

306 where

307

308 toFromStep (from, next, _) = (from, next - from)

309

310 cond (from, next, to) b =

311 let stride = next - from

312 in return

313 $ if next >= from

314 then b <= to + stride / 2

315 else b >= to + stride / 2

316

317-------------------------------------------------------------------------------

318-- Enumeration of Enum types not larger than Int

319-------------------------------------------------------------------------------

320

321-- | Enumerate from given starting Enum value 'from' and to Enum value 'to'

322-- with stride of 1 till to value.

323--

324-- /Internal/

325--

326{-# INLINE enumerateFromToSmall #-}

327enumerateFromToSmall :: (Monad m, Enum a) => Unfold m (a, a) a

328enumerateFromToSmall =

329 fmap toEnum (lmap (bimap fromEnum fromEnum) enumerateFromToIntegral)

330

331-- | Enumerate from given starting Enum value 'from' and then Enum value 'next'

332-- and to Enum value 'to' with stride of (fromEnum next - fromEnum from)

333-- till to value.

334--

335-- /Internal/

336--

337{-# INLINE enumerateFromThenToSmall #-}

338enumerateFromThenToSmall :: (Monad m, Enum a) => Unfold m (a, a, a) a

339enumerateFromThenToSmall =

340 let toInts (x, y, z) = (fromEnum x, fromEnum y, fromEnum z)

341 in fmap toEnum (lmap toInts enumerateFromThenToIntegral)

342

343-------------------------------------------------------------------------------

344-- Bounded Enumeration of Enum types not larger than Int

345-------------------------------------------------------------------------------

346

347-- | Enumerate from given starting Enum value 'from' with stride of 1 till

348-- maxBound

349--

350-- /Internal/

351--

352{-# INLINE enumerateFromSmallBounded #-}

353enumerateFromSmallBounded :: (Monad m, Enum a, Bounded a) => Unfold m a a

354enumerateFromSmallBounded = second maxBound enumerateFromToSmall

355

356-- | Enumerate from given starting Enum value 'from' and next Enum value 'next'

357-- with stride of (fromEnum next - fromEnum from) till maxBound.

358--

359-- /Internal/

360--

361{-# INLINE enumerateFromThenSmallBounded #-}

362enumerateFromThenSmallBounded :: forall m a. (Monad m, Enum a, Bounded a) =>

363 Unfold m (a, a) a

364enumerateFromThenSmallBounded =

365 let adapt (from, next) =

366 let frm = fromEnum from

367 nxt = fromEnum next

368 stride = nxt - frm

369 to = if stride >= 0

370 then fromEnum (maxBound :: a)

371 else fromEnum (minBound :: a)

372 in (frm, nxt, to)

373 in fmap toEnum (lmap adapt enumerateFromThenToIntegral)

374

375-------------------------------------------------------------------------------

376-- Enumerable type class

377-------------------------------------------------------------------------------

378

379-- | Types that can be enumerated as a stream. The operations in this type

380-- class are equivalent to those in the 'Enum' type class, except that these

381-- generate a stream instead of a list. Use the functions in

382-- "Streamly.Internal.Data.Unfold.Enumeration" module to define new instances.

383--

384-- /Pre-release/

385class Enum a => Enumerable a where

386

387 -- | Unfolds @from@ generating a stream starting with the element

388 -- @from@, enumerating up to 'maxBound' when the type is 'Bounded' or

389 -- generating an infinite stream when the type is not 'Bounded'.

390 --

391 -- >>> import qualified Streamly.Prelude as Stream

392 -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold

393 --

394 -- @

395 -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFrom (0 :: Int)

396 -- [0,1,2,3]

397 --

398 -- @

399 --

400 -- For 'Fractional' types, enumeration is numerically stable. However, no

401 -- overflow or underflow checks are performed.

402 --

403 -- @

404 -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFrom 1.1

405 -- [1.1,2.1,3.1,4.1]

406 --

407 -- @

408 --

409 -- /Pre-release/

410 --

411 enumerateFrom :: Monad m => Unfold m a a

412

413 -- | Unfolds @(from, to)@ generating a finite stream starting with the element

414 -- @from@, enumerating the type up to the value @to@. If @to@ is smaller than

415 -- @from@ then an empty stream is returned.

416 --

417 -- >>> import qualified Streamly.Prelude as Stream

418 -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold

419 --

420 -- @

421 -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromTo (0, 4)

422 -- [0,1,2,3,4]

423 --

424 -- @

425 --

426 -- For 'Fractional' types, the last element is equal to the specified @to@

427 -- value after rounding to the nearest integral value.

428 --

429 -- @

430 -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromTo (1.1, 4)

431 -- [1.1,2.1,3.1,4.1]

432 --

433 -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromTo (1.1, 4.6)

434 -- [1.1,2.1,3.1,4.1,5.1]

435 --

436 -- @

437 --

438 -- /Pre-release/

439 enumerateFromTo :: Monad m => Unfold m (a, a) a

440

441 -- | Unfolds @(from, then)@ generating a stream whose first element is

442 -- @from@ and the successive elements are in increments of @then@. Enumeration

443 -- can occur downwards or upwards depending on whether @then@ comes before or

444 -- after @from@. For 'Bounded' types the stream ends when 'maxBound' is

445 -- reached, for unbounded types it keeps enumerating infinitely.

446 --

447 -- >>> import qualified Streamly.Prelude as Stream

448 -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold

449 --

450 -- @

451 -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFromThen (0, 2)

452 -- [0,2,4,6]

453 --

454 -- >>> Stream.toList $ Stream.take 4 $ Stream.unfold Unfold.enumerateFromThen (0,(-2))

455 -- [0,-2,-4,-6]

456 --

457 -- @

458 --

459 -- /Pre-release/

460 enumerateFromThen :: Monad m => Unfold m (a, a) a

461

462 -- | Unfolds @(from, then, to)@ generating a finite stream whose first element

463 -- is @from@ and the successive elements are in increments of @then@ up to

464 -- @to@. Enumeration can occur downwards or upwards depending on whether @then@

465 -- comes before or after @from@.

466 --

467 -- >>> import qualified Streamly.Prelude as Stream

468 -- >>> import qualified Streamly.Internal.Data.Unfold as Unfold

469 --

470 -- @

471 -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromThenTo (0, 2, 6)

472 -- [0,2,4,6]

473 --

474 -- >>> Stream.toList $ Stream.unfold Unfold.enumerateFromThenTo (0, (-2), (-6))

475 -- [0,-2,-4,-6]

476 --

477 -- @

478 --

479 -- /Pre-release/

480 enumerateFromThenTo :: Monad m => Unfold m (a, a, a) a

481

482-------------------------------------------------------------------------------

483-- Enumerable Instances

484-------------------------------------------------------------------------------

485--

486-- For Enum types smaller than or equal to Int size.

487#define ENUMERABLE_BOUNDED_SMALL(SMALL_TYPE) \

488instance Enumerable SMALL_TYPE where { \

489 {-# INLINE enumerateFrom #-}; \

490 enumerateFrom = enumerateFromSmallBounded; \

491 {-# INLINE enumerateFromThen #-}; \

492 enumerateFromThen = enumerateFromThenSmallBounded; \

493 {-# INLINE enumerateFromTo #-}; \

494 enumerateFromTo = enumerateFromToSmall; \

495 {-# INLINE enumerateFromThenTo #-}; \

496 enumerateFromThenTo = enumerateFromThenToSmall }

497

498ENUMERABLE_BOUNDED_SMALL(())

499ENUMERABLE_BOUNDED_SMALL(Bool)

500ENUMERABLE_BOUNDED_SMALL(Ordering)

501ENUMERABLE_BOUNDED_SMALL(Char)

502

503-- For bounded Integral Enum types, may be larger than Int.

504#define ENUMERABLE_BOUNDED_INTEGRAL(INTEGRAL_TYPE) \

505instance Enumerable INTEGRAL_TYPE where { \

506 {-# INLINE enumerateFrom #-}; \

507 enumerateFrom = enumerateFromIntegralBounded; \

508 {-# INLINE enumerateFromThen #-}; \

509 enumerateFromThen = enumerateFromThenIntegralBounded; \

510 {-# INLINE enumerateFromTo #-}; \

511 enumerateFromTo = enumerateFromToIntegralBounded; \

512 {-# INLINE enumerateFromThenTo #-}; \

513 enumerateFromThenTo = enumerateFromThenToIntegralBounded }

514

515ENUMERABLE_BOUNDED_INTEGRAL(Int)

516ENUMERABLE_BOUNDED_INTEGRAL(Int8)

517ENUMERABLE_BOUNDED_INTEGRAL(Int16)

518ENUMERABLE_BOUNDED_INTEGRAL(Int32)

519ENUMERABLE_BOUNDED_INTEGRAL(Int64)

520ENUMERABLE_BOUNDED_INTEGRAL(Word)

521ENUMERABLE_BOUNDED_INTEGRAL(Word8)

522ENUMERABLE_BOUNDED_INTEGRAL(Word16)

523ENUMERABLE_BOUNDED_INTEGRAL(Word32)

524ENUMERABLE_BOUNDED_INTEGRAL(Word64)

525

526-- For unbounded Integral Enum types.

527#define ENUMERABLE_UNBOUNDED_INTEGRAL(INTEGRAL_TYPE) \

528instance Enumerable INTEGRAL_TYPE where { \

529 {-# INLINE enumerateFrom #-}; \

530 enumerateFrom = enumerateFromIntegral; \

531 {-# INLINE enumerateFromThen #-}; \

532 enumerateFromThen = enumerateFromThenIntegral; \

533 {-# INLINE enumerateFromTo #-}; \

534 enumerateFromTo = enumerateFromToIntegral; \

535 {-# INLINE enumerateFromThenTo #-}; \

536 enumerateFromThenTo = enumerateFromThenToIntegral }

537

538ENUMERABLE_UNBOUNDED_INTEGRAL(Integer)

539ENUMERABLE_UNBOUNDED_INTEGRAL(Natural)

540

541#define ENUMERABLE_FRACTIONAL(FRACTIONAL_TYPE,CONSTRAINT) \

542instance (CONSTRAINT) => Enumerable FRACTIONAL_TYPE where { \

543 {-# INLINE enumerateFrom #-}; \

544 enumerateFrom = enumerateFromFractional; \

545 {-# INLINE enumerateFromThen #-}; \

546 enumerateFromThen = enumerateFromThenFractional; \

547 {-# INLINE enumerateFromTo #-}; \

548 enumerateFromTo = enumerateFromToFractional; \

549 {-# INLINE enumerateFromThenTo #-}; \

550 enumerateFromThenTo = enumerateFromThenToFractional }

551

552ENUMERABLE_FRACTIONAL(Float,)

553ENUMERABLE_FRACTIONAL(Double,)

554ENUMERABLE_FRACTIONAL((Fixed a),HasResolution a)

555ENUMERABLE_FRACTIONAL((Ratio a),Integral a)

556

557instance Enumerable a => Enumerable (Identity a) where

558 {-# INLINE enumerateFrom #-}

559 enumerateFrom =

560 map Identity $ lmap runIdentity enumerateFrom

561 {-# INLINE enumerateFromThen #-}

562 enumerateFromThen =

563 map Identity $ lmap (bimap runIdentity runIdentity) enumerateFromThen

564 {-# INLINE enumerateFromTo #-}

565 enumerateFromTo =

566 map Identity $ lmap (bimap runIdentity runIdentity) enumerateFromThen

567 {-# INLINE enumerateFromThenTo #-}

568 enumerateFromThenTo =

569 map Identity $

570 lmap

571 (\(from, next, to) ->

572 (runIdentity from, runIdentity next, runIdentity to))

573 enumerateFromThenTo