Haskell Data.Map lookup AND delete at the same time

Viewed 418

I was recently using the Map type from Data.Map inside a State Monad and so I wanted to write a function, that looks up a value in the Map and also deletes it from the Map inside the State Monad. My current implementation looks like this:

lookupDelete :: (Ord k) => k -> State (Map k v) (Maybe v)
lookupDelete k = do
    m <- get
    put (M.delete k m)
    return $ M.lookup k m

While this works, it feels quite inefficient. With mutable maps in imperative languages, it is not uncommon to find delete functions, that also return the value that was deleted. I couldn't find a function for this, so I would really appreciate if someone knows one (or can explain why there is none)

2 Answers

A simple implementation is in terms of alterF:

lookupDelete :: Ord k => k -> State (Map k v) (Maybe v)
lookupDelete = state . alterF (\x -> (x, Nothing))

The x in alterF's argument is the Maybe value stored at the key given to lookupDelete. This anonymous function returns a (Maybe v, Maybe v). (,) (Maybe v) is a functor, and basically it serves as a "context" through which we can save whatever data we want from x. In this case we just save the whole x. The Nothing in the right element specifies that we want deletion. Once fully applied, alterF then gives us (Maybe v, Map k v), where the context (left element) is whatever we saved in the anonymous function and the right element is the mutated map. Then we wrap this stateful operation in state.

alterF is quite powerful: lots of operations can be built out of it simply by choosing the correct "context" functor. E.g. insert and delete come from using Identity, and lookup comes from using Const (Maybe v). A specialized function for lookupDelete is not necessary when we have alterF. One way to understand why alterF is so powerful is to recognize its type:

flip alterF k :: Functor f => (Maybe a -> f (Maybe a)) -> Map k a -> f (Map k a)

Things with types in this pattern

SomeClass f => (a -> f b) -> s -> f t

are called "optics" (when SomeClass is Functor, they're called "lenses"), and they represent how to "find" and "mutate" and "collate" "fields" inside "structures", because they let us focus on part of a structure, modify it (with the function argument), and save some information through a context (by letting us choose f). See the lens package for other uses of this pattern. (As the docs for alterF note, it's basically at from lens.)

There is no function specifically for "delete and lookup". Instead you use a more general tool: updateLookupWithKey is "lookup and update", where update can be delete or modify.

updateLookupWithKey :: Ord k => 
  (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)

lookupDelete k = do
  (ret, m) <- gets $ updateLookupWithKey (\_ _ -> Nothing) k
  put m
  pure ret
Related