Refactor SrcLoc and SrcSpan

[ghc-hetmet.git] / compiler / utils / Bag.lhs

diff --git a/compiler/utils/Bag.lhs b/compiler/utils/Bag.lhs
index 481dedf..700878a 100644 (file)
--- a/compiler/utils/Bag.lhs
+++ b/compiler/utils/Bag.lhs
@@ -11,16 +11,23 @@ module Bag (
 
         emptyBag, unitBag, unionBags, unionManyBags,
         mapBag,
-        elemBag,
-        filterBag, partitionBag, concatBag, foldBag, foldrBag, foldlBag,
+        elemBag, lengthBag,
+        filterBag, partitionBag, partitionBagWith,
+        concatBag, foldBag, foldrBag, foldlBag,
         isEmptyBag, isSingletonBag, consBag, snocBag, anyBag,
         listToBag, bagToList,
-        mapBagM, mapAndUnzipBagM
+        foldrBagM, foldlBagM, mapBagM, mapBagM_, 
+        flatMapBagM, flatMapBagPairM,
+        mapAndUnzipBagM, mapAccumBagLM
     ) where
 
+#include "Typeable.h"
+
 import Outputable
-import Util ( isSingleton )
+import Util
 
+import MonadUtils
+import Data.Data
 import Data.List ( partition )
 
 infixr 3 `consBag`
@@ -34,6 +41,7 @@ data Bag a
   | UnitBag a
   | TwoBags (Bag a) (Bag a) -- INVARIANT: neither branch is empty
   | ListBag [a]             -- INVARIANT: the list is non-empty
+    deriving Typeable
 
 emptyBag :: Bag a
 emptyBag = EmptyBag
@@ -41,6 +49,12 @@ emptyBag = EmptyBag
 unitBag :: a -> Bag a
 unitBag  = UnitBag
 
+lengthBag :: Bag a -> Int
+lengthBag EmptyBag        = 0
+lengthBag (UnitBag {})    = 1
+lengthBag (TwoBags b1 b2) = lengthBag b1 + lengthBag b2
+lengthBag (ListBag xs)    = length xs
+
 elemBag :: Eq a => a -> Bag a -> Bool
 elemBag _ EmptyBag        = False
 elemBag x (UnitBag y)     = x == y
@@ -106,6 +120,21 @@ partitionBag pred (ListBag vs) = (listToBag sats, listToBag fails)
   where (sats, fails) = partition pred vs
 
 
+partitionBagWith :: (a -> Either b c) -> Bag a 
+                    -> (Bag b {- Left  -},
+                        Bag c {- Right -})
+partitionBagWith _    EmptyBag = (EmptyBag, EmptyBag)
+partitionBagWith pred (UnitBag val)
+    = case pred val of
+         Left a  -> (UnitBag a, EmptyBag) 
+         Right b -> (EmptyBag, UnitBag b)
+partitionBagWith pred (TwoBags b1 b2)
+    = (sat1 `unionBags` sat2, fail1 `unionBags` fail2)
+  where (sat1, fail1) = partitionBagWith pred b1
+        (sat2, fail2) = partitionBagWith pred b2
+partitionBagWith pred (ListBag vs) = (listToBag sats, listToBag fails)
+  where (sats, fails) = partitionWith pred vs
+
 foldBag :: (r -> r -> r) -- Replace TwoBags with this; should be associative
         -> (a -> r)      -- Replace UnitBag with this
         -> r             -- Replace EmptyBag with this
@@ -143,6 +172,17 @@ 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
 
+foldrBagM :: (Monad m) => (a -> b -> m b) -> b -> Bag a -> m b
+foldrBagM _ z EmptyBag        = return z
+foldrBagM k z (UnitBag x)     = k x z
+foldrBagM k z (TwoBags b1 b2) = do { z' <- foldrBagM k z b2; foldrBagM k z' b1 }
+foldrBagM k z (ListBag xs)    = foldrM k z xs
+
+foldlBagM :: (Monad m) => (b -> a -> m b) -> b -> Bag a -> m b
+foldlBagM _ z EmptyBag        = return z
+foldlBagM k z (UnitBag x)     = k z x
+foldlBagM k z (TwoBags b1 b2) = do { z' <- foldlBagM k z b1; foldlBagM k z' b2 }
+foldlBagM k z (ListBag xs)    = foldlM k z xs
 
 mapBag :: (a -> b) -> Bag a -> Bag b
 mapBag _ EmptyBag        = EmptyBag
