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

git.iko.soy/streamly

on 294aa21 help sign in

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

1097 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

51 -- * Transformation Primitives

52 , unShare

54 -- * Deconstruction

55 , uncons

57 -- * Generation

58 -- ** Unfolds

59 , unfoldr

60 , unfoldrM

62 -- ** Specialized Generation

63 , repeat

64 , repeatM

65 , replicate

66 , replicateM

67 , fromIndices

68 , fromIndicesM

69 , iterate

70 , iterateM

72 -- ** Conversions

73 , yield

74 , yieldM

75 , fromFoldable

76 , fromList

77 , fromStreamK

79 -- * foldr/build

80 , foldrS

81 , foldrSM

82 , buildS

83 , buildM

84 , augmentS

85 , augmentSM

87 -- * Elimination

88 -- ** General Folds

89 , foldr

90 , foldr1

91 , foldrM

92 , foldrT

94 , foldl'

95 , foldlM'

96 , foldlS

97 , foldlT

98 , foldlx'

99 , foldlMx'

100

101 -- ** Specialized Folds

102 , drain

103 , null

104 , head

105 , tail

106 , init

107 , elem

108 , notElem

109 , all

110 , any

111 , last

112 , minimum

113 , minimumBy

114 , maximum

115 , maximumBy

116 , findIndices

117 , lookup

118 , findM

119 , find

120 , (!!)

121

122 -- ** Map and Fold

123 , mapM_

124

125 -- ** Conversions

126 , toList

127 , toStreamK

128 , hoist

129

130 -- * Transformation

131 -- ** By folding (scans)

132 , scanl'

133 , scanlx'

134

135 -- ** Filtering

136 , filter

137 , take

138 , takeWhile

139 , drop

140 , dropWhile

141

142 -- ** Mapping

143 , map

144 , mapM

145 , mapMSerial

146 , sequence

147

148 -- ** Inserting

149 , intersperseM

150 , intersperse

151 , insertBy

152

153 -- ** Deleting

154 , deleteBy

155

156 -- ** Reordering

157 , reverse

158

159 -- ** Map and Filter

160 , mapMaybe

161

162 -- ** Zipping

163 , zipWith

164 , zipWithM

165

166 -- ** Merging

167 , mergeBy

168 , mergeByM

169

170 -- ** Nesting

171 , concatMapBy

172 , concatMap

173 , bindWith

174

175 -- ** Transformation comprehensions

176 , the

177

178 -- * Semigroup Style Composition

179 , serial

180

181 -- * Utilities

182 , consMStream

183 , withLocal

184 , mfix

185

186 -- * Deprecated

187 , Streaming -- deprecated

188 , once -- deprecated

189 )

190where

191

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

193import Control.Monad (void, join)

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

195import Data.Function (fix)

196import Prelude

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

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

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

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

201import qualified Prelude

202

203import Streamly.Internal.Data.SVar

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

205

206-------------------------------------------------------------------------------

207-- Deconstruction

208-------------------------------------------------------------------------------

209

