use take and filter from StreamD

src/Streamly/Prelude.hs

@@ -465,26 +465,15 @@ toList m = D.toList $ D.fromStreamK (toStream m)
 --
 -- @since 0.1.0
 {-# INLINE take #-}
-take :: IsStream t => Int -> t m a -> t m a
-take n m = fromStream $ go n (toStream m)
-    where
-    go n1 m1 = K.Stream $ \_ stp sng yld ->
-        let yieldk a r = yld a (go (n1 - 1) r)
-        in if n1 <= 0 then stp else (K.unStream m1) Nothing stp sng yieldk
+take :: (IsStream t, Monad m) => Int -> t m a -> t m a
+take n m = fromStream $ D.toStreamK $ D.take n (D.fromStreamK $ toStream m)
 
 -- | Include only those elements that pass a predicate.
 --
 -- @since 0.1.0
 {-# INLINE filter #-}
-filter :: IsStream t => (a -> Bool) -> t m a -> t m a
-filter p m = fromStream $ go (toStream m)
-    where
-    go m1 = K.Stream $ \_ stp sng yld ->
-        let single a  | p a       = sng a
-                      | otherwise = stp
-            yieldk a r | p a       = yld a (go r)
-                      | otherwise = (K.unStream r) Nothing stp single yieldk
-         in (K.unStream m1) Nothing stp single yieldk
+filter :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a
+filter p m = fromStream $ D.toStreamK $ D.filter p (D.fromStreamK $ toStream m)
 
 -- | End the stream as soon as the predicate fails on an element.
 --

src/Streamly/Streams/StreamD.hs

@@ -76,6 +76,11 @@ module Streamly.Streams.StreamD
     , map
     , mapM
 
+    -- * Filtering
+    , filter
+    , filterM
+    , take
+
     -- * Conversion
     , toStreamK
     , fromStreamK
@@ -83,7 +88,7 @@ module Streamly.Streams.StreamD
 where
 
 import GHC.Types ( SPEC(..) )
-import Prelude hiding (map, mapM, mapM_, repeat, foldr, last)
+import Prelude hiding (map, mapM, mapM_, repeat, foldr, last, take, filter)
 
 import Streamly.SVar (MonadAsync)
 import qualified Streamly.Streams.StreamK as K
@@ -254,6 +259,41 @@ instance Monad m => Functor (Stream m) where
     {-# INLINE fmap #-}
     fmap = map
 
+-------------------------------------------------------------------------------
+-- Filtering
+-------------------------------------------------------------------------------
+
+{-# INLINE_NORMAL take #-}
+take :: Monad m => Int -> Stream m a -> Stream m a
+take n (Stream step state) = n `seq` Stream step' (state, 0)
+  where
+    {-# INLINE_LATE step' #-}
+    step' (st, i) | i < n = do
+        r <- step st
+        return $ case r of
+            Yield x s -> Yield x (s, i + 1)
+            Stop      -> Stop
+    step' (_, _) = return Stop
+
+{-# INLINE_NORMAL filterM #-}
+filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
+filterM f (Stream step state) = Stream step' state
+  where
+    {-# INLINE_LATE step' #-}
+    step' st = do
+        r <- step st
+        case r of
+            Yield x s -> do
+                b <- f x
+                if b
+                then return $ Yield x s
+                else step' s
+            Stop -> return $ Stop
+
+{-# INLINE filter #-}
+filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
+filter f = filterM (return . f)
+
 -------------------------------------------------------------------------------
 -- Elimination
 -------------------------------------------------------------------------------

src/Streamly/Streams/StreamK.hs

@@ -84,6 +84,10 @@ module Streamly.Streams.StreamK
     , mapMaybe
     , sequence
 
+    -- * Filtering
+    , filter
+    , take
+
     -- * Semigroup Style Composition
     , serial
 
@@ -102,7 +106,9 @@ import Control.Monad (void)
 import Control.Monad.Reader.Class  (MonadReader(..))
 import Control.Monad.Trans.Class (MonadTrans(lift))
 import Data.Semigroup (Semigroup(..))
-import Prelude hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence)
+import Prelude
+       hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
+               take, filter)
 import qualified Prelude
 
 import Streamly.SVar
@@ -517,6 +523,29 @@ repeat a = let x = cons a x in x
 fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
 fromFoldable = Prelude.foldr cons nil
 
+-------------------------------------------------------------------------------
+-- Filtering
+-------------------------------------------------------------------------------
+
+{-# INLINE take #-}
+take :: IsStream t => Int -> t m a -> t m a
+take n m = fromStream $ go n (toStream m)
+    where
+    go n1 m1 = Stream $ \_ stp sng yld ->
+        let yieldk a r = yld a (go (n1 - 1) r)
+        in if n1 <= 0 then stp else (unStream m1) Nothing stp sng yieldk
+
+{-# INLINE filter #-}
+filter :: IsStream t => (a -> Bool) -> t m a -> t m a
+filter p m = fromStream $ go (toStream m)
+    where
+    go m1 = Stream $ \_ stp sng yld ->
+        let single a  | p a       = sng a
+                      | otherwise = stp
+            yieldk a r | p a       = yld a (go r)
+                       | otherwise = (unStream r) Nothing stp single yieldk
+         in (unStream m1) Nothing stp single yieldk
+
 ------------------------------------------------------------------------------
 -- Semigroup
 ------------------------------------------------------------------------------