summaryrefslogtreecommitdiff
path: root/compiler/GHC/Data/Bag.hs
diff options
context:
space:
mode:
Diffstat (limited to 'compiler/GHC/Data/Bag.hs')
-rw-r--r--compiler/GHC/Data/Bag.hs13
1 files changed, 12 insertions, 1 deletions
diff --git a/compiler/GHC/Data/Bag.hs b/compiler/GHC/Data/Bag.hs
index 5ace42ba13..a9b8a669de 100644
--- a/compiler/GHC/Data/Bag.hs
+++ b/compiler/GHC/Data/Bag.hs
@@ -19,7 +19,7 @@ module GHC.Data.Bag (
isEmptyBag, isSingletonBag, consBag, snocBag, anyBag, allBag,
listToBag, nonEmptyToBag, bagToList, headMaybe, mapAccumBagL,
concatMapBag, concatMapBagPair, mapMaybeBag, unzipBag,
- mapBagM, mapBagM_,
+ mapBagM, mapBagM_, lookupBag,
flatMapBagM, flatMapBagPairM,
mapAndUnzipBagM, mapAccumBagLM,
anyBagM, filterBagM
@@ -38,6 +38,7 @@ import Data.List ( partition, mapAccumL )
import Data.List.NonEmpty ( NonEmpty(..) )
import qualified Data.List.NonEmpty as NE
import qualified Data.Semigroup ( (<>) )
+import Control.Applicative( Alternative( (<|>) ) )
infixr 3 `consBag`
infixl 3 `snocBag`
@@ -115,6 +116,16 @@ filterBagM pred (ListBag vs) = do
sat <- filterM pred (toList vs)
return (listToBag sat)
+lookupBag :: Eq a => a -> Bag (a,b) -> Maybe b
+lookupBag _ EmptyBag = Nothing
+lookupBag k (UnitBag kv) = lookup_one k kv
+lookupBag k (TwoBags b1 b2) = lookupBag k b1 <|> lookupBag k b2
+lookupBag k (ListBag xs) = foldr ((<|>) . lookup_one k) Nothing xs
+
+lookup_one :: Eq a => a -> (a,b) -> Maybe b
+lookup_one k (k',v) | k==k' = Just v
+ | otherwise = Nothing
+
allBag :: (a -> Bool) -> Bag a -> Bool
allBag _ EmptyBag = True
allBag p (UnitBag v) = p v