Shrinking a large ByteArray# retains all of its memory

Summary

The primop shrinkMutableByteArray# just overwrites the size of its MutableByteArray# argument on the heap. But when that argument is large enough to consist of more than one block, some of its trailing blocks may no longer be part of the shrunken array. However, these blocks are not returned to the RTS until the shrunken array is garbage-collected, which can be very wasteful in extreme situations.

(resizeMutableByteArray# exhibits this problem as well.)

Steps to reproduce

Here's a sample program that can be used to verify this behavior:

{-# LANGUAGE MagicHash, UnboxedTuples, ViewPatterns #-}
import qualified GHC.Exts as Exts
import GHC.IO (IO(..))

import Control.Monad
import Data.Array.Byte
import System.Environment

data ShrinkMethod = Shrink | Resize | Copy
  deriving Read

shrink :: ShrinkMethod -> Exts.MutableByteArray# s -> Exts.Int# -> Exts.State# s
   -> (# Exts.State# s, Exts.MutableByteArray# s #)
shrink p a n s0 = case  p  of
  Shrink -> (# Exts.shrinkMutableByteArray# a n s0, a #)
  Resize -> Exts.resizeMutableByteArray# a n s0
  Copy -> case  Exts.newByteArray# n s0  of
    (# s1, tar #) -> case  Exts.copyMutableByteArray# a 0# tar 0# n s1  of
      s2 -> (# s2, tar #)

mkArr :: Int -> ShrinkMethod -> Int -> IO ByteArray
mkArr (Exts.I# initSize) whichOp (Exts.I# finalSize) = IO $ \s0 ->
  case  Exts.newByteArray# initSize s0  of
  (# s1, marrL #) -> case  shrink whichOp marrL finalSize s1  of
    (# s2, marrS #) -> case  Exts.setByteArray# marrS 0# finalSize 0# s2  of
      s3 -> case  Exts.unsafeFreezeByteArray# marrS s3  of
        (# s4, arr #) -> (# s4, ByteArray arr #)

main :: IO ()
main = do
  [read -> nArrs, read -> initSize, read -> whichOp, read -> finalSize]
    <- getArgs
  li <- replicateM nArrs (mkArr initSize whichOp finalSize)
  -- printing of the first ByteArray is sequenced after
  -- creation of the last ByteArray, so they are all live at once
  mapM_ print li

If I compile and run this program with ./ShrinkDegeneracy 1000 1048576 Shrink 48 +RTS -M1G > /dev/null, I see:

ShrinkDegeneracy: Heap exhausted;
ShrinkDegeneracy: Current maximum heap size is 1073741824 bytes (1024 MB).
ShrinkDegeneracy: Use `+RTS -M<size>' to increase it.

If I run ./ShrinkDegeneracy 1000 1048576 Resize 48 +RTS -M1G > /dev/null, I see the same error.

If I run ./ShrinkDegeneracy 1000 1048576 Copy 48 +RTS -M1G > /dev/null, the program successfully finishes in a small fraction of a second. (This is the desired behavior.)

Environment

• GHC version used: 9.4.4