in-place radix sort with sentinel

radix_sort.rs

use std::{ops::{Add, AddAssign}, convert::TryInto};

/// In-place unstable radix sort with sentinel
///
/// - `L` is the type used for bin counts (must be able to hold `arr.len()`).
/// - `key: F` is a function taking an array item + a level, and returning an alphabet symbol.
/// - `AL` is the alphabet length, i.e. number of bins. it must hold that `key(...) < AL`.
///
/// 0 is the sentinel, meaning that for any item, `key(item, max_level) == 0`.
/// `max_level` may differ among items.
pub fn radix_sort<const AL: usize, L, B, X, F>(key: F, arr: &mut [X])
where
	F: Copy + Fn(&X, B) -> usize,
	B: Copy + Add<B, Output=B>,
	L: Copy + Add<L, Output=L> + AddAssign<L> + Ord + Into<usize>,
	usize: TryInto<L>,
	u8: Into<B> + Into<L>,
{
	assert!(TryInto::<L>::try_into(arr.len()).is_ok());
	radix_sort_step::<AL, L, _, _, _>(key, arr, 0.into());
}

fn radix_sort_step<const AL: usize, L, B, X, F>(key: F, arr: &mut [X], level: B)
where
	F: Copy + Fn(&X, B) -> usize,
	B: Copy + Add<B, Output=B>,
	L: Copy + Add<L, Output=L> + AddAssign<L> + Ord + Into<usize>,
	u8: Into<B> + Into<L>,
{
	if arr.len() <= 1 {
		return
	}

	// first pass: collect frequencies
	let mut ptrs: [L; AL] = [0.into(); AL];
	for x in &*arr {
		ptrs[key(x, level)] += 1.into();
	}

	// accumulate frequencies into bucket pointers
	let mut pos = 0.into();
	let mut endptrs = [0.into(); AL];
	for a in 0..AL {
		endptrs[a] = pos + ptrs[a];
		ptrs[a] = pos;
		pos = endptrs[a];
	}

	// second pass: sort into buckets
	for a in 0..AL {
		while ptrs[a] < endptrs[a] {
			let b = key(&arr[ptrs[a].into()], level);
			arr.swap(ptrs[a].into(), ptrs[b].into());
			ptrs[b] += 1.into();
		}
	}

	// recursively sort each bucket (except the first one, which is sentinel)
	for a in 1..AL {
		let arr = &mut arr[ptrs[a-1].into()..ptrs[a].into()];
		radix_sort_step::<AL, L, _, _, _>(key, arr, level + 1.into());
	}
}