Rustの二次元配列の要素のswap

ボローチェッカーは、GCに頼ることなくメモリ安全性を保障してくれたりデータレースの不在を保障してくれたりするRustのユニークで有用な機能ですが、現在のボローチェッカーはずいぶん保守的な型検査を行うので、普通に考えるとどう考えても安全なコードが型検査を通せなかったりすることがちょくちょくあります。そんなものの一つとして、今日は二次元配列の要素のswapを考えてみようと思います。

二次元配列の要素のswap

二次元配列 v に対して、これの要素 v[i1][j1] と v[i2][j2] を入れ替える操作を考えます。要するに

fn mat_swap<T>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    // ....
}

のようなインターフェースを持つ関数を実装するとします。
他の言語だと、一体これの何が難しいのだ？という話になると思います。たとえばC++だと、

template <typename T>
void mat_swap(vector<vector<T>> &v, size_t i1, size_t j1, size_t i2, size_t j2) {
    T t = v[i2][j2];
    v[i2][j2] = v[i1][j1];
    v[i1][j1] = t;
}

と書くことができるでしょう。しかしこれをそのままRustで書くと、

fn mat_swap<T>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    let t = v[i2][j2];
    v[i2][j2] = v[i1][j1];
    v[i1][j1] = t;
}

次のようなエラーになります。

$ cargo build
...
error[E0507]: cannot move out of indexed content
  --> src/main.rs:11:13
   |
11 |     let t = v[i1][j1];
   |             ^^^^^^^^^
   |             |
   |             cannot move out of indexed content
   |             help: consider using a reference instead: `&v[i1][j1]`

Rustの代入はC++とは違い、デフォルトでmoveです。つまり、配列から要素を外の変数に代入するのは、配列から要素を削除することになるので、これはできません（削除した後の要素にアクセスする可能性を考えると、これは当然）。どうしても取り出したければ、要素をcloneすれば良いので、型パラメータTにCopyトレイトを要求すれば、

fn mat_swap_copy<T: Copy>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    let t = v[i2][j2];
    v[i2][j2] = v[i1][j1];
    v[i1][j1] = t;
}

これはコンパイルが通ります。

Rustでは、Cloneのコストが高い型に対してCopyを実装するのは好ましくなく、また、Copyトレイトを実装しているのであればCloneトレイトは必ず実装されているので、少しコードは変わりますが、Cloneのみを要求した方が良いでしょう。

fn mat_swap_clone<T: Clone>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    let t = v[i2][j2].clone();
    v[i2][j2] = v[i1][j1].clone();
    v[i1][j1] = t;
}

しかしこれには、

• 与えられる型の条件が厳しくなる
• cloneのコストが高い型の場合に性能が悪くなる

のような複数の問題があり、特に後者はあまり許容されるものではありません。
逆に考えると、C++でもこのような要件は暗黙的には存在しており、Rustではそういう条件が明示的に顕在化したのだとも言えるかもしれません。ウッカリこういうコードをC++で書いて、Tにコピーが重い型を入れてしまい、知らず知らずのうちに実行コストがかさんでいたことがあるかも…、などと考えると、うすら寒くもなってもきます。

C++であれば、swap()を用いて、

template <typename T>
void mat_swap(vector<vector<T>> &v, size_t i1, size_t j1, size_t i2, size_t j2) {
    swap(v[i1][j1], v[i2][j2]);
}

のように書くのが普通かもしれません。swap()は型Tに対して最適なオーバーロードを見つけて、効率の良い方法で変数を入れ替えてくれるはずだからです。

これをRustで書くとこうなります。

fn mat_swap<T>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    std::mem::swap(&mut v[i1][j1], &mut v[i2][j2]);
}

しかしこれもボローチェッカーを通りません。まあ、通っていたらこのような記事を書くこともなかったのですが。

error[E0499]: cannot borrow `*v` as mutable more than once at a time
 --> src/main.rs:7:41
  |
7 |     std::mem::swap(&mut v[i1][j1], &mut v[i2][j2]);
  |                         -               ^        - first borrow ends here
  |                         |               |
  |                         |               second mutable borrow occurs here
  |                         first mutable borrow occurs here

