mapMaybe for keys in Data.Map

user_1264328 0
Viewed 168

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?

haskell
3 Answers

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 ]
user_8498910

It’s not bad with foldMapWithKey (or foldrWithKey), without going through an intermediate list.

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

mapMaybeKeys f = Map.foldMapWithKey
  (\ k a -> foldMap (\ k' -> M.singleton k' a) (f k))

-- or:
mapMaybeKeys f = Map.foldMapWithKey (flip (foldMap . flip M.singleton) . f)

-- or:
mapMaybeKeys = M.foldMapWithKey . fmap (flip (foldMap . flip M.singleton))

The inner foldMap produces an empty map if the key function returns Nothing. You could add some helper definitions to make this a simpler-looking pipeline of compositions by avoiding the flips. Note that you need the Ord constraint to use b as the keys of the resulting map.

user_2468860

Thank you @Will Ness and @Jon Purdy, I like your solutions, but with some adjustments, I think I like my original solution best:

import Data.Bitraversable (bitraverse)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (catMaybes)

mapMaybeKeys :: (Ord b) => (a -> Maybe b) -> Map a c -> Map b c
mapMaybeKeys f
  = Map.fromList
  . catMaybes
  . fmap (bitraverse f Just)
  . Map.toList
user_12643280
Related