スプレッドソートで配列を並び替える

set -euo pipefail

WORKDIR="$(mktemp -d)"
trap 'rm -rf "$WORKDIR"' EXIT

cat > "$WORKDIR/Dockerfile" <<'EOF'
FROM rust:1.95.0

WORKDIR /app

RUN mkdir -p src

RUN cat > Cargo.toml <<'CARGO'
[package]
name = "rust-benchmark"
version = "0.1.0"
edition = "2021"

[profile.release]
lto = true
codegen-units = 1
panic = "abort"
CARGO

RUN cat > src/main.rs <<'RUST'
use std::{
    env,
    process::Command,
    time::{Duration, Instant},
};
const MIN_POWER: u32 = 8;
const MAX_POWER: u32 = 18;
const RUNS: usize = 8192;
fn insertion_sort(a: &mut [usize]) {
    for i in 1..a.len() {
        let mut j = i;
        while j > 0 && a[j - 1] > a[j] {
            a.swap(j - 1, j);
            j -= 1;
        }
    }
}



const MEAN_BIN_SIZE: usize = 4;
const MIN_BIN_COUNT: usize = 16;

fn rough_log2_size(n: usize) -> u32 {
    if n == 0 {
        0
    } else {
        usize::BITS - 1 - n.leading_zeros()
    }
}

fn get_max_count(log_range: u32, count: usize) -> usize {
    const MAX_SPLITS: u32 = 11;
    const LOG_CONST: u32 = 2;
    const LOG_MEAN_BIN_SIZE: u32 = 2;
    const LOG_MIN_SPLIT_COUNT: u32 = 4;
    let data_size = usize::BITS;

    let log_size = rough_log2_size(count);
    let denom = log_size.min(MAX_SPLITS).max(1);
    let mut relative_width = (LOG_CONST * log_range) / denom;
    if data_size <= relative_width {
        relative_width = data_size - 1;
    }
    let shift = if relative_width < LOG_MEAN_BIN_SIZE + LOG_MIN_SPLIT_COUNT {
        LOG_MEAN_BIN_SIZE + LOG_MIN_SPLIT_COUNT
    } else {
        relative_width
    };
    1usize << shift.min(31)
}

fn spread_bin_index(x: usize, min: usize, max: usize, bin_count: usize) -> usize {
    if max == min {
        0
    } else {
        (((x - min) as u128 * bin_count as u128 / (max - min) as u128) as usize)
            .min(bin_count.saturating_sub(1))
    }
}

fn spreadsort_rec(a: &mut [usize]) {
    let n = a.len();
    let max_count = get_max_count(rough_log2_size(a.iter().max().copied().unwrap_or(0)), n);
    if n < max_count {
        insertion_sort(a);
        return;
    }

    let min = *a.iter().min().unwrap();
    let max = *a.iter().max().unwrap();
    if min == max {
        return;
    }

    let log_range = rough_log2_size(max - min);
    let bin_count = (n / MEAN_BIN_SIZE).max(MIN_BIN_COUNT).min(n);
    let range = max - min;

    if bin_count >= range + 1 {
        insertion_sort(a);
        return;
    }

    let mut count = vec![0usize; bin_count];
    for &x in a.iter() {
        count[spread_bin_index(x, min, max, bin_count)] += 1;
    }

    let mut offset = vec![0usize; bin_count + 1];
    for i in 0..bin_count {
        offset[i + 1] = offset[i] + count[i];
    }

    let mut temp = vec![0usize; n];
    let mut cursor = offset.clone();
    for &x in a.iter() {
        let bin = spread_bin_index(x, min, max, bin_count);
        temp[cursor[bin]] = x;
        cursor[bin] += 1;
    }
    a.copy_from_slice(&temp);

    let fallback = get_max_count(log_range, n);
    for i in 0..bin_count {
        let start = offset[i];
        let end = offset[i + 1];
        let len = end - start;
        if len < 2 {
            continue;
        }
        if len < fallback {
            insertion_sort(&mut a[start..end]);
        } else {
            spreadsort_rec(&mut a[start..end]);
        }
    }
}

