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

git.iko.soy/streamly

on 90ce159 help sign in

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

1080 lines 34 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

71 -- ** Conversions

72 , yield

73 , yieldM

74 , fromFoldable

75 , fromList

76 , fromStreamK

78 -- * foldr/build

79 , foldrS

80 , foldrSM

81 , buildS

82 , buildM

83 , augmentS

84 , augmentSM

86 -- * Elimination

87 -- ** General Folds

88 , foldr

89 , foldr1

90 , foldrM

91 , foldrT

93 , foldl'

94 , foldlM'

95 , foldlS

96 , foldlT

97 , foldlx'

98 , foldlMx'

100 -- ** Specialized Folds

101 , drain

102 , null

103 , head

104 , tail

105 , init

106 , elem

107 , notElem

108 , all

109 , any

110 , last

111 , minimum

112 , minimumBy

113 , maximum

114 , maximumBy

115 , findIndices

116 , lookup

117 , findM

118 , find

119 , (!!)

120

121 -- ** Map and Fold

122 , mapM_

123

124 -- ** Conversions

125 , toList

126 , toStreamK

127 , hoist

128

129 -- * Transformation

130 -- ** By folding (scans)

131 , scanl'

132 , scanlx'

133

134 -- ** Filtering

135 , filter

136 , take

137 , takeWhile

138 , drop

139 , dropWhile

140

141 -- ** Mapping

142 , map

143 , mapM

144 , mapMSerial

145 , sequence

146

147 -- ** Inserting

148 , intersperseM

149 , intersperse

150 , insertBy

151

152 -- ** Deleting

153 , deleteBy

154

155 -- ** Reordering

156 , reverse

157

158 -- ** Map and Filter

159 , mapMaybe

160

161 -- ** Zipping

162 , zipWith

163 , zipWithM

164

165 -- ** Merging

166 , mergeBy

167 , mergeByM

168

169 -- ** Nesting

170 , concatMapBy

171 , concatMap

172 , bindWith

173

174 -- ** Transformation comprehensions

175 , the

176

177 -- * Semigroup Style Composition

178 , serial

179

180 -- * Utilities

181 , consMStream

182 , withLocal

183 , mfix

184

185 -- * Deprecated

186 , Streaming -- deprecated

187 , once -- deprecated

188 )

189where

190

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

192import Control.Monad (void, join)

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

194import Data.Function (fix)

195import Prelude

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

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

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

199 foldr1, (!!), replicate, reverse, concatMap)

200import qualified Prelude

201

202import Streamly.Internal.Data.SVar

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

204

205-------------------------------------------------------------------------------

206-- Deconstruction

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

208

