@@ -180,6 +180,7 @@ module Streamly.Internal.Data.Stream.StreamK
180180 -- * Utilities
181181 , consMStream
182182 , withLocal
183+ , mfix
183184
184185 -- * Deprecated
185186 , Streaming -- deprecated
@@ -190,6 +191,7 @@ where
190191import Control.Monad.Trans (MonadTrans(lift))
191192import Control.Monad (void, join)
192193import Control.Monad.Reader.Class (MonadReader(..))
194+import Data.Function (fix)
193195import Prelude
194196 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
195197 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
@@ -500,6 +502,47 @@ tail m =
500502 yieldk _ r = return $ Just r
501503 in foldStream defState yieldk single stop m
502504
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+
503546{-# INLINE init #-}
504547init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
505548init m = go1 m