@@ -37,6 +37,7 @@ module Streamly.Internal.Data.Unfold.Type
3737 , supplySecond
3838
3939 -- * Trimming
40+ , takeWhileMWithInput
4041 , takeWhileM
4142 , takeWhile
4243
@@ -69,6 +70,7 @@ where
6970import Control.Monad ((>=>))
7071import Data.Void (Void)
7172import Fusion.Plugin.Types (Fuse(..))
73+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
7274import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
7375
7476import Prelude hiding (map, concatMap, zipWith, takeWhile)
@@ -228,12 +230,37 @@ supplySecond b = lmap (, b)
228230-- Filter input
229231------------------------------------------------------------------------------
230232
233+{-# INLINE_NORMAL takeWhileMWithInput #-}
234+takeWhileMWithInput :: Monad m =>
235+ (a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
236+takeWhileMWithInput f (Unfold step1 inject1) = Unfold step inject
237+
238+ where
239+
240+ inject a = do
241+ s <- inject1 a
242+ return $ Tuple' a s
243+
244+ {-# INLINE_LATE step #-}
245+ step (Tuple' a st) = do
246+ r <- step1 st
247+ case r of
248+ Yield x s -> do
249+ b <- f a x
250+ return $ if b then Yield x (Tuple' a s) else Stop
251+ Skip s -> return $ Skip (Tuple' a s)
252+ Stop -> return Stop
253+
231254-- | Same as 'takeWhile' but with a monadic predicate.
232255--
233256-- /Since: 0.8.0/
234257--
235258{-# INLINE_NORMAL takeWhileM #-}
236259takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
260+-- XXX Check if the compiler simplifies the following to the same as the custom
261+-- implementation below (the Tuple' should help eliminate the unused param):
262+--
263+-- takeWhileM f = takeWhileMWithInput (\_ b -> f b)
237264takeWhileM f (Unfold step1 inject1) = Unfold step inject1
238265 where
239266 {-# INLINE_LATE step #-}