mapMaybe for keys in Data.Map

Question

I want to map the following function over the keys of a Map

f :: a -> Maybe b

and discard the Nothing keys and keep the Just keys, but extracted from the Just. Just like Map.mapMaybe, but for keys

mapMaybeKeys :: (a -> Maybe b) -> Map a c -> Map b c

I searched Hoogle for this type signature but didn't find anything.

I could do this:

mapMaybeKeys f
    = Map.toList
    . catMaybes
    . fmap (fmap swap . traverse f . swap)
    . Map.toList

or:

mapMaybeKeys f
   = Map.mapKeys fromJust
   . Map.delete Nothing
   . Map.mapKeys f

Is there a more elegant way?

Answer 1

With list comprehensions,

import Data.Map (Map)
import qualified Data.Map as M
import Control.Arrow (first)

mapMaybeKeys :: Ord b => (a -> Maybe b) -> Map a c -> Map b c
mapMaybeKeys f m = 
   M.fromList [ (b,a) | (Just b, a) <- map (first f) . M.toList $ m ]