要するに、同じ変数へのmutableなborrowが同時に2個存在しているぞと言って怒っているわけです。Rustでは同じ変数に対するmutableなborrowは複数（あるいはimmutableなborrowとmutableなborrowも）同時に存在してはいけないことになっています。そうすることによって、自動的にデータレース（複数のスレッドから同時に同じ変数に書き込まれることによる、データの不整合）の不在が保証されることになり、これは大変重要なRustの機能です。

v[i1][j1]とv[i2][j2]は違う変数やんけ！とぱっと見は思うかもしれませんが、配列の要素のborrowは、配列そのもののborrowを要求しており、残念ながら、これはvへのborrowを2回発生させることになってしまいます。もっと厳密な型を付けろと言いたくもなるかもしれませんが、そもそもの話i1とj1,i2とj2がそれぞれ相異なるとも限らないわけで、型の設計が保守的すぎるとも言いがたいです。

その1. Unsafeを使う

しかしながら、そんなことは分かっている、この場合に限っては、データレースも何も発生しないのだし安全だと分かりきっている。なんと融通の利かない型システムなんだ。と言いたくもなります。そんなときに便利なのがunsafeです。

fn mat_swap_unsafe<T>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    let p1: *mut T = &mut v[i1][j1];
    let p2: *mut T = &mut v[i2][j2];
    unsafe {
        p1.swap(p2);
    }
}

要素に対してmutableな参照ではなく、ポインタを取得します。ポインタはボロー情報を持っていませんので、何でもできます。その代わりほとんどの操作がunsafeになります。

これで目的の動作は得られますし、これで良いんじゃないのかという気もしてきます。しかし、unsafeという不気味なワードが引っかかります。このケースに限っては安全だと分かっているけど、それは自分がそう判断しているだけであって、賢いコンパイラが機械的に確かめてくれたことではない。もしこんなに短いプログラムであっても、自分の判断が万が一間違っていたとしたら…。本当にこれで良かったのか。うーん…うーん……と、なんとなく後ろ髪引かれる思いを断ち切れずに、このときunsafeを使ったことが、ふとした拍子にずっと頭をもたげてくるかもしれません。

やっぱりunsafeは良くない。使わなくて済むのであれば、その方が良いに越したことはありません。

その2. Defaultを要求する

配列の中身をmoveできないのであれば、空いた要素にゴミを詰めておけば良いというのがこの実装の発想です。

fn mat_swap_default<T: Default>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    let mut t = T::default();
    std::mem::swap(&mut v[i1][j1], &mut t);
    std::mem::swap(&mut v[i2][j2], &mut t);
    std::mem::swap(&mut v[i1][j1], &mut t);
}

詰めるためのゴミもどこからともなく作り出さなければならないので、TにDefaultを要求しておきます。そうすれば、何らかの値を作り出すことができます。こうすれば、同時に複数のmutableなボローを作る必要がなくなり、また実行時に膨大な計算時間が消費される心配もなくなりました。しかしやはり、本来必要無いはずのトレイトを要求してしまったのは、少し気になるところです。

その3. split_at_mutを使う

Rustは同じ変数のmutableなborrowを二つ同時に作れません。そのことは、同じ配列の要素に対しても、同様の制約が課されることを意味します。しかしこういうニーズは少ないものではないはずである、というのを見越してか、RustのVecには配列を二つに分割してそれぞれに対してmutableなborrowを作る関数が用意されています。split_at_mutというメソッドがそれです。v.split_at_mut(i)は、Vecを[0, i)と[i, v.len())のスライスに分割して、それぞれのmutableなborrowを返します。これを使えば余計な制約を持ち込まずに目的の操作が実装できそうです。

fn mat_swap_split<T>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    if i1 == i2 {
        v[i1].swap(j1, j2);
    } else if i1 > i2 {
        mat_swap_split(v, i2, j2, i1, j1);
    } else {
        let (v1, v2) = v.split_at_mut(i2);
        std::mem::swap(&mut v1[i1][j1], &mut v2[0][j2]);
    }
}

