core/src/Streamly/Internal/Data/Stream/Enumerate.hs · streamly

git.iko.soy/streamly

on efe0be3 help sign in

streamly/core/src/Streamly/Internal/Data/Stream/Enumerate.hs

560 lines 21 kB history raw

1-- |

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

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.

23-- XXX The Unfold.Enumeration module is more modular, check the differences and

24-- reconcile the two.

26module Streamly.Internal.Data.Stream.Enumerate

27 (

28 Enumerable (..)

30 -- ** Enumerating 'Bounded' 'Enum' Types

31 , enumerate

32 , enumerateTo

33 , enumerateFromBounded

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

36 , enumerateFromToSmall

37 , enumerateFromThenToSmall

38 , enumerateFromThenSmallBounded

40 -- ** Enumerating 'Bounded' 'Integral' Types

41 , enumerateFromIntegral

42 , enumerateFromThenIntegral

44 -- ** Enumerating 'Integral' Types

45 , enumerateFromToIntegral

46 , enumerateFromThenToIntegral

48 -- ** Enumerating unbounded 'Integral' Types

49 , enumerateFromStepIntegral

51 -- ** Enumerating 'Fractional' Types

52 , enumerateFromFractional

53 , enumerateFromToFractional

54 , enumerateFromThenFractional

55 , enumerateFromThenToFractional

56 )

57where

59import Data.Fixed

60import Data.Int

61import Data.Ratio

62import Data.Word

63import Numeric.Natural

64import Data.Functor.Identity (Identity(..))

66import Streamly.Internal.Data.Stream.Type (Stream, fromStreamD)

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

70-- $setup

71-- >>> import Streamly.Data.Fold as Fold

72-- >>> import Streamly.Internal.Data.Stream as Stream

73-- >>> import Streamly.Internal.Data.Stream.Enumerate as Stream

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

76-- Enumeration of Integral types

77-------------------------------------------------------------------------------

78--

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

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

81-- @step@.

82--

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

84-- check for overflow, underflow or bounds.

85--

86-- @

87-- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromStepIntegral 0 2

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

89--

90-- >>> Stream.fold Fold.toList $ Stream.take 3 $ Stream.enumerateFromStepIntegral 0 (-2)

91-- [0,-2,-4]

92--

93-- @

94--

