Implement shrinkSmallMutableArray# and resizeSmallMutableArray#.

This is a part of GHC Proposal #25: "Offer more array resizing primitives".
Resources related to the proposal:

  - Discussion: https://github.com/ghc-proposals/ghc-proposals/pull/121
  - Proposal: https://github.com/ghc-proposals/ghc-proposals/blob/master/proposals/0025-resize-boxed.rst

Only shrinkSmallMutableArray# is implemented as a primop since a
library-space implementation of resizeSmallMutableArray# (in GHC.Exts)
is no less efficient than a primop would be. This may be replaced by
a primop in the future if someone devises a strategy for growing
arrays in-place. The library-space implementation always copies the
array when growing it.

This commit also tweaks the documentation of the deprecated
sizeofMutableByteArray#, removing the mention of concurrency. That
primop is unsound even in single-threaded applications. Additionally,
the non-negativity assertion on the existing shrinkMutableByteArray#
primop has been removed since this predicate is trivially always true.

12 files changed, 189 insertions, 6 deletions

diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index 155cdcbf80..e309d061a8 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -473,6 +473,7 @@ dispatchPrimop dflags = \case
         (bWord dflags))
 
   SizeofSmallMutableArrayOp -> dispatchPrimop dflags SizeofSmallArrayOp
+  GetSizeofSmallMutableArrayOp -> dispatchPrimop dflags SizeofSmallArrayOp
 
 -- IndexXXXoffAddr
 
@@ -1452,6 +1453,7 @@ dispatchPrimop dflags = \case
   ByteArrayIsPinnedOp -> alwaysExternal
   ShrinkMutableByteArrayOp_Char -> alwaysExternal
   ResizeMutableByteArrayOp_Char -> alwaysExternal
+  ShrinkSmallMutableArrayOp_Char -> alwaysExternal
   NewArrayArrayOp -> alwaysExternal
   NewMutVarOp -> alwaysExternal
   AtomicModifyMutVar2Op -> alwaysExternal
diff --git a/compiler/prelude/primops.txt.pp b/compiler/prelude/primops.txt.pp
index 47a78e2c8d..f47880b58d 100644
--- a/compiler/prelude/primops.txt.pp
+++ b/compiler/prelude/primops.txt.pp
@@ -1244,6 +1244,14 @@ primop  NewSmallArrayOp "newSmallArray#" GenPrimOp
 primop  SameSmallMutableArrayOp "sameSmallMutableArray#" GenPrimOp
    SmallMutableArray# s a -> SmallMutableArray# s a -> Int#
 
