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