unboxing-vectorの紹介:newtypeフレンドリーなunboxed vector


要約

最近[いつ?]のHaskellはSafe CoercionやDerivingVia等が登場し、newtype活用の機運が高まっている。

ところで、高性能なHaskellコードを書くには Data.Vector.UnboxedData.Array.Unboxed にあるようなunboxed vector / unboxed arrayの活用が不可欠だが、これらはnewtypeとの相性が悪いという問題を抱えている。

そこで、newtypeフレンドリーなunboxed vectorのライブラリーを作成した。

実態としてはunboxed vectorなのだが、ライブラリーの名前を unboxed-vector にしてしまうと既存の Data.Vector.Unboxed と紛らわしいので、少しもじって unboxing-vector とした。


問題

mod $10^9+7$ で計算するデータ型を作りたいとする。

newtype IntMod = IntMod Int64 deriving (Eq, Show)

instance Num IntMod where
IntMod x + IntMod y = IntMod ((x + y) `mod` (10^9+7))
IntMod x - IntMod y = IntMod ((x - y) `mod` (10^9+7))
...

このデータ型を普通に使ったり、リストに入れるのはいい。普通の Vector に入れるのも良い。

しかし、unboxed vectorに入れようとした時に問題が起こる。

import qualified Data.Vector.Unboxed as U

main = do
print (U.singleton (IntMod 42))
-- No instance for (U.Unbox IntMod) arising from a use of ‘print’

Data.Vector.Unboxed を使うには、要素の型は Unbox クラスのインスタンスである必要があるのだ。

幸い、Haskell (GHC) にはGeneralizedNewtypeDerivingという便利な機能がある。これを使うと、Unbox Int64 のインスタンスを使って Unbox IntMod の定義を自動的に導出してくれるはずだ。やってみよう。

