Scalaのspecializedアノテーションを使いこなすための基礎知識

こんにちは、アドテクエンジニアーのトデス子です。ふだんスカラを使っているのでスカラの話をします。

ScalaはJavaと同様、型パラメータを使用したコードは内部的にObject型を通して使用されます。 そのため、IntやDoubleといったプリミティブ型を指定した場合は boxing/unboxingのオーバーヘッドが発生します。

このオーバヘッドは多くの場合大した問題になりませんが、数値計算などの特定領域においては パフォーマンスのボトルネックになるケースがあります。

Scalaにおいては、@specializedアノテーションを使用することでこのオーバヘッドを軽減する機構があります。 この記事では、この機構の詳細と使用時の注意点などについて紹介します。

はじめに

@specializedアノテーションについては公式なドキュメントがかなり乏しく、また将来的に挙動が変更される可能性があります。 この記事ではScala 2.11.7 時点での動作について解説します(とは言っても、バイナリ互換性の関係から2.11系のうちは有効、 2.12系でも変更されたという噂はきかないのでおそらく有効な知識だと思われます)。

クラスファイルのデコンパイルにはJADを使用しました。

@specializedアノテーション基礎

Scalaの仕様において、@specialiedは次のように解説されています:

@specialized When applied to the definition of a type parameter, this annotation causes the compiler to generate specialized definitions for primitive types. An optional list of primitive types may be given, in which case specialization takes into account only those types.

(コード例略)

Whenever the static type of an expression matches a specialized variant of a definition, the compiler will instead use the specialized version. See the specialization sid for more details of the implementation.

意訳すると、型パラメータに@specializedアノテーションを適用することで、コンパイラはプリミティブ型に特化した定義を生成するようになります。アノテーションにプリミティブ型のリストを渡すことで、特殊化を特定の型に限定することができます。 式の静的な型が特殊化された定義と一致する場合、コンパイラは特殊化バージョンを使用します。詳細はSID #9を見てねといったところでしょうか。

SID #9はScala2.8時代に書かれたspecializeに関する仕様のドラフトですが、基本的な仕様は今と変わっていないようです。 上記の記述に付け加えることとして、特殊化されたクラスの実装についてやメソッドが特殊化される条件 (引数や返り値の型において、specializedされた型パラメータが単独/Array型として出現)について触れられています。

特殊化されたクラスの構造

class NotSpecialized[A](val x: A)

class Specialized[@specialized A](val x: A)

単純な例から見ていきましょう(以下、JavaコードはコンパイルされたクラスファイルをJADにより逆コンパイルしたものです)

public class NotSpecialized
{

    private final Object x;

    public Object x()
    {
        return x;
    }

    public NotSpecialized(Object x)
    {
        this.x = x;
        super();
    }
}

非特殊化バージョンはシンプルですね。型パラメータが消去されてObject型になっています。続いて特殊化バージョンについて:

Specialized.class
Specialized$mcB$sp.class
Specialized$mcC$sp.class
Specialized$mcD$sp.class
Specialized$mcF$sp.class
Specialized$mcI$sp.class
Specialized$mcJ$sp.class
Specialized$mcS$sp.class
Specialized$mcV$sp.class
Specialized$mcZ$sp.class

specializedアノテーションに引数が指定されない場合、 プリミティブ型に対する特殊化が行われます。

Byte, Short, Int, Long, Char, Float, Double, Boolean, Unitの各プリミティブに対して特殊化されたクラスが生成されているのがわかります。 クラス名のB, Cなどのアルファベットは、Javaバイトコードのfield descriptorにおける命名規則に準拠しているようです。

まずは特殊化されていないSpecializedクラス(ややこしい)を見てみましょう:

import scala.runtime.BoxesRunTime;

public class Specialized
{

    public final Object x;

    public Object x()
    {
        return x;
    }

    public boolean x$mcZ$sp()
    {
        return BoxesRunTime.unboxToBoolean(x());
    }

    public byte x$mcB$sp()
    {
        return BoxesRunTime.unboxToByte(x());
    }

    // ...
    // 特殊化されたxの実装がプリミティブの数だけ続くが省略

    public boolean specInstance$()
    {
        return false;
    }

    public Specialized(Object x)
    {
        this.x = x;
        super();
    }
}

Getterメソッドであるxについて、対応するプリミティブの数だけ特殊化バージョンが生成されています。 Objectとして保持されている値を、対応するプリミティブ型としてunboxして返しています。

