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