testsuite: Fix and extend closure_size test

This was previously broken in several ways. This is fixed and it also
now tests arrays. Unfortunately I was unable to find a way to continue
testing PAP and FUN sizes; these simply depend too much upon the
behavior of the simplifier.

I also tried to extend this to test non-empty arrays as well but
unfortunately this was non-trivial as the array card size constant isn't
readily available from haskell.

Fixes #16531.

2 files changed, 84 insertions, 12 deletions

diff --git a/libraries/ghc-heap/tests/all.T b/libraries/ghc-heap/tests/all.T
index 595bd000d4..5db21b39db 100644
--- a/libraries/ghc-heap/tests/all.T
+++ b/libraries/ghc-heap/tests/all.T
@@ -1,11 +1,15 @@
 test('heap_all',
-     [when(have_profiling(),
-      extra_ways(['prof'])),
+     [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('closure_size',
-     omit_ways(['ghci', 'hpc', 'prof']),
+     [ when(have_profiling(), extra_ways(['prof'])),
+      # These ways produce slightly different heap representations.
+      # Currently we don't test them.
+      omit_ways(['hpc'])
+     ],
      compile_and_run, [''])
+
diff --git a/libraries/ghc-heap/tests/closure_size.hs b/libraries/ghc-heap/tests/closure_size.hs
index d77060779e..f381a57521 100644
--- a/libraries/ghc-heap/tests/closure_size.hs
+++ b/libraries/ghc-heap/tests/closure_size.hs
@@ -1,26 +1,69 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE TypeInType #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE CPP #-}
 
 import Control.Monad
 import Type.Reflection
+import GHC.Exts
 import GHC.Stack
+import GHC.IO
 
 import GHC.Exts.Heap.Closures
 
-assertSize :: forall a. (HasCallStack, Typeable a)
-           => a -> Int -> IO ()
-assertSize !x expected = do
-  let !size = closureSize (asBox x)
-  when (size /= expected) $ do
-    putStrLn $ "closureSize ("++show (typeRep @a)++") == "++show size++", expected "++show expected
+assertSize
+  :: forall a. (HasCallStack, Typeable a)
+  => a     -- ^ closure
+  -> Int   -- ^ expected size in words
+  -> IO ()
+assertSize x =
+  assertSizeBox (asBox x) (typeRep @a)
+
+assertSizeUnlifted
+  :: forall (a :: TYPE 'UnliftedRep). (HasCallStack, Typeable a)
+  => a     -- ^ closure
+  -> Int   -- ^ expected size in words
+  -> IO ()
+assertSizeUnlifted x =
+  assertSizeBox (Box (unsafeCoerce# x)) (typeRep @a)
+
+assertSizeBox
+  :: forall a. (HasCallStack)
+  => Box   -- ^ closure
+  -> TypeRep a
+  -> Int   -- ^ expected size in words
+  -> IO ()
+assertSizeBox x ty expected = do
+  let !size = closureSize x
+  when (size /= expected') $ do
+    putStrLn $ "closureSize ("++show ty++") == "++show size++", expected "++show expected'
     putStrLn $ prettyCallStack callStack
+  where expected' = expected + profHeaderSize
 {-# NOINLINE assertSize #-}
 
 pap :: Int -> Char -> Int
 pap x _ = x
 {-# NOINLINE pap #-}
 
+profHeaderSize :: Int
+#if PROFILING
+profHeaderSize = 2
+#else
+profHeaderSize = 0
+#endif
+
+data A = A (Array# Int)
+data MA = MA (MutableArray# RealWorld Int)
+data BA = BA ByteArray#
+data MBA = MBA (MutableByteArray# RealWorld)
+data B = B BCO#
+data APC a = APC a
+
+
 main :: IO ()
 main = do
   assertSize 'a' 2
@@ -28,7 +71,32 @@ main = do
   assertSize (Nothing :: Maybe ()) 2
   assertSize ((1,2) :: (Int,Int)) 3
   assertSize ((1,2,3) :: (Int,Int,Int)) 4
-  assertSize (id :: Int -> Int) 1
-  assertSize (fst :: (Int,Int) -> Int) 1
-  assertSize (pap 1) 2
 
+  -- These depend too much upon the behavior of the simplifier to
+  -- test reliably.
+  --assertSize (id :: Int -> Int) 1
+  --assertSize (fst :: (Int,Int) -> Int) 1
+  --assertSize (pap 1) 2
+
+  MA ma <- IO $ \s ->
+      case newArray# 0# 0 s of
+          (# s1, x #) -> (# s1, MA x #)
+
+  A a <- IO $ \s ->
+      case freezeArray# ma 0# 0# s of
+          (# s1, x #) -> (# s1, A x #)
+
+  MBA mba <- IO $ \s ->
+      case newByteArray# 0# s of
+          (# s1, x #) -> (# s1, MBA x #)
+
+  BA ba <- IO $ \s ->
+      case newByteArray# 0# s of
+          (# s1, x #) ->
+              case unsafeFreezeByteArray# x s1 of
+                  (# s2, y #) -> (# s2, BA y #)
+
+  assertSizeUnlifted ma 3
+  assertSizeUnlifted a 3
+  assertSizeUnlifted mba 2
+  assertSizeUnlifted ba 2