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

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

@@ -34,6 +34,7 @@ module Streamly.Internal.Data.Stream.IsStream.Common
     -- $smapM_Notes
     , take
     , takeWhile
+    , takeEndBy
     , drop
     , findIndices
     , intersperseM
@@ -430,6 +431,10 @@ take n m = fromStreamS $ S.take n $ toStreamS
 takeWhile :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a
 takeWhile p m = fromStreamS $ S.takeWhile p $ toStreamS m
 
+{-# INLINE takeEndBy #-}
+takeEndBy :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a
+takeEndBy p m = fromStreamD $ D.takeEndBy p $ toStreamD m
+
 -- | Discard first 'n' elements from the stream and take the rest.
 --
 -- @since 0.1.0

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

@@ -74,6 +74,8 @@ module Streamly.Internal.Data.Stream.StreamD.Type
     , take
     , takeWhile
     , takeWhileM
+    , takeEndBy
+    , takeEndByM
 
     -- * Nesting
     , ConcatMapUState (..)
@@ -638,6 +640,32 @@ takeWhileM f (Stream step state) = Stream step' state
 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
 takeWhile f = takeWhileM (return . f)
 
+-- Like takeWhile but with an inverted condition and also taking
+-- the matching element.
+
+{-# INLINE_NORMAL takeEndByM #-}
+takeEndByM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
+takeEndByM f (Stream step state) = Stream step' (Just state)
+  where
+    {-# INLINE_LATE step' #-}
+    step' gst (Just st) = do
+        r <- step gst st
+        case r of
+            Yield x s -> do
+                b <- f x
+                return $
+                    if not b
+                    then Yield x (Just s)
+                    else Yield x Nothing
+            Skip s -> return $ Skip (Just s)
+            Stop   -> return Stop
+
+    step' _ Nothing = return Stop
+
+{-# INLINE takeEndBy #-}
+takeEndBy :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
+takeEndBy f = takeEndByM (return . f)
+
 ------------------------------------------------------------------------------
 -- Combine N Streams - concatAp
 ------------------------------------------------------------------------------

test/lib/Streamly/Test/Prelude/Common.hs

@@ -124,7 +124,7 @@ import GHC.Word (Word8)
 import System.Mem (performMajorGC)
 import Test.Hspec.QuickCheck
 import Test.Hspec
-import Test.QuickCheck (Property, choose, forAll, withMaxSuccess)
+import Test.QuickCheck (Property, choose, forAll, listOf, withMaxSuccess)
 import Test.QuickCheck.Monadic (assert, monadicIO, run)
 
 import Streamly.Prelude (SerialT, IsStream, (.:), nil, (|&), fromSerial)
@@ -134,6 +134,7 @@ import Streamly.Prelude (avgRate, rate, maxBuffer, maxThreads)
 import qualified Streamly.Prelude as S
 import qualified Streamly.Data.Fold as FL
 import qualified Streamly.Internal.Data.Stream.IsStream as S
+import qualified Streamly.Internal.Data.Stream.IsStream.Common as IS
 import qualified Streamly.Internal.Data.Unfold as UF
 
 import qualified Data.Map.Strict as Map
@@ -1074,6 +1075,13 @@ transformCombineFromList constr eq listOp t op a b c =
             let list = a <> listOp (b <> c)
             listEquals eq stream list
 
+takeEndBy :: Property
+takeEndBy = forAll (listOf (chooseInt (0, maxStreamLen))) $ \lst -> monadicIO $ do
+    let (s1, s3) = span (<= 200) lst
+    let s4 = [head s3 | not (null s3)]
+    s2 <- run $ S.toList $ IS.takeEndBy (> 200) $ S.fromList lst
+    assert $ s1 ++ s4 == s2
+
 -- XXX add tests for MonadReader and MonadError etc. In case an SVar is
 -- accidentally passed through them.
 --
@@ -1120,6 +1128,8 @@ transformCombineOpsCommon constr desc eq t = do
     prop (desc <> " takeWhileM False") $
         transform (takeWhile (const False)) t (S.takeWhileM (const $ return False))
 
+    prop "takeEndBy" takeEndBy
+
     prop (desc <> " drop maxBound") $
         transform (drop maxBound) t (S.drop maxBound)
     prop (desc <> " drop 0") $ transform (drop 0) t (S.drop 0)