209{-# INLINE uncons #-}

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

211uncons m =

212 let stop = return Nothing

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

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

215 in foldStream defState yieldk single stop m

216

217-------------------------------------------------------------------------------

218-- Generation

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

220

221{-# INLINE unfoldr #-}

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

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

224 let go s =

225 case next s of

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

227 Nothing -> stp

228 in go s0

229

230{-# INLINE unfoldrM #-}

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

232unfoldrM step = go

233 where

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

235 r <- step s

236 case r of

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

238 Nothing -> stp

239

240{-

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

242--

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

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

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

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

247-- pure stream of monadic actions.

248--

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

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

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

252--

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

254unfoldrA = undefined

255-}

256

257-------------------------------------------------------------------------------

258-- Special generation

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

260

261-- | Same as yieldM

262--

263-- @since 0.2.0

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

265{-# INLINE once #-}

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

267once = yieldM

268

269-- |

270-- @

271-- repeatM = fix . cons

272-- repeatM = cycle1 . yield

273-- @

274--

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

276--

277-- /Internal/

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

279repeatM = go

280 where go m = m |: go m

281

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

283--

284-- /Internal/

285{-# INLINE repeat #-}

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

287repeat a = let x = cons a x in x

288

289{-# INLINE replicateM #-}

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

291replicateM n m = go n

292 where

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

294

295{-# INLINE replicate #-}

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

297replicate n a = go n

298 where

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

300

301{-# INLINE fromIndicesM #-}

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

303fromIndicesM gen = go 0

304 where

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

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

307

308{-# INLINE fromIndices #-}

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

310fromIndices gen = go 0

311 where

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

313

314-------------------------------------------------------------------------------

315-- Conversions

316-------------------------------------------------------------------------------

317

318-- |

319-- @

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

321-- @

322--

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

324--

325-- @since 0.2.0

326{-# INLINE fromFoldable #-}

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

328fromFoldable = Prelude.foldr cons nil

329

330{-# INLINE fromList #-}

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

332fromList = fromFoldable

333

334{-# INLINE fromStreamK #-}

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

336fromStreamK = fromStream

337

338-------------------------------------------------------------------------------

339-- Elimination by Folding

340-------------------------------------------------------------------------------

341

342-- | Lazy right associative fold.

343{-# INLINE foldr #-}

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

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

346

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

348{-# INLINE foldrT #-}

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

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

351foldrT step final m = go m

352 where

353 go m1 = do

354 res <- lift $ uncons m1

355 case res of

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

357 Nothing -> final

358

359{-# INLINE foldr1 #-}

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

361foldr1 step m = do

362 r <- uncons m

363 case r of

364 Nothing -> return Nothing

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

366 where

367 go p m1 =

368 let stp = return p

369 single a = return $ step a p

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

371 in foldStream defState yieldk single stp m1

372

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

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

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

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

377--

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

379{-# INLINE foldlx' #-}

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

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

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

383 where

384 {-# NOINLINE get #-}

385 get :: t m x -> m b

386 get m1 =

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

388 -- okay?

389 let single = return . done

390 -- XXX this is foldSingleton. why foldStreamShared?

391 in foldStreamShared undefined undefined single undefined m1

392

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

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

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

396 -- continuations is much more efficient.

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

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

399 let stop = sng acc

400 single a = sng $ step acc a

401 -- XXX this is foldNonEmptyStream

402 yieldk a r = foldStream defState yld sng undefined $

403 go r (step acc a)

404 in foldStream defState yieldk single stop m1

405

406-- | Strict left associative fold.

407{-# INLINE foldl' #-}

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

409foldl' step begin = foldlx' step begin id

410

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

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

413{-# INLINABLE foldlMx' #-}

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

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

416foldlMx' step begin done m = go begin m

417 where

418 go !acc m1 =

419 let stop = acc >>= done

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

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

422 in foldStream defState yieldk single stop m1

423

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

425{-# INLINE foldlM' #-}

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

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

428

429-- | Lazy left fold to a stream.

430{-# INLINE foldlS #-}

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

432foldlS step begin m = go begin m

433 where

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

435 let run x = foldStream st yld sng stp x

436 stop = run acc

437 single a = run $ step acc a

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

439 in foldStream (adaptState st) yieldk single stop rest

440

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

442{-# INLINE foldlT #-}

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

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

445foldlT step begin m = go begin m

446 where

447 go acc m1 = do

448 res <- lift $ uncons m1

449 case res of

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

451 Nothing -> acc

452

453------------------------------------------------------------------------------

454-- Specialized folds

455------------------------------------------------------------------------------

456

457-- XXX use foldrM to implement folds where possible

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

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

460-- an optimization opportunity that we can exploit.

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

462

463-- |

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

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

466{-# INLINE drain #-}

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

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

469{-

470drain = go

471 where

472 go m1 =

473 let stop = return ()

474 single _ = return ()

475 yieldk _ r = go r

476 in foldStream defState yieldk single stop m1

477-}

478

479{-# INLINE null #-}

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

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

482null m =

483 let stop = return True

484 single _ = return False

485 yieldk _ _ = return False

486 in foldStream defState yieldk single stop m

487

488{-# INLINE head #-}

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

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

491head m =

492 let stop = return Nothing

493 single a = return (Just a)

494 yieldk a _ = return (Just a)

495 in foldStream defState yieldk single stop m

496

497{-# INLINE tail #-}

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

499tail m =

500 let stop = return Nothing

501 single _ = return $ Just nil

502 yieldk _ r = return $ Just r

503 in foldStream defState yieldk single stop m

504

505{-# INLINE headPartial #-}

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

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

508

509{-# INLINE tailPartial #-}

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

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

512 let stop = error "tail of nil"

513 single _ = stp

514 yieldk _ r = foldStream st yld sng stp r

515 in foldStream st yieldk single stop m

516

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

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

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

520--

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

522-- example:

523--

524-- @

525-- import Streamly.Internal.Prelude as S

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

527--

528-- main = do

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

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

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

532-- return (a, b)

533-- @

534--

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

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

537

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

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

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

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

542 in foldStream st yieldk single stp xs

543 where xs = fix (f . headPartial)

544 ys = mfix (tailPartial . f)

545

546{-# INLINE init #-}

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

548init m = go1 m

549 where

550 go1 m1 = do

551 r <- uncons m1

552 case r of

553 Nothing -> return Nothing

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

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

556 let single _ = sng p

557 yieldk a x = yld p $ go a x

558 in foldStream defState yieldk single stp m1

559

560{-# INLINE elem #-}

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

562elem e m = go m

563 where

564 go m1 =

565 let stop = return False

566 single a = return (a == e)

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

568 in foldStream defState yieldk single stop m1

569

570{-# INLINE notElem #-}

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

572notElem e m = go m

573 where

574 go m1 =

575 let stop = return True

576 single a = return (a /= e)

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

578 in foldStream defState yieldk single stop m1

579

580{-# INLINABLE all #-}

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

582all p m = go m

583 where

584 go m1 =

585 let single a | p a = return True

586 | otherwise = return False

587 yieldk a r | p a = go r

588 | otherwise = return False

589 in foldStream defState yieldk single (return True) m1

590

591{-# INLINABLE any #-}

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

593any p m = go m

594 where

595 go m1 =

596 let single a | p a = return True

597 | otherwise = return False

598 yieldk a r | p a = return True

599 | otherwise = go r

600 in foldStream defState yieldk single (return False) m1

601

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

603{-# INLINE last #-}

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

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

606

607{-# INLINE minimum #-}

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

609minimum m = go Nothing m

610 where

611 go Nothing m1 =

612 let stop = return Nothing

613 single a = return (Just a)

614 yieldk a r = go (Just a) r

615 in foldStream defState yieldk single stop m1

616

617 go (Just res) m1 =

618 let stop = return (Just res)

619 single a =

620 if res <= a

621 then return (Just res)

622 else return (Just a)

623 yieldk a r =

624 if res <= a

625 then go (Just res) r

626 else go (Just a) r

627 in foldStream defState yieldk single stop m1

628

629{-# INLINE minimumBy #-}

630minimumBy

631 :: (IsStream t, Monad m)

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

633minimumBy cmp m = go Nothing m

634 where

635 go Nothing m1 =

636 let stop = return Nothing

637 single a = return (Just a)

638 yieldk a r = go (Just a) r

639 in foldStream defState yieldk single stop m1

640

641 go (Just res) m1 =

642 let stop = return (Just res)

643 single a = case cmp res a of

644 GT -> return (Just a)

645 _ -> return (Just res)

646 yieldk a r = case cmp res a of

647 GT -> go (Just a) r

648 _ -> go (Just res) r

649 in foldStream defState yieldk single stop m1

650

651{-# INLINE maximum #-}

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

653maximum m = go Nothing m

654 where

655 go Nothing m1 =

656 let stop = return Nothing

657 single a = return (Just a)

658 yieldk a r = go (Just a) r

659 in foldStream defState yieldk single stop m1

660

661 go (Just res) m1 =

662 let stop = return (Just res)

663 single a =

664 if res <= a

665 then return (Just a)

666 else return (Just res)

667 yieldk a r =

668 if res <= a

669 then go (Just a) r

670 else go (Just res) r

671 in foldStream defState yieldk single stop m1

672

673{-# INLINE maximumBy #-}

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

675maximumBy cmp m = go Nothing m

676 where

677 go Nothing m1 =

678 let stop = return Nothing

679 single a = return (Just a)

680 yieldk a r = go (Just a) r

681 in foldStream defState yieldk single stop m1

682

683 go (Just res) m1 =

684 let stop = return (Just res)

685 single a = case cmp res a of

686 GT -> return (Just res)

687 _ -> return (Just a)

688 yieldk a r = case cmp res a of

689 GT -> go (Just res) r

690 _ -> go (Just a) r

691 in foldStream defState yieldk single stop m1

692

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

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

695m !! i = go i m

696 where

697 go n m1 =

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

699 | otherwise = return Nothing

700 yieldk a x | n < 0 = return Nothing

701 | n == 0 = return $ Just a

702 | otherwise = go (n - 1) x

703 in foldStream defState yieldk single (return Nothing) m1

704

705{-# INLINE lookup #-}

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

707lookup e m = go m

708 where

709 go m1 =

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

711 | otherwise = return Nothing

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

713 | otherwise = go x

714 in foldStream defState yieldk single (return Nothing) m1

715

716{-# INLINE findM #-}

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

718findM p m = go m

719 where

720 go m1 =

721 let single a = do

722 b <- p a

723 if b then return $ Just a else return Nothing

724 yieldk a x = do

725 b <- p a

726 if b then return $ Just a else go x

727 in foldStream defState yieldk single (return Nothing) m1

728

729{-# INLINE find #-}

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

731find p = findM (return . p)

732

733{-# INLINE findIndices #-}

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

735findIndices p = go 0

736 where

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

738 let single a | p a = sng offset

739 | otherwise = stp

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

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

742 go (offset + 1) x

743 in foldStream (adaptState st) yieldk single stp m1

744

745------------------------------------------------------------------------------

746-- Map and Fold

747------------------------------------------------------------------------------

748

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

750-- output of the action.

751{-# INLINE mapM_ #-}

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

753mapM_ f m = go m

754 where

755 go m1 =

756 let stop = return ()

757 single a = void (f a)

758 yieldk a r = f a >> go r

759 in foldStream defState yieldk single stop m1

760

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

762-- Converting folds

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

764

765{-# INLINABLE toList #-}

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

767toList = foldr (:) []

768

769{-# INLINE toStreamK #-}

770toStreamK :: Stream m a -> Stream m a

771toStreamK = id

772

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

774{-# INLINE hoist #-}

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

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

777hoist f str =

778 mkStream $ \st yld sng stp ->

779 let single = return . sng

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

781 stop = return stp

782 state = adaptState st

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

784

785-------------------------------------------------------------------------------

786-- Transformation by folding (Scans)

787-------------------------------------------------------------------------------

788

789{-# INLINE scanlx' #-}

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

791scanlx' step begin done m =

792 cons (done begin) $ go m begin

793 where

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

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

796 yieldk a r =

797 let s = step acc a

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

799 in foldStream (adaptState st) yieldk single stp m1

800

801{-# INLINE scanl' #-}

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

803scanl' step begin = scanlx' step begin id

804

805-------------------------------------------------------------------------------

806-- Filtering

807-------------------------------------------------------------------------------

808

809{-# INLINE filter #-}

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

811filter p m = go m

812 where

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

814 let single a | p a = sng a

815 | otherwise = stp

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

817 | otherwise = foldStream st yieldk single stp r

818 in foldStream st yieldk single stp m1

819

820{-# INLINE take #-}

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

822take n m = go n m

823 where

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

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

826 in if n1 <= 0

827 then stp

828 else foldStream st yieldk sng stp m1

829

830{-# INLINE takeWhile #-}

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

832takeWhile 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 = stp

839 in foldStream st yieldk single stp m1

840

841{-# INLINE drop #-}

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

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

844 where

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

846 let single _ = stp

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

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

849 in if n1 <= 0

850 then foldStreamShared st yld sng stp m1

851 else foldStreamShared st yieldk single stp m1

852

853{-# INLINE dropWhile #-}

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

855dropWhile p m = go m

856 where

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

858 let single a | p a = stp

859 | otherwise = sng a

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

861 | otherwise = yld a r

862 in foldStream st yieldk single stp m1

863

864-------------------------------------------------------------------------------

865-- Mapping

866-------------------------------------------------------------------------------

867

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

869-- other stream types to overload it.

870{-# INLINE sequence #-}

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

872sequence m = go m

873 where

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

875 let single ma = ma >>= sng

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

877 in foldStream (adaptState st) yieldk single stp m1

878

879-------------------------------------------------------------------------------

880-- Inserting

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

882

883{-# INLINE intersperseM #-}

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

885intersperseM a m = prependingStart m

886 where

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

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

889 in foldStream st yieldk sng stp m1

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

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

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

893 in foldStream st yieldk single stp m2

894

895{-# INLINE intersperse #-}

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

897intersperse a = intersperseM (return a)

898

899{-# INLINE insertBy #-}

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

901insertBy cmp x m = go m

902 where

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

904 let single a = case cmp x a of

905 GT -> yld a (yield x)

906 _ -> yld x (yield a)

907 stop = yld x nil

908 yieldk a r = case cmp x a of

909 GT -> yld a (go r)

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

911 in foldStream st yieldk single stop m1

912

913------------------------------------------------------------------------------

914-- Deleting

915------------------------------------------------------------------------------

916

917{-# INLINE deleteBy #-}

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

919deleteBy eq x m = go m

920 where

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

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

923 yieldk a r = if eq x a

924 then foldStream st yld sng stp r

925 else yld a (go r)

926 in foldStream st yieldk single stp m1

927

928------------------------------------------------------------------------------

929-- Reordering

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

931

932{-# INLINE reverse #-}

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

934reverse = foldlS (flip cons) nil

935

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

937-- Map and Filter

938-------------------------------------------------------------------------------

939

940{-# INLINE mapMaybe #-}

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

942mapMaybe f m = go m

943 where

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

945 let single a = case f a of

946 Just b -> sng b

947 Nothing -> stp

948 yieldk a r = case f a of

949 Just b -> yld b $ go r

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

951 in foldStream (adaptState st) yieldk single stp m1

952

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

954-- Serial Zipping

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

956

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

958--

959-- @since 0.1.0

960{-# INLINABLE zipWith #-}

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

962zipWith f = go

963 where

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

965 let merge a ra =

966 let single2 b = sng (f a b)

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

968 in foldStream (adaptState st) yield2 single2 stp my

969 let single1 a = merge a nil

970 yield1 = merge

971 foldStream (adaptState st) yield1 single1 stp mx

972

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

974--

975-- @since 0.1.0

976{-# INLINABLE zipWithM #-}

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

978zipWithM f m1 m2 = go m1 m2

979 where

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

981 let merge a ra =

982 let runIt x = foldStream st yld sng stp x

983 single2 b = f a b >>= sng

984 yield2 b rb = f a b >>= \x -> runIt (x `cons` 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------------------------------------------------------------------------------

991-- Merging

992------------------------------------------------------------------------------

993

994{-# INLINE mergeByM #-}

995mergeByM

996 :: (IsStream t, Monad m)

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

998mergeByM cmp = go

999 where

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

1001 let mergeWithY a ra =

1002 let stop2 = foldStream st yld sng stp mx

1003 single2 b = do

1004 r <- cmp a b

1005 case r of

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

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

1008 yield2 b rb = do

1009 r <- cmp a b

1010 case r of

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

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

1013 in foldStream st yield2 single2 stop2 my

1014 let stopX = foldStream st yld sng stp my

1015 singleX a = mergeWithY a nil

1016 yieldX = mergeWithY

1017 foldStream st yieldX singleX stopX mx

1018

1019{-# INLINABLE mergeBy #-}

1020mergeBy

1021 :: (IsStream t, Monad m)

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

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

1024

1025------------------------------------------------------------------------------

1026-- Transformation comprehensions

1027------------------------------------------------------------------------------

1028

1029{-# INLINE the #-}

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

1031the m = do

1032 r <- uncons m

1033 case r of

1034 Nothing -> return Nothing

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

1036 where

1037 go h m1 =

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

1039 | otherwise = return Nothing

1040 yieldk a r | h == a = go h r

1041 | otherwise = return Nothing

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

1043

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

1045-- Alternative & MonadPlus

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

1047

1048_alt :: Stream m a -> Stream m a -> Stream m a

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

1050 let stop = foldStream st yld sng stp m2

1051 in foldStream st yld sng stop m1

1052

1053------------------------------------------------------------------------------

1054-- MonadReader

1055------------------------------------------------------------------------------

1056

1057{-# INLINABLE withLocal #-}

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

1059withLocal f m =

1060 mkStream $ \st yld sng stp ->

1061 let single = local f . sng

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

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

1064

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

1066-- MonadError

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

1068

1069{-

1070-- XXX handle and test cross thread state transfer

1071withCatchError

1072 :: MonadError e m

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

1074withCatchError m h =

1075 mkStream $ \_ stp sng yld ->

1076 let run x = unStream x Nothing stp sng yieldk

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

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

1079 in handle $ run m

1080-}