@@ -76,10 +76,14 @@ module Streamly.Internal.Data.Unfold
7676 , singleton
7777 , identity
7878 , const
79+ , unfoldrM
7980
8081 , fromList
8182 , fromListM
8283
84+ , fromSVar
85+ , fromProducer
86+
8387 -- ** Specialized Generation
8488 -- | Generate a monadic stream from a seed.
8589 , replicateM
@@ -92,6 +96,10 @@ module Streamly.Internal.Data.Unfold
9296 , enumerateFromToIntegral
9397 , enumerateFromIntegral
9498
99+ , enumerateFromStepNum
100+ , numFrom
101+ , enumerateFromToFractional
102+
95103 -- * Transformations
96104 , map
97105 , mapM
@@ -136,22 +144,21 @@ module Streamly.Internal.Data.Unfold
136144 )
137145where
138146
139−import Control.Exception (Exception, mask_)
140−import Control.Monad ((>=>))
147+import Control.Exception (Exception, fromException, mask_)
148+import Control.Monad ((>=>), when)
149+import Control.Monad.Catch (MonadCatch)
141150import Control.Monad.IO.Class (MonadIO(..))
142151import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
152+import Data.Maybe (isNothing)
143153import Data.Void (Void)
154+import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
144155import GHC.Types (SPEC(..))
145−import Prelude
146− hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith
147− , drop, dropWhile)
148−
149156import Fusion.Plugin.Types (Fuse(..))
150157import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
151158import Streamly.Internal.Data.Unfold.Types (Unfold(..))
152159import Streamly.Internal.Data.Fold.Types (Fold(..))
153−import Streamly.Internal.Data.SVar (defState, MonadAsync)
154−import Control.Monad.Catch (MonadCatch)
160+import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
161+import System.Mem (performMajorGC)
155162
156163import qualified Prelude
157164import qualified Control.Monad.Catch as MC
@@ -159,6 +166,11 @@ import qualified Data.Tuple as Tuple
159166import qualified Streamly.Internal.Data.Stream.StreamK as K
160167import qualified Streamly.Internal.Data.Stream.StreamD as D
161168
169+import Streamly.Internal.Data.SVar
170+import Prelude
171+ hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith
172+ , drop, dropWhile)
173+
162174-------------------------------------------------------------------------------
163175-- Input operations
164176-------------------------------------------------------------------------------
@@ -457,6 +469,17 @@ fromListM = Unfold step inject
457469 step (x:xs) = x >>= \r -> return $ Yield r xs
458470 step [] = return Stop
459471
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+
460483------------------------------------------------------------------------------
461484-- Specialized Generation
462485------------------------------------------------------------------------------
@@ -638,6 +661,146 @@ enumerateFromToIntegral to =
638661enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
639662enumerateFromIntegral = enumerateFromToIntegral maxBound
640663
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+
641804-------------------------------------------------------------------------------
642805-- Zipping
643806-------------------------------------------------------------------------------