fn spread_sort(a: &mut [usize]) {
    if !a.is_empty() {
        spreadsort_rec(a);
    }
}


fn benchmark_sort(array: &mut [usize]) {

    spread_sort(array);

}

fn shuffled(size: usize, seed: u64) -> Vec<usize> {
    let mut v: Vec<usize> = (1..=size).collect();

    let mut state = seed;

    for i in (1..size).rev() {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;

        let j = (state as usize) % (i + 1);

        v.swap(i, j);
    }

    v
}

fn memory_usage_kb() -> usize {
    let contents = std::fs::read_to_string("/proc/self/status")
        .unwrap_or_default();

    for line in contents.lines() {
        if let Some(rest) = line.strip_prefix("VmHWM:") {
            let kb = rest
                .split_whitespace()
                .next()
                .unwrap_or("0")
                .parse::<usize>()
                .unwrap_or(0);

            return kb;
        }
    }

    0
}

fn micros(d: Duration) -> u128 {
    d.as_micros()
}

fn run_once(size: usize, seed: usize) -> (u128, usize) {
    let expected: Vec<usize> = (1..=size).collect();
    let mut array = shuffled(size, seed as u64);

    let start = Instant::now();

    benchmark_sort(&mut array);

    let elapsed = start.elapsed();

    if array != expected {
        panic!(
            "sort failed with seed {} for size {}",
            seed,
            size
        );
    }

    (micros(elapsed), memory_usage_kb())
}

fn run_child(args: &[String]) {
    let size = args[2].parse::<usize>().expect("invalid size");
    let seed = args[3].parse::<usize>().expect("invalid seed");
    let (elapsed_us, mem) = run_once(size, seed);
    println!("{} {}", elapsed_us, mem);
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.get(1).is_some_and(|arg| arg == "--run-once") {
        run_child(&args);
        return;
    }

    println!(
        "| {:>10} | {:>15} | {:>15} | {:>15} | {:>15} |",
        "Size",
        "Average time",
        "Maximum time",
        "Average memory",
        "Maximum memory"
    );

    println!(
        "|{:-<11}:|{:-<16}:|{:-<16}:|{:-<16}:|{:-<16}:|",
        "",
        "",
        "",
        "",
        ""
    );

    for power in MIN_POWER..=MAX_POWER {
        let size = 1usize << power;

        let mut total_time: u128 = 0;
        let mut max_time: u128 = 0;

        let mut total_mem: usize = 0;
        let mut max_mem: usize = 0;

        for seed in 1..=RUNS {
            let output = Command::new(env::current_exe().expect("failed to find current executable"))
                .arg("--run-once")
                .arg(size.to_string())
                .arg(seed.to_string())
                .output()
                .expect("failed to run benchmark child process");

            if !output.status.success() {
                panic!(
                    "benchmark child process failed: {}",
                    String::from_utf8_lossy(&output.stderr)
                );
            }

            let stdout = String::from_utf8(output.stdout)
                .expect("child process returned non-UTF-8 output");
            let mut fields = stdout.split_whitespace();
            let elapsed_us = fields
                .next()
                .expect("missing elapsed time")
                .parse::<u128>()
                .expect("invalid elapsed time");
            let mem = fields
                .next()
                .expect("missing memory usage")
                .parse::<usize>()
                .expect("invalid memory usage");

            total_time += elapsed_us;

            if elapsed_us > max_time {
                max_time = elapsed_us;
            }

            total_mem += mem;

            if mem > max_mem {
                max_mem = mem;
            }
        }

        let avg_time = total_time / RUNS as u128;
        let avg_mem = total_mem / RUNS;

        println!(
            "| {:>10} | {:>15} | {:>15} | {:>15} | {:>15} |",
            size,
            format!("{}.{:06}", avg_time / 1_000_000, avg_time % 1_000_000),
            format!("{}.{:06}", max_time / 1_000_000, max_time % 1_000_000),
            avg_mem,
            max_mem
        );
    }
}
RUST

RUN cargo build --release

CMD ["./target/release/rust-benchmark"]
EOF

docker build -t rust-benchmark "$WORKDIR"
docker run --rm --init rust-benchmark