70ba5db · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

use take and filter from StreamD

70ba5dbHarendra Kumar8 years ago

src/Streamly/Prelude.hs+4−15

compare with master

unifiedsplit

@@ -465,26 +465,15 @@ toList m = D.toList $ D.fromStreamK (toStream m)

465--

466-- @since 0.1.0

467{-# INLINE take #-}

468−take :: IsStream t => Int -> t m a -> t m a

468+take :: (IsStream t, Monad m) => Int -> t m a -> t m a

469−take n m = fromStream $ go n (toStream m)

469+take n m = fromStream $ D.toStreamK $ D.take n (D.fromStreamK $ 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

474

470

475-- | Include only those elements that pass a predicate.

471-- | Include only those elements that pass a predicate.

476--

472--

477-- @since 0.1.0

473-- @since 0.1.0

478{-# INLINE filter #-}

474{-# INLINE filter #-}

479−filter :: IsStream t => (a -> Bool) -> t m a -> t m a

475+filter :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a

480−filter p m = fromStream $ go (toStream m)

476+filter p m = fromStream $ D.toStreamK $ D.filter p (D.fromStreamK $ 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

488

477

489-- | End the stream as soon as the predicate fails on an element.

478-- | End the stream as soon as the predicate fails on an element.

490--

479--

src/Streamly/Streams/StreamD.hs+41−1

@@ -76,6 +76,11 @@ module Streamly.Streams.StreamD

76 , map

77 , mapM

79+ -- * Filtering

80+ , filter

81+ , filterM

82+ , take

83+

79 -- * Conversion

84 -- * Conversion

80 , toStreamK

85 , toStreamK

81 , fromStreamK

86 , fromStreamK

@@ -83,7 +88,7 @@ module Streamly.Streams.StreamD

83where

88where

85import GHC.Types ( SPEC(..) )

90import 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)

88import Streamly.SVar (MonadAsync)

93import Streamly.SVar (MonadAsync)

89import qualified Streamly.Streams.StreamK as K

94import qualified Streamly.Streams.StreamK as K

@@ -254,6 +259,41 @@ instance Monad m => Functor (Stream m) where

254 {-# INLINE fmap #-}

259 {-# INLINE fmap #-}

255 fmap = map

260 fmap = map

256

261

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+

257-------------------------------------------------------------------------------

297-------------------------------------------------------------------------------

258-- Elimination

298-- Elimination

259-------------------------------------------------------------------------------

299-------------------------------------------------------------------------------

src/Streamly/Streams/StreamK.hs+30−1

@@ -84,6 +84,10 @@ module Streamly.Streams.StreamK

84 , mapMaybe

85 , sequence

87+ -- * Filtering

88+ , filter

89+ , take

90+

87 -- * Semigroup Style Composition

91 -- * Semigroup Style Composition

88 , serial

92 , serial

@@ -102,7 +106,9 @@ import Control.Monad (void)

102import Control.Monad.Reader.Class (MonadReader(..))

106import Control.Monad.Reader.Class (MonadReader(..))

103import Control.Monad.Trans.Class (MonadTrans(lift))

107import Control.Monad.Trans.Class (MonadTrans(lift))

104import Data.Semigroup (Semigroup(..))

108import Data.Semigroup (Semigroup(..))

109+import Prelude

105−import Prelude hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence)

110+ hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,

111+ take, filter)

106import qualified Prelude

112import qualified Prelude

107

113

108import Streamly.SVar

114import Streamly.SVar

@@ -517,6 +523,29 @@ repeat a = let x = cons a x in x

517fromFoldable :: (IsStream t, Foldable f) => f a -> t m a

523fromFoldable :: (IsStream t, Foldable f) => f a -> t m a

518fromFoldable = Prelude.foldr cons nil

524fromFoldable = Prelude.foldr cons nil

519

525

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+

520------------------------------------------------------------------------------

549------------------------------------------------------------------------------

521-- Semigroup

550-- Semigroup

522------------------------------------------------------------------------------

551------------------------------------------------------------------------------