Rustでマージソートを実装するのが大変だった

背景

Rust練習のために、とりあえずWikipediaに載っているソートアルゴリズム全部実装する！→マージソートは所有権的に実装難しそう？とりあえずやってみよう

結論

Vecのソートアルゴリズムはマージソート(と要素数の少ないときはインサートソート)なので、使えるなら使おう

vec![2, 0, 1].sort();

Rustドキュメント

The current algorithm is an adaptive, iterative merge sort inspired by timsort.

訳：現在のアルゴリズムはTimsortに触発された適応的で反復型のマージソートです。

Rustソースコード

とは言うものの...

やっぱり実装したい！
ので、所有権システムのために3つの選択肢を考えた。

1. Copyトレイトを強制
2. Defaultトレイトを強制
3. unsafeを使う

まず、マージソートは、

1. 左半分をソート
2. 右半分をソート
3. 左右のソート結果を合成

という手順のアルゴリズム。
Pythonで大まかな実装を示す。

import copy

def merge_sort(array):
    if len(array) <= 1:
        return array

    middle = len(array) // 2
    left = merge_sort(array[:middle])
    right = merge_sort(array[middle:])

    # 右のソート済み配列を逆順にして結合
    merge = copy.copy(left + right[::-1])
    l = 0
    r = len(array) - 1
    for i in range(len(array)):
        if merge[l] > merge[r]:
            array[i] = merge[r]
            r -= 1
        else:
            array[i] = merge[l]
            l += 1
    return array

この時、鬼門になるのは
copy.copy(left + right[::-1])
の部分だ。

スライスに「スライスに格納されている値」の所有権があるため、気軽に値を移動させることができない。

1.Copyトレイトを使う

回答はシンプルで、Copyできるように(Copyトレイトの実装を要求)すれば万事解決がその1。

fn merge_sort<T: Ord + Copy>(data: &mut [T]) {
    if data.len() <= 1 {
        return;
    }

    let length = data.len();
    let middle = length / 2;
    merge_sort(&mut data[..middle]);
    merge_sort(&mut data[middle..]);

    let mut work_vec: Vec<T> = Vec::with_capacity(length);
    for i in 0..length {
        work_vec.push(data[i]);
    }

    work_vec[middle..].reverse();

    let mut l = 0;
    let mut r = length - 1;

    for i in 0..length {
        if work_vec[l] > work_vec[r] {
            data[i] = work_vec[r];
            r -= 1;
        } else {
            data[i] = work_vec[l];
            l += 1;
        }
    }
}

2. Defaultトレイトを使う

std::mem::takeを使って所有権を奪う。これが一番好き。Defaultトレイトは実装していることが多いし。Defaultの実装次第では遅いかも。

fn merge_sort<T: Ord + Default>(data: &mut [T]) {
    if data.len() <= 1 {
        return;
    }

    let length = data.len();
    let middle = length / 2;
    merge_sort(&mut data[..middle]);
    merge_sort(&mut data[middle..]);

    let mut work_vec: Vec<T> = Vec::with_capacity(length);
    for i in 0..length {
        work_vec.push(std::mem::take(&mut data[i]));
    }

    work_vec[middle..].reverse();

    let mut l = 0;
    let mut r = length - 1;

    for i in 0..length {
        if work_vec[l] > work_vec[r] {
            data[i] = std::mem::take(&mut work_vec[r]);
            r -= 1;
        } else {
            data[i] = std::mem::take(&mut work_vec[l]);
            l += 1;
        }
    }
}

3. unsafeを使う

正直自信ない...
Rust公式はunsafeを使った実装だった。
copy_to_nonoverlappingを使い、スライスを全てコピーする。
もっとドキュメント読み込んで自信が持てたらこれがベストだと思う。バグあったらすいません

2021/6/25 修正しました

2021/6/25追記

@vain0x さんからの指摘で、コードが未定義動作を引き起こす問題がありました。
問題を修正するまでコードを削除しておきます。

2021/6/25追記

あまり美しくないが、暫定修正版を出しておく。
値をcopy_to_nonoverlapping()でコピーし、その後work_vecの長さを0にすることでwork_vecに格納されている値からまだ使われているメモリが解放されることを防ぐ。

