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         foldlBagM, mapBagM, mapBagM_,
  20         flatMapBagM, flatMapBagPairM,
  21         mapAndUnzipBagM
  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 foldlBagM :: (Monad m) => (b -> a -> m b) -> b -> Bag a -> m b
 175 foldlBagM _ z EmptyBag        = return z
 176 foldlBagM k z (UnitBag x)     = k z x
 177 foldlBagM k z (TwoBags b1 b2) = do { z' <- foldlBagM k z b1; foldlBagM k z' b2 }
 178 foldlBagM k z (ListBag xs)    = foldlM k z xs
 179
 180 mapBag :: (a -> b) -> Bag a -> Bag b
 181 mapBag _ EmptyBag        = EmptyBag
 182 mapBag f (UnitBag x)     = UnitBag (f x)
 183 mapBag f (TwoBags b1 b2) = TwoBags (mapBag f b1) (mapBag f b2)
 184 mapBag f (ListBag xs)    = ListBag (map f xs)
 185
 186 mapBagM :: Monad m => (a -> m b) -> Bag a -> m (Bag b)
 187 mapBagM _ EmptyBag        = return EmptyBag
 188 mapBagM f (UnitBag x)     = do r <- f x
 189                                return (UnitBag r)
 190 mapBagM f (TwoBags b1 b2) = do r1 <- mapBagM f b1
 191                                r2 <- mapBagM f b2
 192                                return (TwoBags r1 r2)
 193 mapBagM f (ListBag    xs) = do rs <- mapM f xs
 194                                return (ListBag rs)
 195
 196 mapBagM_ :: Monad m => (a -> m b) -> Bag a -> m ()
 197 mapBagM_ _ EmptyBag        = return ()
 198 mapBagM_ f (UnitBag x)     = f x >> return ()
 199 mapBagM_ f (TwoBags b1 b2) = mapBagM_ f b1 >> mapBagM_ f b2
 200 mapBagM_ f (ListBag    xs) = mapM_ f xs
 201
 202 flatMapBagM :: Monad m => (a -> m (Bag b)) -> Bag a -> m (Bag b)
 203 flatMapBagM _ EmptyBag        = return EmptyBag
 204 flatMapBagM f (UnitBag x)     = f x
 205 flatMapBagM f (TwoBags b1 b2) = do r1 <- flatMapBagM f b1
 206                                    r2 <- flatMapBagM f b2
 207                                    return (r1 `unionBags` r2)
 208 flatMapBagM f (ListBag    xs) = foldrM k EmptyBag xs
 209   where
 210     k x b2 = do { b1 <- f x; return (b1 `unionBags` b2) }
 211
 212 flatMapBagPairM :: Monad m => (a -> m (Bag b, Bag c)) -> Bag a -> m (Bag b, Bag c)
 213 flatMapBagPairM _ EmptyBag        = return (EmptyBag, EmptyBag)
 214 flatMapBagPairM f (UnitBag x)     = f x
 215 flatMapBagPairM f (TwoBags b1 b2) = do (r1,s1) <- flatMapBagPairM f b1
 216                                        (r2,s2) <- flatMapBagPairM f b2
 217                                        return (r1 `unionBags` r2, s1 `unionBags` s2)
 218 flatMapBagPairM f (ListBag    xs) = foldrM k (EmptyBag, EmptyBag) xs
 219   where
 220     k x (r2,s2) = do { (r1,s1) <- f x
 221                      ; return (r1 `unionBags` r2, s1 `unionBags` s2) }
 222
 223 mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
 224 mapAndUnzipBagM _ EmptyBag        = return (EmptyBag, EmptyBag)
 225 mapAndUnzipBagM f (UnitBag x)     = do (r,s) <- f x
 226                                        return (UnitBag r, UnitBag s)
 227 mapAndUnzipBagM f (TwoBags b1 b2) = do (r1,s1) <- mapAndUnzipBagM f b1
 228                                        (r2,s2) <- mapAndUnzipBagM f b2
 229                                        return (TwoBags r1 r2, TwoBags s1 s2)
 230 mapAndUnzipBagM f (ListBag xs)    = do ts <- mapM f xs
 231                                        let (rs,ss) = unzip ts
 232                                        return (ListBag rs, ListBag ss)
 233
 234 listToBag :: [a] -> Bag a
 235 listToBag [] = EmptyBag
 236 listToBag vs = ListBag vs
 237
 238 bagToList :: Bag a -> [a]
 239 bagToList b = foldrBag (:) [] b
 240 \end{code}
 241
 242 \begin{code}
 243 instance (Outputable a) => Outputable (Bag a) where
 244     ppr bag = braces (pprWithCommas ppr (bagToList bag))
 245
 246 INSTANCE_TYPEABLE1(Bag,bagTc,"Bag")
 247
 248 instance Data a => Data (Bag a) where
 249   gfoldl k z b = z listToBag `k` bagToList b -- traverse abstract type abstractly
 250   toConstr _   = abstractConstr $ "Bag("++show (typeOf (undefined::a))++")"
 251   gunfold _ _  = error "gunfold"
 252   dataTypeOf _ = mkNoRepType "Bag"
 253 \end{code}