Rust でトレイトオブジェクトと enum のディスパッチ速度比較

はじめに

Rust で同一のインターフェイスを持つオブジェクトを単一の型に収めてオブジェクトごとに違う処理を呼び出したいケースを考えます。
そのときに考えられる方法が、トレイトオブジェクトを使う方法と、enum を使う方法です。

トレイトオブジェクトを使うと以下のような感じです。
構造体 A と B を、 Box<dyn Trait> という同じ型に入れることができ、その状態でメソッド f() を呼び出すことができます。
(この記事ではトレイトオブジェクトを Box に入れた場合のみを扱います)

trait Trait {
    fn f(&self) -> usize;
}

struct A(usize);
struct B(usize, usize);

impl Trait for A {
    fn f(&self) -> usize {self.0}
}
impl Trait for B {
    fn f(&self) -> usize {self.0 + self.1}
}

let a: Box<dyn Trait> = Box::new(A(1));
let b: Box<dyn Trait> = Box::new(B(1, 2));
a.f();
b.f();

一方、enum を使うと以下のようになります。

#[derive(Clone)]
enum Enum {
    A(usize),
    B(usize, usize),
}

impl Enum {
    fn f(&self) -> usize {
        match self {
            Enum::A(a) => *a,
            Enum::B(a, b) => a + b,
        }
    }
}

let a: Enum = Enum::A(1);
let b: Enum = Enum::B(1, 2);
a.f();
b.f();

この２つの方法でメソッド呼び出しの速度を比較しました。

計測用コード(トレイトオブジェクト)

リアリティのために構造体を3つ用意しています。

(補助的に Clone を実装しましたが本質でないです。)

trait Trait {
    fn f(&self) -> usize;
    fn box_clone(&self) -> Box<dyn Trait>;
}

impl Clone for Box<dyn Trait> {
    fn clone(&self) -> Box<dyn Trait> {
        self.box_clone()
    }
}

#[derive(Clone)]
struct A(usize);
#[derive(Clone)]
struct B(usize, usize);
#[derive(Clone)]
struct C(usize, usize, usize);

impl Trait for A {
    fn f(&self) -> usize {self.0}
    fn box_clone(&self) -> Box<dyn Trait> {Box::new(self.clone())}
}
impl Trait for B {
    fn f(&self) -> usize {self.0 + self.1}
    fn box_clone(&self) -> Box<dyn Trait> {Box::new(self.clone())}
}
impl Trait for C {
    fn f(&self) -> usize {self.0 + self.1 + self.2}
    fn box_clone(&self) -> Box<dyn Trait> {Box::new(self.clone())}
}

fn main() {
    let v: Vec<Box<dyn Trait>> = vec![Box::new(A(1)) as Box<dyn Trait>, Box::new(B(1, 2)) as Box<dyn Trait>, Box::new(C(1, 2, 3)) as Box<dyn Trait>].into_iter().cycle().take(1000000).collect();
    let t = std::time::Instant::now();
    let mut acc = 0;
    for a in v {
        acc += a.f();
    }
    dbg!(t.elapsed());
    dbg!(acc);
}

計測用コード(enum)

#[derive(Clone)]
enum Enum {
    A(usize),
    B(usize, usize),
    C(usize, usize, usize),
}

impl Enum {
    fn f(&self) -> usize {
        match self {
            Enum::A(a) => *a,
            Enum::B(a, b) => a + b,
            Enum::C(a, b, c) => a + b + c,
        }
    }
}

{
    let v: Vec<Enum> = vec![Enum::A(1), Enum::B(1, 2), Enum::C(1, 2, 3)].into_iter().cycle().take(1000000).collect();
    let t = std::time::Instant::now();
    let mut acc = 0;
    for a in v {
        acc += a.f();
    }
    dbg!(t.elapsed());
    dbg!(acc);
}

結果

リリースビルドで３回ずつ計測しました。

trait object: 13.4895ms
trait object: 13.0611ms
trait object: 12.4761ms
enum:          4.7434ms
enum:          2.7366ms
enum:          2.2118ms

見ての通り enum の圧勝です。

トレイトオブジェクトを使うと関数呼び出しは動的ディスパッチになるため遅いですが、こんなに差があるのですねー。

enum のコードはずるいのでは？

実は enum のコードは最適化によってインライン展開されていたので関数呼び出しになっていません(…多分)。

つまり、この部分は:

    for a in v {
        acc += a.f();
    }

コンパイル時にこう展開されていたのです:

    for a in v {
        acc += match a {
            Enum::A(a) => *a,
            Enum::B(a, b) => a + b,
            Enum::C(a, b, c) => a + b + c,
        };
    }

なんかフェアでない気がします。

インライン展開なしで enum 再計測

コードをこう修正しました。

...

impl Enum {
    #[inline(never)] // <- これを追加
    fn f(&self) -> usize {
        match self {
            Enum::A(a) => *a,
            Enum::B(a, b) => a + b,
            Enum::C(a, b, c) => a + b + c,
        }
    }
}

...

これでインライン展開されないはず。

結果2

enum2 が再計測の結果です。

trait object: 13.4895ms
trait object: 13.0611ms
trait object: 12.4761ms
enum:          4.7434ms
enum:          2.7366ms
enum:          2.2118ms
enum2:         4.1026ms
enum2:         4.1718ms
enum2:         4.1466ms

若干遅くなりました。やはりインライン展開されていたみたいです。

構造体を enum に入れてみる

こんなケースも考えてみました。
トレイトオブジェクトのときに作った構造体 A, B, C を enum に入れてみます。

...

#[derive(Clone)]
enum EnumOfStruct {
    A(A),
    B(B),
    C(C),
}

impl EnumOfStruct {
    //#[inline(never)] は無しで
    fn f(&self) -> usize {
        match self {
            EnumOfStruct::A(a) => a.f(),
            EnumOfStruct::B(b) => b.f(),
            EnumOfStruct::C(c) => c.f(),
        }
    }
}

{
    let v: Vec<EnumOfStruct> = vec![EnumOfStruct::A(A(1)), EnumOfStruct::B(B(1, 2)), EnumOfStruct::C(C(1, 2, 3))].into_iter().cycle().take(1000000).collect();
    let t = std::time::Instant::now();
    let mut acc = 0;
    for a in v {
        acc += a.f();
    }
    dbg!(t.elapsed());
    dbg!(acc);
}

結果3

「構造体を enum に入れてみる」の結果です。

trait object: 13.4895ms
trait object: 13.0611ms
trait object: 12.4761ms
enum:          4.7434ms
enum:          2.7366ms
enum:          2.2118ms
enum2:         4.1026ms
enum2:         4.1718ms
enum2:         4.1466ms
enumOfStruct:  2.4725ms
enumOfStruct:  2.3041ms
enumOfStruct:  2.2413ms

おお、素の enum に匹敵する速さです！

このケースもインライン展開が働いてます。おそらく内部的に enum と同じレベルに最適化されているのではと思います。

〆

動的ディスパッチは遅いよ。