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

git.iko.soy/streamly

on c3dcabc help sign in

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

546 lines 20 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.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 (fromStreamD)

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

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

67import qualified Streamly.Internal.Data.Stream.Serial as Serial

69-------------------------------------------------------------------------------

70-- Enumeration of Integral types

71-------------------------------------------------------------------------------

72--

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

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

75-- @step@.

76--

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

78-- check for overflow, underflow or bounds.

79--

80-- @

81-- > S.toList $ S.take 4 $ S.enumerateFromStepIntegral 0 2

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

83-- > S.toList $ S.take 3 $ S.enumerateFromStepIntegral 0 (-2)

84-- [0,-2,-4]

85-- @

86--

87-- @since 0.6.0

88{-# INLINE enumerateFromStepIntegral #-}

89enumerateFromStepIntegral

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

91 => a -> a -> t m a

92enumerateFromStepIntegral from stride =

93 fromStreamD $ D.enumerateFromStepIntegral from stride

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

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

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

98--

99-- @

100-- > S.toList $ S.take 4 $ S.enumerateFromIntegral (0 :: Int)

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

102-- @

103--

104-- @since 0.6.0

105{-# INLINE enumerateFromIntegral #-}

106enumerateFromIntegral

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

108 => a -> t m a

109enumerateFromIntegral from = fromStreamD $ D.enumerateFromIntegral from

110

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

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

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

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

115--

116-- @

117-- > S.toList $ S.take 4 $ S.enumerateFromThenIntegral (0 :: Int) 2

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

119-- > S.toList $ S.take 4 $ S.enumerateFromThenIntegral (0 :: Int) (-2)

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

121-- @

122--

123-- @since 0.6.0

124{-# INLINE enumerateFromThenIntegral #-}

125enumerateFromThenIntegral

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

127 => a -> a -> t m a

128enumerateFromThenIntegral from next =

129 fromStreamD $ D.enumerateFromThenIntegral from next

130

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

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

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

134-- @to@.

135--

136-- @

137-- > S.toList $ S.enumerateFromToIntegral 0 4

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

139-- @

140--

141-- @since 0.6.0

142{-# INLINE enumerateFromToIntegral #-}

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

144enumerateFromToIntegral from to =

145 fromStreamD $ D.enumerateFromToIntegral from to

146

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

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

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

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

151--

152-- @

153-- > S.toList $ S.enumerateFromThenToIntegral 0 2 6

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

155-- > S.toList $ S.enumerateFromThenToIntegral 0 (-2) (-6)

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

157-- @

158--

159-- @since 0.6.0

160{-# INLINE enumerateFromThenToIntegral #-}

161enumerateFromThenToIntegral

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

163 => a -> a -> a -> t m a

164enumerateFromThenToIntegral from next to =

165 fromStreamD $ D.enumerateFromThenToIntegral from next to

166

167-------------------------------------------------------------------------------

168-- Enumeration of Fractional types

169-------------------------------------------------------------------------------

170--

171-- Even though the underlying implementation of enumerateFromFractional and

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

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

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

175--

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

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

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

179-- or underflow checks are performed.

180--

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

182--

183-- @

184-- > S.toList $ S.take 4 $ S.enumerateFromFractional 1.1

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

186-- @

187--

188--

189-- @since 0.6.0

190{-# INLINE enumerateFromFractional #-}

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

192enumerateFromFractional from = fromStreamD $ D.numFrom from

193

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

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

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

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

198-- performed.

199--

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

201-- example:

202--

203-- @

204-- > S.toList $ S.take 4 $ S.enumerateFromThenFractional 1.1 2.1

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

206-- > S.toList $ S.take 4 $ S.enumerateFromThenFractional 1.1 (-2.1)

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

208-- @

209--

210-- @since 0.6.0

211{-# INLINE enumerateFromThenFractional #-}

212enumerateFromThenFractional

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

214 => a -> a -> t m a

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

216

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

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

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

220-- to @to@.

221--

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

223-- example:

224--

225-- @

226-- > S.toList $ S.enumerateFromToFractional 1.1 4

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

228-- > S.toList $ S.enumerateFromToFractional 1.1 4.6

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

230-- @

231--

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

233-- rounding to the nearest integer.

234--

235-- @since 0.6.0

236{-# INLINE enumerateFromToFractional #-}

237enumerateFromToFractional

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

239 => a -> a -> t m a

240enumerateFromToFractional from to =

241 fromStreamD $ D.enumerateFromToFractional from to

242

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

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

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

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

247--

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

249-- example:

250--

251-- @

252-- > S.toList $ S.enumerateFromThenToFractional 0.1 2 6

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

254-- > S.toList $ S.enumerateFromThenToFractional 0.1 (-2) (-6)

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

256-- @

257--

258--

259-- @since 0.6.0

260{-# INLINE enumerateFromThenToFractional #-}

261enumerateFromThenToFractional

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

263 => a -> a -> a -> t m a

264enumerateFromThenToFractional from next to =

265 fromStreamD $ D.enumerateFromThenToFractional from next to

266

267-------------------------------------------------------------------------------

268-- Enumeration of Enum types not larger than Int

269-------------------------------------------------------------------------------

270--

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

272--

273-- @since 0.6.0

274{-# INLINE enumerateFromToSmall #-}

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

276enumerateFromToSmall from to = Serial.map toEnum $

277 enumerateFromToIntegral (fromEnum from) (fromEnum to)

278

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

280--

281-- @since 0.6.0

282{-# INLINE enumerateFromThenToSmall #-}

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

284 => a -> a -> a -> t m a

285enumerateFromThenToSmall from next to = Serial.map toEnum $

286 enumerateFromThenToIntegral (fromEnum from) (fromEnum next) (fromEnum to)

287

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

289--

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

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

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

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

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

295--

296-- @since 0.6.0

297{-# INLINE enumerateFromThenSmallBounded #-}

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

299 => a -> a -> t m a

300enumerateFromThenSmallBounded from next =

301 case fromEnum next >= fromEnum from of

302 True -> enumerateFromThenTo from next maxBound

303 False -> enumerateFromThenTo from next minBound

304

305-------------------------------------------------------------------------------

306-- Enumerable type class

307-------------------------------------------------------------------------------

308--

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

310-- enumerations like this:

311--

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

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

314--

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

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

317

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

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

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

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

322--

323-- @since 0.6.0

324class Enum a => Enumerable a where

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

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

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

328 --

329 -- @

330 -- > S.toList $ S.take 4 $ S.enumerateFrom (0 :: Int)

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

332 -- @

333 --

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

335 -- overflow or underflow checks are performed.

336 --

337 -- @

338 -- > S.toList $ S.take 4 $ S.enumerateFrom 1.1

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

340 -- @

341 --

342 -- @since 0.6.0

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

344

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

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

347 -- empty stream is returned.

348 --

349 -- @

350 -- > S.toList $ S.enumerateFromTo 0 4

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

352 -- @

353 --

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

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

356 --

357 -- @

358 -- > S.toList $ S.enumerateFromTo 1.1 4

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

360 -- > S.toList $ S.enumerateFromTo 1.1 4.6

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

362 -- @

363 --

364 -- @since 0.6.0

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

366

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

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

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

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

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

372 -- unbounded types it keeps enumerating infinitely.

373 --

374 -- @

375 -- > S.toList $ S.take 4 $ S.enumerateFromThen 0 2

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

377 -- > S.toList $ S.take 4 $ S.enumerateFromThen 0 (-2)

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

379 -- @

380 --

381 -- @since 0.6.0

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

383

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

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

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

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

388 -- after @from@.

389 --

390 -- @

391 -- > S.toList $ S.enumerateFromThenTo 0 2 6

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

393 -- > S.toList $ S.enumerateFromThenTo 0 (-2) (-6)

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

395 -- @

396 --

397 -- @since 0.6.0

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

399

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

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

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

403-- count by Int.

404--

405-- enumerateN

406-- enumerateFromN

407-- enumerateToN

408-- enumerateFromStep

409-- enumerateFromStepN

410

411-------------------------------------------------------------------------------

412-- Convenient functions for bounded types

413-------------------------------------------------------------------------------

414--

415-- |

416-- > enumerate = enumerateFrom minBound

417--

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

419--

420-- @since 0.6.0

421{-# INLINE enumerate #-}

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

423enumerate = enumerateFrom minBound

424

425-- |

426-- > enumerateTo = enumerateFromTo minBound

427--

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

429--

430-- @since 0.6.0

431{-# INLINE enumerateTo #-}

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

433enumerateTo = enumerateFromTo minBound

434

435-- |

436-- > enumerateFromBounded = enumerateFromTo from maxBound

437--

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

439--

440-- @since 0.6.0

441{-# INLINE enumerateFromBounded #-}

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

443 => a -> t m a

444enumerateFromBounded from = enumerateFromTo from maxBound

445

446-------------------------------------------------------------------------------

447-- Enumerable Instances

448-------------------------------------------------------------------------------

449--

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

451#define ENUMERABLE_BOUNDED_SMALL(SMALL_TYPE) \

452instance Enumerable SMALL_TYPE where { \

453 {-# INLINE enumerateFrom #-}; \

454 enumerateFrom = enumerateFromBounded; \

455 {-# INLINE enumerateFromThen #-}; \

456 enumerateFromThen = enumerateFromThenSmallBounded; \

457 {-# INLINE enumerateFromTo #-}; \

458 enumerateFromTo = enumerateFromToSmall; \

459 {-# INLINE enumerateFromThenTo #-}; \

460 enumerateFromThenTo = enumerateFromThenToSmall }

461

462

463ENUMERABLE_BOUNDED_SMALL(())

464ENUMERABLE_BOUNDED_SMALL(Bool)

465ENUMERABLE_BOUNDED_SMALL(Ordering)

466ENUMERABLE_BOUNDED_SMALL(Char)

467

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

469#define ENUMERABLE_BOUNDED_INTEGRAL(INTEGRAL_TYPE) \

470instance Enumerable INTEGRAL_TYPE where { \

471 {-# INLINE enumerateFrom #-}; \

472 enumerateFrom = enumerateFromIntegral; \

473 {-# INLINE enumerateFromThen #-}; \

474 enumerateFromThen = enumerateFromThenIntegral; \

475 {-# INLINE enumerateFromTo #-}; \

476 enumerateFromTo = enumerateFromToIntegral; \

477 {-# INLINE enumerateFromThenTo #-}; \

478 enumerateFromThenTo = enumerateFromThenToIntegral }

479

480ENUMERABLE_BOUNDED_INTEGRAL(Int)

481ENUMERABLE_BOUNDED_INTEGRAL(Int8)

482ENUMERABLE_BOUNDED_INTEGRAL(Int16)

483ENUMERABLE_BOUNDED_INTEGRAL(Int32)

484ENUMERABLE_BOUNDED_INTEGRAL(Int64)

485ENUMERABLE_BOUNDED_INTEGRAL(Word)

486ENUMERABLE_BOUNDED_INTEGRAL(Word8)

487ENUMERABLE_BOUNDED_INTEGRAL(Word16)

488ENUMERABLE_BOUNDED_INTEGRAL(Word32)

489ENUMERABLE_BOUNDED_INTEGRAL(Word64)

490

491-- For unbounded Integral Enum types.

492#define ENUMERABLE_UNBOUNDED_INTEGRAL(INTEGRAL_TYPE) \

493instance Enumerable INTEGRAL_TYPE where { \

494 {-# INLINE enumerateFrom #-}; \

495 enumerateFrom from = enumerateFromStepIntegral from 1; \

496 {-# INLINE enumerateFromThen #-}; \

497 enumerateFromThen from next = \

498 enumerateFromStepIntegral from (next - from); \

499 {-# INLINE enumerateFromTo #-}; \

500 enumerateFromTo = enumerateFromToIntegral; \

501 {-# INLINE enumerateFromThenTo #-}; \

502 enumerateFromThenTo = enumerateFromThenToIntegral }

503

504ENUMERABLE_UNBOUNDED_INTEGRAL(Integer)

505ENUMERABLE_UNBOUNDED_INTEGRAL(Natural)

506

507#define ENUMERABLE_FRACTIONAL(FRACTIONAL_TYPE,CONSTRAINT) \

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

509 {-# INLINE enumerateFrom #-}; \

510 enumerateFrom = enumerateFromFractional; \

511 {-# INLINE enumerateFromThen #-}; \

512 enumerateFromThen = enumerateFromThenFractional; \

513 {-# INLINE enumerateFromTo #-}; \

514 enumerateFromTo = enumerateFromToFractional; \

515 {-# INLINE enumerateFromThenTo #-}; \

516 enumerateFromThenTo = enumerateFromThenToFractional }

517

518ENUMERABLE_FRACTIONAL(Float,)

519ENUMERABLE_FRACTIONAL(Double,)

520ENUMERABLE_FRACTIONAL(Fixed a,HasResolution a)

521ENUMERABLE_FRACTIONAL(Ratio a,Integral a)

522

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

524 {-# INLINE enumerateFrom #-}

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

526 enumerateFrom from

527 {-# INLINE enumerateFromThen #-}

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

529 enumerateFromThen from next

530 {-# INLINE enumerateFromTo #-}

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

532 enumerateFromTo from to

533 {-# INLINE enumerateFromThenTo #-}

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

535 Serial.map Identity $ enumerateFromThenTo from next to

536

537-- TODO

538{-

539instance Enumerable a => Enumerable (Last a)

540instance Enumerable a => Enumerable (First a)

541instance Enumerable a => Enumerable (Max a)

542instance Enumerable a => Enumerable (Min a)

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

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

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

546-}