ポリフェーズマージソートで配列を並び替える

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::{
    alloc::{GlobalAlloc, Layout, System},
    env,
    process::Command,
    sync::atomic::{AtomicUsize, Ordering},
    time::{Duration, Instant},
};

/// Counts live heap bytes and the high-water mark so auxiliary sort buffers
/// (swap Vecs, etc.) are measured as explicit heap growth during the sort.
struct TrackingAllocator;

static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0);
static PEAK_BYTES: AtomicUsize = AtomicUsize::new(0);

fn record_alloc(size: usize) {
    let live = LIVE_BYTES.fetch_add(size, Ordering::Relaxed) + size;
    PEAK_BYTES.fetch_max(live, Ordering::Relaxed);
}

unsafe impl GlobalAlloc for TrackingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ptr = System.alloc(layout);
        if !ptr.is_null() {
            record_alloc(layout.size());
        }
        ptr
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed);
        System.dealloc(ptr, layout);
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ptr = System.alloc_zeroed(layout);
        if !ptr.is_null() {
            record_alloc(layout.size());
        }
        ptr
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let new_ptr = System.realloc(ptr, layout, new_size);
        if !new_ptr.is_null() {
            LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed);
            record_alloc(new_size);
        }
        new_ptr
    }
}

#[global_allocator]
static GLOBAL: TrackingAllocator = TrackingAllocator;
const MIN_POWER: u32 = 8;
const MAX_POWER: u32 = 18;
const RUNS: usize = 8192;


const NUM_TAPES: usize = 3;
const RUN_SIZE: usize = 32;

fn merge_runs(left: &[usize], right: &[usize]) -> Vec<usize> {
    let mut out = Vec::with_capacity(left.len() + right.len());
    let (mut i, mut j) = (0, 0);
    while i < left.len() && j < right.len() {
        if left[i] <= right[j] {
            out.push(left[i]);
            i += 1;
        } else {
            out.push(right[j]);
            j += 1;
        }
    }
    out.extend_from_slice(&left[i..]);
    out.extend_from_slice(&right[j..]);
    out
}

fn create_runs(a: &[usize], run_size: usize) -> Vec<Vec<usize>> {
    let mut runs = Vec::new();
    let mut i = 0;
    while i < a.len() {
        let end = (i + run_size).min(a.len());
        let mut run = a[i..end].to_vec();
        run.sort_unstable();
        runs.push(run);
        i = end;
    }
    runs
}

fn next_fibonacci_at_least(n: usize) -> (usize, usize) {
    let (mut prev, mut curr) = (1usize, 1usize);
    while curr < n {
        let next = prev + curr;
        prev = curr;
        curr = next;
    }
    (prev, curr)
}

fn distribute_fibonacci(runs: Vec<Vec<usize>>) -> [Vec<Vec<usize>>; NUM_TAPES] {
    let mut tapes: [Vec<Vec<usize>>; NUM_TAPES] = [vec![], vec![], vec![]];
    let n = runs.len();
    if n == 0 {
        return tapes;
    }
    if n == 1 {
        tapes[1].push(runs.into_iter().next().unwrap());
        return tapes;
    }

    let (fib_prev, fib_target) = next_fibonacci_at_least(n);
    let dummies = fib_target - n;
    let on_tape2 = fib_prev.saturating_sub(dummies);
    let on_tape1 = n - on_tape2;

    for (idx, run) in runs.into_iter().enumerate() {
        if idx < on_tape1 {
            tapes[1].push(run);
        } else {
            tapes[2].push(run);
        }
    }
    tapes
}

fn count_runs(tapes: &[Vec<Vec<usize>>; NUM_TAPES]) -> usize {
    tapes.iter().map(|t| t.len()).sum()
}

fn rotate_tapes(tapes: &mut [Vec<Vec<usize>>; NUM_TAPES]) {
    tapes.swap(0, 1);
    tapes.swap(1, 2);
}

fn polyphase_pass(tapes: &mut [Vec<Vec<usize>>; NUM_TAPES]) -> bool {
    let mut merged = false;
    while !tapes[1].is_empty() && !tapes[2].is_empty() {
        let left = tapes[1].remove(0);
        let right = tapes[2].remove(0);
        tapes[0].push(merge_runs(&left, &right));
        merged = true;
    }
    merged
}

fn merge_all_remaining(tapes: &mut [Vec<Vec<usize>>; NUM_TAPES]) -> Vec<usize> {
    let mut all: Vec<Vec<usize>> = Vec::new();
    for tape in tapes.iter() {
        all.extend(tape.iter().cloned());
    }
    while all.len() > 1 {
        let a = all.remove(0);
        let b = all.remove(0);
        all.push(merge_runs(&a, &b));
    }
    all.pop().unwrap_or_default()
}

