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