@@ -95,6 +95,9 @@ module Streamly.Internal.Data.Unfold
9595 , take
9696 , filter
9797 , filterM
98+ , drop
99+ , dropWhile
100+ , dropWhileM
98101
99102 -- * Zipping
100103 , zipWithM
@@ -132,7 +135,8 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
132135import Data.Void (Void)
133136import GHC.Types (SPEC(..))
134137import Prelude
135− hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith)
138+ hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith
139+ , drop, dropWhile)
136140
137141import Fusion.Plugin.Types (Fuse(..))
138142import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
@@ -527,6 +531,60 @@ filterM f (Unfold step1 inject1) = Unfold step inject1
527531filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
528532filter f = filterM (return . f)
529533
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+
530588-------------------------------------------------------------------------------
531589-- Enumeration
532590-------------------------------------------------------------------------------