use core::marker::PhantomData;
use core::ops::{BitAnd, BitOr, BitXor, Not, Shl, Shr};
use core::ptr::NonNull;
pub trait RadixKey:
Copy + Ord + Default + core::fmt::Debug + BitAnd<Output = Self> + BitOr<Output = Self> + BitXor<Output = Self> + Not<Output = Self> + Shr<u32, Output = Self> + Shl<u32, Output = Self>
{
const BITS: u32;
const MAX: Self;
const ZERO: Self;
const ONE: Self;
fn to_u64(self) -> u64;
fn from_u64(v: u64) -> Self;
fn as_u32(self) -> u32;
fn wsub(self, o: Self) -> Self;
fn wadd(self, o: Self) -> Self;
fn count_ones(self) -> u32;
fn trailing_zeros(self) -> u32;
fn leading_zeros(self) -> u32;
}
macro_rules! radix_key_impl {
($($t:ty),*) => {$(
impl RadixKey for $t {
const BITS: u32 = <$t>::BITS;
const MAX: Self = <$t>::MAX;
const ZERO: Self = 0;
const ONE: Self = 1;
#[inline(always)] fn to_u64(self) -> u64 { self as u64 }
#[inline(always)] fn from_u64(v: u64) -> Self { v as $t }
#[inline(always)] fn as_u32(self) -> u32 { self as u32 }
#[inline(always)] fn wsub(self, o: Self) -> Self { self.wrapping_sub(o) }
#[inline(always)] fn wadd(self, o: Self) -> Self { self.wrapping_add(o) }
#[inline(always)] fn count_ones(self) -> u32 { <$t>::count_ones(self) }
#[inline(always)] fn trailing_zeros(self) -> u32 { <$t>::trailing_zeros(self) }
#[inline(always)] fn leading_zeros(self) -> u32 { <$t>::leading_zeros(self) }
}
)*};
}
radix_key_impl!(u32, u64);
#[inline(always)]
pub fn highest_bit<K: RadixKey>(k: K) -> u32 {
K::BITS - 1 - k.leading_zeros()
}
#[inline(always)]
pub fn pext_soft<K: RadixKey>(k: K, mask: K) -> K {
let mut out = K::ZERO;
let mut bit = K::ONE;
let mut m = mask;
while m != K::ZERO {
let low = m & (!m).wadd(K::ONE); if k & low != K::ZERO {
out = out | bit;
}
m = m & m.wsub(K::ONE);
bit = bit << 1;
}
out
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum SimdKind {
None,
I32,
I64,
F64,
K64,
K32,
}
pub trait Elem: Copy + Default {
type Key: RadixKey;
const CHUNKED: bool;
const SIMD: SimdKind;
fn less(a: Self, b: Self) -> bool;
fn compare(a: Self, b: Self) -> i32;
fn compare_from(a: Self, b: Self, chunk: i32) -> i32;
fn radix_key(a: Self, chunk: i32) -> Self::Key;
fn chunk_ends(a: Self, chunk: i32) -> bool;
#[inline(always)]
fn report_key_chunk(_a: &Self, _chunk: i32, _sink: &mut impl FnMut(*const u8, usize)) {}
#[inline(always)]
fn report_key_compare(_a: &Self, _b: &Self, _chunk: i32, _sink: &mut impl FnMut(*const u8, usize)) {}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Stat {
Giveups,
Splits,
SplitRetries,
PartSorts,
DictTries,
DictHits,
PlanRetries,
RadixPasses,
PassesSkipped,
}
pub trait Hooks {
const COUNTED: bool;
#[inline(always)]
fn on_read(_p: *const u8, _bytes: usize) {}
#[inline(always)]
fn on_write(_p: *const u8, _bytes: usize) {}
#[inline(always)]
fn on_table_read(_p: *const u8, _bytes: usize) {}
#[inline(always)]
fn on_table_write(_p: *const u8, _bytes: usize) {}
#[inline(always)]
fn on_table_rw(p: *const u8, bytes: usize) {
Self::on_table_read(p, bytes);
Self::on_table_write(p, bytes);
}
#[inline(always)]
fn on_table_sweep(_p: *const u8, _entries: usize, _entry_bytes: usize, _read: bool, _write: bool) {}
#[inline(always)]
fn on_key(_p: *const u8, _bytes: usize) {}
#[inline(always)]
fn on_compare(_less: bool) {}
#[inline(always)]
fn stat(_s: Stat, _n: u64) {}
#[inline(always)]
fn note_route(_r: u8) {}
#[inline(always)]
fn enter_depth() {}
#[inline(always)]
fn leave_depth() {}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct NoHooks;
impl Hooks for NoHooks {
const COUNTED: bool = false;
}
pub struct DepthScope<H: Hooks>(PhantomData<H>);
impl<H: Hooks> DepthScope<H> {
#[inline(always)]
pub fn new() -> Self {
if H::COUNTED {
H::enter_depth();
}
DepthScope(PhantomData)
}
}
impl<H: Hooks> Default for DepthScope<H> {
fn default() -> Self {
Self::new()
}
}
impl<H: Hooks> Drop for DepthScope<H> {
#[inline(always)]
fn drop(&mut self) {
if H::COUNTED {
H::leave_depth();
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct AllocError;
pub trait Alloc {
fn allocate(bytes: usize) -> Result<NonNull<u8>, AllocError>;
unsafe fn deallocate(p: NonNull<u8>, bytes: usize);
}
#[inline]
pub fn alloc_array<A: Alloc, U>(n: usize) -> Result<NonNull<U>, AllocError> {
if core::mem::align_of::<U>() > 16 {
return Err(AllocError); }
let bytes = n.checked_mul(core::mem::size_of::<U>()).ok_or(AllocError)?;
A::allocate(bytes).map(NonNull::cast)
}
#[inline]
pub unsafe fn free_array<A: Alloc, U>(p: NonNull<U>, n: usize) {
unsafe { A::deallocate(p.cast(), n * core::mem::size_of::<U>()) }
}
pub trait Arr: Copy {
type T: Elem;
type H: Hooks;
type A: Alloc;
fn with_range(self, p: *mut Self::T, n: usize) -> Self;
fn size(self) -> usize;
fn data(self) -> *mut Self::T;
fn get(self, i: usize) -> Self::T;
fn set(self, i: usize, v: Self::T);
fn less(self, a: Self::T, b: Self::T) -> bool;
fn compare(self, a: Self::T, b: Self::T) -> i32;
fn compare_from(self, a: Self::T, b: Self::T, chunk: i32) -> i32;
fn key(x: Self::T, chunk: i32) -> <Self::T as Elem>::Key;
#[inline(always)]
fn swap(self, i: usize, j: usize) {
let t = self.get(i);
self.set(i, self.get(j));
self.set(j, t);
}
#[inline(always)]
fn sub(self, off: usize, len: usize) -> Self {
debug_assert!(off + len <= self.size());
self.with_range(unsafe { self.data().add(off) }, len)
}
#[inline(always)]
fn alloc_array<U>(n: usize) -> Result<NonNull<U>, AllocError> {
alloc_array::<Self::A, U>(n)
}
#[inline(always)]
unsafe fn free_array<U>(p: NonNull<U>, n: usize) {
unsafe { free_array::<Self::A, U>(p, n) }
}
}
pub struct View<T, H, A> {
p: *mut T,
n: usize,
_m: PhantomData<fn() -> (H, A)>,
}
impl<T, H, A> Clone for View<T, H, A> {
fn clone(&self) -> Self {
*self
}
}
impl<T, H, A> Copy for View<T, H, A> {}
impl<T, H, A> View<T, H, A> {
#[inline(always)]
pub fn new(p: *mut T, n: usize) -> Self {
View { p, n, _m: PhantomData }
}
}
impl<T: Elem, H: Hooks, A: Alloc> Arr for View<T, H, A> {
type T = T;
type H = H;
type A = A;
#[inline(always)]
fn with_range(self, p: *mut T, n: usize) -> Self {
View::new(p, n)
}
#[inline(always)]
fn size(self) -> usize {
self.n
}
#[inline(always)]
fn data(self) -> *mut T {
self.p
}
#[inline(always)]
fn get(self, i: usize) -> T {
debug_assert!(i < self.n);
if H::COUNTED {
H::on_read(self.p.wrapping_add(i) as *const u8, core::mem::size_of::<T>());
}
unsafe { *self.p.add(i) }
}
#[inline(always)]
fn set(self, i: usize, v: T) {
debug_assert!(i < self.n);
if H::COUNTED {
H::on_write(self.p.wrapping_add(i) as *const u8, core::mem::size_of::<T>());
}
unsafe { *self.p.add(i) = v }
}
#[inline(always)]
fn less(self, a: T, b: T) -> bool {
if H::COUNTED {
let r = T::less(a, b);
T::report_key_compare(&a, &b, 0, &mut |p, n| H::on_key(p, n));
H::on_compare(r);
r
} else {
T::less(a, b)
}
}
#[inline(always)]
fn compare(self, a: T, b: T) -> i32 {
self.compare_from(a, b, 0)
}
#[inline(always)]
fn compare_from(self, a: T, b: T, chunk: i32) -> i32 {
if H::COUNTED {
let c = T::compare_from(a, b, chunk);
T::report_key_compare(&a, &b, chunk, &mut |p, n| H::on_key(p, n));
H::on_compare(c < 0);
c
} else {
T::compare_from(a, b, chunk)
}
}
#[inline(always)]
fn key(x: T, chunk: i32) -> T::Key {
if H::COUNTED {
T::report_key_chunk(&x, chunk, &mut |p, n| H::on_key(p, n));
}
T::radix_key(x, chunk)
}
}
pub struct AuxBuffer<V: Arr> {
p: NonNull<V::T>,
n: usize,
like: V,
}
impl<V: Arr> AuxBuffer<V> {
pub fn new(like: V, n: usize) -> Result<Self, AllocError> {
let p = V::alloc_array::<V::T>(n)?;
Ok(AuxBuffer { p, n, like })
}
#[inline(always)]
pub fn arr(&self) -> V {
self.like.with_range(self.p.as_ptr(), self.n)
}
#[inline(always)]
pub fn size(&self) -> usize {
self.n
}
#[inline(always)]
pub fn data(&self) -> *mut V::T {
self.p.as_ptr()
}
pub fn resize_discard(&mut self, n: usize) -> Result<(), AllocError> {
unsafe { V::free_array(self.p, self.n) };
self.n = 0;
self.p = NonNull::dangling();
self.p = V::alloc_array::<V::T>(n)?;
self.n = n;
Ok(())
}
}
impl<V: Arr> Drop for AuxBuffer<V> {
fn drop(&mut self) {
if self.n > 0 || self.p != NonNull::dangling() {
unsafe { V::free_array(self.p, self.n) };
}
}
}
pub struct AuxRaw<U, V: Arr> {
p: NonNull<U>,
n: usize,
_v: PhantomData<V>,
}
impl<U, V: Arr> AuxRaw<U, V> {
pub fn new(n: usize) -> Result<Self, AllocError> {
Ok(AuxRaw { p: V::alloc_array::<U>(n)?, n, _v: PhantomData })
}
#[inline(always)]
pub fn data(&self) -> *mut U {
self.p.as_ptr()
}
#[inline(always)]
pub fn size(&self) -> usize {
self.n
}
}
impl<U, V: Arr> Drop for AuxRaw<U, V> {
fn drop(&mut self) {
unsafe { V::free_array(self.p, self.n) };
}
}
pub struct IdxView<H: Hooks> {
p: *mut u32,
n: usize,
_h: PhantomData<H>,
}
impl<H: Hooks> Clone for IdxView<H> {
fn clone(&self) -> Self {
*self
}
}
impl<H: Hooks> Copy for IdxView<H> {}
impl<H: Hooks> IdxView<H> {
#[inline(always)]
pub fn new(p: *mut u32, n: usize) -> Self {
IdxView { p, n, _h: PhantomData }
}
#[inline(always)]
pub fn get(self, i: usize) -> u32 {
debug_assert!(i < self.n);
if H::COUNTED {
H::on_read(self.p.wrapping_add(i) as *const u8, 4);
}
unsafe { *self.p.add(i) }
}
#[inline(always)]
pub fn set(self, i: usize, v: u32) {
debug_assert!(i < self.n);
if H::COUNTED {
H::on_write(self.p.wrapping_add(i) as *const u8, 4);
}
unsafe { *self.p.add(i) = v }
}
}
#[inline]
pub fn insertion_sort<V: Arr>(a: V, lo: usize, hi: usize) {
for i in lo + 1..hi {
let v = a.get(i);
let mut j = i;
while j > lo {
let p = a.get(j - 1);
if !a.less(v, p) {
break;
}
a.set(j, p);
j -= 1;
}
a.set(j, v);
}
}
#[inline]
pub fn copy_forward<V: Arr>(src: V, s: usize, dst: V, d: usize, len: usize) {
for i in 0..len {
dst.set(d + i, src.get(s + i));
}
}
#[inline]
pub fn reverse_range<V: Arr>(a: V, mut lo: usize, mut hi: usize) {
while lo + 1 < hi {
hi -= 1;
a.swap(lo, hi);
lo += 1;
}
}
#[inline]
pub fn floor_log2(mut n: usize) -> usize {
let mut r = 0;
while n > 1 {
n >>= 1;
r += 1;
}
r
}
pub const STR_CHUNK_BYTES: u32 = 7;
pub const STR_KEY_MASK: u64 = !0xFFu64;
#[inline(always)]
pub unsafe fn str_chunk_key(ptr: *const u8, len: u32, chunk: i32) -> u64 {
let off = (chunk as u32).wrapping_mul(STR_CHUNK_BYTES);
if off.wrapping_add(8) <= len && off < len {
let w = unsafe { core::ptr::read_unaligned(ptr.add(off as usize) as *const u64) };
return (u64::from_be(w) & STR_KEY_MASK) | STR_CHUNK_BYTES as u64;
}
if off >= len {
return 0;
}
let valid = len - off; let q = unsafe { ptr.add(off as usize) };
if valid >= 4 {
let hi = unsafe { core::ptr::read_unaligned(q as *const u32) };
let lo = unsafe { core::ptr::read_unaligned(q.add(valid as usize - 4) as *const u32) };
let w = ((u32::from_be(hi) as u64) << 32) | ((u32::from_be(lo) as u64) << (8 * (8 - valid)));
return (w & STR_KEY_MASK) | valid as u64;
}
let mut w = 0u64;
for i in 0..valid as usize {
w = (w << 8) | unsafe { *q.add(i) } as u64;
}
w <<= 8 * (STR_CHUNK_BYTES - valid);
(w << 8) | valid as u64
}
#[inline(always)]
pub unsafe fn str_chunk_key_desc(ptr: *const u8, len: u32, chunk: i32) -> u64 {
let k = unsafe { str_chunk_key(ptr, len, chunk) };
(!k & STR_KEY_MASK) | (STR_CHUNK_BYTES as u64 - (k & 0xFF))
}
#[inline(always)]
pub fn str_chunk_ends(len: u32, chunk: i32) -> bool {
(chunk as u32).wrapping_mul(STR_CHUNK_BYTES).wrapping_add(STR_CHUNK_BYTES) > len
}
#[inline(always)]
pub unsafe fn str_compare_from(pa: *const u8, la0: u32, pb: *const u8, lb0: u32, off: u32) -> i32 {
let la = la0.saturating_sub(off);
let lb = lb0.saturating_sub(off);
let m = la.min(lb) as usize;
let (sa, sb) = unsafe {
(
core::slice::from_raw_parts(if m == 0 { NonNull::dangling().as_ptr() } else { pa.add(off as usize) }, m),
core::slice::from_raw_parts(if m == 0 { NonNull::dangling().as_ptr() } else { pb.add(off as usize) }, m),
)
};
match sa.cmp(sb) {
core::cmp::Ordering::Less => -1,
core::cmp::Ordering::Greater => 1,
core::cmp::Ordering::Equal => (la > lb) as i32 - (la < lb) as i32,
}
}
pub unsafe fn str_examined(pa: *const u8, la0: u32, pb: *const u8, lb0: u32, off: u32) -> u32 {
let la = la0.saturating_sub(off);
let lb = lb0.saturating_sub(off);
let m = la.min(lb);
let mut i = 0u32;
while i < m && unsafe { *pa.add((off + i) as usize) == *pb.add((off + i) as usize) } {
i += 1;
}
if i < m { i + 1 } else { m }
}
pub fn str_chunk_bytes(len: u32, chunk: i32) -> u32 {
let off = (chunk as u32).wrapping_mul(STR_CHUNK_BYTES);
if off.wrapping_add(8) <= len && off < len {
return 8;
}
if off < len { len - off } else { 0 }
}