An alternative implementation of `pack` for bytestrings that can fuse.

PackFusion.hs

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

module Main (main) where

import           Test.Tasty.Bench

import qualified Data.ByteString                       as S
import qualified Data.ByteString.Internal              as SI

import Data.Word
import Foreign.Ptr
import Data.Bits
import System.IO.Unsafe
import GHC.Exts
import GHC.IO
import GHC.Word

data MBA = MBA !(MutableByteArray# RealWorld)
data DMBA = DMBA {-# UNPACK #-} !Int {-# UNPACK #-} !MBA

newMBA :: Int -> IO MBA
newMBA (I# n) = IO $ \s -> 
  case newPinnedByteArray# n s of 
    (# s, mba #) -> (# s, MBA mba #)

newDMBA :: Int -> IO DMBA
newDMBA n = DMBA 0 <$> newMBA n

writeMBA :: MBA -> Int -> Word8 -> IO ()
writeMBA (MBA mba) (I# i) (W8# x) = IO $ \s -> (# writeWord8Array# mba i x s, () #)

sizeMBA :: MBA -> IO Int
sizeMBA (MBA mba) = IO $ \s -> 
  case getSizeofMutableByteArray# mba s of
    (# s, size #) -> (# s, I# size #)

resizeMBA :: MBA -> Int -> IO MBA
resizeMBA (MBA mba) (I# n) = IO $ \s ->
  case resizeMutableByteArray# mba n s of
    (# s, mba #) -> (# s, MBA mba #)

growFun cap = cap `unsafeShiftL` 1
-- growFun cap = (3 * cap) `unsafeShiftR` 1

push :: DMBA -> Word8 -> IO DMBA
push (DMBA n mba) x = do
  size <- sizeMBA mba
  if n < size 
    then DMBA (n + 1) mba <$ writeMBA mba n x
    else do
      let size' = growFun size
      mba <- resizeMBA mba size'
      DMBA (n + 1) mba <$ writeMBA mba n x
{-# INLINE push #-}

toByteString :: DMBA -> IO S.ByteString
toByteString (DMBA (I# n) (MBA mba)) = SI.create (I# n) $ \(Ptr addr) ->
  IO $ \s ->
    let s' = shrinkMutableByteArray# mba n s in
    (# copyMutableByteArrayToAddr# mba 0# addr 0# s', () #)

pack :: [Word8] -> S.ByteString
pack xs = unsafeDupablePerformIO $ do
  dmba <- newDMBA 8
  foldr 
    (\x go dmba -> push dmba x >>= go) 
    (\ dmba@(DMBA n _) -> toByteString dmba)
    xs 
    dmba
{-# INLINE pack #-}

mkBench f = [bench ("10^" ++ show i) $ f (10 ^ i) | i <- [2,4,6]]

main :: IO ()
main = defaultMain
  [ bgroup "pack" $ mkBench (\n -> nf S.pack (replicate n 0))
  , bgroup "pack fusable" $ mkBench (\n -> nf (S.pack . replicate n) 0)
  , bgroup "my-pack" $ mkBench (\n -> nf pack (replicate n 0))
  , bgroup "my-pack fusable" $ mkBench (\n -> nf (pack . replicate n) 0)
  ]

Results:

All
  pack
    1000:    OK (0.29s)
      4.32 μs ± 293 ns, 628 B  allocated,   0 B  copied,  34 MB peak memory
    10000:   OK (0.21s)
      48.8 μs ± 2.9 μs,   0 B  allocated,   0 B  copied,  46 MB peak memory
    100000:  OK (0.27s)
      506  μs ±  32 μs,   0 B  allocated,   0 B  copied,  58 MB peak memory
    1000000: OK (0.95s)
      6.81 ms ± 468 μs, 415 KB allocated,  18 B  copied, 148 MB peak memory
  my-pack
    1000:    OK (0.41s)
      6.29 μs ± 563 ns,  25 KB allocated,   0 B  copied, 148 MB peak memory
    10000:   OK (0.41s)
      67.0 μs ± 5.3 μs, 272 KB allocated,   5 B  copied, 148 MB peak memory
    100000:  OK (0.39s)
      550  μs ±  25 μs, 2.6 MB allocated,  28 B  copied, 148 MB peak memory
    1000000: OK (0.25s)
      6.18 ms ± 590 μs,  24 MB allocated, 335 B  copied, 148 MB peak memory
  my-pack fused
    1000:    OK (0.47s)
      870  ns ±  76 ns, 3.0 KB allocated,   0 B  copied, 148 MB peak memory
    10000:   OK (0.43s)
      7.68 μs ± 741 ns,  40 KB allocated,   0 B  copied, 148 MB peak memory
    100000:  OK (0.46s)
      70.7 μs ± 5.1 μs, 338 KB allocated,   3 B  copied, 148 MB peak memory
    1000000: OK (0.52s)
      725  μs ±  69 μs, 2.9 MB allocated,  28 B  copied, 183 MB peak memory