compiler/utils/Bag.lhs

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