Memoize.hs

{-# LANGUAGE DerivingVia #-}

-- | Basic memoization.
--
-- Functions yielded by 'memoize' and 'memoizeRec' may continue to allocate
-- memory without bound as long as they remain in scope. That is, you can keep
-- them around in a single short-lived thread, such as responding to an HTTP
-- request, but if kept at top-level will cause memory leaks. Use 'runMemRec'
-- to free memory as soon as the result is computed (i.e. forced).
module Memoize (Mem, memoize, MemRec, memoizeRec, runMemRec) where

import Data.Coerce
import Data.Functor.Identity
import Data.IORef
import Data.Map as Map

-- | A value of type @'Mem' A@ is a memoized @A@.
newtype Mem a = Mem (IO a)
  deriving (Functor, Applicative, Monad) via IO

-- | Memoize a non-recursive function.
memoize :: Ord a => (a -> Mem b) -> IO (a -> IO b)
memoize f = do
  storeRef <- newIORef mempty
  let recall = makeRecall storeRef f
  return recall

-- | A memoized recursive function, with the recursive call provided as an argument.
type MemRec a b = (a -> Mem b) -> a -> Mem b

-- | Memoize a recursive function.
memoizeRec :: Ord a => MemRec a b -> IO (a -> IO b)
memoizeRec f = do
  storeRef <- newIORef mempty
  let recall = makeRecall storeRef . f $ coerce recall
  return recall

-- | Compute a value from a memoized recursive function,
-- and free memory once the computation yields.
runMemRec :: Ord a => MemRec a b -> a -> IO b
runMemRec f x = memoizeRec f >>= ($ x)

makeRecall :: Ord a => IORef (Map a b) -> (a -> Mem b) -> a -> IO b
makeRecall storeRef f x = do
  store <- readIORef storeRef
  case Map.lookup x store of
    Just y -> return y
    Nothing -> do
      y <- coerce $ f x
      store' <- readIORef storeRef
      writeIORef storeRef $ Map.insert x y store'
      return y

fibRec :: MemRec Integer Integer
fibRec _ 0 = pure 0
fibRec _ 1 = pure 1
fibRec rec n = (+) <$> rec (n - 1) <*> rec (n - 2)

-- creates a function whose memory footprint grows without bound
-- as long as that function stays in scope.
makeFib :: IO (Integer -> IO Integer)
makeFib = memoizeRec fibRec

-- still uses memory (bounded by a function of the input) but frees
-- its memory as soon as the return value is used.
getFib :: Integer -> IO Integer
getFib = runMemRec fibRec