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