on master help sign in
implement mfix for StreamK
33b00a6Harendra Kumar
src/Streamly/Internal/Data/Stream/StreamK.hs+43
@@ -180,6 +180,7 @@ module Streamly.Internal.Data.Stream.StreamK
180 -- * Utilities
180 -- * Utilities
181 , consMStream
181 , consMStream
182 , withLocal
182 , withLocal
183+ , mfix
183
184
184 -- * Deprecated
185 -- * Deprecated
185 , Streaming -- deprecated
186 , Streaming -- deprecated
@@ -190,6 +191,7 @@ where
190import Control.Monad.Trans (MonadTrans(lift))
191import Control.Monad.Trans (MonadTrans(lift))
191import Control.Monad (void, join)
192import Control.Monad (void, join)
192import Control.Monad.Reader.Class (MonadReader(..))
193import Control.Monad.Reader.Class (MonadReader(..))
194+import Data.Function (fix)
193import Prelude
195import Prelude
194 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
196 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
195 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
197 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
@@ -500,6 +502,47 @@ tail m =
500 yieldk _ r = return $ Just r
502 yieldk _ r = return $ Just r
501 in foldStream defState yieldk single stop m
503 in foldStream defState yieldk single stop m
502
504
505+{-# INLINE headPartial #-}
506+headPartial :: (IsStream t, Monad m) => t m a -> m a
507+headPartial = foldrM (\x _ -> return x) (error "head of nil")
508+
509+{-# INLINE tailPartial #-}
510+tailPartial :: IsStream t => t m a -> t m a
511+tailPartial m = mkStream $ \st yld sng stp ->
512+ let stop = error "tail of nil"
513+ single _ = stp
514+ yieldk _ r = foldStream st yld sng stp r
515+ in foldStream st yieldk single stop m
516+
517+-- | Iterate a lazy function `f` of the shape `m a -> t m a` until it gets
518+-- fully defined i.e. becomes independent of its argument action, then return
519+-- the resulting value of the function (`t m a`).
520+--
521+-- It can be used to construct a stream that uses a cyclic definition. For
522+-- example:
523+--
524+-- @
525+-- import Streamly.Internal.Prelude as S
526+-- import System.IO.Unsafe (unsafeInterleaveIO)
527+--
528+-- main = do
529+-- S.mapM_ print $ S.mfix $ \x -> do
530+-- a <- S.fromList [1,2]
531+-- b <- S.fromListM [return 3, unsafeInterleaveIO (fmap fst x)]
532+-- return (a, b)
533+-- @
534+--
535+-- Note that the function `f` must be lazy in its argument, that's why we use
536+-- 'unsafeInterleaveIO' because IO monad is strict.
537+
538+mfix :: (IsStream t, Monad m) => (m a -> t m a) -> t m a
539+mfix f = mkStream $ \st yld sng stp ->
540+ let single a = foldStream st yld sng stp $ a `cons` ys
541+ yieldk a _ = foldStream st yld sng stp $ a `cons` ys
542+ in foldStream st yieldk single stp xs
543+ where xs = fix (f . headPartial)
544+ ys = mfix (tailPartial . f)
545+
503{-# INLINE init #-}
546{-# INLINE init #-}
504init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
547init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
505init m = go1 m
548init m = go1 m
src/Streamly/Internal/Prelude.hs+3
@@ -398,6 +398,9 @@ module Streamly.Internal.Prelude
398 -- * Concurrency
398 -- * Concurrency
399 , mkParallel
399 , mkParallel
400
400
401+ -- * MonadFix
402+ , K.mfix
403+
401 -- * Diagnostics
404 -- * Diagnostics
402 , inspectMode
405 , inspectMode
403
406