+primop  ShrinkSmallMutableArrayOp_Char "shrinkSmallMutableArray#" GenPrimOp
+   SmallMutableArray# s a -> Int# -> State# s -> State# s
+   {Shrink mutable array to new specified size, in
+    the specified state thread. The new size argument must be less than or
+    equal to the current size as reported by {\tt sizeofSmallMutableArray\#}.}
+   with out_of_line = True
+        has_side_effects = True
+
 primop  ReadSmallArrayOp "readSmallArray#" GenPrimOp
    SmallMutableArray# s a -> Int# -> State# s -> (# State# s, a #)
    {Read from specified index of mutable array. Result is not yet evaluated.}
@@ -1264,6 +1272,13 @@ primop  SizeofSmallArrayOp "sizeofSmallArray#" GenPrimOp
 
 primop  SizeofSmallMutableArrayOp "sizeofSmallMutableArray#" GenPrimOp
    SmallMutableArray# s a -> Int#
+   {Return the number of elements in the array. Note that this is deprecated
+   as it is unsafe in the presence of resize operations on the
+   same byte array.}
+   with deprecated_msg = { Use 'getSizeofSmallMutableArray#' instead }
+
+primop  GetSizeofSmallMutableArrayOp "getSizeofSmallMutableArray#" GenPrimOp
+   SmallMutableArray# s a -> State# s -> (# State# s, Int# #)
    {Return the number of elements in the array.}
 
 primop  IndexSmallArrayOp "indexSmallArray#" GenPrimOp
@@ -1463,7 +1478,7 @@ primop  SizeofByteArrayOp "sizeofByteArray#" GenPrimOp
 primop  SizeofMutableByteArrayOp "sizeofMutableByteArray#" GenPrimOp
    MutableByteArray# s -> Int#
    {Return the size of the array in bytes. Note that this is deprecated as it is
-   unsafe in the presence of concurrent resize operations on the same byte
+   unsafe in the presence of resize operations on the same byte
    array.}
    with deprecated_msg = { Use 'getSizeofMutableByteArray#' instead }
 
diff --git a/includes/rts/storage/ClosureMacros.h b/includes/rts/storage/ClosureMacros.h
index 2af50863d0..b5ae2dafc6 100644
--- a/includes/rts/storage/ClosureMacros.h
+++ b/includes/rts/storage/ClosureMacros.h
@@ -573,13 +573,20 @@ EXTERN_INLINE void overwritingClosure (StgClosure *p)
 // be less than or equal to closure_sizeW(p), and usually at least as
 // large as the respective thunk header.
 //
-// Note: As this calls LDV_recordDead() you have to call LDV_RECORD()
+// Note: As this calls LDV_recordDead() you have to call LDV_RECORD_CREATE()
 //       on the final state of the closure at the call-site
 EXTERN_INLINE void overwritingClosureOfs (StgClosure *p, uint32_t offset);
 EXTERN_INLINE void overwritingClosureOfs (StgClosure *p, uint32_t offset)
 {
-    // Set prim = true because only called on ARR_WORDS with the
-    // shrinkMutableByteArray# primop
+    // Set prim = true because overwritingClosureOfs is only
+    // ever called by
+    //   shrinkMutableByteArray# (ARR_WORDS)
+    //   shrinkSmallMutableArray# (SMALL_MUT_ARR_PTRS)
+    // This causes LDV_recordDead to be invoked. We want this
+    // to happen because the implementations of the above
+    // primops both call LDV_RECORD_CREATE after calling this,
+    // effectively replacing the LDV closure biography.
+    // See Note [LDV Profiling when Shrinking Arrays]
     overwritingClosure_(p, offset, closure_sizeW(p), true);
 }
 
diff --git a/includes/stg/MiscClosures.h b/includes/stg/MiscClosures.h
index 7a2ac2ef51..5b2364407f 100644
--- a/includes/stg/MiscClosures.h
+++ b/includes/stg/MiscClosures.h
@@ -367,6 +367,7 @@ RTS_FUN_DECL(stg_isByteArrayPinnedzh);
 RTS_FUN_DECL(stg_isMutableByteArrayPinnedzh);
 RTS_FUN_DECL(stg_shrinkMutableByteArrayzh);
 RTS_FUN_DECL(stg_resizzeMutableByteArrayzh);
+RTS_FUN_DECL(stg_shrinkSmallMutableArrayzh);
 RTS_FUN_DECL(stg_casIntArrayzh);
 RTS_FUN_DECL(stg_newArrayzh);
 RTS_FUN_DECL(stg_newArrayArrayzh);
diff --git a/libraries/base/GHC/Exts.hs b/libraries/base/GHC/Exts.hs
index 622902a673..47392ada3d 100755
--- a/libraries/base/GHC/Exts.hs
+++ b/libraries/base/GHC/Exts.hs
@@ -38,6 +38,14 @@ module GHC.Exts
         -- * Compat wrapper
         atomicModifyMutVar#,
 
+        -- * Resize functions
+        --
+        -- | Resizing arrays of boxed elements is currently handled in
+        -- library space (rather than being a primop) since there is not
+        -- an efficient way to grow arrays. However, resize operations
+        -- may become primops in a future release of GHC.
+        resizeSmallMutableArray#,
+
         -- * Fusion
         build, augment,
 
@@ -248,3 +256,34 @@ atomicModifyMutVar#
 atomicModifyMutVar# mv f s =
   case unsafeCoerce# (atomicModifyMutVar2# mv f s) of
     (# s', _, ~(_, res) #) -> (# s', res #)
+
+-- | Resize a mutable array to new specified size. The returned
+-- 'SmallMutableArray#' is either the original 'SmallMutableArray#'
+-- resized in-place or, if not possible, a newly allocated
+-- 'SmallMutableArray#' with the original content copied over.
+-- 
+-- To avoid undefined behaviour, the original 'SmallMutableArray#' shall
+-- not be accessed anymore after a 'resizeSmallMutableArray#' has been
+-- performed. Moreover, no reference to the old one should be kept in order
+-- to allow garbage collection of the original 'SmallMutableArray#'  in
+-- case a new 'SmallMutableArray#' had to be allocated.
+--
+-- @since 4.14.0.0
+resizeSmallMutableArray#
+  :: SmallMutableArray# s a -- ^ Array to resize
+  -> Int# -- ^ New size of array
+  -> a
+     -- ^ Newly created slots initialized to this element.
+     -- Only used when array is grown.
+  -> State# s
+  -> (# State# s, SmallMutableArray# s a #)
+resizeSmallMutableArray# arr0 szNew a s0 =
+  case getSizeofSmallMutableArray# arr0 s0 of
+    (# s1, szOld #) -> if isTrue# (szNew <# szOld)
+      then case shrinkSmallMutableArray# arr0 szNew s1 of
+        s2 -> (# s2, arr0 #)
+      else if isTrue# (szNew ># szOld)
+        then case newSmallArray# szNew a s1 of
+          (# s2, arr1 #) -> case copySmallMutableArray# arr0 0# arr1 0# szOld s2 of
+            s3 -> (# s3, arr1 #)
+        else (# s1, arr0 #)
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index eeed94327e..3b5dddf693 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -37,7 +37,9 @@
     `System.Posix.Types`.
 
   * Add `Functor`, `Applicative` and `Monad` instances to `(,,) a b`
-    and `(,,,) a b c`
+    and `(,,,) a b c`.
+
+  * Add `resizeSmallMutableArray#` to `GHC.Exts`.
 
 ## 4.13.0.0 *TBA*
   * Bundled with GHC *TBA*
diff --git a/libraries/ghc-prim/changelog.md b/libraries/ghc-prim/changelog.md
index 4967971e67..411d118aa1 100644
--- a/libraries/ghc-prim/changelog.md
+++ b/libraries/ghc-prim/changelog.md
@@ -2,6 +2,15 @@
 
 - Shipped with GHC 8.10.1
 
+- Add primop for shrinking `SmallMutableArray#`
+  to `GHC.Prim`:
+
+        shrinkSmallMutableArray# :: SmallMutableArray# s a -> Int# -> State# s -> State# s
+
+  Note that `resizeSmallMutableArray#` is not included as
+  as primitive. It has been implemented in library space in
+  `GHC.Exts`. See the release notes of `base`.
+
 - Added to `GHC.Prim`:
 
         closureSize# :: a -> Int#
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index b66c561dcb..b5930363a1 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -174,12 +174,13 @@ stg_isMutableByteArrayPinnedzh ( gcptr mba )
 stg_shrinkMutableByteArrayzh ( gcptr mba, W_ new_size )
 // MutableByteArray# s -> Int# -> State# s -> State# s
 {
-   ASSERT(new_size >= 0);
    ASSERT(new_size <= StgArrBytes_bytes(mba));
 
    OVERWRITING_CLOSURE_OFS(mba, (BYTES_TO_WDS(SIZEOF_StgArrBytes) +
                                  ROUNDUP_BYTES_TO_WDS(new_size)));
    StgArrBytes_bytes(mba) = new_size;
+   // See the comments in overwritingClosureOfs for an explanation
+   // of the interaction with LDV profiling.
    LDV_RECORD_CREATE(mba);
 
    return ();
@@ -224,6 +225,22 @@ stg_resizzeMutableByteArrayzh ( gcptr mba, W_ new_size )
    }
 }
 
+// shrink size of SmallMutableArray in-place
+stg_shrinkSmallMutableArrayzh ( gcptr mba, W_ new_size )
+// SmallMutableArray# s -> Int# -> State# s -> State# s
+{
+   ASSERT(new_size <= StgSmallMutArrPtrs_ptrs(mba));
+
+   OVERWRITING_CLOSURE_OFS(mba, (BYTES_TO_WDS(SIZEOF_StgSmallMutArrPtrs) +
+                                 new_size));
+   StgSmallMutArrPtrs_ptrs(mba) = new_size;
+   // See the comments in overwritingClosureOfs for an explanation
+   // of the interaction with LDV profiling.
+   LDV_RECORD_CREATE(mba);
+
+   return ();
+}
+
 // RRN: This one does not use the "ticketing" approach because it
 // deals in unboxed scalars, not heap pointers.
 stg_casIntArrayzh( gcptr arr, W_ ind, W_ old, W_ new )
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 0611de11cc..b2f90a892d 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -716,6 +716,7 @@
       SymI_HasProto(stg_isMutableByteArrayPinnedzh)                     \
       SymI_HasProto(stg_shrinkMutableByteArrayzh)                       \
       SymI_HasProto(stg_resizzeMutableByteArrayzh)                      \
+      SymI_HasProto(stg_shrinkSmallMutableArrayzh)                       \
       SymI_HasProto(newSpark)                                           \
       SymI_HasProto(updateRemembSetPushThunk)                             \
       SymI_HasProto(updateRemembSetPushThunk_)                            \
diff --git a/testsuite/tests/primops/should_run/ShrinkSmallMutableArrayA.hs b/testsuite/tests/primops/should_run/ShrinkSmallMutableArrayA.hs
new file mode 100644
index 0000000000..9202bcfffb
--- /dev/null
+++ b/testsuite/tests/primops/should_run/ShrinkSmallMutableArrayA.hs
@@ -0,0 +1,35 @@
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+
+import Control.Monad (unless)
+import GHC.Exts
+import GHC.Types
+
+-- This test is nearly a copy of T11296. In T11296, it is
+-- shrinkMutableByteArray# that is tested. Here, it is
+-- shrinkSmallMutableArray# that is tested.
+
+data SmallArray = SA (SmallMutableArray# RealWorld Integer)
+
+main :: IO ()
+main = do
+    let element = 42 :: Integer
+    ba# <- IO (\s0 -> case newSmallArray# 256# element s0 of
+                        (# s1, ba# #) -> (# s1, SA ba# #))
+    let go n = do
+            shrink ba# n
+            sz <- getSize ba#
+            unless (sz == n) $ print (sz, n)
+    mapM go [128, 64, 63, 32, 2, 1]
+    return ()
+
+shrink :: SmallArray -> Int -> IO ()
+shrink (SA ba#) (I# n#) = IO (\s ->
+    case shrinkSmallMutableArray# ba# n# s of
+      s' -> (# s', () #))
+
+getSize :: SmallArray -> IO Int
+getSize (SA ba#) = IO (\s ->
+    case getSizeofSmallMutableArray# ba# s of
+      (# s', n# #) -> (# s', I# n# #))
+
diff --git a/testsuite/tests/primops/should_run/ShrinkSmallMutableArrayB.hs b/testsuite/tests/primops/should_run/ShrinkSmallMutableArrayB.hs
new file mode 100644
index 0000000000..c720c9da14
--- /dev/null
+++ b/testsuite/tests/primops/should_run/ShrinkSmallMutableArrayB.hs
@@ -0,0 +1,53 @@
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+
+import Prelude hiding (read)
+import Control.Monad (unless)
+import GHC.Exts
+import GHC.Types
+import System.Mem (performMajorGC)
+
+-- The purpose of this test is to confirm that when the GC
+-- copies (out of the nursery) a SmallMutableArray# that has
+-- been shrunk, the array does not get corrupted.
+
+data SmallArray = SA (SmallMutableArray# RealWorld Integer)
+
+main :: IO ()
+main = do
+    let element = 42 :: Integer
+    arr <- IO (\s0 -> case newSmallArray# 30# element s0 of
+                        (# s1, ba# #) -> (# s1, SA ba# #))
+    write arr 0 100
+    write arr 13 113
+    write arr 14 114
+    write arr 15 115
+    write arr 16 116
+    shrink arr 14
+    performMajorGC
+    newSz <- getSize arr
+    unless (newSz == 14) (fail "Wrong new size")
+    e0 <- read arr 0
+    unless (e0 == 100) $
+      fail ("Wrong element 0: expected 100 but got " ++ show e0)
+    e13 <- read arr 13
+    unless (e13 == 113) $
+      fail ("Wrong element 13: expected 113 but got " ++ show e13)
+
+shrink :: SmallArray -> Int -> IO ()
+shrink (SA ba#) (I# n#) = IO (\s ->
+    case shrinkSmallMutableArray# ba# n# s of
+      s' -> (# s', () #))
+
+getSize :: SmallArray -> IO Int
+getSize (SA ba#) = IO (\s ->
+    case getSizeofSmallMutableArray# ba# s of
+      (# s', n# #) -> (# s', I# n# #))
+
+write :: SmallArray -> Int -> Integer -> IO ()
+write (SA ba#) (I# i#) e = IO (\s ->
+    case writeSmallArray# ba# i# e s of
+      s' -> (# s', () #))
+
+read :: SmallArray -> Int -> IO Integer
+read (SA ba#) (I# i#) = IO (\s -> readSmallArray# ba# i# s)
diff --git a/testsuite/tests/primops/should_run/all.T b/testsuite/tests/primops/should_run/all.T
index 6f5a04c5ed..bbcbdd8f78 100644
--- a/testsuite/tests/primops/should_run/all.T
+++ b/testsuite/tests/primops/should_run/all.T
@@ -26,3 +26,5 @@ test('ArithWord16', omit_ways(['ghci']), compile_and_run, [''])
 
 test('CmpInt16', normal, compile_and_run, [''])
 test('CmpWord16', normal, compile_and_run, [''])
+test('ShrinkSmallMutableArrayA', normal, compile_and_run, [''])
+test('ShrinkSmallMutableArrayB', normal, compile_and_run, [''])