ghc-heap: implement WEAK closure type #16974

6 files changed, 84 insertions, 0 deletions

diff --git a/libraries/ghc-heap/GHC/Exts/Heap.hs b/libraries/ghc-heap/GHC/Exts/Heap.hs
index d3b9097b2d..54481ccbca 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap.hs
@@ -270,6 +270,17 @@ getClosure x = do
 
         --  pure $ OtherClosure itbl pts wds
         --
+
+        WEAK ->
+            pure $ WeakClosure
+                { info = itbl
+                , cfinalizers = pts !! 0
+                , key = pts !! 1
+                , value = pts !! 2
+                , finalizer = pts !! 3
+                , link = pts !! 4
+                }
+
         _ ->
             pure $ UnsupportedClosure itbl
 
diff --git a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
index 025c30aaa1..82d0790f43 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
@@ -253,6 +253,15 @@ data GenClosure b
         , queue      :: !b              -- ^ ??
         }
 
+  | WeakClosure
+        { info        :: !StgInfoTable
+        , cfinalizers :: !b
+        , key         :: !b
+        , value       :: !b
+        , finalizer   :: !b
+        , link        :: !b -- ^ next weak pointer for the capability, can be NULL.
+        }
+
     ------------------------------------------------------------
     -- Unboxed unlifted closures
 
@@ -335,6 +344,7 @@ allClosures (MutVarClosure {..}) = [var]
 allClosures (MVarClosure {..}) = [queueHead,queueTail,value]
 allClosures (FunClosure {..}) = ptrArgs
 allClosures (BlockingQueueClosure {..}) = [link, blackHole, owner, queue]
+allClosures (WeakClosure {..}) = [cfinalizers, key, value, finalizer, link]
 allClosures (OtherClosure {..}) = hvalues
 allClosures _ = []
 
diff --git a/libraries/ghc-heap/tests/all.T b/libraries/ghc-heap/tests/all.T
index 88520e319d..afa224fde7 100644
--- a/libraries/ghc-heap/tests/all.T
+++ b/libraries/ghc-heap/tests/all.T
@@ -6,6 +6,14 @@ test('heap_all',
      ],
      compile_and_run, [''])
 
+test('heap_weak',
+     [when(have_profiling(), extra_ways(['prof'])),
+      # These ways produce slightly different heap representations.
+      # Currently we don't test them.
+      omit_ways(['ghci', 'hpc'])
+     ],
+     compile_and_run, [''])
+
 # Test everything except FUNs and PAPs in all ways.
 test('closure_size',
      [extra_files(['ClosureSizeUtils.hs']),
diff --git a/libraries/ghc-heap/tests/heap_all.hs b/libraries/ghc-heap/tests/heap_all.hs
index 76da037034..1560d4d9e8 100644
--- a/libraries/ghc-heap/tests/heap_all.hs
+++ b/libraries/ghc-heap/tests/heap_all.hs
@@ -14,6 +14,7 @@ import GHC.IORef
 import GHC.MVar
 import GHC.Stack
 import GHC.STRef
+import GHC.Weak
 import GHC.Word
 import System.Environment
 import System.Mem
@@ -147,6 +148,16 @@ exBlockingQClosure = BlockingQueueClosure
     , queue = asBox []
     }
 
+exWeakClosure :: Closure
+exWeakClosure = WeakClosure
+    { info = exItbl{tipe=WEAK}
+    , cfinalizers = asBox []
+    , key = asBox []
+    , value = asBox []
+    , finalizer = asBox []
+    , link = asBox []
+    }
+
 exIntClosure :: Closure
 exIntClosure = IntClosure
     { ptipe = PInt, intVal = 42 }
@@ -287,6 +298,12 @@ main = do
     -- getClosureData (Just 1) >>=
     --    assertClosuresEq exBlockingQClosure
 
+    -- Weak pointer
+    Weak wk <- mkWeak (1 :: Int) (1 :: Int) Nothing
+
+    getClosureData wk >>=
+        assertClosuresEq exWeakClosure
+
     -----------------------------------------------------
     -- Unboxed unlifted types
 
@@ -378,6 +395,7 @@ compareClosures expected actual =
                     MVarClosure{}           -> [ sEq (tipe . info) ]
                     MutVarClosure{}         -> [ sEq (tipe . info) ]
                     BlockingQueueClosure{}  -> [ sEq (tipe . info) ]
+                    WeakClosure{}           -> [ sEq (tipe . info) ]
                     IntClosure{}            -> [ sEq ptipe
                                                , sEq intVal    ]
                     WordClosure{}           -> [ sEq ptipe
diff --git a/libraries/ghc-heap/tests/heap_weak.hs b/libraries/ghc-heap/tests/heap_weak.hs
new file mode 100644
index 0000000000..47784f2297
--- /dev/null
+++ b/libraries/ghc-heap/tests/heap_weak.hs
@@ -0,0 +1,34 @@
+-- The simplifier changes the shapes of closures that we expect.
+{-# OPTIONS_GHC -O0 #-}
+{-# LANGUAGE MagicHash, UnboxedTuples, LambdaCase #-}
+
+import GHC.Exts.Heap
+import GHC.IORef
+import GHC.Weak
+import System.Mem
+
+main :: IO ()
+main = do
+    key <- newIORef "key"
+    let val = "val"
+    wk@(Weak w) <- mkWeak key val Nothing
+
+    getClosureData w >>= \case
+      WeakClosure{} -> putStrLn "OK"
+      _ -> error "Weak is not a WeakClosure"
+
+    deRefWeak wk >>= \case
+      Nothing -> error "Weak dead when key alive"
+      Just _ -> pure ()
+
+    readIORef key >>= putStrLn
+
+    performMajorGC
+
+    deRefWeak wk >>= \case
+      Nothing -> pure ()
+      Just _ -> error "Weak alive when key dead"
+
+    getClosureData w >>= \case
+      ConstrClosure{} -> putStrLn "OK"
+      _ -> error "dead Weak should be a ConstrClosure"
diff --git a/libraries/ghc-heap/tests/heap_weak.stdout b/libraries/ghc-heap/tests/heap_weak.stdout
new file mode 100644
index 0000000000..b4d5739b90
--- /dev/null
+++ b/libraries/ghc-heap/tests/heap_weak.stdout
@@ -0,0 +1,3 @@
+OK
+key
+OK