cf8311d · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor stream generating functions to StreamD.Generate.

cf8311dPranay Sashank6 years ago

src/Streamly/Internal/Data/Stream/IsStream/Generate.hs+4−2

compare with master

unifiedsplit

@@ -80,14 +80,16 @@ import Streamly.Internal.Data.IORef.Prim (Prim, IORef)

81import qualified Streamly.Internal.Data.Stream.Prelude as P

82import qualified Streamly.Internal.Data.Stream.StreamK as K

83−import qualified Streamly.Internal.Data.Stream.StreamD as D

83+import qualified Streamly.Internal.Data.Stream.StreamD as D (toStreamK)

84+import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D

84import qualified Streamly.Internal.Data.Stream.Serial as Serial

85import qualified Streamly.Internal.Data.Stream.Serial as Serial

85import qualified Prelude

86import qualified Prelude

86import qualified System.IO as IO

87import qualified System.IO as IO

87#ifdef USE_STREAMK_ONLY

88#ifdef USE_STREAMK_ONLY

88import qualified Streamly.Internal.Data.Stream.StreamK as S

89import qualified Streamly.Internal.Data.Stream.StreamK as S

89#else

90#else

90−import qualified Streamly.Internal.Data.Stream.StreamD as S

91+import qualified Streamly.Internal.Data.Stream.StreamD as S (toStreamK)

92+import qualified Streamly.Internal.Data.Stream.StreamD.Generate as S

91#endif

93#endif

93import Prelude hiding

95import Prelude hiding

src/Streamly/Internal/Data/Stream/StreamD.hs+5−414

@@ -304,8 +304,8 @@ where

304

305import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)

306import Control.Exception

307− (assert, Exception, SomeException, AsyncException, fromException, mask_)

307+ (assert, Exception, SomeException, AsyncException, mask_)

308−import Control.Monad (void, when, forever)

308+import Control.Monad (void, when)

309import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)

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

311import Control.Monad.Trans.Reader (ReaderT)

@@ -314,26 +314,21 @@ import Control.Monad.Trans.Class (MonadTrans(lift))

314import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)

315import Data.Bits (shiftR, shiftL, (.|.), (.&.))

316import Data.Functor.Identity (Identity(..))

317−import Data.Int (Int64)

318import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)

317import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)

319−import Data.Maybe (fromJust, isJust, isNothing)

318+import Data.Maybe (fromJust, isJust)

320import Data.Word (Word32)

319import Data.Word (Word32)

321import Foreign.Storable (Storable(..))

320import Foreign.Storable (Storable(..))

322import GHC.Exts (SpecConstrAnnotation(..))

321import GHC.Exts (SpecConstrAnnotation(..))

323import GHC.Types (SPEC(..))

322import GHC.Types (SPEC(..))

324−import System.Mem (performMajorGC)

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

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

326−import Streamly.Internal.Data.IORef.Prim (Prim)

327import Streamly.Internal.Data.Time.Units

324import Streamly.Internal.Data.Time.Units

328− (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)

325+ (TimeUnit64, toRelTime64, diffAbsTime64)

329import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

326import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

330import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))

327import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))

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

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

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

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

333import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))

330import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))

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

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

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

336− (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)

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

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

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

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

339

334

@@ -357,6 +352,7 @@ import Prelude hiding

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

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

358 , reverse, iterate, splitAt)

353 , reverse, iterate, splitAt)

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

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

355+import Streamly.Internal.Data.Stream.StreamD.Generate

360import Streamly.Internal.Data.SVar

356import Streamly.Internal.Data.SVar

361

357

362------------------------------------------------------------------------------

358------------------------------------------------------------------------------

@@ -421,89 +417,6 @@ uncons (UnStream step state) = go state

421 Skip s -> go s

417 Skip s -> go s

422 Stop -> return Nothing

418 Stop -> return Nothing

423

419

424−------------------------------------------------------------------------------

425−-- Generation by unfold

426−------------------------------------------------------------------------------

427−

