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

git.iko.soy/streamly

on 585b825 help sign in

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

1130 lines 36 kB history raw

1{-# LANGUAGE UndecidableInstances #-}

3#include "inline.hs"

5-- |

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

8--

9-- License : BSD3

10-- Maintainer : streamly@composewell.com

11-- Stability : experimental

12-- Portability : GHC

13--

14--

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

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

17--

18-- @

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

20-- @

21--

22module Streamly.Internal.Data.Stream.StreamK

23 (

24 -- * A class for streams

25 IsStream (..)

26 , adapt

28 -- * The stream type

29 , Stream(..)

31 -- * Construction Primitives

32 , mkStream

33 , nil

34 , nilM

35 , cons

36 , (.:)

38 -- * Elimination Primitives

39 , foldStream

40 , foldStreamShared

42 -- * Transformation Primitives

43 , unShare

45 -- * Deconstruction

46 , uncons

48 -- * Generation

49 -- ** Unfolds

50 , unfoldr

51 , unfoldrM

53 -- ** Specialized Generation

54 , repeat

55 , repeatM

56 , replicate

57 , replicateM

58 , fromIndices

59 , fromIndicesM

60 , iterate

61 , iterateM

63 -- ** Conversions

64 , yield

65 , yieldM

66 , fromFoldable

67 , fromList

68 , fromStreamK

70 -- * foldr/build

71 , foldrS

72 , foldrSM

73 , buildS

74 , buildM

75 , augmentS

76 , augmentSM

78 -- * Elimination

79 -- ** General Folds

80 , foldr

81 , foldr1

82 , foldrM

83 , foldrT

85 , foldl'

86 , foldlM'

87 , foldlS

88 , foldlT

89 , foldlx'

90 , foldlMx'

92 -- ** Specialized Folds

93 , drain

94 , null

95 , head

96 , tail

97 , init

98 , elem

99 , notElem

100 , all

101 , any

102 , last

103 , minimum

104 , minimumBy

105 , maximum

106 , maximumBy

107 , findIndices

108 , lookup

109 , findM

110 , find

111 , (!!)

112

113 -- ** Map and Fold

114 , mapM_

115

116 -- ** Conversions

117 , toList

118 , toStreamK

119 , hoist

120

121 -- * Transformation

122 -- ** By folding (scans)

123 , scanl'

124 , scanlx'

125

126 -- ** Filtering

127 , filter

128 , take

129 , takeWhile

130 , drop

131 , dropWhile

132

133 -- ** Mapping

134 , map

135 , mapM

136 , mapMSerial

137 , sequence

138

139 -- ** Inserting

140 , intersperseM

141 , intersperse

142 , insertBy

143

144 -- ** Deleting

145 , deleteBy

146

147 -- ** Reordering

148 , reverse

149

150 -- ** Map and Filter

151 , mapMaybe

152

153 -- ** Zipping

154 , zipWith

155 , zipWithM

156

157 -- ** Merging

158 , mergeBy

159 , mergeByM

160

161 -- ** Nesting

162 , concatMapBy

163 , concatMap

164 , bindWith

165 , apWith

166 , apSerial

167 , apSerialDiscardFst

168 , apSerialDiscardSnd

169

170 -- ** Transformation comprehensions

171 , the

172

173 -- * Semigroup Style Composition

174 , serial

175

176 -- * Utilities

177 , consMStream

178 , withLocal

179 , mfix

180

181 -- * Deprecated

182 , Streaming -- deprecated

183 , once -- deprecated

184 )

185where

186

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

188import Control.Monad (void, join)

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

190import Data.Function (fix)

191import Prelude

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

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

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

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

196import qualified Prelude

197

198import Streamly.Internal.Data.SVar

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

200

201-------------------------------------------------------------------------------

202-- Deconstruction

203-------------------------------------------------------------------------------

204

205{-# INLINE uncons #-}

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

207uncons m =

208 let stop = return Nothing

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

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

211 in foldStream defState yieldk single stop m

212

213-------------------------------------------------------------------------------

214-- Generation

215-------------------------------------------------------------------------------

216

217{-# INLINE unfoldr #-}

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

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

220 let go s =

221 case next s of

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

223 Nothing -> stp

224 in go s0

225

226{-# INLINE unfoldrM #-}

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

228unfoldrM step = go

229 where

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

231 r <- step s

232 case r of

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

234 Nothing -> stp

235

236{-

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

238--

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

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

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

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

243-- pure stream of monadic actions.

244--

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

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

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

248--

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

250unfoldrA = undefined

251-}

252

253-------------------------------------------------------------------------------

254-- Special generation

255-------------------------------------------------------------------------------

256

257-- | Same as yieldM

258--

259-- @since 0.2.0

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

261{-# INLINE once #-}

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

263once = yieldM

264

265-- |

266-- @

267-- repeatM = fix . cons

268-- repeatM = cycle1 . yield

269-- @

270--

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

272--

273-- /Internal/

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

275repeatM = go

276 where go m = m |: go m

277

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

279--

280-- /Internal/

281{-# INLINE repeat #-}

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

283repeat a = let x = cons a x in x

284

285{-# INLINE replicateM #-}

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

287replicateM n m = go n

288 where

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

290

291{-# INLINE replicate #-}

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

293replicate n a = go n

294 where

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

296

297{-# INLINE fromIndicesM #-}

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

299fromIndicesM gen = go 0

300 where

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

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

303

304{-# INLINE fromIndices #-}

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

306fromIndices gen = go 0

307 where

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

309

310{-# INLINE iterate #-}

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

312iterate step = fromStream . go

313 where

314 go s = cons s (go (step s))

315

316{-# INLINE iterateM #-}

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

318iterateM step = go

319 where

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

321 next <- s

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

323

324-------------------------------------------------------------------------------

325-- Conversions

326-------------------------------------------------------------------------------

327

328-- |

329-- @

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

331-- @

332--

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

334--

335-- @since 0.2.0

336{-# INLINE fromFoldable #-}

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

338fromFoldable = Prelude.foldr cons nil

339

340{-# INLINE fromList #-}

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

342fromList = fromFoldable

343

344{-# INLINE fromStreamK #-}

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

346fromStreamK = fromStream

347

348-------------------------------------------------------------------------------

349-- Elimination by Folding

350-------------------------------------------------------------------------------

351

352-- | Lazy right associative fold.

353{-# INLINE foldr #-}

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

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

356

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

358{-# INLINE foldrT #-}

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

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

361foldrT step final m = go m

362 where

363 go m1 = do

364 res <- lift $ uncons m1

365 case res of

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

367 Nothing -> final

368

369{-# INLINE foldr1 #-}

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

371foldr1 step m = do

372 r <- uncons m

373 case r of

374 Nothing -> return Nothing

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

376 where

377 go p m1 =

378 let stp = return p

379 single a = return $ step a p

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

381 in foldStream defState yieldk single stp m1

382

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

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

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

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

387--

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

389{-# INLINE foldlx' #-}

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

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

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

393 where

394 {-# NOINLINE get #-}

395 get :: t m x -> m b

396 get m1 =

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

398 -- okay?

399 let single = return . done

400 -- XXX this is foldSingleton. why foldStreamShared?

401 in foldStreamShared undefined undefined single undefined m1

402

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

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

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

406 -- continuations is much more efficient.

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

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

409 let stop = sng acc

410 single a = sng $ step acc a

411 -- XXX this is foldNonEmptyStream

412 yieldk a r = foldStream defState yld sng undefined $

413 go r (step acc a)

414 in foldStream defState yieldk single stop m1

415

416-- | Strict left associative fold.

417{-# INLINE foldl' #-}

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

419foldl' step begin = foldlx' step begin id

420

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

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

423{-# INLINABLE foldlMx' #-}

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

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

426foldlMx' step begin done m = go begin m

427 where

428 go !acc m1 =

429 let stop = acc >>= done

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

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

432 in foldStream defState yieldk single stop m1

433

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

435{-# INLINE foldlM' #-}

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

437foldlM' step begin = foldlMx' step begin return

438

439-- | Lazy left fold to a stream.

440{-# INLINE foldlS #-}

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

442foldlS step begin m = go begin m

443 where

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

445 let run x = foldStream st yld sng stp x

446 stop = run acc

447 single a = run $ step acc a

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

449 in foldStream (adaptState st) yieldk single stop rest

450

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

452{-# INLINE foldlT #-}

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

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

455foldlT step begin m = go begin m

456 where

457 go acc m1 = do

458 res <- lift $ uncons m1

459 case res of

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

461 Nothing -> acc

462

463------------------------------------------------------------------------------

464-- Specialized folds

465------------------------------------------------------------------------------

466

467-- XXX use foldrM to implement folds where possible

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

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

470-- an optimization opportunity that we can exploit.

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

472

473-- |

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

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

476{-# INLINE drain #-}

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

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

479{-

480drain = go

481 where

482 go m1 =

483 let stop = return ()

484 single _ = return ()

485 yieldk _ r = go r

486 in foldStream defState yieldk single stop m1

487-}

488

489{-# INLINE null #-}

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

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

492null m =

493 let stop = return True

494 single _ = return False

495 yieldk _ _ = return False

496 in foldStream defState yieldk single stop m

497

498{-# INLINE head #-}

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

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

501head m =

502 let stop = return Nothing

503 single a = return (Just a)

504 yieldk a _ = return (Just a)

505 in foldStream defState yieldk single stop m

506

507{-# INLINE tail #-}

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

509tail m =

510 let stop = return Nothing

511 single _ = return $ Just nil

512 yieldk _ r = return $ Just r

513 in foldStream defState yieldk single stop m

514

515{-# INLINE headPartial #-}

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

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

518

519{-# INLINE tailPartial #-}

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

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

522 let stop = error "tail of nil"

523 single _ = stp

524 yieldk _ r = foldStream st yld sng stp r

525 in foldStream st yieldk single stop m

526

527-- | We can define cyclic structures using @let@:

528--

529-- >>> let (a, b) = ([1, b], head a) in (a, b)

530-- ([1,1],1)

531--

532-- The function @fix@ defined as:

533--

534-- > fix f = let x = f x in x

535--

536-- ensures that the argument of a function and its output refer to the same

537-- lazy value @x@ i.e. the same location in memory. Thus @x@ can be defined

538-- in terms of itself, creating structures with cyclic references.

539--

540-- >>> f ~(a, b) = ([1, b], head a)

541-- >>> fix f

542-- ([1,1],1)

543--

544-- 'Control.Monad.mfix' is essentially the same as @fix@ but for monadic

545-- values.

546--

547-- Using 'mfix' for streams we can construct a stream in which each element of

548-- the stream is defined in a cyclic fashion. The argument of the function

549-- being fixed represents the current element of the stream which is being

550-- returned by the stream monad. Thus, we can use the argument to construct

551-- itself.

552--

553-- In the following example, the argument @action@ of the function @f@

554-- represents the tuple @(x,y)@ returned by it in a given iteration. We define

555-- the first element of the tuple in terms of the second.

556--

557-- @

558-- import Streamly.Internal.Data.Stream.IsStream as Stream

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

560--

561-- main = do

562-- Stream.mapM_ print $ Stream.mfix f

563--

564-- where

565--

566-- f action = do

567-- let incr n act = fmap ((+n) . snd) $ unsafeInterleaveIO act

568-- x <- Stream.fromListM [incr 1 action, incr 2 action]

569-- y <- Stream.fromList [4,5]

570-- return (x, y)

571-- @

572--

573-- Note: you cannot achieve this by just changing the order of the monad

574-- statements because that would change the order in which the stream elements

575-- are generated.

576--

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

578-- 'unsafeInterleaveIO' on @action@ because IO monad is strict.

579--

580-- /Internal/

581

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

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

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

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

586 in foldStream st yieldk single stp xs

587

588 where

589

590 -- fix the head element of the stream

591 xs = fix (f . headPartial)

592

593 -- now fix the tail recursively

594 ys = mfix (tailPartial . f)

595

596{-# INLINE init #-}

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

598init m = go1 m

599 where

600 go1 m1 = do

601 r <- uncons m1

602 case r of

603 Nothing -> return Nothing

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

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

606 let single _ = sng p

607 yieldk a x = yld p $ go a x

608 in foldStream defState yieldk single stp m1

609

610{-# INLINE elem #-}

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

612elem e m = go m

613 where

614 go m1 =

615 let stop = return False

616 single a = return (a == e)

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

618 in foldStream defState yieldk single stop m1

619

620{-# INLINE notElem #-}

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

622notElem e m = go m

623 where

624 go m1 =

625 let stop = return True

626 single a = return (a /= e)

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

628 in foldStream defState yieldk single stop m1

629

630{-# INLINABLE all #-}

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

632all p m = go m

633 where

634 go m1 =

635 let single a | p a = return True

636 | otherwise = return False

637 yieldk a r | p a = go r

638 | otherwise = return False

639 in foldStream defState yieldk single (return True) m1

640

641{-# INLINABLE any #-}

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

643any p m = go m

644 where

645 go m1 =

646 let single a | p a = return True

647 | otherwise = return False

648 yieldk a r | p a = return True

649 | otherwise = go r

650 in foldStream defState yieldk single (return False) m1

651

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

653{-# INLINE last #-}

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

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

656

657{-# INLINE minimum #-}

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

659minimum m = go Nothing m

660 where

661 go Nothing m1 =

662 let stop = return Nothing

663 single a = return (Just a)

664 yieldk a r = go (Just a) r

665 in foldStream defState yieldk single stop m1

666

667 go (Just res) m1 =

668 let stop = return (Just res)

669 single a =

670 if res <= a

671 then return (Just res)

672 else return (Just a)

673 yieldk a r =

674 if res <= a

675 then go (Just res) r

676 else go (Just a) r

677 in foldStream defState yieldk single stop m1

678

679{-# INLINE minimumBy #-}

680minimumBy

681 :: (IsStream t, Monad m)

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

683minimumBy cmp m = go Nothing m

684 where

685 go Nothing m1 =

686 let stop = return Nothing

687 single a = return (Just a)

688 yieldk a r = go (Just a) r

689 in foldStream defState yieldk single stop m1

690

691 go (Just res) m1 =

692 let stop = return (Just res)

693 single a = case cmp res a of

694 GT -> return (Just a)

695 _ -> return (Just res)

696 yieldk a r = case cmp res a of

697 GT -> go (Just a) r

698 _ -> go (Just res) r

699 in foldStream defState yieldk single stop m1

700

701{-# INLINE maximum #-}

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

703maximum m = go Nothing m

704 where

705 go Nothing m1 =

706 let stop = return Nothing

707 single a = return (Just a)

708 yieldk a r = go (Just a) r

709 in foldStream defState yieldk single stop m1

710

711 go (Just res) m1 =

712 let stop = return (Just res)

713 single a =

714 if res <= a

715 then return (Just a)

716 else return (Just res)

717 yieldk a r =

718 if res <= a

719 then go (Just a) r

720 else go (Just res) r

721 in foldStream defState yieldk single stop m1

722

723{-# INLINE maximumBy #-}

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

725maximumBy cmp m = go Nothing m

726 where

727 go Nothing m1 =

728 let stop = return Nothing

729 single a = return (Just a)

730 yieldk a r = go (Just a) r

731 in foldStream defState yieldk single stop m1

732

733 go (Just res) m1 =

734 let stop = return (Just res)

735 single a = case cmp res a of

736 GT -> return (Just res)

737 _ -> return (Just a)

738 yieldk a r = case cmp res a of

739 GT -> go (Just res) r

740 _ -> go (Just a) r

741 in foldStream defState yieldk single stop m1

742

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

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

745m !! i = go i m

746 where

747 go n m1 =

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

749 | otherwise = return Nothing

750 yieldk a x | n < 0 = return Nothing

751 | n == 0 = return $ Just a

752 | otherwise = go (n - 1) x

753 in foldStream defState yieldk single (return Nothing) m1

754

755{-# INLINE lookup #-}

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

757lookup e m = go m

758 where

759 go m1 =

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

761 | otherwise = return Nothing

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

763 | otherwise = go x

764 in foldStream defState yieldk single (return Nothing) m1

765

766{-# INLINE findM #-}

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

768findM p m = go m

769 where

770 go m1 =

771 let single a = do

772 b <- p a

773 if b then return $ Just a else return Nothing

774 yieldk a x = do

775 b <- p a

776 if b then return $ Just a else go x

777 in foldStream defState yieldk single (return Nothing) m1

778

779{-# INLINE find #-}

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

781find p = findM (return . p)

782

783{-# INLINE findIndices #-}

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

785findIndices p = go 0

786 where

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

788 let single a | p a = sng offset

789 | otherwise = stp

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

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

792 go (offset + 1) x

793 in foldStream (adaptState st) yieldk single stp m1

794

795------------------------------------------------------------------------------

796-- Map and Fold

797------------------------------------------------------------------------------

798

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

800-- output of the action.

801{-# INLINE mapM_ #-}

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

803mapM_ f m = go m

804 where

805 go m1 =

806 let stop = return ()

807 single a = void (f a)

808 yieldk a r = f a >> go r

809 in foldStream defState yieldk single stop m1

810

811------------------------------------------------------------------------------

812-- Converting folds

813------------------------------------------------------------------------------

814

815{-# INLINABLE toList #-}

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

817toList = foldr (:) []

818

819{-# INLINE toStreamK #-}

820toStreamK :: Stream m a -> Stream m a

821toStreamK = id

822

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

824{-# INLINE hoist #-}

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

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

827hoist f str =

828 mkStream $ \st yld sng stp ->

829 let single = return . sng

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

831 stop = return stp

832 state = adaptState st

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

834

835-------------------------------------------------------------------------------

836-- Transformation by folding (Scans)

837-------------------------------------------------------------------------------

838

839{-# INLINE scanlx' #-}

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

841scanlx' step begin done m =

842 cons (done begin) $ go m begin

843 where

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

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

846 yieldk a r =

847 let s = step acc a

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

849 in foldStream (adaptState st) yieldk single stp m1

850

851{-# INLINE scanl' #-}

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

853scanl' step begin = scanlx' step begin id

854

855-------------------------------------------------------------------------------

856-- Filtering

857-------------------------------------------------------------------------------

858

859{-# INLINE filter #-}

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

861filter p m = go m

862 where

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

864 let single a | p a = sng a

865 | otherwise = stp

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

867 | otherwise = foldStream st yieldk single stp r

868 in foldStream st yieldk single stp m1

869

870{-# INLINE take #-}

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

872take n m = go n m

873 where

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

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

876 in if n1 <= 0

877 then stp

878 else foldStream st yieldk sng stp m1

879

880{-# INLINE takeWhile #-}

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

882takeWhile p m = go m

883 where

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

885 let single a | p a = sng a

886 | otherwise = stp

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

888 | otherwise = stp

889 in foldStream st yieldk single stp m1

890

891{-# INLINE drop #-}

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

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

894 where

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

896 let single _ = stp

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

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

899 in if n1 <= 0

900 then foldStreamShared st yld sng stp m1

901 else foldStreamShared st yieldk single stp m1

902

903{-# INLINE dropWhile #-}

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

905dropWhile p m = go m

906 where

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

908 let single a | p a = stp

909 | otherwise = sng a

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

911 | otherwise = yld a r

912 in foldStream st yieldk single stp m1

913

914-------------------------------------------------------------------------------

915-- Mapping

916-------------------------------------------------------------------------------

917

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

919-- other stream types to overload it.

920{-# INLINE sequence #-}

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

922sequence m = go m

923 where

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

925 let single ma = ma >>= sng

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

927 in foldStream (adaptState st) yieldk single stp m1

928

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

930-- Inserting

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

932

933{-# INLINE intersperseM #-}

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

935intersperseM a m = prependingStart m

936 where

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

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

939 in foldStream st yieldk sng stp m1

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

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

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

943 in foldStream st yieldk single stp m2

944

945{-# INLINE intersperse #-}

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

947intersperse a = intersperseM (return a)

948

949{-# INLINE insertBy #-}

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

951insertBy cmp x m = go m

952 where

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

954 let single a = case cmp x a of

955 GT -> yld a (yield x)

956 _ -> yld x (yield a)

957 stop = yld x nil

958 yieldk a r = case cmp x a of

959 GT -> yld a (go r)

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

961 in foldStream st yieldk single stop m1

962

963------------------------------------------------------------------------------

964-- Deleting

965------------------------------------------------------------------------------

966

967{-# INLINE deleteBy #-}

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

969deleteBy eq x m = go m

970 where

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

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

973 yieldk a r = if eq x a

974 then foldStream st yld sng stp r

975 else yld a (go r)

976 in foldStream st yieldk single stp m1

977

978------------------------------------------------------------------------------

979-- Reordering

980------------------------------------------------------------------------------

981

982{-# INLINE reverse #-}

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

984reverse = foldlS (flip cons) nil

985

986-------------------------------------------------------------------------------

987-- Map and Filter

988-------------------------------------------------------------------------------

989

990{-# INLINE mapMaybe #-}

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

992mapMaybe f m = go m

993 where

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

995 let single a = case f a of

996 Just b -> sng b

997 Nothing -> stp

998 yieldk a r = case f a of

999 Just b -> yld b $ go r

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

1001 in foldStream (adaptState st) yieldk single stp m1

1002

1003------------------------------------------------------------------------------

1004-- Serial Zipping

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

1006

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

1008--

1009-- @since 0.1.0

1010{-# INLINABLE zipWith #-}

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

1012zipWith f = go

1013 where

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

1015 let merge a ra =

1016 let single2 b = sng (f a b)

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

1018 in foldStream (adaptState st) yield2 single2 stp my

1019 let single1 a = merge a nil

1020 yield1 = merge

1021 foldStream (adaptState st) yield1 single1 stp mx

1022

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

1024--

1025-- @since 0.1.0

1026{-# INLINABLE zipWithM #-}

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

1028zipWithM f m1 m2 = go m1 m2

1029 where

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

1031 let merge a ra =

1032 let runIt x = foldStream st yld sng stp x

1033 single2 b = f a b >>= sng

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

1035 in foldStream (adaptState st) yield2 single2 stp my

1036 let single1 a = merge a nil

1037 yield1 = merge

1038 foldStream (adaptState st) yield1 single1 stp mx

1039

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

1041-- Merging

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

1043

1044{-# INLINE mergeByM #-}

1045mergeByM

1046 :: (IsStream t, Monad m)

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

1048mergeByM cmp = go

1049 where

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

1051 let mergeWithY a ra =

1052 let stop2 = foldStream st yld sng stp mx

1053 single2 b = do

1054 r <- cmp a b

1055 case r of

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

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

1058 yield2 b rb = do

1059 r <- cmp a b

1060 case r of

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

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

1063 in foldStream st yield2 single2 stop2 my

1064 let stopX = foldStream st yld sng stp my

1065 singleX a = mergeWithY a nil

1066 yieldX = mergeWithY

1067 foldStream st yieldX singleX stopX mx

1068

1069{-# INLINABLE mergeBy #-}

1070mergeBy

1071 :: (IsStream t, Monad m)

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

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

1074

1075------------------------------------------------------------------------------

1076-- Transformation comprehensions

1077------------------------------------------------------------------------------

1078

1079{-# INLINE the #-}

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

1081the m = do

1082 r <- uncons m

1083 case r of

1084 Nothing -> return Nothing

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

1086 where

1087 go h m1 =

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

1089 | otherwise = return Nothing

1090 yieldk a r | h == a = go h r

1091 | otherwise = return Nothing

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

1093

1094------------------------------------------------------------------------------

1095-- Alternative & MonadPlus

1096------------------------------------------------------------------------------

1097

1098_alt :: Stream m a -> Stream m a -> Stream m a

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

1100 let stop = foldStream st yld sng stp m2

1101 in foldStream st yld sng stop m1

1102

1103------------------------------------------------------------------------------

1104-- MonadReader

1105------------------------------------------------------------------------------

1106

1107{-# INLINABLE withLocal #-}

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

1109withLocal f m =

1110 mkStream $ \st yld sng stp ->

1111 let single = local f . sng

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

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

1114

1115------------------------------------------------------------------------------

1116-- MonadError

1117------------------------------------------------------------------------------

1118

1119{-

1120-- XXX handle and test cross thread state transfer

1121withCatchError

1122 :: MonadError e m

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

1124withCatchError m h =

1125 mkStream $ \_ stp sng yld ->

1126 let run x = unStream x Nothing stp sng yieldk

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

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

1129 in handle $ run m

1130-}