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