1% $Id: Memo.lhs,v 1.1 2005/12/16 10:46:05 simonmar Exp $ 2% 3% (c) The GHC Team, 1999 4% 5% Hashing memo tables. 6 7\begin{code} 8 9module Memo2 10 ( memo -- :: (a -> b) -> a -> b 11 , memoSized -- :: Int -> (a -> b) -> a -> b 12 ) 13 where 14 15import System.Mem.StableName ( StableName, makeStableName, hashStableName ) 16import System.Mem.Weak ( Weak, mkWeakPtr, mkWeak, deRefWeak, finalize ) 17import Data.Array.IO ( IOArray, newArray, readArray, writeArray ) 18import System.IO.Unsafe ( unsafePerformIO ) 19import Control.Concurrent.MVar ( MVar, newMVar, putMVar, takeMVar ) 20\end{code} 21 22----------------------------------------------------------------------------- 23Memo table representation. 24 25The representation is this: a fixed-size hash table where each bucket 26is a list of table entries, of the form (key,value). 27 28The key in this case is (StableName key), and we use hashStableName to 29hash it. 30 31It's important that we can garbage collect old entries in the table 32when the key is no longer reachable in the heap. Hence the value part 33of each table entry is (Weak val), where the weak pointer "key" is the 34key for our memo table, and 'val' is the value of this memo table 35entry. When the key becomes unreachable, a finalizer will fire and 36remove this entry from the hash bucket, and further attempts to 37dereference the weak pointer will return Nothing. References from 38'val' to the key are ignored (see the semantics of weak pointers in 39the documentation). 40 41\begin{code} 42type MemoTable key val 43 = MVar ( 44 Int, -- current table size 45 IOArray Int [MemoEntry key val] -- hash table 46 ) 47 48-- a memo table entry: compile with -funbox-strict-fields to eliminate 49-- the boxes around the StableName and Weak fields. 50data MemoEntry key val = MemoEntry !(StableName key) !(Weak val) 51\end{code} 52 53We use an MVar to the hash table, so that several threads may safely 54access it concurrently. This includes the finalization threads that 55remove entries from the table. 56 57ToDo: Can efficiency be improved at all? 58 59\begin{code} 60memo :: (a -> b) -> a -> b 61memo f = memoSized default_table_size f 62 63default_table_size = 1001 64 65-- Our memo functions are *strict*. Lazy memo functions tend to be 66-- less useful because it is less likely you'll get a memo table hit 67-- for a thunk. This change was made to match Hugs's Memo 68-- implementation, and as the result of feedback from Conal Elliot 69-- <conal@microsoft.com>. 70 71memoSized :: Int -> (a -> b) -> a -> b 72memoSized size f = strict (lazyMemoSized size f) 73 74strict = ($!) 75 76lazyMemoSized :: Int -> (a -> b) -> a -> b 77lazyMemoSized size f = 78 let (table,weak) = unsafePerformIO ( 79 do { tbl <- newArray (0,size) [] 80 ; mvar <- newMVar (size,tbl) 81 ; weak <- mkWeakPtr mvar (Just (table_finalizer tbl size)) 82 ; return (mvar,weak) 83 }) 84 in memo' f table weak 85 86table_finalizer :: IOArray Int [MemoEntry key val] -> Int -> IO () 87table_finalizer table size = 88 sequence_ [ finalizeBucket i | i <- [0..size] ] 89 where 90 finalizeBucket i = do 91 bucket <- readArray table i 92 sequence_ [ finalize w | MemoEntry _ w <- bucket ] 93 94memo' :: (a -> b) -> MemoTable a b -> Weak (MemoTable a b) -> a -> b 95memo' f ref weak_ref = \k -> unsafePerformIO $ do 96 stable_key <- makeStableName k 97 (size, table) <- takeMVar ref 98 let hash_key = hashStableName stable_key `mod` size 99 bucket <- readArray table hash_key 100 lkp <- lookupSN stable_key bucket 101 102 case lkp of 103 Just result -> do 104 putMVar ref (size,table) 105 return result 106 Nothing -> do 107 let result = f k 108 weak <- mkWeak k result (Just (finalizer hash_key stable_key weak_ref)) 109 writeArray table hash_key (MemoEntry stable_key weak : bucket) 110 putMVar ref (size,table) 111 return result 112 113finalizer :: Int -> StableName a -> Weak (MemoTable a b) -> IO () 114finalizer hash_key stable_key weak_ref = 115 do r <- deRefWeak weak_ref 116 case r of 117 Nothing -> return () 118 Just mvar -> do 119 (size,table) <- takeMVar mvar 120 bucket <- readArray table hash_key 121 let new_bucket = [ e | e@(MemoEntry sn weak) <- bucket, 122 sn /= stable_key ] 123 writeArray table hash_key new_bucket 124 putMVar mvar (size,table) 125 126lookupSN :: StableName key -> [MemoEntry key val] -> IO (Maybe val) 127lookupSN sn [] = sn `seq` return Nothing -- make it strict in sn 128lookupSN sn (MemoEntry sn' weak : xs) 129 | sn == sn' = do maybe_item <- deRefWeak weak 130 case maybe_item of 131 Nothing -> error ("dead weak pair: " ++ 132 show (hashStableName sn)) 133 Just v -> return (Just v) 134 | otherwise = lookupSN sn xs 135\end{code} 136