fn polyphase_merge_sort(a: &mut [usize]) {
    if a.len() <= 1 {
        return;
    }
    let runs = create_runs(a, RUN_SIZE);
    if runs.len() <= 1 {
        if let Some(r) = runs.first() {
            a.copy_from_slice(r);
        }
        return;
    }

    let mut tapes = distribute_fibonacci(runs);
    let mut idle = 0usize;
    while count_runs(&tapes) > 1 {
        if polyphase_pass(&mut tapes) {
            rotate_tapes(&mut tapes);
            idle = 0;
        } else {
            idle += 1;
            if idle > NUM_TAPES * 4 {
                break;
            }
            rotate_tapes(&mut tapes);
        }
    }

    let result = if count_runs(&tapes) == 1 {
        tapes
            .iter()
            .find_map(|t| t.first().cloned())
            .unwrap_or_default()
    } else {
        merge_all_remaining(&mut tapes)
    };
    a.copy_from_slice(&result);
}


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

    polyphase_merge_sort(array);

}

fn is_non_decreasing(a: &[usize]) -> bool {
    a.windows(2).all(|w| w[0] <= w[1])
}

fn same_multiset(a: &[usize], b: &[usize]) -> bool {
    if a.len() != b.len() {
        return false;
    }

    let mut left = a.to_vec();
    let mut right = b.to_vec();
    left.sort_unstable();
    right.sort_unstable();
    left == right
}

fn check_correctness_case(label: &str, mut input: Vec<usize>) {
    let original = input.clone();

    benchmark_sort(&mut input);

    if !is_non_decreasing(&input) {
        panic!("correctness case {}: output is not sorted", label);
    }

    if !same_multiset(&input, &original) {
        panic!("correctness case {}: elements were lost or added", label);
    }
}

fn few_unique_values(size: usize, unique: usize, seed: u64) -> Vec<usize> {
    let mut state = seed;

    (0..size)
        .map(|_| {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            (state as usize % unique) + 1
        })
        .collect()
}

fn run_correctness_checks() {
    check_correctness_case("empty", vec![]);
    check_correctness_case("single", vec![42]);
    check_correctness_case("duplicates", vec![3, 1, 3, 2, 1, 2]);
    check_correctness_case("sorted", vec![1, 2, 3, 4, 5]);
    check_correctness_case("reverse", vec![5, 4, 3, 2, 1]);
    check_correctness_case("all_equal", vec![7, 7, 7, 7]);
    check_correctness_case("skewed_range", vec![1_000_000, 2, 1_000_001, 1, 999_999]);
    // Static-buffer Grail skips the in-buffer build when key collection is sparse
    // (ideal_buffer = false). Exercising that path catches regressions in buffer gating.
    check_correctness_case(
        "few_keys_len16",
        vec![2, 2, 2, 2, 2, 2, 2, 2, 4, 3, 1, 2, 3, 4, 1, 4],
    );
    // Seed 0 is a fixed point of the xorshift below, so it would degenerate into
    // yet another all-equal case instead of a 4-value mix. Start at 1.
    for seed in 1..=32 {
        check_correctness_case(
            &format!("few_keys_len32_seed_{seed}"),
            few_unique_values(32, 4, seed),
        );
    }
}


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 micros(d: Duration) -> u128 {
    d.as_micros()
}

fn input_array(size: usize, seed: u64) -> Vec<usize> {
    shuffled(size, seed)
}

/// Peak heap growth during `benchmark_sort`, in KiB (explicit buffers such as swap).
fn run_once(size: usize, seed: usize) -> (u128, usize) {
    let mut array = input_array(size, seed as u64);

    let base_bytes = LIVE_BYTES.load(Ordering::Relaxed);
    PEAK_BYTES.store(base_bytes, Ordering::Relaxed);

    let start = Instant::now();

    benchmark_sort(&mut array);

    let elapsed = start.elapsed();
    let peak_bytes = PEAK_BYTES.load(Ordering::Relaxed);
    let aux_kb = peak_bytes.saturating_sub(base_bytes) / 1024;

    let expected: Vec<usize> = (1..=size).collect();
    if array != expected {
        panic!(
            "sort failed with seed {} for size {}",
            seed,
            size
        );
    }

    (micros(elapsed), aux_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;
    }

    run_correctness_checks();

    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 aux_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 += aux_mem;

            if aux_mem > max_mem {
                max_mem = aux_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