fn merge_sort<T: Ord>(data: &mut [T]) {
    if data.len() <= 1 {
        return;
    }

    let length = data.len();
    let middle = length / 2;
    merge_sort(&mut data[..middle]);
    merge_sort(&mut data[middle..]);

    let mut work_vec: Vec<T> = Vec::with_capacity(length);
    unsafe {
        data.as_ptr()
            .copy_to_nonoverlapping(work_vec.as_mut_ptr(), length);
        work_vec.set_len(length);
    }

    work_vec[middle..].reverse();

    let mut l: isize = 0;
    let mut r: isize = length as isize - 1;
    for i in 0..length {
        if work_vec[l as usize] > work_vec[r as usize] {
            unsafe {
                work_vec
                    .as_ptr()
                    .offset(r)
                    .copy_to_nonoverlapping(data.as_mut_ptr().offset(i as isize), 1);
            }
            r -= 1;
        } else {
            unsafe {
                work_vec
                    .as_ptr()
                    .offset(l)
                    .copy_to_nonoverlapping(data.as_mut_ptr().offset(i as isize), 1);
            }
            l += 1;
        }
    }

    unsafe {
        work_vec.set_len(0);
    }
}

※絶対に使わないこと　今までのコードとダメな理由

このコードがダメな理由は未定義動作を引き起こすのと、二重メモリ解放を起こすこと。完全にunsafeの反面教師みたいなコードだった。
※パフォーマンスも悪い(work_vec.remove(0)によって)が、これは今回は置いておく

初心者の推論と検証なので鵜呑みにしないでください

fn merge_sort<T: Ord>(data: &mut [T]) {
    if data.len() <= 1 {
        return;
    }

    let length = data.len();
    let middle = length / 2;
    MergeSort::merge_sort(&mut data[..middle]);
    MergeSort::merge_sort(&mut data[middle..]);

    let mut work_vec: Vec<T> = Vec::with_capacity(length);
    unsafe {
        data.as_ptr()
            .copy_to_nonoverlapping(work_vec.as_mut_ptr(), length);
        work_vec.set_len(length);
    }

    work_vec[middle..].reverse();

    let mut r = length;

    for i in 0..length {
        r -= 1;
        if work_vec[0] > work_vec[r] {
            data[i] = work_vec.remove(r);
        } else {
            data[i] = work_vec.remove(0);
        }
    }
}

このコードで、以下のような型をソートしようとする。

struct BoxedValue {
    value: Box<i32>,
}

impl Ord for BoxedValue {
    // 省略 実際はPartialOrd、PartialEq, Eqもいる
    // 比較はself.value.deref().cmp(&*other.value)のような形
}

まず、dataからwork_vecに値がコピーされる。この時のコピーはCのmemcpyと同じ挙動のため、Shallow Copyになる(Boxの中身はコピーされず、参照がコピーされる)

data.as_ptr().copy_to_nonoverlapping(work_vec.as_mut_ptr(), length);

そして、work_vecの保持している値(dataからコピーしたもの)がdataに束縛される(所有権の移動)

data[i] = work_vec.remove(0);

このとき、まずwork_vec.remove(0)が評価され、work_vec[0]を初期化(メモリ解放)しようとする(Box<i32>を経由してi32を初期化(メモリを解放、free()を呼び出し)する)。
次に、data[i]は既に初期化(メモリ割り当て、値で初期化)済みであることをRustコンパイラは知っているため、これも初期化(メモリ解放)しようとする。
このとき、work_vecとdataの中身は順序は違うものの同一のため、最初のwork_vec.remove()と、data[i]への値束縛時に同じメモリが解放される(計二回)。そして2重メモリ解放の実行時エラーが出る。

参考：Rustnomicon: チェックされない初期化されていないメモリ

初心者の推論と検証なので鵜呑みにしないでください

最後に

ソートならサクサク書けるかな？と思ったけど、全然そんなことはなかった。
今回のでだいぶ勉強になったし、次はそんなに迷わないと信じたい