【Rust】演算子のオーバーロード

2020-06-05
#rust

例えば自分で定義した型（２次元ベクトルなど）を扱う場合に、演算子のオーバーロードが使えると計算式が簡単に書けて便利だ。 Rustでは演算子に関連するトレイトを実装することで実現できる。

演算子オーバーロードの基本

例として、ジェネリクスで任意の型が使える自作の２次元ベクトルベクトル構造体：

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Vector2D<T> {
    pub x: T,
    pub y: T,
}

に対してAddトレイトを実装：

use std::ops::Add;

impl<T: Add<Output = T> + Copy> Add for Vector2D<T> {  // この `+ Copy` は要素 T がコピー可能という指定
    type Output = Vector2D<T>;
    fn add(self, other: Self) -> Self::Output {
        Self::Output { x: self.x + other.x, y: self.y + other.y }
    }
}

することで + 演算子が使用できるようになる：

#[test]
fn test_add() {
    let v1 = Vector2D{x: 1, y: 2};
    let v2 = Vector2D{x: 3, y: 4};
    assert_eq!(Vector2D{x: 4, y: 6},
               v1 + v2);
}

Rust Playgroundで実行

加算以外にも、減算など他の演算子も同様に可能。

参照を受け付けるようにする

演算子をオーバーロードする構造体が Copy トレイトを実装していないと演算子の使用で移動が起こってしまうのと、 Copy を実装している場合にはセマンティックとしてはコピーが発生してしまう（最適化で取り除かれるのかも知れないが）。 それがちょっと気になるので、構造体自体じゃなくて参照に対して演算子を定義してみる：

impl<T: Add<Output = T> + Copy> Add for &Vector2D<T> {  // <= &
    type Output = Vector2D<T>;  // <= 出力型はVector2D<T>のまま
    fn add(self, other: Self) -> Self::Output {
        Self::Output { x: self.x + other.x, y: self.y + other.y }
    }
}

#[test]
fn test_add2() {
    let v1 = Vector2D{x: 1, y: 2};
    let v2 = Vector2D{x: 3, y: 4};
    assert_eq!(Vector2D{x: 4, y: 6},
               &v1 + &v2);  // <= &同士
}

Rust Playgroundで実行

Output は参照じゃなく Vector2D<T> のままにする。

+ 使用時に & をつけてやる必要があるのでちょっと手間になってしまうが、まあ許容範囲でしょう。

別の型との演算

例えばベクトルを定数倍したりする場合、スカラーとの乗算を実装する：

use std::ops::Mul;

impl<T: Mul<Output = T> + Copy> Mul<T> for &Vector2D<T> {
    type Output = Vector2D<T>;
    fn mul(self, rhs: T) -> Self::Output {
        Self::Output { x: self.x * rhs, y: self.y * rhs }
    }
}

上記で &Vector2D<T> * T が使用できる。

これを逆に、 T * &Vector2D<T> も使用できるようにするにはどうしたらいいか。 個別の型に対して定義：

impl Mul<&Vector2D<f64>> for f64 {
    type Output = Vector2D<f64>;
    fn mul(self, rhs: &Vector2D<f64>) -> Self::Output {
        Self::Output { x: self * rhs.x, y: self * rhs.y }
    }
}

すれば一応使える。 これを個別の型ごとに書かなくて済むようにジェネリクスが使えるんじゃないかと、

// 注意：コンパイルエラーになる
impl<T: Mul<Output = T> + Copy> Mul<&Vector2D<T>> for T {
    type Output = Vector2D<T>;
    fn mul(self, rhs: &Vector2D<T>) -> Self::Output {
        Self::Output { x: self * rhs.x, y: self * rhs.y }
    }
}

などとするとエラーが出る：

error[E0210]: type parameter `T` must be used as the type parameter for some local type (e.g., `MyStruct<T>`)
  --> src/vec.rs:52:6
   |
52 | impl<T> Mul<&Vector2D<T>> for T
   |      ^ type parameter `T` must be used as the type parameter for some local type
   |
   = note: implementing a foreign trait is only possible if at least one of the types for which is it implemented is local
   = note: only traits defined in the current crate can be implemented for a type parameter

error: aborting due to previous error

For more information about this error, try `rustc --explain E0210`.
error: could not compile `vec`.

To learn more, run the command again with --verbose.

どうも演算子の実装を定義できるのは自分のcrateで定義している型に対してだけらしい。 （rustc --explain E0210 で説明が表示される）。

というのも、別の型 Foobar が同じく Mul を実装している場合に、 impl<Vector2D<T>> Mul<Vector2D<T>> for Foobar とのどちらも適用可能でぶつかってしまうとのこと。

マクロで定義

ジェネリクスが使えずに型ごとに同じ定義を繰り返し記述しなきゃならないのは面倒なので、いっちょマクロを書いてみる：

macro_rules! defmulvec {
    ( $( $t:ty ),+ ) => {
        $(
            impl Mul<&Vector2D<$t>> for $t {
                type Output = Vector2D<$t>;
                fn mul(self, rhs: &Vector2D<$t>) -> Self::Output {
                    Self::Output { x: self * rhs.x, y: self * rhs.y }
                }
            }
        )+
    };
}

defmulvec![i32, f64];

代入演算子のオーバーロード

代入演算子もオーバーロードできて、例えば加算はAddAssignトレイトを実装する：

use std::ops::AddAssign;

impl<T: AddAssign + Copy> AddAssign for Vector2D<T> {
    fn add_assign(&mut self, other: Self) {
        self.x += other.x;
        self.y += other.y;
    }
}

これの右辺値を実体じゃなく参照にしたかったが、それはできない模様… （ for の後ろを参照型にできないし、右辺値の型を Self と同じ型じゃないとだめらしい）。 AddAssign が型引数がデフォルトで <Rhs = Self> となっているのを、参照を指定すればできた：

use std::ops::AddAssign;

impl<T: AddAssign + Copy> AddAssign<&Vector2D<T>> for Vector2D<T> {
    fn add_assign(&mut self, other: &Vector2D<T>) {
        self.x += other.x;
        self.y += other.y;
    }
}

#[test]
fn test_add_assign() {
    let mut v1 = Vector2D{x: 1, y: 2};
    let v2 = Vector2D{x: 3, y: 4};
    v1 += &v2;
    assert_eq!(Vector2D{x: 4, y: 6}, v1);
}

Rust Playgroundで実行

参考

関連記事

新着記事