続いて、Int型に特化したバージョンを見てみます。

import scala.runtime.BoxesRunTime;

public class Specialized$mcI$sp extends Specialized
{

    public final int x$mcI$sp;

    public int x$mcI$sp()
    {
        return x$mcI$sp;
    }

    public int x()
    {
        return x$mcI$sp();
    }

    public boolean specInstance$()
    {
        return true;
    }

    public volatile Object x()
    {
        return BoxesRunTime.boxToInteger(x());
    }

    public Specialized$mcI$sp(int x$mcI$sp)
    {
        this.x$mcI$sp = x$mcI$sp;
        super(null);
    }
}

• 非特殊化バージョンのSpecialiedクラスを継承している
• Int型に対応するメンバx$mcI$spが追加されている
• 元のメンバObject xは未使用(nullが代入される)
• 非特殊化バージョンのメソッドObject x()は、同名のint x()を通じて特殊化バージョンのメソッドx$mcI$sp()に処理を委譲+boxingしている
  • 余談ですが、このメソッドにvolatileがついているのは フラグの値を混同した表示ミスで、 本当はBridgeメソッドのようです。

ということが読み取れます。

メソッドが特殊化される条件

class Specialized[@specialized(Int) A] {
  // 引数に使用
  def param(x: A): Int = ???

  // 返り値に使用
  def ret(): A = ???

  // Arrayとして引数に使用
  def arrayArg(x: Array[A]): Int = ???

  // Arrayとして返り値に使用
  def arrayRet(): Array[A] = ???

  // ↑ 上のメソッドはすべて期待通り特殊化されます

  // Seqとして引数に使用
  // Seq[Int]は特殊化されていない
  def paramSeq(x: Seq[A]): Int = ???

  // Functionとして引数に使用
  // Function[Int, Int]は特殊化可能
  def paramFunction(f: A => A): Int = { f(ret()); 0 }

  // Tuple2として返り値に使用
  // Tuple2[Int, Int]は特殊化可能
  def retTuple2(): (A, A) = (ret(), ret())

  // 内部的に使用
  def internal(): Unit = {
    val x: A = ret()
    param(x)
  }

  // 内部的に使用+返り値
  def retAndInternal(): A =
    ret()
}

引数や返り値に使用するケースは当然のように特殊化されるので略。paramSeqは当然のように特殊化されません。

興味深いのは、paramFunctionのようにAそのものではないがFunction1[A, A]のように特殊化可能な型を引数にとった場合。

// Int特化バージョン
public class Specialized$mcI$sp extends Specialized
    // (略)

    public int paramFunction$mcI$sp(Function1 f)
    {
        f.apply$mcII$sp(ret$mcI$sp());
        return 0;
    }
}

f.applyの呼び出しにもきちんと特殊化バージョンが使われています。

返り値にTuple2を返すretTuple2、返り値内部でメソッドを呼んで結果を返すretAndInternalも、特殊化されたメソッドが呼び出されています。

// class Specialized$mcI$sp
    public int retAndInternal$mcI$sp()
    {
        return ret$mcI$sp();
    }

    public Tuple2 retTuple2$mcI$sp()
    {
        return new scala.Tuple2.mcII.sp(ret$mcI$sp(), ret$mcI$sp());
    }

一方残念なのは内部的に特殊化可能メソッドを呼んでいるinternal()で、これは特殊化されませんでした。

   public void internal()
    {
        Object x = ret();
        param(x);
    }

特殊化されたクラス/メソッドが使用される条件

コンパイル時に型パラメータの値が確定するケースじゃないとだめです。

object Main {
  def gen[A](): Specialized[A] =
    new Specialized[A]

  def genS[@specialized(Int) A](): Specialized[A] =
    new Specialized[A]

  def main(args: Array[String]): Unit = {
    // 特殊化される
    val a = new Specialized[Int]

    // されない
    val b = gen[Int]()

    // される
    val c = genS[Int]()
  }
}

   public Specialized gen()
    {
        return new Specialized();
    }

    public Specialized genS()
    {
        return new Specialized();
    }

    public Specialized genS$mIc$sp()
    {
        return new Specialized.mcI.sp();
    }

    public void main(String args[])
    {
        Specialized a = new Specialized.mcI.sp();
        Specialized b = gen();
        Specialized c = genS$mIc$sp();
    }

クラスを特殊化したうえでメソッドも特殊化する

こういったケースはどうなるか:

