git.iko.soy/streamly
on master help sign in
Add takeWhileMInput
78547bfHarendra Kumar
src/Streamly/Internal/Data/Unfold/Type.hs+27
compare with master
unifiedsplit
@@ -37,6 +37,7 @@ module Streamly.Internal.Data.Unfold.Type
37 , supplySecond
37 , supplySecond
38
38
39 -- * Trimming
39 -- * Trimming
40+ , takeWhileMWithInput
40 , takeWhileM
41 , takeWhileM
41 , takeWhile
42 , takeWhile
42
43
@@ -69,6 +70,7 @@ where
69import Control.Monad ((>=>))
70import Control.Monad ((>=>))
70import Data.Void (Void)
71import Data.Void (Void)
71import Fusion.Plugin.Types (Fuse(..))
72import Fusion.Plugin.Types (Fuse(..))
73+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
72import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
74import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
73
75
74import Prelude hiding (map, concatMap, zipWith, takeWhile)
76import Prelude hiding (map, concatMap, zipWith, takeWhile)
@@ -228,12 +230,37 @@ supplySecond b = lmap (, b)
228-- Filter input
230-- Filter input
229------------------------------------------------------------------------------
231------------------------------------------------------------------------------
230
232
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+
231-- | Same as 'takeWhile' but with a monadic predicate.
254-- | Same as 'takeWhile' but with a monadic predicate.
232--
255--
233-- /Since: 0.8.0/
256-- /Since: 0.8.0/
234--
257--
235{-# INLINE_NORMAL takeWhileM #-}
258{-# INLINE_NORMAL takeWhileM #-}
236takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
259takeWhileM :: 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)
237takeWhileM f (Unfold step1 inject1) = Unfold step inject1
264takeWhileM f (Unfold step1 inject1) = Unfold step inject1
238 where
265 where
239 {-# INLINE_LATE step #-}
266 {-# INLINE_LATE step #-}