Rustで複雑な複合型をイテレートする

はじめに

Rustで(A, B, Option<C>, Vec<D>)のような複雑な複合型をイテレートする方法を考えます。すなわち

let x = (A {}, B {}, Some(C {}), vec![D {}, D {}]);
for y in x {
    println!("{:?}", y); // A B C D D と表示したい
}         

こういうことをしたい、ということですね。
簡単にできるわけではないですが、いろいろと下準備をすると実現可能です。

イテレータの返す型

当然ですがイテレータは一つの型しか返せません。今返したい型はA,B,C,Dの4つあるので、それらをまとめた型が必要です。

#[derive(Debug)]
struct A;
#[derive(Debug)]
struct B;
#[derive(Debug)]
struct C;
#[derive(Debug)]
struct D;

#[derive(Debug)]
enum Any<'a> {
    A(&'a A),
    B(&'a B),
    C(&'a C),
    D(&'a D),
}

ここでは4つをまとめる型としてAnyを作りました。
#[derive(Debug)]は最後にprintln!するためにつけただけで、本題とは無関係です。

Anyを返すイテレータ

次にイテレータを表す型Iterを作ります。この型はメンバとしてイテレートすべきAnyのリストを持っています。単にlist: Vec<Any<'a>>とせず、Anysを定義した理由は後程説明します。

struct Anys<'a>(Vec<Any<'a>>);

struct Iter<'a> {
    list: Anys<'a>,
}

このIterに対してIteratorトレイトを実装します。
実装としては単にリストの中身をpop()で取り出していくだけです。

impl<'a> Iterator for Iter<'a> {
    type Item = Any<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        self.list.0.pop()
    }
}

イテレータを作る

先ほど定義したイテレータを、冒頭の複合型から作る方法を考えます。まず、単に型Aの場合は簡単です。

impl<'a> From<&'a A> for Anys<'a> {
    fn from(x: &'a A) -> Self {
        Anys(vec![Any::A(x)])
    }
}

このようなFromトレイトを実装すればAからAnysを作れます。さらに

impl<'a> IntoIterator for Anys<'a> {
    type Item = Any<'a>;
    type IntoIter = Iter<'a>;

    fn into_iter(mut self) -> Self::IntoIter {
        self.0.reverse();
        Iter { list: self }
    }
}

Anysに対してIntoIteratorトレイトを実装すれば、Anysを介してAをイテレートすることができます。つまり

let x = A {};
let x: Anys = (&x).into();
for y in x {
    println!("{:?}", y);
}

こういうことができるようになりました。
Aに対してもIntoIteratorトレイトを実装すれば途中のlet x: Anys = (&x).into();も不要になりますが、ここでは省略しています。

複合型への対応

あとはFromトレイトを増やしていくだけです。B,C,DからAnysへのFromトレイトはAと同じなので省略します。
VecとOption、タプルを同様に定義します。

impl<'a, T: 'a> From<&'a Vec<T>> for Anys<'a>
where
    &'a T: Into<Anys<'a>>,
{
    fn from(x: &'a Vec<T>) -> Self {
        let mut ret = Vec::new();
        for x in x {
            ret.append(&mut x.into().0);
        }
        Anys(ret)
    }
}

これはVecの場合です。型制約&'a T: Into<Anys<'a>>を付けているので、Vecの中身がAnysへ変換可能な任意のVecをこれで変換できます。
Optionとタプルも同様です。

impl<'a, T: 'a> From<&'a Option<T>> for Anys<'a>
where
    &'a T: Into<Anys<'a>>,
{
    fn from(x: &'a Option<T>) -> Self {
        let mut ret = Vec::new();
        if let Some(x) = x {
            ret.append(&mut x.into().0);
        }
        Anys(ret)
    }
}

impl<'a, T0: 'a, T1: 'a, T2: 'a, T3: 'a> From<&'a (T0, T1, T2, T3)> for Anys<'a>
where
    &'a T0: Into<Anys<'a>>,
    &'a T1: Into<Anys<'a>>,
    &'a T2: Into<Anys<'a>>,
    &'a T3: Into<Anys<'a>>,
{
    fn from(x: &'a (T0, T1, T2, T3)) -> Self {
        let mut ret = Vec::new();
        let (t0, t1, t2, t3) = x;
        ret.append(&mut t0.into().0);
        ret.append(&mut t1.into().0);
        ret.append(&mut t2.into().0);
        ret.append(&mut t3.into().0);
        Anys(ret)
    }
}

これで、VecとOptionとA,B,C,Dをタプルで組み合わせた任意の複合型をAnysに変換できるようになりました。
（ここでは4つのタプルしか定義していませんが、それ以外の個数も同様なので必要なだけ増やせばよいです）

また、Vec<Any<'a>>の代わりにAnysを定義したのはこのFromトレイトを実装するためです。Rustではあるトレイトを実装する際に「そのトレイトを定義したクレートか、対象の型を定義したクレートでしか実装できない」という制限があります。FromトレイトもVecもこのクレートで定義したものではないので、VecのままだとFromトレイトは実装できません。しかし、Anysはこのクレートで定義した型なのでFromトレイトを実装できます。

まとめ

最終的にこういうことができるようになりました。

fn main() {
    let x = (A {}, B {}, Some(C {}), vec![D {}, D {}]);
    let x: Anys = (&x).into();
    for y in x {
        println!("{:?}", y);
    }

    let x = (Some(D {}), C {}, vec![B {}, B {}, B {}], Some(A {}));
    let x: Anys = (&x).into();
    for y in x {
        println!("{:?}", y);
    }
}