[project @ 2005-03-09 16:58:15 by simonpj]
[ghc-hetmet.git] / ghc / compiler / utils / Bag.lhs
index 3734df5..4ee8b0f 100644 (file)
@@ -1,5 +1,5 @@
 %
-% (c) The GRASP/AQUA Project, Glasgow University, 1992-1995
+% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
 %
 \section[Bags]{@Bag@: an unordered collection with duplicates}
 
@@ -8,56 +8,63 @@ module Bag (
        Bag,    -- abstract type
 
        emptyBag, unitBag, unionBags, unionManyBags,
-#if ! defined(COMPILING_GHC)
+       mapBag,
        elemBag,
-#endif
-       filterBag, partitionBag,
-       isEmptyBag, snocBag, listToBag, bagToList
+       filterBag, partitionBag, concatBag, foldBag, foldrBag, foldlBag,
+       isEmptyBag, isSingletonBag, consBag, snocBag,
+       listToBag, bagToList, 
+       mapBagM, mapAndUnzipBagM
     ) where
 
-#if defined(COMPILING_GHC)
-import Id              ( Id )
+#include "HsVersions.h"
+
 import Outputable
-import Pretty
-import Util
-#endif
+import Util            ( isSingleton )
+import List            ( partition )
+\end{code}
+
 
+\begin{code}
 data Bag a
   = EmptyBag
   | UnitBag    a
-  | TwoBags    (Bag a) (Bag a) -- The ADT guarantees that at least
-                               -- one branch is non-empty.
-  | ListOfBags [Bag a]         -- The list is non-empty
+  | TwoBags    (Bag a) (Bag a) -- INVARIANT: neither branch is empty
+  | ListBag    [a]             -- INVARIANT: the list is non-empty
 
 emptyBag = EmptyBag
 unitBag  = UnitBag
 
-#if ! defined(COMPILING_GHC)
--- not used in GHC
 elemBag :: Eq a => a -> Bag a -> Bool
+
 elemBag x EmptyBag        = False
 elemBag x (UnitBag y)     = x==y
 elemBag x (TwoBags b1 b2) = x `elemBag` b1 || x `elemBag` b2
-elemBag x (ListOfBags bs) = any (x `elemBag`) bs
-#endif
+elemBag x (ListBag ys)    = any (x ==) ys
 
-unionManyBags [] = EmptyBag
-unionManyBags xs = ListOfBags xs
+unionManyBags :: [Bag a] -> Bag a
+unionManyBags xs = foldr unionBags EmptyBag xs
 
 -- This one is a bit stricter! The bag will get completely evaluated.
 
-
+unionBags :: Bag a -> Bag a -> Bag a
 unionBags EmptyBag b = b
 unionBags b EmptyBag = b
 unionBags b1 b2      = TwoBags b1 b2
 
+consBag :: a -> Bag a -> Bag a
 snocBag :: Bag a -> a -> Bag a
+
+consBag elt bag = (unitBag elt) `unionBags` bag
 snocBag bag elt = bag `unionBags` (unitBag elt)
 
-isEmptyBag EmptyBag        = True
-isEmptyBag (TwoBags b1 b2)  = isEmptyBag b1 && isEmptyBag b2   -- Paranoid, but safe
-isEmptyBag (ListOfBags bs)  = all isEmptyBag bs
-isEmptyBag other           = False
+isEmptyBag EmptyBag = True
+isEmptyBag other    = False    -- NB invariants
+
+isSingletonBag :: Bag a -> Bool
+isSingletonBag EmptyBag                = False
+isSingletonBag (UnitBag x)     = True
+isSingletonBag (TwoBags b1 b2)  = False                -- Neither is empty
+isSingletonBag (ListBag xs)     = isSingleton xs
 
 filterBag :: (a -> Bool) -> Bag a -> Bag a
 filterBag pred EmptyBag = EmptyBag
@@ -66,12 +73,15 @@ filterBag pred (TwoBags b1 b2) = sat1 `unionBags` sat2
                               where
                                 sat1 = filterBag pred b1
                                 sat2 = filterBag pred b2
-filterBag pred (ListOfBags bs) = ListOfBags sats
-                               where
-                                sats = [filterBag pred b | b <- bs]
+filterBag pred (ListBag vs)    = listToBag (filter pred vs)
 
+concatBag :: Bag (Bag a) -> Bag a
+concatBag EmptyBag         = EmptyBag
+concatBag (UnitBag b)       = b
+concatBag (TwoBags b1 b2)   = concatBag b1 `unionBags` concatBag b2
+concatBag (ListBag bs)     = unionManyBags bs
 
-partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -}, 
+partitionBag :: (a -> Bool) -> Bag a -> (Bag a {- Satisfy predictate -},
                                         Bag a {- Don't -})
 partitionBag pred EmptyBag = (EmptyBag, EmptyBag)
 partitionBag pred b@(UnitBag val) = if pred val then (b, EmptyBag) else (EmptyBag, b)
@@ -79,32 +89,83 @@ partitionBag pred (TwoBags b1 b2) = (sat1 `unionBags` sat2, fail1 `unionBags` fa
                                  where
                                    (sat1,fail1) = partitionBag pred b1
                                    (sat2,fail2) = partitionBag pred b2
-partitionBag pred (ListOfBags bs) = (ListOfBags sats, ListOfBags fails)
+partitionBag pred (ListBag vs)   = (listToBag sats, listToBag fails)
                                  where
-                                   (sats, fails) = unzip [partitionBag pred b | b <- bs]
-
+                                   (sats,fails) = partition pred vs
+
+
+foldBag :: (r -> r -> r)       -- Replace TwoBags with this; should be associative
+       -> (a -> r)             -- Replace UnitBag with this
+       -> r                    -- Replace EmptyBag with this
+       -> Bag a
+       -> r
+
+{- Standard definition
+foldBag t u e EmptyBag        = e
+foldBag t u e (UnitBag x)     = u x
+foldBag t u e (TwoBags b1 b2) = (foldBag t u e b1) `t` (foldBag t u e b2)
+foldBag t u e (ListBag xs)    = foldr (t.u) e xs
+-}
+
+-- More tail-recursive definition, exploiting associativity of "t"
+foldBag t u e EmptyBag        = e
+foldBag t u e (UnitBag x)     = u x `t` e
+foldBag t u e (TwoBags b1 b2) = foldBag t u (foldBag t u e b2) b1
+foldBag t u e (ListBag xs)    = foldr (t.u) e xs
+
+foldrBag :: (a -> r -> r) -> r
+        -> Bag a
+        -> r
+
+foldrBag k z EmptyBag        = z
+foldrBag k z (UnitBag x)     = k x z
+foldrBag k z (TwoBags b1 b2) = foldrBag k (foldrBag k z b2) b1
+foldrBag k z (ListBag xs)    = foldr k z xs
+
+foldlBag :: (r -> a -> r) -> r
+        -> Bag a
+        -> r
+
+foldlBag k z EmptyBag        = z
+foldlBag k z (UnitBag x)     = k z x
+foldlBag k z (TwoBags b1 b2) = foldlBag k (foldlBag k z b1) b2
+foldlBag k z (ListBag xs)    = foldl k z xs
+
+
+mapBag :: (a -> b) -> Bag a -> Bag b
+mapBag f EmptyBag       = EmptyBag
+mapBag f (UnitBag x)     = UnitBag (f x)
+mapBag f (TwoBags b1 b2) = TwoBags (mapBag f b1) (mapBag f b2) 
+mapBag f (ListBag xs)    = ListBag (map f xs)
+
+mapBagM :: Monad m => (a -> m b) -> Bag a -> m (Bag b)
+mapBagM f EmptyBag       = return EmptyBag
+mapBagM f (UnitBag x)     = do { r <- f x; return (UnitBag r) }
+mapBagM f (TwoBags b1 b2) = do { r1 <- mapBagM f b1; r2 <- mapBagM f b2; return (TwoBags r1 r2) }
+mapBagM f (ListBag    xs) = do { rs <- mapM    f xs; return (ListBag rs) }
+
+mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
+mapAndUnzipBagM f EmptyBag       = return (EmptyBag, EmptyBag)
+mapAndUnzipBagM f (UnitBag x)     = do { (r,s) <- f x; return (UnitBag r, UnitBag s) }
+mapAndUnzipBagM f (TwoBags b1 b2) = do { (r1,s1) <- mapAndUnzipBagM f b1
+                                       ; (r2,s2) <- mapAndUnzipBagM f b2
+                                       ; return (TwoBags r1 r2, TwoBags s1 s2) }
+mapAndUnzipBagM f (ListBag    xs) = do { ts <- mapM f xs
+                                       ; let (rs,ss) = unzip ts
+                                       ; return (ListBag rs, ListBag ss) }
 
 listToBag :: [a] -> Bag a
-listToBag lst = foldr TwoBags EmptyBag (map UnitBag lst)
+listToBag [] = EmptyBag
+listToBag vs = ListBag vs
 
 bagToList :: Bag a -> [a]
-bagToList b = b_to_l b []
-  where
-    -- (b_to_l b xs) flattens b and puts xs on the end.
-    b_to_l EmptyBag       xs = xs
-    b_to_l (UnitBag x)    xs = x:xs
-    b_to_l (TwoBags b1 b2) xs = b_to_l b1 (b_to_l b2 xs)
-    b_to_l (ListOfBags bs) xs = foldr b_to_l xs bs 
+bagToList b = foldrBag (:) [] b
 \end{code}
 
 \begin{code}
-#if defined(COMPILING_GHC)
-
 instance (Outputable a) => Outputable (Bag a) where
-    ppr sty EmptyBag       = ppStr "emptyBag"
-    ppr sty (UnitBag a)     = ppr sty a
-    ppr sty (TwoBags b1 b2) = ppCat [ppr sty b1, pp'SP, ppr sty b2]
-    ppr sty (ListOfBags bs) = ppCat [ppLbrack, interpp'SP sty bs, ppRbrack]
-
-#endif {- COMPILING_GHC -}
+    ppr EmptyBag       = ptext SLIT("emptyBag")
+    ppr (UnitBag a)     = ppr a
+    ppr (TwoBags b1 b2) = hsep [ppr b1 <> comma, ppr b2]
+    ppr (ListBag as)    = interpp'SP as
 \end{code}