210{-# INLINE uncons #-}

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

212uncons m =

213 let stop = return Nothing

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

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

216 in foldStream defState yieldk single stop m

217

218-------------------------------------------------------------------------------

219-- Generation

220-------------------------------------------------------------------------------

221

222{-# INLINE unfoldr #-}

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

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

225 let go s =

226 case next s of

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

228 Nothing -> stp

229 in go s0

230

231{-# INLINE unfoldrM #-}

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

233unfoldrM step = go

234 where

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

236 r <- step s

237 case r of

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

239 Nothing -> stp

240

241{-

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

243--

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

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

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

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

248-- pure stream of monadic actions.

249--

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

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

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

253--

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

255unfoldrA = undefined

256-}

257

258-------------------------------------------------------------------------------

259-- Special generation

260-------------------------------------------------------------------------------

261

262-- | Same as yieldM

263--

264-- @since 0.2.0

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

266{-# INLINE once #-}

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

268once = yieldM

269

270-- |

271-- @

272-- repeatM = fix . cons

273-- repeatM = cycle1 . yield

274-- @

275--

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

277--

278-- /Internal/

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

280repeatM = go

281 where go m = m |: go m

282

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

284--

285-- /Internal/

286{-# INLINE repeat #-}

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

288repeat a = let x = cons a x in x

289

290{-# INLINE replicateM #-}

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

292replicateM n m = go n

293 where

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

295

296{-# INLINE replicate #-}

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

298replicate n a = go n

299 where

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

301

302{-# INLINE fromIndicesM #-}

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

304fromIndicesM gen = go 0

305 where

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

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

308

309{-# INLINE fromIndices #-}

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

311fromIndices gen = go 0

312 where

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

314

315{-# INLINE iterate #-}

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

317iterate step = fromStream . go

318 where

319 go s = cons s (go (step s))

320

321{-# INLINE iterateM #-}

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

323iterateM step = go

324 where

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

326 next <- s

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

328

329-------------------------------------------------------------------------------

330-- Conversions

331-------------------------------------------------------------------------------

332

333-- |

334-- @

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

336-- @

337--

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

339--

340-- @since 0.2.0

341{-# INLINE fromFoldable #-}

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

343fromFoldable = Prelude.foldr cons nil

344

345{-# INLINE fromList #-}

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

347fromList = fromFoldable

348

349{-# INLINE fromStreamK #-}

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

351fromStreamK = fromStream

352

353-------------------------------------------------------------------------------

354-- Elimination by Folding

355-------------------------------------------------------------------------------

356

357-- | Lazy right associative fold.

358{-# INLINE foldr #-}

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

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

361

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

363{-# INLINE foldrT #-}

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

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

366foldrT step final m = go m

367 where

368 go m1 = do

369 res <- lift $ uncons m1

370 case res of

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

372 Nothing -> final

373

374{-# INLINE foldr1 #-}

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

376foldr1 step m = do

377 r <- uncons m

378 case r of

379 Nothing -> return Nothing

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

381 where

382 go p m1 =

383 let stp = return p

384 single a = return $ step a p

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

386 in foldStream defState yieldk single stp m1

387

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

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

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

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

392--

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

394{-# INLINE foldlx' #-}

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

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

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

398 where

399 {-# NOINLINE get #-}

400 get :: t m x -> m b

401 get m1 =

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

403 -- okay?

404 let single = return . done

405 -- XXX this is foldSingleton. why foldStreamShared?

406 in foldStreamShared undefined undefined single undefined m1

407

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

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

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

411 -- continuations is much more efficient.

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

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

414 let stop = sng acc

415 single a = sng $ step acc a

416 -- XXX this is foldNonEmptyStream

417 yieldk a r = foldStream defState yld sng undefined $

418 go r (step acc a)

419 in foldStream defState yieldk single stop m1

420

421-- | Strict left associative fold.

422{-# INLINE foldl' #-}

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

424foldl' step begin = foldlx' step begin id

425

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

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

428{-# INLINABLE foldlMx' #-}

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

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

431foldlMx' step begin done m = go begin m

432 where

433 go !acc m1 =

434 let stop = acc >>= done

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

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

437 in foldStream defState yieldk single stop m1

438

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

440{-# INLINE foldlM' #-}

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

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

443

444-- | Lazy left fold to a stream.

445{-# INLINE foldlS #-}

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

447foldlS step begin m = go begin m

448 where

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

450 let run x = foldStream st yld sng stp x

451 stop = run acc

452 single a = run $ step acc a

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

454 in foldStream (adaptState st) yieldk single stop rest

455

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

457{-# INLINE foldlT #-}

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

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

460foldlT step begin m = go begin m

461 where

462 go acc m1 = do

463 res <- lift $ uncons m1

464 case res of

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

466 Nothing -> acc

467

468------------------------------------------------------------------------------

469-- Specialized folds

470------------------------------------------------------------------------------

471

472-- XXX use foldrM to implement folds where possible

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

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

475-- an optimization opportunity that we can exploit.

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

477

478-- |

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

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

481{-# INLINE drain #-}

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

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

484{-

485drain = go

486 where

487 go m1 =

488 let stop = return ()

489 single _ = return ()

490 yieldk _ r = go r

491 in foldStream defState yieldk single stop m1

492-}

493

494{-# INLINE null #-}

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

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

497null m =

498 let stop = return True

499 single _ = return False

500 yieldk _ _ = return False

501 in foldStream defState yieldk single stop m

502

503{-# INLINE head #-}

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

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

506head m =

507 let stop = return Nothing

508 single a = return (Just a)

509 yieldk a _ = return (Just a)

510 in foldStream defState yieldk single stop m

511

512{-# INLINE tail #-}

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

514tail m =

515 let stop = return Nothing

516 single _ = return $ Just nil

517 yieldk _ r = return $ Just r

518 in foldStream defState yieldk single stop m

519

520{-# INLINE headPartial #-}

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

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

523

524{-# INLINE tailPartial #-}

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

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

527 let stop = error "tail of nil"

528 single _ = stp

529 yieldk _ r = foldStream st yld sng stp r

530 in foldStream st yieldk single stop m

531

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

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

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

535--

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

537-- example:

538--

539-- @

540-- import Streamly.Internal.Prelude as S

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

542--

543-- main = do

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

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

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

547-- return (a, b)

548-- @

549--

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

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

552--

553-- /Internal/

554

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

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

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

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

559 in foldStream st yieldk single stp xs

560 where xs = fix (f . headPartial)

561 ys = mfix (tailPartial . f)

562

563{-# INLINE init #-}

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

565init m = go1 m

566 where

567 go1 m1 = do

568 r <- uncons m1

569 case r of

570 Nothing -> return Nothing

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

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

573 let single _ = sng p

574 yieldk a x = yld p $ go a x

575 in foldStream defState yieldk single stp m1

576

577{-# INLINE elem #-}

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

579elem e m = go m

580 where

581 go m1 =

582 let stop = return False

583 single a = return (a == e)

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

585 in foldStream defState yieldk single stop m1

586

587{-# INLINE notElem #-}

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

589notElem e m = go m

590 where

591 go m1 =

592 let stop = return True

593 single a = return (a /= e)

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

595 in foldStream defState yieldk single stop m1

596

597{-# INLINABLE all #-}

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

599all p m = go m

600 where

601 go m1 =

602 let single a | p a = return True

603 | otherwise = return False

604 yieldk a r | p a = go r

605 | otherwise = return False

606 in foldStream defState yieldk single (return True) m1

607

608{-# INLINABLE any #-}

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

610any p m = go m

611 where

612 go m1 =

613 let single a | p a = return True

614 | otherwise = return False

615 yieldk a r | p a = return True

616 | otherwise = go r

617 in foldStream defState yieldk single (return False) m1

618

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

620{-# INLINE last #-}

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

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

623

624{-# INLINE minimum #-}

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

626minimum m = go Nothing m

627 where

628 go Nothing m1 =

629 let stop = return Nothing

630 single a = return (Just a)

631 yieldk a r = go (Just a) r

632 in foldStream defState yieldk single stop m1

633

634 go (Just res) m1 =

635 let stop = return (Just res)

636 single a =

637 if res <= a

638 then return (Just res)

639 else return (Just a)

640 yieldk a r =

641 if res <= a

642 then go (Just res) r

643 else go (Just a) r

644 in foldStream defState yieldk single stop m1

645

646{-# INLINE minimumBy #-}

647minimumBy

648 :: (IsStream t, Monad m)

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

650minimumBy cmp m = go Nothing m

651 where

652 go Nothing m1 =

653 let stop = return Nothing

654 single a = return (Just a)

655 yieldk a r = go (Just a) r

656 in foldStream defState yieldk single stop m1

657

658 go (Just res) m1 =

659 let stop = return (Just res)

660 single a = case cmp res a of

661 GT -> return (Just a)

662 _ -> return (Just res)

663 yieldk a r = case cmp res a of

664 GT -> go (Just a) r

665 _ -> go (Just res) r

666 in foldStream defState yieldk single stop m1

667

668{-# INLINE maximum #-}

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

670maximum m = go Nothing m

671 where

672 go Nothing m1 =

673 let stop = return Nothing

674 single a = return (Just a)

675 yieldk a r = go (Just a) r

676 in foldStream defState yieldk single stop m1

677

678 go (Just res) m1 =

679 let stop = return (Just res)

680 single a =

681 if res <= a

682 then return (Just a)

683 else return (Just res)

684 yieldk a r =

685 if res <= a

686 then go (Just a) r

687 else go (Just res) r

688 in foldStream defState yieldk single stop m1

689

690{-# INLINE maximumBy #-}

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

692maximumBy cmp m = go Nothing m

693 where

694 go Nothing m1 =

695 let stop = return Nothing

696 single a = return (Just a)

697 yieldk a r = go (Just a) r

698 in foldStream defState yieldk single stop m1

699

700 go (Just res) m1 =

701 let stop = return (Just res)

702 single a = case cmp res a of

703 GT -> return (Just res)

704 _ -> return (Just a)

705 yieldk a r = case cmp res a of

706 GT -> go (Just res) r

707 _ -> go (Just a) r

708 in foldStream defState yieldk single stop m1

709

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

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

712m !! i = go i m

713 where

714 go n m1 =

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

716 | otherwise = return Nothing

717 yieldk a x | n < 0 = return Nothing

718 | n == 0 = return $ Just a

719 | otherwise = go (n - 1) x

720 in foldStream defState yieldk single (return Nothing) m1

721

722{-# INLINE lookup #-}

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

724lookup e m = go m

725 where

726 go m1 =

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

728 | otherwise = return Nothing

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

730 | otherwise = go x

731 in foldStream defState yieldk single (return Nothing) m1

732

733{-# INLINE findM #-}

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

735findM p m = go m

736 where

737 go m1 =

738 let single a = do

739 b <- p a

740 if b then return $ Just a else return Nothing

741 yieldk a x = do

742 b <- p a

743 if b then return $ Just a else go x

744 in foldStream defState yieldk single (return Nothing) m1

745

746{-# INLINE find #-}

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

748find p = findM (return . p)

749

750{-# INLINE findIndices #-}

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

752findIndices p = go 0

753 where

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

755 let single a | p a = sng offset

756 | otherwise = stp

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

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

759 go (offset + 1) x

760 in foldStream (adaptState st) yieldk single stp m1

761

762------------------------------------------------------------------------------

763-- Map and Fold

764------------------------------------------------------------------------------

765

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

767-- output of the action.

768{-# INLINE mapM_ #-}

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

770mapM_ f m = go m

771 where

772 go m1 =

773 let stop = return ()

774 single a = void (f a)

775 yieldk a r = f a >> go r

776 in foldStream defState yieldk single stop m1

777

778------------------------------------------------------------------------------

779-- Converting folds

780------------------------------------------------------------------------------

781

782{-# INLINABLE toList #-}

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

784toList = foldr (:) []

785

786{-# INLINE toStreamK #-}

787toStreamK :: Stream m a -> Stream m a

788toStreamK = id

789

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

791{-# INLINE hoist #-}

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

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

794hoist f str =

795 mkStream $ \st yld sng stp ->

796 let single = return . sng

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

798 stop = return stp

799 state = adaptState st

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

801

802-------------------------------------------------------------------------------

803-- Transformation by folding (Scans)

804-------------------------------------------------------------------------------

805

806{-# INLINE scanlx' #-}

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

808scanlx' step begin done m =

809 cons (done begin) $ go m begin

810 where

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

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

813 yieldk a r =

814 let s = step acc a

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

816 in foldStream (adaptState st) yieldk single stp m1

817

818{-# INLINE scanl' #-}

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

820scanl' step begin = scanlx' step begin id

821

822-------------------------------------------------------------------------------

823-- Filtering

824-------------------------------------------------------------------------------

825

826{-# INLINE filter #-}

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

828filter p m = go m

829 where

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

831 let single a | p a = sng a

832 | otherwise = stp

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

834 | otherwise = foldStream st yieldk single stp r

835 in foldStream st yieldk single stp m1

836

837{-# INLINE take #-}

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

839take n m = go n m

840 where

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

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

843 in if n1 <= 0

844 then stp

845 else foldStream st yieldk sng stp m1

846

847{-# INLINE takeWhile #-}

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

849takeWhile p m = go m

850 where

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

852 let single a | p a = sng a

853 | otherwise = stp

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

855 | otherwise = stp

856 in foldStream st yieldk single stp m1

857

858{-# INLINE drop #-}

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

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

861 where

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

863 let single _ = stp

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

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

866 in if n1 <= 0

867 then foldStreamShared st yld sng stp m1

868 else foldStreamShared st yieldk single stp m1

869

870{-# INLINE dropWhile #-}

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

872dropWhile p m = go m

873 where

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

875 let single a | p a = stp

876 | otherwise = sng a

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

878 | otherwise = yld a r

879 in foldStream st yieldk single stp m1

880

881-------------------------------------------------------------------------------

882-- Mapping

883-------------------------------------------------------------------------------

884

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

886-- other stream types to overload it.

887{-# INLINE sequence #-}

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

889sequence m = go m

890 where

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

892 let single ma = ma >>= sng

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

894 in foldStream (adaptState st) yieldk single stp m1

895

896-------------------------------------------------------------------------------

897-- Inserting

898-------------------------------------------------------------------------------

899

900{-# INLINE intersperseM #-}

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

902intersperseM a m = prependingStart m

903 where

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

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

906 in foldStream st yieldk sng stp m1

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

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

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

910 in foldStream st yieldk single stp m2

911

912{-# INLINE intersperse #-}

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

914intersperse a = intersperseM (return a)

915

916{-# INLINE insertBy #-}

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

918insertBy cmp x m = go m

919 where

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

921 let single a = case cmp x a of

922 GT -> yld a (yield x)

923 _ -> yld x (yield a)

924 stop = yld x nil

925 yieldk a r = case cmp x a of

926 GT -> yld a (go r)

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

928 in foldStream st yieldk single stop m1

929

930------------------------------------------------------------------------------

931-- Deleting

932------------------------------------------------------------------------------

933

934{-# INLINE deleteBy #-}

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

936deleteBy eq x m = go m

937 where

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

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

940 yieldk a r = if eq x a

941 then foldStream st yld sng stp r

942 else yld a (go r)

943 in foldStream st yieldk single stp m1

944

945------------------------------------------------------------------------------

946-- Reordering

947------------------------------------------------------------------------------

948

949{-# INLINE reverse #-}

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

951reverse = foldlS (flip cons) nil

952

953-------------------------------------------------------------------------------

954-- Map and Filter

955-------------------------------------------------------------------------------

956

957{-# INLINE mapMaybe #-}

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

959mapMaybe f m = go m

960 where

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

962 let single a = case f a of

963 Just b -> sng b

964 Nothing -> stp

965 yieldk a r = case f a of

966 Just b -> yld b $ go r

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

968 in foldStream (adaptState st) yieldk single stp m1

969

970------------------------------------------------------------------------------

971-- Serial Zipping

972------------------------------------------------------------------------------

973

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

975--

976-- @since 0.1.0

977{-# INLINABLE zipWith #-}

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

979zipWith f = go

980 where

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

982 let merge a ra =

983 let single2 b = sng (f a b)

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

985 in foldStream (adaptState st) yield2 single2 stp my

986 let single1 a = merge a nil

987 yield1 = merge

988 foldStream (adaptState st) yield1 single1 stp mx

989

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

991--

992-- @since 0.1.0

993{-# INLINABLE zipWithM #-}

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

995zipWithM f m1 m2 = go m1 m2

996 where

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

998 let merge a ra =

999 let runIt x = foldStream st yld sng stp x

1000 single2 b = f a b >>= sng

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

1002 in foldStream (adaptState st) yield2 single2 stp my

1003 let single1 a = merge a nil

1004 yield1 = merge

1005 foldStream (adaptState st) yield1 single1 stp mx

1006

1007------------------------------------------------------------------------------

1008-- Merging

1009------------------------------------------------------------------------------

1010

1011{-# INLINE mergeByM #-}

1012mergeByM

1013 :: (IsStream t, Monad m)

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

1015mergeByM cmp = go

1016 where

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

1018 let mergeWithY a ra =

1019 let stop2 = foldStream st yld sng stp mx

1020 single2 b = do

1021 r <- cmp a b

1022 case r of

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

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

1025 yield2 b rb = do

1026 r <- cmp a b

1027 case r of

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

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

1030 in foldStream st yield2 single2 stop2 my

1031 let stopX = foldStream st yld sng stp my

1032 singleX a = mergeWithY a nil

1033 yieldX = mergeWithY

1034 foldStream st yieldX singleX stopX mx

1035

1036{-# INLINABLE mergeBy #-}

1037mergeBy

1038 :: (IsStream t, Monad m)

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

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

1041

1042------------------------------------------------------------------------------

1043-- Transformation comprehensions

1044------------------------------------------------------------------------------

1045

1046{-# INLINE the #-}

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

1048the m = do

1049 r <- uncons m

1050 case r of

1051 Nothing -> return Nothing

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

1053 where

1054 go h m1 =

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

1056 | otherwise = return Nothing

1057 yieldk a r | h == a = go h r

1058 | otherwise = return Nothing

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

1060

1061------------------------------------------------------------------------------

1062-- Alternative & MonadPlus

1063------------------------------------------------------------------------------

1064

1065_alt :: Stream m a -> Stream m a -> Stream m a

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

1067 let stop = foldStream st yld sng stp m2

1068 in foldStream st yld sng stop m1

1069

1070------------------------------------------------------------------------------

1071-- MonadReader

1072------------------------------------------------------------------------------

1073

1074{-# INLINABLE withLocal #-}

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

1076withLocal f m =

1077 mkStream $ \st yld sng stp ->

1078 let single = local f . sng

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

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

1081

1082------------------------------------------------------------------------------

1083-- MonadError

1084------------------------------------------------------------------------------

1085

1086{-

1087-- XXX handle and test cross thread state transfer

1088withCatchError

1089 :: MonadError e m

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

1091withCatchError m h =

1092 mkStream $ \_ stp sng yld ->

1093 let run x = unStream x Nothing stp sng yieldk

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

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

1096 in handle $ run m

1097-}