5a1efe4 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Move StreamD's transform ops to StreamD.Transform.

5a1efe4Pranay Sashank6 years ago

src/Streamly/Internal/Data/Stream/StreamD.hs+2−1090

compare with master

unifiedsplit

@@ -301,25 +301,6 @@ module Streamly.Internal.Data.Stream.StreamD

301 )

302where

303

304−import Control.Concurrent (myThreadId)

305−import Control.Monad (void)

306−import Control.Monad.IO.Class (MonadIO(..))

307−import Control.Monad.Trans.Reader (ReaderT)

308−import Control.Monad.Trans.State.Strict (StateT)

309−import Control.Monad.Trans.Class (MonadTrans(lift))

310−import Data.Functor.Identity (Identity(..))

311−import Data.Maybe (fromJust, isJust)

312−import Streamly.Internal.Data.Time.Units

313− (TimeUnit64, toRelTime64, diffAbsTime64)

314−import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))

315−import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

316−

317−import qualified Streamly.Internal.Data.Pipe.Types as Pipe

318−import qualified Streamly.Internal.Data.Fold as FL

319−import qualified Streamly.Internal.Data.Stream.StreamK as K

320−import qualified Control.Monad.Trans.Reader as Reader

321−import qualified Control.Monad.Trans.State.Strict as State

322−

323import Prelude hiding

304import Prelude hiding

324 ( map, mapM, mapM_, repeat, foldr, last, take, filter

305 ( map, mapM, mapM_, repeat, foldr, last, take, filter

325 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem

306 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem

@@ -327,1079 +308,10 @@ import Prelude hiding

327 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat

308 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat

328 , reverse, iterate, splitAt)

309 , reverse, iterate, splitAt)

329import Streamly.Internal.Data.Stream.StreamD.Type

310import Streamly.Internal.Data.Stream.StreamD.Type

330−import Streamly.Internal.Data.Stream.StreamD.Common (takeWhileM, takeWhile)

311+import Streamly.Internal.Data.Stream.StreamD.Common (uncons)

331import Streamly.Internal.Data.Stream.StreamD.Generate

312import Streamly.Internal.Data.Stream.StreamD.Generate

332import Streamly.Internal.Data.Stream.StreamD.Eliminate

313import Streamly.Internal.Data.Stream.StreamD.Eliminate

333import Streamly.Internal.Data.Stream.StreamD.Exceptions

314import Streamly.Internal.Data.Stream.StreamD.Exceptions

334import Streamly.Internal.Data.Stream.StreamD.Flatten

315import Streamly.Internal.Data.Stream.StreamD.Flatten

335import Streamly.Internal.Data.Stream.StreamD.SplitGroup

316import Streamly.Internal.Data.Stream.StreamD.SplitGroup

336−import Streamly.Internal.Data.SVar

317+import Streamly.Internal.Data.Stream.StreamD.Transform

337−

338−-------------------------------------------------------------------------------

339−-- Hoisting the inner monad

340−-------------------------------------------------------------------------------

341−

