-- Memoization Monad
-- by Scot Drysdale based on a more general program by Taylor Campbell, 3/2008 

module MemoS where

import qualified Data.Map as Map
import State          -- Could use the library State monad instead.

-- The type of memoizable functions.
type MemoFunction a b = a -> MemoState a b 

-- The type of an action that uses a memoized function (a -> b) and
-- yields a result of type b.
type MemoState a b = State (Map.Map a b) b

callMemo :: Ord a => MemoFunction a b -> a -> MemoState a b 
callMemo function argument
    = do map <- get        -- The state is a map
         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
                  

runMemo :: Ord a => MemoFunction a b -> a -> b
runMemo function a = evalState (function a) Map.empty