Rustで複数のimpl Traitを返す

この記事は [Rustその2 Advent Calendar 2018](https://qiita.com/advent-calendar/2018/rust2) の 21日目の記事です。

## はじめに

先日、[auto_enums]というクレートをリリースしたのですが、このクレートを作った背景を中心に、[auto_enums]がどのような問題を解決できるのかについても書こうと思います。

## 背景

### `impl Trait`について

`impl Trait`はRust 1.26で安定化された機能で、関数の引数と、トレイトメソッド以外の関数の戻り値に使用でき、引数位置で使用された場合は匿名の型引数に、戻り値位置で使用された場合は存在型になります。

`impl Trait`を使用するとクロージャのような匿名の型や、`Iterator`や`Future`などを使用した際に生成される複雑な型を簡潔に書くことができるようになります。また、静的ディスパッチになるので`Box<dyn Trait>`よりもパフォーマンスが良くなります。
この記事ではタイトル通り、戻り値位置での`impl Trait`について扱います。

`impl Trait`の詳細については、以下を参考にしてください。

- RFCs([1522](https://github.com/rust-lang/rfcs/blob/master/text/1522-conservative-impl-trait.md), [1951](https://github.com/rust-lang/rfcs/blob/master/text/1951-expand-impl-trait.md), [2071](https://github.com/rust-lang/rfcs/blob/master/text/2071-impl-trait-existential-types.md))
- [Tracking issue](https://github.com/rust-lang/rust/issues/34511)
- [qnighy氏による解説](https://qnighy.hatenablog.com/entry/2018/01/28/220000)

### `impl Trait`は異なる型を返せない

`impl Trait`は静的ディスパッチなので異なる型を返すことはできません。このコードはコンパイルエラーになります。

```rust

fn foo(x: u32) -> impl Iterator<Item=u32> {
    if x == 0 {
        Some(0).into_iter()
    } else {
        0..x
    }
}
```

また、クロージャは匿名なので、同じことをしていても違う型として扱われます。このコードもコンパイルエラーです。

```rust

fn bar(x: u32) -> impl Iterator<Item=u32> {
    if x == 0 {
        (x..).map(|y| y + 1)
    } else {
        (x..).map(|y| y + 1)
    }
}
```

上の二つの例にあるように、複数の(異なる型の)`impl Trait`はそのままでは返すことが出来ませんが、いくつかの方法でこの問題を解決することができます。

## `Box<dyn Trait>`を使った解決

複数の`impl Trait`を返す方法として、おそらく最も手っ取り早いのはBoxを使用してトレイトオブジェクト(`dyn Trait`)にしてしまうことです。

メリット

- `impl Trait`よりも多くの場所で使用できる
- stableコンパイラで複数のクロージャ(`Fn*` traits)を返す事ができる

デメリット

- `impl Trait`の利点が失われる(パフォーマンスが落ちる)
- ヒープを使用する(stableコンパイラ + `no_std`環境では使用できない)

```rust

fn foo(x: u32) -> Box<dyn Iterator<Item=u32>> {
    if x == 0 {
        Box::new(Some(0).into_iter())
    } else {
        Box::new(0..x)
    }
}
```

トレイトオブジェクトについては[リファレンス](https://doc.rust-lang.org/reference/types/trait-object.html)か[1stエディションのbook](https://doc.rust-lang.org/book/ch17-02-trait-objects.html)を参考にしてください（2ndエディションには[トレイトオブジェクトの説明は含まれない](https://doc.rust-lang.org/book/second-edition/ch17-02-trait-objects.html)ようです）。

## enumを使った解決

Boxを使用する方法は手軽ですが、パフォーマンスの問題があります。enumを使用する方法ではimpl Traitの利点を維持しつつ、複数の型を返すことができます。

Boxを使用する方法(トレイトオブジェクトなので仮想関数テーブルを使用した動的ディスパッチ)では全てを動的ディスパッチになるのに対して、enumを使用する方法ではパターンマッチの部分のみが動的ディスパッチになります。

ちなみに、enumをネストさせても[多くの場合最適化される](https://internals.rust-lang.org/t/pre-rfc-sum-enums/8782/78)ようです。

enumを使った解決には大きく分けて2つのパターンがあります。以下は2つのパターンに共通するメリットとデメリットです。

メリット

- `impl Trait`の利点を損なわない(パフォーマンス、最適化)
- `no_std`で使用できる

デメリット

- [`Fn*` traits](https://github.com/rust-lang/rust/issues/29625)が安定していないので、nightly以外では複数のクロージャを返せない
- ヴァリアントの数(≈型の数)がパフォーマンスに影響する

### `either` クレート (外部クレートが提供するenumを使用する)

一つ目は外部クレートが提供するenumを使用する方法です。この方法では多くの場合[eitherクレート](https://github.com/bluss/either)が使用されます。外部クレートが提供するenumを使用する場合には、前述のメリット・デメリットに加えて以下のメリット・デメリットがあります。

メリット

- 予めいくつかのトレイトを実装している

デメリット

- 予めサポートされているトレイトか、自分で定義したトレイトにしか使用できない<br>
  この例としては、futures v0.3があります。Futureトレイトが標準ライブラリに移動することになった関係で、futures-core クレートで定義されている[Stream](https://github.com/rust-lang-nursery/futures-rs/blob/master/futures-core/src/stream/mod.rs#L88)と[Sink](https://github.com/rust-lang-nursery/futures-rs/blob/04b7233760dd458b87db711b80c8a88e6efe43af/futures-sink/src/lib.rs#L241)は実装できますが、Futureは実装できません(2018/12/20現在では[TODOになっています](https://github.com/rust-lang-nursery/futures-rs/blob/4bfa6137cd282da25cafdc21be96f170f5e12ec6/futures-util/src/future/mod.rs#L348))

```rust

extern crate either;

use either::{Left, Right};

fn foo(x: u32) -> impl Iterator<Item=u32> {
    if x == 0 {
        Left(Some(0).into_iter())
    } else {
        Right(0..x)
    }
}
```

### ユーザー定義のenumを使用する

もう一つは自分でenumを定義して、トレイトを実装するというものです。前述のメリット・デメリットに加えて以下のメリット・デメリットがあります。

メリット

- 対応可能なトレイトが多い

デメリット

- enumの定義に加えてトレイトの実装もする必要があるなど手間がかかる

```rust

enum Enum<A, B> {
    A(A),
    B(B),
}

impl<A, B> Iterator for Enum<A, B>
where
    A: Iterator,
    B: Iterator<Item = A::Item>,
{
    type Item = A::Item;
    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Enum::A(x) => x.next(),
            Enum::B(x) => x.next(),
        }
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        match self {
            Enum::A(x) => x.size_hint(),
            Enum::B(x) => x.size_hint(),
        }
    }
}

fn foo(x: u32) -> impl Iterator<Item=u32> {
    if x == 0 {
        Enum::A(Some(0).into_iter())
    } else {
        Enum::B(0..x)
    }
}
```

これを真面目にやると面倒なので、マクロが使用されることが多いと思います。

宣言的マクロだと[eitherが内部で使用しているマクロ](https://github.com/bluss/either/blob/master/src/lib.rs#L51)や[fuchsia-asyncのFuture用マクロ](https://fuchsia.googlesource.com/garnet/+/sandbox/cramertj/asyncing/public/rust/crates/fuchsia-async/src/lib.rs#46)などが参考になると思います。

手続き的マクロではそれっぽいのがありそうなのですが、この用途に一致するクレートは見つけられなかったので、[derive_utils]というヘルパークレートを作りました~~(`auto_enums`の`#[enum_derive]`手続き的マクロが[derive_utils]を使用して実装されていて、この用途を想定しているのでそっちの方が便利だと思います)~~。

**更新(2018-12-24)**: [derive_utils] が更新で使いやすくなったので追記します。以下のようにトレイト定義を貼り付けるだけでenum用の`derive`を作ることができます。[このderiveが生成するコード](https://github.com/taiki-e/derive_utils#generated-code)は上の例のものとだいたい同じです(リンク先の例ではIterator以外のトレイトも実装しています)。

```rust

extern crate derive_utils;
extern crate proc_macro;

use derive_utils::quick_derive;
use proc_macro::TokenStream;

#[proc_macro_derive(Iterator)]
pub fn derive_iterator(input: TokenStream) -> TokenStream {
    quick_derive! {
        input,
        // トレイト定義
        trait Iterator {
            type Item;
            fn next(&mut self) -> Option<Self::Item>;
            fn size_hint(&self) -> (usize, Option<usize>);
        }
    }
}
```

## 提案されている方法(RFC PR, pre-RFC, etc.)

impl Traitが安定化される前後から既存の方法が抱える問題を解決するためにいくつかの方法が提案されています。
前述の「ユーザー定義のenumによって複数のimpl Traitを返す」というパターンを自動化するようなものや、enumを一般化するようなものが多く提案されています。以下の3つはその中でも特に盛り上がっていた(と私が考える)議論です。

- [RFC for anonymous variant types, a minimal ad-hoc sum type](https://github.com/rust-lang/rfcs/pull/2587)<br>
  `(_ | _)::0`のような構文で匿名の列挙型的なものを作れるようにしたいというPR。
- [Pre-RFC: sum-enums](https://internals.rust-lang.org/t/pre-rfc-sum-enums/8782)<br>
  上のPRの代替案。`enum`キーワードを使用してenumを一般化する感じです。内容自体は下のissueと似ている部分もありますが、新しい属性ではなく新しい構文を使用するものを提案しています。
- [Auto-generated sum types](https://github.com/rust-lang/rfcs/issues/2414)<br>
  上のpre-RFCの元になったissue。様々なアイデアが出ていて、[手続き的マクロによる実験的な実装](https://github.com/Nemo157/impl_sum)などもありました。マーカーになるマクロもしくはキーワードと、属性を使用するというアイデアが多い気がします。また、[auto_enums]はこのissueでの議論を参考に作られています。

他の提案や議論:

- [Extending `impl Trait` to allow multiple return types](https://internals.rust-lang.org/t/extending-impl-trait-to-allow-multiple-return-types/7921)<br>
- [Allowing multiple disparate return types in impl Trait using unions](https://internals.rust-lang.org/t/allowing-multiple-disparate-return-types-in-impl-trait-using-unions/7439)
- [Figure out a way to ergonomically match against impl traits?](https://github.com/rust-lang/rfcs/issues/2261)

## [auto_enums]による解決

[auto_enums]は前述の「ユーザー定義のenumによって複数のimpl Traitを返す」という解決策を自動化するための手続き的マクロを提供します。

```rust

#[macro_use]
extern crate auto_enums;

#[auto_enum(Iterator)]
fn foo(x: u32) -> impl Iterator<Item=u32> {
    if x == 0 {
        Some(0).into_iter()
    } else {
        0..x
    }
}
```

インポート部分(`#[macro_use] extern crate ..`)を除けばたった1行増えただけですが、上の例はコンパイルできます。実際には上の例では以下のようなコードが生成されます。

```rust

#[macro_use]
extern crate auto_enums;

fn foo(x: u32) -> impl Iterator<Item=u32> {
    enum __Enum1<__T1, __T2> {
        __T1(__T1),
        __T2(__T2),
    }

    impl<__T1, __T2> ::std::iter::Iterator for __Enum1<__T1, __T2>
    where
        __T1: ::std::iter::Iterator,
        __T2: ::std::iter::Iterator<Item = <__T1 as ::std::iter::Iterator>::Item>,
    {
        type Item = <__T1 as ::std::iter::Iterator>::Item;
        #[inline]
        fn next(&mut self) -> ::std::option::Option<Self::Item> {
            match self {
                __Enum1::__T1(x) => x.next(),
                __Enum1::__T2(x) => x.next(),
            }
        }
        #[inline]
        fn size_hint(&self) -> (usize, ::std::option::Option<usize>) {
            match self {
                __Enum1::__T1(x) => x.size_hint(),
                __Enum1::__T2(x) => x.size_hint(),
            }
        }
    }

    if x == 0 {
        __Enum1::__T1(Some(0).into_iter())
    } else {
        __Enum1::__T2(0..x)
    }
}
```

実際にはもう少し多くのメソッドを実装する他、enumの名前もコンパイル毎に変化したりするのですが、基本的には「ユーザー定義のenumによって複数の`impl Trait`を返す」方法で行っていた

- enumを定義
- トレイトを実装
- 戻り値をヴァリアントでラッピング

という処理を手続き的マクロが自動的にやってくれます。

また、提案されていた実装の一部とも似たような感じになります。特徴としては新しい構文ではなく属性であることと、トレイトの指定が必要なこと、そして`.into()`のような変換メソッドや`marker!(..)`のようながマーカーが不要なことなどがあります。

### `#[auto_enum]`

前述の生成されたコード例は`#[auto_enum]`が単独で生成しているのではなく、後で説明する`#[enum_derive]`という手続き的マクロと分担して生成しています。
具体的にはトレイトの実装を`#[enum_derive]`が、それ以外は`#[auto_enum]`が担当します。二つに役割が分割されているのは、「構造体に特定のトレイトを実装した複数の型を入れる」などの`#[auto_enum]`がサポートできないケースが存在するためです。

先ほどの例で`#[auto_enum]`が直接生成したのは以下のコードです(残りは`#[enum_derive]`が生成したものです)。

```rust

#[macro_use]
extern crate auto_enums;

fn foo(x: u32) -> impl Iterator<Item=u32> {
    #[enum_derive(Iterator)]
    enum __Enum1<__T1, __T2> {
        __T1(__T1),
        __T2(__T2),
    }

    if x == 0 {
        __Enum1::__T1(Some(0).into_iter())
    } else {
        __Enum1::__T2(0..x)
    }
}
```

### `#[enum_derive]`

`#[enum_derive]`は指定されたトレイトをenumに実装するための手続き的マクロです。`#[derive]`をenum限定で色んなトレイトに使用できるようにした感じです。
サポートされているトレイトについては[ドキュメント](https://docs.rs/auto_enums/0.3/auto_enums/#supported-traits)を見てください。

また、サポートされていないトレイトが指定された場合には、それを`#[derive]`に渡す仕組みになっています。この仕組みによって[auto_enums]は多くのトレイトに対応することができます。

以下は`#[enum_derive]`を単独で使用する例です。

```rust

#[macro_use]
extern crate auto_enums;

#[enum_derive(Iterator)]
enum Enum<A, B> {
    A(A),
    B(B),
}
```

### `#[auto_enum]`の詳細な解説

ここからは`#[auto_enum]`について詳しく説明します。

#### `#[auto_enum]`が使用できる位置

`#[auto_enum]`は以下の3カ所に使用できます。ただし、[`stmt_expr_attributes`](https://github.com/rust-lang/rust/issues/15701)と[`proc_macro_hygiene`](https://github.com/rust-lang/rust/issues/54727)が安定化していないのでnightly以外では関数に対して空の`#[auto_enum]`を使用する必要があります。

- 関数
- 式
- let文

```rust

// 関数
#[auto_enum(Iterator)]
fn func(x: i32) -> impl Iterator<Item=i32> {
    if x == 0 {
        Some(0).into_iter()
    } else {
        0..x
    }
}

// 式
#[auto_enum] // nightlyでは関数への空の属性は不要
fn expr(x: i32) -> impl Iterator<Item=i32> {
    #[auto_enum(Iterator)]
    match x {
        0 => Some(0).into_iter(),
        _ => 0..x,
    }
}

// let文
#[auto_enum] // nightlyでは関数への空の属性は不要
fn let_binding(x: i32) -> impl Iterator<Item=i32> {
    #[auto_enum(Iterator)]
    let iter = match x {
        0 => Some(0).into_iter(),
        _ => 0..x,
    };
    iter
}
```

#### サポートされる構文

**編集(2018-12-24)**: [auto_enums] の新しいドキュメントにあわせて全面的に修正しました。

- `if`と`match`

    各分岐の戻り値をヴァリアントでラップします。

    ```rust
    // match
    #[auto_enum(Iterator)]
    fn expr_match(x: i32) -> impl Iterator<Item=i32> {
        match x {
            0 => Some(0).into_iter(),
            _ => 0..x,
        }
    }

    // if
    #[auto_enum]
    fn expr_if(x: i32) -> impl Iterator<Item=i32> {
        #[auto_enum(Iterator)]
        let iter = if x == 0 {
            Some(0).into_iter()
        } else {
            0..x
        };
        iter
    }
    ```

- `loop`

    バージョン0.3でサポートが追加されました。各`break`のが返す値をラップします。また、ネストされた`loop`やラベル付き`loop`もサポートされています。

    ```rust

    #[auto_enum(Iterator)]
    fn expr_loop(mut x: i32) -> impl Iterator<Item = i32> {
        loop {
            if x < 0 {
                break x..0;
            } else if x % 5 == 0 {
                break 0..=x;
            }
            x -= 1;
        }
    }
    ```

- ブロック、unsafeブロック、メソッド呼び出し

    `if`、`match`、`loop`、またはサポートされていない式が見つかるまで再帰的に探索します。

    現在のバージョン(0.3)では最終的にif、matchのいずれかにならない場合は有効な式が帰って来ないと見做して探索を終了します。
    また、探索中に他の`#[auto_enum]`を発見した場合はそれ以上は探索しません。これにより、複数の`#[auto_enum]`を安全に使用できます。

    ```rust

    // ブロック
    #[auto_enum(Iterator)]
    fn expr_block(x: i32) -> impl Iterator<Item=i32> {
        {
            if x == 0 {
                Some(0).into_iter()
            } else {
                0..x
            }
        }
    }

    // メソッド
    #[auto_enum(Iterator)]
    fn expr_method(x: i32) -> impl Iterator<Item=i32> {
        match x {
            0 => Some(0).into_iter(),
            _ => 0..x,
        }.map(|y| y + 1)
    }
    ```

- `return`

    関数かクロージャに`#[auto_enum]`が使用されている場合にはスコープ内の`return`を解析できます。ただし、以下の条件を満たす必要があります。

    - 関数の場合は戻り値の型が`impl Trait`である必要があります。OptionやResultなどの構造体に`impl Trait`が入っている場合にはこの解析は行われません(`return`で返す値がimpl Traitか判断できないため)。

    - クロージャの場合はブロックなどで囲まれていないクロージャに直接`#[auto_enum]`が使用されている必要があります(クロージャに使用したいのか、クロージャの中の式に使用したいのかわからないため)。

    この解析は他の`#[auto_enum]`を発見した場合にも探索を継続されます。そして、他のクロージャを見つけるとそこで停止します(Rustのreturnのスコープと同じです)。

    ```rust

    // 関数でのreturn
    #[auto_enum(Iterator)]
    fn func(x: i32) -> impl Iterator<Item=i32> {
        if x == 0 {
            return Some(0).into_iter();
        }

        if x > 0 {
            0..x
        } else {
            x..=0
        }
    }

    // クロージャでのreturn
    #[auto_enum]
    fn closure() -> impl Iterator<Item=i32> {
        #[auto_enum(Iterator)]
        let f = |x| {
            if x == 0 {
                return Some(0).into_iter();
            }

            if x > 0 {
                0..x
            } else {
                x..=0
            }
        };
        f(1)
    }
    ```

#### オプション

`#[auto_enum]`には細かい制御を可能にするためのオプションがいくつかあります。

1.`#[rec]`属性でネストした分岐を解析します。

```rust

#[auto_enum(Iterator)]
fn nested(x: i32) -> impl Iterator<Item=i32> {
    match x {
        0 => Some(0).into_iter(),
        #[rec]
        x => match x {
            1 => vec![5, 10].into_iter()
            _ => 0..x
        },
    }
}
```

2.`marker!`マクロは自動的にenumのヴァリアントに置き換えられます。関数やクロージャの最後の式が分岐しない場合などに使用します(現在(v0.3)の`#[auto_enum]`はこれを解析出来ないため)。

```rust

#[auto_enum(Iterator)]
fn marker(x: i32) -> impl Iterator<Item=i32> {
    if x == 0 {
        return Some(0).into_iter();
    }
    marker!(0..x)
}
```

3.`#[never]`属性を使用すると`#[auto_enum]`の前述の解析を(値が返ってこないと認識させて)スキップさせることが可能です。ただし、`#[auto_enum]`もある程度これを解析できます。以下の構文には`#[never]`属性を付ける必要はありません。

- 必ずpanicするマクロ(`panic!(..)`, `unreachable!(..)`)
- 制御フローに関するキーワード(`return`, `break`, `continue`)
- `None?`と`Err(..)?`
- マーカーマクロ(`marker!(..)`).
- アイテム定義(fn, const, struct, etc.)

ただし、`return`に対する`#[never]`属性は関数かクロージャ内で使用する場合には必要です。

```rust

// このケースでは不要
#[auto_enum(Iterator)]
fn simple_panic(x: i32) -> impl Iterator<Item=i32> {
    match x {
        0 => Some(0).into_iter(),
        1 => panic!(),
        _ => 0..x,
    }
}

// このケースでは必要
#[auto_enum(Iterator)]
fn panic_in_loop(x: i32) -> impl Iterator<Item=i32> {
    match x {
        0 => Some(0).into_iter(),
        #[never]
        1 => loop {
            panic!()
        },
        _ => 0..x,
   }
}
```

### 現在の課題

1. ドキュメント

    ~~現在のドキュメントはすごく分かりづらいので、これを改善したいと考えています。~~特にサポートしてるトレイトとクレートfeaturesがやたらと多いのでこの辺りの説明を綺麗にしたいです。また、`#[enum_derive]`が各トレイトに対してどんなコードを生成するかの例も挙げたいと考えています。

    **更新(2018-12-24)**: バージョン0.3で大きく改善されました。まだわかりずらい部分やドキュメント化されていない部分があるので、今後のバージョンでそれも修正したいと考えています。

2. Boxを使用する方法との比較

    現在のベンチマークは[ここ](https://github.com/taiki-e/auto_enums/tree/master/test_suite/benches)にあるのですが，このベンチマークで分かる部分は、enumを使用する方法では一部でno-opが最適化されている(ように見える)ことでしょうか(アセンブラ読んでないので正確なところは分からないです)。
    ちなみにクレートfeatureの`try_trait`を有効にすると`Iterator`をより効率的に実装できるのですが、このfeatureを有効にしてベンチマークを取ると、ヴァリアントの数によってパフォーマンスが低下していた所も良好なパフォーマンスに変化します。
    この辺りで起きている現象を検証したいです。

    また、`Iterator`以外のトレイトでのベンチマークも取りたいと考えています。

3. ~~`loop`内の`break`~~

    ~~ラベルなし`loop`を次のminorバージョンでサポートする予定です。~~

    **更新(2018-12-24)**: バージョン0.3でサポートされました。

4. 関数やクロージャの最後の式が分岐しないケース

    解析方法の見当はついているので次かその次のminorバージョンでサポートしたいと考えています。

TODO: ほかの課題を挙げる

## まとめ

- 複数の`impl Trait`を返す方法にはBoxを使用してトレイトオブジェクトにする方法と、enumを使用してパターンマッチでディスパッチする方法がある。
- enumを使用する方法を書きやすくするためにいくつかの提案がなされている。
- 提案されている方法を参考に[auto_enums]というenumを使用する方法を自動化するクレートを作成した。

[auto_enums]: https://github.com/taiki-e/auto_enums
[derive_utils]: https://github.com/taiki-e/derive_utils