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

git.iko.soy/streamly

on 3feeec2 help sign in

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

571 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-------------------------------------------------------------------------------

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-- >>> import Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream

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

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

84--

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

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

87--

88-- @

89--

90-- @since 0.6.0

91{-# INLINE enumerateFromStepIntegral #-}

92enumerateFromStepIntegral

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

94 => a -> a -> t m a

95enumerateFromStepIntegral from stride =

96 fromStreamD $ D.enumerateFromStepIntegral from stride

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

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

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

101--

102-- @

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

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

105--

106-- @

107--

108-- @since 0.6.0

109{-# INLINE enumerateFromIntegral #-}

110enumerateFromIntegral

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

112 => a -> t m a

113enumerateFromIntegral from = fromStreamD $ D.enumerateFromIntegral from

114

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

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

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

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

119--

120-- @

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

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

123--

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

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

126--

127-- @

128--

129-- @since 0.6.0

130{-# INLINE enumerateFromThenIntegral #-}

131enumerateFromThenIntegral

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

133 => a -> a -> t m a

134enumerateFromThenIntegral from next =

135 fromStreamD $ D.enumerateFromThenIntegral from next

136

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

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

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

140-- @to@.

141--

142-- @

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

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

145--

146-- @

147--

148-- @since 0.6.0

149{-# INLINE enumerateFromToIntegral #-}

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

151enumerateFromToIntegral from to =

152 fromStreamD $ D.enumerateFromToIntegral from to

153

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

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

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

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

158--

159-- @

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

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

162--

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

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

165--

166-- @

167--

168-- @since 0.6.0

169{-# INLINE enumerateFromThenToIntegral #-}

170enumerateFromThenToIntegral

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

172 => a -> a -> a -> t m a

173enumerateFromThenToIntegral from next to =

174 fromStreamD $ D.enumerateFromThenToIntegral from next to

175

176-------------------------------------------------------------------------------

177-- Enumeration of Fractional types

178-------------------------------------------------------------------------------

179--

180-- Even though the underlying implementation of enumerateFromFractional and

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

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

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

184--

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

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

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

188-- or underflow checks are performed.

189--

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

191--

192-- @

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

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

195--

196-- @

197--

198--

199-- @since 0.6.0

200{-# INLINE enumerateFromFractional #-}

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

202enumerateFromFractional from = fromStreamD $ D.numFrom from

203

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

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

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

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

208-- performed.

209--

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

211-- example:

212--

213-- @

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

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

216--

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

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

219--

220-- @

221--

222-- @since 0.6.0

223{-# INLINE enumerateFromThenFractional #-}

224enumerateFromThenFractional

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

226 => a -> a -> t m a

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

228

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

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

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

232-- to @to@.

233--

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

235-- example:

236--

237-- @

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

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

240--

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

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

243--

244-- @

245--

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

247-- rounding to the nearest integer.

248--

249-- @since 0.6.0

250{-# INLINE enumerateFromToFractional #-}

251enumerateFromToFractional

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

253 => a -> a -> t m a

254enumerateFromToFractional from to =

255 fromStreamD $ D.enumerateFromToFractional from to

256

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

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

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

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

261--

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

263-- example:

264--

265-- @

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

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

268--

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

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

271--

272-- @

273--

274--

275-- @since 0.6.0

276{-# INLINE enumerateFromThenToFractional #-}

277enumerateFromThenToFractional

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

279 => a -> a -> a -> t m a

280enumerateFromThenToFractional from next to =

281 fromStreamD $ D.enumerateFromThenToFractional from next to

282

283-------------------------------------------------------------------------------

284-- Enumeration of Enum types not larger than Int

285-------------------------------------------------------------------------------

286--

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

288--

289-- @since 0.6.0

290{-# INLINE enumerateFromToSmall #-}

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

292enumerateFromToSmall from to = Serial.map toEnum $

293 enumerateFromToIntegral (fromEnum from) (fromEnum to)

294

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

296--

297-- @since 0.6.0

298{-# INLINE enumerateFromThenToSmall #-}

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

300 => a -> a -> a -> t m a

301enumerateFromThenToSmall from next to = Serial.map toEnum $

302 enumerateFromThenToIntegral (fromEnum from) (fromEnum next) (fromEnum to)

303

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

305--

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

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

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

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

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

311--

312-- @since 0.6.0

313{-# INLINE enumerateFromThenSmallBounded #-}

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

315 => a -> a -> t m a

316enumerateFromThenSmallBounded from next =

317 if fromEnum next >= fromEnum from

318 then enumerateFromThenTo from next maxBound

319 else enumerateFromThenTo from next minBound

320

321-------------------------------------------------------------------------------

322-- Enumerable type class

323-------------------------------------------------------------------------------

324--

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

326-- enumerations like this:

327--

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

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

330--

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

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

333

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

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

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

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

338--

339-- @since 0.6.0

340class Enum a => Enumerable a where

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

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

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

344 --

345 -- @

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

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

348 --

349 -- @

350 --

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

352 -- overflow or underflow checks are performed.

353 --

354 -- @

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

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

357 --

358 -- @

359 --

360 -- @since 0.6.0

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

362

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

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

365 -- empty stream is returned.

366 --

367 -- @

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

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

370 --

371 -- @

372 --

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

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

375 --

376 -- @

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

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

379 --

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

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

382 --

383 -- @

384 --

385 -- @since 0.6.0

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

387

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

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

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

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

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

393 -- unbounded types it keeps enumerating infinitely.

394 --

395 -- @

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

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

398 --

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

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

401 --

402 -- @

403 --

404 -- @since 0.6.0

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

406

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

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

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

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

411 -- after @from@.

412 --

413 -- @

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

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

416 --

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

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

419 --

420 -- @

421 --

422 -- @since 0.6.0

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

424

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

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

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

428-- count by Int.

429--

430-- enumerateN

431-- enumerateFromN

432-- enumerateToN

433-- enumerateFromStep

434-- enumerateFromStepN

435

436-------------------------------------------------------------------------------

437-- Convenient functions for bounded types

438-------------------------------------------------------------------------------

439--

440-- |

441-- > enumerate = enumerateFrom minBound

442--

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

444--

445-- @since 0.6.0

446{-# INLINE enumerate #-}

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

448enumerate = enumerateFrom minBound

449

450-- |

451-- > enumerateTo = enumerateFromTo minBound

452--

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

454--

455-- @since 0.6.0

456{-# INLINE enumerateTo #-}

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

458enumerateTo = enumerateFromTo minBound

459

460-- |

461-- > enumerateFromBounded = enumerateFromTo from maxBound

462--

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

464--

465-- @since 0.6.0

466{-# INLINE enumerateFromBounded #-}

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

468 => a -> t m a

469enumerateFromBounded from = enumerateFromTo from maxBound

470

471-------------------------------------------------------------------------------

472-- Enumerable Instances

473-------------------------------------------------------------------------------

474--

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

476#define ENUMERABLE_BOUNDED_SMALL(SMALL_TYPE) \

477instance Enumerable SMALL_TYPE where { \

478 {-# INLINE enumerateFrom #-}; \

479 enumerateFrom = enumerateFromBounded; \

480 {-# INLINE enumerateFromThen #-}; \

481 enumerateFromThen = enumerateFromThenSmallBounded; \

482 {-# INLINE enumerateFromTo #-}; \

483 enumerateFromTo = enumerateFromToSmall; \

484 {-# INLINE enumerateFromThenTo #-}; \

485 enumerateFromThenTo = enumerateFromThenToSmall }

486

487

488ENUMERABLE_BOUNDED_SMALL(())

489ENUMERABLE_BOUNDED_SMALL(Bool)

490ENUMERABLE_BOUNDED_SMALL(Ordering)

491ENUMERABLE_BOUNDED_SMALL(Char)

492

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

494#define ENUMERABLE_BOUNDED_INTEGRAL(INTEGRAL_TYPE) \

495instance Enumerable INTEGRAL_TYPE where { \

496 {-# INLINE enumerateFrom #-}; \

497 enumerateFrom = enumerateFromIntegral; \

498 {-# INLINE enumerateFromThen #-}; \

499 enumerateFromThen = enumerateFromThenIntegral; \

500 {-# INLINE enumerateFromTo #-}; \

501 enumerateFromTo = enumerateFromToIntegral; \

502 {-# INLINE enumerateFromThenTo #-}; \

503 enumerateFromThenTo = enumerateFromThenToIntegral }

504

505ENUMERABLE_BOUNDED_INTEGRAL(Int)

506ENUMERABLE_BOUNDED_INTEGRAL(Int8)

507ENUMERABLE_BOUNDED_INTEGRAL(Int16)

508ENUMERABLE_BOUNDED_INTEGRAL(Int32)

509ENUMERABLE_BOUNDED_INTEGRAL(Int64)

510ENUMERABLE_BOUNDED_INTEGRAL(Word)

511ENUMERABLE_BOUNDED_INTEGRAL(Word8)

512ENUMERABLE_BOUNDED_INTEGRAL(Word16)

513ENUMERABLE_BOUNDED_INTEGRAL(Word32)

514ENUMERABLE_BOUNDED_INTEGRAL(Word64)

515

516-- For unbounded Integral Enum types.

517#define ENUMERABLE_UNBOUNDED_INTEGRAL(INTEGRAL_TYPE) \

518instance Enumerable INTEGRAL_TYPE where { \

519 {-# INLINE enumerateFrom #-}; \

520 enumerateFrom from = enumerateFromStepIntegral from 1; \

521 {-# INLINE enumerateFromThen #-}; \

522 enumerateFromThen from next = \

523 enumerateFromStepIntegral from (next - from); \

524 {-# INLINE enumerateFromTo #-}; \

525 enumerateFromTo = enumerateFromToIntegral; \

526 {-# INLINE enumerateFromThenTo #-}; \

527 enumerateFromThenTo = enumerateFromThenToIntegral }

528

529ENUMERABLE_UNBOUNDED_INTEGRAL(Integer)

530ENUMERABLE_UNBOUNDED_INTEGRAL(Natural)

531

532#define ENUMERABLE_FRACTIONAL(FRACTIONAL_TYPE,CONSTRAINT) \

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

534 {-# INLINE enumerateFrom #-}; \

535 enumerateFrom = enumerateFromFractional; \

536 {-# INLINE enumerateFromThen #-}; \

537 enumerateFromThen = enumerateFromThenFractional; \

538 {-# INLINE enumerateFromTo #-}; \

539 enumerateFromTo = enumerateFromToFractional; \

540 {-# INLINE enumerateFromThenTo #-}; \

541 enumerateFromThenTo = enumerateFromThenToFractional }

542

543ENUMERABLE_FRACTIONAL(Float,)

544ENUMERABLE_FRACTIONAL(Double,)

545ENUMERABLE_FRACTIONAL((Fixed a),HasResolution a)

546ENUMERABLE_FRACTIONAL((Ratio a),Integral a)

547

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

549 {-# INLINE enumerateFrom #-}

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

551 enumerateFrom from

552 {-# INLINE enumerateFromThen #-}

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

554 enumerateFromThen from next

555 {-# INLINE enumerateFromTo #-}

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

557 enumerateFromTo from to

558 {-# INLINE enumerateFromThenTo #-}

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

560 Serial.map Identity $ enumerateFromThenTo from next to

561

562-- TODO

563{-

564instance Enumerable a => Enumerable (Last a)

565instance Enumerable a => Enumerable (First a)

566instance Enumerable a => Enumerable (Max a)

567instance Enumerable a => Enumerable (Min a)

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

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

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

571-}