src/Streamly/Internal/Data/Stream/IsStream/Enumeration.hs · streamly

git.iko.soy/streamly

on bc4dba9 help sign in

streamly/src/Streamly/Internal/Data/Stream/IsStream/Enumeration.hs

574 lines 21 kB history raw

1-- |

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

4--

5-- License : BSD3

6-- Maintainer : streamly@composewell.com

7-- Stability : experimental

8-- Portability : GHC

9--

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

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

12-- instances instead.

13--

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

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

16-- perations in the 'Enum' type class, the only difference is that they produce

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

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

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

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

21-- can be used directly.

23module Streamly.Internal.Data.Stream.IsStream.Enumeration

24 (

25 Enumerable (..)

27 -- ** Enumerating 'Bounded' 'Enum' Types

28 , enumerate

29 , enumerateTo

30 , enumerateFromBounded

32 -- ** Enumerating 'Enum' Types not larger than 'Int'

33 , enumerateFromToSmall

34 , enumerateFromThenToSmall

35 , enumerateFromThenSmallBounded

37 -- ** Enumerating 'Bounded' 'Integral' Types

38 , enumerateFromIntegral

39 , enumerateFromThenIntegral

41 -- ** Enumerating 'Integral' Types

42 , enumerateFromToIntegral

43 , enumerateFromThenToIntegral

45 -- ** Enumerating unbounded 'Integral' Types

46 , enumerateFromStepIntegral

48 -- ** Enumerating 'Fractional' Types

49 , enumerateFromFractional

50 , enumerateFromToFractional

51 , enumerateFromThenFractional

52 , enumerateFromThenToFractional

53 )

54where

56import Data.Fixed

57import Data.Int

58import Data.Ratio

59import Data.Word

60import Numeric.Natural

61import Data.Functor.Identity (Identity(..))

63import Streamly.Internal.Data.Stream.StreamD.Type (fromStreamD)

64import Streamly.Internal.Data.Stream.StreamK.Type (IsStream(..))

66import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D

67import qualified Streamly.Internal.Data.Stream.Serial as Serial (map)

69-- $setup

70-- >>> import Streamly.Prelude as Stream

71-- >>> import Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream

73-------------------------------------------------------------------------------

74-- Enumeration of Integral types

75-------------------------------------------------------------------------------

76--

77-- | @enumerateFromStepIntegral from step@ generates an infinite stream whose

78-- first element is @from@ and the successive elements are in increments of

79-- @step@.

80--

81-- CAUTION: This function is not safe for finite integral types. It does not

82-- check for overflow, underflow or bounds.

83--

84-- @

85-- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromStepIntegral 0 2

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

87--

88-- >>> Stream.toList $ Stream.take 3 $ Stream.enumerateFromStepIntegral 0 (-2)

89-- [0,-2,-4]

90--

91-- @

92--

93-- @since 0.6.0

