ghc-heap: Don't Box NULL pointers

Previously we could construct a `Box` of a NULL pointer from the `link`
field of `StgWeak`. Now we take care to avoid ever introducing such
pointers in `collect_pointers` and ensure that the `link` field is
represented as a `Maybe` in the `Closure` type.

Fixes #21622

5 files changed, 15 insertions, 10 deletions

diff --git a/libraries/ghc-heap/GHC/Exts/Heap.hs b/libraries/ghc-heap/GHC/Exts/Heap.hs
index 74b559d19f..874c61b746 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap.hs
@@ -338,14 +338,17 @@ getClosureDataFromHeapRepPrim getConDesc decodeCCS itbl heapRep pts = do
 
         --  pure $ OtherClosure itbl pts rawHeapWords
         --
-        WEAK ->
+        WEAK -> do
             pure $ WeakClosure
                 { info = itbl
                 , cfinalizers = pts !! 0
                 , key = pts !! 1
                 , value = pts !! 2
                 , finalizer = pts !! 3
-                , link = pts !! 4
+                , weakLink = case drop 4 pts of
+                           []  -> Nothing
+                           [p] -> Just p
+                           _   -> error $ "Expected 4 or 5 words in WEAK, found " ++ show (length pts)
                 }
         TSO | [ u_lnk, u_gbl_lnk, tso_stack, u_trec, u_blk_ex, u_bq] <- pts
                 -> withArray rawHeapWords (\ptr -> do
diff --git a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
index 075e2a5b17..13d1ff71c2 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
@@ -41,6 +41,7 @@ import GHC.Exts.Heap.InfoTableProf ()
 import GHC.Exts.Heap.ProfInfo.Types
 
 import Data.Bits
+import Data.Foldable (toList)
 import Data.Int
 import Data.Word
 import GHC.Exts
@@ -228,8 +229,8 @@ data GenClosure b
         , mccPayload :: ![b]            -- ^ Array payload
         }
 
-    -- | An @MVar#@, with a queue of thread state objects blocking on them
-    | MVarClosure
+  -- | An @MVar#@, with a queue of thread state objects blocking on them
+  | MVarClosure
     { info       :: !StgInfoTable
     , queueHead  :: !b              -- ^ Pointer to head of queue
     , queueTail  :: !b              -- ^ Pointer to tail of queue
@@ -265,7 +266,7 @@ data GenClosure b
         , key         :: !b
         , value       :: !b
         , finalizer   :: !b
-        , link        :: !b -- ^ next weak pointer for the capability, can be NULL.
+        , weakLink    :: !(Maybe b) -- ^ next weak pointer for the capability
         }
 
   -- | Representation of StgTSO: A Thread State Object. The values for
@@ -420,7 +421,7 @@ allClosures (MVarClosure {..}) = [queueHead,queueTail,value]
 allClosures (IOPortClosure {..}) = [queueHead,queueTail,value]
 allClosures (FunClosure {..}) = ptrArgs
 allClosures (BlockingQueueClosure {..}) = [link, blackHole, owner, queue]
-allClosures (WeakClosure {..}) = [cfinalizers, key, value, finalizer, link]
+allClosures (WeakClosure {..}) = [cfinalizers, key, value, finalizer] ++ Data.Foldable.toList weakLink
 allClosures (OtherClosure {..}) = hvalues
 allClosures _ = []
 
diff --git a/libraries/ghc-heap/tests/T21622.hs b/libraries/ghc-heap/tests/T21622.hs
index db53ad8722..b1cc8c75c1 100644
--- a/libraries/ghc-heap/tests/T21622.hs
+++ b/libraries/ghc-heap/tests/T21622.hs
@@ -6,6 +6,7 @@ import GHC.Exts.Heap
 import System.Mem
 import GHC.Exts
 import GHC.Weak
+import Data.Foldable
 
 main :: IO ()
 main = do
@@ -19,7 +20,7 @@ collectWeaks :: Weak# v -> IO [Closure]
 collectWeaks = \w -> getClosureData w >>= go []
   where
     go :: [Closure] -> Closure -> IO [Closure]
-    go acc w@(WeakClosure {link})
-      | next <- link = getBoxedClosureData next >>= go (w:acc)
+    go acc w@(WeakClosure {weakLink})
+      | Just next <- weakLink = getBoxedClosureData next >>= go (w:acc)
       | otherwise = return (w:acc)
 
diff --git a/libraries/ghc-heap/tests/all.T b/libraries/ghc-heap/tests/all.T
index 95a5426358..504eb58857 100644
--- a/libraries/ghc-heap/tests/all.T
+++ b/libraries/ghc-heap/tests/all.T
@@ -52,5 +52,5 @@ test('parse_tso_flags',
      ],
      compile_and_run, [''])
 test('T21622',
-     [expect_broken(21622), only_ways(['normal'])],
+     only_ways(['normal']),
      compile_and_run, [''])
diff --git a/libraries/ghc-heap/tests/heap_all.hs b/libraries/ghc-heap/tests/heap_all.hs
index fa536a2d30..ac8953b5d2 100644
--- a/libraries/ghc-heap/tests/heap_all.hs
+++ b/libraries/ghc-heap/tests/heap_all.hs
@@ -155,7 +155,7 @@ exWeakClosure = WeakClosure
     , key = asBox []
     , value = asBox []
     , finalizer = asBox []
-    , link = asBox []
+    , weakLink = Nothing
     }
 
 exIntClosure :: Closure