Memoization with IntMap and State monad in Haskell

Memo.hs

{-# LANGUAGE FlexibleContexts, BangPatterns #-}

import Data.IntMap.Strict
import Data.Maybe
import Data.Function
import Control.Monad.Identity
import Control.Monad.State.Lazy hiding (fix)

import Prelude hiding (lookup)

import Criterion.Main

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)

memoize :: (MonadState (IntMap v) m) => (Int -> m v) -> Int -> m v
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

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

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

fibs = 0 : 1 : zipWith (+) fibs (tail fibs)
tabFib = (fibs !!)

tailFib n = go 0 1 n
  where
    go !x0 !x1 0 = x0
    go !x0 !x1 !n = go x1 (x1+x0) (n-1)

main = defaultMain [
  bgroup "naiveFib" [ bench (show i) $ whnf naiveFib i | i <- range ], 
  bgroup "memoFib" [ bench (show i) $ whnf memoFib i | i <- range ],
  bgroup "tailFib" [ bench (show i) $ whnf tailFib i | i <- range ],
  bgroup "tabFib"  [ bench (show i) $ whnf tabFib i | i <- range  ] 
  ]
    where range = [9, 11..15]