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