{-# LANGUAGE DerivingStrategies #-}

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
import Data.Int
import qualified Data.Vector.Unboxed as U

newtype IntMod = IntMod Int64
deriving (Eq,Show)
deriving newtype (U.Unbox)

すると今度は、

    • No instance for (Data.Vector.Generic.Base.Vector U.Vector IntMod)

arising from the 'deriving' clause of a data type declaration
Possible fix:
use a standalone 'deriving instance' declaration,
so you can specify the instance context yourself

というエラーが出る。

Unboxクラスの定義を見てみよう。

class ( Data.Vector.Generic.Vector Data.Vector.Unboxed.Vector a

, Data.Vector.Generic.Mutable.MVector Data.Vector.Unboxed.Mutable.MVector a
)
=> Unbox a

Unbox クラスは、メソッドやassociated type familyを持たない。つまり、 Unbox a 制約は

( Data.Vector.Generic.Vector Data.Vector.Unboxed.Vector a

, Data.Vector.Generic.Mutable.MVector Data.Vector.Unboxed.Mutable.MVector a
)

制約の別名と言っても良い。

(以下、 Data.Vector.GenericG, Data.Vector.Generic.MutableGM と略記する)

G.Vector クラスや GM.MVector クラスというのは、 Vector 型の要素に対する操作を実装するクラスである。

module Data.Vector.Generic where

class Vector v a where
basicLength :: v a -> Int
basicUnsafeSlice :: Int -> Int -> v a -> v a
basicUnsafeFreeze :: PrimMonad m => Mutable v (PrimState m) a -> m (v a)
basicUnsafeIndexM :: Monad m => v a -> Int -> m a
...

これらのクラスを通して各種 Vector 型を統一的に扱うことで、boxed vectorだのunboxed vectorだのstorable vectorだのを実装する時に mapfold のような各種操作や stream fusion のような機構をいちいち再実装しなくてよくなっている。

じゃあ G.Vector U.Vector Int64 の実装から G.Vector U.Vector IntMod の実装を導出してやればいいじゃん!という発想に至るわけだが、残念なことに G.Vector クラスは GeneralizedNewtypeDerivingが使えない クラスなのである。

実際、

{-# LANGUAGE DerivingStrategies #-}

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}
import Data.Int
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM

newtype instance U.Vector IntMod = V_IntMod (U.Vector Int64)
newtype instance UM.MVector s IntMod = MV_IntMod (UM.MVector s Int64)

newtype IntMod = IntMod Int64
deriving (Eq,Show)
deriving newtype (U.Unbox, G.Vector U.Vector, GM.MVector UM.MVector)

というコードをコンパイルすると

    • Couldn't match representation of type ‘m Int64’

with that of ‘m IntMod’
arising from the coercion of the method ‘G.basicUnsafeIndexM’
from type ‘forall (m :: * -> *).
Monad m =>
U.Vector Int64 -> Int -> m Int64’
to type ‘forall (m :: * -> *).
Monad m =>
U.Vector IntMod -> Int -> m IntMod’
NB: We cannot know what roles the parameters to ‘m’ have;
we must assume that the role is nominal

というエラーが出る(ここに書いたエラーは代表例で、実際は他にもたくさんエラーが出る)。

ここで問題となっている basicUnsafeIndexM のシグニチャーを見ると、

  basicUnsafeIndexM :: Monad m => v a -> Int -> m a

となっている。問題となるのは U.Vector Int64 -> Int -> m Int64 という型の関数を U.Vector IntMod -> Int -> m IntModcoerce できるかという話で、 m のtype roleがわからないとこれは不可能である。

実際のところ、我々が使うモナドの99.9%はtype roleがrepresentationalで、 m Int64m IntModcoerce しても問題ないはずだが、残りの0.1%のモナドはGADTs等によってtype roleがnominalとなっているかもしれないので、任意のモナド m に対して適用できる関数を coerce することはできない。

というわけで、Data.Vector.Unboxedのunboxed vectorはnewtypeと相性が悪い

自分で定義したnewtypeに対してunboxed vectorを使えるようにするには、次のように書く必要がある:

{-# LANGUAGE TypeFamilies #-}

{-# LANGUAGE MultiParamTypeClasses #-}
import Data.Int
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import Data.Coerce

newtype IntMod = IntMod Int64
deriving (Eq,Show)
-- deriving newtype (U.Unbox, G.Vector U.Vector, GM.MVector UM.MVector)

newtype instance U.Vector IntMod = V_IntMod (U.Vector Int64)
newtype instance UM.MVector s IntMod = MV_IntMod (UM.MVector s Int64)

instance GM.MVector UM.MVector IntMod where
basicLength (MV_IntMod mv) = GM.basicLength mv
basicUnsafeSlice i l (MV_IntMod mv) = MV_IntMod (GM.basicUnsafeSlice i l mv)
basicOverlaps (MV_IntMod mv) (MV_IntMod mv') = GM.basicOverlaps mv mv'
basicUnsafeNew l = MV_IntMod <$> GM.basicUnsafeNew l
basicInitialize (MV_IntMod mv) = GM.basicInitialize mv
basicUnsafeReplicate i x = MV_IntMod <$> GM.basicUnsafeReplicate i (coerce x)
basicUnsafeRead (MV_IntMod mv) i = coerce <$> GM.basicUnsafeRead mv i
basicUnsafeWrite (MV_IntMod mv) i x = GM.basicUnsafeWrite mv i (coerce x)
basicClear (MV_IntMod mv) = GM.basicClear mv
basicSet (MV_IntMod mv) x = GM.basicSet mv (coerce x)
basicUnsafeCopy (MV_IntMod mv) (MV_IntMod mv') = GM.basicUnsafeCopy mv mv'
basicUnsafeMove (MV_IntMod mv) (MV_IntMod mv') = GM.basicUnsafeMove mv mv'
basicUnsafeGrow (MV_IntMod mv) n = MV_IntMod <$> GM.basicUnsafeGrow mv n

instance G.Vector U.Vector IntMod where -- needs MultiParamTypeClasses here
basicUnsafeFreeze (MV_IntMod mv) = V_IntMod <$> G.basicUnsafeFreeze mv
basicUnsafeThaw (V_IntMod v) = MV_IntMod <$> G.basicUnsafeThaw v
basicLength (V_IntMod v) = G.basicLength v
basicUnsafeSlice i l (V_IntMod v) = V_IntMod (G.basicUnsafeSlice i l v)
basicUnsafeIndexM (V_IntMod v) i = coerce <$> G.basicUnsafeIndexM v i
basicUnsafeCopy (MV_IntMod mv) (V_IntMod v) = G.basicUnsafeCopy mv v
elemseq (V_IntMod v) x y = G.elemseq v (coerce x) y

instance U.Unbox IntMod

(筆者はHaskellでAtCoderに参戦しているが、「109+7で割った余りを求める」系の問題で Vector を使う必要があった場合はこの30行近いコードをいちいちコピペしている。例:https://atcoder.jp/contests/abc130/submissions/5997298

このインスタンス定義をTemplate Haskellを使って生成する vector-th-unbox というパッケージが存在するが、Template Haskellは使わずに済ませられるなら使わないに越したことはない

Template Haskellを使わないで、自分で定義したデータ型に対してunboxed vectorを使えるようにする簡単な方法はないか?


解決

「コンピューターサイエンスのあらゆる問題はanother level of indirectionをかませば解決できる」との格言の通り、「unboxed vectorとnewtypeの相性が悪い」という問題もanother level of indirectionをかますと解決できる。

つまり、 Data.Vector.Unboxed.Vector をラップする別の Vector 型をうまく作れば、その Vector 型が newtype との相性をよくできる。

unboxing-vector.png

「newtypeとの相性が良い」unboxed vector型の定義は大雑把に書くと次のようになる:

module Data.Vector.Unboxing

( Unboxable(..)
, Vector -- データ構築子 UnboxingVector は公開しない
, ... -- MVectorとか
) where
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G
import Data.Coerce

class (U.Unbox (Rep a), Coercible a (Rep a)) => Unboxable a where
-- Rep a は a の newtype もとを表す型
type Rep a

-- ここで定義する Vector 型は Unboxed.Vector のラッパー
newtype Vector a = UnboxingVector (U.Vector (Rep a))

instance Unboxable a => G.Vector Vector a where
basicLength (UnboxingVector v) = G.basicLength v
basicUnsafeIndexM (UnboxingVector v) i = coerce <$> G.basicUnsafeIndexM v i
-- 以下同様に、 G.Vector U.Vector (Rep a) のインスタンスと coerce を使ってうまいこと変換する

-- MVector も同様に定義する

instance Unboxable Int64 where
type Rep Int64 = Int64
instance Unboxable Double where
type Rep Double = Double
instance Unboxable a => Unboxable (Sum a) where
type Rep (Sum a) = Rep a

使う側は、

newtype IntMod = IntMod Int64

deriving (Eq,Show)

instance Unboxable IntMod where
type Rep IntMod = Int64

とするか、あるいはGNDを使って

{-# LANGUAGE DerivingStrategies #-}

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
import Data.Int
import qualified Data.Vector.Unboxing as U

newtype IntMod = IntMod Int64
deriving (Eq,Show)
deriving newtype (U.Unboxable)

としても良い。GNDを使う場合は残念ながらUndecidableInstances拡張が必要となる。


Coercibleのリークを防ぐ

さっき定義した IntMod 型には不変条件がある。値 IntMod n に入っている整数 n は $0 \le n < 10^9+7$ を満たさなければならない。

IntMod を定義するモジュール自身はこの不変条件を守るのは当然だが、よそのモジュールに IntMod (10^9+100) という値を作られてしまうと不変条件が破られてしまう。

よそのモジュールにも不変条件を守らせるには、 IntMod 型を定義するモジュールが IntMod のデータ構築子を非公開とすれば良い(抽象データ型)。

さて、抽象データ型とunboxed vectorを組み合わせることはできるだろうか。

もしも Unboxed.Vector Int64Unboxing.Vector Int64 から Unboxing.Vector IntMod への変換が可能だったら、 coerce (Unboxed.singleton (10^9+100) :: Unboxed.Vector Int64) :: Unboxing.Vector IntMod という風にして不正な IntMod 値を作ることができてしまう。

それを防ぐには、Data.Vector.Unboxing.Vector のデータ構築子 UnboxingVector を隠す必要がある。さっき書いた Data.Vector.Unboxing のコードは実際にそうしている。

果たして、これで十分だろうか?

実は、先ほどの Unboxable クラスの定義には問題があり、 Unboxable IntMod 制約を使って Int64IntMod をキャストできてしまう。

以下のコードを考えよう:

import qualified Data.Vector.Unboxing as U

import IntMod

-- U.Unboxable a の制約として Coercible a (Rep a) があることに注意
coerceInt64 :: (U.Unboxable a, Int64 ~ U.Rep a) => Int64 -> a
coerceInt64 = coerce

invalidVal :: IntMod
invalidVal = coerceInt64 (10^9+10)

U.Unboxable aCoercible a (Rep a) を制約に持つので、 coerceInt64 からは Coercible a Int64 が見える。

つまり、 IntMod モジュールが IntMod のデータ構築子を公開しなくても、 Unboxable IntMod インスタンスを定義することによって間接的に IntMod のデータ構築子を公開してしまっているということになる。

これはまずいので、 Unboxable の定義を変えなくてはならない。

結論を書くと、次のように変えれば良い:

{-# LANGUAGE MultiParamTypeClasses, TypeFamilies #-}

{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Vector.Unboxing
( Unboxable(Rep)
, Vector -- データ構築子 UnboxingVector は公開しない
, ... -- MVectorとか
) where
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G
import Data.Coerce
import Data.Type.Coercion

class (U.Unbox (Rep a)) => Unboxable a where
type Rep a
coercion :: Coercion a (Rep a)
default coercion :: Coercible a (Rep a) => Coercion a (Rep a)
coercion = Coercion

newtype Vector a = UnboxingVector (U.Vector (Rep a))

instance Unboxable a => G.Vector Vector a where
basicLength (UnboxingVector v) = G.basicLength v
basicUnsafeIndexM (UnboxingVector v) i = case coercion @a of Coercion -> coerce <$> G.basicUnsafeIndexM v i
-- 以下同様に、 G.Vector U.Vector (Rep a) のインスタンスと coerce を使ってうまいこと変換する

-- MVector も同様に定義する

instance Unboxable Int64 where
type Rep Int64 = Int64
instance Unboxable Double where
type Rep Double = Double

Unboxable クラスには Coercible a (Rep a) を制約として持たせるのではなく、 Coercion a (Rep a) 型を持つメソッドとして持たせる。そのメソッドは外部には公開しないので、 Coercible a (Rep a) がリークすることはない。

Data.Type.CoercionCoercion 型というのは、 Coercible 制約を封じ込めた GADTs である)

Unboxable クラスを使って G.Vector クラスを実装する側は、 coercion メソッドを呼んでパターンマッチをすることで、 Coercible a (Rep a) 制約を取り出すことができる。

coercion メソッドが隠されてしまうと Unboxable クラスのインスタンスを定義する側が困るのではないかと思われるかもしれないが、 coercion メソッドのデフォルト実装があるので問題ない。(むしろ、デフォルト実装しか使えない)

なお、新しい Unboxable クラスは相変わらずGeneralizedNewtypeDerivingが適用できる。


newtype以外の型への一般化

newtypeじゃないもっと複雑なデータ型、例えば

data Foo = Foo !Int !Double

みたいなデータ型をunboxed vectorに突っ込みたいとしよう。タプル (Int, Double) はunboxed vectorに突っ込めるのだから、それと同型な Foo 型をunboxed vectorに突っ込めない道理はない。

そこで、さっきの Unboxable クラスの定義を一般化して、newtypeじゃないデータ型にも使えるようにする。

class (U.Unbox (Rep a)) => Unboxable a where

type Rep a
from :: a -> Rep a
to :: Rep a -> a
default from :: Coercible a (Rep a) => a -> Rep a
from = coerce
default to :: Coercible a (Rep a) => Rep a -> a
to = coerce

instance Unboxable Foo where
type Rep Foo = (Int, Double)
from (Foo x y) = (x, y)
to (x, y) = Foo x y

ただ、この定義には2つの問題がある:


  1. 同型なタプルとの変換関数を手書きするのは面倒である。


  2. from/to を公開してしまうと、抽象データ型の中身に外からアクセスできてしまう。


  3. from/to を非公開にすると、 instance Unboxable Foo みたいなやつを書けない。


変換関数の自動化

まず、データ型の定義が与えられた時にそれと同型なタプルとの変換関数を書くというのは、いかにもGenericsの出番という感じがする。

この場合、

class Unboxable' f where

type Rep' f
from' :: f x -> Rep' f
to' :: Rep' f -> f x

genericFrom :: (GHC.Generics.Generic a, U.Unbox (Rep' (GHC.Generics.Rep a)), Unboxable' (GHC.Generics.Rep a))
=> a -> Rep' (GHC.Generics.Rep a)
genericTo :: (GHC.Generics.Generic a, U.Unbox (Rep' (GHC.Generics.Rep a)), Unboxable' (GHC.Generics.Rep a))
=> Rep' (GHC.Generics.Rep a) -> a

という風な関数を書いて

data Foo = Foo Int Double deriving (GHC.Generics.Generic)

instance Unboxable Foo where
type Rep Foo = Rep' (GHC.Generics.Rep Foo)
from = genericFrom
to = genericTo

という風に変換関数の定義を自動化できる。


メソッドを隠すかどうか

from/to を隠すかどうかについて、やり方はいくつか考えられる。


メソッドを隠す方法

ユーザーが Unboxable のインスタンスを定義したい場合には


  • default定義

  • GeneralizedNewtypeDeriving

  • DerivingVia

のいずれかを使ってもらう。

DerivingViaに関しては、例えば

newtype Generics a = Generics a

instance (GHC.Generics.Generic a, U.Unbox (Rep' (GHC.Generics.Rep a)), Unboxable' (GHC.Generics.Rep a)) => Unboxable (Generics a) where
type Rep (Generics a) = Rep' (GHC.Generics.Rep a)
from (Generics x) = genericFrom x
to y = Generics (genericTo y)

というnewtypeを用意して、データ型の定義で

data Foo = ...

deriving Generic
deriving Unboxable via (Generics Foo)

とする。


メソッドを隠さない方法

要は from/to を外から呼べなければ良いので、

module Data.Vector.Unboxing (ConversionFn, mkConversionFn, ...) where

newtype ConversionFn a b = ConversionFn (a -> b) -- データ構築子は隠す(抽象データ型)
mkConversionFn :: (a -> b) -> ConversionFn a b
mkConversionFn = ConversionFn

class Unboxable a where
type Rep a
from :: ConversionFn a (Rep a)
to :: ConversionFn (Rep a) a

とする。


現状のunboxing-vectorでは「メソッドを隠す方法」を採用している。つまり、 Unboxable クラスのインスタンスを書くには


  • default定義(coerce)を使う

  • GeneralizedNewtypeDerivingを使う


    • 対応:GHC 8.2以降(GHC 8.0まではassociated type familyを持つ型クラスのGNDに非対応なので)



  • DerivingViaを使う


    • 対応:GHC 8.6以降



のいずれかを使う必要がある。

古いGHCでもunboxing-vectorを使いたいという要望があれば「メソッドを隠さない方法」を検討したい。要望は筆者に直接連絡するか、GitHub Issuesに書いて欲しい。