-- Simple Memoization Monad
-- by Taylor Campbell, 3/2008

module SimpleMemo where

import qualified Data.Map as Map
import Control.Monad.Identity
import State

type MemoFunction a b = a -> MemoAction a b b

type Memoization a b = Map.Map a b

type MemoAction a b c = State (Memoization a b) c

callMemo :: Ord a => MemoFunction a b -> MemoFunction a b
callMemo function
    = \ argument ->
          do map <- get
             case Map.lookup argument map of
               Just result      -> return result
               Nothing          ->
                   do result    <- function argument
                      map'      <- get
                      put (Map.insert argument result map')
                      return result

runWithMemo :: MemoAction a b c -> Memoization a b -> (c, Memoization a b)
runWithMemo action memo = runState action memo

runMemo :: MemoAction a b c -> (c, Memoization a b)
runMemo action = runWithMemo action Map.empty

type MemoFixable a b = MemoFunction a b -> MemoFunction a b

fixMemo :: Ord a => MemoFixable a b -> (a -> b)
fixMemo functionGenerator argument = fst (runMemo action)
  where
    action      = callMemo function argument
    function    = functionGenerator (\ argument -> callMemo function argument)

type IdentityFixable a b = (a -> Identity b) -> (a -> Identity b)

fixUnmemo :: IdentityFixable a b -> (a -> b)
fixUnmemo functionGenerator = runIdentity . fix functionGenerator