use core::mem::{self, MaybeUninit};
use core::ptr::{self, NonNull};
use core::slice;
use super::counting::{
Count, STACK_BYTES, Stack, bump, counters, exclusive_scan, in_order, in_order_at_first,
key_slots, min_max, places4, range,
};
use super::key::IntegerKey;
use super::key::sealed::Bits;
pub fn pigeonhole<T: IntegerKey>(arr: &mut [T]) {
let len = arr.len();
if len < 2 {
return;
}
if in_order(arr, |x, y| y < x) {
range(arr[0], arr[len - 1], len, SORT);
return;
}
if in_order(arr, |x, y| y > x) {
range(arr[len - 1], arr[0], len, SORT);
arr.reverse();
return;
}
let (min, max) = min_max(arr);
let k = range(min, max, len, SORT);
let none = NonNull::<()>::dangling().as_ptr();
let keys = arr.as_mut_ptr();
unsafe {
if len <= u16::MAX as usize {
place::<T, T, (), u16, false>(keys, none, len, min, k, |_| min);
} else if len <= u32::MAX as usize {
place::<T, T, (), u32, false>(keys, none, len, min, k, |_| min);
} else {
place::<T, T, (), usize, false>(keys, none, len, min, k, |_| min);
}
}
}
pub fn pigeonhole_by_key<T, K, F>(arr: &mut [T], mut key: F)
where
K: IntegerKey,
F: FnMut(&T) -> K,
{
let len = arr.len();
if len < 2 || mem::size_of::<T>() == 0 {
return;
}
let mut stack = Stack([MaybeUninit::uninit(); STACK_BYTES]);
let mut heap = Vec::new();
let slots = key_slots::<K>(len, &mut stack, &mut heap);
let first = key(&arr[0]);
slots[0].write(first);
let (mut min, mut max, mut prev) = (first, first, first);
let (mut descent, mut ascent) = (false, false);
for (slot, x) in slots[1..].iter_mut().zip(&arr[1..]) {
let x = key(x);
slot.write(x);
min = min.min(x);
max = max.max(x);
descent |= x < prev;
ascent |= x > prev;
prev = x;
}
let k = range(min, max, len, SORT);
if !descent {
return;
}
if !ascent {
arr.reverse();
return;
}
let (keys, elems) = (slots.as_mut_ptr().cast::<K>(), arr.as_mut_ptr());
unsafe {
if k <= 1 << 8 {
by_holes::<K, u8, T>(keys, elems, len, min, k);
} else if k <= 1 << 16 {
by_holes::<K, u16, T>(keys, elems, len, min, k);
} else if k as u64 <= 1 << 32 {
by_holes::<K, u32, T>(keys, elems, len, min, k);
} else {
by_holes::<K, usize, T>(keys, elems, len, min, k);
}
}
}
#[inline(always)]
unsafe fn by_holes<K: IntegerKey, H: IntegerKey + Bits, T>(
keys: *mut K,
elems: *mut T,
len: usize,
min: K,
k: usize,
) {
debug_assert!(mem::size_of::<H>() <= mem::size_of::<K>());
unsafe {
if len <= u16::MAX as usize {
place::<K, H, T, u16, true>(keys, elems, len, min, k, H::of);
} else if len <= u32::MAX as usize {
place::<K, H, T, u32, true>(keys, elems, len, min, k, H::of);
} else {
place::<K, H, T, usize, true>(keys, elems, len, min, k, H::of);
}
}
}
const SORT: &str = "pigeonhole sort";
const SWAP_MIN_LEN: usize = 1 << 15;
const SWAP_BYTES_PER_STEP: usize = 40;
const SWAP_HOLES_PER_ELEMENT: usize = 2;
const SWAP_DESCENTS_PART: usize = 16;
const SWAP_MAX_RANGE: usize = 16;
const SWAP_MIN_PER_KEY: usize = 16;
const SWAP_ALWAYS_BYTES: usize = 16;
const POSITION_MIN_PER_ELEMENT: usize = 2;
const SPARSE_PER_ELEMENT: usize = 32;
unsafe fn place<K: IntegerKey, H: IntegerKey, E, C: Count, const NARROW: bool>(
keys: *mut K,
elems: *mut E,
len: usize,
min: K,
k: usize,
hole: impl Fn(usize) -> H,
) {
let tables = if k <= TABLES_MAX_RANGE && len >= TABLES_MIN_PER_KEY * k {
TABLES
} else {
1
};
let mut stack = Stack([MaybeUninit::uninit(); STACK_BYTES]);
let mut heap = Vec::new();
let ends = counters::<C>(tables * k, &mut stack, &mut heap).as_mut_ptr();
let holes = keys.cast::<H>();
let sparse = k > SPARSE_PER_ELEMENT * len;
unsafe {
let descents = if tables == TABLES {
count::<K, H, C, NARROW, TABLES>(keys, len, min, k, ends, &hole)
} else {
let mut descents = 0;
let mut prev = 0;
if in_order_at_first(slice::from_raw_parts(keys, len)) {
let mut i = 0;
while i + 4 <= len {
let h = [0, 1, 2, 3].map(|j| (*keys.add(i + j)).offset(min));
for (j, &h) in h.iter().enumerate() {
if NARROW {
holes.add(i + j).write(hole(h));
}
descents += usize::from(h < prev);
prev = h;
}
let p = places4(h, h.map(|h| (*ends.add(h)).get()));
for j in 0..4 {
*ends.add(h[j]) = C::ZERO.add(p[j] + 1);
}
i += 4;
}
for i in i..len {
let h = (*keys.add(i)).offset(min);
if NARROW {
holes.add(i).write(hole(h));
}
descents += usize::from(h < prev);
prev = h;
bump(ends.add(h));
}
} else {
for i in 0..len {
let h = (*keys.add(i)).offset(min);
if NARROW {
holes.add(i).write(hole(h));
}
descents += usize::from(h < prev);
prev = h;
bump(ends.add(h));
}
}
descents
};
let base = if NARROW {
hole(0)
} else {
ptr::read((&raw const min).cast::<H>())
};
let ordered = descents * SWAP_DESCENTS_PART <= len;
let by_swaps = !ordered
&& (by_swaps::<H, E>(len, k) || (k <= SWAP_MAX_RANGE && len >= SWAP_MIN_PER_KEY * k));
if by_swaps && !sparse {
exclusive_scan(slice::from_raw_parts_mut(ends, k));
return swaps::<H, E, C, true>(holes, elems, base, k, len, ends, ends);
}
let mut next_stack = Stack([MaybeUninit::uninit(); STACK_BYTES]);
let mut next_heap = Vec::new();
let next = key_slots::<C>(k, &mut next_stack, &mut next_heap)
.as_mut_ptr()
.cast::<C>();
let mut sum = 0;
if sparse {
each_hole(ends, k, |h, count| {
*next.add(h) = C::ZERO.add(sum);
sum += count;
*ends.add(h) = C::ZERO.add(sum);
});
} else {
let mut h = 0;
while h + 4 <= k {
let c = [0, 1, 2, 3].map(|j| (*ends.add(h + j)).get());
let (ab, abc) = (c[0] + c[1], c[0] + c[1] + c[2]);
let starts = [sum, sum + c[0], sum + ab, sum + abc];
sum += abc + c[3];
for j in 0..4 {
*next.add(h + j) = C::ZERO.add(starts[j]);
*ends.add(h + j) = C::ZERO.add(starts[j] + c[j]);
}
h += 4;
}
for h in h..k {
*next.add(h) = C::ZERO.add(sum);
sum += (*ends.add(h)).get();
*ends.add(h) = C::ZERO.add(sum);
}
}
if by_swaps {
swaps::<H, E, C, false>(holes, elems, base, k, len, ends, next);
} else {
let walk = if sparse {
Walk::Sparse
} else if ordered || k >= POSITION_MIN_PER_ELEMENT * len {
Walk::Places(len)
} else {
Walk::Holes
};
cycles(holes, elems, base, k, ends, next, walk);
}
}
}
const TABLES_MAX_RANGE: usize = 256;
const TABLES_MIN_PER_KEY: usize = 16;
const TABLES: usize = 8;
#[inline(always)]
unsafe fn count<K: IntegerKey, H, C: Count, const NARROW: bool, const N: usize>(
keys: *mut K,
len: usize,
min: K,
k: usize,
ends: *mut C,
hole: &impl Fn(usize) -> H,
) -> usize {
let holes = keys.cast::<H>();
let mut descents = 0;
let mut prev = 0;
unsafe {
let mut i = 0;
while i + N <= len {
for j in 0..N {
let h = (*keys.add(i + j)).offset(min);
if NARROW {
holes.add(i + j).write(hole(h));
}
descents += usize::from(h < prev);
prev = h;
bump(ends.add(h * N + j));
}
i += N;
}
for i in i..len {
let h = (*keys.add(i)).offset(min);
if NARROW {
holes.add(i).write(hole(h));
}
descents += usize::from(h < prev);
prev = h;
bump(ends.add(h * N));
}
if N > 1 {
for h in 0..k {
let mut sum = 0;
for t in 0..N {
sum += (*ends.add(h * N + t)).get();
}
*ends.add(h) = C::ZERO.add(sum);
}
}
}
descents
}
fn by_swaps<K, E>(len: usize, k: usize) -> bool {
let bytes = mem::size_of::<E>() + mem::size_of::<K>();
len.max(k / SWAP_HOLES_PER_ELEMENT) >= SWAP_MIN_LEN * (bytes / SWAP_BYTES_PER_STEP).max(1)
}
#[inline(always)]
unsafe fn each_hole<C: Count>(counters: *mut C, k: usize, mut f: impl FnMut(usize, usize)) {
unsafe {
let mut h = 0;
while h + 8 <= k {
let block = ptr::read(counters.add(h).cast::<[C; 8]>());
if block.iter().any(|&c| c != C::ZERO) {
for (j, c) in block.iter().enumerate() {
if *c != C::ZERO {
f(h + j, c.get());
}
}
}
h += 8;
}
while h < k {
let c = (*counters.add(h)).get();
if c != 0 {
f(h, c);
}
h += 1;
}
}
}
unsafe fn cycles<H: IntegerKey, E, C: Count>(
holes: *mut H,
elems: *mut E,
base: H,
k: usize,
ends: *mut C,
next: *mut C,
walk: Walk,
) {
unsafe {
match walk {
Walk::Sparse => each_hole(ends, k, |h, end| fill(holes, elems, base, next, h, end)),
Walk::Places(len) => walk_places(holes, elems, base, len, ends, next),
Walk::Holes => {
for h in 0..k {
fill(holes, elems, base, next, h, (*ends.add(h)).get());
}
}
}
}
}
#[derive(Clone, Copy)]
enum Walk {
Sparse,
Holes,
Places(usize),
}
#[inline(always)]
unsafe fn walk_places<H: IntegerKey, E, C: Count>(
holes: *mut H,
elems: *mut E,
base: H,
len: usize,
ends: *mut C,
next: *mut C,
) {
unsafe {
for i in 0..len {
let mut x = *holes.add(i);
let mut d = x.offset(base);
let start = if d == 0 { 0 } else { (*ends.add(d - 1)).get() };
if start <= i && i < (*ends.add(d)).get() {
continue;
}
let mut cur = ptr::read(elems.add(i).cast::<MaybeUninit<E>>());
loop {
let n = next.add(d);
let j = (*n).get().max(i);
*n = C::ZERO.add(j + 1);
if j == i {
break;
}
let y = *holes.add(j);
let taken = ptr::read(elems.add(j).cast::<MaybeUninit<E>>());
*holes.add(j) = x;
ptr::write(elems.add(j).cast::<MaybeUninit<E>>(), cur);
x = y;
cur = taken;
d = x.offset(base);
}
*holes.add(i) = x;
ptr::write(elems.add(i).cast::<MaybeUninit<E>>(), cur);
}
}
}
#[inline(always)]
unsafe fn fill<H: IntegerKey, E, C: Count>(
holes: *mut H,
elems: *mut E,
base: H,
next: *mut C,
h: usize,
end: usize,
) {
unsafe {
let mut i = (*next.add(h)).get();
while i < end {
let mut x = *holes.add(i);
let mut d = x.offset(base);
if d != h {
let mut cur = ptr::read(elems.add(i).cast::<MaybeUninit<E>>());
loop {
let j = bump(next.add(d));
let y = *holes.add(j);
let taken = ptr::read(elems.add(j).cast::<MaybeUninit<E>>());
*holes.add(j) = x;
ptr::write(elems.add(j).cast::<MaybeUninit<E>>(), cur);
x = y;
cur = taken;
d = x.offset(base);
if d == h {
break;
}
}
*holes.add(i) = x;
ptr::write(elems.add(i).cast::<MaybeUninit<E>>(), cur);
}
i += 1;
}
}
}
unsafe fn swaps<H: IntegerKey, E, C: Count, const STARTS: bool>(
holes: *mut H,
elems: *mut E,
base: H,
k: usize,
total: usize,
ends: *mut C,
next: *mut C,
) {
let mut stack = Stack([MaybeUninit::uninit(); STACK_BYTES]);
let mut heap = Vec::new();
let open = key_slots::<(usize, usize)>(k.min(total) + 1, &mut stack, &mut heap)
.as_mut_ptr()
.cast::<(usize, usize)>();
let mut len = 0;
unsafe {
let mut add = |h: usize, end: usize| {
open.add(len).write((h, end));
len += usize::from((*next.add(h)).get() < end);
};
if STARTS {
for h in 1..k {
add(h - 1, (*next.add(h)).get());
}
add(k - 1, total);
} else {
each_hole(ends, k, add);
}
while len > 0 {
let mut kept = 0;
for o in 0..len {
let (h, end) = *open.add(o);
let mut p = (*next.add(h)).get();
while p < end {
let t = bump(next.add((*holes.add(p)).offset(base)));
if mem::size_of::<E>() + mem::size_of::<H>() <= SWAP_ALWAYS_BYTES || t != p {
ptr::swap(holes.add(p), holes.add(t));
ptr::swap(elems.add(p), elems.add(t));
}
p += 1;
}
if (*next.add(h)).get() < end {
*open.add(kept) = (h, end);
kept += 1;
}
}
len = kept;
}
}
}