git.iko.soy/streamly
on 630fa8b help sign in
streamly/src/Streamly/Internal/Data/Stream/Extra.hs
163 lines 5.0 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.Extra
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
9module Streamly.Internal.Data.Stream.Extra
10 (
11 nub
12 , joinInnerMap
13 , joinLeftMap
14 , joinOuterMap
15 )
16where
18#include "inline.hs"
20import Control.Monad.IO.Class (MonadIO)
21import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
22import Streamly.Internal.Data.Stream.Type (Stream)
24import qualified Data.Map.Strict as Map
25import qualified Data.Set as Set
26import qualified Streamly.Data.Fold as Fold
27import qualified Streamly.Data.Stream as Stream
28import qualified Streamly.Internal.Data.Stream.StreamD as D
30-- $setup
31-- >>> :m
32-- >>> import qualified Streamly.Internal.Data.Stream as Stream
34-- | The memory used is proportional to the number of unique elements in the
35-- stream. If we want to limit the memory we can just use "take" to limit the
36-- uniq elements in the stream.
37{-# INLINE_NORMAL nub #-}
38nub :: (Monad m, Ord a) => D.Stream m a -> D.Stream m a
39nub (D.Stream step1 state1) = D.Stream step (Set.empty, state1)
41 where
43 step gst (set, st) = do
44 r <- step1 gst st
45 return
46 $ case r of
47 Yield x s ->
48 if Set.member x set
49 then Skip (set, s)
50 else Yield x (Set.insert x set, s)
51 Skip s -> Skip (set, s)
52 Stop -> Stop
54-- XXX Generate error if a duplicate insertion is attempted?
55toMap :: (Monad m, Ord k) => Stream m (k, v) -> m (Map.Map k v)
56toMap =
57 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
58 in Stream.fold f
60-- If the second stream is too big it can be partitioned based on hashes and
61-- then we can process one parition at a time.
62--
63-- XXX An IntMap may be faster when the keys are Int.
64-- XXX Use hashmap instead of map?
65--
66-- | Like 'joinInner' but uses a 'Map' for efficiency.
67--
68-- If the input streams have duplicate keys, the behavior is undefined.
69--
70-- For space efficiency use the smaller stream as the second stream.
71--
72-- Space: O(n)
73--
74-- Time: O(m + n)
75--
76-- /Pre-release/
77{-# INLINE joinInnerMap #-}
78joinInnerMap :: (Monad m, Ord k) =>
79 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, b)
80joinInnerMap s1 s2 =
81 Stream.concatEffect $ do
82 km <- toMap s2
83 pure $ Stream.mapMaybe (joinAB km) s1
85 where
87 joinAB kvm (k, a) =
88 case k `Map.lookup` kvm of
89 Just b -> Just (k, a, b)
90 Nothing -> Nothing
92-- | Like 'joinLeft' but uses a hashmap for efficiency.
93--
94-- Space: O(n)
95--
96-- Time: O(m + n)
97--
98-- /Pre-release/
99{-# INLINE joinLeftMap #-}
100joinLeftMap :: (Ord k, Monad m) =>
101 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, Maybe b)
102joinLeftMap s1 s2 =
103 Stream.concatEffect $ do
104 km <- toMap s2
105 return $ fmap (joinAB km) s1
107 where
109 joinAB km (k, a) =
110 case k `Map.lookup` km of
111 Just b -> (k, a, Just b)
112 Nothing -> (k, a, Nothing)
114-- Put the b's that have been paired, in another hash or mutate the hash to set
115-- a flag. At the end go through @Stream m b@ and find those that are not in that
116-- hash to return (Nothing, b).
117--
118-- | Like 'joinOuter' but uses a 'Map' for efficiency.
119--
120-- Space: O(m + n)
121--
122-- Time: O(m + n)
123--
124-- /Pre-release/
125{-# INLINE joinOuterMap #-}
126joinOuterMap ::
127 (Ord k, MonadIO m) =>
128 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, Maybe a, Maybe b)
129joinOuterMap s1 s2 =
130 Stream.concatEffect $ do
131 km1 <- kvFold s1
132 km2 <- kvFold s2
134 -- XXX Not sure if toList/fromList would fuse optimally. We may have to
135 -- create a fused Map.toStream function.
136 let res1 = fmap (joinAB km2)
137 $ Stream.fromList $ Map.toList km1
138 where
139 joinAB km (k, a) =
140 case k `Map.lookup` km of
141 Just b -> (k, Just a, Just b)
142 Nothing -> (k, Just a, Nothing)
144 -- XXX We can take advantage of the lookups in the first pass above to
145 -- reduce the number of lookups in this pass. If we keep mutable cells
146 -- in the second Map, we can flag it in the first pass and not do any
147 -- lookup in the second pass if it is flagged.
148 let res2 = Stream.mapMaybe (joinAB km1)
149 $ Stream.fromList $ Map.toList km2
150 where
151 joinAB km (k, b) =
152 case k `Map.lookup` km of
153 Just _ -> Nothing
154 Nothing -> Just (k, Nothing, Just b)
156 return $ Stream.append res1 res2
158 where
160 -- XXX Generate error if a duplicate insertion is attempted?
161 kvFold =
162 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
163 in Stream.fold f