never executed always true always false
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4
5
6 Bag: an unordered collection with duplicates
7 -}
8
9 {-# LANGUAGE ScopedTypeVariables, DeriveFunctor, TypeFamilies #-}
10
11 module GHC.Data.Bag (
12 Bag, -- abstract type
13
14 emptyBag, unitBag, unionBags, unionManyBags,
15 mapBag,
16 elemBag, lengthBag,
17 filterBag, partitionBag, partitionBagWith,
18 concatBag, catBagMaybes, foldBag,
19 isEmptyBag, isSingletonBag, consBag, snocBag, anyBag, allBag,
20 listToBag, nonEmptyToBag, bagToList, headMaybe, mapAccumBagL,
21 concatMapBag, concatMapBagPair, mapMaybeBag,
22 mapBagM, mapBagM_,
23 flatMapBagM, flatMapBagPairM,
24 mapAndUnzipBagM, mapAccumBagLM,
25 anyBagM, filterBagM
26 ) where
27
28 import GHC.Prelude
29
30 import GHC.Exts ( IsList(..) )
31 import GHC.Utils.Outputable
32 import GHC.Utils.Misc
33 import GHC.Utils.Monad
34 import Control.Monad
35 import Data.Data
36 import Data.Maybe( mapMaybe, listToMaybe )
37 import Data.List ( partition, mapAccumL )
38 import Data.List.NonEmpty ( NonEmpty(..) )
39 import qualified Data.Foldable as Foldable
40 import qualified Data.Semigroup ( (<>) )
41
42 infixr 3 `consBag`
43 infixl 3 `snocBag`
44
45 data Bag a
46 = EmptyBag
47 | UnitBag a
48 | TwoBags (Bag a) (Bag a) -- INVARIANT: neither branch is empty
49 | ListBag [a] -- INVARIANT: the list is non-empty
50 deriving (Functor)
51
52 emptyBag :: Bag a
53 emptyBag = EmptyBag
54
55 unitBag :: a -> Bag a
56 unitBag = UnitBag
57
58 lengthBag :: Bag a -> Int
59 lengthBag EmptyBag = 0
60 lengthBag (UnitBag {}) = 1
61 lengthBag (TwoBags b1 b2) = lengthBag b1 + lengthBag b2
62 lengthBag (ListBag xs) = length xs
63
64 elemBag :: Eq a => a -> Bag a -> Bool
65 elemBag _ EmptyBag = False
66 elemBag x (UnitBag y) = x == y
67 elemBag x (TwoBags b1 b2) = x `elemBag` b1 || x `elemBag` b2
68 elemBag x (ListBag ys) = any (x ==) ys
69
70 unionManyBags :: [Bag a] -> Bag a
71 unionManyBags xs = foldr unionBags EmptyBag xs
72
73 -- This one is a bit stricter! The bag will get completely evaluated.
74
75 unionBags :: Bag a -> Bag a -> Bag a
76 unionBags EmptyBag b = b
77 unionBags b EmptyBag = b
78 unionBags b1 b2 = TwoBags b1 b2
79
80 consBag :: a -> Bag a -> Bag a
81 snocBag :: Bag a -> a -> Bag a
82
83 consBag elt bag = (unitBag elt) `unionBags` bag
84 snocBag bag elt = bag `unionBags` (unitBag elt)
85
86 isEmptyBag :: Bag a -> Bool
87 isEmptyBag EmptyBag = True
88 isEmptyBag _ = False -- NB invariants
89
90 isSingletonBag :: Bag a -> Bool
91 isSingletonBag EmptyBag = False
92 isSingletonBag (UnitBag _) = True
93 isSingletonBag (TwoBags _ _) = False -- Neither is empty
94 isSingletonBag (ListBag xs) = isSingleton xs
95
96 filterBag :: (a -> Bool) -> Bag a -> Bag a
97 filterBag _ EmptyBag = EmptyBag
98 filterBag pred b@(UnitBag val) = if pred val then b else EmptyBag
99 filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
100 where sat1 = filterBag pred b1
101 sat2 = filterBag pred b2
102 filterBag pred (ListBag vs) = listToBag (filter pred vs)
103
104 filterBagM :: Monad m => (a -> m Bool) -> Bag a -> m (Bag a)
105 filterBagM _ EmptyBag = return EmptyBag
106 filterBagM pred b@(UnitBag val) = do
107 flag <- pred val
108 if flag then return b
109 else return EmptyBag
110 filterBagM pred (TwoBags b1 b2) = do
111 sat1 <- filterBagM pred b1
112 sat2 <- filterBagM pred b2
113 return (sat1 `unionBags` sat2)
114 filterBagM pred (ListBag vs) = do
115 sat <- filterM pred vs
116 return (listToBag sat)
117
118 allBag :: (a -> Bool) -> Bag a -> Bool
119 allBag _ EmptyBag = True
120 allBag p (UnitBag v) = p v
121 allBag p (TwoBags b1 b2) = allBag p b1 && allBag p b2
122 allBag p (ListBag xs) = all p xs
123
124 anyBag :: (a -> Bool) -> Bag a -> Bool
125 anyBag _ EmptyBag = False
126 anyBag p (UnitBag v) = p v
127 anyBag p (TwoBags b1 b2) = anyBag p b1 || anyBag p b2
128 anyBag p (ListBag xs) = any p xs
129
130 anyBagM :: Monad m => (a -> m Bool) -> Bag a -> m Bool
131 anyBagM _ EmptyBag = return False
132 anyBagM p (UnitBag v) = p v
133 anyBagM p (TwoBags b1 b2) = do flag <- anyBagM p b1
134 if flag then return True
135 else anyBagM p b2
136 anyBagM p (ListBag xs) = anyM p xs
137
138 concatBag :: Bag (Bag a) -> Bag a
139 concatBag bss = foldr add emptyBag bss
140 where
141 add bs rs = bs `unionBags` rs
142
143 catBagMaybes :: Bag (Maybe a) -> Bag a
144 catBagMaybes bs = foldr add emptyBag bs
145 where
146 add Nothing rs = rs
147 add (Just x) rs = x `consBag` rs
148
149 partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
150 Bag a {- Don't -})
151 partitionBag _ EmptyBag = (EmptyBag, EmptyBag)
152 partitionBag pred b@(UnitBag val)
153 = if pred val then (b, EmptyBag) else (EmptyBag, b)
154 partitionBag pred (TwoBags b1 b2)
155 = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
156 where (sat1, fail1) = partitionBag pred b1
157 (sat2, fail2) = partitionBag pred b2
158 partitionBag pred (ListBag vs) = (listToBag sats, listToBag fails)
159 where (sats, fails) = partition pred vs
160
161
162 partitionBagWith :: (a -> Either b c) -> Bag a
163 -> (Bag b {- Left -},
164 Bag c {- Right -})
165 partitionBagWith _ EmptyBag = (EmptyBag, EmptyBag)
166 partitionBagWith pred (UnitBag val)
167 = case pred val of
168 Left a -> (UnitBag a, EmptyBag)
169 Right b -> (EmptyBag, UnitBag b)
170 partitionBagWith pred (TwoBags b1 b2)
171 = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
172 where (sat1, fail1) = partitionBagWith pred b1
173 (sat2, fail2) = partitionBagWith pred b2
174 partitionBagWith pred (ListBag vs) = (listToBag sats, listToBag fails)
175 where (sats, fails) = partitionWith pred vs
176
177 foldBag :: (r -> r -> r) -- Replace TwoBags with this; should be associative
178 -> (a -> r) -- Replace UnitBag with this
179 -> r -- Replace EmptyBag with this
180 -> Bag a
181 -> r
182
183 {- Standard definition
184 foldBag t u e EmptyBag = e
185 foldBag t u e (UnitBag x) = u x
186 foldBag t u e (TwoBags b1 b2) = (foldBag t u e b1) `t` (foldBag t u e b2)
187 foldBag t u e (ListBag xs) = foldr (t.u) e xs
188 -}
189
190 -- More tail-recursive definition, exploiting associativity of "t"
191 foldBag _ _ e EmptyBag = e
192 foldBag t u e (UnitBag x) = u x `t` e
193 foldBag t u e (TwoBags b1 b2) = foldBag t u (foldBag t u e b2) b1
194 foldBag t u e (ListBag xs) = foldr (t.u) e xs
195
196 mapBag :: (a -> b) -> Bag a -> Bag b
197 mapBag = fmap
198
199 concatMapBag :: (a -> Bag b) -> Bag a -> Bag b
200 concatMapBag _ EmptyBag = EmptyBag
201 concatMapBag f (UnitBag x) = f x
202 concatMapBag f (TwoBags b1 b2) = unionBags (concatMapBag f b1) (concatMapBag f b2)
203 concatMapBag f (ListBag xs) = foldr (unionBags . f) emptyBag xs
204
205 concatMapBagPair :: (a -> (Bag b, Bag c)) -> Bag a -> (Bag b, Bag c)
206 concatMapBagPair _ EmptyBag = (EmptyBag, EmptyBag)
207 concatMapBagPair f (UnitBag x) = f x
208 concatMapBagPair f (TwoBags b1 b2) = (unionBags r1 r2, unionBags s1 s2)
209 where
210 (r1, s1) = concatMapBagPair f b1
211 (r2, s2) = concatMapBagPair f b2
212 concatMapBagPair f (ListBag xs) = foldr go (emptyBag, emptyBag) xs
213 where
214 go a (s1, s2) = (unionBags r1 s1, unionBags r2 s2)
215 where
216 (r1, r2) = f a
217
218 mapMaybeBag :: (a -> Maybe b) -> Bag a -> Bag b
219 mapMaybeBag _ EmptyBag = EmptyBag
220 mapMaybeBag f (UnitBag x) = case f x of
221 Nothing -> EmptyBag
222 Just y -> UnitBag y
223 mapMaybeBag f (TwoBags b1 b2) = unionBags (mapMaybeBag f b1) (mapMaybeBag f b2)
224 mapMaybeBag f (ListBag xs) = ListBag (mapMaybe f xs)
225
226 mapBagM :: Monad m => (a -> m b) -> Bag a -> m (Bag b)
227 mapBagM _ EmptyBag = return EmptyBag
228 mapBagM f (UnitBag x) = do r <- f x
229 return (UnitBag r)
230 mapBagM f (TwoBags b1 b2) = do r1 <- mapBagM f b1
231 r2 <- mapBagM f b2
232 return (TwoBags r1 r2)
233 mapBagM f (ListBag xs) = do rs <- mapM f xs
234 return (ListBag rs)
235
236 mapBagM_ :: Monad m => (a -> m b) -> Bag a -> m ()
237 mapBagM_ _ EmptyBag = return ()
238 mapBagM_ f (UnitBag x) = f x >> return ()
239 mapBagM_ f (TwoBags b1 b2) = mapBagM_ f b1 >> mapBagM_ f b2
240 mapBagM_ f (ListBag xs) = mapM_ f xs
241
242 flatMapBagM :: Monad m => (a -> m (Bag b)) -> Bag a -> m (Bag b)
243 flatMapBagM _ EmptyBag = return EmptyBag
244 flatMapBagM f (UnitBag x) = f x
245 flatMapBagM f (TwoBags b1 b2) = do r1 <- flatMapBagM f b1
246 r2 <- flatMapBagM f b2
247 return (r1 `unionBags` r2)
248 flatMapBagM f (ListBag xs) = foldrM k EmptyBag xs
249 where
250 k x b2 = do { b1 <- f x; return (b1 `unionBags` b2) }
251
252 flatMapBagPairM :: Monad m => (a -> m (Bag b, Bag c)) -> Bag a -> m (Bag b, Bag c)
253 flatMapBagPairM _ EmptyBag = return (EmptyBag, EmptyBag)
254 flatMapBagPairM f (UnitBag x) = f x
255 flatMapBagPairM f (TwoBags b1 b2) = do (r1,s1) <- flatMapBagPairM f b1
256 (r2,s2) <- flatMapBagPairM f b2
257 return (r1 `unionBags` r2, s1 `unionBags` s2)
258 flatMapBagPairM f (ListBag xs) = foldrM k (EmptyBag, EmptyBag) xs
259 where
260 k x (r2,s2) = do { (r1,s1) <- f x
261 ; return (r1 `unionBags` r2, s1 `unionBags` s2) }
262
263 mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
264 mapAndUnzipBagM _ EmptyBag = return (EmptyBag, EmptyBag)
265 mapAndUnzipBagM f (UnitBag x) = do (r,s) <- f x
266 return (UnitBag r, UnitBag s)
267 mapAndUnzipBagM f (TwoBags b1 b2) = do (r1,s1) <- mapAndUnzipBagM f b1
268 (r2,s2) <- mapAndUnzipBagM f b2
269 return (TwoBags r1 r2, TwoBags s1 s2)
270 mapAndUnzipBagM f (ListBag xs) = do ts <- mapM f xs
271 let (rs,ss) = unzip ts
272 return (ListBag rs, ListBag ss)
273
274 mapAccumBagL ::(acc -> x -> (acc, y)) -- ^ combining function
275 -> acc -- ^ initial state
276 -> Bag x -- ^ inputs
277 -> (acc, Bag y) -- ^ final state, outputs
278 mapAccumBagL _ s EmptyBag = (s, EmptyBag)
279 mapAccumBagL f s (UnitBag x) = let (s1, x1) = f s x in (s1, UnitBag x1)
280 mapAccumBagL f s (TwoBags b1 b2) = let (s1, b1') = mapAccumBagL f s b1
281 (s2, b2') = mapAccumBagL f s1 b2
282 in (s2, TwoBags b1' b2')
283 mapAccumBagL f s (ListBag xs) = let (s', xs') = mapAccumL f s xs
284 in (s', ListBag xs')
285
286 mapAccumBagLM :: Monad m
287 => (acc -> x -> m (acc, y)) -- ^ combining function
288 -> acc -- ^ initial state
289 -> Bag x -- ^ inputs
290 -> m (acc, Bag y) -- ^ final state, outputs
291 mapAccumBagLM _ s EmptyBag = return (s, EmptyBag)
292 mapAccumBagLM f s (UnitBag x) = do { (s1, x1) <- f s x; return (s1, UnitBag x1) }
293 mapAccumBagLM f s (TwoBags b1 b2) = do { (s1, b1') <- mapAccumBagLM f s b1
294 ; (s2, b2') <- mapAccumBagLM f s1 b2
295 ; return (s2, TwoBags b1' b2') }
296 mapAccumBagLM f s (ListBag xs) = do { (s', xs') <- mapAccumLM f s xs
297 ; return (s', ListBag xs') }
298
299 listToBag :: [a] -> Bag a
300 listToBag [] = EmptyBag
301 listToBag [x] = UnitBag x
302 listToBag vs = ListBag vs
303
304 nonEmptyToBag :: NonEmpty a -> Bag a
305 nonEmptyToBag (x :| []) = UnitBag x
306 nonEmptyToBag (x :| xs) = ListBag (x : xs)
307
308 bagToList :: Bag a -> [a]
309 bagToList b = foldr (:) [] b
310
311 headMaybe :: Bag a -> Maybe a
312 headMaybe EmptyBag = Nothing
313 headMaybe (UnitBag v) = Just v
314 headMaybe (TwoBags b1 _) = headMaybe b1
315 headMaybe (ListBag l) = listToMaybe l
316
317 instance (Outputable a) => Outputable (Bag a) where
318 ppr bag = braces (pprWithCommas ppr (bagToList bag))
319
320 instance Data a => Data (Bag a) where
321 gfoldl k z b = z listToBag `k` bagToList b -- traverse abstract type abstractly
322 toConstr _ = abstractConstr $ "Bag("++show (typeOf (undefined::a))++")"
323 gunfold _ _ = error "gunfold"
324 dataTypeOf _ = mkNoRepType "Bag"
325 dataCast1 x = gcast1 x
326
327 instance Foldable.Foldable Bag where
328 foldr _ z EmptyBag = z
329 foldr k z (UnitBag x) = k x z
330 foldr k z (TwoBags b1 b2) = foldr k (foldr k z b2) b1
331 foldr k z (ListBag xs) = foldr k z xs
332
333 foldl _ z EmptyBag = z
334 foldl k z (UnitBag x) = k z x
335 foldl k z (TwoBags b1 b2) = foldl k (foldl k z b1) b2
336 foldl k z (ListBag xs) = foldl k z xs
337
338 foldl' _ z EmptyBag = z
339 foldl' k z (UnitBag x) = k z x
340 foldl' k z (TwoBags b1 b2) = let r1 = foldl' k z b1 in seq r1 $ foldl' k r1 b2
341 foldl' k z (ListBag xs) = foldl' k z xs
342
343 instance Traversable Bag where
344 traverse _ EmptyBag = pure EmptyBag
345 traverse f (UnitBag x) = UnitBag <$> f x
346 traverse f (TwoBags b1 b2) = TwoBags <$> traverse f b1 <*> traverse f b2
347 traverse f (ListBag xs) = ListBag <$> traverse f xs
348
349 instance IsList (Bag a) where
350 type Item (Bag a) = a
351 fromList = listToBag
352 toList = bagToList
353
354 instance Semigroup (Bag a) where
355 (<>) = unionBags
356
357 instance Monoid (Bag a) where
358 mempty = emptyBag