src/Streamly/Enumeration.hs · streamly

git.iko.soy/streamly

on bcd787e help sign in

streamly/src/Streamly/Enumeration.hs

523 lines 20 kB history raw

1{-# LANGUAGE CPP #-}

3-- |

4-- Module : Streamly.Enumeration

6--

7-- License : BSD3

8-- Maintainer : harendra.kumar@gmail.com

9-- Stability : experimental

10-- Portability : GHC

11--

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

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

14-- instances instead.

15--

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

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

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

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

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

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

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

23-- can be used directly.

25module Streamly.Enumeration

26 (

27 Enumerable (..)

28 , enumerate

30 -- ** Enumerating 'Enum' Types

31 , enumerateFromToSmall

32 , enumerateFromThenToSmall

33 , enumerateFromThenSmallBounded

34 , enumerateFromBounded

36 -- ** Enumerating 'Bounded' 'Integral' Types

37 , enumerateFromIntegral

38 , enumerateFromThenIntegral

40 -- ** Enumerating unbounded 'Integral' Types

41 , enumerateFromToIntegral

42 , enumerateFromThenToIntegral

43 , enumerateFromStepIntegral

45 -- ** Enumerating unbounded 'Fractional' Types

46 , enumerateFromFractional

47 , enumerateFromToFractional

48 , enumerateFromThenFractional

49 , enumerateFromThenToFractional

51 -- ** Enumerating unbounded 'Num' Types

52 , enumerateFromStepNum

53 )

54where

56import Data.Fixed

57import Data.Int

58import Data.Ratio

59import Data.Word

60import Numeric.Natural

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

63import Streamly.Streams.StreamD (fromStreamD)

64import Streamly.Streams.StreamK (IsStream(..))

66import qualified Streamly.Streams.StreamD as D

67import qualified Streamly.Streams.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@. This does not check for overflow or underflow if the 'Integral' type

76-- is bounded.

77--

78-- @

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

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

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

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

83-- @

84--

85-- @since 0.6.0

86{-# INLINE enumerateFromStepIntegral #-}

87enumerateFromStepIntegral

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

89 => a -> a -> t m a

90enumerateFromStepIntegral from stride =

91 fromStreamD $ D.enumerateFromStepIntegral from stride

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

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

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

96--

97-- @

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

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

100-- @

101--

102-- @since 0.6.0

103{-# INLINE enumerateFromIntegral #-}

104enumerateFromIntegral

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

106 => a -> t m a

107enumerateFromIntegral from = fromStreamD $ D.enumerateFromIntegral from

108

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

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

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

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

113--

114-- @

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

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

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

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

119-- @

120--

121-- @since 0.6.0

122{-# INLINE enumerateFromThenIntegral #-}

123enumerateFromThenIntegral

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

125 => a -> a -> t m a

126enumerateFromThenIntegral from next =

127 fromStreamD $ D.enumerateFromThenIntegral from next

128

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

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

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

132-- @to@.

133--

134-- @

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

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

137-- @

138--

139-- @since 0.6.0

140{-# INLINE enumerateFromToIntegral #-}

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

142enumerateFromToIntegral from to =

143 fromStreamD $ D.enumerateFromToIntegral from to

144

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

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

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

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

149--

150-- @

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

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

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

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

155-- @

156--

157-- @since 0.6.0

158{-# INLINE enumerateFromThenToIntegral #-}

159enumerateFromThenToIntegral

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

161 => a -> a -> a -> t m a

162enumerateFromThenToIntegral from next to =

163 fromStreamD $ D.enumerateFromThenToIntegral from next to

164

165-------------------------------------------------------------------------------

166-- Enumeration of Num types

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

168--

169-- | @enumerateFromStepNum from step@ generates an infinite stream whose first

170-- element is @from@ and the successive elements are in increments of @step@.

171-- This is numerically stable but does not check for overflow or underflow for

172-- bounded types. Note, for 'Integral' types 'enumerateFromStepIntegral' is

173-- faster.

174--

175--

176-- @

177-- > S.toList $ S.take 4 $ S.enumerateFromStepNum 0.1 2

178-- [0.1,2.1,4.1,6.1]

179-- > S.toList $ S.take 3 $ S.enumerateFromStepNum 0.1 (-2)

180-- [0.1,-1.9,-3.9,-5.9]

181-- @

182--

183-- @since 0.6.0

184{-# INLINE enumerateFromStepNum #-}

185enumerateFromStepNum :: (IsStream t, Monad m, Num a) => a -> a -> t m a

186enumerateFromStepNum from stride =

187 fromStreamD $ D.enumerateFromStepNum from stride

188

189-------------------------------------------------------------------------------

190-- Enumeration of Fractional types

191-------------------------------------------------------------------------------

192--

193-- Even though the underlying implementation of enumerateFromFractional and

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

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

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

197--

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

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

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

201-- or underflow checks are performed.

202--

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

204--

205-- @

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

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

208-- @

209--

210--

211-- @since 0.6.0

212{-# INLINE enumerateFromFractional #-}

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

214enumerateFromFractional from = fromStreamD $ D.numFrom from

215

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

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

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

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

220-- performed.

221--

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

223-- example:

224--

225-- @

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

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

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

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

230-- @

231--

232-- @since 0.6.0

233{-# INLINE enumerateFromThenFractional #-}

234enumerateFromThenFractional

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

236 => a -> a -> t m a

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

238

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

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

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

242-- to @to@.

243--

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

245-- example:

246--

247-- @

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

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

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

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

252-- @

253--

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

255-- rounding to the nearest integer.

256--

257-- @since 0.6.0

258{-# INLINE enumerateFromToFractional #-}

259enumerateFromToFractional

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

261 => a -> a -> t m a

262enumerateFromToFractional from to =

263 fromStreamD $ D.enumerateFromToFractional from to

264

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

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

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

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

269--

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

271-- example:

272--

273-- @

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

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

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

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

278-- @

279--

280--

281-- @since 0.6.0

282{-# INLINE enumerateFromThenToFractional #-}

283enumerateFromThenToFractional

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

285 => a -> a -> a -> t m a

286enumerateFromThenToFractional from next to =

287 fromStreamD $ D.enumerateFromThenToFractional from next to

288

289-------------------------------------------------------------------------------

290-- Enumeration of Enum types

291-------------------------------------------------------------------------------

292--

293-- | 'enumerateFromTo' for 'Enum' types not larger than 'Int'. This function

294-- does not check the bounds for 'Bounded' types, it assumes that the @to@

295-- parameter is constrained by the type to be within the range representable by

296-- the type.

297{-# INLINE enumerateFromToSmall #-}

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

299enumerateFromToSmall from to = Serial.map toEnum $

300 enumerateFromToIntegral (fromEnum from) (fromEnum to)

301

302-- | 'enumerateFrom' for 'Bounded' 'Enum' types. This is defined in terms of

303-- 'enumerateFromTo', therefore it will work even for types larger than 'Int'

304-- depending on 'enumerateFromTo'.

305{-# INLINE enumerateFromBounded #-}

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

307 => a -> t m a

308enumerateFromBounded from = enumerateFromTo from maxBound

309

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

311-- function does not check the bounds for 'Bounded' types, it assumes that the

312-- @to@ parameter is constrained by the type to be within the range

313-- representable by the type.

314{-# INLINE enumerateFromThenToSmall #-}

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

316 => a -> a -> a -> t m a

317enumerateFromThenToSmall from next to = Serial.map toEnum $

318 enumerateFromThenToIntegral (fromEnum from) (fromEnum next) (fromEnum to)

319

320-- | 'enumerateFromThen' for 'Bounded' 'Enum' types not larger than 'Int'. For

321-- types smaller than 'Int' we know it is bounded though it may not have a

322-- 'Bounded' instance. It is not necessary to require 'Bounded' instance for

323-- all small types, we require it anyway because it is safe that way.

324{-# INLINE enumerateFromThenSmallBounded #-}

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

326 => a -> a -> t m a

327enumerateFromThenSmallBounded from next =

328 case fromEnum next >= fromEnum from of

329 True -> enumerateFromThenTo from next maxBound

330 False -> enumerateFromThenTo from next minBound

331

332-------------------------------------------------------------------------------

333-- Enumerable type class

334-------------------------------------------------------------------------------

335--

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

337-- enumerations like this:

338--

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

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

341--

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

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

344

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

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

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

348-- "Streamly.Enumeration" module to define new instances.

349--

350-- @since 0.6.0

351class Enum a => Enumerable a where

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

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

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

355 --

356 -- @

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

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

359 -- @

360 --

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

362 -- overflow or underflow checks are performed.

363 --

364 -- @

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

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

367 -- @

368 --

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

370

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

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

373 -- empty stream is returned.

374 --

375 -- @

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

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

378 -- @

379 --

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

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

382 --

383 -- @

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

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

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

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

388 -- @

389 --

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

391

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

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

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

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

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

397 -- unbounded types it keeps enumerating infinitely.

398 --

399 -- @

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

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

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

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

404 -- @

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 -- > S.toList $ S.enumerateFromThenTo 0 2 6

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

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

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

418 -- @

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

420

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

422#define ENUMERABLE_BOUNDED_SMALL(SMALL_TYPE) \

423instance Enumerable SMALL_TYPE where { \

424 {-# INLINE enumerateFrom #-}; \

425 enumerateFrom = enumerateFromBounded; \

426 {-# INLINE enumerateFromThen #-}; \

427 enumerateFromThen = enumerateFromThenSmallBounded; \

428 {-# INLINE enumerateFromTo #-}; \

429 enumerateFromTo = enumerateFromToSmall; \

430 {-# INLINE enumerateFromThenTo #-}; \

431 enumerateFromThenTo = enumerateFromThenToSmall }

432

433

434ENUMERABLE_BOUNDED_SMALL(())

435ENUMERABLE_BOUNDED_SMALL(Bool)

436ENUMERABLE_BOUNDED_SMALL(Ordering)

437ENUMERABLE_BOUNDED_SMALL(Char)

438

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

440#define ENUMERABLE_BOUNDED_INTEGRAL(INTEGRAL_TYPE) \

441instance Enumerable INTEGRAL_TYPE where { \

442 {-# INLINE enumerateFrom #-}; \

443 enumerateFrom = enumerateFromIntegral; \

444 {-# INLINE enumerateFromThen #-}; \

445 enumerateFromThen = enumerateFromThenIntegral; \

446 {-# INLINE enumerateFromTo #-}; \

447 enumerateFromTo = enumerateFromToIntegral; \

448 {-# INLINE enumerateFromThenTo #-}; \

449 enumerateFromThenTo = enumerateFromThenToIntegral }

450

451ENUMERABLE_BOUNDED_INTEGRAL(Int)

452ENUMERABLE_BOUNDED_INTEGRAL(Int8)

453ENUMERABLE_BOUNDED_INTEGRAL(Int16)

454ENUMERABLE_BOUNDED_INTEGRAL(Int32)

455ENUMERABLE_BOUNDED_INTEGRAL(Int64)

456ENUMERABLE_BOUNDED_INTEGRAL(Word)

457ENUMERABLE_BOUNDED_INTEGRAL(Word8)

458ENUMERABLE_BOUNDED_INTEGRAL(Word16)

459ENUMERABLE_BOUNDED_INTEGRAL(Word32)

460ENUMERABLE_BOUNDED_INTEGRAL(Word64)

461

462-- For unbounded Integral Enum types.

463#define ENUMERABLE_UNBOUNDED_INTEGRAL(INTEGRAL_TYPE) \

464instance Enumerable INTEGRAL_TYPE where { \

465 {-# INLINE enumerateFrom #-}; \

466 enumerateFrom from = enumerateFromStepIntegral from 1; \

467 {-# INLINE enumerateFromThen #-}; \

468 enumerateFromThen from next = \

469 enumerateFromStepIntegral from (next - from); \

470 {-# INLINE enumerateFromTo #-}; \

471 enumerateFromTo = enumerateFromToIntegral; \

472 {-# INLINE enumerateFromThenTo #-}; \

473 enumerateFromThenTo = enumerateFromThenToIntegral }

474

475ENUMERABLE_UNBOUNDED_INTEGRAL(Integer)

476ENUMERABLE_UNBOUNDED_INTEGRAL(Natural)

477

478#define ENUMERABLE_FRACTIONAL(FRACTIONAL_TYPE,CONSTRAINT) \

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

480 {-# INLINE enumerateFrom #-}; \

481 enumerateFrom = enumerateFromFractional; \

482 {-# INLINE enumerateFromThen #-}; \

483 enumerateFromThen = enumerateFromThenFractional; \

484 {-# INLINE enumerateFromTo #-}; \

485 enumerateFromTo = enumerateFromToFractional; \

486 {-# INLINE enumerateFromThenTo #-}; \

487 enumerateFromThenTo = enumerateFromThenToFractional }

488

489ENUMERABLE_FRACTIONAL(Float,)

490ENUMERABLE_FRACTIONAL(Double,)

491ENUMERABLE_FRACTIONAL(Fixed a,HasResolution a)

492ENUMERABLE_FRACTIONAL(Ratio a,Integral a)

493

494#if __GLASGOW_HASKELL__ >= 800

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

496 {-# INLINE enumerateFrom #-}

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

498 enumerateFrom from

499 {-# INLINE enumerateFromThen #-}

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

501 enumerateFromThen from next

502 {-# INLINE enumerateFromTo #-}

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

504 enumerateFromTo from to

505 {-# INLINE enumerateFromThenTo #-}

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

507 Serial.map Identity $ enumerateFromThenTo from next to

508#endif

509{-

510-- TODO

511instance Enumerable a => Enumerable (Last a)

512instance Enumerable a => Enumerable (First a)

513instance Enumerable a => Enumerable (Max a)

514instance Enumerable a => Enumerable (Min a)

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

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

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

518-}

519

520-- | Enumerate a finite ('Bounded') data type from its 'minBound' to 'maxBound'

521{-# INLINE enumerate #-}

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

523enumerate = enumerateFrom minBound