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     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 \end{code}