on master help sign in
use take and filter from StreamD
commit70ba5db0315c99847c20522554e98d6671cea0f8 authorHarendra Kumar <harendra.kumar@gmail.com> dateMon 25 Jun 2018, 05:53 (8 years ago) parentd0a07bb
3 files changed+75−17
src/Streamly/Prelude.hs+4−15 src/Streamly/Streams/StreamD.hs+41−1 src/Streamly/Streams/StreamK.hs+30−1
src/Streamly/Prelude.hs+4−15
@@ -465,26 +465,15 @@ toList m = D.toList $ D.fromStreamK (toStream m)
465465--
466466-- @since 0.1.0
467467{-# INLINE take #-}
468−take :: IsStream t => Int -> t m a -> t m a
469−take n m = fromStream $ go n (toStream m)
470− where
471− go n1 m1 = K.Stream $ \_ stp sng yld ->
472− let yieldk a r = yld a (go (n1 - 1) r)
473− in if n1 <= 0 then stp else (K.unStream m1) Nothing stp sng yieldk
468+take :: (IsStream t, Monad m) => Int -> t m a -> t m a
469+take n m = fromStream $ D.toStreamK $ D.take n (D.fromStreamK $ toStream m)
474470
475471-- | Include only those elements that pass a predicate.
476472--
477473-- @since 0.1.0
478474{-# INLINE filter #-}
479−filter :: IsStream t => (a -> Bool) -> t m a -> t m a
480−filter p m = fromStream $ go (toStream m)
481− where
482− go m1 = K.Stream $ \_ stp sng yld ->
483− let single a | p a = sng a
484− | otherwise = stp
485− yieldk a r | p a = yld a (go r)
486− | otherwise = (K.unStream r) Nothing stp single yieldk
487− in (K.unStream m1) Nothing stp single yieldk
475+filter :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a
476+filter p m = fromStream $ D.toStreamK $ D.filter p (D.fromStreamK $ toStream m)
488477
489478-- | End the stream as soon as the predicate fails on an element.
490479--
src/Streamly/Streams/StreamD.hs+41−1
@@ -76,6 +76,11 @@ module Streamly.Streams.StreamD
7676 , map
7777 , mapM
7878
79+ -- * Filtering
80+ , filter
81+ , filterM
82+ , take
83+
7984 -- * Conversion
8085 , toStreamK
8186 , fromStreamK
@@ -83,7 +88,7 @@ module Streamly.Streams.StreamD
8388where
8489
8590import GHC.Types ( SPEC(..) )
86−import Prelude hiding (map, mapM, mapM_, repeat, foldr, last)
91+import Prelude hiding (map, mapM, mapM_, repeat, foldr, last, take, filter)
8792
8893import Streamly.SVar (MonadAsync)
8994import qualified Streamly.Streams.StreamK as K
@@ -254,6 +259,41 @@ instance Monad m => Functor (Stream m) where
254259 {-# INLINE fmap #-}
255260 fmap = map
256261
262+-------------------------------------------------------------------------------
263+-- Filtering
264+-------------------------------------------------------------------------------
265+
266+{-# INLINE_NORMAL take #-}
267+take :: Monad m => Int -> Stream m a -> Stream m a
268+take n (Stream step state) = n `seq` Stream step' (state, 0)
269+ where
270+ {-# INLINE_LATE step' #-}
271+ step' (st, i) | i < n = do
272+ r <- step st
273+ return $ case r of
274+ Yield x s -> Yield x (s, i + 1)
275+ Stop -> Stop
276+ step' (_, _) = return Stop
277+
278+{-# INLINE_NORMAL filterM #-}
279+filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
280+filterM f (Stream step state) = Stream step' state
281+ where
282+ {-# INLINE_LATE step' #-}
283+ step' st = do
284+ r <- step st
285+ case r of
286+ Yield x s -> do
287+ b <- f x
288+ if b
289+ then return $ Yield x s
290+ else step' s
291+ Stop -> return $ Stop
292+
293+{-# INLINE filter #-}
294+filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
295+filter f = filterM (return . f)
296+
257297-------------------------------------------------------------------------------
258298-- Elimination
259299-------------------------------------------------------------------------------
src/Streamly/Streams/StreamK.hs+30−1
@@ -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------------------------------------------------------------------------------