@@ -84,6 +84,10 @@ module Streamly.Streams.StreamK
8484 , mapMaybe
8585 , sequence
8686
87+ -- * Filtering
88+ , filter
89+ , take
90+
8791 -- * Semigroup Style Composition
8892 , serial
8993
@@ -102,7 +106,9 @@ import Control.Monad (void)
102106import Control.Monad.Reader.Class (MonadReader(..))
103107import Control.Monad.Trans.Class (MonadTrans(lift))
104108import Data.Semigroup (Semigroup(..))
105−import Prelude hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence)
109+import Prelude
110+ hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
111+ take, filter)
106112import qualified Prelude
107113
108114import Streamly.SVar
@@ -517,6 +523,29 @@ repeat a = let x = cons a x in x
517523fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
518524fromFoldable = Prelude.foldr cons nil
519525
526+-------------------------------------------------------------------------------
527+-- Filtering
528+-------------------------------------------------------------------------------
529+
530+{-# INLINE take #-}
531+take :: IsStream t => Int -> t m a -> t m a
532+take n m = fromStream $ go n (toStream m)
533+ where
534+ go n1 m1 = Stream $ \_ stp sng yld ->
535+ let yieldk a r = yld a (go (n1 - 1) r)
536+ in if n1 <= 0 then stp else (unStream m1) Nothing stp sng yieldk
537+
538+{-# INLINE filter #-}
539+filter :: IsStream t => (a -> Bool) -> t m a -> t m a
540+filter p m = fromStream $ go (toStream m)
541+ where
542+ go m1 = Stream $ \_ stp sng yld ->
543+ let single a | p a = sng a
544+ | otherwise = stp
545+ yieldk a r | p a = yld a (go r)
546+ | otherwise = (unStream r) Nothing stp single yieldk
547+ in (unStream m1) Nothing stp single yieldk
548+
520549------------------------------------------------------------------------------
521550-- Semigroup
522551------------------------------------------------------------------------------