extern crate alloc;
use alloc::vec::Vec;
use crate::allocator::power;
use crate::allocator::StabilityRefusal;
use crate::cascade::{self, CascadeRefusal, MAX_LENS_MAGNITUDE};
use crate::fixed::{NonNegativeFixed, SignedFixed};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EscortRefusal {
EmptyInput,
UnsupportedLens { lens: SignedFixed },
NumericFault { index: usize, error_code: u32 },
ExactPathRefused {
index: usize,
reason: CascadeRefusal,
},
DegenerateNormalization,
}
impl core::fmt::Display for EscortRefusal {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
core::fmt::Debug::fmt(self, f)
}
}
#[cfg(feature = "std")]
impl std::error::Error for EscortRefusal {}
#[inline]
fn exact_integer_lens(q: SignedFixed) -> Option<i32> {
if q.to_bits() & 0xFFFF == 0 {
Some(q.to_num())
} else {
None
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PathConfidence {
Exact,
Approximate { max_relative_error_bps: u32 },
}
#[inline]
fn approximate_error_bound_bps(abs_q_bits: u32) -> u32 {
const Q_0_25: u32 = 1 << 14; const Q_1: u32 = 1 << 16;
const Q_2: u32 = 2 << 16;
const Q_4: u32 = 4 << 16;
const Q_8: u32 = 8 << 16;
if abs_q_bits <= Q_0_25 {
50
} else if abs_q_bits <= Q_1 {
150
} else if abs_q_bits <= Q_2 {
350
} else if abs_q_bits <= Q_4 {
760
} else if abs_q_bits <= Q_8 {
1640
} else {
3690
}
}
#[inline]
fn exact_floor_share(w_bits: u32, sum_bits: u32) -> Result<NonNegativeFixed, u32> {
let q = ((w_bits as u64) << 16) / u64::from(sum_bits);
if q > u32::MAX as u64 {
return Err(StabilityRefusal::NumericRangeExceeded as u32);
}
Ok(NonNegativeFixed::from_bits(q as u32))
}
#[allow(clippy::missing_errors_doc)]
pub fn escort_distribution(
masses: &[NonNegativeFixed],
q: SignedFixed,
) -> Result<Vec<NonNegativeFixed>, EscortRefusal> {
escort_distribution_with_confidence(masses, q).map(|(values, _confidence)| values)
}
#[allow(clippy::missing_errors_doc)]
pub fn escort_distribution_with_confidence(
masses: &[NonNegativeFixed],
q: SignedFixed,
) -> Result<(Vec<NonNegativeFixed>, PathConfidence), EscortRefusal> {
if masses.is_empty() {
return Err(EscortRefusal::EmptyInput);
}
let abs_q_bits = q.to_bits().unsigned_abs();
if abs_q_bits > MAX_LENS_MAGNITUDE << 16 {
return Err(EscortRefusal::UnsupportedLens { lens: q });
}
let exact_lens = exact_integer_lens(q);
let confidence = match exact_lens {
Some(_) => PathConfidence::Exact,
None => PathConfidence::Approximate {
max_relative_error_bps: approximate_error_bound_bps(abs_q_bits),
},
};
let mut weighted: Vec<NonNegativeFixed> = Vec::with_capacity(masses.len());
for (index, &mass) in masses.iter().enumerate() {
let w = if let Some(lens) = exact_lens {
cascade::escort_weight(mass, lens, index)
.map_err(|reason| EscortRefusal::ExactPathRefused { index, reason })?
} else {
let w = power(mass, q);
if w.err != u32::MAX {
return Err(EscortRefusal::NumericFault {
index,
error_code: w.err,
});
}
w
};
weighted.push(w);
}
let mut sum = NonNegativeFixed::ZERO;
for &w in &weighted {
sum += w;
}
if sum.err != u32::MAX || sum.to_bits() == 0 {
return Err(EscortRefusal::DegenerateNormalization);
}
let sum_bits = sum.to_bits();
for (index, w) in weighted.iter_mut().enumerate() {
match exact_floor_share(w.to_bits(), sum_bits) {
Ok(share) => *w = share,
Err(error_code) => {
return Err(EscortRefusal::NumericFault { index, error_code });
}
}
}
Ok((weighted, confidence))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::allocator::StabilityRefusal;
fn approx_eq(a: NonNegativeFixed, b: NonNegativeFixed, tol_bits: i64) -> bool {
(a.to_bits() as i64 - b.to_bits() as i64).abs() < tol_bits
}
fn mass(x: f32) -> NonNegativeFixed {
NonNegativeFixed::from_bits((x * 65536.0).round() as u32)
}
fn q(x: f32) -> SignedFixed {
SignedFixed::from_bits((x * 65536.0).round() as i32)
}
#[test]
fn q_one_reproduces_a_normalized_input_distribution() {
let p = [mass(0.2), mass(0.3), mass(0.5)];
let escort = escort_distribution(&p, SignedFixed::ONE).unwrap();
for (input, output) in p.iter().zip(escort.iter()) {
assert!(approx_eq(*input, *output, 200), "{input:?} vs {output:?}");
}
}
#[test]
fn q_zero_yields_uniform_distribution() {
let p = [mass(0.1), mass(0.4), mass(0.5), mass(10.0)];
let escort = escort_distribution(&p, SignedFixed::ZERO).unwrap();
let expected = NonNegativeFixed::ONE / NonNegativeFixed::from_num(4);
for value in &escort {
assert!(approx_eq(*value, expected, 50), "{value:?}");
}
}
#[test]
fn output_sums_to_approximately_one() {
let p = [mass(1.0), mass(2.0), mass(3.0), mass(4.0), mass(5.0)];
for exponent in [-2.0f32, -0.5, 0.0, 0.5, 1.0, 2.0, 5.0] {
let escort = escort_distribution(&p, q(exponent)).unwrap();
let mut sum = NonNegativeFixed::ZERO;
for v in &escort {
sum += *v;
}
assert!(
approx_eq(sum, NonNegativeFixed::ONE, 500),
"q={exponent} sum={sum:?}"
);
}
}
#[test]
fn higher_q_concentrates_mass_on_the_largest_input() {
let p = [mass(1.0), mass(2.0), mass(10.0)];
let escort_low = escort_distribution(&p, q(1.0)).unwrap();
let escort_high = escort_distribution(&p, q(4.0)).unwrap();
assert!(escort_high[2].to_bits() > escort_low[2].to_bits());
}
#[test]
fn rejects_empty_input() {
assert!(matches!(
escort_distribution(&[], SignedFixed::ONE),
Err(EscortRefusal::EmptyInput)
));
}
#[test]
fn degenerate_normalization_is_refused_not_silently_zero() {
let p = [NonNegativeFixed::ZERO, NonNegativeFixed::ZERO];
assert!(matches!(
escort_distribution(&p, SignedFixed::ONE),
Err(EscortRefusal::DegenerateNormalization)
));
}
#[test]
fn power_disagrees_with_the_exact_path_at_a_measured_bound() {
let p = [mass(1.0), mass(2.0), mass(3.0), mass(4.0)];
let lens: i32 = 3;
let power_weights: alloc::vec::Vec<NonNegativeFixed> =
p.iter().map(|&m| power(m, q(lens as f32))).collect();
let mut power_sum = NonNegativeFixed::ZERO;
for &w in &power_weights {
power_sum += w;
}
let via_power: alloc::vec::Vec<NonNegativeFixed> =
power_weights.into_iter().map(|w| w / power_sum).collect();
let exact_weights: alloc::vec::Vec<NonNegativeFixed> = p
.iter()
.enumerate()
.map(|(node, &m)| crate::cascade::escort_weight(m, lens, node).unwrap())
.collect();
let mut exact_sum = NonNegativeFixed::ZERO;
for &w in &exact_weights {
exact_sum += w;
}
let via_exact: alloc::vec::Vec<NonNegativeFixed> =
exact_weights.into_iter().map(|w| w / exact_sum).collect();
for (index, (approx, exact)) in via_power.iter().zip(via_exact.iter()).enumerate() {
assert!(
approx_eq(*approx, *exact, 900),
"index={index} approx={approx:?} exact={exact:?}"
);
}
}
#[test]
fn exact_lens_never_reaches_the_approximate_path() {
let p = [mass(1.0), mass(2.0), mass(3.0), mass(4.0)];
let lens: i32 = 3;
let via_dispatch = escort_distribution(&p, q(lens as f32)).unwrap();
let exact_weights: alloc::vec::Vec<NonNegativeFixed> = p
.iter()
.enumerate()
.map(|(node, &m)| crate::cascade::escort_weight(m, lens, node).unwrap())
.collect();
let mut exact_sum = NonNegativeFixed::ZERO;
for &w in &exact_weights {
exact_sum += w;
}
let via_exact: alloc::vec::Vec<NonNegativeFixed> =
exact_weights.into_iter().map(|w| w / exact_sum).collect();
assert_eq!(
via_dispatch, via_exact,
"an integer lens must reach the exact path bit-for-bit, not just within tolerance"
);
}
#[test]
fn lens_beyond_the_declared_magnitude_is_refused() {
let p = [mass(1.0), mass(2.0)];
let too_large = SignedFixed::from_num((MAX_LENS_MAGNITUDE + 1) as i32);
assert_eq!(
escort_distribution(&p, too_large),
Err(EscortRefusal::UnsupportedLens { lens: too_large })
);
}
#[test]
fn lens_at_the_declared_magnitude_boundary_is_admitted() {
let p = [mass(0.9), mass(1.0)];
let at_bound = SignedFixed::from_num(MAX_LENS_MAGNITUDE as i32);
assert!(escort_distribution(&p, at_bound).is_ok());
}
#[test]
fn zero_mass_under_fractional_negative_q_is_refused() {
let p = [mass(0.0), mass(1.0), mass(2.0)];
let negative_fractional_q = q(-2.5);
assert!(
exact_integer_lens(negative_fractional_q).is_none(),
"q(-2.5) must be genuinely fractional so this test exercises the `power` path, \
not cascade::escort_weight's already-refusing exact-integer path"
);
match escort_distribution(&p, negative_fractional_q) {
Err(EscortRefusal::NumericFault { index, error_code }) => {
assert_eq!(index, 0, "the zero mass is at index 0");
assert_eq!(
error_code,
StabilityRefusal::UnsupportedDomain as u32,
"expected power's UnsupportedDomain refusal for 0^(negative)"
);
}
other => panic!(
"expected EscortRefusal::NumericFault for the zero-mass element under a \
negative q, got {other:?}"
),
}
}
#[test]
fn caller_can_distinguish_high_error_fractional_q_from_low_error_fractional_q() {
let p = [mass(0.8), mass(0.9), mass(1.0), mass(1.1)];
let (_low_values, low_confidence) =
escort_distribution_with_confidence(&p, q(0.1)).unwrap();
let (_high_values, high_confidence) =
escort_distribution_with_confidence(&p, q(15.9)).unwrap();
let low_bound = match low_confidence {
PathConfidence::Approximate {
max_relative_error_bps,
} => max_relative_error_bps,
PathConfidence::Exact => panic!("q=0.1 is genuinely fractional, must not be Exact"),
};
let high_bound = match high_confidence {
PathConfidence::Approximate {
max_relative_error_bps,
} => max_relative_error_bps,
PathConfidence::Exact => panic!("q=15.9 is genuinely fractional, must not be Exact"),
};
assert!(
high_bound > low_bound,
"high-|q| call must report a larger error bound than low-|q|: \
low={low_bound}bps high={high_bound}bps"
);
assert_eq!(low_bound, 50, "q=0.1 falls in the |q| <= 0.25 bucket");
assert_eq!(
high_bound, 3690,
"q=15.9 falls in the final |q| <= 16 bucket"
);
}
#[test]
fn exact_integer_lens_reports_exact_confidence() {
let p = [mass(1.0), mass(2.0), mass(3.0)];
let (_values, confidence) = escort_distribution_with_confidence(&p, q(3.0)).unwrap();
assert_eq!(confidence, PathConfidence::Exact);
}
#[test]
fn saturating_div_is_exact_floor_division_guard() {
let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
let mut next = || {
state ^= state >> 12;
state ^= state << 25;
state ^= state >> 27;
state.wrapping_mul(0x2545_F491_4F6C_DD1D)
};
let edges = [
0u32,
1,
2,
0xFFFF,
0x10000,
0x10001,
0x7FFF_FFFF,
0x8000_0000,
u32::MAX,
];
let mut checked: u64 = 0;
for &n in &edges {
for &d in &edges {
if d == 0 {
continue;
}
let a = NonNegativeFixed::from_bits(n);
let b = NonNegativeFixed::from_bits(d);
let q = a.saturating_div(b);
let exact = ((n as u64) << 16) / u64::from(d);
if exact <= u32::MAX as u64 {
assert_eq!(
q.to_bits(),
exact as u32,
"saturating_div({n:#x}, {d:#x}) is not exact floor division"
);
assert_eq!(q.err, u32::MAX, "clean pair must stay err-clean");
} else {
assert_eq!(
q.to_bits(),
u32::MAX,
"saturating_div({n:#x}, {d:#x}) must saturate like the fast path"
);
assert_ne!(q.err, u32::MAX);
assert_eq!(
q.err,
crate::allocator::StabilityRefusal::NumericRangeExceeded as u32
);
}
checked += 1;
}
}
for _ in 0..1_000_000u64 {
let r = next();
let n = (r >> 32) as u32;
let d = (r as u32).max(1);
let a = NonNegativeFixed::from_bits(n);
let b = NonNegativeFixed::from_bits(d);
let q = a.saturating_div(b);
let exact = ((n as u64) << 16) / u64::from(d);
if exact <= u32::MAX as u64 {
assert_eq!(
q.to_bits(),
exact as u32,
"saturating_div({n:#x}, {d:#x}) is not exact floor division"
);
assert_eq!(q.err, u32::MAX);
} else {
assert_eq!(q.to_bits(), u32::MAX);
assert_eq!(
q.err,
crate::allocator::StabilityRefusal::NumericRangeExceeded as u32
);
}
checked += 1;
}
for d in 1u32..=4096 {
for _ in 0..64 {
let n = (next() >> 32) as u32;
let q =
NonNegativeFixed::from_bits(n).saturating_div(NonNegativeFixed::from_bits(d));
let exact = ((n as u64) << 16) / u64::from(d);
if exact <= u32::MAX as u64 {
assert_eq!(
q.to_bits(),
exact as u32,
"saturating_div({n:#x}, {d:#x}) is not exact floor division"
);
assert_eq!(q.err, u32::MAX);
} else {
assert_eq!(q.to_bits(), u32::MAX);
assert_eq!(
q.err,
crate::allocator::StabilityRefusal::NumericRangeExceeded as u32
);
}
checked += 1;
}
}
assert!(checked > 1_000_000, "sweep unexpectedly small: {checked}");
}
}