ghc/compiler/utils/Bag.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1995
   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 #if ! defined(COMPILING_GHC)
  12         elemBag,
  13 #endif
  14         filterBag, partitionBag,
  15         isEmptyBag, snocBag, listToBag, bagToList
  16     ) where
  17
  18 #if defined(COMPILING_GHC)
  19 import Id               ( Id )
  20 import Outputable
  21 import Pretty
  22 import Util
  23 #endif
  24
  25 data Bag a
  26   = EmptyBag
  27   | UnitBag     a
  28   | TwoBags     (Bag a) (Bag a) -- The ADT guarantees that at least
  29                                 -- one branch is non-empty.
  30   | ListOfBags  [Bag a]         -- The list is non-empty
  31
  32 emptyBag = EmptyBag
  33 unitBag  = UnitBag
  34
  35 #if ! defined(COMPILING_GHC)
  36 -- not used in GHC
  37 elemBag :: Eq a => a -> Bag a -> Bool
  38 elemBag x EmptyBag        = False
  39 elemBag x (UnitBag y)     = x==y
  40 elemBag x (TwoBags b1 b2) = x `elemBag` b1 || x `elemBag` b2
  41 elemBag x (ListOfBags bs) = any (x `elemBag`) bs
  42 #endif
  43
  44 unionManyBags [] = EmptyBag
  45 unionManyBags xs = ListOfBags xs
  46
  47 -- This one is a bit stricter! The bag will get completely evaluated.
  48
  49
  50 unionBags EmptyBag b = b
  51 unionBags b EmptyBag = b
  52 unionBags b1 b2      = TwoBags b1 b2
  53
  54 snocBag :: Bag a -> a -> Bag a
  55 snocBag bag elt = bag `unionBags` (unitBag elt)
  56
  57 isEmptyBag EmptyBag         = True
  58 isEmptyBag (TwoBags b1 b2)  = isEmptyBag b1 && isEmptyBag b2    -- Paranoid, but safe
  59 isEmptyBag (ListOfBags bs)  = all isEmptyBag bs
  60 isEmptyBag other            = False
  61
  62 filterBag :: (a -> Bool) -> Bag a -> Bag a
  63 filterBag pred EmptyBag = EmptyBag
  64 filterBag pred b@(UnitBag val) = if pred val then b else EmptyBag
  65 filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
  66                                where
  67                                  sat1 = filterBag pred b1
  68                                  sat2 = filterBag pred b2
  69 filterBag pred (ListOfBags bs) = ListOfBags sats
  70                                 where
  71                                  sats = [filterBag pred b | b <- bs]
  72
  73
  74 partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
  75                                          Bag a {- Don't -})
  76 partitionBag pred EmptyBag = (EmptyBag, EmptyBag)
  77 partitionBag pred b@(UnitBag val) = if pred val then (b, EmptyBag) else (EmptyBag, b)
  78 partitionBag pred (TwoBags b1 b2) = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
  79                                   where
  80                                     (sat1,fail1) = partitionBag pred b1
  81                                     (sat2,fail2) = partitionBag pred b2
  82 partitionBag pred (ListOfBags bs) = (ListOfBags sats, ListOfBags fails)
  83                                   where
  84                                     (sats, fails) = unzip [partitionBag pred b | b <- bs]
  85
  86
  87 listToBag :: [a] -> Bag a
  88 listToBag lst = foldr TwoBags EmptyBag (map UnitBag lst)
  89
  90 bagToList :: Bag a -> [a]
  91 bagToList b = b_to_l b []
  92   where
  93     -- (b_to_l b xs) flattens b and puts xs on the end.
  94     b_to_l EmptyBag        xs = xs
  95     b_to_l (UnitBag x)     xs = x:xs
  96     b_to_l (TwoBags b1 b2) xs = b_to_l b1 (b_to_l b2 xs)
  97     b_to_l (ListOfBags bs) xs = foldr b_to_l xs bs
  98 \end{code}
  99
 100 \begin{code}
 101 #if defined(COMPILING_GHC)
 102
 103 instance (Outputable a) => Outputable (Bag a) where
 104     ppr sty EmptyBag        = ppStr "emptyBag"
 105     ppr sty (UnitBag a)     = ppr sty a
 106     ppr sty (TwoBags b1 b2) = ppCat [ppr sty b1, pp'SP, ppr sty b2]
 107     ppr sty (ListOfBags bs) = ppCat [ppLbrack, interpp'SP sty bs, ppRbrack]
 108
 109 #endif {- COMPILING_GHC -}
 110 \end{code}