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, concatBag, foldBag, foldrBag, foldlBag,
  16         isEmptyBag, isSingletonBag, consBag, snocBag, anyBag,
  17         listToBag, bagToList,
  18         mapBagM, mapAndUnzipBagM
  19     ) where
  20
  21 import Outputable
  22 import Util ( isSingleton )
  23
  24 import Data.List ( partition )
  25
  26 infixr 3 `consBag`
  27 infixl 3 `snocBag`
  28 \end{code}
  29
  30
  31 \begin{code}
  32 data Bag a
  33   = EmptyBag
  34   | UnitBag a
  35   | TwoBags (Bag a) (Bag a) -- INVARIANT: neither branch is empty
  36   | ListBag [a]             -- INVARIANT: the list is non-empty
  37
  38 emptyBag :: Bag a
  39 emptyBag = EmptyBag
  40
  41 unitBag :: a -> Bag a
  42 unitBag  = UnitBag
  43
  44 lengthBag :: Bag a -> Int
  45 lengthBag EmptyBag        = 0
  46 lengthBag (UnitBag {})    = 1
  47 lengthBag (TwoBags b1 b2) = lengthBag b1 + lengthBag b2
  48 lengthBag (ListBag xs)    = length xs
  49
  50 elemBag :: Eq a => a -> Bag a -> Bool
  51 elemBag _ EmptyBag        = False
  52 elemBag x (UnitBag y)     = x == y
  53 elemBag x (TwoBags b1 b2) = x `elemBag` b1 || x `elemBag` b2
  54 elemBag x (ListBag ys)    = any (x ==) ys
  55
  56 unionManyBags :: [Bag a] -> Bag a
  57 unionManyBags xs = foldr unionBags EmptyBag xs
  58
  59 -- This one is a bit stricter! The bag will get completely evaluated.
  60
  61 unionBags :: Bag a -> Bag a -> Bag a
  62 unionBags EmptyBag b = b
  63 unionBags b EmptyBag = b
  64 unionBags b1 b2      = TwoBags b1 b2
  65
  66 consBag :: a -> Bag a -> Bag a
  67 snocBag :: Bag a -> a -> Bag a
  68
  69 consBag elt bag = (unitBag elt) `unionBags` bag
  70 snocBag bag elt = bag `unionBags` (unitBag elt)
  71
  72 isEmptyBag :: Bag a -> Bool
  73 isEmptyBag EmptyBag = True
  74 isEmptyBag _        = False -- NB invariants
  75
  76 isSingletonBag :: Bag a -> Bool
  77 isSingletonBag EmptyBag      = False
  78 isSingletonBag (UnitBag _)   = True
  79 isSingletonBag (TwoBags _ _) = False          -- Neither is empty
  80 isSingletonBag (ListBag xs)  = isSingleton xs
  81
  82 filterBag :: (a -> Bool) -> Bag a -> Bag a
  83 filterBag _    EmptyBag = EmptyBag
  84 filterBag pred b@(UnitBag val) = if pred val then b else EmptyBag
  85 filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
  86     where sat1 = filterBag pred b1
  87           sat2 = filterBag pred b2
  88 filterBag pred (ListBag vs)    = listToBag (filter pred vs)
  89
  90 anyBag :: (a -> Bool) -> Bag a -> Bool
  91 anyBag _ EmptyBag        = False
  92 anyBag p (UnitBag v)     = p v
  93 anyBag p (TwoBags b1 b2) = anyBag p b1 || anyBag p b2
  94 anyBag p (ListBag xs)    = any p xs
  95
  96 concatBag :: Bag (Bag a) -> Bag a
  97 concatBag EmptyBag        = EmptyBag
  98 concatBag (UnitBag b)     = b
  99 concatBag (TwoBags b1 b2) = concatBag b1 `unionBags` concatBag b2
 100 concatBag (ListBag bs)    = unionManyBags bs
 101
 102 partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
 103                                          Bag a {- Don't -})
 104 partitionBag _    EmptyBag = (EmptyBag, EmptyBag)
 105 partitionBag pred b@(UnitBag val)
 106     = if pred val then (b, EmptyBag) else (EmptyBag, b)
 107 partitionBag pred (TwoBags b1 b2)
 108     = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
 109   where (sat1, fail1) = partitionBag pred b1
 110         (sat2, fail2) = partitionBag pred b2
 111 partitionBag pred (ListBag vs) = (listToBag sats, listToBag fails)
 112   where (sats, fails) = partition pred vs
 113
 114
 115 foldBag :: (r -> r -> r) -- Replace TwoBags with this; should be associative
 116         -> (a -> r)      -- Replace UnitBag with this
 117         -> r             -- Replace EmptyBag with this
 118         -> Bag a
 119         -> r
 120
 121 {- Standard definition
 122 foldBag t u e EmptyBag        = e
 123 foldBag t u e (UnitBag x)     = u x
 124 foldBag t u e (TwoBags b1 b2) = (foldBag t u e b1) `t` (foldBag t u e b2)
 125 foldBag t u e (ListBag xs)    = foldr (t.u) e xs
 126 -}
 127
 128 -- More tail-recursive definition, exploiting associativity of "t"
 129 foldBag _ _ e EmptyBag        = e
 130 foldBag t u e (UnitBag x)     = u x `t` e
 131 foldBag t u e (TwoBags b1 b2) = foldBag t u (foldBag t u e b2) b1
 132 foldBag t u e (ListBag xs)    = foldr (t.u) e xs
 133
 134 foldrBag :: (a -> r -> r) -> r
 135          -> Bag a
 136          -> r
 137
 138 foldrBag _ z EmptyBag        = z
 139 foldrBag k z (UnitBag x)     = k x z
 140 foldrBag k z (TwoBags b1 b2) = foldrBag k (foldrBag k z b2) b1
 141 foldrBag k z (ListBag xs)    = foldr k z xs
 142
 143 foldlBag :: (r -> a -> r) -> r
 144          -> Bag a
 145          -> r
 146
 147 foldlBag _ z EmptyBag        = z
 148 foldlBag k z (UnitBag x)     = k z x
 149 foldlBag k z (TwoBags b1 b2) = foldlBag k (foldlBag k z b1) b2
 150 foldlBag k z (ListBag xs)    = foldl k z xs
 151
 152
 153 mapBag :: (a -> b) -> Bag a -> Bag b
 154 mapBag _ EmptyBag        = EmptyBag
 155 mapBag f (UnitBag x)     = UnitBag (f x)
 156 mapBag f (TwoBags b1 b2) = TwoBags (mapBag f b1) (mapBag f b2)
 157 mapBag f (ListBag xs)    = ListBag (map f xs)
 158
 159 mapBagM :: Monad m => (a -> m b) -> Bag a -> m (Bag b)
 160 mapBagM _ EmptyBag        = return EmptyBag
 161 mapBagM f (UnitBag x)     = do r <- f x
 162                                return (UnitBag r)
 163 mapBagM f (TwoBags b1 b2) = do r1 <- mapBagM f b1
 164                                r2 <- mapBagM f b2
 165                                return (TwoBags r1 r2)
 166 mapBagM f (ListBag    xs) = do rs <- mapM f xs
 167                                return (ListBag rs)
 168
 169 mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
 170 mapAndUnzipBagM _ EmptyBag        = return (EmptyBag, EmptyBag)
 171 mapAndUnzipBagM f (UnitBag x)     = do (r,s) <- f x
 172                                        return (UnitBag r, UnitBag s)
 173 mapAndUnzipBagM f (TwoBags b1 b2) = do (r1,s1) <- mapAndUnzipBagM f b1
 174                                        (r2,s2) <- mapAndUnzipBagM f b2
 175                                        return (TwoBags r1 r2, TwoBags s1 s2)
 176 mapAndUnzipBagM f (ListBag xs)    = do ts <- mapM f xs
 177                                        let (rs,ss) = unzip ts
 178                                        return (ListBag rs, ListBag ss)
 179
 180 listToBag :: [a] -> Bag a
 181 listToBag [] = EmptyBag
 182 listToBag vs = ListBag vs
 183
 184 bagToList :: Bag a -> [a]
 185 bagToList b = foldrBag (:) [] b
 186 \end{code}
 187
 188 \begin{code}
 189 instance (Outputable a) => Outputable (Bag a) where
 190     ppr bag = braces (pprWithCommas ppr (bagToList bag))
 191 \end{code}