Concise way to conditionally update map in State monad

Question

Below is the code from an answer regarding memoization, showing a memoization function used in the State monad, where the state is updated with the result of the passed function if the key is not already in the map.

type MyMemo a b = State (Map.Map a b) b

myMemo :: Ord a => (a -> MyMemo a b) -> a -> MyMemo a b
myMemo f x = do
  map <- get
  case Map.lookup x map of
    Just y  -> return y
    Nothing -> do
      y <- f x
      modify $ \map' -> Map.insert x y map'
      return y

It doesn't seem like idiomatic Haskell: it feels very imperative, with not really that much going on per line.

Is there a way to do the above, but in a more concise/functional style? I've had a look around the functions available at http://hackage.haskell.org/package/transformers-0.5.4.0/docs/Control-Monad-Trans-State-Lazy.html#v:state, but nothing really seems helpful.

Answer 1

I think your code is in functional style, but you can bit simplify it.

myMemo f x = maybe work return =<< gets (Map.lookup x)
  where
    work = do
        y <- f x
        modify $ Map.insert x y
        return y

Answer 2

This is alternative that uses mapState, as well as >>= and maybe from https://stackoverflow.com/a/44515364/1319998, that avoids all the do notation

myMemo f x = gets (Map.lookup x) >>= maybe y' return
  where
    y' = mapState (\(y, map) -> (y, Map.insert x y map)) $ f x 

Answer 3

This is an alternative that expands on https://stackoverflow.com/a/44515364/1319998, using more >>= that avoids all the do notation

myMemo :: Ord a => (a -> MyMemo a b) -> a -> MyMemo a b
myMemo f x = gets (Map.lookup x) >>= maybe y' return
  where
    y' = f x >>= \y -> state $ \map -> (y, Map.insert x y map)

Answer 4

This is an alternative that expands on https://stackoverflow.com/a/44515364/1319998, essentially de-sugaring the do-notation

myMemo f x = gets (Map.lookup x) >>= maybe y' return
  where
    y' = f x >>= \y -> modify (Map.insert x y) >> return y