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 #include "HsVersions.h"
  22
  23 import Outputable
  24 import Util             ( isSingleton )
  25
  26 import Data.List        ( partition )
  27 \end{code}
  28
  29
  30 \begin{code}
  31 data Bag a
  32   = EmptyBag
  33   | UnitBag     a
  34   | TwoBags     (Bag a) (Bag a) -- INVARIANT: neither branch is empty
  35   | ListBag     [a]             -- INVARIANT: the list is non-empty
  36
  37 emptyBag = EmptyBag
  38 unitBag  = UnitBag
  39
  40 elemBag :: Eq a => a -> Bag a -> Bool
  41
  42 elemBag x 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 EmptyBag = True
  64 isEmptyBag other    = False     -- NB invariants
  65
  66 isSingletonBag :: Bag a -> Bool
  67 isSingletonBag EmptyBag         = False
  68 isSingletonBag (UnitBag x)      = True
  69 isSingletonBag (TwoBags b1 b2)  = False         -- Neither is empty
  70 isSingletonBag (ListBag xs)     = isSingleton xs
  71
  72 filterBag :: (a -> Bool) -> Bag a -> Bag a
  73 filterBag pred EmptyBag = EmptyBag
  74 filterBag pred b@(UnitBag val) = if pred val then b else EmptyBag
  75 filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
  76                                where
  77                                  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 p 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 pred EmptyBag = (EmptyBag, EmptyBag)
  96 partitionBag pred b@(UnitBag val) = if pred val then (b, EmptyBag) else (EmptyBag, b)
  97 partitionBag pred (TwoBags b1 b2) = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
  98                                   where
  99                                     (sat1,fail1) = partitionBag pred b1
 100                                     (sat2,fail2) = partitionBag pred b2
 101 partitionBag pred (ListBag vs)    = (listToBag sats, listToBag fails)
 102                                   where
 103                                     (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 t u 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 k 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 k 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 f 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 f EmptyBag        = return EmptyBag
 152 mapBagM f (UnitBag x)     = do { r <- f x; return (UnitBag r) }
 153 mapBagM f (TwoBags b1 b2) = do { r1 <- mapBagM f b1; r2 <- mapBagM f b2; return (TwoBags r1 r2) }
 154 mapBagM f (ListBag    xs) = do { rs <- mapM    f xs; return (ListBag rs) }
 155
 156 mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
 157 mapAndUnzipBagM f EmptyBag        = return (EmptyBag, EmptyBag)
 158 mapAndUnzipBagM f (UnitBag x)     = do { (r,s) <- f x; return (UnitBag r, UnitBag s) }
 159 mapAndUnzipBagM f (TwoBags b1 b2) = do  { (r1,s1) <- mapAndUnzipBagM f b1
 160                                         ; (r2,s2) <- mapAndUnzipBagM f b2
 161                                         ; return (TwoBags r1 r2, TwoBags s1 s2) }
 162 mapAndUnzipBagM f (ListBag    xs) = do  { ts <- mapM f xs
 163                                         ; let (rs,ss) = unzip ts
 164                                         ; return (ListBag rs, ListBag ss) }
 165
 166 listToBag :: [a] -> Bag a
 167 listToBag [] = EmptyBag
 168 listToBag vs = ListBag vs
 169
 170 bagToList :: Bag a -> [a]
 171 bagToList b = foldrBag (:) [] b
 172 \end{code}
 173
 174 \begin{code}
 175 instance (Outputable a) => Outputable (Bag a) where
 176     ppr bag = char '<' <> pprWithCommas ppr (bagToList bag) <> char '>'
 177 \end{code}