ByteString あれこれ

Haskell で高速なプログラムを書くには ByteString の深い知識が必要となる。鍵となるのは、Data.ByteString.Internal というモジュールである。このモジュールは公開されているが、ドキュメントは隠されているので、詳しく知るためにはソースを読まないといけない。

定義

ByteString の定義は以下の通り。

data ByteString = PS {-# UNPACK #-} !(ForeignPtr Word8) -- payload
                     {-# UNPACK #-} !Int                -- offset
                     {-# UNPACK #-} !Int                -- length

なぜ構成子が BS ではなく、PS なのだろう? (追記:元々 PackedString という型名だったからと @ma0e さんに教えて頂きました。)

ForeignPtr Word8 は、いわゆるバッファーを指してる。このバッファーは複数の ByteString から共有される。バッファー中のどこを使っているかを示すために、メタ情報として offset と length がある。

Haskell ポインタープログラミングを読んだ人なら、バッファーを読み書きする方法は分かるだろう。以下のようにするのが定石である。

import Data.ByteString.Internal (ByteString(..))
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (plusPtr)

foo :: ByteString -> IO 適当な型
foo (PS fptr off len) = withForeignPtr fptr $ \ptr -> do
    let beg = ptr `plusPtr` off
        end = beg `plusPtr` len
    -- beg から end の間(endは含まない)を読み書きする

書き換えてバッファー溢れを起こさないのは自己責任だし、書き換えるバッファーが共有されてないことを保証するのも自己責任である。

大きさの分かっている ByteString を作る

バッファー操作とともに ByteString を作りたい場合、あらかじめできあがる ByteString の大きさが分かっているなら、create が使える。

create :: Int -> (Ptr Word8 -> IO ()) -> IO ByteString
create l f = do
    fp <- mallocByteString l
    withForeignPtr fp $ \p -> f p
    return $! PS fp 0 l

Int は、バッファーのバイト数である。Ptr Word8 -> IO () という関数で、バッファー全体(オフセット 0 から指定したバイト数)を初期化すればよい。

副作用がないと確信できるなら、unsafeCreate を使おう。

unsafeCreate :: Int -> (Ptr Word8 -> IO ()) -> ByteString
unsafeCreate l f = unsafeDupablePerformIO (create l f)

大きさが分からない ByteString を作る

できあがる ByteString の大きさが分からない場合、どうすればいいだろうか? 3つぐらい方法がある。

createAndTrim

一つ目は createAndTrim を使う方法だ。十分な大きさのバッファーを用意してデータを作り、大きさが分かった時点で、新しい ByteString を作る。渡す関数は、できあがったデータのバイト数を返さないといけない。

以下のソースコードが示すように、はじめに用意したバッファーぴったりのデータができたときに限り、コピーが発生しない。

createAndTrim :: Int -> (Ptr Word8 -> IO Int) -> IO ByteString
createAndTrim l f = do
    fp <- mallocByteString l
    withForeignPtr fp $ \p -> do
        l' <- f p
        if assert (l' <= l) $ l' >= l
            then return $! PS fp 0 l
            else create l' $ \p' -> memcpy p' p l'
リスト

二つ目は中間データとしてリストを作る方法。たとえば、[Word8]というリストを作っておいて、pack すればいい。

[Word8]を作る際は、リストが後ろに伸びて行くだろう。逆順のリストを作って reverse してもよいが、差分リストを使う手もある。いずれの場合も、作成中に大きさを管理できるかもしれない。

[Word8]の大きさがあらかじめ分かっているなら、pack を使うよりも、create して自前でリストの内容をバッファーへコピーする方がいいだろう。

Builder を使う

三つ目は、Builder を使う方法である。

二つ目の方法は、単一の Word8 という型だけを使う場合に有効であった。一方、Builder は、さまざまな型のデータから作れる。

詳しい理由は省略するが、

  • GHC 7.6.x 以下では blaze-builder の Builder
  • GHC 7.8.x 以上では bytestring 自体の Builder

を使うといいだろう。

Builder は差分リストの一種と言われているが、それは誤解を招きやすい。Builder とは、与えられたデータをバッファーへコピーする関数であり、Builder をくっつけることは、関数合成をすることである。この関数は継続を用いて実装されているので、最終的に走らせた場合、「左から右へ」バファーのコピー関数が走る。(差分リストのように、右結合になる訳ではない。)

Builder には、最終的な大きさを予測できないという問題がある。だから、以下のように使われる。

  • バッファーを1つ用いて、出力をする。すなわち、バッファーが溢れそうになると書き出し、そのバッファーを再利用することを繰り返す。具体的には、blaze-builder では toByteStringIO を使えばよい。
  • バッファーが溢れたそうになったら、新たなバッファーを用意して利用する。Strict な ByteString が複数できることになり、それはすなわち Lazy ByteString である。具体的には、blaze-builder では toLazyByteString を使う。

mallocByteString

create などで使われている mallocByteString を調べてみよう。

mallocByteString :: Int -> IO (ForeignPtr a)
mallocByteString = mallocPlainForeignPtrBytes

mallocPlainForeignPtrBytes は、GHC.GHC.ForeignPtr にある。

mallocPlainForeignPtrBytes :: Int -> IO (ForeignPtr a)
mallocPlainForeignPtrBytes (I# size) = IO $ \s ->
    case newPinnedByteArray# size s      of { (# s', mbarr# #) ->
       (# s', ForeignPtr (byteArrayContents# (unsafeCoerce# mbarr#))
                         (PlainPtr mbarr#) #)
     }

newPinnedByteArray# は、その名が示すようにピン止めされたデータを割り当てる。コピーCGで動かされることのないデータである。つまり、ByteString のバッファーは常に固定されているので、GC を気にせずに読み書きしてよい。

newPinnedByteArray# は、rts/PrimOps.cmm:stg_newPinnedByteArrayzh のことであり、allocatePinned を呼んでいる。

allocatePinned は rts/sm/Storage.c にあり、以下のように C で実装されている。

StgPtr
allocatePinned (Capability *cap, W_ n)
{
    StgPtr p;
    bdescr *bd;

    // If the request is for a large object, then allocate()
    // will give us a pinned object anyway.
    if (n >= LARGE_OBJECT_THRESHOLD/sizeof(W_)) {
        p = allocate(cap, n);
        Bdescr(p)->flags |= BF_PINNED;
        return p;
    }
    ...

LARGE_OBJECT_THRESHOLD/sizeof(W_) 以上の大きさの場合、allocate を呼び出すことが分かる。allocate は同じファイルに存在し、コードを読むと必ずロックするのが分かる。

StgPtr
allocate (Capability *cap, W_ n) {
...
        ACQUIRE_SM_LOCK
        bd = allocGroup(req_blocks);
        dbl_link_onto(bd, &g0->large_objects);
        g0->n_large_blocks += bd->blocks; // might be larger than req_blocks
        g0->n_new_large_words += n;
        RELEASE_SM_LOCK;
...

つまり、ある大きさ以上のバッファーを持つ ByteString を作成する場合は、必ずグローバルなロックが利用されるのだ。このことは、特に +RTS -N を指定して複数のコアを使うようにしたアプリケーションの性能に著しく影響を及ぼす。

定義をたくさん調べて LARGE_OBJECT_THRESHOLD/sizeof(W_) を計算してみると、以下のようになる。

注:Simon Marlow 先生の指摘により数値を修正。

  • 32ビット環境では 3276 バイト(819ワード)以上の ByteString を作成する場合グローバルロックが使われる
  • 64ビット環境では 3272 バイト(409ワード)以上の ByteString を作成する場合グローバルロックが使われる

大きな ByteString の作成には慎重になること。createAndTrim に 4,096 を渡したのに、できあがった ByteString の大きさは 3,300 なんてのは最悪だ。そんな関数の例としては、Network.Socket.ByteString.recv がある。