From ca9f14b7371efd7c2baa4d53d50becd03db46ea4 Mon Sep 17 00:00:00 2001
From: Adithya Kumar <adithya@composewell.com>
Date: Wed, 4 Dec 2024 19:41:06 +0530
Subject: [PATCH] Update MutByteArray.new to autopin if size >
 LARGE_BLOCK_THREASHOLD

BLOCK_SIZE = 4 * 1024 = 4096 bytes
LARGE_BLOCK_THREASHOLD = BLOCK_SIZE * 8 / 10 = 3276 bytes
---
 .../Internal/Data/MutByteArray/Type.hs        | 37 ++++++++++++++-----
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/core/src/Streamly/Internal/Data/MutByteArray/Type.hs b/core/src/Streamly/Internal/Data/MutByteArray/Type.hs
index 2042b3902d..8c4cbccf36 100644
--- a/core/src/Streamly/Internal/Data/MutByteArray/Type.hs
+++ b/core/src/Streamly/Internal/Data/MutByteArray/Type.hs
@@ -155,14 +155,20 @@ empty = unsafePerformIO $ new 0
 nil :: MutByteArray
 nil = empty
 
--- XXX add "newRoundedUp" to round up the large size to the next page boundary
--- and return the allocated size.
-{-# INLINE new #-}
-new :: Int -> IO MutByteArray
-new nbytes | nbytes < 0 =
-  errorWithoutStackTrace "newByteArray: size must be >= 0"
-new (I# nbytes) = IO $ \s ->
-    case newByteArray# nbytes s of
+-- 4000
+{-# INLINE _BLOCK_SIZE #-}
+_BLOCK_SIZE :: Int
+_BLOCK_SIZE = 4 * 1024
+
+-- 3276
+{-# INLINE _LARGE_BLOCK_THRESHOLD #-}
+_LARGE_BLOCK_THRESHOLD :: Int
+_LARGE_BLOCK_THRESHOLD = (_BLOCK_SIZE * 8) `div` 10
+
+{-# INLINE pinnedNewRaw #-}
+pinnedNewRaw :: Int -> IO MutByteArray
+pinnedNewRaw (I# nbytes) = IO $ \s ->
+    case newPinnedByteArray# nbytes s of
         (# s', mbarr# #) ->
            let c = MutByteArray mbarr#
             in (# s', c #)
@@ -171,8 +177,19 @@ new (I# nbytes) = IO $ \s ->
 pinnedNew :: Int -> IO MutByteArray
 pinnedNew nbytes | nbytes < 0 =
   errorWithoutStackTrace "pinnedNew: size must be >= 0"
-pinnedNew (I# nbytes) = IO $ \s ->
-    case newPinnedByteArray# nbytes s of
+pinnedNew nbytes = pinnedNewRaw nbytes
+
+-- XXX add "newRoundedUp" to round up the large size to the next page boundary
+-- and return the allocated size.
+-- Uses the pinned version of allocated if the size required is >
+-- _LARGE_BLOCK_THRESHOLD
+{-# INLINE new #-}
+new :: Int -> IO MutByteArray
+new nbytes | nbytes > _LARGE_BLOCK_THRESHOLD = pinnedNewRaw nbytes
+new nbytes | nbytes < 0 =
+  errorWithoutStackTrace "newByteArray: size must be >= 0"
+new (I# nbytes) = IO $ \s ->
+    case newByteArray# nbytes s of
         (# s', mbarr# #) ->
            let c = MutByteArray mbarr#
             in (# s', c #)