【Rust】collect関数って何してんの？

はじめに

Rustを書いていると頻繁に呼び出す関数があります。
イテレータの最後に呼び出すcollect関数です。

let filtered: Vec<i32> = (0..10).filter(|x| x % 2 == 0).collect();
println!("{:?}", filtered); // [0, 2, 4, 6, 8]

今までは深く考えずに呼び出していましたが
ふと内部の動作が気になったので、今回調べることにしました。

collect関数の内部

collect関数はコアのソースコードのiterator.rsに定義されています。
実装は以下の通りです。

iterator.rs

    fn collect<B: FromIterator<Self::Item>>(self) -> B
    where
        Self: Sized,
    {
        FromIterator::from_iter(self)
    }

FromIteratorというトレイトの関数を呼び出しているだけのシンプルな処理でした。
次はこのFromIteratorトレイトについて説明していきます。

FromIterator

FromIteratorトレイトはイテレータから別の型に変換するためのトレイトです。
VecやHashMap等はこのトレイトを実装しているのでcollect関数で変換できます。

FromIteratorトレイトは、collect.rsに以下の様に定義されています。

collect.rs

pub trait FromIterator<A>: Sized {
    /// Creates a value from an iterator.
    ///
    /// See the [module-level documentation] for more.
    ///
    /// [module-level documentation]: crate::iter
    ///
    /// # Examples
    ///
    /// ```
    /// let five_fives = std::iter::repeat(5).take(5);
    ///
    /// let v = Vec::from_iter(five_fives);
    ///
    /// assert_eq!(v, vec![5, 5, 5, 5, 5]);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    #[rustc_diagnostic_item = "from_iter_fn"]
    fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> Self;
}

上記にあるIntoIteratorトレイトとは、ある型からイテレータへの変換を行うトレイトです。
FromIteratorとは逆の変換を行うトレイトです。

実装例

実際にFromIteratorトレイトを実装してみましょう。
以下の様に整数コレクションの計算結果を格納するCalcResult構造体を定義します。
sumには整数の総和を、productには総乗を格納します。

struct CalcResult {
    pub sum: i32,
    pub product: i32,
}

次に整数のイテレータからCalcResultに変換するコードを書きます。
当然ですが、この時点ではコンパイルエラーになります。

// FromIteratorトレイトが実装されていないので変換できない
let result: CalcResult = (0..10).filter(|x| x % 2 == 1).collect();
println!("{} {}", result.sum, result.product);

CalcResult構造体にFromIteratorトレイトを実装します。
fold関数を使って総和と総乗を一度に計算します。

impl FromIterator<i32> for CalcResult {
    fn from_iter<T: IntoIterator<Item = i32>>(iter: T) -> Self {
        let result = iter
            .into_iter()
            .fold((0, 1), |result, x| (result.0 + x, result.1 * x));
        let (sum, product) = result;
        CalcResult { sum, product }
    }
}

するとcollect関数内で上記のfrom_iter関数が呼び出されます。
これでコンパイルが通る様になり実行結果が表示されます。

25 945

総和は1 + 3 + 5 + 7 + 9 = 25で、総乗は1 * 3 * 5 * 7 * 9 = 945です。
このようにFromIteratorトレイトを実装すると、イテレータを独自の構造体に変換することが可能です。

コンパイルエラーの原因

collect関数内の仕組みが分かると、なぜ型を明示しないとコンパイルエラーになるかが理解できます。

// このコードはコンパイルエラーになる
let filtered = (0..10).filter(|x| x % 2 == 0).collect();

collect関数内でどの型のFromIterator::from_iter関数を呼び出すべきかを
コンパイラが推論できないのでエラーになっていたということです。

以下の方法で型を明示する必要があります。

// 変数に型注釈をつける
let filtered: Vec<i32> = (0..10).filter(|x| x % 2 == 0).collect();
// collect関数に型引数を渡す
let filtered = (0..10).filter(|x| x % 2 == 0).collect::<Vec<i32>>();

参考文献