動的配列の無難なHaskell実装

qiita.com

C++、Rust、Pythonなど、他の言語では当たり前のように多用される動的配列だが、Haskell実装は(開発を始めた時点では)見当たらなかったので作ってみたお話*1

動的配列とはミュータブルな配列の一種で、通常の配列操作だけでなく、末尾への要素の追加・削除が定数時間で行える構造である。確保しておいた領域がいっぱいになったら、その2倍の大きさの領域を確保するという方法によって、漸近的には要素の追加は定数時間となる。

内部の配列には、デファクトスタンダードであるvectorパッケージを用いる。Vectorには無印(boxed)、Unboxed、Storableの三種類の変種があり、それぞれ以下のように使い分ける。

  • 無印(Data.Vector): サンク含め、任意のHaskellのオブジェクトを格納できる。Traversableなどのインスタンスであり、汎用性が高い
  • Unboxed (Data.Vector.Unboxed): ポインタを経由せず、要素がメモリ上に連続的に配置される。タプルはfst側の配列とsnd側の配列の二つが内部で作られる。効率が求められるときに使う
  • Storable (Data.Vector.Storable): Storableクラスを通してメモリ上の表現と対応させる。外部のライブラリと連携したい際に有用

それぞれにミュータブル版(.Mutable)のモジュールも存在する。幸運なことに、これらのAPIを共通化するData.Vector.Genericというモジュールがあるので、それを土台として実装していく。完成したのがこれだ。

github.com

hackage.haskell.org

実装のポイント

vectorAPIは、primitiveパッケージによって抽象化されており、IOとSTの違いを吸収できる。こちらもそれに合わせて設計する。

data GVState v s a = GVState !Int !(v s a)

-- | 'Growable' is a dynamic vector based on mutable vector @v@.
newtype Growable v s a = Growable (MutVar s (GVState v s a))

Vectorの実装を表すパラメータv、世界の型を表すパラメータs、要素の型aを引数に持つGrowableを定義する。IORefの一般化であるMutVarに格納し、配列を入れ替える挙動を可能にする。 配列の末尾に要素を追加する操作pushを定義してみよう。

-- | Append an element to the vector (not atomic).
push :: (PrimMonad m, MVector v a) => Growable v (PrimState m) a -> a -> m ()
push (Growable ref) val = do
  GVState len vec <- readMutVar ref
  vec' <- if MV.length vec == len
    then MV.unsafeGrow vec (max 1 len)
    else pure vec
  writeMutVar ref $ GVState (len + 1) vec'
  MV.write vec' len val
{-# INLINE push #-}

定義について特筆すべきことはないが、ここではINLINE化が重要になってくる。インライン化しないと、PrimMonadMVectorも辞書渡しになってしまい、実行時のパフォーマンスに大きな影響を及ぼす。プリミティブな動作を定義するコードでは特に影響を受けやすいため、辞書渡しされないようしっかり対策を忘れないようにしたい。

まとめ

Haskellのエコシステムは、ありそうなものがないということが多々ある。自作する際も、ベストプラクティスをしっかりと押さえ丁寧に実装していきたい。

なお、当初はアトミック(複数のスレッドが同時に実行しても干渉せず正しく動く)なpushを実装する予定だった。アルゴリズムはさほど難しいものではないが*2、BoxedおよびUnboxed Vectorに対するCASを実装する必要があり、時間がなかったため断念した。しかし、役に立つ場面は必ず出てくるはずなのでいずれ実装したい。

追記 バージョン0.0.1で、ロックフリー版のatomicPushとatomicPopを追加した。

*1:改めて確認すると、grow-vector: Mutable vector with efficient appendsというパッケージが最近アップロードされたようだ

*2:Damian Dechev, Peter Pirkelbauer, and Bjarne Stroustrup - Lock-free Dynamically Resizable Arrays