WeakMap.hs

module Main where

import Control.Monad
import Data.IORef
--import Debug.Trace
import System.Mem.Weak
import qualified Data.Map as M

import Prelude hiding (lookup)

class Copy a where
  copy :: a -> a

instance Copy Int where
  copy x = fromIntegral $ toInteger x -- I'm sure there is a better way to force a copy

data WeakMap k a = WeakMap (IORef (M.Map k a))

empty :: IO (WeakMap k a)
empty = liftM WeakMap $ newIORef M.empty

lookup :: (Ord k) => k -> WeakMap k a -> IO (Maybe a)
lookup k (WeakMap ref) = liftM (M.lookup k) $ readIORef ref

delete :: (Ord t) => t -> WeakMap t t1 -> IO ()
delete k (WeakMap ref) = modifyIORef ref (M.delete k)

insert :: (Copy a, Ord a) => a -> t -> WeakMap a t -> IO ()
insert k v wm@(WeakMap ref) = do
  let k' = copy k
  modifyIORef ref (M.insert k' v)
  addFinalizer k (delete k' wm) -- do I need to use the key copy here?
  --addFinalizer k (trace ("finalizer called on: " ++ show k') delete k' wm) -- do I need to have the key copy here?


main :: IO ()
main = do
    wm <- empty
    let myKey = 1
    insert myKey "myStuff" wm
    mapM (testInsert wm) ([2..100000]::[Int])
    lookup myKey wm >>= print -- this key stayed around
    lookup 6 wm >>= print -- this key was GC'd so isn't in the map
  where
    testInsert wm k = insert k "asdf" wm