Classic Ord-generic Memoisation

Memo.hs

{-# LANGUAGE GHC2021, BlockArguments, ExplicitNamespaces, DerivingVia #-}

module Memo (

  -- * Memoisation Functions
  memo, memoFix,

  -- * Memoised Function Type
  type (-->),
  toMemo, ($$),
  (~),

) where

-- base
import Data.Function (fix)
import Data.Functor ((<&>))
import Data.Monoid (Ap(..))
import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef')
import System.IO.Unsafe (unsafePerformIO)

-- containers
import Data.Map.Lazy (Map, empty, (!?), insert, intersectionWith, compose)


memo :: Ord a => (a -> b) -> a -> b
memo f = \x -> mf $$ x
 where
  {-# NOINLINE mf #-}
  !mf = toMemo f

memoFix :: Ord a => ((a -> b) -> a -> b) -> a -> b
memoFix recf = fix \f -> memo (recf f)


data a --> b = Memo
  { cache :: {-# UNPACK #-} !(IORef (Map a b))
  , comp  :: a -> b
  } deriving (Semigroup, Monoid) via Ap ((-->) a) b

toMemo :: (a -> b) -> a --> b
toMemo 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)

instance Ord a => Functor ((-->) a) where
  fmap f m = unsafePerformIO do
    cached <- readIORef (cache m)
    cache  <- newIORef (fmap f cached)
    pure Memo
      { cache
      , comp = \w -> f (m $$ w)
      }

instance Ord a => Applicative ((-->) a) where
  pure = toMemo . pure
  liftA2 f mx my = unsafePerformIO do
    cachedx <- readIORef (cache mx)
    cachedy <- readIORef (cache my)
    cache   <- newIORef (intersectionWith f cachedx cachedy)
    pure Memo
      { cache
      , comp = \w -> f (mx $$ w) (my $$ w)
      }

infixr 8 ~
(~) :: (Ord a, Ord b) => b --> c -> a --> b -> a --> c
bc ~ ab = unsafePerformIO do
  cachedbc <- readIORef (cache bc)
  cachedab <- readIORef (cache ab)
  cache    <- newIORef (compose cachedbc cachedab)
  pure Memo
    { cache
    , comp = \a -> bc $$ (ab $$ a)
    }