ghc/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, consBag, snocBag,
  15         listToBag, bagToList
  16     ) where
  17
  18 #include "HsVersions.h"
  19
  20 import Outputable
  21 import List             ( partition )
  22 \end{code}
  23
  24
  25 \begin{code}
  26 data Bag a
  27   = EmptyBag
  28   | UnitBag     a
  29   | TwoBags     (Bag a) (Bag a) -- The ADT guarantees that at least
  30                                 -- one branch is non-empty
  31   | ListBag     [a]             -- The list is non-empty
  32   | ListOfBags  [Bag a]         -- 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 elemBag x (ListOfBags bs) = any (x `elemBag`) bs
  44
  45 unionManyBags [] = EmptyBag
  46 unionManyBags xs = ListOfBags xs
  47
  48 -- This one is a bit stricter! The bag will get completely evaluated.
  49
  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 (UnitBag x)      = False
  62 isEmptyBag (TwoBags b1 b2)  = isEmptyBag b1 && isEmptyBag b2    -- Paranoid, but safe
  63 isEmptyBag (ListBag xs)     = null xs                           -- Paranoid, but safe
  64 isEmptyBag (ListOfBags bs)  = all isEmptyBag bs
  65
  66 filterBag :: (a -> Bool) -> Bag a -> Bag a
  67 filterBag pred EmptyBag = EmptyBag
  68 filterBag pred b@(UnitBag val) = if pred val then b else EmptyBag
  69 filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
  70                                where
  71                                  sat1 = filterBag pred b1
  72                                  sat2 = filterBag pred b2
  73 filterBag pred (ListBag vs)    = listToBag (filter pred vs)
  74 filterBag pred (ListOfBags bs) = ListOfBags sats
  75                                 where
  76                                  sats = [filterBag pred b | b <- bs]
  77
  78 concatBag :: Bag (Bag a) -> Bag a
  79
  80 concatBag EmptyBag          = EmptyBag
  81 concatBag (UnitBag b)       = b
  82 concatBag (TwoBags b1 b2)   = concatBag b1 `TwoBags` concatBag b2
  83 concatBag (ListBag bs)      = ListOfBags bs
  84 concatBag (ListOfBags bbs)  = ListOfBags (map concatBag bbs)
  85
  86 partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
  87                                          Bag a {- Don't -})
  88 partitionBag pred EmptyBag = (EmptyBag, EmptyBag)
  89 partitionBag pred b@(UnitBag val) = if pred val then (b, EmptyBag) else (EmptyBag, b)
  90 partitionBag pred (TwoBags b1 b2) = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
  91                                   where
  92                                     (sat1,fail1) = partitionBag pred b1
  93                                     (sat2,fail2) = partitionBag pred b2
  94 partitionBag pred (ListBag vs)    = (listToBag sats, listToBag fails)
  95                                   where
  96                                     (sats,fails) = partition pred vs
  97 partitionBag pred (ListOfBags bs) = (ListOfBags sats, ListOfBags fails)
  98                                   where
  99                                     (sats, fails) = unzip [partitionBag pred b | b <- bs]
 100
 101
 102 foldBag :: (r -> r -> r)        -- Replace TwoBags with this; should be associative
 103         -> (a -> r)             -- Replace UnitBag with this
 104         -> r                    -- Replace EmptyBag with this
 105         -> Bag a
 106         -> r
 107
 108 {- Standard definition
 109 foldBag t u e EmptyBag        = e
 110 foldBag t u e (UnitBag x)     = u x
 111 foldBag t u e (TwoBags b1 b2) = (foldBag t u e b1) `t` (foldBag t u e b2)
 112 foldBag t u e (ListBag xs)    = foldr (t.u) e xs
 113 foldBag t u e (ListOfBags bs) = foldr (\b r -> foldBag e u t b `t` r) e bs
 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 foldBag t u e (ListOfBags bs) = foldr (\b r -> foldBag t u r b) e bs
 122
 123 foldrBag :: (a -> r -> r) -> r
 124          -> Bag a
 125          -> r
 126
 127 foldrBag k z EmptyBag        = z
 128 foldrBag k z (UnitBag x)     = k x z
 129 foldrBag k z (TwoBags b1 b2) = foldrBag k (foldrBag k z b2) b1
 130 foldrBag k z (ListBag xs)    = foldr k z xs
 131 foldrBag k z (ListOfBags bs) = foldr (\b r -> foldrBag k r b) z bs
 132
 133 foldlBag :: (r -> a -> r) -> r
 134          -> Bag a
 135          -> r
 136
 137 foldlBag k z EmptyBag        = z
 138 foldlBag k z (UnitBag x)     = k z x
 139 foldlBag k z (TwoBags b1 b2) = foldlBag k (foldlBag k z b1) b2
 140 foldlBag k z (ListBag xs)    = foldl k z xs
 141 foldlBag k z (ListOfBags bs) = foldl (\r b -> foldlBag k r b) z bs
 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 mapBag f (ListOfBags bs) = ListOfBags (map (mapBag f) bs)
 150
 151
 152 listToBag :: [a] -> Bag a
 153 listToBag [] = EmptyBag
 154 listToBag vs = ListBag vs
 155
 156 bagToList :: Bag a -> [a]
 157 bagToList b = foldrBag (:) [] b
 158 \end{code}
 159
 160 \begin{code}
 161 instance (Outputable a) => Outputable (Bag a) where
 162     ppr EmptyBag        = ptext SLIT("emptyBag")
 163     ppr (UnitBag a)     = ppr a
 164     ppr (TwoBags b1 b2) = hsep [ppr b1 <> comma, ppr b2]
 165     ppr (ListBag as)    = interpp'SP as
 166     ppr (ListOfBags bs) = brackets (interpp'SP bs)
 167
 168 \end{code}