Add kind equalities to GHC.
[ghc.git] / compiler / utils / Bag.hs
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 DeriveDataTypeable, ScopedTypeVariables, CPP #-}
10
11 module Bag (
12 Bag, -- abstract type
13
14 emptyBag, unitBag, unionBags, unionManyBags,
15 mapBag,
16 elemBag, lengthBag,
17 filterBag, partitionBag, partitionBagWith,
18 concatBag, catBagMaybes, foldBag, foldrBag, foldlBag,
19 isEmptyBag, isSingletonBag, consBag, snocBag, anyBag,
20 listToBag, bagToList,
21 foldrBagM, foldlBagM, mapBagM, mapBagM_,
22 flatMapBagM, flatMapBagPairM,
23 mapAndUnzipBagM, mapAccumBagLM,
24 anyBagM, filterBagM
25 ) where
26
27 import Outputable
28 import Util
29
30 import MonadUtils
31 import Control.Monad
32 import Data.Data
33 import Data.List ( partition )
34 import qualified Data.Foldable as Foldable
35
36 infixr 3 `consBag`
37 infixl 3 `snocBag`
38
39 data Bag a
40 = EmptyBag
41 | UnitBag a
42 | TwoBags (Bag a) (Bag a) -- INVARIANT: neither branch is empty
43 | ListBag [a] -- INVARIANT: the list is non-empty
44 deriving Typeable
45
46 emptyBag :: Bag a
47 emptyBag = EmptyBag
48
49 unitBag :: a -> Bag a
50 unitBag = UnitBag
51
52 lengthBag :: Bag a -> Int
53 lengthBag EmptyBag = 0
54 lengthBag (UnitBag {}) = 1
55 lengthBag (TwoBags b1 b2) = lengthBag b1 + lengthBag b2
56 lengthBag (ListBag xs) = length xs
57
58 elemBag :: Eq a => a -> Bag a -> Bool
59 elemBag _ EmptyBag = False
60 elemBag x (UnitBag y) = x == y
61 elemBag x (TwoBags b1 b2) = x `elemBag` b1 || x `elemBag` b2
62 elemBag x (ListBag ys) = any (x ==) ys
63
64 unionManyBags :: [Bag a] -> Bag a
65 unionManyBags xs = foldr unionBags EmptyBag xs
66
67 -- This one is a bit stricter! The bag will get completely evaluated.
68
69 unionBags :: Bag a -> Bag a -> Bag a
70 unionBags EmptyBag b = b
71 unionBags b EmptyBag = b
72 unionBags b1 b2 = TwoBags b1 b2
73
74 consBag :: a -> Bag a -> Bag a
75 snocBag :: Bag a -> a -> Bag a
76
77 consBag elt bag = (unitBag elt) `unionBags` bag
78 snocBag bag elt = bag `unionBags` (unitBag elt)
79
80 isEmptyBag :: Bag a -> Bool
81 isEmptyBag EmptyBag = True
82 isEmptyBag _ = False -- NB invariants
83
84 isSingletonBag :: Bag a -> Bool
85 isSingletonBag EmptyBag = False
86 isSingletonBag (UnitBag _) = True
87 isSingletonBag (TwoBags _ _) = False -- Neither is empty
88 isSingletonBag (ListBag xs) = isSingleton xs
89
90 filterBag :: (a -> Bool) -> Bag a -> Bag a
91 filterBag _ EmptyBag = EmptyBag
92 filterBag pred b@(UnitBag val) = if pred val then b else EmptyBag
93 filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
94 where sat1 = filterBag pred b1
95 sat2 = filterBag pred b2
96 filterBag pred (ListBag vs) = listToBag (filter pred vs)
97
98 filterBagM :: Monad m => (a -> m Bool) -> Bag a -> m (Bag a)
99 filterBagM _ EmptyBag = return EmptyBag
100 filterBagM pred b@(UnitBag val) = do
101 flag <- pred val
102 if flag then return b
103 else return EmptyBag
104 filterBagM pred (TwoBags b1 b2) = do
105 sat1 <- filterBagM pred b1
106 sat2 <- filterBagM pred b2
107 return (sat1 `unionBags` sat2)
108 filterBagM pred (ListBag vs) = do
109 sat <- filterM pred vs
110 return (listToBag sat)
111
112 anyBag :: (a -> Bool) -> Bag a -> Bool
113 anyBag _ EmptyBag = False
114 anyBag p (UnitBag v) = p v
115 anyBag p (TwoBags b1 b2) = anyBag p b1 || anyBag p b2
116 anyBag p (ListBag xs) = any p xs
117
118 anyBagM :: Monad m => (a -> m Bool) -> Bag a -> m Bool
119 anyBagM _ EmptyBag = return False
120 anyBagM p (UnitBag v) = p v
121 anyBagM p (TwoBags b1 b2) = do flag <- anyBagM p b1
122 if flag then return True
123 else anyBagM p b2
124 anyBagM p (ListBag xs) = anyM p xs
125
126 concatBag :: Bag (Bag a) -> Bag a
127 concatBag bss = foldrBag add emptyBag bss
128 where
129 add bs rs = bs `unionBags` rs
130
131 catBagMaybes :: Bag (Maybe a) -> Bag a
132 catBagMaybes bs = foldrBag add emptyBag bs
133 where
134 add Nothing rs = rs
135 add (Just x) rs = x `consBag` rs
136
137 partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
138 Bag a {- Don't -})
139 partitionBag _ EmptyBag = (EmptyBag, EmptyBag)
140 partitionBag pred b@(UnitBag val)
141 = if pred val then (b, EmptyBag) else (EmptyBag, b)
142 partitionBag pred (TwoBags b1 b2)
143 = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
144 where (sat1, fail1) = partitionBag pred b1
145 (sat2, fail2) = partitionBag pred b2
146 partitionBag pred (ListBag vs) = (listToBag sats, listToBag fails)
147 where (sats, fails) = partition pred vs
148
149
150 partitionBagWith :: (a -> Either b c) -> Bag a
151 -> (Bag b {- Left -},
152 Bag c {- Right -})
153 partitionBagWith _ EmptyBag = (EmptyBag, EmptyBag)
154 partitionBagWith pred (UnitBag val)
155 = case pred val of
156 Left a -> (UnitBag a, EmptyBag)
157 Right b -> (EmptyBag, UnitBag b)
158 partitionBagWith pred (TwoBags b1 b2)
159 = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
160 where (sat1, fail1) = partitionBagWith pred b1
161 (sat2, fail2) = partitionBagWith pred b2
162 partitionBagWith pred (ListBag vs) = (listToBag sats, listToBag fails)
163 where (sats, fails) = partitionWith pred vs
164
165 foldBag :: (r -> r -> r) -- Replace TwoBags with this; should be associative
166 -> (a -> r) -- Replace UnitBag with this
167 -> r -- Replace EmptyBag with this
168 -> Bag a
169 -> r
170
171 {- Standard definition
172 foldBag t u e EmptyBag = e
173 foldBag t u e (UnitBag x) = u x
174 foldBag t u e (TwoBags b1 b2) = (foldBag t u e b1) `t` (foldBag t u e b2)
175 foldBag t u e (ListBag xs) = foldr (t.u) e xs
176 -}
177
178 -- More tail-recursive definition, exploiting associativity of "t"
179 foldBag _ _ e EmptyBag = e
180 foldBag t u e (UnitBag x) = u x `t` e
181 foldBag t u e (TwoBags b1 b2) = foldBag t u (foldBag t u e b2) b1
182 foldBag t u e (ListBag xs) = foldr (t.u) e xs
183
184 foldrBag :: (a -> r -> r) -> r
185 -> Bag a
186 -> r
187
188 foldrBag _ z EmptyBag = z
189 foldrBag k z (UnitBag x) = k x z
190 foldrBag k z (TwoBags b1 b2) = foldrBag k (foldrBag k z b2) b1
191 foldrBag k z (ListBag xs) = foldr k z xs
192
193 foldlBag :: (r -> a -> r) -> r
194 -> Bag a
195 -> r
196
197 foldlBag _ z EmptyBag = z
198 foldlBag k z (UnitBag x) = k z x
199 foldlBag k z (TwoBags b1 b2) = foldlBag k (foldlBag k z b1) b2
200 foldlBag k z (ListBag xs) = foldl k z xs
201
202 foldrBagM :: (Monad m) => (a -> b -> m b) -> b -> Bag a -> m b
203 foldrBagM _ z EmptyBag = return z
204 foldrBagM k z (UnitBag x) = k x z
205 foldrBagM k z (TwoBags b1 b2) = do { z' <- foldrBagM k z b2; foldrBagM k z' b1 }
206 foldrBagM k z (ListBag xs) = foldrM k z xs
207
208 foldlBagM :: (Monad m) => (b -> a -> m b) -> b -> Bag a -> m b
209 foldlBagM _ z EmptyBag = return z
210 foldlBagM k z (UnitBag x) = k z x
211 foldlBagM k z (TwoBags b1 b2) = do { z' <- foldlBagM k z b1; foldlBagM k z' b2 }
212 foldlBagM k z (ListBag xs) = foldlM k z xs
213
214 mapBag :: (a -> b) -> Bag a -> Bag b
215 mapBag _ EmptyBag = EmptyBag
216 mapBag f (UnitBag x) = UnitBag (f x)
217 mapBag f (TwoBags b1 b2) = TwoBags (mapBag f b1) (mapBag f b2)
218 mapBag f (ListBag xs) = ListBag (map f xs)
219
220 mapBagM :: Monad m => (a -> m b) -> Bag a -> m (Bag b)
221 mapBagM _ EmptyBag = return EmptyBag
222 mapBagM f (UnitBag x) = do r <- f x
223 return (UnitBag r)
224 mapBagM f (TwoBags b1 b2) = do r1 <- mapBagM f b1
225 r2 <- mapBagM f b2
226 return (TwoBags r1 r2)
227 mapBagM f (ListBag xs) = do rs <- mapM f xs
228 return (ListBag rs)
229
230 mapBagM_ :: Monad m => (a -> m b) -> Bag a -> m ()
231 mapBagM_ _ EmptyBag = return ()
232 mapBagM_ f (UnitBag x) = f x >> return ()
233 mapBagM_ f (TwoBags b1 b2) = mapBagM_ f b1 >> mapBagM_ f b2
234 mapBagM_ f (ListBag xs) = mapM_ f xs
235
236 flatMapBagM :: Monad m => (a -> m (Bag b)) -> Bag a -> m (Bag b)
237 flatMapBagM _ EmptyBag = return EmptyBag
238 flatMapBagM f (UnitBag x) = f x
239 flatMapBagM f (TwoBags b1 b2) = do r1 <- flatMapBagM f b1
240 r2 <- flatMapBagM f b2
241 return (r1 `unionBags` r2)
242 flatMapBagM f (ListBag xs) = foldrM k EmptyBag xs
243 where
244 k x b2 = do { b1 <- f x; return (b1 `unionBags` b2) }
245
246 flatMapBagPairM :: Monad m => (a -> m (Bag b, Bag c)) -> Bag a -> m (Bag b, Bag c)
247 flatMapBagPairM _ EmptyBag = return (EmptyBag, EmptyBag)
248 flatMapBagPairM f (UnitBag x) = f x
249 flatMapBagPairM f (TwoBags b1 b2) = do (r1,s1) <- flatMapBagPairM f b1
250 (r2,s2) <- flatMapBagPairM f b2
251 return (r1 `unionBags` r2, s1 `unionBags` s2)
252 flatMapBagPairM f (ListBag xs) = foldrM k (EmptyBag, EmptyBag) xs
253 where
254 k x (r2,s2) = do { (r1,s1) <- f x
255 ; return (r1 `unionBags` r2, s1 `unionBags` s2) }
256
257 mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
258 mapAndUnzipBagM _ EmptyBag = return (EmptyBag, EmptyBag)
259 mapAndUnzipBagM f (UnitBag x) = do (r,s) <- f x
260 return (UnitBag r, UnitBag s)
261 mapAndUnzipBagM f (TwoBags b1 b2) = do (r1,s1) <- mapAndUnzipBagM f b1
262 (r2,s2) <- mapAndUnzipBagM f b2
263 return (TwoBags r1 r2, TwoBags s1 s2)
264 mapAndUnzipBagM f (ListBag xs) = do ts <- mapM f xs
265 let (rs,ss) = unzip ts
266 return (ListBag rs, ListBag ss)
267
268 mapAccumBagLM :: Monad m
269 => (acc -> x -> m (acc, y)) -- ^ combining funcction
270 -> acc -- ^ initial state
271 -> Bag x -- ^ inputs
272 -> m (acc, Bag y) -- ^ final state, outputs
273 mapAccumBagLM _ s EmptyBag = return (s, EmptyBag)
274 mapAccumBagLM f s (UnitBag x) = do { (s1, x1) <- f s x; return (s1, UnitBag x1) }
275 mapAccumBagLM f s (TwoBags b1 b2) = do { (s1, b1') <- mapAccumBagLM f s b1
276 ; (s2, b2') <- mapAccumBagLM f s1 b2
277 ; return (s2, TwoBags b1' b2') }
278 mapAccumBagLM f s (ListBag xs) = do { (s', xs') <- mapAccumLM f s xs
279 ; return (s', ListBag xs') }
280
281 listToBag :: [a] -> Bag a
282 listToBag [] = EmptyBag
283 listToBag vs = ListBag vs
284
285 bagToList :: Bag a -> [a]
286 bagToList b = foldrBag (:) [] b
287
288 instance (Outputable a) => Outputable (Bag a) where
289 ppr bag = braces (pprWithCommas ppr (bagToList bag))
290
291 instance Data a => Data (Bag a) where
292 gfoldl k z b = z listToBag `k` bagToList b -- traverse abstract type abstractly
293 toConstr _ = abstractConstr $ "Bag("++show (typeOf (undefined::a))++")"
294 gunfold _ _ = error "gunfold"
295 dataTypeOf _ = mkNoRepType "Bag"
296 dataCast1 x = gcast1 x
297
298 instance Foldable.Foldable Bag where
299 foldr = foldrBag