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

git.iko.soy/streamly

on 511fe85 help sign in

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

1096 lines 35 kB history raw

1{-# LANGUAGE BangPatterns #-}

2{-# LANGUAGE CPP #-}

3{-# LANGUAGE ConstraintKinds #-}

4{-# LANGUAGE FlexibleContexts #-}

5{-# LANGUAGE FlexibleInstances #-}

6{-# LANGUAGE InstanceSigs #-}

7{-# LANGUAGE MultiParamTypeClasses #-}

8{-# LANGUAGE RankNTypes #-}

9{-# LANGUAGE ScopedTypeVariables #-}

10{-# LANGUAGE UndecidableInstances #-} -- XXX

12#include "inline.hs"

14-- |

15-- Module : Streamly.Internal.Data.Stream.StreamK

17--

18-- License : BSD3

19-- Maintainer : streamly@composewell.com

20-- Stability : experimental

21-- Portability : GHC

22--

23--

24-- Continuation passing style (CPS) stream implementation. The symbol 'K' below

25-- denotes a function as well as a Kontinuation.

26--

27-- @

28-- import qualified Streamly.Internal.Data.Stream.StreamK as K

29-- @

30--

31module Streamly.Internal.Data.Stream.StreamK

32 (

33 -- * A class for streams

34 IsStream (..)

35 , adapt

37 -- * The stream type

38 , Stream

40 -- * Construction Primitives

41 , mkStream

42 , nil

43 , nilM

44 , cons

45 , (.:)

47 -- * Elimination Primitives

48 , foldStream

49 , foldStreamShared

50 , foldStreamSVar

52 -- * Transformation Primitives

53 , unShare

55 -- * Deconstruction

56 , uncons

58 -- * Generation

59 -- ** Unfolds

60 , unfoldr

61 , unfoldrM

63 -- ** Specialized Generation

64 , repeat

65 , repeatM

66 , replicate

67 , replicateM

68 , fromIndices

69 , fromIndicesM

70 , iterate

71 , iterateM

73 -- ** Conversions

74 , yield

75 , yieldM

76 , fromFoldable

77 , fromList

78 , fromStreamK

80 -- * foldr/build

81 , foldrS

82 , foldrSM

83 , buildS

84 , buildM

85 , augmentS

86 , augmentSM

88 -- * Elimination

89 -- ** General Folds

90 , foldr

91 , foldr1

92 , foldrM

93 , foldrT

95 , foldl'

96 , foldlM'

97 , foldlS

98 , foldlT

99 , foldlx'

100 , foldlMx'

101

102 -- ** Specialized Folds

103 , drain

104 , null

105 , head

106 , tail

107 , init

108 , elem

109 , notElem

110 , all

111 , any

112 , last

113 , minimum

114 , minimumBy

115 , maximum

116 , maximumBy

117 , findIndices

118 , lookup

119 , findM

120 , find

121 , (!!)

122

123 -- ** Map and Fold

124 , mapM_

125

126 -- ** Conversions

127 , toList

128 , toStreamK

129 , hoist

130

131 -- * Transformation

132 -- ** By folding (scans)

133 , scanl'

134 , scanlx'

135

136 -- ** Filtering

137 , filter

138 , take

139 , takeWhile

140 , drop

141 , dropWhile

142

143 -- ** Mapping

144 , map

145 , mapM

146 , mapMSerial

147 , sequence

148

149 -- ** Inserting

150 , intersperseM

151 , intersperse

152 , insertBy

153

154 -- ** Deleting

155 , deleteBy

156

157 -- ** Reordering

158 , reverse

159

160 -- ** Map and Filter

161 , mapMaybe

162

163 -- ** Zipping

164 , zipWith

165 , zipWithM

166

167 -- ** Merging

168 , mergeBy

169 , mergeByM

170

171 -- ** Nesting

172 , concatMapBy

173 , concatMap

174 , bindWith

175

176 -- ** Transformation comprehensions

177 , the

178

179 -- * Semigroup Style Composition

180 , serial

181

182 -- * Utilities

183 , consMStream

184 , withLocal

185 , mfix

186

187 -- * Deprecated

188 , Streaming -- deprecated

189 , once -- deprecated

190 )

191where

192

193import Control.Monad.Trans (MonadTrans(lift))

194import Control.Monad (void, join)

195import Control.Monad.Reader.Class (MonadReader(..))

196import Data.Function (fix)

197import Prelude

198 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,

199 take, filter, all, any, takeWhile, drop, dropWhile, minimum,

200 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,

201 foldr1, (!!), replicate, reverse, concatMap, iterate)

202import qualified Prelude

203

204import Streamly.Internal.Data.SVar

205import Streamly.Internal.Data.Stream.StreamK.Type

206

207-------------------------------------------------------------------------------

208-- Deconstruction

209-------------------------------------------------------------------------------

210

211{-# INLINE uncons #-}

212uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))

213uncons m =

214 let stop = return Nothing

215 single a = return (Just (a, nil))

216 yieldk a r = return (Just (a, r))

217 in foldStream defState yieldk single stop m

218

219-------------------------------------------------------------------------------

220-- Generation

221-------------------------------------------------------------------------------

222

223{-# INLINE unfoldr #-}

224unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a

225unfoldr next s0 = build $ \yld stp ->

226 let go s =

227 case next s of

228 Just (a, b) -> yld a (go b)

229 Nothing -> stp

230 in go s0

231

232{-# INLINE unfoldrM #-}

233unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a

234unfoldrM step = go

235 where

236 go s = sharedM $ \yld _ stp -> do

237 r <- step s

238 case r of

239 Just (a, b) -> yld a (go b)

240 Nothing -> stp

241

242{-

243-- Generalization of concurrent streams/SVar via unfoldr.

244--

245-- Unfold a value into monadic actions and then run the resulting monadic

246-- actions to generate a stream. Since the step of generating the monadic

247-- action and running them are decoupled we can run the monadic actions

248-- cooncurrently. For example, the seed could be a list of monadic actions or a

249-- pure stream of monadic actions.

250--

251-- We can have different flavors of this depending on the stream type t. The

252-- concurrent version could be async or ahead etc. Depending on how we queue

253-- back the feedback portion b, it could be DFS or BFS style.

254--

255unfoldrA :: (IsStream t, MonadAsync m) => (b -> Maybe (m a, b)) -> b -> t m a

256unfoldrA = undefined

257-}

258

259-------------------------------------------------------------------------------

260-- Special generation

261-------------------------------------------------------------------------------

262

263-- | Same as yieldM

264--

265-- @since 0.2.0

266{-# DEPRECATED once "Please use yieldM instead." #-}

267{-# INLINE once #-}

268once :: (Monad m, IsStream t) => m a -> t m a

269once = yieldM

270

271-- |

272-- @

273-- repeatM = fix . cons

274-- repeatM = cycle1 . yield

275-- @

276--

277-- Generate an infinite stream by repeating a monadic value.

278--

279-- /Internal/

280repeatM :: (IsStream t, MonadAsync m) => m a -> t m a

281repeatM = go

282 where go m = m |: go m

283

284-- Generate an infinite stream by repeating a pure value.

285--

286-- /Internal/

287{-# INLINE repeat #-}

288repeat :: IsStream t => a -> t m a

289repeat a = let x = cons a x in x

290

291{-# INLINE replicateM #-}

292replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a

293replicateM n m = go n

294 where

295 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)

296

297{-# INLINE replicate #-}

298replicate :: IsStream t => Int -> a -> t m a

299replicate n a = go n

300 where

301 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)

302

303{-# INLINE fromIndicesM #-}

304fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a

305fromIndicesM gen = go 0

306 where

307 go i = mkStream $ \st stp sng yld -> do

308 foldStreamShared st stp sng yld (gen i |: go (i + 1))

309

310{-# INLINE fromIndices #-}

311fromIndices :: IsStream t => (Int -> a) -> t m a

312fromIndices gen = go 0

313 where

314 go n = (gen n) `cons` go (n + 1)

315

316{-# INLINE iterate #-}

317iterate :: IsStream t => (a -> a) -> a -> t m a

318iterate step = fromStream . go

319 where

320 go s = cons s (go (step s))

321

322{-# INLINE iterateM #-}

323iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a

324iterateM step = go

325 where

326 go s = mkStream $ \st stp sng yld -> do

327 next <- s

328 foldStreamShared st stp sng yld (return next |: go (step next))

329

330-------------------------------------------------------------------------------

331-- Conversions

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

333

334-- |

335-- @

336-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'

337-- @

338--

339-- Construct a stream from a 'Foldable' containing pure values:

340--

341-- @since 0.2.0

342{-# INLINE fromFoldable #-}

343fromFoldable :: (IsStream t, Foldable f) => f a -> t m a

344fromFoldable = Prelude.foldr cons nil

345

346{-# INLINE fromList #-}

347fromList :: IsStream t => [a] -> t m a

348fromList = fromFoldable

349

350{-# INLINE fromStreamK #-}

351fromStreamK :: IsStream t => Stream m a -> t m a

352fromStreamK = fromStream

353

354-------------------------------------------------------------------------------

355-- Elimination by Folding

356-------------------------------------------------------------------------------

357

358-- | Lazy right associative fold.

359{-# INLINE foldr #-}

360foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b

361foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)

362

363-- | Right associative fold to an arbitrary transformer monad.

364{-# INLINE foldrT #-}

365foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)

366 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b

367foldrT step final m = go m

368 where

369 go m1 = do

370 res <- lift $ uncons m1

371 case res of

372 Just (h, t) -> step h (go t)

373 Nothing -> final

374

375{-# INLINE foldr1 #-}

376foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)

377foldr1 step m = do

378 r <- uncons m

379 case r of

380 Nothing -> return Nothing

381 Just (h, t) -> fmap Just (go h t)

382 where

383 go p m1 =

384 let stp = return p

385 single a = return $ step a p

386 yieldk a r = fmap (step p) (go a r)

387 in foldStream defState yieldk single stp m1

388

389-- | Strict left fold with an extraction function. Like the standard strict

390-- left fold, but applies a user supplied extraction function (the third

391-- argument) to the folded value at the end. This is designed to work with the

392-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.

393--

394-- Note that the accumulator is always evaluated including the initial value.

395{-# INLINE foldlx' #-}

396foldlx' :: forall t m a b x. (IsStream t, Monad m)

397 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b

398foldlx' step begin done m = get $ go m begin

399 where

400 {-# NOINLINE get #-}

401 get :: t m x -> m b

402 get m1 =

403 -- XXX we are not strictly evaluating the accumulator here. Is this

404 -- okay?

405 let single = return . done

406 -- XXX this is foldSingleton. why foldStreamShared?

407 in foldStreamShared undefined undefined single undefined m1

408

409 -- Note, this can be implemented by making a recursive call to "go",

410 -- however that is more expensive because of unnecessary recursion

411 -- that cannot be tail call optimized. Unfolding recursion explicitly via

412 -- continuations is much more efficient.

413 go :: t m a -> x -> t m x

414 go m1 !acc = mkStream $ \_ yld sng _ ->

415 let stop = sng acc

416 single a = sng $ step acc a

417 -- XXX this is foldNonEmptyStream

418 yieldk a r = foldStream defState yld sng undefined $

419 go r (step acc a)

420 in foldStream defState yieldk single stop m1

421

422-- | Strict left associative fold.

423{-# INLINE foldl' #-}

424foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b

425foldl' step begin = foldlx' step begin id

426

427-- XXX replace the recursive "go" with explicit continuations.

428-- | Like 'foldx', but with a monadic step function.

429{-# INLINABLE foldlMx' #-}

430foldlMx' :: (IsStream t, Monad m)

431 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b

432foldlMx' step begin done m = go begin m

433 where

434 go !acc m1 =

435 let stop = acc >>= done

436 single a = acc >>= \b -> step b a >>= done

437 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r

438 in foldStream defState yieldk single stop m1

439

440-- | Like 'foldl'' but with a monadic step function.

441{-# INLINE foldlM' #-}

442foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> m b

443foldlM' step begin = foldlMx' step (return begin) return

444

445-- | Lazy left fold to a stream.

446{-# INLINE foldlS #-}

447foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b

448foldlS step begin m = go begin m

449 where

450 go acc rest = mkStream $ \st yld sng stp ->

451 let run x = foldStream st yld sng stp x

452 stop = run acc

453 single a = run $ step acc a

454 yieldk a r = run $ go (step acc a) r

455 in foldStream (adaptState st) yieldk single stop rest

456

457-- | Lazy left fold to an arbitrary transformer monad.

458{-# INLINE foldlT #-}

459foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)

460 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b

461foldlT step begin m = go begin m

462 where

463 go acc m1 = do

464 res <- lift $ uncons m1

465 case res of

466 Just (h, t) -> go (step acc h) t

467 Nothing -> acc

468

469------------------------------------------------------------------------------

470-- Specialized folds

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

472

473-- XXX use foldrM to implement folds where possible

474-- XXX This (commented) definition of drain and mapM_ perform much better on

475-- some benchmarks but worse on others. Need to investigate why, may there is

476-- an optimization opportunity that we can exploit.

477-- drain = foldrM (\_ xs -> return () >> xs) (return ())

478

479-- |

480-- > drain = foldl' (\_ _ -> ()) ()

481-- > drain = mapM_ (\_ -> return ())

482{-# INLINE drain #-}

483drain :: (Monad m, IsStream t) => t m a -> m ()

484drain = foldrM (\_ xs -> xs) (return ())

485{-

486drain = go

487 where

488 go m1 =

489 let stop = return ()

490 single _ = return ()

491 yieldk _ r = go r

492 in foldStream defState yieldk single stop m1

493-}

494

495{-# INLINE null #-}

496null :: (IsStream t, Monad m) => t m a -> m Bool

497-- null = foldrM (\_ _ -> return True) (return False)

498null m =

499 let stop = return True

500 single _ = return False

501 yieldk _ _ = return False

502 in foldStream defState yieldk single stop m

503

504{-# INLINE head #-}

505head :: (IsStream t, Monad m) => t m a -> m (Maybe a)

506-- head = foldrM (\x _ -> return $ Just x) (return Nothing)

507head m =

508 let stop = return Nothing

509 single a = return (Just a)

510 yieldk a _ = return (Just a)

511 in foldStream defState yieldk single stop m

512

513{-# INLINE tail #-}

514tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))

515tail m =

516 let stop = return Nothing

517 single _ = return $ Just nil

518 yieldk _ r = return $ Just r

519 in foldStream defState yieldk single stop m

520

521{-# INLINE headPartial #-}

522headPartial :: (IsStream t, Monad m) => t m a -> m a

523headPartial = foldrM (\x _ -> return x) (error "head of nil")

524

525{-# INLINE tailPartial #-}

526tailPartial :: IsStream t => t m a -> t m a

527tailPartial m = mkStream $ \st yld sng stp ->

528 let stop = error "tail of nil"

529 single _ = stp

530 yieldk _ r = foldStream st yld sng stp r

531 in foldStream st yieldk single stop m

532

533-- | Iterate a lazy function `f` of the shape `m a -> t m a` until it gets

534-- fully defined i.e. becomes independent of its argument action, then return

535-- the resulting value of the function (`t m a`).

536--

537-- It can be used to construct a stream that uses a cyclic definition. For

538-- example:

539--

540-- @

541-- import Streamly.Internal.Prelude as S

542-- import System.IO.Unsafe (unsafeInterleaveIO)

543--

544-- main = do

545-- S.mapM_ print $ S.mfix $ \x -> do

546-- a <- S.fromList [1,2]

547-- b <- S.fromListM [return 3, unsafeInterleaveIO (fmap fst x)]

548-- return (a, b)

549-- @

550--

551-- Note that the function `f` must be lazy in its argument, that's why we use

552-- 'unsafeInterleaveIO' because IO monad is strict.

553

554mfix :: (IsStream t, Monad m) => (m a -> t m a) -> t m a

555mfix f = mkStream $ \st yld sng stp ->

556 let single a = foldStream st yld sng stp $ a `cons` ys

557 yieldk a _ = foldStream st yld sng stp $ a `cons` ys

558 in foldStream st yieldk single stp xs

559 where xs = fix (f . headPartial)

560 ys = mfix (tailPartial . f)

561

562{-# INLINE init #-}

563init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))

564init m = go1 m

565 where

566 go1 m1 = do

567 r <- uncons m1

568 case r of

569 Nothing -> return Nothing

570 Just (h, t) -> return . Just $ go h t

571 go p m1 = mkStream $ \_ yld sng stp ->

572 let single _ = sng p

573 yieldk a x = yld p $ go a x

574 in foldStream defState yieldk single stp m1

575

576{-# INLINE elem #-}

577elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool

578elem e m = go m

579 where

580 go m1 =

581 let stop = return False

582 single a = return (a == e)

583 yieldk a r = if a == e then return True else go r

584 in foldStream defState yieldk single stop m1

585

586{-# INLINE notElem #-}

587notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool

588notElem e m = go m

589 where

590 go m1 =

591 let stop = return True

592 single a = return (a /= e)

593 yieldk a r = if a == e then return False else go r

594 in foldStream defState yieldk single stop m1

595

596{-# INLINABLE all #-}

597all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool

598all p m = go m

599 where

600 go m1 =

601 let single a | p a = return True

602 | otherwise = return False

603 yieldk a r | p a = go r

604 | otherwise = return False

605 in foldStream defState yieldk single (return True) m1

606

607{-# INLINABLE any #-}

608any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool

609any p m = go m

610 where

611 go m1 =

612 let single a | p a = return True

613 | otherwise = return False

614 yieldk a r | p a = return True

615 | otherwise = go r

616 in foldStream defState yieldk single (return False) m1

617

618-- | Extract the last element of the stream, if any.

619{-# INLINE last #-}

620last :: (IsStream t, Monad m) => t m a -> m (Maybe a)

621last = foldlx' (\_ y -> Just y) Nothing id

622

623{-# INLINE minimum #-}

624minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)

625minimum m = go Nothing m

626 where

627 go Nothing m1 =

628 let stop = return Nothing

629 single a = return (Just a)

630 yieldk a r = go (Just a) r

631 in foldStream defState yieldk single stop m1

632

633 go (Just res) m1 =

634 let stop = return (Just res)

635 single a =

636 if res <= a

637 then return (Just res)

638 else return (Just a)

639 yieldk a r =

640 if res <= a

641 then go (Just res) r

642 else go (Just a) r

643 in foldStream defState yieldk single stop m1

644

645{-# INLINE minimumBy #-}

646minimumBy

647 :: (IsStream t, Monad m)

648 => (a -> a -> Ordering) -> t m a -> m (Maybe a)

649minimumBy cmp m = go Nothing m

650 where

651 go Nothing m1 =

652 let stop = return Nothing

653 single a = return (Just a)

654 yieldk a r = go (Just a) r

655 in foldStream defState yieldk single stop m1

656

657 go (Just res) m1 =

658 let stop = return (Just res)

659 single a = case cmp res a of

660 GT -> return (Just a)

661 _ -> return (Just res)

662 yieldk a r = case cmp res a of

663 GT -> go (Just a) r

664 _ -> go (Just res) r

665 in foldStream defState yieldk single stop m1

666

667{-# INLINE maximum #-}

668maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)

669maximum m = go Nothing m

670 where

671 go Nothing m1 =

672 let stop = return Nothing

673 single a = return (Just a)

674 yieldk a r = go (Just a) r

675 in foldStream defState yieldk single stop m1

676

677 go (Just res) m1 =

678 let stop = return (Just res)

679 single a =

680 if res <= a

681 then return (Just a)

682 else return (Just res)

683 yieldk a r =

684 if res <= a

685 then go (Just a) r

686 else go (Just res) r

687 in foldStream defState yieldk single stop m1

688

689{-# INLINE maximumBy #-}

690maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)

691maximumBy cmp m = go Nothing m

692 where

693 go Nothing m1 =

694 let stop = return Nothing

695 single a = return (Just a)

696 yieldk a r = go (Just a) r

697 in foldStream defState yieldk single stop m1

698

699 go (Just res) m1 =

700 let stop = return (Just res)

701 single a = case cmp res a of

702 GT -> return (Just res)

703 _ -> return (Just a)

704 yieldk a r = case cmp res a of

705 GT -> go (Just res) r

706 _ -> go (Just a) r

707 in foldStream defState yieldk single stop m1

708

709{-# INLINE (!!) #-}

710(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)

711m !! i = go i m

712 where

713 go n m1 =

714 let single a | n == 0 = return $ Just a

715 | otherwise = return Nothing

716 yieldk a x | n < 0 = return Nothing

717 | n == 0 = return $ Just a

718 | otherwise = go (n - 1) x

719 in foldStream defState yieldk single (return Nothing) m1

720

721{-# INLINE lookup #-}

722lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)

723lookup e m = go m

724 where

725 go m1 =

726 let single (a, b) | a == e = return $ Just b

727 | otherwise = return Nothing

728 yieldk (a, b) x | a == e = return $ Just b

729 | otherwise = go x

730 in foldStream defState yieldk single (return Nothing) m1

731

732{-# INLINE findM #-}

733findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)

734findM p m = go m

735 where

736 go m1 =

737 let single a = do

738 b <- p a

739 if b then return $ Just a else return Nothing

740 yieldk a x = do

741 b <- p a

742 if b then return $ Just a else go x

743 in foldStream defState yieldk single (return Nothing) m1

744

745{-# INLINE find #-}

746find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)

747find p = findM (return . p)

748

749{-# INLINE findIndices #-}

750findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int

751findIndices p = go 0

752 where

753 go offset m1 = mkStream $ \st yld sng stp ->

754 let single a | p a = sng offset

755 | otherwise = stp

756 yieldk a x | p a = yld offset $ go (offset + 1) x

757 | otherwise = foldStream (adaptState st) yld sng stp $

758 go (offset + 1) x

759 in foldStream (adaptState st) yieldk single stp m1

760

761------------------------------------------------------------------------------

762-- Map and Fold

763------------------------------------------------------------------------------

764

765-- | Apply a monadic action to each element of the stream and discard the

766-- output of the action.

767{-# INLINE mapM_ #-}

768mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()

769mapM_ f m = go m

770 where

771 go m1 =

772 let stop = return ()

773 single a = void (f a)

774 yieldk a r = f a >> go r

775 in foldStream defState yieldk single stop m1

776

777------------------------------------------------------------------------------

778-- Converting folds

779------------------------------------------------------------------------------

780

781{-# INLINABLE toList #-}

782toList :: (IsStream t, Monad m) => t m a -> m [a]

783toList = foldr (:) []

784

785{-# INLINE toStreamK #-}

786toStreamK :: Stream m a -> Stream m a

787toStreamK = id

788

789-- Based on suggestions by David Feuer and Pranay Sashank

790{-# INLINE hoist #-}

791hoist :: (IsStream t, Monad m, Monad n)

792 => (forall x. m x -> n x) -> t m a -> t n a

793hoist f str =

794 mkStream $ \st yld sng stp ->

795 let single = return . sng

796 yieldk a s = return $ yld a (hoist f s)

797 stop = return stp

798 state = adaptState st

799 in join . f $ foldStreamShared state yieldk single stop str

800

801-------------------------------------------------------------------------------

802-- Transformation by folding (Scans)

803-------------------------------------------------------------------------------

804

805{-# INLINE scanlx' #-}

806scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

807scanlx' step begin done m =

808 cons (done begin) $ go m begin

809 where

810 go m1 !acc = mkStream $ \st yld sng stp ->

811 let single a = sng (done $ step acc a)

812 yieldk a r =

813 let s = step acc a

814 in yld (done s) (go r s)

815 in foldStream (adaptState st) yieldk single stp m1

816

817{-# INLINE scanl' #-}

818scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b

819scanl' step begin = scanlx' step begin id

820

821-------------------------------------------------------------------------------

822-- Filtering

823-------------------------------------------------------------------------------

824

825{-# INLINE filter #-}

826filter :: IsStream t => (a -> Bool) -> t m a -> t m a

827filter p m = go m

828 where

829 go m1 = mkStream $ \st yld sng stp ->

830 let single a | p a = sng a

831 | otherwise = stp

832 yieldk a r | p a = yld a (go r)

833 | otherwise = foldStream st yieldk single stp r

834 in foldStream st yieldk single stp m1

835

836{-# INLINE take #-}

837take :: IsStream t => Int -> t m a -> t m a

838take n m = go n m

839 where

840 go n1 m1 = mkStream $ \st yld sng stp ->

841 let yieldk a r = yld a (go (n1 - 1) r)

842 in if n1 <= 0

843 then stp

844 else foldStream st yieldk sng stp m1

845

846{-# INLINE takeWhile #-}

847takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a

848takeWhile p m = go m

849 where

850 go m1 = mkStream $ \st yld sng stp ->

851 let single a | p a = sng a

852 | otherwise = stp

853 yieldk a r | p a = yld a (go r)

854 | otherwise = stp

855 in foldStream st yieldk single stp m1

856

857{-# INLINE drop #-}

858drop :: IsStream t => Int -> t m a -> t m a

859drop n m = fromStream $ unShare (go n (toStream m))

860 where

861 go n1 m1 = mkStream $ \st yld sng stp ->

862 let single _ = stp

863 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r

864 -- Somehow "<=" check performs better than a ">"

865 in if n1 <= 0

866 then foldStreamShared st yld sng stp m1

867 else foldStreamShared st yieldk single stp m1

868

869{-# INLINE dropWhile #-}

870dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a

871dropWhile p m = go m

872 where

873 go m1 = mkStream $ \st yld sng stp ->

874 let single a | p a = stp

875 | otherwise = sng a

876 yieldk a r | p a = foldStream st yieldk single stp r

877 | otherwise = yld a r

878 in foldStream st yieldk single stp m1

879

880-------------------------------------------------------------------------------

881-- Mapping

882-------------------------------------------------------------------------------

883

884-- Be careful when modifying this, this uses a consM (|:) deliberately to allow

885-- other stream types to overload it.

886{-# INLINE sequence #-}

887sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a

888sequence m = go m

889 where

890 go m1 = mkStream $ \st yld sng stp ->

891 let single ma = ma >>= sng

892 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r

893 in foldStream (adaptState st) yieldk single stp m1

894

895-------------------------------------------------------------------------------

896-- Inserting

897-------------------------------------------------------------------------------

898

899{-# INLINE intersperseM #-}

900intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a

901intersperseM a m = prependingStart m

902 where

903 prependingStart m1 = mkStream $ \st yld sng stp ->

904 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x

905 in foldStream st yieldk sng stp m1

906 go m2 = mkStream $ \st yld sng stp ->

907 let single i = foldStreamShared st yld sng stp $ a |: yield i

908 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x

909 in foldStream st yieldk single stp m2

910

911{-# INLINE intersperse #-}

912intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a

913intersperse a = intersperseM (return a)

914

915{-# INLINE insertBy #-}

916insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a

917insertBy cmp x m = go m

918 where

919 go m1 = mkStream $ \st yld _ _ ->

920 let single a = case cmp x a of

921 GT -> yld a (yield x)

922 _ -> yld x (yield a)

923 stop = yld x nil

924 yieldk a r = case cmp x a of

925 GT -> yld a (go r)

926 _ -> yld x (a `cons` r)

927 in foldStream st yieldk single stop m1

928

929------------------------------------------------------------------------------

930-- Deleting

931------------------------------------------------------------------------------

932

933{-# INLINE deleteBy #-}

934deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a

935deleteBy eq x m = go m

936 where

937 go m1 = mkStream $ \st yld sng stp ->

938 let single a = if eq x a then stp else sng a

939 yieldk a r = if eq x a

940 then foldStream st yld sng stp r

941 else yld a (go r)

942 in foldStream st yieldk single stp m1

943

944------------------------------------------------------------------------------

945-- Reordering

946------------------------------------------------------------------------------

947

948{-# INLINE reverse #-}

949reverse :: IsStream t => t m a -> t m a

950reverse = foldlS (flip cons) nil

951

952-------------------------------------------------------------------------------

953-- Map and Filter

954-------------------------------------------------------------------------------

955

956{-# INLINE mapMaybe #-}

957mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b

958mapMaybe f m = go m

959 where

960 go m1 = mkStream $ \st yld sng stp ->

961 let single a = case f a of

962 Just b -> sng b

963 Nothing -> stp

964 yieldk a r = case f a of

965 Just b -> yld b $ go r

966 Nothing -> foldStream (adaptState st) yieldk single stp r

967 in foldStream (adaptState st) yieldk single stp m1

968

969------------------------------------------------------------------------------

970-- Serial Zipping

971------------------------------------------------------------------------------

972

973-- | Zip two streams serially using a pure zipping function.

974--

975-- @since 0.1.0

976{-# INLINABLE zipWith #-}

977zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c

978zipWith f = go

979 where

980 go mx my = mkStream $ \st yld sng stp -> do

981 let merge a ra =

982 let single2 b = sng (f a b)

983 yield2 b rb = yld (f a b) (go ra rb)

984 in foldStream (adaptState st) yield2 single2 stp my

985 let single1 a = merge a nil

986 yield1 = merge

987 foldStream (adaptState st) yield1 single1 stp mx

988

989-- | Zip two streams serially using a monadic zipping function.

990--

991-- @since 0.1.0

992{-# INLINABLE zipWithM #-}

993zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c

994zipWithM f m1 m2 = go m1 m2

995 where

996 go mx my = mkStream $ \st yld sng stp -> do

997 let merge a ra =

998 let runIt x = foldStream st yld sng stp x

999 single2 b = f a b >>= sng

1000 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)

1001 in foldStream (adaptState st) yield2 single2 stp my

1002 let single1 a = merge a nil

1003 yield1 = merge

1004 foldStream (adaptState st) yield1 single1 stp mx

1005

1006------------------------------------------------------------------------------

1007-- Merging

1008------------------------------------------------------------------------------

1009

1010{-# INLINE mergeByM #-}

1011mergeByM

1012 :: (IsStream t, Monad m)

1013 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a

1014mergeByM cmp = go

1015 where

1016 go mx my = mkStream $ \st yld sng stp -> do

1017 let mergeWithY a ra =

1018 let stop2 = foldStream st yld sng stp mx

1019 single2 b = do

1020 r <- cmp a b

1021 case r of

1022 GT -> yld b (go (a `cons` ra) nil)

1023 _ -> yld a (go ra (b `cons` nil))

1024 yield2 b rb = do

1025 r <- cmp a b

1026 case r of

1027 GT -> yld b (go (a `cons` ra) rb)

1028 _ -> yld a (go ra (b `cons` rb))

1029 in foldStream st yield2 single2 stop2 my

1030 let stopX = foldStream st yld sng stp my

1031 singleX a = mergeWithY a nil

1032 yieldX = mergeWithY

1033 foldStream st yieldX singleX stopX mx

1034

1035{-# INLINABLE mergeBy #-}

1036mergeBy

1037 :: (IsStream t, Monad m)

1038 => (a -> a -> Ordering) -> t m a -> t m a -> t m a

1039mergeBy cmp = mergeByM (\a b -> return $ cmp a b)

1040

1041------------------------------------------------------------------------------

1042-- Transformation comprehensions

1043------------------------------------------------------------------------------

1044

1045{-# INLINE the #-}

1046the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)

1047the m = do

1048 r <- uncons m

1049 case r of

1050 Nothing -> return Nothing

1051 Just (h, t) -> go h t

1052 where

1053 go h m1 =

1054 let single a | h == a = return $ Just h

1055 | otherwise = return Nothing

1056 yieldk a r | h == a = go h r

1057 | otherwise = return Nothing

1058 in foldStream defState yieldk single (return $ Just h) m1

1059

1060------------------------------------------------------------------------------

1061-- Alternative & MonadPlus

1062------------------------------------------------------------------------------

1063

1064_alt :: Stream m a -> Stream m a -> Stream m a

1065_alt m1 m2 = mkStream $ \st yld sng stp ->

1066 let stop = foldStream st yld sng stp m2

1067 in foldStream st yld sng stop m1

1068

1069------------------------------------------------------------------------------

1070-- MonadReader

1071------------------------------------------------------------------------------

1072

1073{-# INLINABLE withLocal #-}

1074withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a

1075withLocal f m =

1076 mkStream $ \st yld sng stp ->

1077 let single = local f . sng

1078 yieldk a r = local f $ yld a (withLocal f r)

1079 in foldStream st yieldk single (local f stp) m

1080

1081------------------------------------------------------------------------------

1082-- MonadError

1083------------------------------------------------------------------------------

1084

1085{-

1086-- XXX handle and test cross thread state transfer

1087withCatchError

1088 :: MonadError e m

1089 => Stream m a -> (e -> Stream m a) -> Stream m a

1090withCatchError m h =

1091 mkStream $ \_ stp sng yld ->

1092 let run x = unStream x Nothing stp sng yieldk

1093 handle r = r `catchError` \e -> run $ h e

1094 yieldk a r = yld a (withCatchError r h)

1095 in handle $ run m

1096-}