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

git.iko.soy/streamly

on 70fb992 help sign in

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

1101 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 , apWith

175 , apSerial

176 , apSerialDiscardFst

177 , apSerialDiscardSnd

178

179 -- ** Transformation comprehensions

180 , the

181

182 -- * Semigroup Style Composition

183 , serial

184

185 -- * Utilities

186 , consMStream

187 , withLocal

188 , mfix

189

190 -- * Deprecated

191 , Streaming -- deprecated

192 , once -- deprecated

193 )

194where

195

196import Control.Monad.Trans.Class (MonadTrans(lift))

197import Control.Monad (void, join)

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

199import Data.Function (fix)

200import Prelude

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

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

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

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

205import qualified Prelude

206

207import Streamly.Internal.Data.SVar

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

209

210-------------------------------------------------------------------------------

211-- Deconstruction

212-------------------------------------------------------------------------------

213

214{-# INLINE uncons #-}

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

216uncons m =

217 let stop = return Nothing

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

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

220 in foldStream defState yieldk single stop m

221

222-------------------------------------------------------------------------------

223-- Generation

224-------------------------------------------------------------------------------

225

226{-# INLINE unfoldr #-}

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

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

229 let go s =

230 case next s of

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

232 Nothing -> stp

233 in go s0

234

235{-# INLINE unfoldrM #-}

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

237unfoldrM step = go

238 where

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

240 r <- step s

241 case r of

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

243 Nothing -> stp

244

245{-

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

247--

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

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

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

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

252-- pure stream of monadic actions.

253--

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

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

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

257--

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

259unfoldrA = undefined

260-}

261

262-------------------------------------------------------------------------------

263-- Special generation

264-------------------------------------------------------------------------------

265

266-- | Same as yieldM

267--

268-- @since 0.2.0

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

270{-# INLINE once #-}

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

272once = yieldM

273

274-- |

275-- @

276-- repeatM = fix . cons

277-- repeatM = cycle1 . yield

278-- @

279--

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

281--

282-- /Internal/

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

284repeatM = go

285 where go m = m |: go m

286

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

288--

289-- /Internal/

290{-# INLINE repeat #-}

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

292repeat a = let x = cons a x in x

293

294{-# INLINE replicateM #-}

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

296replicateM n m = go n

297 where

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

299

300{-# INLINE replicate #-}

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

302replicate n a = go n

303 where

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

305

306{-# INLINE fromIndicesM #-}

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

308fromIndicesM gen = go 0

309 where

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

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

312

313{-# INLINE fromIndices #-}

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

315fromIndices gen = go 0

316 where

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

318

319{-# INLINE iterate #-}

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

321iterate step = fromStream . go

322 where

323 go s = cons s (go (step s))

324

325{-# INLINE iterateM #-}

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

327iterateM step = go

328 where

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

330 next <- s

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

332

333-------------------------------------------------------------------------------

334-- Conversions

335-------------------------------------------------------------------------------

336

337-- |

338-- @

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

340-- @

341--

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

343--

344-- @since 0.2.0

345{-# INLINE fromFoldable #-}

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

347fromFoldable = Prelude.foldr cons nil

348

349{-# INLINE fromList #-}

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

351fromList = fromFoldable

352

353{-# INLINE fromStreamK #-}

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

355fromStreamK = fromStream

356

357-------------------------------------------------------------------------------

358-- Elimination by Folding

359-------------------------------------------------------------------------------

360

361-- | Lazy right associative fold.

362{-# INLINE foldr #-}

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

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

365

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

367{-# INLINE foldrT #-}

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

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

370foldrT step final m = go m

371 where

372 go m1 = do

373 res <- lift $ uncons m1

374 case res of

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

376 Nothing -> final

377

378{-# INLINE foldr1 #-}

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

380foldr1 step m = do

381 r <- uncons m

382 case r of

383 Nothing -> return Nothing

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

385 where

386 go p m1 =

387 let stp = return p

388 single a = return $ step a p

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

390 in foldStream defState yieldk single stp m1

391

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

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

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

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

396--

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

398{-# INLINE foldlx' #-}

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

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

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

402 where

403 {-# NOINLINE get #-}

404 get :: t m x -> m b

405 get m1 =

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

407 -- okay?

408 let single = return . done

409 -- XXX this is foldSingleton. why foldStreamShared?

410 in foldStreamShared undefined undefined single undefined m1

411

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

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

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

415 -- continuations is much more efficient.

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

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

418 let stop = sng acc

419 single a = sng $ step acc a

420 -- XXX this is foldNonEmptyStream

421 yieldk a r = foldStream defState yld sng undefined $

422 go r (step acc a)

423 in foldStream defState yieldk single stop m1

424

425-- | Strict left associative fold.

426{-# INLINE foldl' #-}

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

428foldl' step begin = foldlx' step begin id

429

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

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

432{-# INLINABLE foldlMx' #-}

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

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

435foldlMx' step begin done m = go begin m

436 where

437 go !acc m1 =

438 let stop = acc >>= done

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

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

441 in foldStream defState yieldk single stop m1

442

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

444{-# INLINE foldlM' #-}

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

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

447

448-- | Lazy left fold to a stream.

449{-# INLINE foldlS #-}

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

451foldlS step begin m = go begin m

452 where

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

454 let run x = foldStream st yld sng stp x

455 stop = run acc

456 single a = run $ step acc a

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

458 in foldStream (adaptState st) yieldk single stop rest

459

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

461{-# INLINE foldlT #-}

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

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

464foldlT step begin m = go begin m

465 where

466 go acc m1 = do

467 res <- lift $ uncons m1

468 case res of

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

470 Nothing -> acc

471

472------------------------------------------------------------------------------

473-- Specialized folds

474------------------------------------------------------------------------------

475

476-- XXX use foldrM to implement folds where possible

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

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

479-- an optimization opportunity that we can exploit.

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

481

482-- |

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

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

485{-# INLINE drain #-}

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

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

488{-

489drain = go

490 where

491 go m1 =

492 let stop = return ()

493 single _ = return ()

494 yieldk _ r = go r

495 in foldStream defState yieldk single stop m1

496-}

497

498{-# INLINE null #-}

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

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

501null m =

502 let stop = return True

503 single _ = return False

504 yieldk _ _ = return False

505 in foldStream defState yieldk single stop m

506

507{-# INLINE head #-}

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

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

510head m =

511 let stop = return Nothing

512 single a = return (Just a)

513 yieldk a _ = return (Just a)

514 in foldStream defState yieldk single stop m

515

516{-# INLINE tail #-}

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

518tail m =

519 let stop = return Nothing

520 single _ = return $ Just nil

521 yieldk _ r = return $ Just r

522 in foldStream defState yieldk single stop m

523

524{-# INLINE headPartial #-}

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

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

527

528{-# INLINE tailPartial #-}

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

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

531 let stop = error "tail of nil"

532 single _ = stp

533 yieldk _ r = foldStream st yld sng stp r

534 in foldStream st yieldk single stop m

535

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

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

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

539--

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

541-- example:

542--

543-- @

544-- import Streamly.Internal.Prelude as S

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

546--

547-- main = do

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

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

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

551-- return (a, b)

552-- @

553--

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

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

556--

557-- /Internal/

558

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

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

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

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

563 in foldStream st yieldk single stp xs

564 where xs = fix (f . headPartial)

565 ys = mfix (tailPartial . f)

566

567{-# INLINE init #-}

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

569init m = go1 m

570 where

571 go1 m1 = do

572 r <- uncons m1

573 case r of

574 Nothing -> return Nothing

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

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

577 let single _ = sng p

578 yieldk a x = yld p $ go a x

579 in foldStream defState yieldk single stp m1

580

581{-# INLINE elem #-}

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

583elem e m = go m

584 where

585 go m1 =

586 let stop = return False

587 single a = return (a == e)

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

589 in foldStream defState yieldk single stop m1

590

591{-# INLINE notElem #-}

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

593notElem e m = go m

594 where

595 go m1 =

596 let stop = return True

597 single a = return (a /= e)

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

599 in foldStream defState yieldk single stop m1

600

601{-# INLINABLE all #-}

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

603all p m = go m

604 where

605 go m1 =

606 let single a | p a = return True

607 | otherwise = return False

608 yieldk a r | p a = go r

609 | otherwise = return False

610 in foldStream defState yieldk single (return True) m1

611

612{-# INLINABLE any #-}

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

614any p m = go m

615 where

616 go m1 =

617 let single a | p a = return True

618 | otherwise = return False

619 yieldk a r | p a = return True

620 | otherwise = go r

621 in foldStream defState yieldk single (return False) m1

622

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

624{-# INLINE last #-}

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

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

627

628{-# INLINE minimum #-}

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

630minimum m = go Nothing m

631 where

632 go Nothing m1 =

633 let stop = return Nothing

634 single a = return (Just a)

635 yieldk a r = go (Just a) r

636 in foldStream defState yieldk single stop m1

637

638 go (Just res) m1 =

639 let stop = return (Just res)

640 single a =

641 if res <= a

642 then return (Just res)

643 else return (Just a)

644 yieldk a r =

645 if res <= a

646 then go (Just res) r

647 else go (Just a) r

648 in foldStream defState yieldk single stop m1

649

650{-# INLINE minimumBy #-}

651minimumBy

652 :: (IsStream t, Monad m)

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

654minimumBy cmp m = go Nothing m

655 where

656 go Nothing m1 =

657 let stop = return Nothing

658 single a = return (Just a)

659 yieldk a r = go (Just a) r

660 in foldStream defState yieldk single stop m1

661

662 go (Just res) m1 =

663 let stop = return (Just res)

664 single a = case cmp res a of

665 GT -> return (Just a)

666 _ -> return (Just res)

667 yieldk a r = case cmp res a of

668 GT -> go (Just a) r

669 _ -> go (Just res) r

670 in foldStream defState yieldk single stop m1

671

672{-# INLINE maximum #-}

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

674maximum m = go Nothing m

675 where

676 go Nothing m1 =

677 let stop = return Nothing

678 single a = return (Just a)

679 yieldk a r = go (Just a) r

680 in foldStream defState yieldk single stop m1

681

682 go (Just res) m1 =

683 let stop = return (Just res)

684 single a =

685 if res <= a

686 then return (Just a)

687 else return (Just res)

688 yieldk a r =

689 if res <= a

690 then go (Just a) r

691 else go (Just res) r

692 in foldStream defState yieldk single stop m1

693

694{-# INLINE maximumBy #-}

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

696maximumBy cmp m = go Nothing m

697 where

698 go Nothing m1 =

699 let stop = return Nothing

700 single a = return (Just a)

701 yieldk a r = go (Just a) r

702 in foldStream defState yieldk single stop m1

703

704 go (Just res) m1 =

705 let stop = return (Just res)

706 single a = case cmp res a of

707 GT -> return (Just res)

708 _ -> return (Just a)

709 yieldk a r = case cmp res a of

710 GT -> go (Just res) r

711 _ -> go (Just a) r

712 in foldStream defState yieldk single stop m1

713

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

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

716m !! i = go i m

717 where

718 go n m1 =

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

720 | otherwise = return Nothing

721 yieldk a x | n < 0 = return Nothing

722 | n == 0 = return $ Just a

723 | otherwise = go (n - 1) x

724 in foldStream defState yieldk single (return Nothing) m1

725

726{-# INLINE lookup #-}

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

728lookup e m = go m

729 where

730 go m1 =

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

732 | otherwise = return Nothing

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

734 | otherwise = go x

735 in foldStream defState yieldk single (return Nothing) m1

736

737{-# INLINE findM #-}

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

739findM p m = go m

740 where

741 go m1 =

742 let single a = do

743 b <- p a

744 if b then return $ Just a else return Nothing

745 yieldk a x = do

746 b <- p a

747 if b then return $ Just a else go x

748 in foldStream defState yieldk single (return Nothing) m1

749

750{-# INLINE find #-}

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

752find p = findM (return . p)

753

754{-# INLINE findIndices #-}

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

756findIndices p = go 0

757 where

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

759 let single a | p a = sng offset

760 | otherwise = stp

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

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

763 go (offset + 1) x

764 in foldStream (adaptState st) yieldk single stp m1

765

766------------------------------------------------------------------------------

767-- Map and Fold

768------------------------------------------------------------------------------

769

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

771-- output of the action.

772{-# INLINE mapM_ #-}

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

774mapM_ f m = go m

775 where

776 go m1 =

777 let stop = return ()

778 single a = void (f a)

779 yieldk a r = f a >> go r

780 in foldStream defState yieldk single stop m1

781

782------------------------------------------------------------------------------

783-- Converting folds

784------------------------------------------------------------------------------

785

786{-# INLINABLE toList #-}

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

788toList = foldr (:) []

789

790{-# INLINE toStreamK #-}

791toStreamK :: Stream m a -> Stream m a

792toStreamK = id

793

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

795{-# INLINE hoist #-}

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

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

798hoist f str =

799 mkStream $ \st yld sng stp ->

800 let single = return . sng

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

802 stop = return stp

803 state = adaptState st

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

805

806-------------------------------------------------------------------------------

807-- Transformation by folding (Scans)

808-------------------------------------------------------------------------------

809

810{-# INLINE scanlx' #-}

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

812scanlx' step begin done m =

813 cons (done begin) $ go m begin

814 where

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

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

817 yieldk a r =

818 let s = step acc a

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

820 in foldStream (adaptState st) yieldk single stp m1

821

822{-# INLINE scanl' #-}

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

824scanl' step begin = scanlx' step begin id

825

826-------------------------------------------------------------------------------

827-- Filtering

828-------------------------------------------------------------------------------

829

830{-# INLINE filter #-}

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

832filter p m = go m

833 where

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

835 let single a | p a = sng a

836 | otherwise = stp

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

838 | otherwise = foldStream st yieldk single stp r

839 in foldStream st yieldk single stp m1

840

841{-# INLINE take #-}

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

843take n m = go n m

844 where

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

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

847 in if n1 <= 0

848 then stp

849 else foldStream st yieldk sng stp m1

850

851{-# INLINE takeWhile #-}

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

853takeWhile p m = go m

854 where

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

856 let single a | p a = sng a

857 | otherwise = stp

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

859 | otherwise = stp

860 in foldStream st yieldk single stp m1

861

862{-# INLINE drop #-}

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

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

865 where

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

867 let single _ = stp

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

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

870 in if n1 <= 0

871 then foldStreamShared st yld sng stp m1

872 else foldStreamShared st yieldk single stp m1

873

874{-# INLINE dropWhile #-}

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

876dropWhile p m = go m

877 where

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

879 let single a | p a = stp

880 | otherwise = sng a

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

882 | otherwise = yld a r

883 in foldStream st yieldk single stp m1

884

885-------------------------------------------------------------------------------

886-- Mapping

887-------------------------------------------------------------------------------

888

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

890-- other stream types to overload it.

891{-# INLINE sequence #-}

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

893sequence m = go m

894 where

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

896 let single ma = ma >>= sng

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

898 in foldStream (adaptState st) yieldk single stp m1

899

900-------------------------------------------------------------------------------

901-- Inserting

902-------------------------------------------------------------------------------

903

904{-# INLINE intersperseM #-}

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

906intersperseM a m = prependingStart m

907 where

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

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

910 in foldStream st yieldk sng stp m1

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

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

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

914 in foldStream st yieldk single stp m2

915

916{-# INLINE intersperse #-}

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

918intersperse a = intersperseM (return a)

919

920{-# INLINE insertBy #-}

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

922insertBy cmp x m = go m

923 where

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

925 let single a = case cmp x a of

926 GT -> yld a (yield x)

927 _ -> yld x (yield a)

928 stop = yld x nil

929 yieldk a r = case cmp x a of

930 GT -> yld a (go r)

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

932 in foldStream st yieldk single stop m1

933

934------------------------------------------------------------------------------

935-- Deleting

936------------------------------------------------------------------------------

937

938{-# INLINE deleteBy #-}

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

940deleteBy eq x m = go m

941 where

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

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

944 yieldk a r = if eq x a

945 then foldStream st yld sng stp r

946 else yld a (go r)

947 in foldStream st yieldk single stp m1

948

949------------------------------------------------------------------------------

950-- Reordering

951------------------------------------------------------------------------------

952

953{-# INLINE reverse #-}

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

955reverse = foldlS (flip cons) nil

956

957-------------------------------------------------------------------------------

958-- Map and Filter

959-------------------------------------------------------------------------------

960

961{-# INLINE mapMaybe #-}

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

963mapMaybe f m = go m

964 where

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

966 let single a = case f a of

967 Just b -> sng b

968 Nothing -> stp

969 yieldk a r = case f a of

970 Just b -> yld b $ go r

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

972 in foldStream (adaptState st) yieldk single stp m1

973

974------------------------------------------------------------------------------

975-- Serial Zipping

976------------------------------------------------------------------------------

977

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

979--

980-- @since 0.1.0

981{-# INLINABLE zipWith #-}

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

983zipWith f = go

984 where

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

986 let merge a ra =

987 let single2 b = sng (f a b)

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

989 in foldStream (adaptState st) yield2 single2 stp my

990 let single1 a = merge a nil

991 yield1 = merge

992 foldStream (adaptState st) yield1 single1 stp mx

993

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

995--

996-- @since 0.1.0

997{-# INLINABLE zipWithM #-}

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

999zipWithM f m1 m2 = go m1 m2

1000 where

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

1002 let merge a ra =

1003 let runIt x = foldStream st yld sng stp x

1004 single2 b = f a b >>= sng

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

1006 in foldStream (adaptState st) yield2 single2 stp my

1007 let single1 a = merge a nil

1008 yield1 = merge

1009 foldStream (adaptState st) yield1 single1 stp mx

1010

1011------------------------------------------------------------------------------

1012-- Merging

1013------------------------------------------------------------------------------

1014

1015{-# INLINE mergeByM #-}

1016mergeByM

1017 :: (IsStream t, Monad m)

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

1019mergeByM cmp = go

1020 where

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

1022 let mergeWithY a ra =

1023 let stop2 = foldStream st yld sng stp mx

1024 single2 b = do

1025 r <- cmp a b

1026 case r of

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

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

1029 yield2 b rb = do

1030 r <- cmp a b

1031 case r of

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

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

1034 in foldStream st yield2 single2 stop2 my

1035 let stopX = foldStream st yld sng stp my

1036 singleX a = mergeWithY a nil

1037 yieldX = mergeWithY

1038 foldStream st yieldX singleX stopX mx

1039

1040{-# INLINABLE mergeBy #-}

1041mergeBy

1042 :: (IsStream t, Monad m)

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

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

1045

1046------------------------------------------------------------------------------

1047-- Transformation comprehensions

1048------------------------------------------------------------------------------

1049

1050{-# INLINE the #-}

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

1052the m = do

1053 r <- uncons m

1054 case r of

1055 Nothing -> return Nothing

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

1057 where

1058 go h m1 =

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

1060 | otherwise = return Nothing

1061 yieldk a r | h == a = go h r

1062 | otherwise = return Nothing

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

1064

1065------------------------------------------------------------------------------

1066-- Alternative & MonadPlus

1067------------------------------------------------------------------------------

1068

1069_alt :: Stream m a -> Stream m a -> Stream m a

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

1071 let stop = foldStream st yld sng stp m2

1072 in foldStream st yld sng stop m1

1073

1074------------------------------------------------------------------------------

1075-- MonadReader

1076------------------------------------------------------------------------------

1077

1078{-# INLINABLE withLocal #-}

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

1080withLocal f m =

1081 mkStream $ \st yld sng stp ->

1082 let single = local f . sng

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

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

1085

1086------------------------------------------------------------------------------

1087-- MonadError

1088------------------------------------------------------------------------------

1089

1090{-

1091-- XXX handle and test cross thread state transfer

1092withCatchError

1093 :: MonadError e m

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

1095withCatchError m h =

1096 mkStream $ \_ stp sng yld ->

1097 let run x = unStream x Nothing stp sng yieldk

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

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

1100 in handle $ run m

1101-}