[Rust] 構造体要素を持つenumにはIterator系メソッドを実装すると便利

前置き: 構造体要素を持つenumとは

enumはC系の言語でおなじみの列挙型です。Rustの場合は以下のように書きます。

enum MyEnum {
    ItemA,
    ItemB,
}

Rustの場合、以下のように構造体をもたせることも出来ます。

enum MyEnum {
    ItemA(i64),
    ItemB { x: i64, y: i64 },
}

これは、以下のように使います。

let mut e: MyEnum = MyEnum::Item(123);
// ...
if let MyEnum::ItemA(i) = e {
    println!("{}", i);
} else if let MyEnum::ItemB{x: x, y: y} = e {
    println!("x = {}, y = {}", x, y);
}

本題

例えばRustを使って構文木を扱いたい場合、以下のようなenumを宣言するでしょう。

enum Expression {
    SomeValue,
    PrefixOperator(Box<Expression>),
    BinaryOperator(Box<Expression>, Box<Expression>),
}

この時、不便なのがExpressionの入れ子になったExpressionにアクセスする時です。通常、これらにアクセスするためには

if let Expression::PrefixOperator(ee) = e {
    // ...
}

のように、if letやmatch等を使わなければなりません。eがExpression::PrefixOperatorである時に特有の処理を書きたいのであればそれでも良いのですが、場合によっては、Expressionの種類に関わらず、すべての子要素に対して特定の処理を実行したいという場合もあるでしょう。

そのような場合、ExpressionにIterator的な関数を実装しておくとすごく便利です。

map関数

例えば、RustのOption型はmap関数を持っています。これを使うと、Option<T>をある関数f: T -> Uを用いてOption<U>に変換する処理は

// input: Option<T>
let output = if let Some(o) = input {
    Some(func(o)) // func: T -> U
} else {
    None
}

// 上と同じことは下のようにかける
let output = input.map(func);

と書くことができます。これにならい、Expressionにもmap関数を実装しましょう。

impl Expression {
    pub fn map<F: Copy + FnMut(Expression) -> Expression>(self, mut f: F) -> Expression {
        match self {
            Expression::SomeValue => Expression::SomeValue,
            Expression::BinaryOperator(e1, e2) => {
                Expression::BinaryOperator(Box::new(f(*e1)), Box::new(f(*e2)))
            }
            Expression::PrefixOperator(e) => Expression::PrefixOperator(Box::new(f(*e))),
        }
    }
}

注意点としては、e1等はBox<Expression>なので、*演算子を用いて中身を取り出しています。

map関数を用いれば、例えばExpression::SomeValueに対して再帰的に最適化処理する関数は以下のように掛けるでしょう。

fn func(e: Expression) -> Expression {
    if let Expression::SomeValue = e {
        // ここにSomeValueの最適化処理をかく
        Expression::SomeValue
    } else {
        e.map(|e| func(e))
    }
}

Iteratorを返す関数

上のmap関数はExpressionをExpressionに変換する関数でした。一方、子要素に対して繰り返し処理を行いたい場合はIteratorを返す関数を実装すると便利です。

以下ではiter(), iter_mut(), into_iter()を実装しています。(手抜きのためstd::vec::IntoIterを援用しています。)

impl Expression {
    pub fn iter<'a>(&'a self) -> std::vec::IntoIter<&'a Expression> {
        let v = match self {
            Expression::SomeValue => Vec::new(),
            Expression::BinaryOperator(e1, e2) => vec![e1.as_ref(), e2.as_ref()],
            Expression::PrefixOperator(e) => vec![e.as_ref()],
        };
        v.into_iter()
    }

    pub fn iter_mut<'a>(&'a mut self) -> std::vec::IntoIter<&'a mut Expression> {
        let v = match self {
            Expression::SomeValue => Vec::new(),
            Expression::BinaryOperator(e1, e2) => vec![e1.as_mut(), e2.as_mut()],
            Expression::PrefixOperator(e) => vec![e.as_mut()],
        };
        v.into_iter()
    }
}

impl IntoIterator for Expression {
    type Item = Expression;
    type IntoIter = std::vec::IntoIter<Expression>;

    fn into_iter(self) -> Self::IntoIter {
        let v = match self {
            Expression::SomeValue => Vec::new(),
            Expression::BinaryOperator(e1, e2) => vec![*e1, *e2],
            Expression::PrefixOperator(e) => vec![*e],
        };
        v.into_iter()
    }
}

これを使うと、例えばDebugトレイトを実装したくなった時に

use std::fmt;
impl fmt::Debug for Expression {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Expression::SomeValue => {
                f.write_str("SomeValue")?;
                return Ok(());
            }
            Expression::BinaryOperator(_, _) => f.write_str("BinaryOperator(")?,
            Expression::PrefixOperator(_) => f.write_str("PrefixOperator(")?,
        }
        let mut is_first = true;
        for e in self.iter() {
            if !is_first {
                f.write_str(", ")?;
            }
            is_first = false;
            f.write_fmt(format_args!("{:?}", e))?;
        }
        f.write_str(")")?;
        Ok(())
    }
}

と簡単に実装できるようになります。