IOUArray 到 ByteSring，越快越好

Question

我需要快速改变 Word8 的固定大小数组中的元素。为此，我使用了 IOUArray。我需要通过 websocket 连接发送这个数组。 websockets 包中的函数 sendBinaryData 需要一个 ByteString。我需要从一种表示形式转换为另一种表示形式。我目前正在使用这个功能：

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

此函数将数组的元素转换为 [Word8]，然后将该列表打包为字节串，从分析中我可以看出它非常慢。我想知道是否有办法加快这个功能，或者直接通过 websocket 连接发送数组？

我目前使用的数组是：

size = 1000;
numBytes = size * size * 4

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

以及性能报告中的一个例外：

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

理想情况下 arrayToBS 会比创建数组快得多。如果我将 size 更改为 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

Answer 1

好的，感谢 user2407038 我有 一些东西（请注意，我以前从未玩过基元或未装箱的类型）：

import Control.Monad.ST
import qualified Data.ByteString as BS
import Data.Word
import Data.Array.ST
import Data.Array.Base
import Data.ByteString.Internal
import GHC.Prim
import GHC.Exts
import GHC.ForeignPtr

bs2Addr# :: BS.ByteString -> Addr#
bs2Addr# (PS fptr offset len) = case fptr of
  (ForeignPtr addr _ ) -> addr

arrayPrim (STUArray _ _ _ x) = x

unbox :: Int -> Int#
unbox (I# n#) = n#

copy :: Int -> IO BS.ByteString
copy len = do
  -- Get the length as unboxed
  let len# = unbox len

  -- Bytestring to copy to, filled with 0s initially
  let bs = BS.pack (replicate len 0)

  -- Create a new STUArray. I don't know why it needs to be length * 2.
  arr <- stToIO (newArray (0, len * 2) 255 :: ST s (STUArray s Int Word8))

  -- MutableByteArray#
  let mArrPrim = arrayPrim arr

  -- Addr#
  let addr = bs2Addr# bs

  -- I don't know what the 2nd and 4th Int# arguments are suppose to be.
  let _ = copyMutableByteArrayToAddr# mArrPrim len# addr len# realWorld#
  return bs

我现在在这里使用 STUArray 而不是 IOUArray 因为我找不到 IOUArray 构造函数。

使用 4000000 个元素数组分析此代码的结果：

    Sun Aug 20 20:49 2017 Time and Allocation Profiling Report  (Final)

       shoot-exe +RTS -N -p -RTS

    total time  =        0.05 secs   (47 ticks @ 1000 us, 1 processor)
    total alloc = 204,067,640 bytes  (excludes profiling overheads)

COST CENTRE MODULE SRC                        %time %alloc

copy.bs     Lib    src/Lib.hs:32:7-36          66.0   96.0
copy        Lib    src/Lib.hs:(27,1)-(45,11)   34.0    3.9

这是我与之比较的代码：

arrayToBS :: (STUArray s Int Word8) -> ST s (BS.ByteString)
arrayToBS = (fmap BS.pack) . getElems

slowCopy :: Int -> IO BS.ByteString
slowCopy len = do
  arr <- stToIO (newArray (0, len - 1) 255 :: ST s (STUArray s Int Word8))
  stToIO $ arrayToBS arr

及其分析报告：

    Sun Aug 20 20:48 2017 Time and Allocation Profiling Report  (Final)

       shoot-exe +RTS -N -p -RTS

    total time  =        0.55 secs   (548 ticks @ 1000 us, 1 processor)
    total alloc = 1,604,073,872 bytes  (excludes profiling overheads)

COST CENTRE MODULE SRC                        %time %alloc

arrayToBS   Lib    src/Lib.hs:48:1-37          98.2   99.7
slowCopy    Lib    src/Lib.hs:(51,1)-(53,24)    1.6    0.2

好的，新版本更快。它们都产生相同的输出。但是，我仍然想知道 copyMutableByteArrayToAddr# 的 #Int 参数是什么，以及为什么我必须将快速版本中的数组长度乘以 2。我会多玩一些并更新如果我发现了这个答案。

更新：Alec 的回答

对于那些好奇的人，这是分析 Alec 的回答的结果：

    Sun Aug 20 21:13 2017 Time and Allocation Profiling Report  (Final)

       shoot-exe +RTS -N -p -RTS

    total time  =        0.01 secs   (7 ticks @ 1000 us, 1 processor)
    total alloc =   8,067,696 bytes  (excludes profiling overheads)

COST CENTRE   MODULE SRC                          %time %alloc

newBuffer     Other  src/Other.hs:23:1-33          85.7   49.6
arrayToBS.\.\ Other  src/Other.hs:19:5-69          14.3    0.0
arrayToBS     Other  src/Other.hs:(16,1)-(20,21)    0.0   49.6

看起来就是这样。

Answer 2

免责声明：我对这些低级原语不是很熟悉，所以在某些情况下这可能不安全。

您至少需要将数据复制一次，因为正如@user2407038 所说，存储在 IOUArray 中的基础数据是一个未固定的数组，因此我们不能指望 GHC 不移动数组。然而，没有副本也可以反向（ByteString 到 IOArray）。

{-# 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)

这是对此工作的测试（请记住 'A' 的 ascii 代码是 65）：

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

IOUArray 到 ByteSring，越快越好

IOUArray to ByteSring, as quickly as possible

arrays

haskell

websocket

ghc

bytestring

更新：Alec 的回答