95{-# INLINE enumerateFromStepIntegral #-}

96enumerateFromStepIntegral

97 :: (Monad m, Integral a)

98 => a -> a -> Stream m a

99enumerateFromStepIntegral from stride =

100 fromStreamD $ D.enumerateFromStepIntegral from stride

101

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

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

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

105--

106-- @

107-- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromIntegral (0 :: Int)

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

109--

110-- @

111--

112{-# INLINE enumerateFromIntegral #-}

113enumerateFromIntegral

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

115 => a -> Stream 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.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromThenIntegral (0 :: Int) 2

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

126--

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

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

129--

130-- @

131--

132{-# INLINE enumerateFromThenIntegral #-}

133enumerateFromThenIntegral

134 :: (Monad m, Integral a, Bounded a)

135 => a -> a -> Stream m a

136enumerateFromThenIntegral from next =

137 fromStreamD $ D.enumerateFromThenIntegral from next

138

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

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

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

142-- @to@.

143--

144-- @

145-- >>> Stream.fold Fold.toList $ Stream.enumerateFromToIntegral 0 4

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

147--

148-- @

149--

150{-# INLINE enumerateFromToIntegral #-}

151enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a

152enumerateFromToIntegral from to =

153 fromStreamD $ D.enumerateFromToIntegral from to

154

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

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

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

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

159--

160-- @

161-- >>> Stream.fold Fold.toList $ Stream.enumerateFromThenToIntegral 0 2 6

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

163--

164-- >>> Stream.fold Fold.toList $ Stream.enumerateFromThenToIntegral 0 (-2) (-6)

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

166--

167-- @

168--

169{-# INLINE enumerateFromThenToIntegral #-}

170enumerateFromThenToIntegral

171 :: (Monad m, Integral a)

172 => a -> a -> a -> Stream 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.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromFractional 1.1

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

195--

196-- @

197--

198--

199{-# INLINE enumerateFromFractional #-}

200enumerateFromFractional :: (Monad m, Fractional a) => a -> Stream m a

201enumerateFromFractional from = fromStreamD $ D.enumerateFromNum from

202

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

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

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

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

207-- performed.

208--

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

210-- example:

211--

212-- @

213-- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromThenFractional 1.1 2.1

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

215--

216-- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromThenFractional 1.1 (-2.1)

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

218--

219-- @

220--

221{-# INLINE enumerateFromThenFractional #-}

222enumerateFromThenFractional

223 :: (Monad m, Fractional a)

224 => a -> a -> Stream m a

225enumerateFromThenFractional from next = fromStreamD $ D.enumerateFromThenNum from next

226

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

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

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

230-- to @to@.

231--

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

233-- example:

234--

235-- @

236-- >>> Stream.fold Fold.toList $ Stream.enumerateFromToFractional 1.1 4

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

238--

239-- >>> Stream.fold Fold.toList $ Stream.enumerateFromToFractional 1.1 4.6

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

241--

242-- @

243--

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

245-- rounding to the nearest integer.

246--

247{-# INLINE enumerateFromToFractional #-}

248enumerateFromToFractional

249 :: (Monad m, Fractional a, Ord a)

250 => a -> a -> Stream m a

251enumerateFromToFractional from to =

252 fromStreamD $ D.enumerateFromToFractional from to

253

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

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

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

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

258--

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

260-- example:

261--

262-- @

263-- >>> Stream.fold Fold.toList $ Stream.enumerateFromThenToFractional 0.1 2 6

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

265--

266-- >>> Stream.fold Fold.toList $ Stream.enumerateFromThenToFractional 0.1 (-2) (-6)

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

268--

269-- @

270--

271--

272{-# INLINE enumerateFromThenToFractional #-}

273enumerateFromThenToFractional

274 :: (Monad m, Fractional a, Ord a)

275 => a -> a -> a -> Stream m a

276enumerateFromThenToFractional from next to =

277 fromStreamD $ D.enumerateFromThenToFractional from next to

278

279-------------------------------------------------------------------------------

280-- Enumeration of Enum types not larger than Int

281-------------------------------------------------------------------------------

282--

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

284--

285{-# INLINE enumerateFromToSmall #-}

286enumerateFromToSmall :: (Monad m, Enum a) => a -> a -> Stream m a

287enumerateFromToSmall from to =

288 fmap toEnum

289 $ enumerateFromToIntegral (fromEnum from) (fromEnum to)

290

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

292--

293{-# INLINE enumerateFromThenToSmall #-}

294enumerateFromThenToSmall :: (Monad m, Enum a)

295 => a -> a -> a -> Stream m a

296enumerateFromThenToSmall from next to =

297 fmap toEnum

298 $ enumerateFromThenToIntegral

299 (fromEnum from) (fromEnum next) (fromEnum to)

300

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

302--

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

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

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

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

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

308--

309{-# INLINE enumerateFromThenSmallBounded #-}

310enumerateFromThenSmallBounded :: (Monad m, Enumerable a, Bounded a)

311 => a -> a -> Stream m a

312enumerateFromThenSmallBounded from next =

313 if fromEnum next >= fromEnum from

314 then enumerateFromThenTo from next maxBound

315 else enumerateFromThenTo from next minBound

316

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

318-- Enumerable type class

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

320--

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

322-- enumerations like this:

323--

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

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

326--

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

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

329

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

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

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

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

334--

335class Enum a => Enumerable a where

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

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

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

339 --

340 -- @

341 -- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFrom (0 :: Int)

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

343 --

344 -- @

345 --

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

347 -- overflow or underflow checks are performed.

348 --

349 -- @

350 -- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFrom 1.1

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

352 --

353 -- @

354 --

355 enumerateFrom :: (Monad m) => a -> Stream m a

356

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

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

359 -- empty stream is returned.

360 --

361 -- @

362 -- >>> Stream.fold Fold.toList $ Stream.enumerateFromTo 0 4

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

364 --

365 -- @

366 --

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

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

369 --

370 -- @

371 -- >>> Stream.fold Fold.toList $ Stream.enumerateFromTo 1.1 4

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

373 --

374 -- >>> Stream.fold Fold.toList $ Stream.enumerateFromTo 1.1 4.6

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

376 --

377 -- @

378 --

379 enumerateFromTo :: (Monad m) => a -> a -> Stream m a

380

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

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

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

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

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

386 -- unbounded types it keeps enumerating infinitely.

387 --

388 -- @

389 -- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromThen 0 2

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

391 --

392 -- >>> Stream.fold Fold.toList $ Stream.take 4 $ Stream.enumerateFromThen 0 (-2)

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

394 --

395 -- @

396 --

397 enumerateFromThen :: (Monad m) => a -> a -> Stream m a

398

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

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

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

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

403 -- after @from@.

404 --

405 -- @

406 -- >>> Stream.fold Fold.toList $ Stream.enumerateFromThenTo 0 2 6

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

408 --

409 -- >>> Stream.fold Fold.toList $ Stream.enumerateFromThenTo 0 (-2) (-6)

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

411 --

412 -- @

413 --

414 enumerateFromThenTo :: (Monad m) => a -> a -> a -> Stream m a

415

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

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

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

419-- count by Int.

420--

421-- enumerateN

422-- enumerateFromN

423-- enumerateToN

424-- enumerateFromStep

425-- enumerateFromStepN

426

427-------------------------------------------------------------------------------

428-- Convenient functions for bounded types

429-------------------------------------------------------------------------------

430--

431-- |

432-- > enumerate = enumerateFrom minBound

433--

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

435--

436{-# INLINE enumerate #-}

437enumerate :: (Monad m, Bounded a, Enumerable a) => Stream m a

438enumerate = enumerateFrom minBound

439

440-- |

441-- > enumerateTo = enumerateFromTo minBound

442--

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

444--

445{-# INLINE enumerateTo #-}

446enumerateTo :: (Monad m, Bounded a, Enumerable a) => a -> Stream m a

447enumerateTo = enumerateFromTo minBound

448

449-- |

450-- > enumerateFromBounded = enumerateFromTo from maxBound

451--

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

453--

454{-# INLINE enumerateFromBounded #-}

455enumerateFromBounded :: (Monad m, Enumerable a, Bounded a)

456 => a -> Stream m a

457enumerateFromBounded from = enumerateFromTo from maxBound

458

459-------------------------------------------------------------------------------

460-- Enumerable Instances

461-------------------------------------------------------------------------------

462--

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

464#define ENUMERABLE_BOUNDED_SMALL(SMALL_TYPE) \

465instance Enumerable SMALL_TYPE where { \

466 {-# INLINE enumerateFrom #-}; \

467 enumerateFrom = enumerateFromBounded; \

468 {-# INLINE enumerateFromThen #-}; \

469 enumerateFromThen = enumerateFromThenSmallBounded; \

470 {-# INLINE enumerateFromTo #-}; \

471 enumerateFromTo = enumerateFromToSmall; \

472 {-# INLINE enumerateFromThenTo #-}; \

473 enumerateFromThenTo = enumerateFromThenToSmall }

474

475

476ENUMERABLE_BOUNDED_SMALL(())

477ENUMERABLE_BOUNDED_SMALL(Bool)

478ENUMERABLE_BOUNDED_SMALL(Ordering)

479ENUMERABLE_BOUNDED_SMALL(Char)

480

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

482#define ENUMERABLE_BOUNDED_INTEGRAL(INTEGRAL_TYPE) \

483instance Enumerable INTEGRAL_TYPE where { \

484 {-# INLINE enumerateFrom #-}; \

485 enumerateFrom = enumerateFromIntegral; \

486 {-# INLINE enumerateFromThen #-}; \

487 enumerateFromThen = enumerateFromThenIntegral; \

488 {-# INLINE enumerateFromTo #-}; \

489 enumerateFromTo = enumerateFromToIntegral; \

490 {-# INLINE enumerateFromThenTo #-}; \

491 enumerateFromThenTo = enumerateFromThenToIntegral }

492

493ENUMERABLE_BOUNDED_INTEGRAL(Int)

494ENUMERABLE_BOUNDED_INTEGRAL(Int8)

495ENUMERABLE_BOUNDED_INTEGRAL(Int16)

496ENUMERABLE_BOUNDED_INTEGRAL(Int32)

497ENUMERABLE_BOUNDED_INTEGRAL(Int64)

498ENUMERABLE_BOUNDED_INTEGRAL(Word)

499ENUMERABLE_BOUNDED_INTEGRAL(Word8)

500ENUMERABLE_BOUNDED_INTEGRAL(Word16)

501ENUMERABLE_BOUNDED_INTEGRAL(Word32)

502ENUMERABLE_BOUNDED_INTEGRAL(Word64)

503

504-- For unbounded Integral Enum types.

505#define ENUMERABLE_UNBOUNDED_INTEGRAL(INTEGRAL_TYPE) \

506instance Enumerable INTEGRAL_TYPE where { \

507 {-# INLINE enumerateFrom #-}; \

508 enumerateFrom from = enumerateFromStepIntegral from 1; \

509 {-# INLINE enumerateFromThen #-}; \

510 enumerateFromThen from next = \

511 enumerateFromStepIntegral from (next - from); \

512 {-# INLINE enumerateFromTo #-}; \

513 enumerateFromTo = enumerateFromToIntegral; \

514 {-# INLINE enumerateFromThenTo #-}; \

515 enumerateFromThenTo = enumerateFromThenToIntegral }

516

517ENUMERABLE_UNBOUNDED_INTEGRAL(Integer)

518ENUMERABLE_UNBOUNDED_INTEGRAL(Natural)

519

520#define ENUMERABLE_FRACTIONAL(FRACTIONAL_TYPE,CONSTRAINT) \

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

522 {-# INLINE enumerateFrom #-}; \

523 enumerateFrom = enumerateFromFractional; \

524 {-# INLINE enumerateFromThen #-}; \

525 enumerateFromThen = enumerateFromThenFractional; \

526 {-# INLINE enumerateFromTo #-}; \

527 enumerateFromTo = enumerateFromToFractional; \

528 {-# INLINE enumerateFromThenTo #-}; \

529 enumerateFromThenTo = enumerateFromThenToFractional }

530

531ENUMERABLE_FRACTIONAL(Float,)

532ENUMERABLE_FRACTIONAL(Double,)

533ENUMERABLE_FRACTIONAL((Fixed a),HasResolution a)

534ENUMERABLE_FRACTIONAL((Ratio a),Integral a)

535

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

537 {-# INLINE enumerateFrom #-}

538 enumerateFrom (Identity from) =

539 fmap Identity $ enumerateFrom from

540 {-# INLINE enumerateFromThen #-}

541 enumerateFromThen (Identity from) (Identity next) =

542 fmap Identity $ enumerateFromThen from next

543 {-# INLINE enumerateFromTo #-}

544 enumerateFromTo (Identity from) (Identity to) =

545 fmap Identity $ enumerateFromTo from to

546 {-# INLINE enumerateFromThenTo #-}

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

548 fmap Identity

549 $ enumerateFromThenTo from next to

550

551-- TODO

552{-

553instance Enumerable a => Enumerable (Last a)

554instance Enumerable a => Enumerable (First a)

555instance Enumerable a => Enumerable (Max a)

556instance Enumerable a => Enumerable (Min a)

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

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

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

560-}