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