342−{-# INLINE_NORMAL hoist #-}

343−hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a

344−hoist f (Stream step state) = (Stream step' state)

345− where

346− {-# INLINE_LATE step' #-}

347− step' gst st = do

348− r <- f $ step (adaptState gst) st

349− return $ case r of

350− Yield x s -> Yield x s

351− Skip s -> Skip s

352− Stop -> Stop

353−

354−{-# INLINE generally #-}

355−generally :: Monad m => Stream Identity a -> Stream m a

356−generally = hoist (return . runIdentity)

357−

358−{-# INLINE_NORMAL liftInner #-}

359−liftInner :: (Monad m, MonadTrans t, Monad (t m))

360− => Stream m a -> Stream (t m) a

361−liftInner (Stream step state) = Stream step' state

362− where

363− {-# INLINE_LATE step' #-}

364− step' gst st = do

365− r <- lift $ step (adaptState gst) st

366− return $ case r of

367− Yield x s -> Yield x s

368− Skip s -> Skip s

369− Stop -> Stop

370−

371−{-# INLINE_NORMAL runReaderT #-}

372−runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a

373−runReaderT env (Stream step state) = Stream step' (state, env)

374− where

375− {-# INLINE_LATE step' #-}

376− step' gst (st, action) = do

377− sv <- action

378− r <- Reader.runReaderT (step (adaptState gst) st) sv

379− return $ case r of

380− Yield x s -> Yield x (s, return sv)

381− Skip s -> Skip (s, return sv)

382− Stop -> Stop

383−

384−{-# INLINE_NORMAL evalStateT #-}

385−evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a

386−evalStateT initial (Stream step state) = Stream step' (state, initial)

387− where

388− {-# INLINE_LATE step' #-}

389− step' gst (st, action) = do

390− sv <- action

391− (r, sv') <- State.runStateT (step (adaptState gst) st) sv

392− return $ case r of

393− Yield x s -> Yield x (s, return sv')

394− Skip s -> Skip (s, return sv')

395− Stop -> Stop

396−

397−{-# INLINE_NORMAL runStateT #-}

398−runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)

399−runStateT initial (Stream step state) = Stream step' (state, initial)

400− where

401− {-# INLINE_LATE step' #-}

402− step' gst (st, action) = do

403− sv <- action

404− (r, sv') <- State.runStateT (step (adaptState gst) st) sv

405− return $ case r of

406− Yield x s -> Yield (sv', x) (s, return sv')

407− Skip s -> Skip (s, return sv')

408− Stop -> Stop

409−

410−------------------------------------------------------------------------------

411−-- Substreams

412−------------------------------------------------------------------------------

413−

414−{-# INLINE_NORMAL isPrefixOf #-}

415−isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool

416−isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)

417− where

418− go (sa, sb, Nothing) = do

419− r <- stepa defState sa

420− case r of

421− Yield x sa' -> go (sa', sb, Just x)

422− Skip sa' -> go (sa', sb, Nothing)

423− Stop -> return True

424−

425− go (sa, sb, Just x) = do

426− r <- stepb defState sb

427− case r of

428− Yield y sb' ->

429− if x == y

430− then go (sa, sb', Nothing)

431− else return False

432− Skip sb' -> go (sa, sb', Just x)

433− Stop -> return False

434−

435−{-# INLINE_NORMAL isSubsequenceOf #-}

436−isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool

437−isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)

438− where

439− go (sa, sb, Nothing) = do

440− r <- stepa defState sa

441− case r of

442− Yield x sa' -> go (sa', sb, Just x)

443− Skip sa' -> go (sa', sb, Nothing)

444− Stop -> return True

445−

446− go (sa, sb, Just x) = do

447− r <- stepb defState sb

448− case r of

449− Yield y sb' ->

450− if x == y

451− then go (sa, sb', Nothing)

452− else go (sa, sb', Just x)

453− Skip sb' -> go (sa, sb', Just x)

454− Stop -> return False

455−

456−{-# INLINE_NORMAL stripPrefix #-}

457−stripPrefix

458− :: (Eq a, Monad m)

459− => Stream m a -> Stream m a -> m (Maybe (Stream m a))

460−stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)

461− where

462− go (sa, sb, Nothing) = do

463− r <- stepa defState sa

464− case r of

465− Yield x sa' -> go (sa', sb, Just x)

466− Skip sa' -> go (sa', sb, Nothing)

467− Stop -> return $ Just (Stream stepb sb)

468−

469− go (sa, sb, Just x) = do

470− r <- stepb defState sb

471− case r of

472− Yield y sb' ->

473− if x == y

474− then go (sa, sb', Nothing)

475− else return Nothing

476− Skip sb' -> go (sa, sb', Just x)

477− Stop -> return Nothing

478−

479−------------------------------------------------------------------------------

480−-- Map and Fold

481−------------------------------------------------------------------------------

482−

483−-- | Execute a monadic action for each element of the 'Stream'

484−{-# INLINE_NORMAL mapM_ #-}

485−mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()

486−mapM_ m = drain . mapM m

487−

488−-------------------------------------------------------------------------------

489−-- General transformation

490−-------------------------------------------------------------------------------

491−

492−{-# INLINE_NORMAL transform #-}

493−transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b

494−transform (Pipe pstep1 pstep2 pstate) (Stream step state) =

495− Stream step' (Consume pstate, state)

496−

497− where

498−

499− {-# INLINE_LATE step' #-}

500−

501− step' gst (Consume pst, st) = pst `seq` do

502− r <- step (adaptState gst) st

503− case r of

504− Yield x s -> do

505− res <- pstep1 pst x

506− case res of

507− Pipe.Yield b pst' -> return $ Yield b (pst', s)

508− Pipe.Continue pst' -> return $ Skip (pst', s)

509− Skip s -> return $ Skip (Consume pst, s)

510− Stop -> return Stop

511−

512− step' _ (Produce pst, st) = pst `seq` do

513− res <- pstep2 pst

514− case res of

515− Pipe.Yield b pst' -> return $ Yield b (pst', st)

516− Pipe.Continue pst' -> return $ Skip (pst', st)

517−

518−------------------------------------------------------------------------------

519−-- Transformation by Folding (Scans)

520−------------------------------------------------------------------------------

521−

522−------------------------------------------------------------------------------

523−-- Prescans

524−------------------------------------------------------------------------------

525−

526−-- XXX Is a prescan useful, discarding the last step does not sound useful? I

527−-- am not sure about the utility of this function, so this is implemented but

528−-- not exposed. We can expose it if someone provides good reasons why this is

529−-- useful.

530−--

531−-- XXX We have to execute the stream one step ahead to know that we are at the

532−-- last step. The vector implementation of prescan executes the last fold step

533−-- but does not yield the result. This means we have executed the effect but

534−-- discarded value. This does not sound right. In this implementation we are

535−-- not executing the last fold step.

536−{-# INLINE_NORMAL prescanlM' #-}

537−prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

538−prescanlM' f mz (Stream step state) = Stream step' (state, mz)

539− where

540− {-# INLINE_LATE step' #-}

541− step' gst (st, prev) = do

542− r <- step (adaptState gst) st

543− case r of

544− Yield x s -> do

545− acc <- prev

546− return $ Yield acc (s, f acc x)

547− Skip s -> return $ Skip (s, prev)

548− Stop -> return Stop

549−

550−{-# INLINE prescanl' #-}

551−prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

552−prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)

553−

554−------------------------------------------------------------------------------

555−-- Monolithic postscans (postscan followed by a map)

556−------------------------------------------------------------------------------

557−

558−-- The performance of a modular postscan followed by a map seems to be

559−-- equivalent to this monolithic scan followed by map therefore we may not need

560−-- this implementation. We just have it for performance comparison and in case

561−-- modular version does not perform well in some situation.

562−--

563−{-# INLINE_NORMAL postscanlMx' #-}

564−postscanlMx' :: Monad m

565− => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b

566−postscanlMx' fstep begin done (Stream step state) = do

567− Stream step' (state, begin)

568− where

569− {-# INLINE_LATE step' #-}

570− step' gst (st, acc) = do

571− r <- step (adaptState gst) st

572− case r of

573− Yield x s -> do

574− old <- acc

575− y <- fstep old x

576− v <- done y

577− v `seq` y `seq` return (Yield v (s, return y))

578− Skip s -> return $ Skip (s, acc)

579− Stop -> return Stop

580−

581−{-# INLINE_NORMAL postscanlx' #-}

582−postscanlx' :: Monad m

583− => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b

584−postscanlx' fstep begin done s =

585− postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s

586−

587−-- XXX do we need consM strict to evaluate the begin value?

588−{-# INLINE scanlMx' #-}

589−scanlMx' :: Monad m

590− => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b

591−scanlMx' fstep begin done s =

592− (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s

593−

594−

595−data PostScanState s f = PostScan s !f

596−

597−{-# INLINE_NORMAL postscanOnce #-}

598−postscanOnce :: Monad m => FL.Fold m a b -> Stream m a -> Stream m b

599−postscanOnce (FL.Fold fstep begin done) (Stream step state) =

600− Stream step' (PostScan state begin)

601−

602− where

603−

604− {-# INLINE_LATE step' #-}

605− step' gst (PostScan st acc) = do

606− r <- step (adaptState gst) st

607− case r of

608− Yield x s -> do

609− old <- acc

610− res <- fstep old x

611− case res of

612− FL.Partial fs -> do

613− !v <- done fs

614− return $ Yield v $ PostScan s (return fs)

615− FL.Done _ -> return Stop

616− Skip s -> return $ Skip $ PostScan s acc

617− Stop -> return Stop

618−

619−{-# INLINE scanOnce #-}

620−scanOnce :: Monad m

621− => FL.Fold m a b -> Stream m a -> Stream m b

622−scanOnce fld@(FL.Fold _ begin done) s =

623− (begin >>= \x -> x `seq` done x) `consM` postscanOnce fld s

624−

625−{-# INLINE scanlx' #-}

626−scanlx' :: Monad m

627− => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b

628−scanlx' fstep begin done s =

629− scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s

630−

631−------------------------------------------------------------------------------

632−-- postscans

633−------------------------------------------------------------------------------

634−

635−{-# INLINE_NORMAL postscanlM' #-}

636−postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

637−postscanlM' fstep begin (Stream step state) =

638− Stream step' Nothing

639− where

640− {-# INLINE_LATE step' #-}

641− step' _ Nothing = do

642− !x <- begin

643− return $ Skip (Just (state, x))

644−

645− step' gst (Just (st, acc)) = do

646− r <- step (adaptState gst) st

647− case r of

648− Yield x s -> do

649− !y <- fstep acc x

650− return $ Yield y (Just (s, y))

651− Skip s -> return $ Skip (Just (s, acc))

652− Stop -> return Stop

653−

654−{-# INLINE_NORMAL postscanl' #-}

655−postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

656−postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)

657−

658−-- We can possibly have the "done" function as a Maybe to provide an option to

659−-- emit or not emit the accumulator when the stream stops.

660−--

661−-- TBD: use a single Yield point

662−--

663−{-# INLINE_NORMAL postscanlMAfter' #-}

664−postscanlMAfter' :: Monad m

665− => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

666−postscanlMAfter' fstep initial done (Stream step1 state1) = do

667− Stream step (Just (state1, initial))

668−

669− where

670−

671− {-# INLINE_LATE step #-}

672− step gst (Just (st, acc)) = do

673− r <- step1 (adaptState gst) st

674− case r of

675− Yield x s -> do

676− old <- acc

677− y <- fstep old x

678− y `seq` return (Yield y (Just (s, return y)))

679− Skip s -> return $ Skip $ Just (s, acc)

680− Stop -> acc >>= done >>= \res -> return (Yield res Nothing)

681− step _ Nothing = return Stop

682−

683−{-# INLINE_NORMAL postscanlM #-}

684−postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

685−postscanlM fstep begin (Stream step state) = Stream step' Nothing

686− where

687− {-# INLINE_LATE step' #-}

688− step' _ Nothing = do

689− r <- begin

690− return $ Skip (Just (state, r))

691−

692− step' gst (Just (st, acc)) = do

693− r <- step (adaptState gst) st

694− case r of

695− Yield x s -> do

696− y <- fstep acc x

697− return (Yield y (Just (s, y)))

698− Skip s -> return $ Skip (Just (s, acc))

699− Stop -> return Stop

700−

701−{-# INLINE_NORMAL postscanl #-}

702−postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

703−postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)

704−

705−{-# INLINE_NORMAL scanlM' #-}

706−scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

707−scanlM' fstep begin (Stream step state) = Stream step' Nothing

708− where

709− {-# INLINE_LATE step' #-}

710− step' _ Nothing = do

711− !x <- begin

712− return $ Yield x (Just (state, x))

713− step' gst (Just (st, acc)) = do

714− r <- step (adaptState gst) st

715− case r of

716− Yield x s -> do

717− !y <- fstep acc x

718− return $ Yield y (Just (s, y))

719− Skip s -> return $ Skip (Just (s, acc))

720− Stop -> return Stop

721−

722−{-# INLINE scanlMAfter' #-}

723−scanlMAfter' :: Monad m

724− => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

725−scanlMAfter' fstep initial done s =

726− (initial >>= \x -> x `seq` return x) `consM`

727− postscanlMAfter' fstep initial done s

728−

729−{-# INLINE scanl' #-}

730−scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

731−scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)

732−

733−{-# INLINE_NORMAL scanlM #-}

734−scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

735−scanlM fstep begin (Stream step state) = Stream step' Nothing

736− where

737− {-# INLINE_LATE step' #-}

738− step' _ Nothing = do

739− x <- begin

740− return $ Yield x (Just (state, x))

741− step' gst (Just (st, acc)) = do

742− r <- step (adaptState gst) st

743− case r of

744− Yield x s -> do

745− y <- fstep acc x

746− return $ Yield y (Just (s, y))

747− Skip s -> return $ Skip (Just (s, acc))

748− Stop -> return $ Stop

749−

750−{-# INLINE scanl #-}

751−scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

752−scanl f seed = scanlM (\a b -> return (f a b)) (return seed)

753−

754−{-# INLINE_NORMAL scanl1M #-}

755−scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a

756−scanl1M fstep (Stream step state) = Stream step' (state, Nothing)

757− where

758− {-# INLINE_LATE step' #-}

759− step' gst (st, Nothing) = do

760− r <- step gst st

761− case r of

762− Yield x s -> return $ Yield x (s, Just x)

763− Skip s -> return $ Skip (s, Nothing)

764− Stop -> return Stop

765−

766− step' gst (st, Just acc) = do

767− r <- step gst st

768− case r of

769− Yield y s -> do

770− z <- fstep acc y

771− return $ Yield z (s, Just z)

772− Skip s -> return $ Skip (s, Just acc)

773− Stop -> return Stop

774−

775−{-# INLINE scanl1 #-}

776−scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a

777−scanl1 f = scanl1M (\x y -> return (f x y))

778−

779−{-# INLINE_NORMAL scanl1M' #-}

780−scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a

781−scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)

782− where

783− {-# INLINE_LATE step' #-}

784− step' gst (st, Nothing) = do

785− r <- step gst st

786− case r of

787− Yield x s -> x `seq` return $ Yield x (s, Just x)

788− Skip s -> return $ Skip (s, Nothing)

789− Stop -> return Stop

790−

791− step' gst (st, Just acc) = acc `seq` do

792− r <- step gst st

793− case r of

794− Yield y s -> do

795− z <- fstep acc y

796− z `seq` return $ Yield z (s, Just z)

797− Skip s -> return $ Skip (s, Just acc)

798− Stop -> return Stop

799−

800−{-# INLINE scanl1' #-}

801−scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a

802−scanl1' f = scanl1M' (\x y -> return (f x y))

803−

804−------------------------------------------------------------------------------

805−-- Stateful map/scan

806−------------------------------------------------------------------------------

807−

808−data RollingMapState s a = RollingMapInit s | RollingMapGo s a

809−

810−{-# INLINE rollingMapM #-}

811−rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b

812−rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)

813− where

814− step gst (RollingMapInit st) = do

815− r <- step1 (adaptState gst) st

816− return $ case r of

817− Yield x s -> Skip $ RollingMapGo s x

818− Skip s -> Skip $ RollingMapInit s

819− Stop -> Stop

820−

821− step gst (RollingMapGo s1 x1) = do

822− r <- step1 (adaptState gst) s1

823− case r of

824− Yield x s -> do

825− !res <- f x x1

826− return $ Yield res $ RollingMapGo s x

827− Skip s -> return $ Skip $ RollingMapGo s x1

828− Stop -> return $ Stop

829−

830−{-# INLINE rollingMap #-}

831−rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b

832−rollingMap f = rollingMapM (\x y -> return $ f x y)

833−

834−-------------------------------------------------------------------------------

835−-- Filtering

836−-------------------------------------------------------------------------------

837−

838−{-# INLINE_NORMAL drop #-}

839−drop :: Monad m => Int -> Stream m a -> Stream m a

840−drop n (Stream step state) = Stream step' (state, Just n)

841− where

842− {-# INLINE_LATE step' #-}

843− step' gst (st, Just i)

844− | i > 0 = do

845− r <- step gst st

846− return $

847− case r of

848− Yield _ s -> Skip (s, Just (i - 1))

849− Skip s -> Skip (s, Just i)

850− Stop -> Stop

851− | otherwise = return $ Skip (st, Nothing)

852−

853− step' gst (st, Nothing) = do

854− r <- step gst st

855− return $

856− case r of

857− Yield x s -> Yield x (s, Nothing)

858− Skip s -> Skip (s, Nothing)

859− Stop -> Stop

860−

861−data DropWhileState s a

862− = DropWhileDrop s

863− | DropWhileYield a s

864− | DropWhileNext s

865−

866−{-# INLINE_NORMAL dropWhileM #-}

867−dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a

868−dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)

869− where

870− {-# INLINE_LATE step' #-}

871− step' gst (DropWhileDrop st) = do

872− r <- step gst st

873− case r of

874− Yield x s -> do

875− b <- f x

876− if b

877− then return $ Skip (DropWhileDrop s)

878− else return $ Skip (DropWhileYield x s)

879− Skip s -> return $ Skip (DropWhileDrop s)

880− Stop -> return Stop

881−

882− step' gst (DropWhileNext st) = do

883− r <- step gst st

884− case r of

885− Yield x s -> return $ Skip (DropWhileYield x s)

886− Skip s -> return $ Skip (DropWhileNext s)

887− Stop -> return Stop

888−

889− step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)

890−

891−{-# INLINE dropWhile #-}

892−dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

893−dropWhile f = dropWhileM (return . f)

894−

895−{-# INLINE_NORMAL filterM #-}

896−filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a

897−filterM f (Stream step state) = Stream step' state

898− where

899− {-# INLINE_LATE step' #-}

900− step' gst st = do

901− r <- step gst st

902− case r of

903− Yield x s -> do

904− b <- f x

905− return $ if b

906− then Yield x s

907− else Skip s

908− Skip s -> return $ Skip s

909− Stop -> return Stop

910−

911−{-# INLINE filter #-}

912−filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

913−filter f = filterM (return . f)

914−

915−{-# INLINE_NORMAL uniq #-}

916−uniq :: (Eq a, Monad m) => Stream m a -> Stream m a

917−uniq (Stream step state) = Stream step' (Nothing, state)

918− where

919− {-# INLINE_LATE step' #-}

920− step' gst (Nothing, st) = do

921− r <- step gst st

922− case r of

923− Yield x s -> return $ Yield x (Just x, s)

924− Skip s -> return $ Skip (Nothing, s)

925− Stop -> return Stop

926− step' gst (Just x, st) = do

927− r <- step gst st

928− case r of

929− Yield y s | x == y -> return $ Skip (Just x, s)

930− | otherwise -> return $ Yield y (Just y, s)

931− Skip s -> return $ Skip (Just x, s)

932− Stop -> return Stop

933−

934−------------------------------------------------------------------------------

935−-- Transformation by Mapping

936−------------------------------------------------------------------------------

937−

938−{-# INLINE_NORMAL sequence #-}

939−sequence :: Monad m => Stream m (m a) -> Stream m a

940−sequence (Stream step state) = Stream step' state

941− where

942− {-# INLINE_LATE step' #-}

943− step' gst st = do

944− r <- step (adaptState gst) st

945− case r of

946− Yield x s -> x >>= \a -> return (Yield a s)

947− Skip s -> return $ Skip s

948− Stop -> return Stop

949−

950−------------------------------------------------------------------------------

951−-- Inserting

952−------------------------------------------------------------------------------

953−

954−data LoopState x s = FirstYield s

955− | InterspersingYield s

956− | YieldAndCarry x s

957−

958−{-# INLINE_NORMAL intersperseM #-}

959−intersperseM :: Monad m => m a -> Stream m a -> Stream m a

960−intersperseM m (Stream step state) = Stream step' (FirstYield state)

961− where

962− {-# INLINE_LATE step' #-}

963− step' gst (FirstYield st) = do

964− r <- step gst st

965− return $

966− case r of

967− Yield x s -> Skip (YieldAndCarry x s)

968− Skip s -> Skip (FirstYield s)

969− Stop -> Stop

970−

971− step' gst (InterspersingYield st) = do

972− r <- step gst st

973− case r of

974− Yield x s -> do

975− a <- m

976− return $ Yield a (YieldAndCarry x s)

977− Skip s -> return $ Skip $ InterspersingYield s

978− Stop -> return Stop

979−

980− step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)

981−

982−{-# INLINE_NORMAL intersperseM_ #-}

983−intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a

984−intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))

985− where

986− {-# INLINE_LATE step #-}

987− step gst (Left (eff, st)) = do

988− r <- step1 gst st

989− case r of

990− Yield x s -> eff >> return (Yield x (Right s))

991− Skip s -> return $ Skip (Left (eff, s))

992− Stop -> return Stop

993−

994− step _ (Right st) = return $ Skip $ Left (void m, st)

995−

996−data SuffixState s a

997− = SuffixElem s

998− | SuffixSuffix s

999− | SuffixYield a (SuffixState s a)

1000−

1001−{-# INLINE_NORMAL intersperseSuffix #-}

1002−intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a

1003−intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)

1004− where

1005− {-# INLINE_LATE step' #-}

1006− step' gst (SuffixElem st) = do

1007− r <- step gst st

1008− return $ case r of

1009− Yield x s -> Skip (SuffixYield x (SuffixSuffix s))

1010− Skip s -> Skip (SuffixElem s)

1011− Stop -> Stop

1012−

1013− step' _ (SuffixSuffix st) = do

1014− action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))

1015−

1016− step' _ (SuffixYield x next) = return $ Yield x next

1017−

1018−{-# INLINE_NORMAL intersperseSuffix_ #-}

1019−intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a

1020−intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)

1021− where

1022− {-# INLINE_LATE step #-}

1023− step gst (Left st) = do

1024− r <- step1 gst st

1025− case r of

1026− Yield x s -> return $ Yield x (Right s)

1027− Skip s -> return $ Skip $ Left s

1028− Stop -> return Stop

1029−

1030− step _ (Right st) = m >> return (Skip (Left st))

1031−

1032−data SuffixSpanState s a

1033− = SuffixSpanElem s Int

1034− | SuffixSpanSuffix s

1035− | SuffixSpanYield a (SuffixSpanState s a)

1036− | SuffixSpanLast

1037− | SuffixSpanStop

1038−

1039−-- | intersperse after every n items

1040−{-# INLINE_NORMAL intersperseSuffixBySpan #-}

1041−intersperseSuffixBySpan :: forall m a. Monad m

1042− => Int -> m a -> Stream m a -> Stream m a

1043−intersperseSuffixBySpan n action (Stream step state) =

1044− Stream step' (SuffixSpanElem state n)

1045− where

1046− {-# INLINE_LATE step' #-}

1047− step' gst (SuffixSpanElem st i) | i > 0 = do

1048− r <- step gst st

1049− return $ case r of

1050− Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))

1051− Skip s -> Skip (SuffixSpanElem s i)

1052− Stop -> if i == n then Stop else Skip SuffixSpanLast

1053− step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)

1054−

1055− step' _ (SuffixSpanSuffix st) = do

1056− action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))

1057−

1058− step' _ (SuffixSpanLast) = do

1059− action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)

1060−

1061− step' _ (SuffixSpanYield x next) = return $ Yield x next

1062−

1063− step' _ (SuffixSpanStop) = return Stop

1064−

1065−{-# INLINE intersperse #-}

1066−intersperse :: Monad m => a -> Stream m a -> Stream m a

1067−intersperse a = intersperseM (return a)

1068−

1069−{-# INLINE_NORMAL insertBy #-}

1070−insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a

1071−insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)

1072− where

1073− {-# INLINE_LATE step' #-}

1074− step' gst (st, False, _) = do

1075− r <- step gst st

1076− case r of

1077− Yield x s -> case cmp a x of

1078− GT -> return $ Yield x (s, False, Nothing)

1079− _ -> return $ Yield a (s, True, Just x)

1080− Skip s -> return $ Skip (s, False, Nothing)

1081− Stop -> return $ Yield a (st, True, Nothing)

1082−

1083− step' _ (_, True, Nothing) = return Stop

1084−

1085− step' gst (st, True, Just prev) = do

1086− r <- step gst st

1087− case r of

1088− Yield x s -> return $ Yield prev (s, True, Just x)

1089− Skip s -> return $ Skip (s, True, Just prev)

1090− Stop -> return $ Yield prev (st, True, Nothing)

1091−

1092−------------------------------------------------------------------------------

1093−-- Deleting

1094−------------------------------------------------------------------------------

1095−

1096−{-# INLINE_NORMAL deleteBy #-}

1097−deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a

1098−deleteBy eq x (Stream step state) = Stream step' (state, False)

1099− where

1100− {-# INLINE_LATE step' #-}

1101− step' gst (st, False) = do

1102− r <- step gst st

1103− case r of

1104− Yield y s -> return $

1105− if eq x y then Skip (s, True) else Yield y (s, False)

1106− Skip s -> return $ Skip (s, False)

1107− Stop -> return Stop

1108−

1109− step' gst (st, True) = do

1110− r <- step gst st

1111− case r of

1112− Yield y s -> return $ Yield y (s, True)

1113− Skip s -> return $ Skip (s, True)

1114− Stop -> return Stop

1115−

1116−------------------------------------------------------------------------------

1117−-- Transformation by Map and Filter

1118−------------------------------------------------------------------------------

1119−

1120−-- XXX Will this always fuse properly?

1121−{-# INLINE_NORMAL mapMaybe #-}

1122−mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b

1123−mapMaybe f = fmap fromJust . filter isJust . map f

1124−

1125−{-# INLINE_NORMAL mapMaybeM #-}

1126−mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b

1127−mapMaybeM f = fmap fromJust . filter isJust . mapM f

1128−

1129−------------------------------------------------------------------------------

1130−-- Zipping

1131−------------------------------------------------------------------------------

1132−

1133−{-# INLINE_NORMAL indexed #-}

1134−indexed :: Monad m => Stream m a -> Stream m (Int, a)

1135−indexed (Stream step state) = Stream step' (state, 0)

1136− where

1137− {-# INLINE_LATE step' #-}

1138− step' gst (st, i) = i `seq` do

1139− r <- step (adaptState gst) st

1140− case r of

1141− Yield x s -> return $ Yield (i, x) (s, i+1)

1142− Skip s -> return $ Skip (s, i)

1143− Stop -> return Stop

1144−

1145−{-# INLINE_NORMAL indexedR #-}

1146−indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)

1147−indexedR m (Stream step state) = Stream step' (state, m)

1148− where

1149− {-# INLINE_LATE step' #-}

1150− step' gst (st, i) = i `seq` do

1151− r <- step (adaptState gst) st

1152− case r of

1153− Yield x s -> let i' = i - 1

1154− in return $ Yield (i, x) (s, i')

1155− Skip s -> return $ Skip (s, i)

1156− Stop -> return Stop

1157−

1158−{-# INLINE_NORMAL zipWithM #-}

1159−zipWithM :: Monad m

1160− => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c

1161−zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)

1162− where

1163− {-# INLINE_LATE step #-}

1164− step gst (sa, sb, Nothing) = do

1165− r <- stepa (adaptState gst) sa

1166− return $

1167− case r of

1168− Yield x sa' -> Skip (sa', sb, Just x)

1169− Skip sa' -> Skip (sa', sb, Nothing)

1170− Stop -> Stop

1171−

1172− step gst (sa, sb, Just x) = do

1173− r <- stepb (adaptState gst) sb

1174− case r of

1175− Yield y sb' -> do

1176− z <- f x y

1177− return $ Yield z (sa, sb', Nothing)

1178− Skip sb' -> return $ Skip (sa, sb', Just x)

1179− Stop -> return Stop

1180−

1181−#if __GLASGOW_HASKELL__ >= 801

1182−{-# RULES "zipWithM xs xs"

1183− forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}

1184−#endif

1185−

1186−{-# INLINE zipWith #-}

1187−zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c

1188−zipWith f = zipWithM (\a b -> return (f a b))

1189−

1190−------------------------------------------------------------------------------

1191−-- Merging

1192−------------------------------------------------------------------------------

1193−

1194−{-# INLINE_NORMAL mergeByM #-}

1195−mergeByM

1196− :: (Monad m)

1197− => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a

1198−mergeByM cmp (Stream stepa ta) (Stream stepb tb) =

1199− Stream step (Just ta, Just tb, Nothing, Nothing)

1200− where

1201− {-# INLINE_LATE step #-}

1202−

1203− -- one of the values is missing, and the corresponding stream is running

1204− step gst (Just sa, sb, Nothing, b) = do

1205− r <- stepa gst sa

1206− return $ case r of

1207− Yield a sa' -> Skip (Just sa', sb, Just a, b)

1208− Skip sa' -> Skip (Just sa', sb, Nothing, b)

1209− Stop -> Skip (Nothing, sb, Nothing, b)

1210−

1211− step gst (sa, Just sb, a, Nothing) = do

1212− r <- stepb gst sb

1213− return $ case r of

1214− Yield b sb' -> Skip (sa, Just sb', a, Just b)

1215− Skip sb' -> Skip (sa, Just sb', a, Nothing)

1216− Stop -> Skip (sa, Nothing, a, Nothing)

1217−

1218− -- both the values are available

1219− step _ (sa, sb, Just a, Just b) = do

1220− res <- cmp a b

1221− return $ case res of

1222− GT -> Yield b (sa, sb, Just a, Nothing)

1223− _ -> Yield a (sa, sb, Nothing, Just b)

1224−

1225− -- one of the values is missing, corresponding stream is done

1226− step _ (Nothing, sb, Nothing, Just b) =

1227− return $ Yield b (Nothing, sb, Nothing, Nothing)

1228−

1229− step _ (sa, Nothing, Just a, Nothing) =

1230− return $ Yield a (sa, Nothing, Nothing, Nothing)

1231−

1232− step _ (Nothing, Nothing, Nothing, Nothing) = return Stop

1233−

1234−{-# INLINE mergeBy #-}

1235−mergeBy

1236− :: (Monad m)

1237− => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a

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

1239−

1240−------------------------------------------------------------------------------

1241−-- Transformation comprehensions

1242−------------------------------------------------------------------------------

1243−

1244−{-# INLINE_NORMAL the #-}

1245−the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)

1246−the (Stream step state) = go state

1247− where

1248− go st = do

1249− r <- step defState st

1250− case r of

1251− Yield x s -> go' x s

1252− Skip s -> go s

1253− Stop -> return Nothing

1254− go' n st = do

1255− r <- step defState st

1256− case r of

1257− Yield x s | x == n -> go' n s

1258− | otherwise -> return Nothing

1259− Skip s -> go' n s

1260− Stop -> return (Just n)

1261−

1262−-------------------------------------------------------------------------------

1263−-- Concurrent application and fold

1264−-------------------------------------------------------------------------------

1265−

1266−{-# INLINE_NORMAL mkParallelD #-}

1267−mkParallelD :: MonadAsync m => Stream m a -> Stream m a

1268−mkParallelD m = Stream step Nothing

1269− where

1270−

1271− step gst Nothing = do

1272− sv <- newParallelVar StopNone gst

1273− toSVarParallel gst sv m

1274− -- XXX use unfold instead?

1275− return $ Skip $ Just $ fromSVar sv

1276−

1277− step gst (Just (UnStream step1 st)) = do

1278− r <- step1 gst st

1279− return $ case r of

1280− Yield a s -> Yield a (Just $ Stream step1 s)

1281− Skip s -> Skip (Just $ Stream step1 s)

1282− Stop -> Stop

1283−

1284−-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a

1285−-- parallel style SVar for evaluation. Currently, parallel style cannot use

1286−-- rate control whereas Async style can use rate control. In async style SVar

1287−-- the worker thread terminates when the buffer is full whereas in Parallel

1288−-- style it blocks.

1289−--

1290−-- | Make the stream producer and consumer run concurrently by introducing a

1291−-- buffer between them. The producer thread evaluates the input stream until

1292−-- the buffer fills, it blocks if the buffer is full until there is space in

1293−-- the buffer. The consumer consumes the stream lazily from the buffer.

1294−--

1295−-- /Internal/

1296−--

1297−{-# INLINE_NORMAL mkParallel #-}

1298−mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a

1299−mkParallel = fromStreamD . mkParallelD . toStreamD

1300−

1301−-- Note: we can use another API with two callbacks stop and yield if we want

1302−-- the callback to be able to indicate end of stream.

1303−--

1304−-- | Generates a callback and a stream pair. The callback returned is used to

1305−-- queue values to the stream. The stream is infinite, there is no way for the

1306−-- callback to indicate that it is done now.

1307−--

1308−-- /Internal/

1309−--

1310−{-# INLINE_NORMAL newCallbackStream #-}

1311−newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)

1312−newCallbackStream = do

1313− sv <- newParallelVar StopNone defState

1314−

1315− -- XXX Add our own thread-id to the SVar as we can not know the callback's

1316− -- thread-id and the callback is not run in a managed worker. We need to

1317− -- handle this better.

1318− liftIO myThreadId >>= modifyThread sv

1319−

1320− let callback a = liftIO $ void $ send sv (ChildYield a)

1321− -- XXX we can return an SVar and then the consumer can unfold from the

1322− -- SVar?

1323− return (callback, fromStreamD (fromSVar sv))

1324−

1325−------------------------------------------------------------------------------

1326−-- Time related

1327−------------------------------------------------------------------------------

1328−

1329−-- XXX using getTime in the loop can be pretty expensive especially for

1330−-- computations where iterations are lightweight. We have the following

1331−-- options:

1332−--

1333−-- 1) Run a timeout thread updating a flag asynchronously and check that

1334−-- flag here, that way we can have a cheap termination check.

1335−--

1336−-- 2) Use COARSE clock to get time with lower resolution but more efficiently.

1337−--

1338−-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the

1339−-- termination value of rdtsc in the beginning and then in each iteration just

1340−-- get rdtsc and check if we should terminate.

1341−--

1342−data TakeByTime st s

1343− = TakeByTimeInit st

1344− | TakeByTimeCheck st s

1345− | TakeByTimeYield st s

1346−

1347−{-# INLINE_NORMAL takeByTime #-}

1348−takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a

1349−takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)

1350− where

1351−

1352− lim = toRelTime64 duration

1353−

1354− {-# INLINE_LATE step #-}

1355− step _ (TakeByTimeInit _) | lim == 0 = return Stop

1356− step _ (TakeByTimeInit st) = do

1357− t0 <- liftIO $ getTime Monotonic

1358− return $ Skip (TakeByTimeYield st t0)

1359− step _ (TakeByTimeCheck st t0) = do

1360− t <- liftIO $ getTime Monotonic

1361− return $

1362− if diffAbsTime64 t t0 > lim

1363− then Stop

1364− else Skip (TakeByTimeYield st t0)

1365− step gst (TakeByTimeYield st t0) = do

1366− r <- step1 gst st

1367− return $ case r of

1368− Yield x s -> Yield x (TakeByTimeCheck s t0)

1369− Skip s -> Skip (TakeByTimeCheck s t0)

1370− Stop -> Stop

1371−

1372−data DropByTime st s x

1373− = DropByTimeInit st

1374− | DropByTimeGen st s

1375− | DropByTimeCheck st s x

1376− | DropByTimeYield st

1377−

1378−{-# INLINE_NORMAL dropByTime #-}

1379−dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a

1380−dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)

1381− where

1382−

1383− lim = toRelTime64 duration

1384−

1385− {-# INLINE_LATE step #-}

1386− step _ (DropByTimeInit st) = do

1387− t0 <- liftIO $ getTime Monotonic

1388− return $ Skip (DropByTimeGen st t0)

1389− step gst (DropByTimeGen st t0) = do

1390− r <- step1 gst st

1391− return $ case r of

1392− Yield x s -> Skip (DropByTimeCheck s t0 x)

1393− Skip s -> Skip (DropByTimeGen s t0)

1394− Stop -> Stop

1395− step _ (DropByTimeCheck st t0 x) = do

1396− t <- liftIO $ getTime Monotonic

1397− if diffAbsTime64 t t0 <= lim

1398− then return $ Skip $ DropByTimeGen st t0

1399− else return $ Yield x $ DropByTimeYield st

1400− step gst (DropByTimeYield st) = do

1401− r <- step1 gst st

1402− return $ case r of

1403− Yield x s -> Yield x (DropByTimeYield s)

1404− Skip s -> Skip (DropByTimeYield s)

1405− Stop -> Stop

src/Streamly/Internal/Data/Stream/StreamD/Common.hs+210−11

@@ -12,31 +12,39 @@

13module Streamly.Internal.Data.Stream.StreamD.Common

14 (

15− FromSVarState (..)

16− , fromProducer

15+ fromProducer

17 , fromPrimIORef

16 , fromPrimIORef

17+ , fromSVar

18+ , foldOnce

19+ , toSVarParallel

20+ , toStreamD

18 , uncons

21 , uncons

19 , nilM

22 , nilM

23+ , consM

20 , takeWhileM

24 , takeWhileM

21 , takeWhile

25 , takeWhile

22 )

26 )

23where

27where

29+

30+import Control.Exception (fromException)

25import Control.Monad (when)

31import Control.Monad (when)

32+import Control.Monad.Catch (throwM)

26import Control.Monad.IO.Class (MonadIO(..))

33import Control.Monad.IO.Class (MonadIO(..))

27−import Data.IORef (writeIORef)

34+import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

35+import Data.Maybe (isNothing)

36+import GHC.Types (SPEC(..))

37+import System.Mem (performMajorGC)

38+import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

28import Streamly.Internal.Data.IORef.Prim (Prim)

39import Streamly.Internal.Data.IORef.Prim (Prim)

40+import Streamly.Internal.Data.Fold.Types (Fold(..))

29import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

41import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

43+import qualified Streamly.Internal.Data.Stream.StreamK as K

44+import qualified Streamly.Internal.Data.Fold as FL

31import qualified Streamly.Internal.Data.IORef.Prim as Prim

45import qualified Streamly.Internal.Data.IORef.Prim as Prim

32−import qualified Prelude

46+

33−

34−import Prelude hiding

47+import Prelude hiding (takeWhile)

35− ( map, mapM, mapM_, repeat, foldr, last, take, filter

36− , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem

37− , notElem, null, head, tail, zipWith, lookup, foldr1, sequence

38− , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat

39− , reverse, iterate, splitAt)

40import Streamly.Internal.Data.Stream.StreamD.Type

48import Streamly.Internal.Data.Stream.StreamD.Type

41import Streamly.Internal.Data.SVar

49import Streamly.Internal.Data.SVar

@@ -65,12 +73,108 @@ uncons (UnStream step state) = go state

65nilM :: Monad m => m b -> Stream m a

73nilM :: Monad m => m b -> Stream m a

66nilM m = Stream (\_ _ -> m >> return Stop) ()

74nilM m = Stream (\_ _ -> m >> return Stop) ()

76+{-# INLINE_NORMAL consM #-}

77+consM :: Monad m => m a -> Stream m a -> Stream m a

78+consM m (Stream step state) = Stream step1 Nothing

79+ where

80+ {-# INLINE_LATE step1 #-}

81+ step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)

82+ step1 gst (Just st) = do

83+ r <- step gst st

84+ return $

85+ case r of

86+ Yield a s -> Yield a (Just s)

87+ Skip s -> Skip (Just s)

88+ Stop -> Stop

89+

90+{-# INLINE toStreamD #-}

91+toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a

92+toStreamD = fromStreamK . K.toStream

68data FromSVarState t m a =

93data FromSVarState t m a =

69 FromSVarInit

94 FromSVarInit

70 | FromSVarRead (SVar t m a)

95 | FromSVarRead (SVar t m a)

71 | FromSVarLoop (SVar t m a) [ChildEvent a]

96 | FromSVarLoop (SVar t m a) [ChildEvent a]

72 | FromSVarDone (SVar t m a)

97 | FromSVarDone (SVar t m a)

99+-------------------------------------------------------------------------------

100+-- Generation from SVar

101+-------------------------------------------------------------------------------

102+

103+{-# INLINE_NORMAL fromSVar #-}

104+fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a

105+fromSVar svar = Stream step FromSVarInit

106+ where

107+

108+ {-# INLINE_LATE step #-}

109+ step _ FromSVarInit = do

110+ ref <- liftIO $ newIORef ()

111+ _ <- liftIO $ mkWeakIORef ref hook

112+ -- when this copy of svar gets garbage collected "ref" will get

113+ -- garbage collected and our GC hook will be called.

114+ let sv = svar{svarRef = Just ref}

115+ return $ Skip (FromSVarRead sv)

116+

117+ where

118+

119+ {-# NOINLINE hook #-}

120+ hook = do

121+ when (svarInspectMode svar) $ do

122+ r <- liftIO $ readIORef (svarStopTime (svarStats svar))

123+ when (isNothing r) $

124+ printSVar svar "SVar Garbage Collected"

125+ cleanupSVar svar

126+ -- If there are any SVars referenced by this SVar a GC will prompt

127+ -- them to be cleaned up quickly.

128+ when (svarInspectMode svar) performMajorGC

129+

130+ step _ (FromSVarRead sv) = do

131+ list <- readOutputQ sv

132+ -- Reversing the output is important to guarantee that we process the

133+ -- outputs in the same order as they were generated by the constituent

134+ -- streams.

135+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)

136+

137+ step _ (FromSVarLoop sv []) = do

138+ done <- postProcess sv

139+ return $ Skip $ if done

140+ then (FromSVarDone sv)

141+ else (FromSVarRead sv)

142+

143+ step _ (FromSVarLoop sv (ev : es)) = do

144+ case ev of

145+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)

146+ ChildStop tid e -> do

147+ accountThread sv tid

148+ case e of

149+ Nothing -> do

150+ stop <- shouldStop tid

151+ if stop

152+ then do

153+ liftIO (cleanupSVar sv)

154+ return $ Skip (FromSVarDone sv)

155+ else return $ Skip (FromSVarLoop sv es)

156+ Just ex ->

157+ case fromException ex of

158+ Just ThreadAbort ->

159+ return $ Skip (FromSVarLoop sv es)

160+ Nothing -> liftIO (cleanupSVar sv) >> throwM ex

161+ where

162+

163+ shouldStop tid =

164+ case svarStopStyle sv of

165+ StopNone -> return False

166+ StopAny -> return True

167+ StopBy -> do

168+ sid <- liftIO $ readIORef (svarStopBy sv)

169+ return $ if tid == sid then True else False

170+

171+ step _ (FromSVarDone sv) = do

172+ when (svarInspectMode sv) $ do

173+ t <- liftIO $ getTime Monotonic

174+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)

175+ liftIO $ printSVar sv "SVar Done"

176+ return Stop

177+

74{-# INLINE_NORMAL fromPrimIORef #-}

178{-# INLINE_NORMAL fromPrimIORef #-}

75fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a

179fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a

76fromPrimIORef var = Stream step ()

180fromPrimIORef var = Stream step ()

@@ -116,6 +220,101 @@ fromProducer svar = Stream step (FromSVarRead svar)

116

220

117 step _ FromSVarInit = undefined

221 step _ FromSVarInit = undefined

118

222

223+------------------------------------------------------------------------------

224+-- Left Folds

225+------------------------------------------------------------------------------

226+

227+{-# INLINE_NORMAL foldOnce #-}

228+foldOnce :: (Monad m) => Fold m a b -> Stream m a -> m b

229+foldOnce (Fold fstep begin done) (Stream step state) =

230+ begin >>= \x -> go SPEC x state

231+

232+ where

233+

234+ {-# INLINE go #-}

235+ go !_ !fs st = do

236+ r <- step defState st

237+ case r of

238+ Yield x s -> do

239+ res <- fstep fs x

240+ case res of

241+ FL.Done b -> return b

242+ FL.Partial fs1 -> go SPEC fs1 s

243+ Skip s -> go SPEC fs s

244+ Stop -> done fs

245+

246+-------------------------------------------------------------------------------

247+-- Concurrent application and fold

248+-------------------------------------------------------------------------------

249+

250+-- XXX These functions should be moved to Stream/Parallel.hs

251+--

252+-- Using StreamD the worker stream producing code can fuse with the code to

253+-- queue output to the SVar giving some perf boost.

254+--

255+-- Note that StreamD can only be used in limited situations, specifically, we

256+-- cannot implement joinStreamVarPar using this.

257+--

258+-- XXX make sure that the SVar passed is a Parallel style SVar.

259+

260+-- | Fold the supplied stream to the SVar asynchronously using Parallel

261+-- concurrency style.

262+-- {-# INLINE_NORMAL toSVarParallel #-}

263+{-# INLINE toSVarParallel #-}

264+toSVarParallel :: MonadAsync m

265+ => State t m a -> SVar t m a -> Stream m a -> m ()

266+toSVarParallel st sv xs =

267+ if svarInspectMode sv

268+ then forkWithDiag

269+ else do

270+ tid <-

271+ case getYieldLimit st of

272+ Nothing -> doFork (work Nothing)

273+ (svarMrun sv)

274+ (handleChildException sv)

275+ Just _ -> doFork (workLim Nothing)

276+ (svarMrun sv)

277+ (handleChildException sv)

278+ modifyThread sv tid

279+

280+ where

281+

282+ {-# NOINLINE work #-}

283+ work info = (foldOnce (FL.toParallelSVar sv info) xs)

284+

285+ {-# NOINLINE workLim #-}

286+ workLim info = foldOnce (FL.toParallelSVarLimited sv info) xs

287+

288+ {-# NOINLINE forkWithDiag #-}

289+ forkWithDiag = do

290+ -- We do not use workerCount in case of ParallelVar but still there is

291+ -- no harm in maintaining it correctly.

292+ liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1

293+ recordMaxWorkers sv

294+ -- This allocation matters when significant number of workers are being

295+ -- sent. We allocate it only when needed. The overhead increases by 4x.

296+ winfo <-

297+ case yieldRateInfo sv of

298+ Nothing -> return Nothing

299+ Just _ -> liftIO $ do

300+ cntRef <- newIORef 0

301+ t <- getTime Monotonic

302+ lat <- newIORef (0, t)

303+ return $ Just WorkerInfo

304+ { workerYieldMax = 0

305+ , workerYieldCount = cntRef

306+ , workerLatencyStart = lat

307+ }

308+ tid <-

309+ case getYieldLimit st of

310+ Nothing -> doFork (work winfo)

311+ (svarMrun sv)

312+ (handleChildException sv)

313+ Just _ -> doFork (workLim winfo)

314+ (svarMrun sv)

315+ (handleChildException sv)

316+ modifyThread sv tid

317+

119-------------------------------------------------------------------------------

318-------------------------------------------------------------------------------

120-- Filtering

319-- Filtering

121-------------------------------------------------------------------------------

320-------------------------------------------------------------------------------

src/Streamly/Internal/Data/Stream/StreamD/Eliminate.hs+173−93

@@ -64,6 +64,24 @@ module Streamly.Internal.Data.Stream.StreamD.Eliminate

64 , find

65 , (!!)

66 , toSVarParallel

67+

68+ -- ** Substreams

69+ , isPrefixOf

70+ , isSubsequenceOf

71+ , stripPrefix

72+

73+ -- ** Map and Fold

74+ , mapM_

75+

76+ -- ** Conversions

77+ -- | Transform a stream into another type.

78+ , hoist

79+ , generally

80+

81+ , liftInner

82+ , runReaderT

83+ , evalStateT

84+ , runStateT

67 )

85 )

68where

86where

@@ -79,11 +97,12 @@ import Foreign.Storable (Storable(..))

79import GHC.Exts (SpecConstrAnnotation(..))

97import GHC.Exts (SpecConstrAnnotation(..))

80import GHC.Types (SPEC(..))

98import GHC.Types (SPEC(..))

81import Fusion.Plugin.Types (Fuse(..))

99import Fusion.Plugin.Types (Fuse(..))

82−import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

83import Streamly.Internal.Data.Fold.Types (Fold(..))

100import Streamly.Internal.Data.Fold.Types (Fold(..))

84import Streamly.Internal.Data.Parser (ParseError(..))

101import Streamly.Internal.Data.Parser (ParseError(..))

85−import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

86import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)

102import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)

103+import Control.Monad.Trans.Reader (ReaderT)

104+import Control.Monad.Trans.State.Strict (StateT)

105+import Data.Functor.Identity (Identity(..))

106

88import qualified Streamly.Internal.Data.IORef.Prim as Prim

107import qualified Streamly.Internal.Data.IORef.Prim as Prim

89import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A

108import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A

@@ -92,6 +111,8 @@ import qualified Streamly.Internal.Data.Stream.StreamK as K

92import qualified Streamly.Internal.Data.Parser as PR

111import qualified Streamly.Internal.Data.Parser as PR

93import qualified Streamly.Internal.Data.Parser.ParserD as PRD

112import qualified Streamly.Internal.Data.Parser.ParserD as PRD

94import qualified Control.Monad.Catch as MC

113import qualified Control.Monad.Catch as MC

114+import qualified Control.Monad.Trans.Reader as Reader

115+import qualified Control.Monad.Trans.State.Strict as State

95import qualified Prelude

116import qualified Prelude

117

97import Prelude hiding

118import Prelude hiding

@@ -125,25 +146,6 @@ foldr1 f m = do

125-- Left Folds

146-- Left Folds

126------------------------------------------------------------------------------

147------------------------------------------------------------------------------

127

148

128−{-# INLINE_NORMAL foldOnce #-}

129−foldOnce :: (Monad m) => Fold m a b -> Stream m a -> m b

130−foldOnce (Fold fstep begin done) (Stream step state) =

131− begin >>= \x -> go SPEC x state

132−

133− where

134−

135− {-# INLINE go #-}

136− go !_ !fs st = do

137− r <- step defState st

138− case r of

139− Yield x s -> do

140− res <- fstep fs x

141− case res of

142− FL.Done b -> return b

143− FL.Partial fs1 -> go SPEC fs1 s

144− Skip s -> go SPEC fs s

145− Stop -> done fs

146−

147{-# INLINE_NORMAL foldlT #-}

149{-# INLINE_NORMAL foldlT #-}

148foldlT :: (Monad m, Monad (s m), MonadTrans s)

150foldlT :: (Monad m, Monad (s m), MonadTrans s)

149 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b

151 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b

@@ -755,78 +757,6 @@ reverse' m =

755 $ toStreamK

757 $ toStreamK

756 $ A.fromStreamDArraysOf A.defaultChunkSize m

758 $ A.fromStreamDArraysOf A.defaultChunkSize m

757

759

758−-------------------------------------------------------------------------------

759−-- Concurrent application and fold

760−-------------------------------------------------------------------------------

761−

762−-- XXX These functions should be moved to Stream/Parallel.hs

763−--

764−-- Using StreamD the worker stream producing code can fuse with the code to

765−-- queue output to the SVar giving some perf boost.

766−--

767−-- Note that StreamD can only be used in limited situations, specifically, we

768−-- cannot implement joinStreamVarPar using this.

769−--

770−-- XXX make sure that the SVar passed is a Parallel style SVar.

771−

772−-- | Fold the supplied stream to the SVar asynchronously using Parallel

773−-- concurrency style.

774−-- {-# INLINE_NORMAL toSVarParallel #-}

775−{-# INLINE toSVarParallel #-}

776−toSVarParallel :: MonadAsync m

777− => State t m a -> SVar t m a -> Stream m a -> m ()

778−toSVarParallel st sv xs =

779− if svarInspectMode sv

780− then forkWithDiag

781− else do

782− tid <-

783− case getYieldLimit st of

784− Nothing -> doFork (work Nothing)

785− (svarMrun sv)

786− (handleChildException sv)

787− Just _ -> doFork (workLim Nothing)

788− (svarMrun sv)

789− (handleChildException sv)

790− modifyThread sv tid

791−

792− where

793−

794− {-# NOINLINE work #-}

795− work info = (foldOnce (FL.toParallelSVar sv info) xs)

796−

797− {-# NOINLINE workLim #-}

798− workLim info = foldOnce (FL.toParallelSVarLimited sv info) xs

799−

800− {-# NOINLINE forkWithDiag #-}

801− forkWithDiag = do

802− -- We do not use workerCount in case of ParallelVar but still there is

803− -- no harm in maintaining it correctly.

804− liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1

805− recordMaxWorkers sv

806− -- This allocation matters when significant number of workers are being

807− -- sent. We allocate it only when needed. The overhead increases by 4x.

808− winfo <-

809− case yieldRateInfo sv of

810− Nothing -> return Nothing

811− Just _ -> liftIO $ do

812− cntRef <- newIORef 0

813− t <- getTime Monotonic

814− lat <- newIORef (0, t)

815− return $ Just WorkerInfo

816− { workerYieldMax = 0

817− , workerYieldCount = cntRef

818− , workerLatencyStart = lat

819− }

820− tid <-

821− case getYieldLimit st of

822− Nothing -> doFork (work winfo)

823− (svarMrun sv)

824− (handleChildException sv)

825− Just _ -> doFork (workLim winfo)

826− (svarMrun sv)

827− (handleChildException sv)

828− modifyThread sv tid

829−

830------------------------------------------------------------------------------

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

831-- Tapping/Distributing

761-- Tapping/Distributing

832------------------------------------------------------------------------------

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

@@ -1057,3 +987,153 @@ tapAsync f (Stream step1 state1) = Stream step TapInit

1057 Yield a s -> Yield a (TapDone s)

987 Yield a s -> Yield a (TapDone s)

1058 Skip s -> Skip (TapDone s)

988 Skip s -> Skip (TapDone s)

1059 Stop -> Stop

989 Stop -> Stop

990+

991+-------------------------------------------------------------------------------

992+-- Hoisting the inner monad

993+-------------------------------------------------------------------------------

994+

995+{-# INLINE_NORMAL hoist #-}

996+hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a

997+hoist f (Stream step state) = (Stream step' state)

998+ where

999+ {-# INLINE_LATE step' #-}

1000+ step' gst st = do

1001+ r <- f $ step (adaptState gst) st

1002+ return $ case r of

1003+ Yield x s -> Yield x s

1004+ Skip s -> Skip s

1005+ Stop -> Stop

1006+

1007+{-# INLINE generally #-}

1008+generally :: Monad m => Stream Identity a -> Stream m a

1009+generally = hoist (return . runIdentity)

1010+

1011+{-# INLINE_NORMAL liftInner #-}

1012+liftInner :: (Monad m, MonadTrans t, Monad (t m))

1013+ => Stream m a -> Stream (t m) a

1014+liftInner (Stream step state) = Stream step' state

1015+ where

1016+ {-# INLINE_LATE step' #-}

1017+ step' gst st = do

1018+ r <- lift $ step (adaptState gst) st

1019+ return $ case r of

1020+ Yield x s -> Yield x s

1021+ Skip s -> Skip s

1022+ Stop -> Stop

1023+

1024+{-# INLINE_NORMAL runReaderT #-}

1025+runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a

1026+runReaderT env (Stream step state) = Stream step' (state, env)

1027+ where

1028+ {-# INLINE_LATE step' #-}

1029+ step' gst (st, action) = do

1030+ sv <- action

1031+ r <- Reader.runReaderT (step (adaptState gst) st) sv

1032+ return $ case r of

1033+ Yield x s -> Yield x (s, return sv)

1034+ Skip s -> Skip (s, return sv)

1035+ Stop -> Stop

1036+

1037+{-# INLINE_NORMAL evalStateT #-}

1038+evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a

1039+evalStateT initial (Stream step state) = Stream step' (state, initial)

1040+ where

1041+ {-# INLINE_LATE step' #-}

1042+ step' gst (st, action) = do

1043+ sv <- action

1044+ (r, sv') <- State.runStateT (step (adaptState gst) st) sv

1045+ return $ case r of

1046+ Yield x s -> Yield x (s, return sv')

1047+ Skip s -> Skip (s, return sv')

1048+ Stop -> Stop

1049+

1050+{-# INLINE_NORMAL runStateT #-}

1051+runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)

1052+runStateT initial (Stream step state) = Stream step' (state, initial)

1053+ where

1054+ {-# INLINE_LATE step' #-}

1055+ step' gst (st, action) = do

1056+ sv <- action

1057+ (r, sv') <- State.runStateT (step (adaptState gst) st) sv

1058+ return $ case r of

1059+ Yield x s -> Yield (sv', x) (s, return sv')

1060+ Skip s -> Skip (s, return sv')

1061+ Stop -> Stop

1062+

1063+------------------------------------------------------------------------------

1064+-- Substreams

1065+------------------------------------------------------------------------------

1066+

1067+{-# INLINE_NORMAL isPrefixOf #-}

1068+isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool

1069+isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)

1070+ where

1071+ go (sa, sb, Nothing) = do

1072+ r <- stepa defState sa

1073+ case r of

1074+ Yield x sa' -> go (sa', sb, Just x)

1075+ Skip sa' -> go (sa', sb, Nothing)

1076+ Stop -> return True

1077+

1078+ go (sa, sb, Just x) = do

1079+ r <- stepb defState sb

1080+ case r of

1081+ Yield y sb' ->

1082+ if x == y

1083+ then go (sa, sb', Nothing)

1084+ else return False

1085+ Skip sb' -> go (sa, sb', Just x)

1086+ Stop -> return False

1087+

1088+{-# INLINE_NORMAL isSubsequenceOf #-}

1089+isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool

1090+isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)

1091+ where

1092+ go (sa, sb, Nothing) = do

1093+ r <- stepa defState sa

1094+ case r of

1095+ Yield x sa' -> go (sa', sb, Just x)

1096+ Skip sa' -> go (sa', sb, Nothing)

1097+ Stop -> return True

1098+

1099+ go (sa, sb, Just x) = do

1100+ r <- stepb defState sb

1101+ case r of

1102+ Yield y sb' ->

1103+ if x == y

1104+ then go (sa, sb', Nothing)

1105+ else go (sa, sb', Just x)

1106+ Skip sb' -> go (sa, sb', Just x)

1107+ Stop -> return False

1108+

1109+{-# INLINE_NORMAL stripPrefix #-}

1110+stripPrefix

1111+ :: (Eq a, Monad m)

1112+ => Stream m a -> Stream m a -> m (Maybe (Stream m a))

1113+stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)

1114+ where

1115+ go (sa, sb, Nothing) = do

1116+ r <- stepa defState sa

1117+ case r of

1118+ Yield x sa' -> go (sa', sb, Just x)

1119+ Skip sa' -> go (sa', sb, Nothing)

1120+ Stop -> return $ Just (Stream stepb sb)

1121+

1122+ go (sa, sb, Just x) = do

1123+ r <- stepb defState sb

1124+ case r of

1125+ Yield y sb' ->

1126+ if x == y

1127+ then go (sa, sb', Nothing)

1128+ else return Nothing

1129+ Skip sb' -> go (sa, sb', Just x)

1130+ Stop -> return Nothing

1131+

1132+------------------------------------------------------------------------------

1133+-- Map and Fold

1134+------------------------------------------------------------------------------

1135+

1136+-- | Execute a monadic action for each element of the 'Stream'

1137+{-# INLINE_NORMAL mapM_ #-}

1138+mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()

1139+mapM_ m = drain . mapM m

src/Streamly/Internal/Data/Stream/StreamD/Generate.hs+3−113

@@ -68,15 +68,9 @@ module Streamly.Internal.Data.Stream.StreamD.Generate

68where

70import Control.Concurrent (threadDelay)

71−import Control.Exception

71+import Control.Monad (forever)

72− (fromException)

73−import Control.Monad (when, forever)

74−import Control.Monad.Catch (throwM)

75import Control.Monad.IO.Class (MonadIO(..))

72import Control.Monad.IO.Class (MonadIO(..))

76import Data.Int (Int64)

73import Data.Int (Int64)

77−import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

78−import Data.Maybe (isNothing)

79−import System.Mem (performMajorGC)

80import Streamly.Internal.Data.Time.Units

74import Streamly.Internal.Data.Time.Units

81 (toRelTime64, RelTime64)

75 (toRelTime64, RelTime64)

82import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

76import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

@@ -85,15 +79,8 @@ import Streamly.Internal.Data.Time.Units

85import Streamly.Internal.Data.Unfold.Types (Unfold(..))

79import Streamly.Internal.Data.Unfold.Types (Unfold(..))

87import qualified Streamly.Internal.Data.IORef.Prim as Prim

81import qualified Streamly.Internal.Data.IORef.Prim as Prim

88−import qualified Streamly.Internal.Data.Stream.StreamK as K

82+

89−import qualified Prelude

90−

91−import Prelude hiding

83+import Prelude hiding (iterate, repeat, replicate, takeWhile)

92− ( map, mapM, mapM_, repeat, foldr, last, take, filter

93− , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem

94− , notElem, null, head, tail, zipWith, lookup, foldr1, sequence

95− , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat

96− , reverse, iterate, splitAt)

97import Streamly.Internal.Data.Stream.StreamD.Type

84import Streamly.Internal.Data.Stream.StreamD.Type

98import Streamly.Internal.Data.Stream.StreamD.Common

85import Streamly.Internal.Data.Stream.StreamD.Common

99import Streamly.Internal.Data.SVar

86import Streamly.Internal.Data.SVar

@@ -107,20 +94,6 @@ import Streamly.Internal.Data.SVar

107nil :: Monad m => Stream m a

94nil :: Monad m => Stream m a

108nil = Stream (\_ _ -> return Stop) ()

95nil = Stream (\_ _ -> return Stop) ()

109

110−{-# INLINE_NORMAL consM #-}

111−consM :: Monad m => m a -> Stream m a -> Stream m a

112−consM m (Stream step state) = Stream step1 Nothing

113− where

114− {-# INLINE_LATE step1 #-}

115− step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)

116− step1 gst (Just st) = do

117− r <- step gst st

118− return $

119− case r of

120− Yield a s -> Yield a (Just s)

121− Skip s -> Skip (Just s)

122− Stop -> Stop

123−

124-- XXX implement in terms of consM?

97-- XXX implement in terms of consM?

125-- cons x = consM (return x)

98-- cons x = consM (return x)

126--

99--

@@ -440,86 +413,3 @@ fromListM = Stream step

440 {-# INLINE_LATE step #-}

413 {-# INLINE_LATE step #-}

441 step _ (m:ms) = m >>= \x -> return $ Yield x ms

414 step _ (m:ms) = m >>= \x -> return $ Yield x ms

442 step _ [] = return Stop

415 step _ [] = return Stop

443−

444−{-# INLINE toStreamD #-}

445−toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a

446−toStreamD = fromStreamK . K.toStream

447−

448−-------------------------------------------------------------------------------

449−-- Generation from SVar

450−-------------------------------------------------------------------------------

451−

452−{-# INLINE_NORMAL fromSVar #-}

453−fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a

454−fromSVar svar = Stream step FromSVarInit

455− where

456−

457− {-# INLINE_LATE step #-}

458− step _ FromSVarInit = do

459− ref <- liftIO $ newIORef ()

460− _ <- liftIO $ mkWeakIORef ref hook

461− -- when this copy of svar gets garbage collected "ref" will get

462− -- garbage collected and our GC hook will be called.

463− let sv = svar{svarRef = Just ref}

464− return $ Skip (FromSVarRead sv)

465−

466− where

467−

468− {-# NOINLINE hook #-}

469− hook = do

470− when (svarInspectMode svar) $ do

471− r <- liftIO $ readIORef (svarStopTime (svarStats svar))

472− when (isNothing r) $

473− printSVar svar "SVar Garbage Collected"

474− cleanupSVar svar

475− -- If there are any SVars referenced by this SVar a GC will prompt

476− -- them to be cleaned up quickly.

477− when (svarInspectMode svar) performMajorGC

478−

479− step _ (FromSVarRead sv) = do

480− list <- readOutputQ sv

481− -- Reversing the output is important to guarantee that we process the

482− -- outputs in the same order as they were generated by the constituent

483− -- streams.

484− return $ Skip $ FromSVarLoop sv (Prelude.reverse list)

485−

486− step _ (FromSVarLoop sv []) = do

487− done <- postProcess sv

488− return $ Skip $ if done

489− then (FromSVarDone sv)

490− else (FromSVarRead sv)

491−

492− step _ (FromSVarLoop sv (ev : es)) = do

493− case ev of

494− ChildYield a -> return $ Yield a (FromSVarLoop sv es)

495− ChildStop tid e -> do

496− accountThread sv tid

497− case e of

498− Nothing -> do

499− stop <- shouldStop tid

500− if stop

501− then do

502− liftIO (cleanupSVar sv)

503− return $ Skip (FromSVarDone sv)

504− else return $ Skip (FromSVarLoop sv es)

505− Just ex ->

506− case fromException ex of

507− Just ThreadAbort ->

508− return $ Skip (FromSVarLoop sv es)

509− Nothing -> liftIO (cleanupSVar sv) >> throwM ex

510− where

511−

512− shouldStop tid =

513− case svarStopStyle sv of

514− StopNone -> return False

515− StopAny -> return True

516− StopBy -> do

517− sid <- liftIO $ readIORef (svarStopBy sv)

518− return $ if tid == sid then True else False

519−

520− step _ (FromSVarDone sv) = do

521− when (svarInspectMode sv) $ do

522− t <- liftIO $ getTime Monotonic

523− liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)

524− liftIO $ printSVar sv "SVar Done"

525− return Stop

src/Streamly/Internal/Data/Stream/StreamD/Transform.hsnew file+1040

@@ -0,0 +1,1040 @@

1+#include "inline.hs"

3+-- |

4+-- Module : Streamly.Internal.Data.Stream.StreamD.Transform

6+-- License : BSD-3-Clause

7+-- Maintainer : streamly@composewell.com

8+-- Stability : experimental

9+-- Portability : GHC

10+

11+module Streamly.Internal.Data.Stream.StreamD.Transform

12+ (

13+ -- * Transformation

14+ transform

15+

16+ -- ** By folding (scans)

17+ , scanlM'

18+ , scanlMAfter'

19+ , scanl'

20+ , scanlM

21+ , scanl

22+ , scanl1M'

23+ , scanl1'

24+ , scanl1M

25+ , scanl1

26+

27+ , prescanl'

28+ , prescanlM'

29+

30+ , postscanl

31+ , postscanlM

32+ , postscanl'

33+ , postscanlM'

34+ , postscanlMAfter'

35+

36+ , postscanlx'

37+ , postscanlMx'

38+ , scanlMx'

39+ , scanlx'

40+ , postscanOnce

41+ , scanOnce

42+

43+ -- * Filtering

44+ , filter

45+ , filterM

46+ , uniq

47+ , take

48+ , takeByTime

49+ , takeWhile

50+ , takeWhileM

51+ , drop

52+ , dropByTime

53+ , dropWhile

54+ , dropWhileM

55+

56+ -- * Mapping

57+ , map

58+ , mapM

59+ , sequence

60+ , rollingMap

61+ , rollingMapM

62+

63+ -- * Inserting

64+ , intersperseM

65+ , intersperseM_

66+ , intersperse

67+ , intersperseSuffix

68+ , intersperseSuffix_

69+ , intersperseSuffixBySpan

70+ , insertBy

71+

72+ -- * Deleting

73+ , deleteBy

74+

75+ -- ** Map and Filter

76+ , mapMaybe

77+ , mapMaybeM

78+

79+ -- * Zipping

80+ , indexed

81+ , indexedR

82+ , zipWith

83+ , zipWithM

84+

85+ -- * Comparisons

86+ , eqBy

87+ , cmpBy

88+

89+ -- * Merging

90+ , mergeBy

91+ , mergeByM

92+

93+ -- * Transformation comprehensions

94+ , the

95+

96+ -- * Concurrent Application

97+ , mkParallel

98+ , mkParallelD

99+ , newCallbackStream

100+ )

101+where

102+

103+import Control.Concurrent (myThreadId)

104+import Control.Monad (void)

105+import Control.Monad.IO.Class (MonadIO(..))

106+import Data.Functor.Identity (Identity(..))

107+import Data.Maybe (fromJust, isJust)

108+import Streamly.Internal.Data.Time.Units

109+ (TimeUnit64, toRelTime64, diffAbsTime64)

110+import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))

111+import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

112+

113+import qualified Streamly.Internal.Data.Pipe.Types as Pipe

114+import qualified Streamly.Internal.Data.Fold as FL

115+import qualified Streamly.Internal.Data.Stream.StreamK as K

116+

117+import Prelude hiding

118+ ( drop, dropWhile, filter, map, mapM

119+ , scanl, scanl1, sequence, take, takeWhile, zipWith)

120+import Streamly.Internal.Data.Stream.StreamD.Type

121+import Streamly.Internal.Data.Stream.StreamD.Common

122+import Streamly.Internal.Data.SVar

123+

124+-------------------------------------------------------------------------------

125+-- General transformation

126+-------------------------------------------------------------------------------

127+

128+{-# INLINE_NORMAL transform #-}

129+transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b

130+transform (Pipe pstep1 pstep2 pstate) (Stream step state) =

131+ Stream step' (Consume pstate, state)

132+

133+ where

134+

135+ {-# INLINE_LATE step' #-}

136+

137+ step' gst (Consume pst, st) = pst `seq` do

138+ r <- step (adaptState gst) st

139+ case r of

140+ Yield x s -> do

141+ res <- pstep1 pst x

142+ case res of

143+ Pipe.Yield b pst' -> return $ Yield b (pst', s)

144+ Pipe.Continue pst' -> return $ Skip (pst', s)

145+ Skip s -> return $ Skip (Consume pst, s)

146+ Stop -> return Stop

147+

148+ step' _ (Produce pst, st) = pst `seq` do

149+ res <- pstep2 pst

150+ case res of

151+ Pipe.Yield b pst' -> return $ Yield b (pst', st)

152+ Pipe.Continue pst' -> return $ Skip (pst', st)

153+

154+------------------------------------------------------------------------------

155+-- Transformation by Folding (Scans)

156+------------------------------------------------------------------------------

157+

158+------------------------------------------------------------------------------

159+-- Prescans

160+------------------------------------------------------------------------------

161+

162+-- XXX Is a prescan useful, discarding the last step does not sound useful? I

163+-- am not sure about the utility of this function, so this is implemented but

164+-- not exposed. We can expose it if someone provides good reasons why this is

165+-- useful.

166+--

167+-- XXX We have to execute the stream one step ahead to know that we are at the

168+-- last step. The vector implementation of prescan executes the last fold step

169+-- but does not yield the result. This means we have executed the effect but

170+-- discarded value. This does not sound right. In this implementation we are

171+-- not executing the last fold step.

172+{-# INLINE_NORMAL prescanlM' #-}

173+prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

174+prescanlM' f mz (Stream step state) = Stream step' (state, mz)

175+ where

176+ {-# INLINE_LATE step' #-}

177+ step' gst (st, prev) = do

178+ r <- step (adaptState gst) st

179+ case r of

180+ Yield x s -> do

181+ acc <- prev

182+ return $ Yield acc (s, f acc x)

183+ Skip s -> return $ Skip (s, prev)

184+ Stop -> return Stop

185+

186+{-# INLINE prescanl' #-}

187+prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

188+prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)

189+

190+------------------------------------------------------------------------------

191+-- Monolithic postscans (postscan followed by a map)

192+------------------------------------------------------------------------------

193+

194+-- The performance of a modular postscan followed by a map seems to be

195+-- equivalent to this monolithic scan followed by map therefore we may not need

196+-- this implementation. We just have it for performance comparison and in case

197+-- modular version does not perform well in some situation.

198+--

199+{-# INLINE_NORMAL postscanlMx' #-}

200+postscanlMx' :: Monad m

201+ => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b

202+postscanlMx' fstep begin done (Stream step state) = do

203+ Stream step' (state, begin)

204+ where

205+ {-# INLINE_LATE step' #-}

206+ step' gst (st, acc) = do

207+ r <- step (adaptState gst) st

208+ case r of

209+ Yield x s -> do

210+ old <- acc

211+ y <- fstep old x

212+ v <- done y

213+ v `seq` y `seq` return (Yield v (s, return y))

214+ Skip s -> return $ Skip (s, acc)

215+ Stop -> return Stop

216+

217+{-# INLINE_NORMAL postscanlx' #-}

218+postscanlx' :: Monad m

219+ => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b

220+postscanlx' fstep begin done s =

221+ postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s

222+

223+-- XXX do we need consM strict to evaluate the begin value?

224+{-# INLINE scanlMx' #-}

225+scanlMx' :: Monad m

226+ => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b

227+scanlMx' fstep begin done s =

228+ (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s

229+

230+data PostScanState s f = PostScan s !f

231+

232+{-# INLINE_NORMAL postscanOnce #-}

233+postscanOnce :: Monad m => FL.Fold m a b -> Stream m a -> Stream m b

234+postscanOnce (FL.Fold fstep begin done) (Stream step state) =

235+ Stream step' (PostScan state begin)

236+

237+ where

238+

239+ {-# INLINE_LATE step' #-}

240+ step' gst (PostScan st acc) = do

241+ r <- step (adaptState gst) st

242+ case r of

243+ Yield x s -> do

244+ old <- acc

245+ res <- fstep old x

246+ case res of

247+ FL.Partial fs -> do

248+ !v <- done fs

249+ return $ Yield v $ PostScan s (return fs)

250+ FL.Done _ -> return Stop

251+ Skip s -> return $ Skip $ PostScan s acc

252+ Stop -> return Stop

253+

254+{-# INLINE scanOnce #-}

255+scanOnce :: Monad m

256+ => FL.Fold m a b -> Stream m a -> Stream m b

257+scanOnce fld@(FL.Fold _ begin done) s =

258+ (begin >>= \x -> x `seq` done x) `consM` postscanOnce fld s

259+

260+{-# INLINE scanlx' #-}

261+scanlx' :: Monad m

262+ => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b

263+scanlx' fstep begin done s =

264+ scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s

265+

266+------------------------------------------------------------------------------

267+-- postscans

268+------------------------------------------------------------------------------

269+

270+{-# INLINE_NORMAL postscanlM' #-}

271+postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

272+postscanlM' fstep begin (Stream step state) =

273+ Stream step' Nothing

274+ where

275+ {-# INLINE_LATE step' #-}

276+ step' _ Nothing = do

277+ !x <- begin

278+ return $ Skip (Just (state, x))

279+

280+ step' gst (Just (st, acc)) = do

281+ r <- step (adaptState gst) st

282+ case r of

283+ Yield x s -> do

284+ !y <- fstep acc x

285+ return $ Yield y (Just (s, y))

286+ Skip s -> return $ Skip (Just (s, acc))

287+ Stop -> return Stop

288+

289+{-# INLINE_NORMAL postscanl' #-}

290+postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

291+postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)

292+

293+-- We can possibly have the "done" function as a Maybe to provide an option to

294+-- emit or not emit the accumulator when the stream stops.

295+--

296+-- TBD: use a single Yield point

297+--

298+{-# INLINE_NORMAL postscanlMAfter' #-}

299+postscanlMAfter' :: Monad m

300+ => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

301+postscanlMAfter' fstep initial done (Stream step1 state1) = do

302+ Stream step (Just (state1, initial))

303+

304+ where

305+

306+ {-# INLINE_LATE step #-}

307+ step gst (Just (st, acc)) = do

308+ r <- step1 (adaptState gst) st

309+ case r of

310+ Yield x s -> do

311+ old <- acc

312+ y <- fstep old x

313+ y `seq` return (Yield y (Just (s, return y)))

314+ Skip s -> return $ Skip $ Just (s, acc)

315+ Stop -> acc >>= done >>= \res -> return (Yield res Nothing)

316+ step _ Nothing = return Stop

317+

318+{-# INLINE_NORMAL postscanlM #-}

319+postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

320+postscanlM fstep begin (Stream step state) = Stream step' Nothing

321+ where

322+ {-# INLINE_LATE step' #-}

323+ step' _ Nothing = do

324+ r <- begin

325+ return $ Skip (Just (state, r))

326+

327+ step' gst (Just (st, acc)) = do

328+ r <- step (adaptState gst) st

329+ case r of

330+ Yield x s -> do

331+ y <- fstep acc x

332+ return (Yield y (Just (s, y)))

333+ Skip s -> return $ Skip (Just (s, acc))

334+ Stop -> return Stop

335+

336+{-# INLINE_NORMAL postscanl #-}

337+postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

338+postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)

339+

340+{-# INLINE_NORMAL scanlM' #-}

341+scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

342+scanlM' fstep begin (Stream step state) = Stream step' Nothing

343+ where

344+ {-# INLINE_LATE step' #-}

345+ step' _ Nothing = do

346+ !x <- begin

347+ return $ Yield x (Just (state, x))

348+ step' gst (Just (st, acc)) = do

349+ r <- step (adaptState gst) st

350+ case r of

351+ Yield x s -> do

352+ !y <- fstep acc x

353+ return $ Yield y (Just (s, y))

354+ Skip s -> return $ Skip (Just (s, acc))

355+ Stop -> return Stop

356+

357+{-# INLINE scanlMAfter' #-}

358+scanlMAfter' :: Monad m

359+ => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

360+scanlMAfter' fstep initial done s =

361+ (initial >>= \x -> x `seq` return x) `consM`

362+ postscanlMAfter' fstep initial done s

363+

364+{-# INLINE scanl' #-}

365+scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

366+scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)

367+

368+{-# INLINE_NORMAL scanlM #-}

369+scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b

370+scanlM fstep begin (Stream step state) = Stream step' Nothing

371+ where

372+ {-# INLINE_LATE step' #-}

373+ step' _ Nothing = do

374+ x <- begin

375+ return $ Yield x (Just (state, x))

376+ step' gst (Just (st, acc)) = do

377+ r <- step (adaptState gst) st

378+ case r of

379+ Yield x s -> do

380+ y <- fstep acc x

381+ return $ Yield y (Just (s, y))

382+ Skip s -> return $ Skip (Just (s, acc))

383+ Stop -> return $ Stop

384+

385+{-# INLINE scanl #-}

386+scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

387+scanl f seed = scanlM (\a b -> return (f a b)) (return seed)

388+

389+{-# INLINE_NORMAL scanl1M #-}

390+scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a

391+scanl1M fstep (Stream step state) = Stream step' (state, Nothing)

392+ where

393+ {-# INLINE_LATE step' #-}

394+ step' gst (st, Nothing) = do

395+ r <- step gst st

396+ case r of

397+ Yield x s -> return $ Yield x (s, Just x)

398+ Skip s -> return $ Skip (s, Nothing)

399+ Stop -> return Stop

400+

401+ step' gst (st, Just acc) = do

402+ r <- step gst st

403+ case r of

404+ Yield y s -> do

405+ z <- fstep acc y

406+ return $ Yield z (s, Just z)

407+ Skip s -> return $ Skip (s, Just acc)

408+ Stop -> return Stop

409+

410+{-# INLINE scanl1 #-}

411+scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a

412+scanl1 f = scanl1M (\x y -> return (f x y))

413+

414+{-# INLINE_NORMAL scanl1M' #-}

415+scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a

416+scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)

417+ where

418+ {-# INLINE_LATE step' #-}

419+ step' gst (st, Nothing) = do

420+ r <- step gst st

421+ case r of

422+ Yield x s -> x `seq` return $ Yield x (s, Just x)

423+ Skip s -> return $ Skip (s, Nothing)

424+ Stop -> return Stop

425+

426+ step' gst (st, Just acc) = acc `seq` do

427+ r <- step gst st

428+ case r of

429+ Yield y s -> do

430+ z <- fstep acc y

431+ z `seq` return $ Yield z (s, Just z)

432+ Skip s -> return $ Skip (s, Just acc)

433+ Stop -> return Stop

434+

435+{-# INLINE scanl1' #-}

436+scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a

437+scanl1' f = scanl1M' (\x y -> return (f x y))

438+

439+------------------------------------------------------------------------------

440+-- Stateful map/scan

441+------------------------------------------------------------------------------

442+

443+data RollingMapState s a = RollingMapInit s | RollingMapGo s a

444+

445+{-# INLINE rollingMapM #-}

446+rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b

447+rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)

448+ where

449+ step gst (RollingMapInit st) = do

450+ r <- step1 (adaptState gst) st

451+ return $ case r of

452+ Yield x s -> Skip $ RollingMapGo s x

453+ Skip s -> Skip $ RollingMapInit s

454+ Stop -> Stop

455+

456+ step gst (RollingMapGo s1 x1) = do

457+ r <- step1 (adaptState gst) s1

458+ case r of

459+ Yield x s -> do

460+ !res <- f x x1

461+ return $ Yield res $ RollingMapGo s x

462+ Skip s -> return $ Skip $ RollingMapGo s x1

463+ Stop -> return $ Stop

464+

465+{-# INLINE rollingMap #-}

466+rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b

467+rollingMap f = rollingMapM (\x y -> return $ f x y)

468+

469+-------------------------------------------------------------------------------

470+-- Filtering

471+-------------------------------------------------------------------------------

472+

473+{-# INLINE_NORMAL drop #-}

474+drop :: Monad m => Int -> Stream m a -> Stream m a

475+drop n (Stream step state) = Stream step' (state, Just n)

476+ where

477+ {-# INLINE_LATE step' #-}

478+ step' gst (st, Just i)

479+ | i > 0 = do

480+ r <- step gst st

481+ return $

482+ case r of

483+ Yield _ s -> Skip (s, Just (i - 1))

484+ Skip s -> Skip (s, Just i)

485+ Stop -> Stop

486+ | otherwise = return $ Skip (st, Nothing)

487+

488+ step' gst (st, Nothing) = do

489+ r <- step gst st

490+ return $

491+ case r of

492+ Yield x s -> Yield x (s, Nothing)

493+ Skip s -> Skip (s, Nothing)

494+ Stop -> Stop

495+

496+data DropWhileState s a

497+ = DropWhileDrop s

498+ | DropWhileYield a s

499+ | DropWhileNext s

500+

501+{-# INLINE_NORMAL dropWhileM #-}

502+dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a

503+dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)

504+ where

505+ {-# INLINE_LATE step' #-}

506+ step' gst (DropWhileDrop st) = do

507+ r <- step gst st

508+ case r of

509+ Yield x s -> do

510+ b <- f x

511+ if b

512+ then return $ Skip (DropWhileDrop s)

513+ else return $ Skip (DropWhileYield x s)

514+ Skip s -> return $ Skip (DropWhileDrop s)

515+ Stop -> return Stop

516+

517+ step' gst (DropWhileNext st) = do

518+ r <- step gst st

519+ case r of

520+ Yield x s -> return $ Skip (DropWhileYield x s)

521+ Skip s -> return $ Skip (DropWhileNext s)

522+ Stop -> return Stop

523+

524+ step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)

525+

526+{-# INLINE dropWhile #-}

527+dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

528+dropWhile f = dropWhileM (return . f)

529+

530+{-# INLINE_NORMAL filterM #-}

531+filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a

532+filterM f (Stream step state) = Stream step' state

533+ where

534+ {-# INLINE_LATE step' #-}

535+ step' gst st = do

536+ r <- step gst st

537+ case r of

538+ Yield x s -> do

539+ b <- f x

540+ return $ if b

541+ then Yield x s

542+ else Skip s

543+ Skip s -> return $ Skip s

544+ Stop -> return Stop

545+

546+{-# INLINE filter #-}

547+filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

548+filter f = filterM (return . f)

549+

550+{-# INLINE_NORMAL uniq #-}

551+uniq :: (Eq a, Monad m) => Stream m a -> Stream m a

552+uniq (Stream step state) = Stream step' (Nothing, state)

553+ where

554+ {-# INLINE_LATE step' #-}

555+ step' gst (Nothing, st) = do

556+ r <- step gst st

557+ case r of

558+ Yield x s -> return $ Yield x (Just x, s)

559+ Skip s -> return $ Skip (Nothing, s)

560+ Stop -> return Stop

561+ step' gst (Just x, st) = do

562+ r <- step gst st

563+ case r of

564+ Yield y s | x == y -> return $ Skip (Just x, s)

565+ | otherwise -> return $ Yield y (Just y, s)

566+ Skip s -> return $ Skip (Just x, s)

567+ Stop -> return Stop

568+

569+------------------------------------------------------------------------------

570+-- Transformation by Mapping

571+------------------------------------------------------------------------------

572+

573+{-# INLINE_NORMAL sequence #-}

574+sequence :: Monad m => Stream m (m a) -> Stream m a

575+sequence (Stream step state) = Stream step' state

576+ where

577+ {-# INLINE_LATE step' #-}

578+ step' gst st = do

579+ r <- step (adaptState gst) st

580+ case r of

581+ Yield x s -> x >>= \a -> return (Yield a s)

582+ Skip s -> return $ Skip s

583+ Stop -> return Stop

584+

585+------------------------------------------------------------------------------

586+-- Inserting

587+------------------------------------------------------------------------------

588+

589+data LoopState x s = FirstYield s

590+ | InterspersingYield s

591+ | YieldAndCarry x s

592+

593+{-# INLINE_NORMAL intersperseM #-}

594+intersperseM :: Monad m => m a -> Stream m a -> Stream m a

595+intersperseM m (Stream step state) = Stream step' (FirstYield state)

596+ where

597+ {-# INLINE_LATE step' #-}

598+ step' gst (FirstYield st) = do

599+ r <- step gst st

600+ return $

601+ case r of

602+ Yield x s -> Skip (YieldAndCarry x s)

603+ Skip s -> Skip (FirstYield s)

604+ Stop -> Stop

605+

606+ step' gst (InterspersingYield st) = do

607+ r <- step gst st

608+ case r of

609+ Yield x s -> do

610+ a <- m

611+ return $ Yield a (YieldAndCarry x s)

612+ Skip s -> return $ Skip $ InterspersingYield s

613+ Stop -> return Stop

614+

615+ step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)

616+

617+{-# INLINE_NORMAL intersperseM_ #-}

618+intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a

619+intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))

620+ where

621+ {-# INLINE_LATE step #-}

622+ step gst (Left (eff, st)) = do

623+ r <- step1 gst st

624+ case r of

625+ Yield x s -> eff >> return (Yield x (Right s))

626+ Skip s -> return $ Skip (Left (eff, s))

627+ Stop -> return Stop

628+

629+ step _ (Right st) = return $ Skip $ Left (void m, st)

630+

631+data SuffixState s a

632+ = SuffixElem s

633+ | SuffixSuffix s

634+ | SuffixYield a (SuffixState s a)

635+

636+{-# INLINE_NORMAL intersperseSuffix #-}

637+intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a

638+intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)

639+ where

640+ {-# INLINE_LATE step' #-}

641+ step' gst (SuffixElem st) = do

642+ r <- step gst st

643+ return $ case r of

644+ Yield x s -> Skip (SuffixYield x (SuffixSuffix s))

645+ Skip s -> Skip (SuffixElem s)

646+ Stop -> Stop

647+

648+ step' _ (SuffixSuffix st) = do

649+ action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))

650+

651+ step' _ (SuffixYield x next) = return $ Yield x next

652+

653+{-# INLINE_NORMAL intersperseSuffix_ #-}

654+intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a

655+intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)

656+ where

657+ {-# INLINE_LATE step #-}

658+ step gst (Left st) = do

659+ r <- step1 gst st

660+ case r of

661+ Yield x s -> return $ Yield x (Right s)

662+ Skip s -> return $ Skip $ Left s

663+ Stop -> return Stop

664+

665+ step _ (Right st) = m >> return (Skip (Left st))

666+

667+data SuffixSpanState s a

668+ = SuffixSpanElem s Int

669+ | SuffixSpanSuffix s

670+ | SuffixSpanYield a (SuffixSpanState s a)

671+ | SuffixSpanLast

672+ | SuffixSpanStop

673+

674+-- | intersperse after every n items

675+{-# INLINE_NORMAL intersperseSuffixBySpan #-}

676+intersperseSuffixBySpan :: forall m a. Monad m

677+ => Int -> m a -> Stream m a -> Stream m a

678+intersperseSuffixBySpan n action (Stream step state) =

679+ Stream step' (SuffixSpanElem state n)

680+ where

681+ {-# INLINE_LATE step' #-}

682+ step' gst (SuffixSpanElem st i) | i > 0 = do

683+ r <- step gst st

684+ return $ case r of

685+ Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))

686+ Skip s -> Skip (SuffixSpanElem s i)

687+ Stop -> if i == n then Stop else Skip SuffixSpanLast

688+ step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)

689+

690+ step' _ (SuffixSpanSuffix st) = do

691+ action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))

692+

693+ step' _ (SuffixSpanLast) = do

694+ action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)

695+

696+ step' _ (SuffixSpanYield x next) = return $ Yield x next

697+

698+ step' _ (SuffixSpanStop) = return Stop

699+

700+{-# INLINE intersperse #-}

701+intersperse :: Monad m => a -> Stream m a -> Stream m a

702+intersperse a = intersperseM (return a)

703+

704+{-# INLINE_NORMAL insertBy #-}

705+insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a

706+insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)

707+ where

708+ {-# INLINE_LATE step' #-}

709+ step' gst (st, False, _) = do

710+ r <- step gst st

711+ case r of

712+ Yield x s -> case cmp a x of

713+ GT -> return $ Yield x (s, False, Nothing)

714+ _ -> return $ Yield a (s, True, Just x)

715+ Skip s -> return $ Skip (s, False, Nothing)

716+ Stop -> return $ Yield a (st, True, Nothing)

717+

718+ step' _ (_, True, Nothing) = return Stop

719+

720+ step' gst (st, True, Just prev) = do

721+ r <- step gst st

722+ case r of

723+ Yield x s -> return $ Yield prev (s, True, Just x)

724+ Skip s -> return $ Skip (s, True, Just prev)

725+ Stop -> return $ Yield prev (st, True, Nothing)

726+

727+------------------------------------------------------------------------------

728+-- Deleting

729+------------------------------------------------------------------------------

730+

731+{-# INLINE_NORMAL deleteBy #-}

732+deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a

733+deleteBy eq x (Stream step state) = Stream step' (state, False)

734+ where

735+ {-# INLINE_LATE step' #-}

736+ step' gst (st, False) = do

737+ r <- step gst st

738+ case r of

739+ Yield y s -> return $

740+ if eq x y then Skip (s, True) else Yield y (s, False)

741+ Skip s -> return $ Skip (s, False)

742+ Stop -> return Stop

743+

744+ step' gst (st, True) = do

745+ r <- step gst st

746+ case r of

747+ Yield y s -> return $ Yield y (s, True)

748+ Skip s -> return $ Skip (s, True)

749+ Stop -> return Stop

750+

751+------------------------------------------------------------------------------

752+-- Transformation by Map and Filter

753+------------------------------------------------------------------------------

754+

755+-- XXX Will this always fuse properly?

756+{-# INLINE_NORMAL mapMaybe #-}

757+mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b

758+mapMaybe f = fmap fromJust . filter isJust . map f

759+

760+{-# INLINE_NORMAL mapMaybeM #-}

761+mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b

762+mapMaybeM f = fmap fromJust . filter isJust . mapM f

763+

764+------------------------------------------------------------------------------

765+-- Zipping

766+------------------------------------------------------------------------------

767+

768+{-# INLINE_NORMAL indexed #-}

769+indexed :: Monad m => Stream m a -> Stream m (Int, a)

770+indexed (Stream step state) = Stream step' (state, 0)

771+ where

772+ {-# INLINE_LATE step' #-}

773+ step' gst (st, i) = i `seq` do

774+ r <- step (adaptState gst) st

775+ case r of

776+ Yield x s -> return $ Yield (i, x) (s, i+1)

777+ Skip s -> return $ Skip (s, i)

778+ Stop -> return Stop

779+

780+{-# INLINE_NORMAL indexedR #-}

781+indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)

782+indexedR m (Stream step state) = Stream step' (state, m)

783+ where

784+ {-# INLINE_LATE step' #-}

785+ step' gst (st, i) = i `seq` do

786+ r <- step (adaptState gst) st

787+ case r of

788+ Yield x s -> let i' = i - 1

789+ in return $ Yield (i, x) (s, i')

790+ Skip s -> return $ Skip (s, i)

791+ Stop -> return Stop

792+

793+{-# INLINE_NORMAL zipWithM #-}

794+zipWithM :: Monad m

795+ => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c

796+zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)

797+ where

798+ {-# INLINE_LATE step #-}

799+ step gst (sa, sb, Nothing) = do

800+ r <- stepa (adaptState gst) sa

801+ return $

802+ case r of

803+ Yield x sa' -> Skip (sa', sb, Just x)

804+ Skip sa' -> Skip (sa', sb, Nothing)

805+ Stop -> Stop

806+

807+ step gst (sa, sb, Just x) = do

808+ r <- stepb (adaptState gst) sb

809+ case r of

810+ Yield y sb' -> do

811+ z <- f x y

812+ return $ Yield z (sa, sb', Nothing)

813+ Skip sb' -> return $ Skip (sa, sb', Just x)

814+ Stop -> return Stop

815+

816+#if __GLASGOW_HASKELL__ >= 801

817+{-# RULES "zipWithM xs xs"

818+ forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}

819+#endif

820+

821+{-# INLINE zipWith #-}

822+zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c

823+zipWith f = zipWithM (\a b -> return (f a b))

824+

825+------------------------------------------------------------------------------

826+-- Merging

827+------------------------------------------------------------------------------

828+

829+{-# INLINE_NORMAL mergeByM #-}

830+mergeByM

831+ :: (Monad m)

832+ => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a

833+mergeByM cmp (Stream stepa ta) (Stream stepb tb) =

834+ Stream step (Just ta, Just tb, Nothing, Nothing)

835+ where

836+ {-# INLINE_LATE step #-}

837+

838+ -- one of the values is missing, and the corresponding stream is running

839+ step gst (Just sa, sb, Nothing, b) = do

840+ r <- stepa gst sa

841+ return $ case r of

842+ Yield a sa' -> Skip (Just sa', sb, Just a, b)

843+ Skip sa' -> Skip (Just sa', sb, Nothing, b)

844+ Stop -> Skip (Nothing, sb, Nothing, b)

845+

846+ step gst (sa, Just sb, a, Nothing) = do

847+ r <- stepb gst sb

848+ return $ case r of

849+ Yield b sb' -> Skip (sa, Just sb', a, Just b)

850+ Skip sb' -> Skip (sa, Just sb', a, Nothing)

851+ Stop -> Skip (sa, Nothing, a, Nothing)

852+

853+ -- both the values are available

854+ step _ (sa, sb, Just a, Just b) = do

855+ res <- cmp a b

856+ return $ case res of

857+ GT -> Yield b (sa, sb, Just a, Nothing)

858+ _ -> Yield a (sa, sb, Nothing, Just b)

859+

860+ -- one of the values is missing, corresponding stream is done

861+ step _ (Nothing, sb, Nothing, Just b) =

862+ return $ Yield b (Nothing, sb, Nothing, Nothing)

863+

864+ step _ (sa, Nothing, Just a, Nothing) =

865+ return $ Yield a (sa, Nothing, Nothing, Nothing)

866+

867+ step _ (Nothing, Nothing, Nothing, Nothing) = return Stop

868+

869+{-# INLINE mergeBy #-}

870+mergeBy

871+ :: (Monad m)

872+ => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a

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

874+

875+------------------------------------------------------------------------------

876+-- Transformation comprehensions

877+------------------------------------------------------------------------------

878+

879+{-# INLINE_NORMAL the #-}

880+the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)

881+the (Stream step state) = go state

882+ where

883+ go st = do

884+ r <- step defState st

885+ case r of

886+ Yield x s -> go' x s

887+ Skip s -> go s

888+ Stop -> return Nothing

889+ go' n st = do

890+ r <- step defState st

891+ case r of

892+ Yield x s | x == n -> go' n s

893+ | otherwise -> return Nothing

894+ Skip s -> go' n s

895+ Stop -> return (Just n)

896+

897+-------------------------------------------------------------------------------

898+-- Concurrent application and fold

899+-------------------------------------------------------------------------------

900+

901+{-# INLINE_NORMAL mkParallelD #-}

902+mkParallelD :: MonadAsync m => Stream m a -> Stream m a

903+mkParallelD m = Stream step Nothing

904+ where

905+

906+ step gst Nothing = do

907+ sv <- newParallelVar StopNone gst

908+ toSVarParallel gst sv m

909+ -- XXX use unfold instead?

910+ return $ Skip $ Just $ fromSVar sv

911+

912+ step gst (Just (UnStream step1 st)) = do

913+ r <- step1 gst st

914+ return $ case r of

915+ Yield a s -> Yield a (Just $ Stream step1 s)

916+ Skip s -> Skip (Just $ Stream step1 s)

917+ Stop -> Stop

918+

919+-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a

920+-- parallel style SVar for evaluation. Currently, parallel style cannot use

921+-- rate control whereas Async style can use rate control. In async style SVar

922+-- the worker thread terminates when the buffer is full whereas in Parallel

923+-- style it blocks.

924+--

925+-- | Make the stream producer and consumer run concurrently by introducing a

926+-- buffer between them. The producer thread evaluates the input stream until

927+-- the buffer fills, it blocks if the buffer is full until there is space in

928+-- the buffer. The consumer consumes the stream lazily from the buffer.

929+--

930+-- /Internal/

931+--

932+{-# INLINE_NORMAL mkParallel #-}

933+mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a

934+mkParallel = fromStreamD . mkParallelD . toStreamD

935+

936+-- Note: we can use another API with two callbacks stop and yield if we want

937+-- the callback to be able to indicate end of stream.

938+--

939+-- | Generates a callback and a stream pair. The callback returned is used to

940+-- queue values to the stream. The stream is infinite, there is no way for the

941+-- callback to indicate that it is done now.

942+--

943+-- /Internal/

944+--

945+{-# INLINE_NORMAL newCallbackStream #-}

946+newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)

947+newCallbackStream = do

948+ sv <- newParallelVar StopNone defState

949+

950+ -- XXX Add our own thread-id to the SVar as we can not know the callback's

951+ -- thread-id and the callback is not run in a managed worker. We need to

952+ -- handle this better.

953+ liftIO myThreadId >>= modifyThread sv

954+

955+ let callback a = liftIO $ void $ send sv (ChildYield a)

956+ -- XXX we can return an SVar and then the consumer can unfold from the

957+ -- SVar?

958+ return (callback, fromStreamD (fromSVar sv))

959+

960+------------------------------------------------------------------------------

961+-- Time related

962+------------------------------------------------------------------------------

963+

964+-- XXX using getTime in the loop can be pretty expensive especially for

965+-- computations where iterations are lightweight. We have the following

966+-- options:

967+--

968+-- 1) Run a timeout thread updating a flag asynchronously and check that

969+-- flag here, that way we can have a cheap termination check.

970+--

971+-- 2) Use COARSE clock to get time with lower resolution but more efficiently.

972+--

973+-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the

974+-- termination value of rdtsc in the beginning and then in each iteration just

975+-- get rdtsc and check if we should terminate.

976+--

977+data TakeByTime st s

978+ = TakeByTimeInit st

979+ | TakeByTimeCheck st s

980+ | TakeByTimeYield st s

981+

982+{-# INLINE_NORMAL takeByTime #-}

983+takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a

984+takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)

985+ where

986+

987+ lim = toRelTime64 duration

988+

989+ {-# INLINE_LATE step #-}

990+ step _ (TakeByTimeInit _) | lim == 0 = return Stop

991+ step _ (TakeByTimeInit st) = do

992+ t0 <- liftIO $ getTime Monotonic

993+ return $ Skip (TakeByTimeYield st t0)

994+ step _ (TakeByTimeCheck st t0) = do

995+ t <- liftIO $ getTime Monotonic

996+ return $

997+ if diffAbsTime64 t t0 > lim

998+ then Stop

999+ else Skip (TakeByTimeYield st t0)

1000+ step gst (TakeByTimeYield st t0) = do

1001+ r <- step1 gst st

1002+ return $ case r of

1003+ Yield x s -> Yield x (TakeByTimeCheck s t0)

1004+ Skip s -> Skip (TakeByTimeCheck s t0)

1005+ Stop -> Stop

1006+

1007+data DropByTime st s x

1008+ = DropByTimeInit st

1009+ | DropByTimeGen st s

1010+ | DropByTimeCheck st s x

1011+ | DropByTimeYield st

1012+

1013+{-# INLINE_NORMAL dropByTime #-}

1014+dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a

1015+dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)

1016+ where

1017+

1018+ lim = toRelTime64 duration

1019+

1020+ {-# INLINE_LATE step #-}

1021+ step _ (DropByTimeInit st) = do

1022+ t0 <- liftIO $ getTime Monotonic

1023+ return $ Skip (DropByTimeGen st t0)

1024+ step gst (DropByTimeGen st t0) = do

1025+ r <- step1 gst st

1026+ return $ case r of

1027+ Yield x s -> Skip (DropByTimeCheck s t0 x)

1028+ Skip s -> Skip (DropByTimeGen s t0)

1029+ Stop -> Stop

1030+ step _ (DropByTimeCheck st t0 x) = do

1031+ t <- liftIO $ getTime Monotonic

1032+ if diffAbsTime64 t t0 <= lim

1033+ then return $ Skip $ DropByTimeGen st t0

1034+ else return $ Yield x $ DropByTimeYield st

1035+ step gst (DropByTimeYield st) = do

1036+ r <- step1 gst st

1037+ return $ case r of

1038+ Yield x s -> Yield x (DropByTimeYield s)

1039+ Skip s -> Skip (DropByTimeYield s)

1040+ Stop -> Stop

streamly.cabal+1

@@ -457,6 +457,7 @@ library

457 , Streamly.Internal.Data.Stream.StreamD.SplitGroup

458 , Streamly.Internal.Data.Stream.StreamD.Eliminate

459 , Streamly.Internal.Data.Stream.StreamD.Flatten

460+ , Streamly.Internal.Data.Stream.StreamD.Transform

460 , Streamly.Internal.Data.Stream.StreamD.Exceptions

461 , Streamly.Internal.Data.Stream.StreamD.Exceptions

461 , Streamly.Internal.Data.Stream.StreamDK.Type

462 , Streamly.Internal.Data.Stream.StreamDK.Type

462 , Streamly.Internal.Data.Stream.StreamDK

463 , Streamly.Internal.Data.Stream.StreamDK