class Specialized[@specialized(Int) A, @specialized(Double) B] {
  def foo[@specialized(Char) X, @specialized(Boolean) Y](a: A, b: B, x: X, y: Y) = null
}

public class Specialized
{
    // 実装とbridge methodは略
    public Nothing$ foo(int a, double b, Object x, Object y)

    public Nothing$ foo$mcID$sp(int a, double b, Object x, Object y)

    public Nothing$ foo$mCZc$sp(int a, double b, char x, boolean y)

    public Nothing$ foo$mCZcID$sp(int a, double b, char x, boolean y)
}

• クラスもメソッドも特殊化されていないバージョン: bar
• クラスの型パラメータのみ特殊化されたバージョン: bar$mcID$sp
• クラスとメソッドの型パラメータが特殊化されたバージョンが二つ
  • foo$mCZc$sp
  • foo$mCZcID$sp

メソッドの命名は、{base_name}$m{method_types}c{class_types}$sp形式になっていると推測されます。よく出てくる$mcXXっていうのはこういう意味だったのか！

クラス・メソッド両方について特殊化されたバージョンが二つあるのは謎です。呼び分けについては、以下のように 「クラスとメソッド両方が特殊化可能な場合のみ特殊化メソッドが使われる」という実装になっているようになってるようです。 foo$mCZcID$sp以外は無駄なメソッドが生成されているように見えるのですが、意味があるのかどうかは不明です。

  def main(args: Array[String]): Unit = {
    val sp = new Specialized[Int, Double]
    val nsp = new Specialized[Any, Any]

    sp.foo[Any, Any](1, 1.0, null, null)
    sp.foo[Char, Boolean](1, 1.0, 'c', false)

    nsp.foo[Any, Any](1, 1.0, null, null)
    nsp.foo[Char, Boolean](1, 1.0, 'c', false)
  }

   public void main(String args[])
    {
        Specialized sp = new Specialized.mcID.sp();
        Specialized nsp = new Specialized();
        sp.foo(BoxesRunTime.boxToInteger(1), BoxesRunTime.boxToDouble(1.0D), null, null);
        sp.foo$mCZcID$sp(1, 1.0D, 'c', false);
        nsp.foo(BoxesRunTime.boxToInteger(1), BoxesRunTime.boxToDouble(1.0D), null, null);
        nsp.foo(BoxesRunTime.boxToInteger(1), BoxesRunTime.boxToDouble(1.0D), BoxesRunTime.boxToCharacter('c'), BoxesRunTime.boxToBoolean(false));
    }

複数の型パラメータを特殊化する

class Specialized[@specialized(Int) A, @specialized(Int) B]

注意すべき点として、「一部の型パラメータだけ特殊化されている」というケースには対応していません。オールオアナッシングです。

val a = new Specialized[Int, Int] // 特殊化バージョンが使われる
val b = new Specialized[Int, String] // 通常バージョンが使われる

非primitive型の特殊化

@specializedの引数にはprimitive型のほかにAnyRefも指定できます。

引数未指定の場合はプリミティブ型のみが使われるので、意図せず非特殊化バージョンが使用されるという罠が！

class Specialized[@specialized(Int, AnyRef) A, @specialized(Int) B] {
}

object Main {
  def main(args: Array[String]): Unit = {
    val sp = new Specialized[String, Int]
  }
}

   public void main(String args[])
    {
        Specialized sp = new Specialized.mcLI.sp();
    }

まとめ

• @specializedパラメータを付けることで特定の型に特殊化したクラス/メソッドを生成できる
• 特殊化はオールオアナッシング。一部の型パラメータだけ特殊化したバージョンは生成されない。
• アノテーションでAnyRefを明示的に指定しないとプリミティブ型のみが特殊化の対象になる
• 特殊化されたメソッドは、特殊化されたクラスの上で有効になる
• 特殊化されたバージョンが使用されるのは、特殊化する型パラメータの値が静的に決定可能なとき
• 命名規則は$m{メソッドの型パラメータ}c{クラスの型パラメータ}$sp、型名はJVMのfield descriptor準拠
• 疑問を持ったらクラスファイルを逆コンパイルしろ
• ボクシングのコストが問題になることは稀なので、ふつうはこの機能つかわなくていいです(まずはボトルネックを計測せよ。Integer.valueOfなどが頻繁に出現するようなら検討の余地がある)

命名規則などの詳細は普通知る必要はないのですが、この知識とクラスの自動生成を組み合わせるとさらなる最適化のオポチュニティがあります(次回の予告です)。