git.iko.soy/streamly
on master help sign in
Add additional generation unfolds
c004bd7adithyaov
src/Streamly/Internal/Data/Unfold.hs+171−8
compare with master
unifiedsplit
@@ -76,10 +76,14 @@ module Streamly.Internal.Data.Unfold
76 , singleton
76 , singleton
77 , identity
77 , identity
78 , const
78 , const
79+ , unfoldrM
79
80
80 , fromList
81 , fromList
81 , fromListM
82 , fromListM
82
83
84+ , fromSVar
85+ , fromProducer
86+
83 -- ** Specialized Generation
87 -- ** Specialized Generation
84 -- | Generate a monadic stream from a seed.
88 -- | Generate a monadic stream from a seed.
85 , replicateM
89 , replicateM
@@ -92,6 +96,10 @@ module Streamly.Internal.Data.Unfold
92 , enumerateFromToIntegral
96 , enumerateFromToIntegral
93 , enumerateFromIntegral
97 , enumerateFromIntegral
94
98
99+ , enumerateFromStepNum
100+ , numFrom
101+ , enumerateFromToFractional
102+
95 -- * Transformations
103 -- * Transformations
96 , map
104 , map
97 , mapM
105 , mapM
@@ -136,22 +144,21 @@ module Streamly.Internal.Data.Unfold
136 )
144 )
137where
145where
138
146
139−import Control.Exception (Exception, mask_)
147+import Control.Exception (Exception, fromException, mask_)
140−import Control.Monad ((>=>))
148+import Control.Monad ((>=>), when)
149+import Control.Monad.Catch (MonadCatch)
141import Control.Monad.IO.Class (MonadIO(..))
150import Control.Monad.IO.Class (MonadIO(..))
142import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
151import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
152+import Data.Maybe (isNothing)
143import Data.Void (Void)
153import Data.Void (Void)
154+import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
144import GHC.Types (SPEC(..))
155import GHC.Types (SPEC(..))
145−import Prelude
146− hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith
147− , drop, dropWhile)
148−
149import Fusion.Plugin.Types (Fuse(..))
156import Fusion.Plugin.Types (Fuse(..))
150import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
157import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
151import Streamly.Internal.Data.Unfold.Types (Unfold(..))
158import Streamly.Internal.Data.Unfold.Types (Unfold(..))
152import Streamly.Internal.Data.Fold.Types (Fold(..))
159import Streamly.Internal.Data.Fold.Types (Fold(..))
153−import Streamly.Internal.Data.SVar (defState, MonadAsync)
160+import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
154−import Control.Monad.Catch (MonadCatch)
161+import System.Mem (performMajorGC)
155
162
156import qualified Prelude
163import qualified Prelude
157import qualified Control.Monad.Catch as MC
164import qualified Control.Monad.Catch as MC
@@ -159,6 +166,11 @@ import qualified Data.Tuple as Tuple
159import qualified Streamly.Internal.Data.Stream.StreamK as K
166import qualified Streamly.Internal.Data.Stream.StreamK as K
160import qualified Streamly.Internal.Data.Stream.StreamD as D
167import qualified Streamly.Internal.Data.Stream.StreamD as D
161
168
169+import Streamly.Internal.Data.SVar
170+import Prelude
171+ hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith
172+ , drop, dropWhile)
173+
162-------------------------------------------------------------------------------
174-------------------------------------------------------------------------------
163-- Input operations
175-- Input operations
164-------------------------------------------------------------------------------
176-------------------------------------------------------------------------------
@@ -457,6 +469,17 @@ fromListM = Unfold step inject
457 step (x:xs) = x >>= \r -> return $ Yield r xs
469 step (x:xs) = x >>= \r -> return $ Yield r xs
458 step [] = return Stop
470 step [] = return Stop
459
471
472+{-# INLINE unfoldrM #-}
473+unfoldrM :: Monad m => (a -> m (Maybe (b, a))) -> Unfold m a b
474+unfoldrM next = Unfold step return
475+ where
476+ {-# INLINE_LATE step #-}
477+ step st = do
478+ r <- next st
479+ return $ case r of
480+ Just (x, s) -> Yield x s
481+ Nothing -> Stop
482+
460------------------------------------------------------------------------------
483------------------------------------------------------------------------------
461-- Specialized Generation
484-- Specialized Generation
462------------------------------------------------------------------------------
485------------------------------------------------------------------------------
@@ -638,6 +661,146 @@ enumerateFromToIntegral to =
638enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
661enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
639enumerateFromIntegral = enumerateFromToIntegral maxBound
662enumerateFromIntegral = enumerateFromToIntegral maxBound
640
663
664+{-# INLINE enumerateFromStepNum #-}
665+enumerateFromStepNum :: (Monad m, Num a) => a -> Unfold m a a
666+enumerateFromStepNum stride = Unfold step return
667+ where
668+ {-# INLINE_LATE step #-}
669+ step !s = return $ (Yield $! s) $! (s + stride)
670+
671+{-# INLINE_NORMAL numFrom #-}
672+numFrom :: (Monad m, Num a) => Unfold m a a
673+numFrom = enumerateFromStepNum 1
674+
675+{-# INLINE_NORMAL enumerateFromToFractional #-}
676+enumerateFromToFractional :: (Monad m, Fractional a, Ord a) => a -> Unfold m a a
677+enumerateFromToFractional to =
678+ takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
679+
680+-------------------------------------------------------------------------------
681+-- Generation from SVar
682+-------------------------------------------------------------------------------
683+
684+data FromSVarState t m a =
685+ FromSVarInit (SVar t m a)
686+ | FromSVarRead (SVar t m a)
687+ | FromSVarLoop (SVar t m a) [ChildEvent a]
688+ | FromSVarDone (SVar t m a)
689+
690+{-# INLINE_NORMAL fromSVar #-}
691+fromSVar :: MonadAsync m => Unfold m (SVar t m a) a
692+fromSVar = Unfold step (return . FromSVarInit)
693+ where
694+
695+ {-# INLINE_LATE step #-}
696+ step (FromSVarInit svar) = do
697+ ref <- liftIO $ newIORef ()
698+ _ <- liftIO $ mkWeakIORef ref hook
699+ -- when this copy of svar gets garbage collected "ref" will get
700+ -- garbage collected and our GC hook will be called.
701+ let sv = svar{svarRef = Just ref}
702+ return $ Skip (FromSVarRead sv)
703+
704+ where
705+
706+ {-# NOINLINE hook #-}
707+ hook = do
708+ when (svarInspectMode svar) $ do
709+ r <- liftIO $ readIORef (svarStopTime (svarStats svar))
710+ when (isNothing r) $
711+ printSVar svar "SVar Garbage Collected"
712+ cleanupSVar svar
713+ -- If there are any SVars referenced by this SVar a GC will prompt
714+ -- them to be cleaned up quickly.
715+ when (svarInspectMode svar) performMajorGC
716+
717+ step (FromSVarRead sv) = do
718+ list <- readOutputQ sv
719+ -- Reversing the output is important to guarantee that we process the
720+ -- outputs in the same order as they were generated by the constituent
721+ -- streams.
722+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
723+
724+ step (FromSVarLoop sv []) = do
725+ done <- postProcess sv
726+ return $ Skip $ if done
727+ then FromSVarDone sv
728+ else FromSVarRead sv
729+
730+ step (FromSVarLoop sv (ev : es)) = do
731+ case ev of
732+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)
733+ ChildStop tid e -> do
734+ accountThread sv tid
735+ case e of
736+ Nothing -> do
737+ stop <- shouldStop tid
738+ if stop
739+ then do
740+ liftIO (cleanupSVar sv)
741+ return $ Skip (FromSVarDone sv)
742+ else return $ Skip (FromSVarLoop sv es)
743+ Just ex ->
744+ case fromException ex of
745+ Just ThreadAbort ->
746+ return $ Skip (FromSVarLoop sv es)
747+ Nothing -> liftIO (cleanupSVar sv) >> MC.throwM ex
748+ where
749+
750+ shouldStop tid =
751+ case svarStopStyle sv of
752+ StopNone -> return False
753+ StopAny -> return True
754+ StopBy -> do
755+ sid <- liftIO $ readIORef (svarStopBy sv)
756+ return $ tid == sid
757+
758+ step (FromSVarDone sv) = do
759+ when (svarInspectMode sv) $ do
760+ t <- liftIO $ getTime Monotonic
761+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
762+ liftIO $ printSVar sv "SVar Done"
763+ return Stop
764+
765+-------------------------------------------------------------------------------
766+-- Process events received by a fold consumer from a stream producer
767+-------------------------------------------------------------------------------
768+
769+-- XXX Why don't we do the same thing for "FromSVarInit" that we did in "fromSVar"?
770+{-# INLINE_NORMAL fromProducer #-}
771+fromProducer :: MonadAsync m => Unfold m (SVar t m a) a
772+fromProducer = Unfold step (return . FromSVarRead)
773+ where
774+
775+ {-# INLINE_LATE step #-}
776+ step (FromSVarRead sv) = do
777+ list <- readOutputQ sv
778+ -- Reversing the output is important to guarantee that we process the
779+ -- outputs in the same order as they were generated by the constituent
780+ -- streams.
781+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
782+
783+ step (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
784+ step (FromSVarLoop sv (ev : es)) = do
785+ case ev of
786+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)
787+ ChildStop tid e -> do
788+ accountThread sv tid
789+ case e of
790+ Nothing -> do
791+ sendStopToProducer sv
792+ return $ Skip (FromSVarDone sv)
793+ Just _ -> error "Bug: fromProducer: received exception"
794+
795+ step (FromSVarDone sv) = do
796+ when (svarInspectMode sv) $ do
797+ t <- liftIO $ getTime Monotonic
798+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
799+ liftIO $ printSVar sv "SVar Done"
800+ return Stop
801+
802+ step (FromSVarInit _) = undefined
803+
641-------------------------------------------------------------------------------
804-------------------------------------------------------------------------------
642-- Zipping
805-- Zipping
643-------------------------------------------------------------------------------
806-------------------------------------------------------------------------------