428−{-# INLINE_NORMAL unfoldrM #-}

429−unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a

430−unfoldrM next state = Stream step state

431− where

432− {-# INLINE_LATE step #-}

433− step _ st = do

434− r <- next st

435− return $ case r of

436− Just (x, s) -> Yield x s

437− Nothing -> Stop

438−

439−{-# INLINE_LATE unfoldr #-}

440−unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a

441−unfoldr f = unfoldrM (return . f)

442−

443−data UnfoldState s = UnfoldNothing | UnfoldJust s

444−

445−-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.

446−--

447−{-# INLINE_NORMAL unfold #-}

448−unfold :: Monad m => Unfold m a b -> a -> Stream m b

449−unfold (Unfold ustep inject) seed = Stream step UnfoldNothing

450− where

451− {-# INLINE_LATE step #-}

452− step _ UnfoldNothing = inject seed >>= return . Skip . UnfoldJust

453− step _ (UnfoldJust st) = do

454− r <- ustep st

455− return $ case r of

456− Yield x s -> Yield x (UnfoldJust s)

457− Skip s -> Skip (UnfoldJust s)

458− Stop -> Stop

459−

460−------------------------------------------------------------------------------

461−-- Specialized Generation

462−------------------------------------------------------------------------------

463−

464−{-# INLINE_NORMAL repeatM #-}

465−repeatM :: Monad m => m a -> Stream m a

466−repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()

467−

468−{-# INLINE_NORMAL repeat #-}

469−repeat :: Monad m => a -> Stream m a

470−repeat x = Stream (\_ _ -> return $ Yield x ()) ()

471−

472−{-# INLINE_NORMAL iterateM #-}

473−iterateM :: Monad m => (a -> m a) -> m a -> Stream m a

474−iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))

475−

476−{-# INLINE_NORMAL iterate #-}

477−iterate :: Monad m => (a -> a) -> a -> Stream m a

478−iterate step st = iterateM (return . step) (return st)

479−

480−{-# INLINE_NORMAL replicateM #-}

481−replicateM :: forall m a. Monad m => Int -> m a -> Stream m a

482−replicateM n p = Stream step n

483− where

484− {-# INLINE_LATE step #-}

485− step _ (i :: Int)

486− | i <= 0 = return Stop

487− | otherwise = do

488− x <- p

489− return $ Yield x (i - 1)

490−

491−{-# INLINE_NORMAL replicate #-}

492−replicate :: Monad m => Int -> a -> Stream m a

493−replicate n x = replicateM n (return x)

494−

495−-- This would not work properly for floats, therefore we put an Integral

496−-- constraint.

497−-- | Can be used to enumerate unbounded integrals. This does not check for

498−-- overflow or underflow for bounded integrals.

499−{-# INLINE_NORMAL enumerateFromStepIntegral #-}

500−enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a

501−enumerateFromStepIntegral from stride =

502− from `seq` stride `seq` Stream step from

503− where

504− {-# INLINE_LATE step #-}

505− step _ !x = return $ Yield x $! (x + stride)

506−

507-- We are assuming that "to" is constrained by the type to be within

420-- We are assuming that "to" is constrained by the type to be within

508-- max/min bounds.

421-- max/min bounds.

509{-# INLINE enumerateFromToIntegral #-}

422{-# INLINE enumerateFromToIntegral #-}

@@ -515,100 +428,6 @@ enumerateFromToIntegral from to =

515enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a

428enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a

516enumerateFromIntegral from = enumerateFromToIntegral from maxBound

429enumerateFromIntegral from = enumerateFromToIntegral from maxBound

517

430

518−data EnumState a = EnumInit | EnumYield a a a | EnumStop

519−

520−{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}

521−enumerateFromThenToIntegralUp

522− :: (Monad m, Integral a)

523− => a -> a -> a -> Stream m a

524−enumerateFromThenToIntegralUp from next to = Stream step EnumInit

525− where

526− {-# INLINE_LATE step #-}

527− step _ EnumInit =

528− return $

529− if to < next

530− then if to < from

531− then Stop

532− else Yield from EnumStop

533− else -- from <= next <= to

534− let stride = next - from

535− in Skip $ EnumYield from stride (to - stride)

536−

537− step _ (EnumYield x stride toMinus) =

538− return $

539− if x > toMinus

540− then Yield x EnumStop

541− else Yield x $ EnumYield (x + stride) stride toMinus

542−

543− step _ EnumStop = return Stop

544−

545−{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}

546−enumerateFromThenToIntegralDn

547− :: (Monad m, Integral a)

548− => a -> a -> a -> Stream m a

549−enumerateFromThenToIntegralDn from next to = Stream step EnumInit

550− where

551− {-# INLINE_LATE step #-}

552− step _ EnumInit =

553− return $ if to > next

554− then if to > from

555− then Stop

556− else Yield from EnumStop

557− else -- from >= next >= to

558− let stride = next - from

559− in Skip $ EnumYield from stride (to - stride)

560−

561− step _ (EnumYield x stride toMinus) =

562− return $

563− if x < toMinus

564− then Yield x EnumStop

565− else Yield x $ EnumYield (x + stride) stride toMinus

566−

567− step _ EnumStop = return Stop

568−

569−{-# INLINE_NORMAL enumerateFromThenToIntegral #-}

570−enumerateFromThenToIntegral

571− :: (Monad m, Integral a)

572− => a -> a -> a -> Stream m a

573−enumerateFromThenToIntegral from next to

574− | next >= from = enumerateFromThenToIntegralUp from next to

575− | otherwise = enumerateFromThenToIntegralDn from next to

576−

577−{-# INLINE_NORMAL enumerateFromThenIntegral #-}

578−enumerateFromThenIntegral

579− :: (Monad m, Integral a, Bounded a)

580− => a -> a -> Stream m a

581−enumerateFromThenIntegral from next =

582− if next > from

583− then enumerateFromThenToIntegralUp from next maxBound

584− else enumerateFromThenToIntegralDn from next minBound

585−

586−-- For floating point numbers if the increment is less than the precision then

587−-- it just gets lost. Therefore we cannot always increment it correctly by just

588−-- repeated addition.

589−-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15

590−-- 9007199254740992 + 2 :: Double => 9.007199254740994e15

591−

592−-- Instead we accumulate the increment counter and compute the increment

593−-- every time before adding it to the starting number.

594−--

595−-- This works for Integrals as well as floating point numbers, but

596−-- enumerateFromStepIntegral is faster for integrals.

597−{-# INLINE_NORMAL enumerateFromStepNum #-}

598−enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a

599−enumerateFromStepNum from stride = Stream step 0

600− where

601− {-# INLINE_LATE step #-}

602− step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)

603−

604−{-# INLINE_NORMAL numFrom #-}

605−numFrom :: (Monad m, Num a) => a -> Stream m a

606−numFrom from = enumerateFromStepNum from 1

607−

608−{-# INLINE_NORMAL numFromThen #-}

609−numFromThen :: (Monad m, Num a) => a -> a -> Stream m a

610−numFromThen from next = enumerateFromStepNum from (next - from)

611−

612-- We cannot write a general function for Num. The only way to write code

431-- We cannot write a general function for Num. The only way to write code

613-- portable between the two is to use a 'Real' constraint and convert between

432-- portable between the two is to use a 'Real' constraint and convert between

614-- Fractional and Integral using fromRational which is horribly slow.

433-- Fractional and Integral using fromRational which is horribly slow.

@@ -630,181 +449,6 @@ enumerateFromThenToFractional from next to =

630 predicate | next >= from = (<= to + mid)

449 predicate | next >= from = (<= to + mid)

631 | otherwise = (>= to + mid)

450 | otherwise = (>= to + mid)

632

451

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

634−-- Generation by Conversion

635−-------------------------------------------------------------------------------

636−

637−{-# INLINE_NORMAL fromIndicesM #-}

638−fromIndicesM :: Monad m => (Int -> m a) -> Stream m a

639−fromIndicesM gen = Stream step 0

640− where

641− {-# INLINE_LATE step #-}

642− step _ i = do

643− x <- gen i

644− return $ Yield x (i + 1)

645−

646−{-# INLINE fromIndices #-}

647−fromIndices :: Monad m => (Int -> a) -> Stream m a

648−fromIndices gen = fromIndicesM (return . gen)

649−

650−{-# INLINE_NORMAL generateM #-}

651−generateM :: Monad m => Int -> (Int -> m a) -> Stream m a

652−generateM n gen = n `seq` Stream step 0

653− where

654− {-# INLINE_LATE step #-}

655− step _ i | i < n = do

656− x <- gen i

657− return $ Yield x (i + 1)

658− | otherwise = return Stop

659−

660−{-# INLINE generate #-}

661−generate :: Monad m => Int -> (Int -> a) -> Stream m a

662−generate n gen = generateM n (return . gen)

663−

664−-- XXX we need the MonadAsync constraint because of a rewrite rule.

665−-- | Convert a list of monadic actions to a 'Stream'

666−{-# INLINE_LATE fromListM #-}

667−fromListM :: MonadAsync m => [m a] -> Stream m a

668−fromListM = Stream step

669− where

670− {-# INLINE_LATE step #-}

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

672− step _ [] = return Stop

673−

674−{-# INLINE toStreamD #-}

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

676−toStreamD = fromStreamK . K.toStream

677−

678−{-# INLINE_NORMAL fromPrimIORef #-}

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

680−fromPrimIORef var = Stream step ()

681− where

682− {-# INLINE_LATE step #-}

683− step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()

684−

685−-------------------------------------------------------------------------------

686−-- Generation from SVar

687−-------------------------------------------------------------------------------

688−

689−data FromSVarState t m a =

690− FromSVarInit

691− | FromSVarRead (SVar t m a)

692− | FromSVarLoop (SVar t m a) [ChildEvent a]

693− | FromSVarDone (SVar t m a)

694−

695−{-# INLINE_NORMAL fromSVar #-}

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

697−fromSVar svar = Stream step FromSVarInit

698− where

699−

700− {-# INLINE_LATE step #-}

701− step _ FromSVarInit = do

702− ref <- liftIO $ newIORef ()

703− _ <- liftIO $ mkWeakIORef ref hook

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

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

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

707− return $ Skip (FromSVarRead sv)

708−

709− where

710−

711− {-# NOINLINE hook #-}

712− hook = do

713− when (svarInspectMode svar) $ do

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

715− when (isNothing r) $

716− printSVar svar "SVar Garbage Collected"

717− cleanupSVar svar

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

719− -- them to be cleaned up quickly.

720− when (svarInspectMode svar) performMajorGC

721−

722− step _ (FromSVarRead sv) = do

723− list <- readOutputQ sv

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

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

726− -- streams.

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

728−

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

730− done <- postProcess sv

731− return $ Skip $ if done

732− then (FromSVarDone sv)

733− else (FromSVarRead sv)

734−

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

736− case ev of

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

738− ChildStop tid e -> do

739− accountThread sv tid

740− case e of

741− Nothing -> do

742− stop <- shouldStop tid

743− if stop

744− then do

745− liftIO (cleanupSVar sv)

746− return $ Skip (FromSVarDone sv)

747− else return $ Skip (FromSVarLoop sv es)

748− Just ex ->

749− case fromException ex of

750− Just ThreadAbort ->

751− return $ Skip (FromSVarLoop sv es)

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

753− where

754−

755− shouldStop tid =

756− case svarStopStyle sv of

757− StopNone -> return False

758− StopAny -> return True

759− StopBy -> do

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

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

762−

763− step _ (FromSVarDone sv) = do

764− when (svarInspectMode sv) $ do

765− t <- liftIO $ getTime Monotonic

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

767− liftIO $ printSVar sv "SVar Done"

768− return Stop

769−

770−-------------------------------------------------------------------------------

771−-- Process events received by a fold consumer from a stream producer

772−-------------------------------------------------------------------------------

773−

774−{-# INLINE_NORMAL fromProducer #-}

775−fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a

776−fromProducer svar = Stream step (FromSVarRead svar)

777− where

778−

779− {-# INLINE_LATE step #-}

780− step _ (FromSVarRead sv) = do

781− list <- readOutputQ sv

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

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

784− -- streams.

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

786−

787− step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv

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

789− case ev of

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

791− ChildStop tid e -> do

792− accountThread sv tid

793− case e of

794− Nothing -> do

795− sendStopToProducer sv

796− return $ Skip (FromSVarDone sv)

797− Just _ -> error "Bug: fromProducer: received exception"

798−

799− step _ (FromSVarDone sv) = do

800− when (svarInspectMode sv) $ do

801− t <- liftIO $ getTime Monotonic

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

803− liftIO $ printSVar sv "SVar Done"

804− return Stop

805−

806− step _ FromSVarInit = undefined

807−

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

452-------------------------------------------------------------------------------

809-- Hoisting the inner monad

453-- Hoisting the inner monad

810-------------------------------------------------------------------------------

454-------------------------------------------------------------------------------

@@ -4896,56 +4540,3 @@ dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)

4896 Yield x s -> Yield x (DropByTimeYield s)

4540 Yield x s -> Yield x (DropByTimeYield s)

4897 Skip s -> Skip (DropByTimeYield s)

4541 Skip s -> Skip (DropByTimeYield s)

4898 Stop -> Stop

4542 Stop -> Stop

4899−

4900−{-# INLINE updateTimeVar #-}

4901−updateTimeVar :: Prim.IORef Int64 -> IO ()

4902−updateTimeVar timeVar = do

4903− MicroSecond64 t <- fromAbsTime <$> getTime Monotonic

4904− Prim.modifyIORef' timeVar (const t)

4905−

4906−{-# INLINE updateWithDelay #-}

4907−updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()

4908−updateWithDelay precision timeVar = do

4909− threadDelay (delayTime precision)

4910− updateTimeVar timeVar

4911−

4912− where

4913−

4914− -- Keep the minimum at least a millisecond to avoid high CPU usage

4915− {-# INLINE delayTime #-}

4916− delayTime g

4917− | g' >= fromIntegral (maxBound :: Int) = maxBound

4918− | g' < 1000 = 1000

4919− | otherwise = round g'

4920−

4921− where

4922−

4923− g' = g * 10 ^ (6 :: Int)

4924−

4925−-- XXX we should move this to stream generation section of this file. Also, the

4926−-- take/drop combinators above should be moved to filtering section.

4927−

4928−{-# INLINE_NORMAL times #-}

4929−times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)

4930−times g = Stream step Nothing

4931−

4932− where

4933−

4934− {-# INLINE_LATE step #-}

4935− step _ Nothing = do

4936− -- XXX note that this is safe only on a 64-bit machine. On a 32-bit

4937− -- machine a 64-bit 'Var' cannot be read consistently without a lock

4938− -- while another thread is writing to it.

4939− timeVar <- liftIO $ Prim.newIORef (0 :: Int64)

4940− liftIO $ updateTimeVar timeVar

4941− tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)

4942− a <- liftIO $ Prim.readIORef timeVar

4943− return $ Skip $ Just (timeVar, tid, a)

4944−

4945− step _ s@(Just (timeVar, _, t0)) = do

4946− a <- liftIO $ Prim.readIORef timeVar

4947− -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for

4948− -- efficiency. or maybe we can use a representation using Double for

4949− -- floating precision time

4950− return $ Yield (toAbsTime (MicroSecond64 t0),

4951− (toRelTime64 (MicroSecond64 (a - t0)))) s

src/Streamly/Internal/Data/Stream/StreamD/Generate.hsnew file+494

@@ -0,0 +1,494 @@

1+#include "inline.hs"

3+-- |

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

6+-- License : BSD-3-Clause

7+-- Maintainer : streamly@composewell.com

8+-- Stability : experimental

9+-- Portability : GHC

10+

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

12+ (

13+ -- * Generation

14+ -- ** Unfolds

15+ unfoldr

16+ , unfoldrM

17+ , unfold

18+

19+ -- ** Specialized Generation

20+ -- | Generate a monadic stream from a seed.

21+ , repeat

22+ , repeatM

23+ , replicate

24+ , replicateM

25+ , fromIndices

26+ , fromIndicesM

27+ , generate

28+ , generateM

29+ , iterate

30+ , iterateM

31+

32+ -- ** Enumerations

33+ , enumerateFromStepIntegral

34+ , enumerateFromThenIntegral

35+ , enumerateFromThenToIntegral

36+

37+ , enumerateFromStepNum

38+ , numFrom

39+ , numFromThen

40+

41+ -- ** Time

42+ , times

43+

44+ -- ** Conversions

45+ -- | Transform an input structure into a stream.

46+ -- | Direct style stream does not support @fromFoldable@.

47+ , yield

48+ , yieldM

49+ , fromList

50+ , fromListM

51+ , fromStreamK

52+ , fromStreamD

53+ , fromPrimIORef

54+ , fromProducer

55+ , fromSVar

56+ , toStreamD

57+ )

58+where

59+

60+import Control.Concurrent (threadDelay)

61+import Control.Exception

62+ (fromException)

63+import Control.Monad (when, forever)

64+import Control.Monad.Catch (throwM)

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

66+import Data.Int (Int64)

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

68+import Data.Maybe (isNothing)

69+import System.Mem (performMajorGC)

70+import Streamly.Internal.Data.IORef.Prim (Prim)

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

72+ (toRelTime64, RelTime64)

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

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

75+ (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)

76+import Streamly.Internal.Data.Unfold.Types (Unfold(..))

77+

78+import qualified Streamly.Internal.Data.IORef.Prim as Prim

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

80+import qualified Prelude

81+

82+import Prelude hiding

83+ ( map, mapM, mapM_, repeat, foldr, last, take, filter

84+ , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem

85+ , notElem, null, head, tail, zipWith, lookup, foldr1, sequence

86+ , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat

87+ , reverse, iterate, splitAt)

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

89+import Streamly.Internal.Data.SVar

90+

91+

92+------------------------------------------------------------------------------

93+-- Generation by unfold

94+------------------------------------------------------------------------------

95+

96+{-# INLINE_NORMAL unfoldrM #-}

97+unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a

98+unfoldrM next state = Stream step state

99+ where

100+ {-# INLINE_LATE step #-}

101+ step _ st = do

102+ r <- next st

103+ return $ case r of

104+ Just (x, s) -> Yield x s

105+ Nothing -> Stop

106+

107+{-# INLINE_LATE unfoldr #-}

108+unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a

109+unfoldr f = unfoldrM (return . f)

110+

111+data UnfoldState s = UnfoldNothing | UnfoldJust s

112+

113+-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.

114+--

115+{-# INLINE_NORMAL unfold #-}

116+unfold :: Monad m => Unfold m a b -> a -> Stream m b

117+unfold (Unfold ustep inject) seed = Stream step UnfoldNothing

118+ where

119+ {-# INLINE_LATE step #-}

120+ step _ UnfoldNothing = inject seed >>= return . Skip . UnfoldJust

121+ step _ (UnfoldJust st) = do

122+ r <- ustep st

123+ return $ case r of

124+ Yield x s -> Yield x (UnfoldJust s)

125+ Skip s -> Skip (UnfoldJust s)

126+ Stop -> Stop

127+

128+------------------------------------------------------------------------------

129+-- Specialized Generation

130+------------------------------------------------------------------------------

131+

132+{-# INLINE_NORMAL repeatM #-}

133+repeatM :: Monad m => m a -> Stream m a

134+repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()

135+

136+{-# INLINE_NORMAL repeat #-}

137+repeat :: Monad m => a -> Stream m a

138+repeat x = Stream (\_ _ -> return $ Yield x ()) ()

139+

140+{-# INLINE_NORMAL iterateM #-}

141+iterateM :: Monad m => (a -> m a) -> m a -> Stream m a

142+iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))

143+

144+{-# INLINE_NORMAL iterate #-}

145+iterate :: Monad m => (a -> a) -> a -> Stream m a

146+iterate step st = iterateM (return . step) (return st)

147+

148+{-# INLINE_NORMAL replicateM #-}

149+replicateM :: forall m a. Monad m => Int -> m a -> Stream m a

150+replicateM n p = Stream step n

151+ where

152+ {-# INLINE_LATE step #-}

153+ step _ (i :: Int)

154+ | i <= 0 = return Stop

155+ | otherwise = do

156+ x <- p

157+ return $ Yield x (i - 1)

158+

159+{-# INLINE_NORMAL replicate #-}

160+replicate :: Monad m => Int -> a -> Stream m a

161+replicate n x = replicateM n (return x)

162+

163+-- This would not work properly for floats, therefore we put an Integral

164+-- constraint.

165+-- | Can be used to enumerate unbounded integrals. This does not check for

166+-- overflow or underflow for bounded integrals.

167+{-# INLINE_NORMAL enumerateFromStepIntegral #-}

168+enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a

169+enumerateFromStepIntegral from stride =

170+ from `seq` stride `seq` Stream step from

171+ where

172+ {-# INLINE_LATE step #-}

173+ step _ !x = return $ Yield x $! (x + stride)

174+

175+data EnumState a = EnumInit | EnumYield a a a | EnumStop

176+

177+{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}

178+enumerateFromThenToIntegralUp

179+ :: (Monad m, Integral a)

180+ => a -> a -> a -> Stream m a

181+enumerateFromThenToIntegralUp from next to = Stream step EnumInit

182+ where

183+ {-# INLINE_LATE step #-}

184+ step _ EnumInit =

185+ return $

186+ if to < next

187+ then if to < from

188+ then Stop

189+ else Yield from EnumStop

190+ else -- from <= next <= to

191+ let stride = next - from

192+ in Skip $ EnumYield from stride (to - stride)

193+

194+ step _ (EnumYield x stride toMinus) =

195+ return $

196+ if x > toMinus

197+ then Yield x EnumStop

198+ else Yield x $ EnumYield (x + stride) stride toMinus

199+

200+ step _ EnumStop = return Stop

201+

202+{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}

203+enumerateFromThenToIntegralDn

204+ :: (Monad m, Integral a)

205+ => a -> a -> a -> Stream m a

206+enumerateFromThenToIntegralDn from next to = Stream step EnumInit

207+ where

208+ {-# INLINE_LATE step #-}

209+ step _ EnumInit =

210+ return $ if to > next

211+ then if to > from

212+ then Stop

213+ else Yield from EnumStop

214+ else -- from >= next >= to

215+ let stride = next - from

216+ in Skip $ EnumYield from stride (to - stride)

217+

218+ step _ (EnumYield x stride toMinus) =

219+ return $

220+ if x < toMinus

221+ then Yield x EnumStop

222+ else Yield x $ EnumYield (x + stride) stride toMinus

223+

224+ step _ EnumStop = return Stop

225+

226+{-# INLINE_NORMAL enumerateFromThenToIntegral #-}

227+enumerateFromThenToIntegral

228+ :: (Monad m, Integral a)

229+ => a -> a -> a -> Stream m a

230+enumerateFromThenToIntegral from next to

231+ | next >= from = enumerateFromThenToIntegralUp from next to

232+ | otherwise = enumerateFromThenToIntegralDn from next to

233+

234+{-# INLINE_NORMAL enumerateFromThenIntegral #-}

235+enumerateFromThenIntegral

236+ :: (Monad m, Integral a, Bounded a)

237+ => a -> a -> Stream m a

238+enumerateFromThenIntegral from next =

239+ if next > from

240+ then enumerateFromThenToIntegralUp from next maxBound

241+ else enumerateFromThenToIntegralDn from next minBound

242+

243+-- For floating point numbers if the increment is less than the precision then

244+-- it just gets lost. Therefore we cannot always increment it correctly by just

245+-- repeated addition.

246+-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15

247+-- 9007199254740992 + 2 :: Double => 9.007199254740994e15

248+

249+-- Instead we accumulate the increment counter and compute the increment

250+-- every time before adding it to the starting number.

251+--

252+-- This works for Integrals as well as floating point numbers, but

253+-- enumerateFromStepIntegral is faster for integrals.

254+{-# INLINE_NORMAL enumerateFromStepNum #-}

255+enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a

256+enumerateFromStepNum from stride = Stream step 0

257+ where

258+ {-# INLINE_LATE step #-}

259+ step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)

260+

261+{-# INLINE_NORMAL numFrom #-}

262+numFrom :: (Monad m, Num a) => a -> Stream m a

263+numFrom from = enumerateFromStepNum from 1

264+

265+{-# INLINE_NORMAL numFromThen #-}

266+numFromThen :: (Monad m, Num a) => a -> a -> Stream m a

267+numFromThen from next = enumerateFromStepNum from (next - from)

268+

269+{-# INLINE updateTimeVar #-}

270+updateTimeVar :: Prim.IORef Int64 -> IO ()

271+updateTimeVar timeVar = do

272+ MicroSecond64 t <- fromAbsTime <$> getTime Monotonic

273+ Prim.modifyIORef' timeVar (const t)

274+

275+{-# INLINE updateWithDelay #-}

276+updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()

277+updateWithDelay precision timeVar = do

278+ threadDelay (delayTime precision)

279+ updateTimeVar timeVar

280+

281+ where

282+

283+ -- Keep the minimum at least a millisecond to avoid high CPU usage

284+ {-# INLINE delayTime #-}

285+ delayTime g

286+ | g' >= fromIntegral (maxBound :: Int) = maxBound

287+ | g' < 1000 = 1000

288+ | otherwise = round g'

289+

290+ where

291+

292+ g' = g * 10 ^ (6 :: Int)

293+

294+-- The take/drop combinators above should be moved to filtering section.

295+

296+{-# INLINE_NORMAL times #-}

297+times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)

298+times g = Stream step Nothing

299+

300+ where

301+

302+ {-# INLINE_LATE step #-}

303+ step _ Nothing = do

304+ -- XXX note that this is safe only on a 64-bit machine. On a 32-bit

305+ -- machine a 64-bit 'Var' cannot be read consistently without a lock

306+ -- while another thread is writing to it.

307+ timeVar <- liftIO $ Prim.newIORef (0 :: Int64)

308+ liftIO $ updateTimeVar timeVar

309+ tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)

310+ a <- liftIO $ Prim.readIORef timeVar

311+ return $ Skip $ Just (timeVar, tid, a)

312+

313+ step _ s@(Just (timeVar, _, t0)) = do

314+ a <- liftIO $ Prim.readIORef timeVar

315+ -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for

316+ -- efficiency. or maybe we can use a representation using Double for

317+ -- floating precision time

318+ return $ Yield (toAbsTime (MicroSecond64 t0),

319+ (toRelTime64 (MicroSecond64 (a - t0)))) s

320+

321+-------------------------------------------------------------------------------

322+-- Generation by Conversion

323+-------------------------------------------------------------------------------

324+

325+{-# INLINE_NORMAL fromIndicesM #-}

326+fromIndicesM :: Monad m => (Int -> m a) -> Stream m a

327+fromIndicesM gen = Stream step 0

328+ where

329+ {-# INLINE_LATE step #-}

330+ step _ i = do

331+ x <- gen i

332+ return $ Yield x (i + 1)

333+

334+{-# INLINE fromIndices #-}

335+fromIndices :: Monad m => (Int -> a) -> Stream m a

336+fromIndices gen = fromIndicesM (return . gen)

337+

338+{-# INLINE_NORMAL generateM #-}

339+generateM :: Monad m => Int -> (Int -> m a) -> Stream m a

340+generateM n gen = n `seq` Stream step 0

341+ where

342+ {-# INLINE_LATE step #-}

343+ step _ i | i < n = do

344+ x <- gen i

345+ return $ Yield x (i + 1)

346+ | otherwise = return Stop

347+

348+{-# INLINE generate #-}

349+generate :: Monad m => Int -> (Int -> a) -> Stream m a

350+generate n gen = generateM n (return . gen)

351+

352+-- XXX we need the MonadAsync constraint because of a rewrite rule.

353+-- | Convert a list of monadic actions to a 'Stream'

354+{-# INLINE_LATE fromListM #-}

355+fromListM :: MonadAsync m => [m a] -> Stream m a

356+fromListM = Stream step

357+ where

358+ {-# INLINE_LATE step #-}

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

360+ step _ [] = return Stop

361+

362+{-# INLINE toStreamD #-}

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

364+toStreamD = fromStreamK . K.toStream

365+

366+{-# INLINE_NORMAL fromPrimIORef #-}

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

368+fromPrimIORef var = Stream step ()

369+ where

370+ {-# INLINE_LATE step #-}

371+ step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()

372+

373+-------------------------------------------------------------------------------

374+-- Generation from SVar

375+-------------------------------------------------------------------------------

376+

377+data FromSVarState t m a =

378+ FromSVarInit

379+ | FromSVarRead (SVar t m a)

380+ | FromSVarLoop (SVar t m a) [ChildEvent a]

381+ | FromSVarDone (SVar t m a)

382+

383+{-# INLINE_NORMAL fromSVar #-}

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

385+fromSVar svar = Stream step FromSVarInit

386+ where

387+

388+ {-# INLINE_LATE step #-}

389+ step _ FromSVarInit = do

390+ ref <- liftIO $ newIORef ()

391+ _ <- liftIO $ mkWeakIORef ref hook

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

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

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

395+ return $ Skip (FromSVarRead sv)

396+

397+ where

398+

399+ {-# NOINLINE hook #-}

400+ hook = do

401+ when (svarInspectMode svar) $ do

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

403+ when (isNothing r) $

404+ printSVar svar "SVar Garbage Collected"

405+ cleanupSVar svar

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

407+ -- them to be cleaned up quickly.

408+ when (svarInspectMode svar) performMajorGC

409+

410+ step _ (FromSVarRead sv) = do

411+ list <- readOutputQ sv

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

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

414+ -- streams.

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

416+

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

418+ done <- postProcess sv

419+ return $ Skip $ if done

420+ then (FromSVarDone sv)

421+ else (FromSVarRead sv)

422+

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

424+ case ev of

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

426+ ChildStop tid e -> do

427+ accountThread sv tid

428+ case e of

429+ Nothing -> do

430+ stop <- shouldStop tid

431+ if stop

432+ then do

433+ liftIO (cleanupSVar sv)

434+ return $ Skip (FromSVarDone sv)

435+ else return $ Skip (FromSVarLoop sv es)

436+ Just ex ->

437+ case fromException ex of

438+ Just ThreadAbort ->

439+ return $ Skip (FromSVarLoop sv es)

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

441+ where

442+

443+ shouldStop tid =

444+ case svarStopStyle sv of

445+ StopNone -> return False

446+ StopAny -> return True

447+ StopBy -> do

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

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

450+

451+ step _ (FromSVarDone sv) = do

452+ when (svarInspectMode sv) $ do

453+ t <- liftIO $ getTime Monotonic

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

455+ liftIO $ printSVar sv "SVar Done"

456+ return Stop

457+

458+-------------------------------------------------------------------------------

459+-- Process events received by a fold consumer from a stream producer

460+-------------------------------------------------------------------------------

461+

462+{-# INLINE_NORMAL fromProducer #-}

463+fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a

464+fromProducer svar = Stream step (FromSVarRead svar)

465+ where

466+

467+ {-# INLINE_LATE step #-}

468+ step _ (FromSVarRead sv) = do

469+ list <- readOutputQ sv

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

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

472+ -- streams.

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

474+

475+ step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv

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

477+ case ev of

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

479+ ChildStop tid e -> do

480+ accountThread sv tid

481+ case e of

482+ Nothing -> do

483+ sendStopToProducer sv

484+ return $ Skip (FromSVarDone sv)

485+ Just _ -> error "Bug: fromProducer: received exception"

486+

487+ step _ (FromSVarDone sv) = do

488+ when (svarInspectMode sv) $ do

489+ t <- liftIO $ getTime Monotonic

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

491+ liftIO $ printSVar sv "SVar Done"

492+ return Stop

493+

494+ step _ FromSVarInit = undefined

src/Streamly/Internal/Data/Stream/StreamD/Type.hs+2

@@ -567,6 +567,8 @@ cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2

567 Skip s2' -> cmp_null s2'

568 Stop -> return EQ

569

570+-- XXX The take/drop combinators above should be moved to filtering section.

571+

570{-# INLINE_NORMAL take #-}

572{-# INLINE_NORMAL take #-}

571take :: Monad m => Int -> Stream m a -> Stream m a

573take :: Monad m => Int -> Stream m a -> Stream m a

572take n (Stream step state) = n `seq` Stream step' (state, 0)

574take n (Stream step state) = n `seq` Stream step' (state, 0)

streamly.cabal+1

@@ -452,6 +452,7 @@ library

452 , Streamly.Internal.Data.Stream.StreamK

453 , Streamly.Internal.Data.Stream.StreamD.Type

454 , Streamly.Internal.Data.Stream.StreamD

455+ , Streamly.Internal.Data.Stream.StreamD.Generate

455 , Streamly.Internal.Data.Stream.StreamDK.Type

456 , Streamly.Internal.Data.Stream.StreamDK.Type

456 , Streamly.Internal.Data.Stream.StreamDK

457 , Streamly.Internal.Data.Stream.StreamDK

457 , Streamly.Internal.Data.Stream.Enumeration

458 , Streamly.Internal.Data.Stream.Enumeration