{-# LANGUAGE BlockArguments #-}

import Data.Functor ((<&>))
import System.IO.Unsafe (unsafePerformIO)
import Data.IORef (IORef, newIORef, atomicModifyIORef')
import Data.Map.Lazy


data a --> b = Memo
  { cache :: !(IORef (Map a b))
  , comp  ::  (a -> b)
  }

mkMemo :: (a -> b) -> a --> b
mkMemo comp = unsafePerformIO $ newIORef empty <&> \cache ->
  Memo{cache,comp}

($$) :: Ord a => (a --> b) -> a -> b
Memo{cache,comp} $$ x = unsafePerformIO do
  atomicModifyIORef' cache \c -> case c !? x of
    Nothing -> (insert x y c, y)
      where y = comp x
    Just y  -> (           c, y)

memoFix :: Ord a => ((a -> b) -> a -> b) -> a -> b
memoFix withF = f where
  f x = memoF $$ x
  memoF = mkMemo (withF f)


memoFib :: Integer -> Integer
memoFib = memoFix \fib n -> if n <= 1
  then n
  else fib (n-2) + fib (n-1)

main :: IO ()
main = do
  print (memoFib 0)
  print (memoFib 1)
  print (memoFib 2)
  print (memoFib 3)
  print (memoFib 4)
  print (memoFib 5)
  print (memoFib 40)
  print (memoFib 100)
  print (memoFib 200)
  print (memoFib 300)
  print (memoFib 400)