@@ -21,42 +21,42 @@ module Streamly.Internal.Data.Stream.Top
2121 -- * Nesting
2222 -- ** Set like operations
2323 -- | These are not exactly set operations because streams are not
24− -- necessarily sets, they may have duplicated elements.
25− , intersectBy
26− , intersectBySorted
27− , differenceBy
28− , mergeDifferenceBy
29− , unionBy
30− , mergeUnionBy
24+ -- necessarily sets, they may have duplicated elements. These operations
25+ -- are generic i.e. they work on streams of unconstrained types, therefore,
26+ -- they have quadratic performance characterstics. For better performance
27+ -- using Set structures see the Streamly.Internal.Data.Stream.Container
28+ -- module.
29+ , filterInStreamGenericBy
30+ , deleteInStreamGenericBy
31+ , unionWithStreamGenericBy
32+
33+ -- ** Set like operations on sorted streams
34+ , filterInStreamAscBy
35+ , deleteInStreamAscBy
36+ , unionWithStreamAscBy
3137
3238 -- ** Join operations
33− , crossJoin
34− , joinInner
35− , joinInnerMerge
36− , joinLeft
37− , mergeLeftJoin
38− , joinOuter
39− , mergeOuterJoin
39+ , joinInnerGeneric
40+
41+ -- * Joins on sorted stream
42+ , joinInnerAscBy
43+ , joinLeftAscBy
44+ , joinOuterAscBy
4045 )
4146where
4247
4348#include "inline.hs"
4449
4550import Control.Monad.IO.Class (MonadIO(..))
46−import Control.Monad.Trans.State.Strict (get, put)
47−import Data.Function ((&))
4851import Data.IORef (newIORef, readIORef, modifyIORef')
49−import Data.Maybe (isJust)
52+import Streamly.Internal.Data.Fold.Type (Fold)
5053import Streamly.Internal.Data.Stream.Common ()
51−import Streamly.Internal.Data.Stream.Cross (CrossStream (..))
52−import Streamly.Internal.Data.Stream.Type (Stream, fromStreamD, toStreamD)
54+import Streamly.Internal.Data.Stream.Type
55+ (Stream, fromStreamD, toStreamD, cross)
5356
5457import qualified Data.List as List
55−import qualified Streamly.Internal.Data.Array.Generic as Array
56−import qualified Streamly.Internal.Data.Array.Mut.Type as MA
5758import qualified Streamly.Internal.Data.Fold as Fold
5859import qualified Streamly.Internal.Data.Parser as Parser
59−import qualified Streamly.Internal.Data.Stream.Transformer as Stream
6060import qualified Streamly.Internal.Data.Stream.Eliminate as Stream
6161import qualified Streamly.Internal.Data.Stream.Generate as Stream
6262import qualified Streamly.Internal.Data.Stream.Expand as Stream
@@ -140,59 +140,37 @@ sortBy cmp =
140140-- source changes during iterations. We can also use an Unfold instead of
141141-- stream. We probably need a way to distinguish streams that can be read
142142-- mutliple times without any interference (e.g. unfolding a stream using an
143−-- immutable handle would work i.e. using pread/pwrite instead of maintianing
143+-- immutable handle would work i.e. using pread/pwrite instead of maintaining
144144-- an offset in the handle).
145145
146−-- XXX We can do this concurrently.
147−--
148−-- | This is the same as 'Streamly.Internal.Data.Unfold.outerProduct' but less
149−-- efficient.
150−--
151−-- The second stream is evaluated multiple times. If the second stream is
152−-- consume-once stream then it can be cached in an 'Data.Array.Array' before
153−-- calling this function. Caching may also improve performance if the stream is
154−-- expensive to evaluate.
155−--
156−-- Time: O(m x n)
157−--
158−-- /Pre-release/
159−{-# INLINE crossJoin #-}
160−crossJoin :: Monad m => Stream m a -> Stream m b -> Stream m (a, b)
161−crossJoin s1 s2 = unCrossStream $ do
162− -- XXX use concatMap instead?
163− a <- CrossStream s1
164− b <- CrossStream s2
165− return (a, b)
166−
167146-- XXX We can do this concurrently.
168147-- XXX If the second stream is sorted and passed as an Array we could use
169148-- binary search if we have an Ord instance or Ordering returning function. The
170149-- time complexity would then become (m x log n).
150+
151+-- | Like 'cross' but emits only those tuples where @a == b@ using the
152+-- supplied equality predicate.
171153--
172−-- | For all elements in @Stream m a@, for all elements in @Stream m b@ if @a@ and @b@
173−-- are equal by the given equality pedicate then return the tuple (a, b).
174−--
175−-- The second stream is evaluated multiple times. If the stream is a
176−-- consume-once stream then the caller should cache it (e.g. in a
177−-- 'Data.Array.Array') before calling this function. Caching may also improve
178−-- performance if the stream is expensive to evaluate.
154+-- Definition:
179155--
180−-- For space efficiency use the smaller stream as the second stream.
156+-- >>> joinInnerGeneric eq s1 s2 = Stream.filter (\(a, b) -> a `eq` b) $ Stream.cross s1 s2
181157--
182−-- You should almost always use joinInnerMap instead of joinInner. joinInnerMap
183−-- is an order of magnitude faster. joinInner may be used when the second
184−-- stream is generated from a seed, therefore, need not be stored in memory and
185−-- the amount of memory it takes is a concern.
158+-- You should almost always prefer @joinInnerOrd@ over 'joinInnerGeneric' if
159+-- possible. @joinInnerOrd@ is an order of magnitude faster but may take more
160+-- space for caching the second stream.
186161--
187−-- Space: O(n) assuming the second stream is cached in memory.
162+-- See 'Streamly.Internal.Data.Unfold.joinInnerGeneric' for a much faster fused
163+-- alternative.
188164--
189165-- Time: O(m x n)
190166--
191167-- /Pre-release/
192−{-# INLINE joinInner #-}
193−joinInner :: Monad m =>
168+{-# INLINE joinInnerGeneric #-}
169+joinInnerGeneric :: Monad m =>
194170 (a -> b -> Bool) -> Stream m a -> Stream m b -> Stream m (a, b)
195−joinInner eq s1 s2 = do
171+joinInnerGeneric eq s1 s2 = Stream.filter (\(a, b) -> a `eq` b) $ cross s1 s2
172+{-
173+joinInnerGeneric eq s1 s2 = do
196174 -- ConcatMap works faster than bind
197175 Stream.concatMap (\a ->
198176 Stream.concatMap (\b ->
@@ -201,165 +179,46 @@ joinInner eq s1 s2 = do
201179 else Stream.nil
202180 ) s2
203181 ) s1
182+-}
204183
205−-- | Like 'joinInner' but works only on sorted streams.
184+-- | A more efficient 'joinInner' for sorted streams.
206185--
207186-- Space: O(1)
208187--
209188-- Time: O(m + n)
210189--
211190-- /Unimplemented/
212−{-# INLINE joinInnerMerge #-}
213−joinInnerMerge ::
191+{-# INLINE joinInnerAscBy #-}
192+joinInnerAscBy ::
214193 (a -> b -> Ordering) -> Stream m a -> Stream m b -> Stream m (a, b)
215−joinInnerMerge = undefined
194+joinInnerAscBy = undefined
216195
217−-- XXX We can do this concurrently.
218−-- XXX If the second stream is sorted and passed as an Array or a seek capable
219−-- stream then we could use binary search if we have an Ord instance or
220−-- Ordering returning function. The time complexity would then become (m x log
221−-- n).
222−--
223−-- | For all elements in @Stream m a@, for all elements in @Stream m b@ if @a@ and @b@
224−-- are equal then return the tuple @(a, Just b)@. If @a@ is not present in @t
225−-- m b@ then return @(a, Nothing)@.
226−--
227−-- The second stream is evaluated multiple times. If the stream is a
228−-- consume-once stream then the caller should cache it in an 'Data.Array.Array'
229−-- before calling this function. Caching may also improve performance if the
230−-- stream is expensive to evaluate.
231−--
232−-- @
233−-- rightJoin = flip joinLeft
234−-- @
235−--
236−-- Space: O(n) assuming the second stream is cached in memory.
237−--
238−-- Time: O(m x n)
239−--
240−-- /Unimplemented/
241−{-# INLINE joinLeft #-}
242−joinLeft :: Monad m =>
243− (a -> b -> Bool) -> Stream m a -> Stream m b -> Stream m (a, Maybe b)
244−joinLeft eq s1 s2 = Stream.evalStateT (return False) $ unCrossStream $ do
245− a <- CrossStream (Stream.liftInner s1)
246− -- XXX should we use StreamD monad here?
247− -- XXX Is there a better way to perform some action at the end of a loop
248− -- iteration?
249− CrossStream (Stream.fromEffect $ put False)
250− let final = Stream.concatEffect $ do
251− r <- get
252− if r
253− then pure Stream.nil
254− else pure (Stream.fromPure Nothing)
255− b <- CrossStream (fmap Just (Stream.liftInner s2) <> final)
256− case b of
257− Just b1 ->
258− if a `eq` b1
259− then do
260− CrossStream (Stream.fromEffect $ put True)
261− return (a, Just b1)
262− else CrossStream Stream.nil
263− Nothing -> return (a, Nothing)
264−
265−-- | Like 'joinLeft' but works only on sorted streams.
196+-- | A more efficient 'joinLeft' for sorted streams.
266197--
267198-- Space: O(1)
268199--
269200-- Time: O(m + n)
270201--
271202-- /Unimplemented/
272−{-# INLINE mergeLeftJoin #-}
273−mergeLeftJoin :: -- Monad m =>
203+{-# INLINE joinLeftAscBy #-}
204+joinLeftAscBy :: -- Monad m =>
274205 (a -> b -> Ordering) -> Stream m a -> Stream m b -> Stream m (a, Maybe b)
275−mergeLeftJoin _eq _s1 _s2 = undefined
206+joinLeftAscBy _eq _s1 _s2 = undefined
276207
277−-- XXX We can do this concurrently.
278−--
279−-- | For all elements in @t m a@, for all elements in @t m b@ if @a@ and @b@
280−-- are equal by the given equality pedicate then return the tuple (Just a, Just
281−-- b). If @a@ is not found in @t m b@ then return (a, Nothing), return
282−-- (Nothing, b) for vice-versa.
283−--
284−-- For space efficiency use the smaller stream as the second stream.
285−--
286−-- Space: O(n)
287−--
288−-- Time: O(m x n)
289−--
290−-- /Pre-release/
291−{-# INLINE joinOuter #-}
292−joinOuter :: MonadIO m =>
293− (a -> b -> Bool)
294− -> Stream m a
295− -> Stream m b
296− -> Stream m (Maybe a, Maybe b)
297−joinOuter eq s1 s =
298− Stream.concatEffect $ do
299− inputArr <- Array.fromStream s
300− let len = Array.length inputArr
301− foundArr <-
302− Stream.fold
303− (MA.writeN len)
304− (Stream.fromList (Prelude.replicate len False))
305− return $ go inputArr foundArr <> leftOver inputArr foundArr
306−
307− where
308−
309− leftOver inputArr foundArr =
310− let stream1 = Array.read inputArr
311− stream2 = Stream.unfold MA.reader foundArr
312− in Stream.filter
313− isJust
314− ( Stream.zipWith (\x y ->
315− if y
316− then Nothing
317− else Just (Nothing, Just x)
318− ) stream1 stream2
319− ) & Stream.catMaybes
320−
321− evalState = Stream.evalStateT (return False) . unCrossStream
322−
323− go inputArr foundArr = evalState $ do
324− a <- CrossStream (Stream.liftInner s1)
325− -- XXX should we use StreamD monad here?
326− -- XXX Is there a better way to perform some action at the end of a loop
327− -- iteration?
328− CrossStream (Stream.fromEffect $ put False)
329− let final = Stream.concatEffect $ do
330− r <- get
331− if r
332− then pure Stream.nil
333− else pure (Stream.fromPure Nothing)
334− (i, b) <-
335− let stream = Array.read inputArr
336− in CrossStream
337− (Stream.indexed $ fmap Just (Stream.liftInner stream) <> final)
338−
339− case b of
340− Just b1 ->
341− if a `eq` b1
342− then do
343− CrossStream (Stream.fromEffect $ put True)
344− MA.putIndex i True foundArr
345− return (Just a, Just b1)
346− else CrossStream Stream.nil
347− Nothing -> return (Just a, Nothing)
348−
349−-- | Like 'joinOuter' but works only on sorted streams.
208+-- | A more efficient 'joinOuter' for sorted streams.
350209--
351210-- Space: O(1)
352211--
353212-- Time: O(m + n)
354213--
355214-- /Unimplemented/
356−{-# INLINE mergeOuterJoin #-}
357−mergeOuterJoin :: -- Monad m =>
215+{-# INLINE joinOuterAscBy #-}
216+joinOuterAscBy :: -- Monad m =>
358217 (a -> b -> Ordering)
359218 -> Stream m a
360219 -> Stream m b
361220 -> Stream m (Maybe a, Maybe b)
362−mergeOuterJoin _eq _s1 _s2 = undefined
221+joinOuterAscBy _eq _s1 _s2 = undefined
363222
364223------------------------------------------------------------------------------
365224-- Set operations (special joins)
@@ -369,19 +228,39 @@ mergeOuterJoin _eq _s1 _s2 = undefined
369228-- the best would be to use an Array with linear search. If the second stream
370229-- is sorted we can also use a binary search, using Ord constraint.
371230
372−-- | 'intersectBy' is essentially a filtering operation that retains only those
373−-- elements in the first stream that are present in the second stream.
231+-- | Keep only those elements in the second stream that are present in the
232+-- first stream too. The first stream is folded to a container using the
233+-- supplied fold and then the elements in the container are looked up using the
234+-- supplied lookup function.
374235--
375−-- >>> Stream.fold Fold.toList $ Stream.intersectBy (==) (Stream.fromList [1,2,2,4]) (Stream.fromList [2,1,1,3])
376−-- [1,2,2]
236+-- The first stream must be finite and must not block.
237+{-# INLINE filterStreamWith #-}
238+filterStreamWith :: Monad m =>
239+ Fold m a (f a)
240+ -> (a -> f a -> Bool)
241+ -> Stream m a
242+ -> Stream m a
243+ -> Stream m a
244+filterStreamWith fld member s1 s2 =
245+ Stream.concatEffect
246+ $ do
247+ xs <- Stream.fold fld s1
248+ return $ Stream.filter (`member` xs) s2
249+
250+-- | 'filterInStreamGenericBy' retains only those elements in the second stream that
251+-- are present in the first stream.
377252--
378−-- >>> Stream.fold Fold.toList $ Stream.intersectBy (==) (Stream.fromList [2,1,1,3]) (Stream.fromList [1,2,2,4])
253+-- >>> Stream.fold Fold.toList $ Stream.filterInStreamGenericBy (==) (Stream.fromList [1,2,2,4]) (Stream.fromList [2,1,1,3])
379254-- [2,1,1]
380255--
381−-- 'intersectBy' is similar to but not the same as 'joinInner':
256+-- >>> Stream.fold Fold.toList $ Stream.filterInStreamGenericBy (==) (Stream.fromList [2,1,1,3]) (Stream.fromList [1,2,2,4])
257+-- [1,2,2]
258+--
259+-- Similar to the list intersectBy operation but with the stream argument order
260+-- flipped.
382261--
383−-- >>> Stream.fold Fold.toList $ fmap fst $ Stream.joinInner (==) (Stream.fromList [1,2,2,4]) (Stream.fromList [2,1,1,3])
384−-- [1,1,2,2]
262+-- The first stream must be finite and must not block. Second stream is
263+-- processed only after the first stream is fully realized.
385264--
386265-- Space: O(n) where @n@ is the number of elements in the second stream.
387266--
@@ -389,48 +268,49 @@ mergeOuterJoin _eq _s1 _s2 = undefined
389268-- @n@ is the number of elements in the second stream.
390269--
391270-- /Pre-release/
392−{-# INLINE intersectBy #-}
393−intersectBy :: Monad m =>
271+{-# INLINE filterInStreamGenericBy #-}
272+filterInStreamGenericBy :: Monad m =>
394273 (a -> a -> Bool) -> Stream m a -> Stream m a -> Stream m a
395−intersectBy eq s1 s2 =
396− Stream.concatEffect
397− $ do
398− -- This may work well when s2 is small
399− xs <- Stream.fold Fold.toListRev $ Stream.uniqBy eq s2
400− return $ Stream.filter (\x -> List.any (eq x) xs) s1
274+filterInStreamGenericBy eq =
275+ -- XXX Use an (unboxed) array instead.
276+ filterStreamWith
277+ (Fold.scanMaybe (Fold.uniqBy eq) Fold.toListRev)
278+ (List.any . eq)
401279
402−-- | Like 'intersectBy' but works only on streams sorted in ascending order.
280+-- | Like 'filterInStreamGenericBy' but assumes that the input streams are sorted in
281+-- ascending order. To use it on streams sorted in descending order pass an
282+-- inverted comparison function returning GT for less than and LT for greater
283+-- than.
403284--
404285-- Space: O(1)
405286--
406287-- Time: O(m+n)
407288--
408289-- /Pre-release/
409−{-# INLINE intersectBySorted #-}
410−intersectBySorted :: (Monad m) =>
290+{-# INLINE filterInStreamAscBy #-}
291+filterInStreamAscBy :: Monad m =>
411292 (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
412−intersectBySorted eq s1 =
293+filterInStreamAscBy eq s1 s2 =
413294 fromStreamD
414− . StreamD.intersectBySorted eq (toStreamD s1)
415− . toStreamD
416−
417−-- Roughly joinLeft s1 s2 = s1 `difference` s2 + s1 `intersection` s2
295+ $ StreamD.intersectBySorted eq (toStreamD s2)
296+ $ toStreamD s1
418297
419−-- | Delete first occurrences of those elements from the first stream that are
420−-- present in the second stream. If an element occurs multiple times in the
421−-- second stream as many occurrences of it are deleted from the first stream.
298+-- | Delete all elements of the first stream from the seconds stream. If an
299+-- element occurs multiple times in the first stream as many occurrences of it
300+-- are deleted from the second stream.
422301--
423−-- >>> Stream.fold Fold.toList $ Stream.differenceBy (==) (Stream.fromList [1,2,2]) (Stream.fromList [1,2,3])
302+-- >>> Stream.fold Fold.toList $ Stream.deleteInStreamGenericBy (==) (Stream.fromList [1,2,3]) (Stream.fromList [1,2,2])
424303-- [2]
425304--
426305-- The following laws hold:
427306--
428−-- @
429−-- (s1 `serial` s2) `differenceBy eq` s1 === s2
430−-- (s1 `wSerial` s2) `differenceBy eq` s1 === s2
431−-- @
307+-- > deleteInStreamGenericBy (==) s1 (s1 `append` s2) === s2
308+-- > deleteInStreamGenericBy (==) s1 (s1 `interleave` s2) === s2
309+--
310+-- Same as the list 'Data.List.//' operation but with argument order flipped.
432311--
433−-- Same as the list 'Data.List.//' operation.
312+-- The first stream must be finite and must not block. Second stream is
313+-- processed only after the first stream is fully realized.
434314--
435315-- Space: O(m) where @m@ is the number of elements in the first stream.
436316--
@@ -438,10 +318,10 @@ intersectBySorted eq s1 =
438318-- @n@ is the number of elements in the second stream.
439319--
440320-- /Pre-release/
441−{-# INLINE differenceBy #-}
442−differenceBy :: (Monad m) =>
321+{-# INLINE deleteInStreamGenericBy #-}
322+deleteInStreamGenericBy :: Monad m =>
443323 (a -> a -> Bool) -> Stream m a -> Stream m a -> Stream m a
444−differenceBy eq s1 s2 =
324+deleteInStreamGenericBy eq s1 s2 =
445325 Stream.concatEffect
446326 $ do
447327 -- This may work well if s1 is small
@@ -449,47 +329,48 @@ differenceBy eq s1 s2 =
449329 -- not emitting an element if it was successfully deleted from s2.
450330 -- we will need a deleteBy that can return whether the element was
451331 -- deleted or not.
452− xs <- Stream.fold Fold.toList s1
332+ xs <- Stream.fold Fold.toList s2
453333 let f = Fold.foldl' (flip (List.deleteBy eq)) xs
454− fmap Stream.fromList $ Stream.fold f s2
334+ fmap Stream.fromList $ Stream.fold f s1
455335
456−-- | Like 'differenceBy' but works only on sorted streams.
336+-- | A more efficient 'deleteInStreamGenericBy' for streams sorted in ascending order.
457337--
458338-- Space: O(1)
459339--
460340-- /Unimplemented/
461−{-# INLINE mergeDifferenceBy #-}
462−mergeDifferenceBy :: -- (Monad m) =>
341+{-# INLINE deleteInStreamAscBy #-}
342+deleteInStreamAscBy :: -- (Monad m) =>
463343 (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
464−mergeDifferenceBy _eq _s1 _s2 = undefined
344+deleteInStreamAscBy _eq _s1 _s2 = undefined
465345
466−-- | This is essentially an append operation that appends all the extra
467−-- occurrences of elements from the second stream that are not already present
468−-- in the first stream.
346+-- XXX Remove the MonadIO constraint. We can just cache one stream and then
347+-- implement using differenceEqBy.
348+
349+-- | This essentially appends to the second stream all the occurrences of
350+-- elements in the first stream that are not already present in the second
351+-- stream.
469352--
470−-- >>> Stream.fold Fold.toList $ Stream.unionBy (==) (Stream.fromList [1,2,2,4]) (Stream.fromList [1,1,2,3])
471−-- [1,2,2,4,3]
353+-- Equivalent to the following except that @s2@ is evaluated only once:
472354--
473−-- Equivalent to the following except that @s1@ is evaluated only once:
355+-- >>> unionWithStreamGenericBy eq s1 s2 = s2 `Stream.append` (Stream.deleteInStreamGenericBy eq s2 s1)
474356--
475−-- @
476−-- unionBy eq s1 s2 = s1 \`serial` (s2 `differenceBy eq` s1)
477−-- @
357+-- Example:
478358--
479−-- Similar to 'joinOuter' but not the same.
359+-- >>> Stream.fold Fold.toList $ Stream.unionWithStreamGenericBy (==) (Stream.fromList [1,1,2,3]) (Stream.fromList [1,2,2,4])
360+-- [1,2,2,4,3]
480361--
481362-- Space: O(n)
482363--
483364-- Time: O(m x n)
484365--
485366-- /Pre-release/
486−{-# INLINE unionBy #-}
487−unionBy :: MonadIO m =>
367+{-# INLINE unionWithStreamGenericBy #-}
368+unionWithStreamGenericBy :: MonadIO m =>
488369 (a -> a -> Bool) -> Stream m a -> Stream m a -> Stream m a
489−unionBy eq s1 s2 =
370+unionWithStreamGenericBy eq s1 s2 =
490371 Stream.concatEffect
491372 $ do
492− xs <- Stream.fold Fold.toList s2
373+ xs <- Stream.fold Fold.toList s1
493374 -- XXX we can use postscanlMAfter' instead of IORef
494375 ref <- liftIO $ newIORef $! List.nubBy eq xs
495376 let f x = do
@@ -499,14 +380,14 @@ unionBy eq s1 s2 =
499380 $ do
500381 xs1 <- liftIO $ readIORef ref
501382 return $ Stream.fromList xs1
502− return $ Stream.mapM f s1 <> s3
383+ return $ Stream.mapM f s2 <> s3
503384
504−-- | Like 'unionBy' but works only on sorted streams.
385+-- | A more efficient 'unionWithStreamGenericBy' for sorted streams.
505386--
506387-- Space: O(1)
507388--
508389-- /Unimplemented/
509−{-# INLINE mergeUnionBy #-}
510−mergeUnionBy :: -- (Monad m) =>
390+{-# INLINE unionWithStreamAscBy #-}
391+unionWithStreamAscBy :: -- (Monad m) =>
511392 (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
512−mergeUnionBy _eq _s1 _s2 = undefined
393+unionWithStreamAscBy _eq _s1 _s2 = undefined