xp.memo / code / Memo / SNMap / SNMap.hs

{-# LANGUAGE StandaloneDeriving #-}
module Memo.SNMap.SNMap ( memo ) where

import Data.HashMap as M

import System.Mem.Weak
import System.Mem.StableName

import System.IO.Unsafe ( unsafePerformIO )




-------------------------------------------------------------------------------- 
-- SNMap Implementation
type SNMap k v = HashMap (StableName k) v

deriving instance Ord (StableName k)

newSNMap :: IO (SNMap k v)
newSNMap = return empty

lookupSNMap :: SNMap k v -> StableName k -> IO (Maybe v)
lookupSNMap map sn = return (M.lookup sn map)

insertSNMap :: SNMap k v -> StableName k -> v -> IO ()
insertSNMap = undefined

removeSNMap :: SNMap k v -> StableName k -> IO ()
removeSNMap = undefined

snMapElems :: SNMap k v -> IO [(k,v)]
snMapElems = undefined

-------------------------------------------------------------------------------- 
-- Memoisation with finalisers

type MemoTable a b = SNMap a (Weak b)

memo :: (a -> b) -> a -> b
memo f =
  let (table, weak) = unsafePerformIO (do { table  <- newSNMap
                                          ; weak   <- mkWeak table table (Just (table_finalizer table))
                                          ; return (table , weak)
                                          })
   in memo' f table weak
  
table_finalizer :: MemoTable a b -> IO ()
table_finalizer table =
  do { pairs <- snMapElems table; sequence_ [ finalize w | (_ , w)  <- pairs] }

memo' :: (a -> b) -> MemoTable a b -> Weak (MemoTable a b) -> a -> b
memo' f table weak_table arg = unsafePerformIO (
  do {  sn   <- makeStableName arg
     ;  lkp  <- lookupSNMap table sn
     ;  case lkp of 
          Nothing  -> not_found sn
          Just v   -> do  { maybe_val <- deRefWeak v
                          ; case maybe_val of
                              Nothing   -> not_found sn
                              Just val  -> return val
                          }
     })
  where val = f arg
        not_found sn = do  {  weak <- mkWeak arg val (Just (finalizer sn weak_table))
                           ;  insertSNMap table sn weak
                           ;  return val
                           }
                       
finalizer :: StableName a -> Weak (MemoTable a b) -> IO ()
finalizer sn weak_table = do  {  r <- deRefWeak weak_table
                              ;  case r of
                                   Nothing    -> return ()
                                   Just mvar  -> removeSNMap mvar sn
                              }
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.