@@ -144,6 +144,7 @@ module Streamly.Internal.Data.Stream.StreamD.Nesting
144144 , splitInnerBySuffix
145145 , intersectBySorted
146146 , unionBySorted
147+ , differenceBySorted
147148 )
148149where
149150
@@ -3038,3 +3039,67 @@ unionBySorted cmp (Stream stepa ta) (Stream stepb tb) =
30383039 )
30393040
30403041 step _ (_, _, _, _, _, _, _) = return Stop
3042+
3043+ -------------------------------------------------------------------------------
3044+ -- Difference of sorted streams -----------------------------------------------
3045+ -------------------------------------------------------------------------------
3046+ {-# INLINE_NORMAL differenceBySorted #-}
3047+ differenceBySorted :: (Monad m ) =>
3048+ (a -> a -> Ordering ) -> Stream m a -> Stream m a -> Stream m a
3049+ differenceBySorted cmp (Stream stepa ta) (Stream stepb tb) =
3050+ Stream step (Just ta, Just tb, Nothing , Nothing , Nothing )
3051+
3052+ where
3053+ {-# INLINE_LATE step #-}
3054+
3055+ -- one of the values is missing, and the corresponding stream is running
3056+ step gst (Just sa, sb, Nothing , b, Nothing ) = do
3057+ r <- stepa gst sa
3058+ return $ case r of
3059+ Yield a sa' -> Skip (Just sa', sb, Just a, b, Nothing )
3060+ Skip sa' -> Skip (Just sa', sb, Nothing , b, Nothing )
3061+ Stop -> Skip (Nothing , sb, Nothing , b, Nothing )
3062+
3063+ step gst (sa, Just sb, a, Nothing , Nothing ) = do
3064+ r <- stepb gst sb
3065+ return $ case r of
3066+ Yield b sb' -> Skip (sa, Just sb', a, Just b, Nothing )
3067+ Skip sb' -> Skip (sa, Just sb', a, Nothing , Nothing )
3068+ Stop -> Skip (sa, Nothing , a, Nothing , Nothing )
3069+
3070+ -- Matching element
3071+ step gst (Just sa, Just sb, Nothing , _, Just _) = do
3072+ r1 <- stepa gst sa
3073+ r2 <- stepb gst sb
3074+ return $ case r1 of
3075+ Yield a sa' ->
3076+ case r2 of
3077+ Yield c sb' ->
3078+ Skip (Just sa', Just sb', Just a, Just c, Nothing )
3079+ Skip sb' ->
3080+ Skip (Just sa', Just sb', Just a, Just a, Nothing )
3081+ Stop ->
3082+ Yield a (Just sa', Just sb, Nothing , Nothing , Just a)
3083+ Skip sa' ->
3084+ case r2 of
3085+ Yield c sb' ->
3086+ Skip (Just sa', Just sb', Just c, Just c, Nothing )
3087+ Skip sb' ->
3088+ Skip (Just sa', Just sb', Nothing , Nothing , Nothing )
3089+ Stop ->
3090+ Stop
3091+ Stop ->
3092+ Stop
3093+
3094+ -- both the values are available
3095+ step _ (sa, sb, Just a, Just b, Nothing ) = do
3096+ let res = cmp a b
3097+ return $ case res of
3098+ GT -> Skip (sa, sb, Just a, Nothing , Nothing )
3099+ LT -> Yield a (sa, sb, Nothing , Just b, Nothing )
3100+ EQ -> Skip (sa, sb, Nothing , Just b, Just b)
3101+
3102+ -- one of the values is missing, corresponding stream is done
3103+ step _ (sa, Nothing , Just a, Nothing , Nothing ) =
3104+ return $ Yield a (sa, Nothing , Nothing , Nothing , Nothing )
3105+ step _ (_, _, _, _, _) = return Stop
0 commit comments