実装に関してですが、まず、i1 == i2のケースを分けてやる必要があります。このケースでは、同じrowの要素を入れ替えることになるので、split_at_mutでは扱えません。同一配列内であれば、これまたおあつらえ向きにswapという要素を入れ替えるメソッドがあるので、これが使えます。

次にi1とi2の大小関係を調べます。配列を分割する都合上大小関係が分かっている必要があります。今回はi1 < i2を満たすようにしておきます。

最後に、複数得られたmutableな参照に対して、それぞれ要素をborrowしてstd::mem::swapに渡せば目的の処理が完了です。

ついに、インターフェースを変えず、定数時間で終わることが保障された実装ができました！しかし、こんな単純なことをするだけのコードにしては長い！本質的に同じ変数のmutable参照を取ることはできないので、同じにならないかのチェックをこちらでやってやることが必要になると言うわけです。確かに理論的にはおかしくはない。しかしRust的にはこれが正解なのか？こんな簡単なことをする為にこんな場合分けが必要になるのか？本当にそれでいいのか？疑問は深まります。

その4. 配列をいじくり回す

直接的な方法は先ほどの通りですが、配列をいじくり回せばsplit_at_mutを使わずとも実装できます。配列から要素を取り出すことはできないと言いましたが、実は特別なケースでは取り出すことができます。その一つが、末尾からの要素のpopです。理屈としては、取り出して長さを短くしてしまうのであれば、配列からmoveして空白ができてしまうことがないという話です。

fn mat_swap_pop_back<T>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    let n = v[i1].len();
    v[i1].swap(j1, n - 1);
    if i1 == i2 && j2 == n - 1 {
        return;
    }
    let mut e1 = v[i1].pop().unwrap();
    std::mem::swap(&mut v[i2][j2], &mut e1);
    v[i1].push(e1);
    v[i1].swap(j1, n - 1);
}

実装は、

• まず、入れ替えたい要素v[i1][j1]を配列v[i1]の末尾に持っていく
• 次に、その要素をpopして取り出す。
• 取り出した要素をv[i2][j2]とswapする
• swapした要素を、v[i1]の末尾に入れ直す
• 入れた要素を、元あった場所と入れ替えて元に戻す

という手順になります。i1 == i2かつ、j2 == v[i1].len() - 1の場合に、存在しない要素へのアクセスが発生してしまうことに気を付ける必要があります（同じ配列へのアクセスが発生する可能性があるということが罠になりうるということが、ここでも顕在化しているということなんだろうか？）。この場合は、既に入れ替えが完了していることになるので、ここでreturnしておきます。

配列としての特性を使っている分だけ、こちらの方が若干シンプルな気もしますが、どっちが良いかというと、どうなんでしょうか。少なくとも、こちらはVec<Vec<T>>を前提とすることになるので、例えば、スライスに対して同様の方法を用いることはできません。

その4.1. swap_removeを使う

実はRustのVecには、こういうユースケースを見越してか、配列のある位置の要素を、末尾と入れ替えた後にそれをpopする、swap_removeというメソッドがあります。なんとおあつらえ向きなんでしょうか。これを使えば、先ほどの手順のswapとpopが一つになって、僅かに短くなります。

fn mat_swap_swap_remove<T>(v: &mut Vec<Vec<T>>, i1: usize, j1: usize, i2: usize, j2: usize) {
    if i1 == i2 {
        v[i1].swap(j1, j2);
        return;
    }
    let n = v[i1].len();
    let mut e1 = v[i1].swap_remove(j1);
    std::mem::swap(&mut v[i2][j2], &mut e1);
    v[i1].push(e1);
    v[i1].swap(j1, n - 1);
}

なお、逆の、pushしてswapという操作はないようなので、全体としてはそこまで短くもなりません。

まとめ

というわけで、Rustの日常に潜む落とし穴的な、どうしてこれがこんなに面倒になるんだという問題を一つ取り上げて考えてみましたが、どれも決定的で合理的な解法とも言えない気がします。もし、これよりもシンプルで素直な実装ができたと言う方がいらっしゃいましたら、是非ともお教えください。