e65c5db · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add a takeEndBy transformation

e65c5dbRanjeet Kumar Ranjan5 years ago

src/Streamly/Internal/Data/Stream/IsStream/Common.hs+5

compare with master

unifiedsplit

@@ -34,6 +34,7 @@ module Streamly.Internal.Data.Stream.IsStream.Common

34 -- $smapM_Notes

35 , take

36 , takeWhile

37+ , takeEndBy

37 , drop

38 , drop

38 , findIndices

39 , findIndices

39 , intersperseM

40 , intersperseM

@@ -430,6 +431,10 @@ take n m = fromStreamS $ S.take n $ toStreamS

430takeWhile :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a

431takeWhile :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a

431takeWhile p m = fromStreamS $ S.takeWhile p $ toStreamS m

432takeWhile p m = fromStreamS $ S.takeWhile p $ toStreamS m

432

433

434+{-# INLINE takeEndBy #-}

435+takeEndBy :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a

436+takeEndBy p m = fromStreamD $ D.takeEndBy p $ toStreamD m

437+

433-- | Discard first 'n' elements from the stream and take the rest.

438-- | Discard first 'n' elements from the stream and take the rest.

434--

439--

435-- @since 0.1.0

440-- @since 0.1.0

src/Streamly/Internal/Data/Stream/StreamD/Type.hs+28

@@ -74,6 +74,8 @@ module Streamly.Internal.Data.Stream.StreamD.Type

74 , take

75 , takeWhile

76 , takeWhileM

77+ , takeEndBy

78+ , takeEndByM

78 -- * Nesting

80 -- * Nesting

79 , ConcatMapUState (..)

81 , ConcatMapUState (..)

@@ -638,6 +640,32 @@ takeWhileM f (Stream step state) = Stream step' state

638takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

640takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

639takeWhile f = takeWhileM (return . f)

641takeWhile f = takeWhileM (return . f)

640

642

643+-- Like takeWhile but with an inverted condition and also taking

644+-- the matching element.

645+

646+{-# INLINE_NORMAL takeEndByM #-}

647+takeEndByM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a

648+takeEndByM f (Stream step state) = Stream step' (Just state)

649+ where

650+ {-# INLINE_LATE step' #-}

651+ step' gst (Just st) = do

652+ r <- step gst st

653+ case r of

654+ Yield x s -> do

655+ b <- f x

656+ return $

657+ if not b

658+ then Yield x (Just s)

659+ else Yield x Nothing

660+ Skip s -> return $ Skip (Just s)

661+ Stop -> return Stop

662+

663+ step' _ Nothing = return Stop

664+

665+{-# INLINE takeEndBy #-}

666+takeEndBy :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

667+takeEndBy f = takeEndByM (return . f)

668+

641------------------------------------------------------------------------------

669------------------------------------------------------------------------------

642-- Combine N Streams - concatAp

670-- Combine N Streams - concatAp

643------------------------------------------------------------------------------

671------------------------------------------------------------------------------

test/lib/Streamly/Test/Prelude/Common.hs+11−1

@@ -124,7 +124,7 @@ import GHC.Word (Word8)

124import System.Mem (performMajorGC)

125import Test.Hspec.QuickCheck

126import Test.Hspec

127−import Test.QuickCheck (Property, choose, forAll, withMaxSuccess)

127+import Test.QuickCheck (Property, choose, forAll, listOf, withMaxSuccess)

128import Test.QuickCheck.Monadic (assert, monadicIO, run)

129

130import Streamly.Prelude (SerialT, IsStream, (.:), nil, (|&), fromSerial)

@@ -134,6 +134,7 @@ import Streamly.Prelude (avgRate, rate, maxBuffer, maxThreads)

134import qualified Streamly.Prelude as S

135import qualified Streamly.Data.Fold as FL

136import qualified Streamly.Internal.Data.Stream.IsStream as S

137+import qualified Streamly.Internal.Data.Stream.IsStream.Common as IS

137import qualified Streamly.Internal.Data.Unfold as UF

138import qualified Streamly.Internal.Data.Unfold as UF

138

139

139import qualified Data.Map.Strict as Map

140import qualified Data.Map.Strict as Map

@@ -1074,6 +1075,13 @@ transformCombineFromList constr eq listOp t op a b c =

1074 let list = a <> listOp (b <> c)

1075 let list = a <> listOp (b <> c)

1075 listEquals eq stream list

1076 listEquals eq stream list

1076

1077

1078+takeEndBy :: Property

1079+takeEndBy = forAll (listOf (chooseInt (0, maxStreamLen))) $ \lst -> monadicIO $ do

1080+ let (s1, s3) = span (<= 200) lst

1081+ let s4 = [head s3 | not (null s3)]

1082+ s2 <- run $ S.toList $ IS.takeEndBy (> 200) $ S.fromList lst

1083+ assert $ s1 ++ s4 == s2

1084+

1077-- XXX add tests for MonadReader and MonadError etc. In case an SVar is

1085-- XXX add tests for MonadReader and MonadError etc. In case an SVar is

1078-- accidentally passed through them.

1086-- accidentally passed through them.

1079--

1087--

@@ -1120,6 +1128,8 @@ transformCombineOpsCommon constr desc eq t = do

1120 prop (desc <> " takeWhileM False") $

1128 prop (desc <> " takeWhileM False") $

1121 transform (takeWhile (const False)) t (S.takeWhileM (const $ return False))

1129 transform (takeWhile (const False)) t (S.takeWhileM (const $ return False))

1122

1130

1131+ prop "takeEndBy" takeEndBy

1132+

1123 prop (desc <> " drop maxBound") $

1133 prop (desc <> " drop maxBound") $

1124 transform (drop maxBound) t (S.drop maxBound)

1134 transform (drop maxBound) t (S.drop maxBound)

1125 prop (desc <> " drop 0") $ transform (drop 0) t (S.drop 0)

1135 prop (desc <> " drop 0") $ transform (drop 0) t (S.drop 0)