ghc/compiler/utils/Bag.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1996
   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         elemBag, mapBag,
  12         filterBag, partitionBag, concatBag, foldBag,
  13         isEmptyBag, consBag, snocBag,
  14         listToBag, bagToList, bagToList_append
  15     ) where
  16
  17 #ifdef COMPILING_GHC
  18 import Ubiq{-uitous-}
  19
  20 import Outputable       ( interpp'SP )
  21 import Pretty
  22 #endif
  23
  24 data Bag a
  25   = EmptyBag
  26   | UnitBag     a
  27   | TwoBags     (Bag a) (Bag a) -- The ADT guarantees that at least
  28                                 -- one branch is non-empty
  29   | ListBag     [a]             -- The list is non-empty
  30   | ListOfBags  [Bag a]         -- The list is non-empty
  31
  32 emptyBag = EmptyBag
  33 unitBag  = UnitBag
  34
  35 elemBag :: Eq a => a -> Bag a -> Bool
  36
  37 elemBag x EmptyBag        = False
  38 elemBag x (UnitBag y)     = x==y
  39 elemBag x (TwoBags b1 b2) = x `elemBag` b1 || x `elemBag` b2
  40 elemBag x (ListBag ys)    = any (x ==) ys
  41 elemBag x (ListOfBags bs) = any (x `elemBag`) bs
  42
  43 unionManyBags [] = EmptyBag
  44 unionManyBags xs = ListOfBags xs
  45
  46 -- This one is a bit stricter! The bag will get completely evaluated.
  47
  48 unionBags EmptyBag b = b
  49 unionBags b EmptyBag = b
  50 unionBags b1 b2      = TwoBags b1 b2
  51
  52 consBag :: a -> Bag a -> Bag a
  53 consBag elt bag = (unitBag elt) `unionBags` bag
  54 snocBag :: Bag a -> a -> Bag a
  55 snocBag bag elt = bag `unionBags` (unitBag elt)
  56
  57 isEmptyBag EmptyBag         = True
  58 isEmptyBag (UnitBag x)      = False
  59 isEmptyBag (TwoBags b1 b2)  = isEmptyBag b1 && isEmptyBag b2    -- Paranoid, but safe
  60 isEmptyBag (ListBag xs)     = null xs                           -- Paranoid, but safe
  61 isEmptyBag (ListOfBags bs)  = all isEmptyBag bs
  62
  63 filterBag :: (a -> Bool) -> Bag a -> Bag a
  64 filterBag pred EmptyBag = EmptyBag
  65 filterBag pred b@(UnitBag val) = if pred val then b else EmptyBag
  66 filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
  67                                where
  68                                  sat1 = filterBag pred b1
  69                                  sat2 = filterBag pred b2
  70 filterBag pred (ListBag vs)    = listToBag (filter pred vs)
  71 filterBag pred (ListOfBags bs) = ListOfBags sats
  72                                 where
  73                                  sats = [filterBag pred b | b <- bs]
  74
  75 concatBag :: Bag (Bag a) -> Bag a
  76
  77 concatBag EmptyBag          = EmptyBag
  78 concatBag (UnitBag b)       = b
  79 concatBag (TwoBags b1 b2)   = concatBag b1 `TwoBags` concatBag b2
  80 concatBag (ListBag bs)      = ListOfBags bs
  81 concatBag (ListOfBags bbs)  = ListOfBags (map concatBag bbs)
  82
  83 partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
  84                                          Bag a {- Don't -})
  85 partitionBag pred EmptyBag = (EmptyBag, EmptyBag)
  86 partitionBag pred b@(UnitBag val) = if pred val then (b, EmptyBag) else (EmptyBag, b)
  87 partitionBag pred (TwoBags b1 b2) = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
  88                                   where
  89                                     (sat1,fail1) = partitionBag pred b1
  90                                     (sat2,fail2) = partitionBag pred b2
  91 partitionBag pred (ListBag vs)    = (listToBag sats, listToBag fails)
  92                                   where
  93                                     (sats,fails) = partition pred vs
  94 partitionBag pred (ListOfBags bs) = (ListOfBags sats, ListOfBags fails)
  95                                   where
  96                                     (sats, fails) = unzip [partitionBag pred b | b <- bs]
  97
  98
  99 foldBag :: (r -> r -> r)        -- Replace TwoBags with this; should be associative
 100         -> (a -> r)             -- Replace UnitBag with this
 101         -> r                    -- Replace EmptyBag with this
 102         -> Bag a
 103         -> r
 104
 105 {- Standard definition
 106 foldBag t u e EmptyBag        = e
 107 foldBag t u e (UnitBag x)     = u x
 108 foldBag t u e (TwoBags b1 b2) = (foldBag t u e b1) `t` (foldBag t u e b2)
 109 foldBag t u e (ListBag xs)    = foldr (t.u) e xs
 110 foldBag t u e (ListOfBags bs) = foldr (\b r -> foldBag e u t b `t` r) e bs
 111 -}
 112
 113 -- More tail-recursive definition, exploiting associativity of "t"
 114 foldBag t u e EmptyBag        = e
 115 foldBag t u e (UnitBag x)     = u x `t` e
 116 foldBag t u e (TwoBags b1 b2) = foldBag t u (foldBag t u e b2) b1
 117 foldBag t u e (ListBag xs)    = foldr (t.u) e xs
 118 foldBag t u e (ListOfBags bs) = foldr (\b r -> foldBag t u r b) e bs
 119
 120
 121 mapBag :: (a -> b) -> Bag a -> Bag b
 122 mapBag f EmptyBag        = EmptyBag
 123 mapBag f (UnitBag x)     = UnitBag (f x)
 124 mapBag f (TwoBags b1 b2) = TwoBags (mapBag f b1) (mapBag f b2)
 125 mapBag f (ListBag xs)    = ListBag (map f xs)
 126 mapBag f (ListOfBags bs) = ListOfBags (map (mapBag f) bs)
 127
 128
 129 listToBag :: [a] -> Bag a
 130 listToBag [] = EmptyBag
 131 listToBag vs = ListBag vs
 132
 133 bagToList :: Bag a -> [a]
 134 bagToList EmptyBag     = []
 135 bagToList (ListBag vs) = vs
 136 bagToList b = bagToList_append b []
 137
 138     -- (bagToList_append b xs) flattens b and puts xs on the end.
 139 bagToList_append EmptyBag        xs = xs
 140 bagToList_append (UnitBag x)     xs = x:xs
 141 bagToList_append (TwoBags b1 b2) xs = bagToList_append b1 (bagToList_append b2 xs)
 142 bagToList_append (ListBag xx)    xs = xx++xs
 143 bagToList_append (ListOfBags bs) xs = foldr bagToList_append xs bs
 144 \end{code}
 145
 146 \begin{code}
 147 #ifdef COMPILING_GHC
 148
 149 instance (Outputable a) => Outputable (Bag a) where
 150     ppr sty EmptyBag        = ppStr "emptyBag"
 151     ppr sty (UnitBag a)     = ppr sty a
 152     ppr sty (TwoBags b1 b2) = ppCat [ppr sty b1, pp'SP, ppr sty b2]
 153     ppr sty (ListBag as)    = interpp'SP sty as
 154     ppr sty (ListOfBags bs) = ppCat [ppLbrack, interpp'SP sty bs, ppRbrack]
 155
 156 #endif {- COMPILING_GHC -}
 157 \end{code}