Implement {resize,shrink}MutableByteArray# primops

The two new primops with the type-signatures

  resizeMutableByteArray# :: MutableByteArray# s -> Int#
                          -> State# s -> (# State# s, MutableByteArray# s #)

  shrinkMutableByteArray# :: MutableByteArray# s -> Int#
                          -> State# s -> State# s

allow to resize MutableByteArray#s in-place (when possible), and are useful
for algorithms where memory is temporarily over-allocated. The motivating
use-case is for implementing integer backends, where the final target size of
the result is either N or N+1, and only known after the operation has been
performed.

A future commit will implement a stateful variant of the
`sizeofMutableByteArray#` operation (see #9447 for details), since now the
size of a `MutableByteArray#` may change over its lifetime (i.e before
it gets frozen or GCed).

Test Plan: ./validate --slow

Reviewers: ezyang, austin, simonmar

Reviewed By: austin, simonmar

Differential Revision: https://phabricator.haskell.org/D133

3 files changed, 38 insertions, 0 deletions

diff --git a/includes/Cmm.h b/includes/Cmm.h
index 24bdda30c5..e62e96fcc0 100644
--- a/includes/Cmm.h
+++ b/includes/Cmm.h
@@ -600,8 +600,11 @@
 
 #if defined(PROFILING) || (!defined(THREADED_RTS) && defined(DEBUG))
 #define OVERWRITING_CLOSURE(c) foreign "C" overwritingClosure(c "ptr")
+#define OVERWRITING_CLOSURE_OFS(c,n) \
+    foreign "C" overwritingClosureOfs(c "ptr", n)
 #else
 #define OVERWRITING_CLOSURE(c) /* nothing */
+#define OVERWRITING_CLOSURE_OFS(c,n) /* nothing */
 #endif
 
 #ifdef THREADED_RTS
diff --git a/includes/rts/storage/ClosureMacros.h b/includes/rts/storage/ClosureMacros.h
index 3407b716c8..2a0f197218 100644
--- a/includes/rts/storage/ClosureMacros.h
+++ b/includes/rts/storage/ClosureMacros.h
@@ -504,8 +504,11 @@ INLINE_HEADER StgWord8 *mutArrPtrsCard (StgMutArrPtrs *a, W_ n)
 
 #if ZERO_SLOP_FOR_LDV_PROF || ZERO_SLOP_FOR_SANITY_CHECK
 #define OVERWRITING_CLOSURE(c) overwritingClosure(c)
+#define OVERWRITING_CLOSURE_OFS(c,n) \
+    overwritingClosureOfs(c,n)
 #else
 #define OVERWRITING_CLOSURE(c) /* nothing */
+#define OVERWRITING_CLOSURE_OFS(c,n) /* nothing */
 #endif
 
 #ifdef PROFILING
@@ -534,4 +537,34 @@ EXTERN_INLINE void overwritingClosure (StgClosure *p)
     }
 }
 
+// Version of 'overwritingClosure' which overwrites only a suffix of a
+// closure.  The offset is expressed in words relative to 'p' and shall
+// 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()
+//       on the final state of the closure at the call-site
+EXTERN_INLINE void overwritingClosureOfs (StgClosure *p, nat offset);
+EXTERN_INLINE void overwritingClosureOfs (StgClosure *p, nat offset)
+{
+    nat size, i;
+
+#if ZERO_SLOP_FOR_LDV_PROF && !ZERO_SLOP_FOR_SANITY_CHECK
+    // see Note [zeroing slop], also #8402
+    if (era <= 0) return;
+#endif
+
+    size = closure_sizeW(p);
+
+    ASSERT(offset <= size);
+
+    // For LDV profiling, we need to record the closure as dead
+#if defined(PROFILING)
+    LDV_recordDead(p, size);
+#endif
+
+    for (i = offset; i < size; i++)
+        ((StgWord *)p)[i] = 0;
+}
+
 #endif /* RTS_STORAGE_CLOSUREMACROS_H */
diff --git a/includes/stg/MiscClosures.h b/includes/stg/MiscClosures.h
index ee5a119aa1..d2b933deb0 100644
--- a/includes/stg/MiscClosures.h
+++ b/includes/stg/MiscClosures.h
@@ -347,6 +347,8 @@ RTS_FUN_DECL(stg_casArrayzh);
 RTS_FUN_DECL(stg_newByteArrayzh);
 RTS_FUN_DECL(stg_newPinnedByteArrayzh);
 RTS_FUN_DECL(stg_newAlignedPinnedByteArrayzh);
+RTS_FUN_DECL(stg_shrinkMutableByteArrayzh);
+RTS_FUN_DECL(stg_resizzeMutableByteArrayzh);
 RTS_FUN_DECL(stg_casIntArrayzh);
 RTS_FUN_DECL(stg_newArrayzh);
 RTS_FUN_DECL(stg_newArrayArrayzh);