94{-# INLINE enumerateFromStepIntegral #-}

95enumerateFromStepIntegral

96 :: (IsStream t, Monad m, Integral a)

97 => a -> a -> t m a

98enumerateFromStepIntegral from stride =

99 fromStreamD $ D.enumerateFromStepIntegral from stride

100

101-- | Enumerate an 'Integral' type. @enumerateFromIntegral from@ generates a

102-- stream whose first element is @from@ and the successive elements are in

103-- increments of @1@. The stream is bounded by the size of the 'Integral' type.

104--

105-- @

106-- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromIntegral (0 :: Int)

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

108--

109-- @

110--

111-- @since 0.6.0

112{-# INLINE enumerateFromIntegral #-}

113enumerateFromIntegral

114 :: (IsStream t, Monad m, Integral a, Bounded a)

115 => a -> t m a

116enumerateFromIntegral from = fromStreamD $ D.enumerateFromIntegral from

117

118-- | Enumerate an 'Integral' type in steps. @enumerateFromThenIntegral from

119-- then@ generates a stream whose first element is @from@, the second element

120-- is @then@ and the successive elements are in increments of @then - from@.

121-- The stream is bounded by the size of the 'Integral' type.

122--

123-- @

124-- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromThenIntegral (0 :: Int) 2

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

126--

127-- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromThenIntegral (0 :: Int) (-2)

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

129--

130-- @

131--

132-- @since 0.6.0

133{-# INLINE enumerateFromThenIntegral #-}

134enumerateFromThenIntegral

135 :: (IsStream t, Monad m, Integral a, Bounded a)

136 => a -> a -> t m a

137enumerateFromThenIntegral from next =

138 fromStreamD $ D.enumerateFromThenIntegral from next

139

140-- | Enumerate an 'Integral' type up to a given limit.

141-- @enumerateFromToIntegral from to@ generates a finite stream whose first

142-- element is @from@ and successive elements are in increments of @1@ up to

143-- @to@.

144--

145-- @

146-- >>> Stream.toList $ Stream.enumerateFromToIntegral 0 4

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

148--

149-- @

150--

151-- @since 0.6.0

152{-# INLINE enumerateFromToIntegral #-}

153enumerateFromToIntegral :: (IsStream t, Monad m, Integral a) => a -> a -> t m a

154enumerateFromToIntegral from to =

155 fromStreamD $ D.enumerateFromToIntegral from to

156

157-- | Enumerate an 'Integral' type in steps up to a given limit.

158-- @enumerateFromThenToIntegral from then to@ generates a finite stream whose

159-- first element is @from@, the second element is @then@ and the successive

160-- elements are in increments of @then - from@ up to @to@.

161--

162-- @

163-- >>> Stream.toList $ Stream.enumerateFromThenToIntegral 0 2 6

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

165--

166-- >>> Stream.toList $ Stream.enumerateFromThenToIntegral 0 (-2) (-6)

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

168--

169-- @

170--

171-- @since 0.6.0

172{-# INLINE enumerateFromThenToIntegral #-}

173enumerateFromThenToIntegral

174 :: (IsStream t, Monad m, Integral a)

175 => a -> a -> a -> t m a

176enumerateFromThenToIntegral from next to =

177 fromStreamD $ D.enumerateFromThenToIntegral from next to

178

179-------------------------------------------------------------------------------

180-- Enumeration of Fractional types

181-------------------------------------------------------------------------------

182--

183-- Even though the underlying implementation of enumerateFromFractional and

184-- enumerateFromThenFractional works for any 'Num' we have restricted these to

185-- 'Fractional' because these do not perform any bounds check, in contrast to

186-- integral versions and are therefore not equivalent substitutes for those.

187--

188-- | Numerically stable enumeration from a 'Fractional' number in steps of size

189-- @1@. @enumerateFromFractional from@ generates a stream whose first element

190-- is @from@ and the successive elements are in increments of @1@. No overflow

191-- or underflow checks are performed.

192--

193-- This is the equivalent to 'enumFrom' for 'Fractional' types. For example:

194--

195-- @

196-- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromFractional 1.1

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

198--

199-- @

200--

201--

202-- @since 0.6.0

203{-# INLINE enumerateFromFractional #-}

204enumerateFromFractional :: (IsStream t, Monad m, Fractional a) => a -> t m a

205enumerateFromFractional from = fromStreamD $ D.numFrom from

206

207-- | Numerically stable enumeration from a 'Fractional' number in steps.

208-- @enumerateFromThenFractional from then@ generates a stream whose first

209-- element is @from@, the second element is @then@ and the successive elements

210-- are in increments of @then - from@. No overflow or underflow checks are

211-- performed.

212--

213-- This is the equivalent of 'enumFromThen' for 'Fractional' types. For

214-- example:

215--

216-- @

217-- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromThenFractional 1.1 2.1

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

219--

220-- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromThenFractional 1.1 (-2.1)

221-- [1.1,-2.1,-5.300000000000001,-8.500000000000002]

222--

223-- @

224--

225-- @since 0.6.0

226{-# INLINE enumerateFromThenFractional #-}

227enumerateFromThenFractional

228 :: (IsStream t, Monad m, Fractional a)

229 => a -> a -> t m a

230enumerateFromThenFractional from next = fromStreamD $ D.numFromThen from next

231

232-- | Numerically stable enumeration from a 'Fractional' number to a given

233-- limit. @enumerateFromToFractional from to@ generates a finite stream whose

234-- first element is @from@ and successive elements are in increments of @1@ up

235-- to @to@.

236--

237-- This is the equivalent of 'enumFromTo' for 'Fractional' types. For

238-- example:

239--

240-- @

241-- >>> Stream.toList $ Stream.enumerateFromToFractional 1.1 4

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

243--

244-- >>> Stream.toList $ Stream.enumerateFromToFractional 1.1 4.6

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

246--

247-- @

248--

249-- Notice that the last element is equal to the specified @to@ value after

250-- rounding to the nearest integer.

251--

252-- @since 0.6.0

253{-# INLINE enumerateFromToFractional #-}

254enumerateFromToFractional

255 :: (IsStream t, Monad m, Fractional a, Ord a)

256 => a -> a -> t m a

257enumerateFromToFractional from to =

258 fromStreamD $ D.enumerateFromToFractional from to

259

260-- | Numerically stable enumeration from a 'Fractional' number in steps up to a

261-- given limit. @enumerateFromThenToFractional from then to@ generates a

262-- finite stream whose first element is @from@, the second element is @then@

263-- and the successive elements are in increments of @then - from@ up to @to@.

264--

265-- This is the equivalent of 'enumFromThenTo' for 'Fractional' types. For

266-- example:

267--

268-- @

269-- >>> Stream.toList $ Stream.enumerateFromThenToFractional 0.1 2 6

270-- [0.1,2.0,3.9,5.799999999999999]

271--

272-- >>> Stream.toList $ Stream.enumerateFromThenToFractional 0.1 (-2) (-6)

273-- [0.1,-2.0,-4.1000000000000005,-6.200000000000001]

274--

275-- @

276--

277--

278-- @since 0.6.0

279{-# INLINE enumerateFromThenToFractional #-}

280enumerateFromThenToFractional

281 :: (IsStream t, Monad m, Fractional a, Ord a)

282 => a -> a -> a -> t m a

283enumerateFromThenToFractional from next to =

284 fromStreamD $ D.enumerateFromThenToFractional from next to

285

286-------------------------------------------------------------------------------

287-- Enumeration of Enum types not larger than Int

288-------------------------------------------------------------------------------

289--

290-- | 'enumerateFromTo' for 'Enum' types not larger than 'Int'.

291--

292-- @since 0.6.0

293{-# INLINE enumerateFromToSmall #-}

294enumerateFromToSmall :: (IsStream t, Monad m, Enum a) => a -> a -> t m a

295enumerateFromToSmall from to = Serial.map toEnum $

296 enumerateFromToIntegral (fromEnum from) (fromEnum to)

297

298-- | 'enumerateFromThenTo' for 'Enum' types not larger than 'Int'.

299--

300-- @since 0.6.0

301{-# INLINE enumerateFromThenToSmall #-}

302enumerateFromThenToSmall :: (IsStream t, Monad m, Enum a)

303 => a -> a -> a -> t m a

304enumerateFromThenToSmall from next to = Serial.map toEnum $

305 enumerateFromThenToIntegral (fromEnum from) (fromEnum next) (fromEnum to)

306

307-- | 'enumerateFromThen' for 'Enum' types not larger than 'Int'.

308--

309-- Note: We convert the 'Enum' to 'Int' and enumerate the 'Int'. If a

310-- type is bounded but does not have a 'Bounded' instance then we can go on

311-- enumerating it beyond the legal values of the type, resulting in the failure

312-- of 'toEnum' when converting back to 'Enum'. Therefore we require a 'Bounded'

313-- instance for this function to be safely used.

314--

315-- @since 0.6.0

316{-# INLINE enumerateFromThenSmallBounded #-}

317enumerateFromThenSmallBounded :: (IsStream t, Monad m, Enumerable a, Bounded a)

318 => a -> a -> t m a

319enumerateFromThenSmallBounded from next =

320 if fromEnum next >= fromEnum from

321 then enumerateFromThenTo from next maxBound

322 else enumerateFromThenTo from next minBound

323

324-------------------------------------------------------------------------------

325-- Enumerable type class

326-------------------------------------------------------------------------------

327--

328-- NOTE: We would like to rewrite calls to fromList [1..] etc. to stream

329-- enumerations like this:

330--

331-- {-# RULES "fromList enumFrom" [1]

332-- forall (a :: Int). D.fromList (enumFrom a) = D.enumerateFromIntegral a #-}

333--

334-- But this does not work because enumFrom is a class method and GHC rewrites

335-- it quickly, so we do not get a chance to have our rule fired.

336

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

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

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

340-- "Streamly.Internal.Data.Stream.Enumeration" module to define new instances.

341--

342-- @since 0.6.0

343class Enum a => Enumerable a where

344 -- | @enumerateFrom from@ generates a stream starting with the element

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

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

347 --

348 -- @

349 -- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFrom (0 :: Int)

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

351 --

352 -- @

353 --

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

355 -- overflow or underflow checks are performed.

356 --

357 -- @

358 -- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFrom 1.1

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

360 --

361 -- @

362 --

363 -- @since 0.6.0

364 enumerateFrom :: (IsStream t, Monad m) => a -> t m a

365

366 -- | Generate a finite stream starting with the element @from@, enumerating

367 -- the type up to the value @to@. If @to@ is smaller than @from@ then an

368 -- empty stream is returned.

369 --

370 -- @

371 -- >>> Stream.toList $ Stream.enumerateFromTo 0 4

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

373 --

374 -- @

375 --

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

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

378 --

379 -- @

380 -- >>> Stream.toList $ Stream.enumerateFromTo 1.1 4

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

382 --

383 -- >>> Stream.toList $ Stream.enumerateFromTo 1.1 4.6

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

385 --

386 -- @

387 --

388 -- @since 0.6.0

389 enumerateFromTo :: (IsStream t, Monad m) => a -> a -> t m a

390

391 -- | @enumerateFromThen from then@ generates a stream whose first element

392 -- is @from@, the second element is @then@ and the successive elements are

393 -- in increments of @then - from@. Enumeration can occur downwards or

394 -- upwards depending on whether @then@ comes before or after @from@. For

395 -- 'Bounded' types the stream ends when 'maxBound' is reached, for

396 -- unbounded types it keeps enumerating infinitely.

397 --

398 -- @

399 -- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromThen 0 2

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

401 --

402 -- >>> Stream.toList $ Stream.take 4 $ Stream.enumerateFromThen 0 (-2)

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

404 --

405 -- @

406 --

407 -- @since 0.6.0

408 enumerateFromThen :: (IsStream t, Monad m) => a -> a -> t m a

409

410 -- | @enumerateFromThenTo from then to@ generates a finite stream whose

411 -- first element is @from@, the second element is @then@ and the successive

412 -- elements are in increments of @then - from@ up to @to@. Enumeration can

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

414 -- after @from@.

415 --

416 -- @

417 -- >>> Stream.toList $ Stream.enumerateFromThenTo 0 2 6

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

419 --

420 -- >>> Stream.toList $ Stream.enumerateFromThenTo 0 (-2) (-6)

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

422 --

423 -- @

424 --

425 -- @since 0.6.0

426 enumerateFromThenTo :: (IsStream t, Monad m) => a -> a -> a -> t m a

427

428-- MAYBE: Sometimes it is more convenient to know the count rather then the

429-- ending or starting element. For those cases we can define the folllowing

430-- APIs. All of these will work only for bounded types if we represent the

431-- count by Int.

432--

433-- enumerateN

434-- enumerateFromN

435-- enumerateToN

436-- enumerateFromStep

437-- enumerateFromStepN

438

439-------------------------------------------------------------------------------

440-- Convenient functions for bounded types

441-------------------------------------------------------------------------------

442--

443-- |

444-- > enumerate = enumerateFrom minBound

445--

446-- Enumerate a 'Bounded' type from its 'minBound' to 'maxBound'

447--

448-- @since 0.6.0

449{-# INLINE enumerate #-}

450enumerate :: (IsStream t, Monad m, Bounded a, Enumerable a) => t m a

451enumerate = enumerateFrom minBound

452

453-- |

454-- > enumerateTo = enumerateFromTo minBound

455--

456-- Enumerate a 'Bounded' type from its 'minBound' to specified value.

457--

458-- @since 0.6.0

459{-# INLINE enumerateTo #-}

460enumerateTo :: (IsStream t, Monad m, Bounded a, Enumerable a) => a -> t m a

461enumerateTo = enumerateFromTo minBound

462

463-- |

464-- > enumerateFromBounded = enumerateFromTo from maxBound

465--

466-- 'enumerateFrom' for 'Bounded' 'Enum' types.

467--

468-- @since 0.6.0

469{-# INLINE enumerateFromBounded #-}

470enumerateFromBounded :: (IsStream t, Monad m, Enumerable a, Bounded a)

471 => a -> t m a

472enumerateFromBounded from = enumerateFromTo from maxBound

473

474-------------------------------------------------------------------------------

475-- Enumerable Instances

476-------------------------------------------------------------------------------

477--

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

479#define ENUMERABLE_BOUNDED_SMALL(SMALL_TYPE) \

480instance Enumerable SMALL_TYPE where { \

481 {-# INLINE enumerateFrom #-}; \

482 enumerateFrom = enumerateFromBounded; \

483 {-# INLINE enumerateFromThen #-}; \

484 enumerateFromThen = enumerateFromThenSmallBounded; \

485 {-# INLINE enumerateFromTo #-}; \

486 enumerateFromTo = enumerateFromToSmall; \

487 {-# INLINE enumerateFromThenTo #-}; \

488 enumerateFromThenTo = enumerateFromThenToSmall }

489

490

491ENUMERABLE_BOUNDED_SMALL(())

492ENUMERABLE_BOUNDED_SMALL(Bool)

493ENUMERABLE_BOUNDED_SMALL(Ordering)

494ENUMERABLE_BOUNDED_SMALL(Char)

495

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

497#define ENUMERABLE_BOUNDED_INTEGRAL(INTEGRAL_TYPE) \

498instance Enumerable INTEGRAL_TYPE where { \

499 {-# INLINE enumerateFrom #-}; \

500 enumerateFrom = enumerateFromIntegral; \

501 {-# INLINE enumerateFromThen #-}; \

502 enumerateFromThen = enumerateFromThenIntegral; \

503 {-# INLINE enumerateFromTo #-}; \

504 enumerateFromTo = enumerateFromToIntegral; \

505 {-# INLINE enumerateFromThenTo #-}; \

506 enumerateFromThenTo = enumerateFromThenToIntegral }

507

508ENUMERABLE_BOUNDED_INTEGRAL(Int)

509ENUMERABLE_BOUNDED_INTEGRAL(Int8)

510ENUMERABLE_BOUNDED_INTEGRAL(Int16)

511ENUMERABLE_BOUNDED_INTEGRAL(Int32)

512ENUMERABLE_BOUNDED_INTEGRAL(Int64)

513ENUMERABLE_BOUNDED_INTEGRAL(Word)

514ENUMERABLE_BOUNDED_INTEGRAL(Word8)

515ENUMERABLE_BOUNDED_INTEGRAL(Word16)

516ENUMERABLE_BOUNDED_INTEGRAL(Word32)

517ENUMERABLE_BOUNDED_INTEGRAL(Word64)

518

519-- For unbounded Integral Enum types.

520#define ENUMERABLE_UNBOUNDED_INTEGRAL(INTEGRAL_TYPE) \

521instance Enumerable INTEGRAL_TYPE where { \

522 {-# INLINE enumerateFrom #-}; \

523 enumerateFrom from = enumerateFromStepIntegral from 1; \

524 {-# INLINE enumerateFromThen #-}; \

525 enumerateFromThen from next = \

526 enumerateFromStepIntegral from (next - from); \

527 {-# INLINE enumerateFromTo #-}; \

528 enumerateFromTo = enumerateFromToIntegral; \

529 {-# INLINE enumerateFromThenTo #-}; \

530 enumerateFromThenTo = enumerateFromThenToIntegral }

531

532ENUMERABLE_UNBOUNDED_INTEGRAL(Integer)

533ENUMERABLE_UNBOUNDED_INTEGRAL(Natural)

534

535#define ENUMERABLE_FRACTIONAL(FRACTIONAL_TYPE,CONSTRAINT) \

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

537 {-# INLINE enumerateFrom #-}; \

538 enumerateFrom = enumerateFromFractional; \

539 {-# INLINE enumerateFromThen #-}; \

540 enumerateFromThen = enumerateFromThenFractional; \

541 {-# INLINE enumerateFromTo #-}; \

542 enumerateFromTo = enumerateFromToFractional; \

543 {-# INLINE enumerateFromThenTo #-}; \

544 enumerateFromThenTo = enumerateFromThenToFractional }

545

546ENUMERABLE_FRACTIONAL(Float,)

547ENUMERABLE_FRACTIONAL(Double,)

548ENUMERABLE_FRACTIONAL((Fixed a),HasResolution a)

549ENUMERABLE_FRACTIONAL((Ratio a),Integral a)

550

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

552 {-# INLINE enumerateFrom #-}

553 enumerateFrom (Identity from) = Serial.map Identity $

554 enumerateFrom from

555 {-# INLINE enumerateFromThen #-}

556 enumerateFromThen (Identity from) (Identity next) = Serial.map Identity $

557 enumerateFromThen from next

558 {-# INLINE enumerateFromTo #-}

559 enumerateFromTo (Identity from) (Identity to) = Serial.map Identity $

560 enumerateFromTo from to

561 {-# INLINE enumerateFromThenTo #-}

562 enumerateFromThenTo (Identity from) (Identity next) (Identity to) =

563 Serial.map Identity $ enumerateFromThenTo from next to

564

565-- TODO

566{-

567instance Enumerable a => Enumerable (Last a)

568instance Enumerable a => Enumerable (First a)

569instance Enumerable a => Enumerable (Max a)

570instance Enumerable a => Enumerable (Min a)

571instance Enumerable a => Enumerable (Const a b)

572instance Enumerable (f a) => Enumerable (Alt f a)

573instance Enumerable (f a) => Enumerable (Ap f a)

574-}