blockの巣

【Rust】Vec3<T>のTが何であってもlength()が浮動小数点型を返すようにする

2024/09/05 00:05 公開
Rust

やりたいこと

以下のような構造体Vec3<T>があるとします。Tは整数型も浮動小数点型も取りうるとします。

struct Vec3<T> {
    x: T,
    y: T,
    z: T,
}

これに対して長さを取得する関数length()を実装したいです。
しかし長さを求めるにはには平方根を計算する必要があり、整数型の場合は平方根を取ると多くの場合浮動小数点型になってしまいます。
そのためfn length(&self) -> Tとすることはできません。
そこでT = f32T = f64ならばfn length(&self) -> Tとなり、T = i32T = i64ならばfn length(&self) -> f64となるような実装を行います。

コード

trait VectorElement {
    type FloatCalcType: VectorElement + num::Float;

    fn as_float_type(&self) -> Self::FloatCalcType;
}

impl VectorElement for i32 {
    type FloatCalcType = f64;
    fn as_float_type(&self) -> Self::FloatCalcType {
        *self as Self::FloatCalcType
    }
}

impl VectorElement for f64 {
    type FloatCalcType = f64;
    fn as_float_type(&self) -> Self::FloatCalcType {
        *self
    }
}

struct Vec3<T: VectorElement> {
    x: T,
    y: T,
    z: T,
}

impl<T: VectorElement> Vec3<T> {
    fn length(&self) -> T::FloatCalcType {
        use num::Float as _;
        (self.x.as_float_type().powi(2)
            + self.y.as_float_type().powi(2)
            + self.z.as_float_type().powi(2))
        .sqrt()
    }
}

fn main() {
    use std::any::{Any, TypeId};
    let f64_vec = Vec3::<f64> {
        x: 4.0,
        y: 2.0,
        z: 4.0,
    };
    let f64_vec_len = f64_vec.length();
    assert_eq!(f64_vec_len, 6.0);
    assert_eq!(f64_vec_len.type_id(), TypeId::of::<f64>());

    let i32_vec = Vec3::<i32> { x: 4, y: 2, z: 4 };
    let i32_vec_len = i32_vec.length();
    assert_eq!(i32_vec_len, 6.0);
    assert_eq!(i32_vec_len.type_id(), TypeId::of::<f64>());
}

Rust Playgroundで実行

解説

まず浮動小数点が絡む計算をどの型で行うかを決めるために、VectorElementトレイトを定義します。
これは浮動小数点型を満たしていて欲しいのと、Vec3の要素としても使用したいためVectorElement + num::Floatとしています。

trait VectorElement {
    type FloatCalcType: VectorElement + num::Float;

    fn as_float_type(&self) -> Self::FloatCalcType;
}

as_float_typeTT::FloatCalcTypeに変換する関数ですがnum::NumCastなどで代替可能なので必須ではないです。

次にVec3の要素として使用する型に対してVectorElementトレイトを実装します。ここではi32f64に対して実装しています。
i32の場合はFloatCalcTypef64にしています。

impl VectorElement for i32 {
    type FloatCalcType = f64;
    fn as_float_type(&self) -> Self::FloatCalcType {
        *self as Self::FloatCalcType
    }
}

impl VectorElement for f64 {
    type FloatCalcType = f64;
    fn as_float_type(&self) -> Self::FloatCalcType {
        *self
    }
}

最後にVec3に対してlength()を実装します。
length()T::FloatCalcTypeを返すようにします。

impl<T: VectorElement> Vec3<T> {
    fn length(&self) -> T::FloatCalcType {
        use num::Float as _;
        (self.x.as_float_type().powi(2)
            + self.y.as_float_type().powi(2)
            + self.z.as_float_type().powi(2))
        .sqrt()
    }
}

これでVec3の要素が整数型でも浮動小数点型でもlength()が正しく計算できるようになりました。

おわりに

正直これくらいであれば

impl Vec<f64> {
    // length実装
}
impl Vec<i32> {
    // length実装
}

としてやっても良いか思いますが、今回の方法であれば

impl VectorElement for i16 { /* ... */ }
impl VectorElement for f32 { /* ... */ }

のように要素型に対してトレイトを実装するだけVec3のコードを変更することなく対応要素型を増やすことができます。

おまけ

normalize(正規化)実装

戻り値を要素が浮動小数点型のVec3にすることもできます。

impl<T: VectorElement> Vec3<T> {
    fn normalized(&self) -> Vec3<<T as VectorElement>::FloatCalcType> {
        let length = self.length();
        Vec3 {
            x: self.x.as_float_type() / length,
            y: self.y.as_float_type() / length,
            z: self.z.as_float_type() / length,
        }
    }
}

feature

外部から利用するときにfeatureによってFloatCalcTypeを変更するということもできそうな気がしますね。