git.iko.soy/streamly
on 79df417 help sign in
streamly/core/src/Streamly/Internal/Data/Stream/Container.hs
302 lines 9.9 kB history raw
1{-# LANGUAGE CPP #-}
2-- |
3-- Module : Streamly.Internal.Data.Stream.Container
4-- Copyright : (c) 2019 Composewell Technologies
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- Stream operations that require transformers or containers like Set or Map.
12module Streamly.Internal.Data.Stream.Container
13 (
14 nub
16 -- * Joins for unconstrained types
17 , joinLeftGeneric
18 , joinOuterGeneric
20 -- * Joins with Ord constraint
21 , joinInner
22 , joinLeft
23 , joinOuter
24 )
25where
27#include "inline.hs"
29import Control.Monad.IO.Class (MonadIO)
30import Control.Monad.Trans.State.Strict (get, put)
31import Data.Function ((&))
32import Data.Maybe (isJust)
33import Streamly.Internal.Data.Stream.Step (Step(..))
34import Streamly.Internal.Data.Stream.Type
35 (Stream(..), mkCross, unCross)
37import qualified Data.Map.Strict as Map
38import qualified Data.Set as Set
39import qualified Streamly.Data.Fold as Fold
40import qualified Streamly.Internal.Data.Array.Generic as Array
41import qualified Streamly.Internal.Data.MutArray.Type as MA
42import qualified Streamly.Internal.Data.Stream.Type as Stream
43import qualified Streamly.Internal.Data.Stream.Nesting as Stream
44import qualified Streamly.Internal.Data.Stream.Generate as Stream
45import qualified Streamly.Internal.Data.Stream.Transform as Stream
46import qualified Streamly.Internal.Data.Stream.Transformer as Stream
48#include "DocTestDataStream.hs"
50-- | The memory used is proportional to the number of unique elements in the
51-- stream. If we want to limit the memory we can just use "take" to limit the
52-- uniq elements in the stream.
53{-# INLINE_NORMAL nub #-}
54nub :: (Monad m, Ord a) => Stream m a -> Stream m a
55nub (Stream step1 state1) = Stream step (Set.empty, state1)
57 where
59 step gst (set, st) = do
60 r <- step1 gst st
61 return
62 $ case r of
63 Yield x s ->
64 if Set.member x set
65 then Skip (set, s)
66 else Yield x (Set.insert x set, s)
67 Skip s -> Skip (set, s)
68 Stop -> Stop
70-- XXX Generate error if a duplicate insertion is attempted?
71toMap :: (Monad m, Ord k) => Stream m (k, v) -> m (Map.Map k v)
72toMap =
73 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
74 in Stream.fold f
76-- If the second stream is too big it can be partitioned based on hashes and
77-- then we can process one parition at a time.
78--
79-- XXX An IntMap may be faster when the keys are Int.
80-- XXX Use hashmap instead of map?
81--
82-- | Like 'joinInner' but uses a 'Map' for efficiency.
83--
84-- If the input streams have duplicate keys, the behavior is undefined.
85--
86-- For space efficiency use the smaller stream as the second stream.
87--
88-- Space: O(n)
89--
90-- Time: O(m + n)
91--
92-- /Pre-release/
93{-# INLINE joinInner #-}
94joinInner :: (Monad m, Ord k) =>
95 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, b)
96joinInner s1 s2 =
97 Stream.concatEffect $ do
98 km <- toMap s2
99 pure $ Stream.mapMaybe (joinAB km) s1
101 where
103 joinAB kvm (k, a) =
104 case k `Map.lookup` kvm of
105 Just b -> Just (k, a, b)
106 Nothing -> Nothing
108-- XXX We can do this concurrently.
109-- XXX If the second stream is sorted and passed as an Array or a seek capable
110-- stream then we could use binary search if we have an Ord instance or
111-- Ordering returning function. The time complexity would then become (m x log
112-- n).
114-- XXX Check performance of StreamD vs StreamK
116-- | Like 'joinInner' but emit @(a, Just b)@, and additionally, for those @a@'s
117-- that are not equal to any @b@ emit @(a, Nothing)@.
118--
119-- The second stream is evaluated multiple times. If the stream is a
120-- consume-once stream then the caller should cache it in an 'Data.Array.Array'
121-- before calling this function. Caching may also improve performance if the
122-- stream is expensive to evaluate.
123--
124-- >>> joinRightGeneric eq = flip (Stream.joinLeftGeneric eq)
125--
126-- Space: O(n) assuming the second stream is cached in memory.
127--
128-- Time: O(m x n)
129--
130-- /Unimplemented/
131{-# INLINE joinLeftGeneric #-}
132joinLeftGeneric :: Monad m =>
133 (a -> b -> Bool) -> Stream m a -> Stream m b -> Stream m (a, Maybe b)
134joinLeftGeneric eq s1 s2 = Stream.evalStateT (return False) $ unCross $ do
135 a <- mkCross (Stream.liftInner s1)
136 -- XXX should we use StreamD monad here?
137 -- XXX Is there a better way to perform some action at the end of a loop
138 -- iteration?
139 mkCross (Stream.fromEffect $ put False)
140 let final = Stream.concatEffect $ do
141 r <- get
142 if r
143 then pure Stream.nil
144 else pure (Stream.fromPure Nothing)
145 b <- mkCross (fmap Just (Stream.liftInner s2) `Stream.append` final)
146 case b of
147 Just b1 ->
148 if a `eq` b1
149 then do
150 mkCross (Stream.fromEffect $ put True)
151 return (a, Just b1)
152 else mkCross Stream.nil
153 Nothing -> return (a, Nothing)
155-- XXX rename to joinLeftOrd?
157-- | A more efficient 'joinLeft' using a hashmap for efficiency.
158--
159-- Space: O(n)
160--
161-- Time: O(m + n)
162--
163-- /Pre-release/
164{-# INLINE joinLeft #-}
165joinLeft :: (Ord k, Monad m) =>
166 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, Maybe b)
167joinLeft s1 s2 =
168 Stream.concatEffect $ do
169 km <- toMap s2
170 return $ fmap (joinAB km) s1
172 where
174 joinAB km (k, a) =
175 case k `Map.lookup` km of
176 Just b -> (k, a, Just b)
177 Nothing -> (k, a, Nothing)
179-- XXX We can do this concurrently.
181-- XXX Check performance of StreamD vs StreamK
183-- | Like 'joinLeft' but emits a @(Just a, Just b)@. Like 'joinLeft', for those
184-- @a@'s that are not equal to any @b@ emit @(Just a, Nothing)@, but
185-- additionally, for those @b@'s that are not equal to any @a@ emit @(Nothing,
186-- Just b)@.
187--
188-- For space efficiency use the smaller stream as the second stream.
189--
190-- Space: O(n)
191--
192-- Time: O(m x n)
193--
194-- /Pre-release/
195{-# INLINE joinOuterGeneric #-}
196joinOuterGeneric :: MonadIO m =>
197 (a -> b -> Bool)
198 -> Stream m a
199 -> Stream m b
200 -> Stream m (Maybe a, Maybe b)
201joinOuterGeneric eq s1 s =
202 Stream.concatEffect $ do
203 inputArr <- Array.fromStream s
204 let len = Array.length inputArr
205 foundArr <-
206 Stream.fold
207 (MA.writeN len)
208 (Stream.fromList (Prelude.replicate len False))
209 return $ go inputArr foundArr `Stream.append` leftOver inputArr foundArr
211 where
213 leftOver inputArr foundArr =
214 let stream1 = Array.read inputArr
215 stream2 = Stream.unfold MA.reader foundArr
216 in Stream.filter
217 isJust
218 ( Stream.zipWith (\x y ->
219 if y
220 then Nothing
221 else Just (Nothing, Just x)
222 ) stream1 stream2
223 ) & Stream.catMaybes
225 evalState = Stream.evalStateT (return False) . unCross
227 go inputArr foundArr = evalState $ do
228 a <- mkCross (Stream.liftInner s1)
229 -- XXX should we use StreamD monad here?
230 -- XXX Is there a better way to perform some action at the end of a loop
231 -- iteration?
232 mkCross (Stream.fromEffect $ put False)
233 let final = Stream.concatEffect $ do
234 r <- get
235 if r
236 then pure Stream.nil
237 else pure (Stream.fromPure Nothing)
238 (i, b) <-
239 let stream = Array.read inputArr
240 in mkCross
241 (Stream.indexed $ fmap Just (Stream.liftInner stream) `Stream.append` final)
243 case b of
244 Just b1 ->
245 if a `eq` b1
246 then do
247 mkCross (Stream.fromEffect $ put True)
248 MA.putIndex i foundArr True
249 return (Just a, Just b1)
250 else mkCross Stream.nil
251 Nothing -> return (Just a, Nothing)
253-- Put the b's that have been paired, in another hash or mutate the hash to set
254-- a flag. At the end go through @Stream m b@ and find those that are not in that
255-- hash to return (Nothing, b).
257-- | Like 'joinOuter' but uses a 'Map' for efficiency.
258--
259-- Space: O(m + n)
260--
261-- Time: O(m + n)
262--
263-- /Pre-release/
264{-# INLINE joinOuter #-}
265joinOuter ::
266 (Ord k, MonadIO m) =>
267 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, Maybe a, Maybe b)
268joinOuter s1 s2 =
269 Stream.concatEffect $ do
270 km1 <- kvFold s1
271 km2 <- kvFold s2
273 -- XXX Not sure if toList/fromList would fuse optimally. We may have to
274 -- create a fused Map.toStream function.
275 let res1 = fmap (joinAB km2)
276 $ Stream.fromList $ Map.toList km1
277 where
278 joinAB km (k, a) =
279 case k `Map.lookup` km of
280 Just b -> (k, Just a, Just b)
281 Nothing -> (k, Just a, Nothing)
283 -- XXX We can take advantage of the lookups in the first pass above to
284 -- reduce the number of lookups in this pass. If we keep mutable cells
285 -- in the second Map, we can flag it in the first pass and not do any
286 -- lookup in the second pass if it is flagged.
287 let res2 = Stream.mapMaybe (joinAB km1)
288 $ Stream.fromList $ Map.toList km2
289 where
290 joinAB km (k, b) =
291 case k `Map.lookup` km of
292 Just _ -> Nothing
293 Nothing -> Just (k, Nothing, Just b)
295 return $ Stream.append res1 res2
297 where
299 -- XXX Generate error if a duplicate insertion is attempted?
300 kvFold =
301 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
302 in Stream.fold f