如何加速Haskell中rotBitcount函数的位块旋转操作？

我分别用Rust和Haskell实现了一款全消息加密程序，程序会读取整个文件、加密后写入输出。处理1 MiB文件时，Rust代码耗时2.4秒，而Haskell代码耗时18.7秒。通过性能分析发现，36.6%的时间与内存分配消耗在rotBitcount函数上。该函数的功能是将一个字节数组复制到另一个数组，并按数组中1的位数对整个数组进行旋转。

以下是rotBitcount的相关代码：

module Cryptography.WringTwistree.RotBitcount
  ( rotBitcount
  ) where

{- This module is used in both Wring and Twistree.
 - It rotates an array of bytes by a multiple of its bitcount,
 - producing another array of the same size. As long as the multiplier
 - is relatively prime to the number of bits in the array, this
 - operation satisfies the strict avalanche criterion. Changing *two*
 - bits, however, has half a chance of changing only two bits in
 - the output.
 -
 - Bit 0 of byte 0 is bit 0 of the array. Bit 0 of byte 1 is bit 8 of the array.
 - e1 00 00 00 00 00 00 00, rotated by its bitcount (4), becomes
 - 10 0e 00 00 00 00 00 00.
 -}

import Data.Bits
import Data.Word
import Data.Array.Unboxed

rotBitcount :: (Integral a,Ix a,Bits a) => UArray a Word8 -> a -> UArray a Word8
-- The type a may be signed or unsigned, but the array index must begin at 0.
-- a should hold the square of eight times the bounds; so if the bounds are
-- (0..31), Word16 is adequate, but Int16 and Word8 are not.
rotBitcount src mult = array bnd
  [ (i, (src ! ((i+len-byte)   `mod` len) `shift` bit) .|.
        (src ! ((i+len-byte-1) `mod` len) `shift` (bit-8))) | i <- [0..(len-1)]]
  where
    bnd = bounds src
    len = (snd bnd) +1
    multmod = mult `mod` (len * 8)
    bitcount = fromIntegral $ sum $ map popCount $ elems src
    rotcount = (bitcount * multmod) `mod` (len * 8)
    byte = rotcount `shift` (-3)
    bit = fromIntegral (rotcount .&. 7)

我使用stack run进行编译（性能分析时除外），Stack版本为2.11.1，请问如何提升该函数的运行速度？

1. 替换列表推导式为ST monad数组构造 #

当前实现用列表推导式生成结果数组，每个元素需两次调用src ! idx，且惰性列表会带来额外内存分配与遍历开销。改用STUArray在ST monad中直接写入数组，可避免中间列表生成，减少索引重复计算：

import Control.Monad.ST
import Data.Array.ST
import Data.Bits
import Data.Word
import Data.Array.Unboxed

rotBitcount :: UArray Int Word8 -> Int -> UArray Int Word8
rotBitcount src mult = runSTUArray $ do
  let bnd@(0, lenMinus1) = bounds src
      len = lenMinus1 + 1
      totalBits = len * 8
      multmod = mult `mod` totalBits
      bitcount = fromIntegral $ sum $ map popCount $ elems src
      rotcount = (bitcount * multmod) `mod` totalBits
      byteShift = rotcount `shiftR` 3
      bitShift = rotcount .&. 7
  dst <- newArray bnd 0
  let srcIdx i = (i + len - byteShift) `mod` len
      srcIdxPrev i = (i + len - byteShift - 1) `mod` len
  forM_ [0..lenMinus1] $ \i -> do
    let s1 = src ! srcIdx i
        s2 = src ! srcIdxPrev i
        val = (s1 `shiftL` bitShift) .|. (s2 `shiftR` (8 - bitShift))
    writeArray dst i val
  return dst

2. 优化bitcount计算效率 #

当前bitcount通过sum $ map popCount $ elems src计算，elems生成的中间列表会增加内存开销。改用Data.Vector.Unboxed替代UArray，其foldl'可严格累加位计数结果，遍历性能更优：

import Data.Vector.Unboxed as V
import Data.Bits
import Data.Word

rotBitcount :: Vector Word8 -> Int -> Vector Word8
rotBitcount src mult = V.generate len $ \i ->
  let srcIdx = (i + len - byteShift) `mod` len
      srcIdxPrev = (i + len - byteShift - 1) `mod` len
      s1 = src V.! srcIdx
      s2 = src V.! srcIdxPrev
  in (s1 `shiftL` bitShift) .|. (s2 `shiftR` (8 - bitShift))
  where
    len = V.length src
    totalBits = len * 8
    multmod = mult `mod` totalBits
    bitcount = V.foldl' (\acc b -> acc + popCount b) 0 src
    rotcount = (bitcount * multmod) `mod` totalBits
    byteShift = rotcount `shiftR` 3
    bitShift = rotcount .&. 7

3. 减少模运算开销 #

多次mod len操作可针对性优化：

• 若程序处理的文件大小为2的幂（如1MiB），mod len可替换为(.&.) (len-1)，位运算远快于模运算；
• 提前将len转为无符号类型（如Word），减少符号处理带来的额外开销；
• 利用(i - byteShift) mod len等价于(i + len - byteShift) mod len的特性，简化计算逻辑。

4. 启用编译优化选项 #

默认stack run的优化级别不足，需在项目配置中启用高级优化：

• 在package.yaml设置optimization: 2（对应GHC -O2）；
• 若安装LLVM后端，添加-fllvm选项，LLVM代码生成对数值计算与数组操作更高效；
• 额外添加-funbox-strict-fields、-fspec-constr选项，帮助GHC生成更激进的优化代码。

修改stack.yaml中的ghc-options：

ghc-options:
  - -O2
  - -fllvm
  - -funbox-strict-fields

5. 去掉多态类型开销 #

当前函数的多态类型会带来字典传递与查找开销，实际使用中数组索引几乎均为Int，可指定具体类型签名，让GHC生成单态高效代码：

rotBitcount :: UArray Int Word8 -> Int -> UArray Int Word8

6. 批量处理字节减少循环次数 #

将输入视为连续位流，旋转后重新分组为字节，可采用块处理方式：比如预计算所有输入索引的映射关系，或一次处理多个字节，减少循环迭代次数，提升缓存命中率。

内容的提问来源于stack exchange，提问作者Pierre Abbat