git.iko.soy/streamly
on master help sign in
Add drop, dropWhile & dropWhileM to Internal/Data/Unfold
240a082adithyaov
src/Streamly/Internal/Data/Unfold.hs+59−1
compare with master
unifiedsplit
@@ -95,6 +95,9 @@ module Streamly.Internal.Data.Unfold
95 , take
95 , take
96 , filter
96 , filter
97 , filterM
97 , filterM
98+ , drop
99+ , dropWhile
100+ , dropWhileM
98
101
99 -- * Zipping
102 -- * Zipping
100 , zipWithM
103 , zipWithM
@@ -132,7 +135,8 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
132import Data.Void (Void)
135import Data.Void (Void)
133import GHC.Types (SPEC(..))
136import GHC.Types (SPEC(..))
134import Prelude
137import Prelude
135− hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith)
138+ hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith
139+ , drop, dropWhile)
136
140
137import Fusion.Plugin.Types (Fuse(..))
141import Fusion.Plugin.Types (Fuse(..))
138import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
142import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
@@ -527,6 +531,60 @@ filterM f (Unfold step1 inject1) = Unfold step inject1
527filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
531filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
528filter f = filterM (return . f)
532filter f = filterM (return . f)
529
533
534+{-# INLINE_NORMAL drop #-}
535+drop :: Monad m => Int -> Unfold m a b -> Unfold m a b
536+drop n (Unfold step inject) = Unfold step' inject'
537+ where
538+ inject' a = do
539+ b <- inject a
540+ return (b, n)
541+ {-# INLINE_LATE step' #-}
542+ step' (st, i)
543+ | i > 0 = do
544+ r <- step st
545+ return $
546+ case r of
547+ Yield _ s -> Skip (s, i - 1)
548+ Skip s -> Skip (s, i)
549+ Stop -> Stop
550+ | otherwise = do
551+ r <- step st
552+ return $
553+ case r of
554+ Yield x s -> Yield x (s, 0)
555+ Skip s -> Skip (s, 0)
556+ Stop -> Stop
557+
558+{-# INLINE_NORMAL dropWhileM #-}
559+dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
560+dropWhileM f (Unfold step inject) = Unfold step' inject'
561+ where
562+ inject' a = do
563+ b <- inject a
564+ return $ Left b
565+ {-# INLINE_LATE step' #-}
566+ step' (Left st) = do
567+ r <- step st
568+ case r of
569+ Yield x s -> do
570+ b <- f x
571+ if b
572+ then return $ Skip (Left s)
573+ else return $ Yield x (Right s)
574+ Skip s -> return $ Skip (Left s)
575+ Stop -> return Stop
576+
577+ step' (Right st) = do
578+ r <- step st
579+ case r of
580+ Yield x s -> return $ Yield x (Right s)
581+ Skip s -> return $ Skip (Right s)
582+ Stop -> return Stop
583+
584+{-# INLINE dropWhile #-}
585+dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
586+dropWhile f = dropWhileM (return . f)
587+
530-------------------------------------------------------------------------------
588-------------------------------------------------------------------------------
531-- Enumeration
589-- Enumeration
532-------------------------------------------------------------------------------
590-------------------------------------------------------------------------------