fib-memo.hs

{-# LANGUAGE LambdaCase #-}

import Data.IORef
import Data.Map (Map)
import qualified Data.Map as Map
import Numeric.Natural

type Cache a b = IORef (Map a b)

newCache :: IO (Cache a b)
newCache = newIORef Map.empty

cacheRead :: Ord a => Cache a b -> a -> IO (Maybe b)
cacheRead cache a =
  Map.lookup a <$> readIORef cache

cacheWrite :: Ord a => Cache a b -> a -> b -> IO ()
cacheWrite cache a b =
  modifyIORef cache (Map.insert a b)

doAndCache :: Ord a => Cache a b -> (a -> IO b) -> (a -> IO b)
doAndCache cache f a = do
  b <- f a
  cacheWrite cache a b
  return b

memoize :: Ord a => Cache a b -> (a -> IO b) -> (a -> IO b)
memoize cache f a =
  cacheRead cache a >>= \case
    Just b -> return b
    Nothing -> doAndCache cache f a

main :: IO ()
main =
  do
    cache <- newCache

    let
      fib :: Natural -> IO Natural
      fib = memoize cache $ \case
        0 -> return 0
        1 -> return 1
        n -> (+) <$> fib (n - 1) <*> fib (n - 2)

    x <- fib 30
    print x