`filterM` for containers like `Data.Map.Map`, or `Data.Set.Set`

Question

In short: How would you filter elements of a Map, or Set on a monadic predicate in Haskell?

I could think of two possible ways:

a) Round-trip through a list and filterM (probably not very efficient):

filterMapM1 :: (Monad m, Ord k) => (v -> m Bool) -> M.Map k v -> m (M.Map k v)
filterMapM1 f m = liftM M.fromList $ filterM (f.snd) $ M.toList m

b) If the predicate is not inherently monadic, but e.g. is a comparison to the state in a State monad; then we can use Data.Map.filter (quite a special case):

filterMapM2 :: (Monad m, Ord k) => (v -> v -> Bool) -> M.Map k v -> StateT v m (M.Map k v)
filterMapM2 f m = do
    s <- get
    return $ M.filter (f s) m

Is there a better way to do this?

---

Here's a little example program for demonstration.

import Control.Applicative
import Control.Monad
import Control.Monad.State
import qualified Data.Map as M

-- | filterM for M.Map. (Round trip through a list and filterM)
filterMapM1 :: (Monad m, Ord k) => (v -> m Bool) -> M.Map k v -> m (M.Map k v)
filterMapM1 f m = liftM M.fromList $ filterM (f.snd) $ M.toList m

-- | filterM for M.Map. (Uses M.filter to filter on comparison to state)
filterMapM2 :: (Monad m, Ord k) => (v -> v -> Bool) -> M.Map k v -> StateT v m (M.Map k v)
filterMapM2 f m = do
    s <- get
    return $ M.filter (f s) m

-- | Inherently monadic predicate: Result depends on user-input.
askUser :: Int -> IO Bool
askUser n = do
    liftIO $ putStrLn $ "Do you like the number " ++ show n ++ "?"
    liftIO $ (=="yes") <$> getLine

main :: IO ()
main = do
    let m = M.fromList $ take 6 $ zip ['a'..] [1..]
    -- Use inherently monadic predicate
    print =<< filterMapM1 askUser m
    -- Compare to state
    (`evalStateT` 4) $ do
        filt2 <- filterMapM2 (/=) m
        liftIO $ print filt2

---

Update: I did a benchmark between different implementations of filterMapM. It turns out that the round-trip through a list is actually pretty good. Surprisingly, it did even better than an implementation right on the internal structure of Map. The code and data are available here.

Answer 1

The two approaches have quite different semantics and implications for efficiency.

When you round-trip through a list you are allowing the filtering to be affected by all of the previous comparisons, and so the final result could be affected by the order in which the elements are visited. In the second case, however, the filtering is a pure function, so the answer will be the same no matter in what order the comparisons are made.

For instance, the user answering the questions might want to keep the number of even and odd numbers roughly the same, and thus whether the user likes a particular number will depend on all of the numbers that have been presented before.

On the other hand, here is the code for M.filter:

-- | /O(n)/. Filter all elements that satisfy the predicate.
filter :: (a -> Bool) -> Set a -> Set a
filter _ Tip = Tip
filter p (Bin _ x l r)
    | p x       = link x (filter p l) (filter p r)
    | otherwise = merge (filter p l) (filter p r)

The important thing is note that shape of the code - the resulting tree structure is greatly influenced by the structure of the original tree. Perhaps only a small amount of rebalancing will be needed. This could be a lot more efficient than rebuilding the tree from zero knowledge using M.fromList.

The upshot is that in the filterM1 case you should be concerned about the order in which the comparisons are made. Perhaps M.toList gives an acceptable ordering, or perhaps you want reverse . M.toList, or ...

In the second case you don't need to care so you can let M.filter do all of the work and take advantage of its knowledge of the data structure.

Update: I just noticed the M.toAscList and M.fromAscList functions, so perhaps this version of filterMapM1 is slightly more efficient:

filterMapM1 f m = liftM M.fromAscList $ filterM (f.snd) $ M.toAscList m

Answer 2

Currently I don't think there is a more general abstraction that would directly support filterM. Traversable doesn't, because traverse can't change the shape of the structure. I think it is technically possible to do this with lens, however the documentation suggests that you really shouldn't do so and I think it round-trips through another structure anyway.

You could use something like this (uncompilable because filter is not a class member):

filterM :: (Traversable c, Monad m, Applicative m) => (a -> m Bool) -> c a -> m (c a)
filterM p = fmap snd . filter fst . traverse p'
  where p' x = (,x) <$> p x

whether this is more or less efficient than round-tripping through a list probably depends on the structure. It's arguably less clean.