Strict versions of mapMaybe*

Author: jcpetruzza

Data/HashMap/Base.hs

@@ -84,6 +84,7 @@ module Data.HashMap.Base
     , update16M
     , update16With'
     , updateOrConcatWith
+    , filterMapAux
     ) where
 
 #if __GLASGOW_HASKELL__ >= 709
@@ -884,6 +885,7 @@ filterWithKey pred = filterMapAux onLeaf onColl
 
         onColl el@(L k v) | pred k v = Just el
         onColl _ = Nothing
+{-# INLINE filterWithKey #-}
 
 
 -- | Common implementation for 'filterWithKey' and 'mapMaybeWithKey',
@@ -953,7 +955,7 @@ filterMapAux onLeaf onColl = go
             | Just el <- onColl (A.index ary i)
                 = A.write mary j el >> step ary mary (i+1) (j+1) n
             | otherwise = step ary mary (i+1) j n
-{-# INLINE filterWithKey #-}
+{-# INLINE filterMapAux #-}
 
 -- | /O(n)/ Filter this map by retaining only elements which values
 -- satisfy a predicate.

Data/HashMap/Strict.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE BangPatterns, CPP #-}
+{-# LANGUAGE BangPatterns, CPP, PatternGuards #-}
 
 #if __GLASGOW_HASKELL__ >= 702
 {-# LANGUAGE Trustworthy #-}
@@ -97,7 +97,7 @@ import qualified Data.HashMap.Array as A
 import qualified Data.HashMap.Base as HM
 import Data.HashMap.Base hiding (
     alter, adjust, fromList, fromListWith, insert, insertWith, intersectionWith,
-    map, mapWithKey, singleton, update, unionWith)
+    map, mapWithKey, mapMaybe, mapMaybeWithKey, singleton, update, unionWith)
 import Data.HashMap.Unsafe (runST)
 
 -- $strictness
@@ -357,6 +357,28 @@ map :: (v1 -> v2) -> HashMap k v1 -> HashMap k v2
 map f = mapWithKey (const f)
 {-# INLINE map #-}
 
+
+------------------------------------------------------------------------
+-- * Filter
+
+-- | /O(n)/ Transform this map by applying a function to every value
+--   and retaining only some of them.
+mapMaybeWithKey :: (k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
+mapMaybeWithKey f = filterMapAux onLeaf onColl
+  where onLeaf (Leaf h (L k v)) | Just v' <- f k v = Just (leaf h k v')
+        onLeaf _ = Nothing
+
+        onColl (L k v) | Just v' <- f k v = Just (L k v')
+                       | otherwise = Nothing
+{-# INLINE mapMaybeWithKey #-}
+
+-- | /O(n)/ Transform this map by applying a function to every value
+--   and retaining only some of them.
+mapMaybe :: (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
+mapMaybe f = mapMaybeWithKey (const f)
+{-# INLINE mapMaybe #-}
+
+
 -- TODO: Should we add a strict traverseWithKey?
 
 ------------------------------------------------------------------------