Haskell memoization monadically

Memo.hs

{-# LANGUAGE FlexibleContexts #-}

import           Control.Monad.Identity
import           Control.Monad.State.Lazy
import           Data.Map
import           Prelude                  hiding (lookup)

memoize :: (MonadState (Map k a) m, Ord k) =>
             (k -> m a) -> k -> m a
memoize f x = do
  v <- gets $ lookup x
  case v of
    Just y -> return y
    _      -> do
      y <- f x
      modify $ insert x y
      return y

fib :: (Monad m) => (Int -> m Integer) -> (Int -> m Integer)
fib _ 0 = return 0
fib _ 1 = return 1
fib f n = (+) <$> f (n-1) <*> f (n-2)

slowFib :: Int -> Integer
slowFib n = runIdentity (fix fib n)

memoFib :: Int -> Integer
memoFib n = evalState (fix (memoize . fib) n) empty