【Rust】異なるミュータビリティの参照を返す関数

※Rustのバージョンは1.33.0 (stable)を使用しています。

Rustでは、&Sと&mut Sは別の型扱いになるため、似たような処理をする場合でも、戻り値が&Sの場合と&mut Sの場合では別の関数を書かなければなりません。
この2つの関数をゼロコスト抽象化で共通化できないかというのが、この記事の主題です。

こんな関数を共通化してみる

/// テスト用構造体
struct MyStruct<V> {
    pivot: i32,
    left: V,
    right: V
}

impl<V> MyStruct<V> {

    /// key >= pivotならright, key < pivotならleftを返す
    /// immutable版
    fn left_or_right_ref(&self, key:i32) -> &V {
        if key >= self.pivot {
            &self.right
        } else {
            &self.left
        }
    }

    /// key >= pivotならright, key < pivotならleftを返す
    /// mutable版
    fn left_or_right_mut(&mut self, key:i32) -> &mut V {
        if key >= self.pivot {
            &mut self.right
        } else {
            &mut self.left
        }
    }

}

immutable版とmutable版では中身はほぼ同じですが、戻り値が別の型となるため、通常は2つのメソッドを用意する必要があります。この2つメソッドの内部処理を共通化してみます。

完成形

まずは、完成した形を載せます。詳細は後述。

完成形

use std::borrow::Borrow;

/// テスト用構造体
struct MyStruct<V> {
    pivot: i32,
    left: V,
    right: V
}

/// テスト用構造体のプロパティ用トレイト
trait MyStructRef<V> : Borrow<MyStruct<V>> {

    /// pivotの型
    type PivotValue;
    /// left, rightの型
    type Value;

    fn fields(self) -> (Self::PivotValue, Self::Value, Self::Value);

}

/// &MyStruct<V>のプロパティ用トレイト実装
impl<'a, V> MyStructRef<V> for &'a MyStruct<V> {
    type PivotValue = &'a i32;
    type Value = &'a V;

    fn fields(self) -> (Self::PivotValue, Self::Value, Self::Value) {
        (&self.pivot, &self.left, &self.right) 
    }
}

/// &mut MyStruct<V>のプロパティ用トレイト実装
impl<'a, V> MyStructRef<V> for &'a mut MyStruct<V> {

    type PivotValue = &'a mut i32;
    type Value = &'a mut V;

    fn fields(self) -> (Self::PivotValue, Self::Value, Self::Value) {
        (&mut self.pivot, &mut self.left, &mut self.right) 
    }

}

impl<V> MyStruct<V> {

    /// key >= pivotならright, key < pivotならleftを返す
    fn left_or_right<T>(my_refmut:T, key:i32) -> T::Value
    where T: MyStructRef<V> {
        // &mut MyStruct<V>と&MyStruct<V>を&MyStruct<V>に揃える
        // &MyStruct<V>で型を確定すると、フィールドにアクセスできる
        let my_ref = my_refmut.borrow();
        if key >= my_ref.pivot {
            let (_, _, right) = my_refmut.fields();
            // T = &mut MyStruct<V>なら&mut V
            // T = &MyStruct<V>なら&V
            right
        } else {
            let (_, left, _) = my_refmut.fields();
            // T = &mut MyStruct<V>なら&mut V
            // T = &MyStruct<V>なら&V
            left
        }
    }

    fn left_or_right_ref(&self, key:i32) -> &V {
        Self::left_or_right(self, key)
    }

    fn left_or_right_mut(&mut self, key:i32) -> &mut V {
        Self::left_or_right(self, key)
    }

}

順を追ってやってみる

方針

まずは、2つのメソッドから呼び出される、関連関数（selfを引数に取らない関数）の定義を考えます。

fn left_or_right<T>(my_refmut:T, key:i32) -> &V か &mut V
where T: &MyStruct<V>と&mut MyStruct<V>で実装のあるトレイト

こんなイメージの関数があれば、共通化できそうです。
&Sと&mut Sで実装のあるトレイトは、標準ライブラリリファレンスのreferenceに記述があります。その中ではstd::borrow::Borrowが使えそうです。が、Borrowをトレイト境界にしただけでは、MyStructのフィールドをイミュータブルでしか扱えず、ミュータブルで戻す処理が書けません。

fn left_or_right<T>(my_refmut:T, key:i32) -> &V か &mut V
where T: Borrow<MyStruct<V>> {
    // &mut MyStruct<V>と&MyStruct<V>を&MyStruct<V>に揃える
    // &MyStruct<V>で型を確定すると、フィールドにアクセスできる
    let my_ref = my_refmut.borrow();
    if key >= my_ref.pivot {
        // my_refmutはジェネリック型T
        // ジェネリック型からはフィールドが見えないので下記の処理はエラー
        ☓ my_refmut.right

        // my_refは&MyStruct<V>
        // &my_ref.rightは可能だが、&mut my_ref.rightはエラー
        ☓ &my_ref.right
    } else {
        ☓ my_refmut.left
        ☓ &my_ref.left
    }
}

