Memoization.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DefaultSignatures #-}

{- | Stochastic memoization for the `LazyPPL` library.
Stochastic memoization is a useful primitive for programming non-parametric models. 
It has type @(a -> Prob b) -> Prob (a -> b)@, and can be thought of as simultaneously sampling results for all possible arguments, in a lazy way. 

When @a@ is enumerable, this amounts to converting a stream of probabilities to a random stream, which we can do by sampling each probability once. 

This module provides:

* A general type-class `MonadMemo` for monads @m@ that support memoization at certain argument types @a@. 

* A default trie-based implementation when @a@ is enumerable, and a curry-based implementation when @a@ is a pair type. 

* A general implementation `generalmemoize` for probability, using memo-tables.

* A memoized recursion combinator, `memrec`. 


For illustrations, see the [graph example](https://lazyppl-team.github.io/GraphDemo.html), [clustering](https://lazyppl-team.github.io/ClusteringDemo.html), [additive clustering](https://lazyppl-team.github.io/AdditiveClusteringDemo.html), or the [infinite relational model](https://lazyppl-team.github.io/IrmDemo.html). 
-}

module LazyPPL.Distributions.Memoization (MonadMemo, memoize, generalmemoize, memrec) where

import Control.Monad
import Data.IORef
import Data.Map (empty, insert, keys, lookup, size)
import LazyPPL
import System.IO.Unsafe

{-| Type class for memoizable argument types @a@ under a monad @m@ -}
class (Monad m) => MonadMemo m a where
  memoize :: (a -> m b) -> m (a -> b)
  default memoize :: (Enum a) => (a -> m b) -> m (a -> b)
  memoize f = 
    do
      t <- ini 0 (f . toEnum)
      return $ \x -> look t (fromEnum x)

{-- Basic trie based integer-indexed memo table.
    NB Currently ignores negative integers --}
data BinTree a = Branch a (BinTree a) (BinTree a)

ini :: (Monad m) => Int -> (Int -> m a) -> m (BinTree a)
ini n f = do x <- f n; l <- ini (2 * n + 1) f; r <- ini (2 * n + 2) f; return $ Branch x l r

look :: BinTree a -> Int -> a
look (Branch x l r) 0 = x
look (Branch _ l r) n = if even n then look r (n `div` 2) else look l (n `div` 2)

{-| Implementation for enumerable types using tries -}
instance (Monad m) => MonadMemo m Int

{-| Implementation for pair types using currying -}
instance (Monad m, MonadMemo m a, MonadMemo m b) => MonadMemo m (a, b) where
  memoize f = fmap uncurry $ memoize $ \x -> memoize $ \y -> f (x, y)

{-| A general memoization method when @m@ is a probability monad.

    We use unsafePerformIO to maintain
    a table of calls that have already been made.
    If @a@ is finite, we could just sample all values of @a@ in advance
    and avoid unsafePerformIO. If @a@ is enumerable, we can use the trie method. 
-}
generalmemoize :: Ord a => (a -> Prob b) -> Prob (a -> b)
generalmemoize f = Prob $ \(Tree _ gs) ->
  unsafePerformIO $ do
    ref <- newIORef Data.Map.empty
    return $ \x -> unsafePerformIO $ do
      m <- liftM (Data.Map.lookup x) (readIORef ref)
      case m of
        Just y -> return y
        Nothing -> do
          n <- readIORef ref
          let y = runProb (f x) (gs !! (1 + size n))
          modifyIORef' ref (Data.Map.insert x y)
          return y

{-- Stochastic memoization for recursive functions.
    Applying 'memoize' to a recursively defined function only memoizes at the
    top-level: recursive calls are calls to the non-memoized function.
    'memrec' is an alternative implementation which resolves recursion and
    memoization at the same time, so that recursive calls are also memoized.
--}
memrec :: Ord a => Show a => ((a -> b) -> (a -> Prob b)) -> Prob (a -> b)
memrec f =
  Prob $ \(Tree _ gs) ->
    unsafePerformIO $ do
      ref <- newIORef Data.Map.empty
      let memoized_fixpoint = \x -> unsafePerformIO $ do
            m <- liftM (Data.Map.lookup x) (readIORef ref)
            case m of
              Just y -> return y
              Nothing -> do
                n <- readIORef ref
                let fix = f memoized_fixpoint
                let y = runProb (fix x) (gs !! (1 + size n))
                modifyIORef' ref (Data.Map.insert x y)
                return y
      return memoized_fixpoint