diff --git a/src/Streamly/Internal/Data/Stream/IsStream/Top.hs b/src/Streamly/Internal/Data/Stream/IsStream/Top.hs index 2f624e8250..36624d947d 100644 --- a/src/Streamly/Internal/Data/Stream/IsStream/Top.hs +++ b/src/Streamly/Internal/Data/Stream/IsStream/Top.hs @@ -30,7 +30,7 @@ module Streamly.Internal.Data.Stream.IsStream.Top , intersectBy , intersectBySorted , differenceBy - , mergeDifferenceBy + , differenceBySorted , unionBy , unionBySorted @@ -633,11 +633,14 @@ differenceBy eq s1 s2 = -- -- Space: O(1) -- --- /Unimplemented/ -{-# INLINE mergeDifferenceBy #-} -mergeDifferenceBy :: -- (IsStream t, Monad m) => +-- /Pre-release/ +{-# INLINE differenceBySorted #-} +differenceBySorted :: (IsStream t, MonadIO m) => (a -> a -> Ordering) -> t m a -> t m a -> t m a -mergeDifferenceBy _eq _s1 _s2 = undefined +differenceBySorted eq s1 = + IsStream.fromStreamD + . StreamD.differenceBySorted eq (IsStream.toStreamD s1) + . IsStream.toStreamD -- | This is essentially an append operation that appends all the extra -- occurrences of elements from the second stream that are not already present diff --git a/src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs b/src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs index b8ea4771c0..f3a8b2ddc2 100644 --- a/src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs +++ b/src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs @@ -144,6 +144,7 @@ module Streamly.Internal.Data.Stream.StreamD.Nesting , splitInnerBySuffix , intersectBySorted , unionBySorted + , differenceBySorted ) where @@ -3039,3 +3040,67 @@ unionBySorted cmp (Stream stepa ta) (Stream stepb tb) = ) step _ (_, _, _, _, _, _, _) = return Stop + +------------------------------------------------------------------------------- +-- Difference of sorted streams ----------------------------------------------- +------------------------------------------------------------------------------- +{-# INLINE_NORMAL differenceBySorted #-} +differenceBySorted :: (Monad m) => + (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a +differenceBySorted cmp (Stream stepa ta) (Stream stepb tb) = + Stream step (Just ta, Just tb, Nothing, Nothing, Nothing) + + where + {-# INLINE_LATE step #-} + + -- one of the values is missing, and the corresponding stream is running + step gst (Just sa, sb, Nothing, b, Nothing) = do + r <- stepa gst sa + return $ case r of + Yield a sa' -> Skip (Just sa', sb, Just a, b, Nothing) + Skip sa' -> Skip (Just sa', sb, Nothing, b, Nothing) + Stop -> Skip (Nothing, sb, Nothing, b, Nothing) + + step gst (sa, Just sb, a, Nothing, Nothing) = do + r <- stepb gst sb + return $ case r of + Yield b sb' -> Skip (sa, Just sb', a, Just b, Nothing) + Skip sb' -> Skip (sa, Just sb', a, Nothing, Nothing) + Stop -> Skip (sa, Nothing, a, Nothing, Nothing) + + -- Matching element + step gst (Just sa, Just sb, Nothing, _, Just _) = do + r1 <- stepa gst sa + r2 <- stepb gst sb + return $ case r1 of + Yield a sa' -> + case r2 of + Yield c sb' -> + Skip (Just sa', Just sb', Just a, Just c, Nothing) + Skip sb' -> + Skip (Just sa', Just sb', Just a, Just a, Nothing) + Stop -> + Yield a (Just sa', Just sb, Nothing, Nothing, Just a) + Skip sa' -> + case r2 of + Yield c sb' -> + Skip (Just sa', Just sb', Just c, Just c, Nothing) + Skip sb' -> + Skip (Just sa', Just sb', Nothing, Nothing, Nothing) + Stop -> + Stop + Stop -> + Stop + + -- both the values are available + step _ (sa, sb, Just a, Just b, Nothing) = do + let res = cmp a b + return $ case res of + GT -> Skip (sa, sb, Just a, Nothing, Nothing) + LT -> Yield a (sa, sb, Nothing, Just b, Nothing) + EQ -> Skip (sa, sb, Nothing, Just b, Just b) + + -- one of the values is missing, corresponding stream is done + step _ (sa, Nothing, Just a, Nothing, Nothing) = + return $ Yield a (sa, Nothing, Nothing, Nothing , Nothing) + step _ (_, _, _, _, _) = return Stop diff --git a/test/Streamly/Test/Data/Stream/Top.hs b/test/Streamly/Test/Data/Stream/Top.hs index fbdd621a43..78684ede35 100644 --- a/test/Streamly/Test/Data/Stream/Top.hs +++ b/test/Streamly/Test/Data/Stream/Top.hs @@ -1,7 +1,7 @@ module Main (main) where -import Data.List (intersect, sort, union) +import Data.List (intersect, sort, union, (\\)) import Test.QuickCheck ( Gen , Property @@ -64,6 +64,26 @@ unionBySorted = (S.fromList ls1) let v2 = sort $ union ls0 ls1 assert (v1 == v2) + +differenceBySorted :: Property +differenceBySorted = + forAll (listOf (chooseInt (min_value, max_value))) $ \ls0 -> + forAll (listOf (chooseInt (min_value, max_value))) $ \ls1 -> + monadicIO $ action (sort ls0) (sort ls1) + + where + + action ls0 ls1 = do + v1 <- + run + $ S.toList + $ Top.differenceBySorted + compare + (S.fromList ls0) + (S.fromList ls1) + let v2 = ls0 \\ ls1 + assert (v1 == sort v2) + ------------------------------------------------------------------------------- moduleName :: String moduleName = "Data.Stream.Top" @@ -74,3 +94,4 @@ main = hspec $ do -- intersect prop "intersectBySorted" Main.intersectBySorted prop "unionBySorted" Main.unionBySorted + prop "differenceBySorted" Main.differenceBySorted