95feaed9f88483984d12a219ae23e988a5e64174
[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, 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 EmptyBag = EmptyBag
103 concatBag (UnitBag b) = b
104 concatBag (TwoBags b1 b2) = concatBag b1 `unionBags` concatBag b2
105 concatBag (ListBag bs) = unionManyBags bs
106
107 partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
108 Bag a {- Don't -})
109 partitionBag _ EmptyBag = (EmptyBag, EmptyBag)
110 partitionBag pred b@(UnitBag val)
111 = if pred val then (b, EmptyBag) else (EmptyBag, b)
112 partitionBag pred (TwoBags b1 b2)
113 = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
114 where (sat1, fail1) = partitionBag pred b1
115 (sat2, fail2) = partitionBag pred b2
116 partitionBag pred (ListBag vs) = (listToBag sats, listToBag fails)
117 where (sats, fails) = partition pred vs
118
119
120 partitionBagWith :: (a -> Either b c) -> Bag a
121 -> (Bag b {- Left -},
122 Bag c {- Right -})
123 partitionBagWith _ EmptyBag = (EmptyBag, EmptyBag)
124 partitionBagWith pred (UnitBag val)
125 = case pred val of
126 Left a -> (UnitBag a, EmptyBag)
127 Right b -> (EmptyBag, UnitBag b)
128 partitionBagWith pred (TwoBags b1 b2)
129 = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
130 where (sat1, fail1) = partitionBagWith pred b1
131 (sat2, fail2) = partitionBagWith pred b2
132 partitionBagWith pred (ListBag vs) = (listToBag sats, listToBag fails)
133 where (sats, fails) = partitionWith pred vs
134
135 foldBag :: (r -> r -> r) -- Replace TwoBags with this; should be associative
136 -> (a -> r) -- Replace UnitBag with this
137 -> r -- Replace EmptyBag with this
138 -> Bag a
139 -> r
140
141 {- Standard definition
142 foldBag t u e EmptyBag = e
143 foldBag t u e (UnitBag x) = u x
144 foldBag t u e (TwoBags b1 b2) = (foldBag t u e b1) `t` (foldBag t u e b2)
145 foldBag t u e (ListBag xs) = foldr (t.u) e xs
146 -}
147
148 -- More tail-recursive definition, exploiting associativity of "t"
149 foldBag _ _ e EmptyBag = e
150 foldBag t u e (UnitBag x) = u x `t` e
151 foldBag t u e (TwoBags b1 b2) = foldBag t u (foldBag t u e b2) b1
152 foldBag t u e (ListBag xs) = foldr (t.u) e xs
153
154 foldrBag :: (a -> r -> r) -> r
155 -> Bag a
156 -> r
157
158 foldrBag _ z EmptyBag = z
159 foldrBag k z (UnitBag x) = k x z
160 foldrBag k z (TwoBags b1 b2) = foldrBag k (foldrBag k z b2) b1
161 foldrBag k z (ListBag xs) = foldr k z xs
162
163 foldlBag :: (r -> a -> r) -> r
164 -> Bag a
165 -> r
166
167 foldlBag _ z EmptyBag = z
168 foldlBag k z (UnitBag x) = k z x
169 foldlBag k z (TwoBags b1 b2) = foldlBag k (foldlBag k z b1) b2
170 foldlBag k z (ListBag xs) = foldl k z xs
171
172 foldrBagM :: (Monad m) => (a -> b -> m b) -> b -> Bag a -> m b
173 foldrBagM _ z EmptyBag = return z
174 foldrBagM k z (UnitBag x) = k x z
175 foldrBagM k z (TwoBags b1 b2) = do { z' <- foldrBagM k z b2; foldrBagM k z' b1 }
176 foldrBagM k z (ListBag xs) = foldrM k z xs
177
178 foldlBagM :: (Monad m) => (b -> a -> m b) -> b -> Bag a -> m b
179 foldlBagM _ z EmptyBag = return z
180 foldlBagM k z (UnitBag x) = k z x
181 foldlBagM k z (TwoBags b1 b2) = do { z' <- foldlBagM k z b1; foldlBagM k z' b2 }
182 foldlBagM k z (ListBag xs) = foldlM k z xs
183
184 mapBag :: (a -> b) -> Bag a -> Bag b
185 mapBag _ EmptyBag = EmptyBag
186 mapBag f (UnitBag x) = UnitBag (f x)
187 mapBag f (TwoBags b1 b2) = TwoBags (mapBag f b1) (mapBag f b2)
188 mapBag f (ListBag xs) = ListBag (map f xs)
189
190 mapBagM :: Monad m => (a -> m b) -> Bag a -> m (Bag b)
191 mapBagM _ EmptyBag = return EmptyBag
192 mapBagM f (UnitBag x) = do r <- f x
193 return (UnitBag r)
194 mapBagM f (TwoBags b1 b2) = do r1 <- mapBagM f b1
195 r2 <- mapBagM f b2
196 return (TwoBags r1 r2)
197 mapBagM f (ListBag xs) = do rs <- mapM f xs
198 return (ListBag rs)
199
200 mapBagM_ :: Monad m => (a -> m b) -> Bag a -> m ()
201 mapBagM_ _ EmptyBag = return ()
202 mapBagM_ f (UnitBag x) = f x >> return ()
203 mapBagM_ f (TwoBags b1 b2) = mapBagM_ f b1 >> mapBagM_ f b2
204 mapBagM_ f (ListBag xs) = mapM_ f xs
205
206 flatMapBagM :: Monad m => (a -> m (Bag b)) -> Bag a -> m (Bag b)
207 flatMapBagM _ EmptyBag = return EmptyBag
208 flatMapBagM f (UnitBag x) = f x
209 flatMapBagM f (TwoBags b1 b2) = do r1 <- flatMapBagM f b1
210 r2 <- flatMapBagM f b2
211 return (r1 `unionBags` r2)
212 flatMapBagM f (ListBag xs) = foldrM k EmptyBag xs
213 where
214 k x b2 = do { b1 <- f x; return (b1 `unionBags` b2) }
215
216 flatMapBagPairM :: Monad m => (a -> m (Bag b, Bag c)) -> Bag a -> m (Bag b, Bag c)
217 flatMapBagPairM _ EmptyBag = return (EmptyBag, EmptyBag)
218 flatMapBagPairM f (UnitBag x) = f x
219 flatMapBagPairM f (TwoBags b1 b2) = do (r1,s1) <- flatMapBagPairM f b1
220 (r2,s2) <- flatMapBagPairM f b2
221 return (r1 `unionBags` r2, s1 `unionBags` s2)
222 flatMapBagPairM f (ListBag xs) = foldrM k (EmptyBag, EmptyBag) xs
223 where
224 k x (r2,s2) = do { (r1,s1) <- f x
225 ; return (r1 `unionBags` r2, s1 `unionBags` s2) }
226
227 mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
228 mapAndUnzipBagM _ EmptyBag = return (EmptyBag, EmptyBag)
229 mapAndUnzipBagM f (UnitBag x) = do (r,s) <- f x
230 return (UnitBag r, UnitBag s)
231 mapAndUnzipBagM f (TwoBags b1 b2) = do (r1,s1) <- mapAndUnzipBagM f b1
232 (r2,s2) <- mapAndUnzipBagM f b2
233 return (TwoBags r1 r2, TwoBags s1 s2)
234 mapAndUnzipBagM f (ListBag xs) = do ts <- mapM f xs
235 let (rs,ss) = unzip ts
236 return (ListBag rs, ListBag ss)
237
238 mapAccumBagLM :: Monad m
239 => (acc -> x -> m (acc, y)) -- ^ combining funcction
240 -> acc -- ^ initial state
241 -> Bag x -- ^ inputs
242 -> m (acc, Bag y) -- ^ final state, outputs
243 mapAccumBagLM _ s EmptyBag = return (s, EmptyBag)
244 mapAccumBagLM f s (UnitBag x) = do { (s1, x1) <- f s x; return (s1, UnitBag x1) }
245 mapAccumBagLM f s (TwoBags b1 b2) = do { (s1, b1') <- mapAccumBagLM f s b1
246 ; (s2, b2') <- mapAccumBagLM f s1 b2
247 ; return (s2, TwoBags b1' b2') }
248 mapAccumBagLM f s (ListBag xs) = do { (s', xs') <- mapAccumLM f s xs
249 ; return (s', ListBag xs') }
250
251 listToBag :: [a] -> Bag a
252 listToBag [] = EmptyBag
253 listToBag vs = ListBag vs
254
255 bagToList :: Bag a -> [a]
256 bagToList b = foldrBag (:) [] b
257
258 instance (Outputable a) => Outputable (Bag a) where
259 ppr bag = braces (pprWithCommas ppr (bagToList bag))
260
261 instance Data a => Data (Bag a) where
262 gfoldl k z b = z listToBag `k` bagToList b -- traverse abstract type abstractly
263 toConstr _ = abstractConstr $ "Bag("++show (typeOf (undefined::a))++")"
264 gunfold _ _ = error "gunfold"
265 dataTypeOf _ = mkNoRepType "Bag"
266 dataCast1 x = gcast1 x