use core::alloc::Layout;
use core::hint::select_unpredictable;
use core::mem::{self, MaybeUninit};
use core::ptr;
pub fn tree<T: Ord>(arr: &mut [T]) {
if arr.len() < 2 || mem::size_of::<T>() == 0 {
return;
}
let (size, align) = (mem::size_of::<T>(), mem::align_of::<T>());
if size > 4 && size <= 8 && align <= 8 {
lanes::<T, [MaybeUninit<u64>; 1]>(arr);
} else if size > 16 && size <= 24 && align <= 8 {
lanes::<T, [MaybeUninit<u64>; 3]>(arr);
} else {
lanes::<T, ()>(arr);
}
}
fn lanes<T: Ord, V>(arr: &mut [T]) {
let len = arr.len();
if len >= u32::MAX as usize {
sort::<T, usize, V, 8>(arr);
} else if len >= LANES_8_MIN {
sort::<T, u32, V, 8>(arr);
} else if len >= LANES_4_MIN {
sort::<T, u32, V, 4>(arr);
} else {
sort::<T, u32, V, 2>(arr);
}
}
const LANES_4_MIN: usize = 1024;
const LANES_8_MIN: usize = 1 << 18;
const HEIGHT: usize = 96;
const NONE: usize = usize::MAX;
trait Link: Copy + Eq {
const NIL: Self;
fn new(i: usize) -> Self;
fn get(self) -> usize;
unsafe fn pick(p: *const [Self; 2], right: bool) -> Self;
}
impl Link for u32 {
const NIL: Self = u32::MAX;
#[inline(always)]
fn new(i: usize) -> Self {
i as u32
}
#[inline(always)]
fn get(self) -> usize {
self as usize
}
#[inline(always)]
unsafe fn pick(p: *const [Self; 2], right: bool) -> Self {
let pair = unsafe { p.cast::<u64>().read_unaligned() };
let (l, r) = if cfg!(target_endian = "little") {
(pair as u32, (pair >> 32) as u32)
} else {
((pair >> 32) as u32, pair as u32)
};
select_unpredictable(right, r, l)
}
}
impl Link for usize {
const NIL: Self = usize::MAX;
#[inline(always)]
fn new(i: usize) -> Self {
i
}
#[inline(always)]
fn get(self) -> usize {
self
}
#[inline(always)]
unsafe fn pick(p: *const [Self; 2], right: bool) -> Self {
let [l, r] = unsafe { *p };
select_unpredictable(right, r, l)
}
}
#[repr(C)]
struct Node<L, V> {
link: [L; 2],
val: V,
}
type Path = [MaybeUninit<usize>; HEIGHT];
struct Tree<T, L, V> {
nodes: *mut Node<L, V>,
bal: *mut i8,
keys: *const T,
root: usize,
}
impl<T: Ord, L: Link, V> Tree<T, L, V> {
const EMBED: bool = mem::size_of::<V>() != 0;
#[inline(always)]
unsafe fn key(&self, i: usize) -> *const T {
unsafe {
if Self::EMBED {
(&raw const (*self.nodes.add(i)).val).cast::<T>()
} else {
self.keys.add(i)
}
}
}
#[inline(always)]
unsafe fn link(&self, i: usize) -> *mut [L; 2] {
unsafe { &raw mut (*self.nodes.add(i)).link }
}
#[inline(always)]
unsafe fn child(&self, i: usize, dir: usize) -> L {
unsafe { *self.link(i).cast::<L>().add(dir) }
}
#[inline(always)]
unsafe fn set_child(&self, i: usize, dir: usize, c: L) {
unsafe { *self.link(i).cast::<L>().add(dir) = c }
}
#[inline(always)]
unsafe fn bal(&self, i: usize) -> i8 {
unsafe { *self.bal.add(i) }
}
#[inline(always)]
unsafe fn set_bal(&self, i: usize, b: i8) {
unsafe { *self.bal.add(i) = b }
}
#[inline(always)]
unsafe fn search_from(&self, x: &T, path: &mut Path, mut depth: usize, mut c: usize) -> usize {
unsafe {
loop {
let right = !(*x < *self.key(c));
let next = L::pick(self.link(c), right);
path[depth].write(c << 1 | usize::from(right));
depth += 1;
if next == L::NIL {
return depth;
}
c = next.get();
}
}
}
#[inline(always)]
unsafe fn search(&self, x: &T, path: &mut Path) -> usize {
unsafe { self.search_from(x, path, 0, self.root) }
}
#[inline(always)]
unsafe fn resume(&self, x: &T, path: &mut Path, from: usize) -> usize {
unsafe {
if from == 0 {
return self.search(x, path);
}
let e = path[from - 1].assume_init();
let c = self.child(e >> 1, e & 1);
if c == L::NIL {
return from;
}
self.search_from(x, path, from, c.get())
}
}
#[inline(always)]
unsafe fn search_lanes<const LANES: usize>(
&self,
xs: [&T; LANES],
paths: &mut [Path; LANES],
) -> [usize; LANES] {
unsafe {
let mut c = [self.root; LANES];
let mut depth = [0; LANES];
let mut active = [true; LANES];
loop {
for j in 0..LANES {
if active[j] {
let right = !(*xs[j] < *self.key(c[j]));
let next = L::pick(self.link(c[j]), right);
paths[j][depth[j]].write(c[j] << 1 | usize::from(right));
depth[j] += 1;
active[j] = next != L::NIL;
c[j] = next.get();
}
}
if active == [false; LANES] {
return depth;
}
}
}
}
#[inline(always)]
unsafe fn beside(&self, node: usize, dir: usize, path: &mut Path, mut depth: usize) -> usize {
unsafe {
path[depth].write(node << 1 | dir);
depth += 1;
let mut c = self.child(node, dir);
while c != L::NIL {
path[depth].write(c.get() << 1 | (dir ^ 1));
depth += 1;
c = self.child(c.get(), dir ^ 1);
}
depth
}
}
#[inline(always)]
unsafe fn attach(
&mut self,
q: usize,
path: &mut Path,
depth: usize,
keep: bool,
) -> (usize, usize, usize) {
unsafe {
let last = path[depth - 1].assume_init();
self.set_child(last >> 1, last & 1, L::new(q));
let mut k = depth;
while k > 0 {
k -= 1;
let e = path[k].assume_init();
let (p, dir) = (e >> 1, e & 1);
let s: i8 = if dir == 1 { 1 } else { -1 };
let b = self.bal(p);
if b == 0 {
self.set_bal(p, s);
continue;
}
if b == -s {
self.set_bal(p, 0);
return (0, HEIGHT, depth);
}
let c = self.child(p, dir).get();
let (r, ancestors);
if self.bal(c) == s {
self.set_child(p, dir, self.child(c, 1 - dir));
self.set_child(c, 1 - dir, L::new(p));
self.set_bal(p, 0);
self.set_bal(c, 0);
r = c;
if keep {
shift_down(path, k + 1, depth);
}
ancestors = depth - 1;
} else {
let g = self.child(c, 1 - dir).get();
self.set_child(c, 1 - dir, self.child(g, dir));
self.set_child(p, dir, self.child(g, 1 - dir));
self.set_child(g, dir, L::new(c));
self.set_child(g, 1 - dir, L::new(p));
let gb = self.bal(g);
self.set_bal(p, if gb == s { -s } else { 0 });
self.set_bal(c, if gb == -s { s } else { 0 });
self.set_bal(g, 0);
r = g;
if k + 2 == depth {
ancestors = k;
} else if keep {
let dg = path[k + 2].assume_init() & 1;
if dg == dir {
path[k].write(g << 1 | dir);
path[k + 1].write(c << 1 | (1 - dir));
} else {
path[k].write(g << 1 | (1 - dir));
path[k + 1].write(p << 1 | dir);
}
shift_down(path, k + 3, depth);
ancestors = depth - 1;
} else {
ancestors = depth - 1;
}
}
if k == 0 {
self.root = r;
} else {
let e = path[k - 1].assume_init();
self.set_child(e >> 1, e & 1, L::new(r));
}
return (p, k, ancestors);
}
(0, HEIGHT, depth)
}
}
}
#[inline(always)]
fn shift_down(path: &mut Path, from: usize, to: usize) {
match to - from {
0 => {}
1 => path[from - 1] = path[from],
2 => {
path[from - 1] = path[from];
path[from] = path[from + 1];
}
_ => path.copy_within(from..to, from - 1),
}
}
#[inline(always)]
fn neighbours(path: &Path, depth: usize) -> (usize, usize) {
let (mut pred, mut succ) = (NONE, NONE);
for k in (0..depth).rev() {
let e = unsafe { path[k].assume_init() };
if e & 1 == 1 {
if pred == NONE {
pred = e >> 1;
if succ != NONE {
break;
}
}
} else if succ == NONE {
succ = e >> 1;
if pred != NONE {
break;
}
}
}
(pred, succ)
}
struct Guard<T, L, V> {
nodes: *mut Node<L, V>,
arr: *mut T,
len: usize,
}
impl<T, L, V> Drop for Guard<T, L, V> {
fn drop(&mut self) {
if mem::size_of::<V>() == 0 {
return;
}
for i in 0..self.len {
unsafe {
ptr::copy_nonoverlapping(
(&raw const (*self.nodes.add(i)).val).cast::<T>(),
self.arr.add(i),
1,
)
}
}
}
}
const STACK_BYTES: usize = 4096;
#[repr(C)]
struct StackBuffer<A> {
_align: [A; 0],
bytes: [MaybeUninit<u8>; STACK_BYTES],
}
fn regions<T, N>(len: usize, embed: bool) -> Option<(Layout, usize, usize)> {
let mut layout = Layout::array::<N>(len).ok()?;
let mut gather = 0;
if !embed {
(layout, gather) = layout.extend(Layout::array::<T>(len).ok()?).ok()?;
}
let (layout, balance) = layout.extend(Layout::array::<i8>(len).ok()?).ok()?;
Some((layout, gather, balance))
}
fn sort<T: Ord, L: Link, V, const LANES: usize>(arr: &mut [T]) {
let len = arr.len();
let embed = Tree::<T, L, V>::EMBED;
let (layout, gather, balance) =
regions::<T, Node<L, V>>(len, embed).expect("capacity overflow");
let mut stack = StackBuffer::<(Node<L, V>, T)> {
_align: [],
bytes: [MaybeUninit::uninit(); STACK_BYTES],
};
let mut heap: Vec<MaybeUninit<(Node<L, V>, T)>> = Vec::new();
let region: *mut u8 = if layout.size() <= STACK_BYTES {
stack.bytes.as_mut_ptr().cast()
} else {
heap.reserve_exact(layout.size().div_ceil(mem::size_of::<(Node<L, V>, T)>()));
heap.as_mut_ptr().cast()
};
let base = arr.as_mut_ptr();
let mut t = unsafe {
Tree::<T, L, V> {
nodes: region.cast(),
bal: region.add(balance).cast(),
keys: base,
root: 0,
}
};
let out: *mut T = unsafe {
ptr::write_bytes(t.bal, 0, len);
if embed {
base
} else {
region.add(gather).cast()
}
};
let mut guard = Guard::<T, L, V> {
nodes: t.nodes,
arr: base,
len: 0,
};
unsafe {
for i in 0..len {
let n = t.nodes.add(i);
ptr::write(&raw mut (*n).link, [L::NIL; 2]);
if embed {
ptr::copy_nonoverlapping(base.add(i), (&raw mut (*n).val).cast::<T>(), 1);
}
}
guard.len = len;
let mut paths: [Path; LANES] = [[MaybeUninit::uninit(); HEIGHT]; LANES];
let mut last = 0;
let mut f = 0;
let mut ancestors = 0;
let (mut pred, mut succ) = (NONE, NONE);
let mut finger = true;
let mut q = 1;
while q < len {
let x = &*t.key(q);
if finger {
let path = &mut paths[f];
if !(*x < *t.key(last)) {
if succ == NONE || *x < *t.key(succ) {
let depth = t.beside(last, 1, path, ancestors);
ancestors = t.attach(q, path, depth, true).2;
pred = last;
last = q;
q += 1;
continue;
}
} else if pred == NONE || !(*x < *t.key(pred)) {
let depth = t.beside(last, 0, path, ancestors);
ancestors = t.attach(q, path, depth, true).2;
succ = last;
last = q;
q += 1;
continue;
}
finger = false;
}
if q + LANES > len {
let path = &mut paths[0];
let depth = t.search(x, path);
if path[depth - 1].assume_init() >> 1 == last {
finger = true;
(pred, succ) = neighbours(path, depth);
}
ancestors = t.attach(q, path, depth, finger).2;
f = 0;
last = q;
q += 1;
continue;
}
let xs: [&T; LANES] = core::array::from_fn(|j| &*t.key(q + j));
let depths = t.search_lanes(xs, &mut paths);
let mut slots = [0usize; LANES];
let mut rotated = [(0usize, HEIGHT); LANES];
for j in 0..LANES {
let path = &mut paths[j];
let mut depth = depths[j];
let slot = path[depth - 1].assume_init();
let mut from = depth;
let mut taken = false;
for i in 0..j {
let (z, k) = rotated[i];
taken |= slots[i] == slot;
if k < from && path[k].assume_init() >> 1 == z {
from = k;
}
}
if from < depth || taken {
depth = t.resume(xs[j], path, from);
}
let slot = path[depth - 1].assume_init();
slots[j] = slot;
let keep = j == LANES - 1 && slot >> 1 == q + j - 1;
if keep {
finger = true;
(pred, succ) = neighbours(path, depth);
}
let (z, k, a) = t.attach(q + j, path, depth, keep);
rotated[j] = (z, k);
ancestors = a;
}
f = LANES - 1;
last = q + LANES - 1;
q += LANES;
}
mem::forget(guard);
let pending = &mut paths[0];
let mut sp = 0;
let mut cur = L::new(t.root);
let mut k = 0;
loop {
while cur != L::NIL {
pending[sp].write(cur.get());
sp += 1;
cur = t.child(cur.get(), 0);
}
if sp == 0 {
break;
}
sp -= 1;
let n = pending[sp].assume_init();
ptr::copy_nonoverlapping(t.key(n), out.add(k), 1);
k += 1;
cur = t.child(n, 1);
}
if !embed {
ptr::copy_nonoverlapping(out, base, len);
}
}
}