e65c5db · streamly · git.iko.soy

git.iko.soy/streamly

Add a takeEndBy transformation

Fix review comments Fix review comments Use fromStreamD and toStreamD instead of fromStreamS and add a test case Move test case to Common Fix hlint

commite65c5db97e635f7c5ef3cd617a67c9074ebbb4a1 authorRanjeet Kumar Ranjan <ranjeetr@composewell.com> dateWed 2 Feb 2022, 22:24 (5 years ago) parentbabb151

3 files changed+44−1

compare with master

unifiedsplit

src/Streamly/Internal/Data/Stream/IsStream/Common.hs+5 src/Streamly/Internal/Data/Stream/StreamD/Type.hs+28 test/lib/Streamly/Test/Prelude/Common.hs+11−1

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

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

3434 -- $smapM_Notes

3535 , take

3636 , takeWhile

37+ , takeEndBy

3738 , drop

3839 , findIndices

3940 , intersperseM

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

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

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

432433

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+

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

434439--

435440-- @since 0.1.0

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

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

7474 , take

7575 , takeWhile

7676 , takeWhileM

77+ , takeEndBy

78+ , takeEndByM

7779

7880 -- * Nesting

7981 , ConcatMapUState (..)

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

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

639641takeWhile f = takeWhileM (return . f)

640642

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+

641669------------------------------------------------------------------------------

642670-- Combine N Streams - concatAp

643671------------------------------------------------------------------------------

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

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

124124import System.Mem (performMajorGC)

125125import Test.Hspec.QuickCheck

126126import Test.Hspec

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

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

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

129129

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

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

134134import qualified Streamly.Prelude as S

135135import qualified Streamly.Data.Fold as FL

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

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

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

138139

139140import qualified Data.Map.Strict as Map

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

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

10751076 listEquals eq stream list

10761077

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+

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

10781086-- accidentally passed through them.

10791087--

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

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

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

11221130

1131+ prop "takeEndBy" takeEndBy

1132+

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

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

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