IOUArray to ByteSring, as quickly as possible

I need to mutate elements in a fixed size array of Word8 very quickly. For this purpose I am using an IOUArray. I need to send this array over a websocket connection. The function sendBinaryData from the websockets package requires a ByteString. I need to convert from one representation to the other. I am using this function currently:

arrayToBS :: IOUArray Int Word8 -> IO (BS.ByteString)
arrayToBS = (fmap BS.pack) . getElems

This function converts the elements of an array to a [Word8] before packing that list into a bytestring, and from profiling I can see that it is quite slow. I was wondering if there a way to speed up this function, or possibly send the array over a websocket connection directly?

The array that I am using currently is:

size = 1000;
numBytes = size * size * 4

newBuffer :: IO (IOUArray Int Word8)
newBuffer = newArray (0, numBytes) 200 :: IO (IOUArray Int Word8)

and an except from the performance report:

COST CENTRE MODULE SRC                        %time %alloc

arrayToBS   Lib    src/Lib.hs:28:1-37          88.1   99.0
newBuffer   Lib    src/Lib.hs:(23,1)-(25,12)    9.9    0.8

Ideally arrayToBS would be much faster than creating the array. If I change the size to 100:

COST CENTRE         MODULE                          SRC                                                %time %alloc

arrayToBS           Lib                             src/Lib.hs:21:1-37                           100.0   86.1
mkEncodeTable.table Data.ByteString.Base64.Internal Data/ByteString/Base64/Internal.hs:105:5-75    0.0    8.0
mkEncodeTable.ix    Data.ByteString.Base64.Internal Data/ByteString/Base64/Internal.hs:104:5-43    0.0    1.1

Solution

Disclaimer: I'm not very familiar with these low level primitives so this might be unsafe in some cases.

You will at the very least need to copy the data over once since, as @user2407038 remarks, the underlying data stored in an IOUArray is an unpinned array, so we can't count on GHC not moving the array around. The reverse direction (ByteString to IOArray) is however possible without a copy.

{-# LANGUAGE UnboxedTuples, MagicHash #-}

import Data.ByteString.Internal (ByteString(..))
import Data.Array.IO.Internals  (IOUArray(..))
import Data.Array.Base          (STUArray(..))
import Data.Word                (Word8)

import Foreign.ForeignPtr (mallocForeignPtrBytes, withForeignPtr)
import GHC.IO             (IO(..))
import GHC.Exts           (copyMutableByteArrayToAddr#, Ptr(..), Int(..))

arrayToBS :: IOUArray Int Word8 -> IO ByteString
arrayToBS (IOUArray (STUArray _ _ n@(I# n') mutByteArr)) = do
  bytes <- mallocForeignPtrBytes n
  withForeignPtr bytes $ \(Ptr addr) -> IO $ \state ->
    (# copyMutableByteArrayToAddr# mutByteArr 0# addr n' state, () #)
  pure (PS bytes 0 n)

Here is a test of this working (remember that the ascii code for 'A' is 65):

ghci> iou <- newListArray (-2,9) [65,67..] :: IO (IOUArray Int Word8)
ghci> arrayToBS iou
"ACEGIKMOQSUW"