|
|
1
|
+module T15907A
|
|
|
2
|
+ ( memo
|
|
|
3
|
+ , memoSized
|
|
|
4
|
+ ) where
|
|
|
5
|
+
|
|
|
6
|
+import Control.Concurrent.MVar (MVar, newMVar, putMVar, takeMVar)
|
|
|
7
|
+import Data.Array.IO (IOArray, newArray, readArray, writeArray)
|
|
|
8
|
+import System.IO.Unsafe (unsafePerformIO)
|
|
|
9
|
+import System.Mem.StableName (StableName, hashStableName, makeStableName)
|
|
|
10
|
+import System.Mem.Weak (Weak, deRefWeak, finalize, mkWeak, mkWeakPtr)
|
|
|
11
|
+
|
|
|
12
|
+type MemoTable key val =
|
|
|
13
|
+ MVar
|
|
|
14
|
+ ( Int
|
|
|
15
|
+ , IOArray Int [MemoEntry key val]
|
|
|
16
|
+ )
|
|
|
17
|
+
|
|
|
18
|
+data MemoEntry key val = MemoEntry !(StableName key) !(Weak val)
|
|
|
19
|
+
|
|
|
20
|
+memo :: (a -> b) -> a -> b
|
|
|
21
|
+memo f = memoSized defaultTableSize f
|
|
|
22
|
+
|
|
|
23
|
+defaultTableSize :: Int
|
|
|
24
|
+defaultTableSize = 1001
|
|
|
25
|
+
|
|
|
26
|
+memoSized :: Int -> (a -> b) -> a -> b
|
|
|
27
|
+memoSized size f = strict (lazyMemoSized size f)
|
|
|
28
|
+
|
|
|
29
|
+strict :: (a -> b) -> a -> b
|
|
|
30
|
+strict = ($!)
|
|
|
31
|
+
|
|
|
32
|
+lazyMemoSized :: Int -> (a -> b) -> a -> b
|
|
|
33
|
+lazyMemoSized size f =
|
|
|
34
|
+ let (table, weak) =
|
|
|
35
|
+ unsafePerformIO $ do
|
|
|
36
|
+ tbl <- newArray (0, size) []
|
|
|
37
|
+ mvar <- newMVar (size, tbl)
|
|
|
38
|
+ weak <- mkWeakPtr mvar (Just (tableFinalizer tbl size))
|
|
|
39
|
+ pure (mvar, weak)
|
|
|
40
|
+ in memo' f table weak
|
|
|
41
|
+
|
|
|
42
|
+tableFinalizer :: IOArray Int [MemoEntry key val] -> Int -> IO ()
|
|
|
43
|
+tableFinalizer table size =
|
|
|
44
|
+ sequence_ [finalizeBucket i | i <- [0 .. size]]
|
|
|
45
|
+ where
|
|
|
46
|
+ finalizeBucket i = do
|
|
|
47
|
+ bucket <- readArray table i
|
|
|
48
|
+ sequence_ [finalize w | MemoEntry _ w <- bucket]
|
|
|
49
|
+
|
|
|
50
|
+memo' :: (a -> b) -> MemoTable a b -> Weak (MemoTable a b) -> a -> b
|
|
|
51
|
+memo' f ref weakRef = \k -> unsafePerformIO $ do
|
|
|
52
|
+ stableKey <- makeStableName k
|
|
|
53
|
+ (size, table) <- takeMVar ref
|
|
|
54
|
+ let hashKey = hashStableName stableKey `mod` size
|
|
|
55
|
+ bucket <- readArray table hashKey
|
|
|
56
|
+ lkp <- lookupSN stableKey bucket
|
|
|
57
|
+ case lkp of
|
|
|
58
|
+ Just result -> do
|
|
|
59
|
+ putMVar ref (size, table)
|
|
|
60
|
+ pure result
|
|
|
61
|
+ Nothing -> do
|
|
|
62
|
+ let result = f k
|
|
|
63
|
+ weak <- mkWeak k result (Just (finalizer hashKey stableKey weakRef))
|
|
|
64
|
+ writeArray table hashKey (MemoEntry stableKey weak : bucket)
|
|
|
65
|
+ putMVar ref (size, table)
|
|
|
66
|
+ pure result
|
|
|
67
|
+
|
|
|
68
|
+finalizer :: Int -> StableName a -> Weak (MemoTable a b) -> IO ()
|
|
|
69
|
+finalizer hashKey stableKey weakRef = do
|
|
|
70
|
+ r <- deRefWeak weakRef
|
|
|
71
|
+ case r of
|
|
|
72
|
+ Nothing -> pure ()
|
|
|
73
|
+ Just mvar -> do
|
|
|
74
|
+ (size, table) <- takeMVar mvar
|
|
|
75
|
+ bucket <- readArray table hashKey
|
|
|
76
|
+ let newBucket =
|
|
|
77
|
+ [ e
|
|
|
78
|
+ | e@(MemoEntry sn _) <- bucket
|
|
|
79
|
+ , sn /= stableKey
|
|
|
80
|
+ ]
|
|
|
81
|
+ writeArray table hashKey newBucket
|
|
|
82
|
+ putMVar mvar (size, table)
|
|
|
83
|
+
|
|
|
84
|
+lookupSN :: StableName key -> [MemoEntry key val] -> IO (Maybe val)
|
|
|
85
|
+lookupSN sn [] = sn `seq` pure Nothing
|
|
|
86
|
+lookupSN sn (MemoEntry sn' weak : xs)
|
|
|
87
|
+ | sn == sn' = do
|
|
|
88
|
+ maybeItem <- deRefWeak weak
|
|
|
89
|
+ case maybeItem of
|
|
|
90
|
+ Nothing -> error ("dead weak pair: " ++ show (hashStableName sn))
|
|
|
91
|
+ Just v -> pure (Just v)
|
|
|
92
|
+ | otherwise = lookupSN sn xs |