mod simd;
use core::mem::MaybeUninit;
use crate::Buffer;
const SIGBITS: u32 = 5;
const RSHIFT: u32 = 8 - SIGBITS;
const HISTO_LEVELS: usize = 1 << SIGBITS; const HISTO_SIZE: usize = 1 << (3 * SIGBITS); const MAX_ITERATIONS: usize = 1000;
const MAX_BOXES: usize = 256;
pub(crate) struct BoxArena {
data: [MaybeUninit<VBox>; MAX_BOXES],
len: usize,
}
#[allow(unsafe_code)]
impl BoxArena {
pub const fn new() -> Self {
Self {
data: [const { MaybeUninit::uninit() }; MAX_BOXES],
len: 0,
}
}
pub fn len(&self) -> usize {
self.len
}
pub fn push(&mut self, v: VBox) -> bool {
if self.len >= MAX_BOXES {
return false;
}
self.data[self.len].write(v);
self.len += 1;
true
}
pub fn pop(&mut self) -> Option<VBox> {
if self.len == 0 {
return None;
}
self.len -= 1;
Some(unsafe { self.data[self.len].assume_init_read() })
}
pub fn as_mut_slice(&mut self) -> &mut [VBox] {
unsafe { core::slice::from_raw_parts_mut(self.data.as_mut_ptr() as *mut VBox, self.len) }
}
pub fn clear(&mut self) {
while self.pop().is_some() {}
}
pub fn drain_into(&mut self, dst: &mut Self) {
while let Some(v) = self.pop() {
let _ = dst.push(v);
}
}
}
impl Drop for BoxArena {
fn drop(&mut self) {
self.clear();
}
}
#[inline]
fn histo_index(r: u32, g: u32, b: u32) -> usize {
((r << (2 * SIGBITS)) + (g << SIGBITS) + b) as usize
}
#[derive(Clone)]
pub(crate) struct VBox {
r1: u32,
r2: u32,
g1: u32,
g2: u32,
b1: u32,
b2: u32,
count_cache: Option<u32>,
avg_cache: Option<[u8; 3]>,
}
impl VBox {
fn volume(&self) -> u32 {
(self.r2 - self.r1 + 1) * (self.g2 - self.g1 + 1) * (self.b2 - self.b1 + 1)
}
fn count(&mut self, histo: &[u32]) -> u32 {
if let Some(c) = self.count_cache {
return c;
}
let mut npix: u32 = 0;
for r in self.r1..=self.r2 {
for g in self.g1..=self.g2 {
let lo = histo_index(r, g, self.b1);
let hi = histo_index(r, g, self.b2);
let row_sum = simd::sum_u32_slice(&histo[lo..=hi]);
npix = npix.saturating_add(row_sum);
}
}
self.count_cache = Some(npix);
npix
}
fn avg(&mut self, histo: &[u32]) -> [u8; 3] {
if let Some(a) = self.avg_cache {
return a;
}
let mult = 1u32 << RSHIFT;
let mut ntot: u64 = 0;
let mut rsum: f64 = 0.0;
let mut gsum: f64 = 0.0;
let mut bsum: f64 = 0.0;
for r in self.r1..=self.r2 {
for g in self.g1..=self.g2 {
for b in self.b1..=self.b2 {
let pop = histo[histo_index(r, g, b)] as u64;
if pop == 0 {
continue;
}
ntot += pop;
let popf = pop as f64;
rsum += popf * (r as f64 + 0.5) * mult as f64;
gsum += popf * (g as f64 + 0.5) * mult as f64;
bsum += popf * (b as f64 + 0.5) * mult as f64;
}
}
}
let out = if ntot > 0 {
let n = ntot as f64;
[(rsum / n) as u8, (gsum / n) as u8, (bsum / n) as u8]
} else {
let center = |a: u32, b: u32| -> u8 { ((mult * (a + b + 1)) / 2).min(255) as u8 };
[
center(self.r1, self.r2),
center(self.g1, self.g2),
center(self.b1, self.b2),
]
};
self.avg_cache = Some(out);
out
}
}
fn build_histogram<I: Iterator<Item = [u8; 3]>>(pixels: I, histo: &mut [u32; HISTO_SIZE]) {
histo.fill(0);
for [r, g, b] in pixels {
let rv = (r as u32) >> RSHIFT;
let gv = (g as u32) >> RSHIFT;
let bv = (b as u32) >> RSHIFT;
histo[histo_index(rv, gv, bv)] = histo[histo_index(rv, gv, bv)].saturating_add(1);
}
}
fn initial_vbox(histo: &[u32]) -> Option<VBox> {
let mut rmin = u32::MAX;
let mut rmax = 0;
let mut gmin = u32::MAX;
let mut gmax = 0;
let mut bmin = u32::MAX;
let mut bmax = 0;
let mut any = false;
for r in 0..HISTO_LEVELS as u32 {
for g in 0..HISTO_LEVELS as u32 {
for b in 0..HISTO_LEVELS as u32 {
if histo[histo_index(r, g, b)] > 0 {
any = true;
if r < rmin {
rmin = r;
}
if r > rmax {
rmax = r;
}
if g < gmin {
gmin = g;
}
if g > gmax {
gmax = g;
}
if b < bmin {
bmin = b;
}
if b > bmax {
bmax = b;
}
}
}
}
}
if !any {
return None;
}
Some(VBox {
r1: rmin,
r2: rmax,
g1: gmin,
g2: gmax,
b1: bmin,
b2: bmax,
count_cache: None,
avg_cache: None,
})
}
#[derive(Clone, Copy)]
enum Axis {
R,
G,
B,
}
fn median_cut(vbox: &VBox, histo: &[u32]) -> Option<(VBox, Option<VBox>)> {
let mut probe = vbox.clone();
let count = probe.count(histo);
if count == 0 {
return None;
}
if count == 1 {
return Some((vbox.clone(), None));
}
let rw = vbox.r2 - vbox.r1 + 1;
let gw = vbox.g2 - vbox.g1 + 1;
let bw = vbox.b2 - vbox.b1 + 1;
if rw == 1 && gw == 1 && bw == 1 {
return Some((vbox.clone(), None));
}
let maxw = rw.max(gw).max(bw);
let axis = if maxw == rw {
Axis::R
} else if maxw == gw {
Axis::G
} else {
Axis::B
};
let (lo, hi) = match axis {
Axis::R => (vbox.r1, vbox.r2),
Axis::G => (vbox.g1, vbox.g2),
Axis::B => (vbox.b1, vbox.b2),
};
let mut partialsum = [0u32; HISTO_LEVELS];
let mut total: u32 = 0;
for i in lo..=hi {
let mut sum: u32 = 0;
match axis {
Axis::R => {
for g in vbox.g1..=vbox.g2 {
for b in vbox.b1..=vbox.b2 {
sum = sum.saturating_add(histo[histo_index(i, g, b)]);
}
}
}
Axis::G => {
for r in vbox.r1..=vbox.r2 {
for b in vbox.b1..=vbox.b2 {
sum = sum.saturating_add(histo[histo_index(r, i, b)]);
}
}
}
Axis::B => {
for r in vbox.r1..=vbox.r2 {
for g in vbox.g1..=vbox.g2 {
sum = sum.saturating_add(histo[histo_index(r, g, i)]);
}
}
}
}
total = total.saturating_add(sum);
partialsum[i as usize] = total;
}
let lookaheadsum: [u32; HISTO_LEVELS] =
core::array::from_fn(|i| total.saturating_sub(partialsum[i]));
for i in lo..=hi {
if partialsum[i as usize] <= total / 2 {
continue;
}
let left = i - lo;
let right = hi - i;
let d2_initial: i64 = if left <= right {
let candidate = i as i64 + (right / 2) as i64;
candidate.min(hi as i64 - 1)
} else {
let candidate = i as i64 - 1 - (left / 2) as i64;
candidate.max(lo as i64)
};
let mut d2 = d2_initial.clamp(lo as i64, hi as i64) as u32;
while d2 < hi && partialsum[d2 as usize] == 0 {
d2 += 1;
}
while d2 > lo && lookaheadsum[d2 as usize] == 0 && partialsum[(d2 - 1) as usize] != 0 {
d2 -= 1;
}
let d2 = d2.min(hi.saturating_sub(1)).max(lo);
let mut left_box = vbox.clone();
let mut right_box = vbox.clone();
match axis {
Axis::R => {
left_box.r2 = d2;
right_box.r1 = d2 + 1;
}
Axis::G => {
left_box.g2 = d2;
right_box.g1 = d2 + 1;
}
Axis::B => {
left_box.b2 = d2;
right_box.b1 = d2 + 1;
}
}
left_box.count_cache = None;
left_box.avg_cache = None;
right_box.count_cache = None;
right_box.avg_cache = None;
return Some((left_box, Some(right_box)));
}
Some((vbox.clone(), None))
}
fn iterate_split<F>(boxes: &mut BoxArena, target: usize, histo: &[u32], score: F)
where
F: Fn(&mut VBox, &[u32]) -> u64,
{
let mut iters = 0;
let mut exhausted = BoxArena::new();
while boxes.len() + exhausted.len() < target && iters < MAX_ITERATIONS {
iters += 1;
boxes.as_mut_slice().sort_unstable_by_key(|b| {
let mut probe = b.clone();
score(&mut probe, histo)
});
let mut top = match boxes.pop() {
Some(b) => b,
None => break,
};
if top.count(histo) == 0 {
break;
}
match median_cut(&top, histo) {
Some((mut left, Some(mut right))) => {
let l_pop = left.count(histo);
let r_pop = right.count(histo);
match (l_pop, r_pop) {
(0, 0) => {
let _ = exhausted.push(top);
}
(0, _) => {
let _ = boxes.push(right);
}
(_, 0) => {
let _ = boxes.push(left);
}
(_, _) => {
let _ = boxes.push(left);
let _ = boxes.push(right);
}
}
}
Some((only, None)) => {
let _ = exhausted.push(only);
}
None => {
let _ = exhausted.push(top);
}
}
}
exhausted.drain_into(boxes);
}
pub struct Mmcq {
histogram: [u32; HISTO_SIZE],
boxes: BoxArena,
}
impl Mmcq {
#[cfg_attr(not(tarpaulin), inline(always))]
pub const fn new() -> Self {
Self {
histogram: [0u32; HISTO_SIZE],
boxes: BoxArena::new(),
}
}
#[cfg(any(feature = "alloc", feature = "std"))]
#[cfg_attr(docsrs, doc(cfg(any(feature = "alloc", feature = "std"))))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub fn new_boxed() -> std::boxed::Box<Self> {
#[allow(unsafe_code)]
unsafe {
std::boxed::Box::<Self>::new_zeroed().assume_init()
}
}
pub fn extract<I, B>(&mut self, pixels: I, count: u8, algo: crate::Algorithm, out: &mut B)
where
I: Iterator<Item = [u8; 3]>,
B: Buffer<crate::Dominant>,
{
if count == 0 {
return;
}
self.boxes.clear();
build_histogram(pixels, &mut self.histogram);
let total_pixels: u32 = self
.histogram
.iter()
.fold(0u32, |a, b| a.saturating_add(*b));
let inv_total_pct: f32 = if total_pixels == 0 {
0.0
} else {
100.0 / total_pixels as f32
};
let initial = match initial_vbox(&self.histogram) {
Some(b) => b,
None => return,
};
let _ = self.boxes.push(initial);
let target = (count as usize).clamp(2, 256);
let phase1_target = (3 * target).div_ceil(4);
iterate_split(&mut self.boxes, phase1_target, &self.histogram, |b, h| {
b.count(h) as u64
});
iterate_split(&mut self.boxes, target, &self.histogram, |b, h| {
(b.count(h) as u64) * (b.volume() as u64)
});
let histo: &[u32] = &self.histogram;
self.boxes.as_mut_slice().sort_unstable_by_key(|b| {
let mut probe = b.clone();
core::cmp::Reverse(probe.count(histo))
});
let mut written: usize = 0;
let max_out = count as usize;
for vbox in self.boxes.as_mut_slice().iter_mut() {
if written >= max_out {
break;
}
let pop = vbox.count(histo);
if pop == 0 {
continue;
}
let rgb = vbox.avg(histo);
let dominant = crate::Dominant {
rgb,
color: algo.extract(rgb),
population: pop,
percentage: pop as f32 * inv_total_pct,
};
if out.try_push(dominant).is_some() {
break;
}
written += 1;
}
}
}
#[cfg(feature = "std")]
std::thread_local! {
static MMCQ_TLS: core::cell::RefCell<std::boxed::Box<Mmcq>> =
core::cell::RefCell::new(Mmcq::new_boxed());
}
#[cfg(feature = "std")]
pub(crate) fn quantize<I: Iterator<Item = [u8; 3]>>(
pixels: I,
count: u8,
algo: crate::Algorithm,
) -> std::vec::Vec<crate::Dominant> {
MMCQ_TLS.with(|cell| {
let mut mmcq = cell.borrow_mut();
let mut out: std::vec::Vec<crate::Dominant> = std::vec::Vec::new();
mmcq.extract(pixels, count, algo, &mut out);
out
})
}
#[cfg(all(feature = "alloc", not(feature = "std")))]
pub(crate) fn quantize<I: Iterator<Item = [u8; 3]>>(
pixels: I,
count: u8,
algo: crate::Algorithm,
) -> std::vec::Vec<crate::Dominant> {
let mut mmcq = Mmcq::new_boxed();
let mut out: std::vec::Vec<crate::Dominant> = std::vec::Vec::new();
mmcq.extract(pixels, count, algo, &mut out);
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::RgbFrame;
fn make_frame(width: u32, height: u32, pixels: &[[u8; 3]]) -> Vec<u8> {
assert_eq!(pixels.len() as u32, width * height);
let mut buf = Vec::with_capacity(pixels.len() * 3);
for p in pixels {
buf.extend_from_slice(p);
}
buf
}
#[test]
fn solid_red_frame_yields_a_red_dominant() {
let pixels = vec![[255, 0, 0]; 16];
let buf = make_frame(4, 4, &pixels);
let frame = RgbFrame::try_new(&buf, 4, 4, 12).expect("frame");
let dominants = quantize(frame.pixels(), 5, crate::Algorithm::default());
assert!(!dominants.is_empty(), "MMCQ produced zero dominants");
let top = &dominants[0];
assert!(top.rgb[0] > 200, "expected R>200, got {:?}", top.rgb);
assert!(top.rgb[1] < 30, "expected G<30, got {:?}", top.rgb);
assert!(top.rgb[2] < 30, "expected B<30, got {:?}", top.rgb);
}
#[test]
fn solid_white_recovered_at_bin_center() {
let pixels = vec![[255u8, 255, 255]; 16];
let buf = make_frame(4, 4, &pixels);
let frame = RgbFrame::try_new(&buf, 4, 4, 12).expect("frame");
let dominants = quantize(frame.pixels(), 5, crate::Algorithm::default());
let top = &dominants[0];
assert_eq!(top.rgb, [252, 252, 252], "expected bin-center white");
}
#[test]
fn solid_black_recovered_at_bin_center() {
let pixels = vec![[0u8, 0, 0]; 16];
let buf = make_frame(4, 4, &pixels);
let frame = RgbFrame::try_new(&buf, 4, 4, 12).expect("frame");
let dominants = quantize(frame.pixels(), 5, crate::Algorithm::default());
let top = &dominants[0];
assert_eq!(top.rgb, [4, 4, 4], "expected bin-center black");
}
#[test]
fn bin_edge_inputs_collapse_to_bin_center() {
for value in [8u8, 15u8] {
let pixels = vec![[value; 3]; 16];
let buf = make_frame(4, 4, &pixels);
let frame = RgbFrame::try_new(&buf, 4, 4, 12).expect("frame");
let dominants = quantize(frame.pixels(), 5, crate::Algorithm::default());
let top = &dominants[0];
assert_eq!(
top.rgb,
[12, 12, 12],
"input [{value};3] should collapse to bin-center [12;3], got {:?}",
top.rgb
);
}
}
#[test]
fn checkerboard_red_blue_yields_two_dominants() {
let red = [255, 0, 0];
let blue = [0, 0, 255];
let mut pixels = Vec::with_capacity(16);
for i in 0..16 {
pixels.push(if i % 2 == 0 { red } else { blue });
}
let buf = make_frame(4, 4, &pixels);
let frame = RgbFrame::try_new(&buf, 4, 4, 12).expect("frame");
let dominants = quantize(frame.pixels(), 5, crate::Algorithm::default());
assert!(
dominants.len() >= 2,
"expected at least 2 dominants, got {}",
dominants.len()
);
let has_red = dominants.iter().any(|d| d.rgb[0] > 200 && d.rgb[2] < 50);
let has_blue = dominants.iter().any(|d| d.rgb[2] > 200 && d.rgb[0] < 50);
assert!(
has_red && has_blue,
"expected red AND blue dominants, got {:?}",
dominants.iter().map(|d| d.rgb).collect::<Vec<_>>()
);
}
#[test]
fn padded_stride_is_respected() {
let mut buf = Vec::new();
buf.extend_from_slice(&[255, 0, 0, 255, 0, 0, 0xFF, 0xFF]);
buf.extend_from_slice(&[255, 0, 0, 255, 0, 0, 0xFF, 0xFF]);
let frame = RgbFrame::try_new(&buf, 2, 2, 8).expect("frame with padding");
let dominants = quantize(frame.pixels(), 5, crate::Algorithm::default());
let top = &dominants[0];
assert!(
top.rgb[0] > 200 && top.rgb[1] < 30 && top.rgb[2] < 30,
"padding leaked into the histogram; top dominant was {:?}",
top.rgb
);
}
#[test]
fn extract_with_empty_pixel_iterator_yields_nothing() {
let mut mmcq = Mmcq::new_boxed();
let mut out: Vec<crate::Dominant> = Vec::new();
mmcq.extract(
core::iter::empty::<[u8; 3]>(),
5,
crate::Algorithm::default(),
&mut out,
);
assert!(out.is_empty(), "empty input must produce no dominants");
}
#[test]
fn mmcq_extract_count_zero_does_no_work() {
let mut mmcq = Mmcq::new_boxed();
let buf = vec![255u8, 0, 0, 0, 255, 0, 0, 0, 255]; let frame = RgbFrame::try_new(&buf, 3, 1, 9).expect("frame");
let mut out: Vec<crate::Dominant> = Vec::new();
mmcq.extract(frame.pixels(), 0, crate::Algorithm::default(), &mut out);
assert!(out.is_empty(), "count=0 must produce no dominants");
}
#[test]
fn mmcq_extract_stops_when_output_buffer_full() {
let pixels: Vec<[u8; 3]> = (0..16)
.map(|i| if i < 8 { [255, 0, 0] } else { [0, 0, 255] })
.collect();
let buf = make_frame(4, 4, &pixels);
let frame = RgbFrame::try_new(&buf, 4, 4, 12).expect("frame");
let mut out: [Option<crate::Dominant>; 1] = [const { None }; 1];
let mut mmcq = Mmcq::new_boxed();
mmcq.extract(frame.pixels(), 5, crate::Algorithm::default(), &mut out);
assert!(out[0].is_some(), "first slot must be filled");
}
#[test]
fn mmcq_new_const_is_callable() {
static MMCQ: Mmcq = Mmcq::new();
assert_eq!(MMCQ.histogram[0], 0);
assert_eq!(MMCQ.boxes.len(), 0);
let runtime = std::boxed::Box::new(Mmcq::new());
assert_eq!(runtime.histogram[0], 0);
}
#[test]
fn mmcq_reuse_via_direct_extract_resets_state() {
let mut mmcq = Mmcq::new_boxed();
let red_pixels = vec![[200u8, 30, 30]; 16];
let blue_pixels = vec![[30u8, 30, 200]; 16];
let red_buf = make_frame(4, 4, &red_pixels);
let blue_buf = make_frame(4, 4, &blue_pixels);
let red_frame = RgbFrame::try_new(&red_buf, 4, 4, 12).expect("frame");
let blue_frame = RgbFrame::try_new(&blue_buf, 4, 4, 12).expect("frame");
let mut red_out: Vec<crate::Dominant> = Vec::new();
mmcq.extract(
red_frame.pixels(),
3,
crate::Algorithm::default(),
&mut red_out,
);
let mut blue_out: Vec<crate::Dominant> = Vec::new();
mmcq.extract(
blue_frame.pixels(),
3,
crate::Algorithm::default(),
&mut blue_out,
);
let red_top = red_out[0].rgb;
let blue_top = blue_out[0].rgb;
assert!(red_top[0] > 100 && red_top[2] < 100);
assert!(blue_top[2] > 100 && blue_top[0] < 100);
}
#[test]
fn box_arena_push_rejects_when_full() {
let mut arena = BoxArena::new();
let template = VBox {
r1: 0,
r2: 0,
g1: 0,
g2: 0,
b1: 0,
b2: 0,
count_cache: Some(1),
avg_cache: Some([0, 0, 0]),
};
for _ in 0..MAX_BOXES {
assert!(arena.push(template.clone()));
}
assert!(!arena.push(template), "push must reject when at capacity");
}
#[test]
fn vbox_avg_cache_hit_returns_same_value() {
let pixels = vec![[120u8, 80, 200]; 16];
let buf = make_frame(4, 4, &pixels);
let frame = RgbFrame::try_new(&buf, 4, 4, 12).expect("frame");
let mut histo = [0u32; HISTO_SIZE];
build_histogram(frame.pixels(), &mut histo);
let mut vbox = initial_vbox(&histo).expect("non-empty histogram");
let first = vbox.avg(&histo);
let second = vbox.avg(&histo);
assert_eq!(first, second, "avg must be stable across calls");
}
#[test]
fn vbox_avg_on_empty_box_returns_geometric_center() {
let histo = [0u32; HISTO_SIZE]; let mut empty = VBox {
r1: 0,
r2: 7,
g1: 0,
g2: 7,
b1: 0,
b2: 7,
count_cache: None,
avg_cache: None,
};
let avg = empty.avg(&histo);
assert_eq!(avg[0], avg[1]);
assert_eq!(avg[1], avg[2]);
assert!(avg[0] > 0);
}
}