@@ -160,6 +200,33 @@ mapBagM f (TwoBags b1 b2) = do r1 <- mapBagM f b1
 mapBagM f (ListBag    xs) = do rs <- mapM f xs
                                return (ListBag rs)
 
+mapBagM_ :: Monad m => (a -> m b) -> Bag a -> m ()
+mapBagM_ _ EmptyBag        = return ()
+mapBagM_ f (UnitBag x)     = f x >> return ()
+mapBagM_ f (TwoBags b1 b2) = mapBagM_ f b1 >> mapBagM_ f b2
+mapBagM_ f (ListBag    xs) = mapM_ f xs
+
+flatMapBagM :: Monad m => (a -> m (Bag b)) -> Bag a -> m (Bag b)
+flatMapBagM _ EmptyBag        = return EmptyBag
+flatMapBagM f (UnitBag x)     = f x
+flatMapBagM f (TwoBags b1 b2) = do r1 <- flatMapBagM f b1
+                                   r2 <- flatMapBagM f b2
+                                   return (r1 `unionBags` r2)
+flatMapBagM f (ListBag    xs) = foldrM k EmptyBag xs
+  where
+    k x b2 = do { b1 <- f x; return (b1 `unionBags` b2) }
+
+flatMapBagPairM :: Monad m => (a -> m (Bag b, Bag c)) -> Bag a -> m (Bag b, Bag c)
+flatMapBagPairM _ EmptyBag        = return (EmptyBag, EmptyBag)
+flatMapBagPairM f (UnitBag x)     = f x
+flatMapBagPairM f (TwoBags b1 b2) = do (r1,s1) <- flatMapBagPairM f b1
+                                       (r2,s2) <- flatMapBagPairM f b2
+                                       return (r1 `unionBags` r2, s1 `unionBags` s2)
+flatMapBagPairM f (ListBag    xs) = foldrM k (EmptyBag, EmptyBag) xs
+  where
+    k x (r2,s2) = do { (r1,s1) <- f x
+                     ; return (r1 `unionBags` r2, s1 `unionBags` s2) }
+
 mapAndUnzipBagM :: Monad m => (a -> m (b,c)) -> Bag a -> m (Bag b, Bag c)
 mapAndUnzipBagM _ EmptyBag        = return (EmptyBag, EmptyBag)
 mapAndUnzipBagM f (UnitBag x)     = do (r,s) <- f x
@@ -171,6 +238,19 @@ mapAndUnzipBagM f (ListBag xs)    = do ts <- mapM f xs
                                        let (rs,ss) = unzip ts
                                        return (ListBag rs, ListBag ss)
 
+mapAccumBagLM :: Monad m
+            => (acc -> x -> m (acc, y)) -- ^ combining funcction
+            -> acc                      -- ^ initial state
+            -> Bag x                    -- ^ inputs
+            -> m (acc, Bag y)           -- ^ final state, outputs
+mapAccumBagLM _ s EmptyBag        = return (s, EmptyBag)
+mapAccumBagLM f s (UnitBag x)     = do { (s1, x1) <- f s x; return (s1, UnitBag x1) }
+mapAccumBagLM f s (TwoBags b1 b2) = do { (s1, b1') <- mapAccumBagLM f s  b1
+                                       ; (s2, b2') <- mapAccumBagLM f s1 b2
+                                       ; return (s2, TwoBags b1' b2') }
+mapAccumBagLM f s (ListBag xs)    = do { (s', xs') <- mapAccumLM f s xs
+                                       ; return (s', ListBag xs') }
+
 listToBag :: [a] -> Bag a
 listToBag [] = EmptyBag
 listToBag vs = ListBag vs
@@ -182,4 +262,10 @@ bagToList b = foldrBag (:) [] b
 \begin{code}
 instance (Outputable a) => Outputable (Bag a) where
     ppr bag = braces (pprWithCommas ppr (bagToList bag))
+
+instance Data a => Data (Bag a) where
+  gfoldl k z b = z listToBag `k` bagToList b -- traverse abstract type abstractly
+  toConstr _   = abstractConstr $ "Bag("++show (typeOf (undefined::a))++")"
+  gunfold _ _  = error "gunfold"
+  dataTypeOf _ = mkNoRepType "Bag"
 \end{code}