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

git.iko.soy/streamly

on f6fd95a help sign in

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

1095 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

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

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

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

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

557 in foldStream st yieldk single stp xs

558 where xs = fix (f . headPartial)

559 ys = mfix (tailPartial . f)

560

561{-# INLINE init #-}

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

563init m = go1 m

564 where

565 go1 m1 = do

566 r <- uncons m1

567 case r of

568 Nothing -> return Nothing

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

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

571 let single _ = sng p

572 yieldk a x = yld p $ go a x

573 in foldStream defState yieldk single stp m1

574

575{-# INLINE elem #-}

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

577elem e m = go m

578 where

579 go m1 =

580 let stop = return False

581 single a = return (a == e)

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

583 in foldStream defState yieldk single stop m1

584

585{-# INLINE notElem #-}

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

587notElem e m = go m

588 where

589 go m1 =

590 let stop = return True

591 single a = return (a /= e)

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

593 in foldStream defState yieldk single stop m1

594

595{-# INLINABLE all #-}

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

597all p m = go m

598 where

599 go m1 =

600 let single a | p a = return True

601 | otherwise = return False

602 yieldk a r | p a = go r

603 | otherwise = return False

604 in foldStream defState yieldk single (return True) m1

605

606{-# INLINABLE any #-}

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

608any p m = go m

609 where

610 go m1 =

611 let single a | p a = return True

612 | otherwise = return False

613 yieldk a r | p a = return True

614 | otherwise = go r

615 in foldStream defState yieldk single (return False) m1

616

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

618{-# INLINE last #-}

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

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

621

622{-# INLINE minimum #-}

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

624minimum m = go Nothing m

625 where

626 go Nothing m1 =

627 let stop = return Nothing

628 single a = return (Just a)

629 yieldk a r = go (Just a) r

630 in foldStream defState yieldk single stop m1

631

632 go (Just res) m1 =

633 let stop = return (Just res)

634 single a =

635 if res <= a

636 then return (Just res)

637 else return (Just a)

638 yieldk a r =

639 if res <= a

640 then go (Just res) r

641 else go (Just a) r

642 in foldStream defState yieldk single stop m1

643

644{-# INLINE minimumBy #-}

645minimumBy

646 :: (IsStream t, Monad m)

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

648minimumBy cmp m = go Nothing m

649 where

650 go Nothing m1 =

651 let stop = return Nothing

652 single a = return (Just a)

653 yieldk a r = go (Just a) r

654 in foldStream defState yieldk single stop m1

655

656 go (Just res) m1 =

657 let stop = return (Just res)

658 single a = case cmp res a of

659 GT -> return (Just a)

660 _ -> return (Just res)

661 yieldk a r = case cmp res a of

662 GT -> go (Just a) r

663 _ -> go (Just res) r

664 in foldStream defState yieldk single stop m1

665

666{-# INLINE maximum #-}

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

668maximum m = go Nothing m

669 where

670 go Nothing m1 =

671 let stop = return Nothing

672 single a = return (Just a)

673 yieldk a r = go (Just a) r

674 in foldStream defState yieldk single stop m1

675

676 go (Just res) m1 =

677 let stop = return (Just res)

678 single a =

679 if res <= a

680 then return (Just a)

681 else return (Just res)

682 yieldk a r =

683 if res <= a

684 then go (Just a) r

685 else go (Just res) r

686 in foldStream defState yieldk single stop m1

687

688{-# INLINE maximumBy #-}

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

690maximumBy cmp m = go Nothing m

691 where

692 go Nothing m1 =

693 let stop = return Nothing

694 single a = return (Just a)

695 yieldk a r = go (Just a) r

696 in foldStream defState yieldk single stop m1

697

698 go (Just res) m1 =

699 let stop = return (Just res)

700 single a = case cmp res a of

701 GT -> return (Just res)

702 _ -> return (Just a)

703 yieldk a r = case cmp res a of

704 GT -> go (Just res) r

705 _ -> go (Just a) r

706 in foldStream defState yieldk single stop m1

707

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

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

710m !! i = go i m

711 where

712 go n m1 =

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

714 | otherwise = return Nothing

715 yieldk a x | n < 0 = return Nothing

716 | n == 0 = return $ Just a

717 | otherwise = go (n - 1) x

718 in foldStream defState yieldk single (return Nothing) m1

719

720{-# INLINE lookup #-}

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

722lookup e m = go m

723 where

724 go m1 =

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

726 | otherwise = return Nothing

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

728 | otherwise = go x

729 in foldStream defState yieldk single (return Nothing) m1

730

731{-# INLINE findM #-}

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

733findM p m = go m

734 where

735 go m1 =

736 let single a = do

737 b <- p a

738 if b then return $ Just a else return Nothing

739 yieldk a x = do

740 b <- p a

741 if b then return $ Just a else go x

742 in foldStream defState yieldk single (return Nothing) m1

743

744{-# INLINE find #-}

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

746find p = findM (return . p)

747

748{-# INLINE findIndices #-}

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

750findIndices p = go 0

751 where

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

753 let single a | p a = sng offset

754 | otherwise = stp

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

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

757 go (offset + 1) x

758 in foldStream (adaptState st) yieldk single stp m1

759

760------------------------------------------------------------------------------

761-- Map and Fold

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

763

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

765-- output of the action.

766{-# INLINE mapM_ #-}

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

768mapM_ f m = go m

769 where

770 go m1 =

771 let stop = return ()

772 single a = void (f a)

773 yieldk a r = f a >> go r

774 in foldStream defState yieldk single stop m1

775

776------------------------------------------------------------------------------

777-- Converting folds

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

779

780{-# INLINABLE toList #-}

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

782toList = foldr (:) []

783

784{-# INLINE toStreamK #-}

785toStreamK :: Stream m a -> Stream m a

786toStreamK = id

787

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

789{-# INLINE hoist #-}

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

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

792hoist f str =

793 mkStream $ \st yld sng stp ->

794 let single = return . sng

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

796 stop = return stp

797 state = adaptState st

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

799

800-------------------------------------------------------------------------------

801-- Transformation by folding (Scans)

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

803

804{-# INLINE scanlx' #-}

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

806scanlx' step begin done m =

807 cons (done begin) $ go m begin

808 where

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

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

811 yieldk a r =

812 let s = step acc a

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

814 in foldStream (adaptState st) yieldk single stp m1

815

816{-# INLINE scanl' #-}

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

818scanl' step begin = scanlx' step begin id

819

820-------------------------------------------------------------------------------

821-- Filtering

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

823

824{-# INLINE filter #-}

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

826filter p m = go m

827 where

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

829 let single a | p a = sng a

830 | otherwise = stp

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

832 | otherwise = foldStream st yieldk single stp r

833 in foldStream st yieldk single stp m1

834

835{-# INLINE take #-}

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

837take n m = go n m

838 where

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

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

841 in if n1 <= 0

842 then stp

843 else foldStream st yieldk sng stp m1

844

845{-# INLINE takeWhile #-}

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

847takeWhile p m = go m

848 where

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

850 let single a | p a = sng a

851 | otherwise = stp

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

853 | otherwise = stp

854 in foldStream st yieldk single stp m1

855

856{-# INLINE drop #-}

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

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

859 where

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

861 let single _ = stp

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

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

864 in if n1 <= 0

865 then foldStreamShared st yld sng stp m1

866 else foldStreamShared st yieldk single stp m1

867

868{-# INLINE dropWhile #-}

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

870dropWhile p m = go m

871 where

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

873 let single a | p a = stp

874 | otherwise = sng a

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

876 | otherwise = yld a r

877 in foldStream st yieldk single stp m1

878

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

880-- Mapping

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

882

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

884-- other stream types to overload it.

885{-# INLINE sequence #-}

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

887sequence m = go m

888 where

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

890 let single ma = ma >>= sng

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

892 in foldStream (adaptState st) yieldk single stp m1

893

894-------------------------------------------------------------------------------

895-- Inserting

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

897

898{-# INLINE intersperseM #-}

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

900intersperseM a m = prependingStart m

901 where

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

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

904 in foldStream st yieldk sng stp m1

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

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

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

908 in foldStream st yieldk single stp m2

909

910{-# INLINE intersperse #-}

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

912intersperse a = intersperseM (return a)

913

914{-# INLINE insertBy #-}

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

916insertBy cmp x m = go m

917 where

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

919 let single a = case cmp x a of

920 GT -> yld a (yield x)

921 _ -> yld x (yield a)

922 stop = yld x nil

923 yieldk a r = case cmp x a of

924 GT -> yld a (go r)

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

926 in foldStream st yieldk single stop m1

927

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

929-- Deleting

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

931

932{-# INLINE deleteBy #-}

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

934deleteBy eq x m = go m

935 where

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

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

938 yieldk a r = if eq x a

939 then foldStream st yld sng stp r

940 else yld a (go r)

941 in foldStream st yieldk single stp m1

942

943------------------------------------------------------------------------------

944-- Reordering

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

946

947{-# INLINE reverse #-}

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

949reverse = foldlS (flip cons) nil

950

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

952-- Map and Filter

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

954

955{-# INLINE mapMaybe #-}

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

957mapMaybe f m = go m

958 where

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

960 let single a = case f a of

961 Just b -> sng b

962 Nothing -> stp

963 yieldk a r = case f a of

964 Just b -> yld b $ go r

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

966 in foldStream (adaptState st) yieldk single stp m1

967

968------------------------------------------------------------------------------

969-- Serial Zipping

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

971

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

973--

974-- @since 0.1.0

975{-# INLINABLE zipWith #-}

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

977zipWith f = go

978 where

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

980 let merge a ra =

981 let single2 b = sng (f a b)

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

983 in foldStream (adaptState st) yield2 single2 stp my

984 let single1 a = merge a nil

985 yield1 = merge

986 foldStream (adaptState st) yield1 single1 stp mx

987

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

989--

990-- @since 0.1.0

991{-# INLINABLE zipWithM #-}

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

993zipWithM f m1 m2 = go m1 m2

994 where

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

996 let merge a ra =

997 let runIt x = foldStream st yld sng stp x

998 single2 b = f a b >>= sng

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

1000 in foldStream (adaptState st) yield2 single2 stp my

1001 let single1 a = merge a nil

1002 yield1 = merge

1003 foldStream (adaptState st) yield1 single1 stp mx

1004

1005------------------------------------------------------------------------------

1006-- Merging

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

1008

1009{-# INLINE mergeByM #-}

1010mergeByM

1011 :: (IsStream t, Monad m)

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

1013mergeByM cmp = go

1014 where

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

1016 let mergeWithY a ra =

1017 let stop2 = foldStream st yld sng stp mx

1018 single2 b = do

1019 r <- cmp a b

1020 case r of

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

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

1023 yield2 b rb = do

1024 r <- cmp a b

1025 case r of

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

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

1028 in foldStream st yield2 single2 stop2 my

1029 let stopX = foldStream st yld sng stp my

1030 singleX a = mergeWithY a nil

1031 yieldX = mergeWithY

1032 foldStream st yieldX singleX stopX mx

1033

1034{-# INLINABLE mergeBy #-}

1035mergeBy

1036 :: (IsStream t, Monad m)

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

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

1039

1040------------------------------------------------------------------------------

1041-- Transformation comprehensions

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

1043

1044{-# INLINE the #-}

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

1046the m = do

1047 r <- uncons m

1048 case r of

1049 Nothing -> return Nothing

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

1051 where

1052 go h m1 =

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

1054 | otherwise = return Nothing

1055 yieldk a r | h == a = go h r

1056 | otherwise = return Nothing

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

1058

1059------------------------------------------------------------------------------

1060-- Alternative & MonadPlus

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

1062

1063_alt :: Stream m a -> Stream m a -> Stream m a

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

1065 let stop = foldStream st yld sng stp m2

1066 in foldStream st yld sng stop m1

1067

1068------------------------------------------------------------------------------

1069-- MonadReader

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

1071

1072{-# INLINABLE withLocal #-}

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

1074withLocal f m =

1075 mkStream $ \st yld sng stp ->

1076 let single = local f . sng

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

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

1079

1080------------------------------------------------------------------------------

1081-- MonadError

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

1083

1084{-

1085-- XXX handle and test cross thread state transfer

1086withCatchError

1087 :: MonadError e m

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

1089withCatchError m h =

1090 mkStream $ \_ stp sng yld ->

1091 let run x = unStream x Nothing stp sng yieldk

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

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

1094 in handle $ run m

1095-}