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