というわけで、標準ライブラリだけでの実装は諦め、トレイトを自作します。

トレイトの作成

上記の方針でだめだったポイントは、「ミュータビリティを保持したままフィールドにアクセスできない」点なので、フィールドにアクセスできるように、フィールドを返すトレイトを定義します。

/// テスト用構造体のプロパティ用トレイト
trait MyStructRef<V> : Borrow<MyStruct<V>> {
    /// pivotの型
    type PivotValue;
    /// left, rightの型
    type Value;

    fn fields(self) -> (Self::PivotValue, Self::Value, Self::Value);
}

Borrow<MyStruct<V>>を継承してますが、必ず必要というわけではありません。&MyStructとして扱いたい場面は多いと思うので、付けておいたほうが便利かと思います。
各フィールドの戻り値は、&Vか&mut Vにするので、ジェネリック（関連型）でそれぞれ定義しておきます。
今回は関連型にトレイト境界は使ってませんが、もし戻り値に対して必要な操作がある場合は、

type Value: Borrow<V>;

のように制約を入れると、トレイトのメソッドが使えるようになります。
また、fieldsは所有権を要求している点に注意してください。このメソッドを使用後は、元の構造体の参照にはアクセスできなくなります。(所有権を要求しない方法は、unsafeを使用しなければないと思っていますが…どうでしょう。)
では、&MyStruct<V>と&mut MyStruct<V>に対して、このトレイトを実装します。

/// &MyStruct<V>のプロパティ用トレイト実装
impl<'a, V> MyStructRef<V> for &'a MyStruct<V> {
    type PivotValue = &'a i32;
    type Value = &'a V;

    fn fields(self) -> (Self::PivotValue, Self::Value, Self::Value) {
        (&self.pivot, &self.left, &self.right) 
    }
}

/// &mut MyStruct<V>のプロパティ用トレイト実装
impl<'a, V> MyStructRef<V> for &'a mut MyStruct<V> {

    type PivotValue = &'a mut i32;
    type Value = &'a mut V;

    fn fields(self) -> (Self::PivotValue, Self::Value, Self::Value) {
        (&mut self.pivot, &mut self.left, &mut self.right) 
    }

}

共通関数の実装

あとは、先程の方針通りに共通関数を実装します。

impl<V> MyStruct<V> {

    /// key >= pivotならright, key < pivotならleftを返す
    fn left_or_right<T>(my_refmut:T, key:i32) -> T::Value
    where T: MyStructRef<V> {
        // &mut MyStruct<V>と&MyStruct<V>を&MyStruct<V>に揃える
        // &MyStruct<V>で型を確定すると、フィールドにアクセスできる
        let my_ref = my_refmut.borrow();
        if key >= my_ref.pivot {
            let (_, _, right) = my_refmut.fields();
            // T = &mut MyStruct<V>なら&mut V
            // T = &MyStruct<V>なら&V
            right
        } else {
            let (_, left, _) = my_refmut.fields();
            // T = &mut MyStruct<V>なら&mut V
            // T = &MyStruct<V>なら&V
            left
        }
    }

    fn left_or_right_ref(&self, key:i32) -> &V {
        Self::left_or_right(self, key)
    }

    fn left_or_right_mut(&mut self, key:i32) -> &mut V {
        Self::left_or_right(self, key)
    }

}

戻り値が&Vか&mut Vとなるのは、MyStructRefトレイトの関連型Valueなので、戻り値はT::Valueと指定します。この型はコンパイル時に必ず一意に決まります。
また、前述の通り、my_refmut.fields()は所有権を奪っているため、呼び出した後はmy_refmutやmy_refは使用できません。
以上で実装が完了しました。

試しになんか動かしてみる

検証コード

fn main() {
    let mut a = MyStruct { pivot: 0, left:"left".to_string(), right:"right".to_string() };

    let rmut = a.left_or_right_mut(0);
    println!("{:?}", rmut);

    *rmut = "俺様参上!!".to_string();

    let rref = a.left_or_right_ref(0);
    println!("{:?}", rref);
}

実行結果

"right"
"俺様参上!!"

最後に

専用のフィールドを返すトレイトを実装するのは、割とめんどくさい。&Sと&mut Sのenumを作って場合分けでいいんじゃないかと思ってしまいます。
もっといい方法があれば、コメントにてご教示願います。