Skip to main content

maple_render_core/
quantize.rs

1//! Median-cut color quantization with Floyd-Steinberg dithering.
2use image::RgbaImage;
3
4const HIST_C0_BITS: usize = 5; // R
5const HIST_C1_BITS: usize = 6; // G
6const HIST_C2_BITS: usize = 5; // B
7
8const HIST_C0_ELEMS: usize = 1 << HIST_C0_BITS;
9const HIST_C1_ELEMS: usize = 1 << HIST_C1_BITS;
10const HIST_C2_ELEMS: usize = 1 << HIST_C2_BITS;
11const HIST_ELEMS: usize = HIST_C0_ELEMS * HIST_C1_ELEMS * HIST_C2_ELEMS;
12
13const C0_SHIFT: usize = 8 - HIST_C0_BITS;
14const C1_SHIFT: usize = 8 - HIST_C1_BITS;
15const C2_SHIFT: usize = 8 - HIST_C2_BITS;
16
17// Distance scale factors (G > R > B perceptual weighting). /shrug
18const C0_SCALE: i32 = 2;
19const C1_SCALE: i32 = 3;
20const C2_SCALE: i32 = 1;
21
22// Box subdivision parameters
23const BOX_C0_LOG: usize = HIST_C0_BITS - 3;
24const BOX_C1_LOG: usize = HIST_C1_BITS - 3;
25const BOX_C2_LOG: usize = HIST_C2_BITS - 3;
26
27const BOX_C0_ELEMS: usize = 1 << BOX_C0_LOG;
28const BOX_C1_ELEMS: usize = 1 << BOX_C1_LOG;
29const BOX_C2_ELEMS: usize = 1 << BOX_C2_LOG;
30
31/// Number of histogram cells one `fill_inverse_cmap` call resolves.
32const BOX_ELEMS: usize = BOX_C0_ELEMS * BOX_C1_ELEMS * BOX_C2_ELEMS;
33
34const BOX_C0_SHIFT: usize = C0_SHIFT + BOX_C0_LOG;
35const BOX_C1_SHIFT: usize = C1_SHIFT + BOX_C1_LOG;
36const BOX_C2_SHIFT: usize = C2_SHIFT + BOX_C2_LOG;
37
38const MAXJSAMPLE: i32 = 255;
39const MAXNUMCOLORS: usize = 256;
40
41#[derive(Clone)]
42pub struct Palette {
43    pub red: [u8; 256],
44    pub green: [u8; 256],
45    pub blue: [u8; 256],
46    pub alpha: [u8; 256],
47    pub colors_total: usize,
48}
49
50impl Default for Palette {
51    fn default() -> Self {
52        Palette {
53            red: [0; 256],
54            green: [0; 256],
55            blue: [0; 256],
56            alpha: [255; 256],
57            colors_total: 0,
58        }
59    }
60}
61
62impl Palette {
63    pub fn new() -> Self {
64        Self::default()
65    }
66
67    pub fn get(&self, idx: usize) -> (u8, u8, u8, u8) {
68        (self.red[idx], self.green[idx], self.blue[idx], self.alpha[idx])
69    }
70
71    pub fn set(&mut self, idx: usize, r: u8, g: u8, b: u8) {
72        self.red[idx] = r;
73        self.green[idx] = g;
74        self.blue[idx] = b;
75        self.alpha[idx] = 255;
76    }
77}
78
79/// Bitmap of the occupied histogram cells, one `u32` per `(c0, c1)` row.
80///
81/// `HIST_C2_ELEMS` is exactly 32, so a whole row of the histogram fits in a
82/// single word and the 128 KB histogram condenses to 8 KB that stays in L1 for
83/// the entire median cut. The histogram is immutable while boxes are being
84/// split, so the bitmap is built once and stays valid for every split.
85struct Occupancy {
86    rows: Box<[u32; HIST_C0_ELEMS * HIST_C1_ELEMS]>,
87}
88
89const _: () = assert!(HIST_C2_ELEMS == u32::BITS as usize);
90const _: () = assert!(HIST_C0_ELEMS <= u32::BITS as usize);
91const _: () = assert!(HIST_C1_ELEMS <= u64::BITS as usize);
92
93impl Occupancy {
94    fn from_histogram(histogram: &[u16; HIST_ELEMS]) -> Self {
95        let mut rows = Box::new([0u32; HIST_C0_ELEMS * HIST_C1_ELEMS]);
96
97        for (bits, cells) in rows.iter_mut().zip(histogram.chunks_exact(HIST_C2_ELEMS)) {
98            let mut row = 0u32;
99            for (c2, &count) in cells.iter().enumerate() {
100                row |= ((count != 0) as u32) << c2;
101            }
102            *bits = row;
103        }
104
105        Occupancy { rows }
106    }
107
108    #[inline(always)]
109    fn row(&self, c0: i32, c1: i32) -> u32 {
110        self.rows[c0 as usize * HIST_C1_ELEMS + c1 as usize]
111    }
112
113    /// Mask of the `c2` bits a box spans.
114    #[inline(always)]
115    fn c2_mask(c2min: i32, c2max: i32) -> u32 {
116        let width = (c2max - c2min + 1) as u32;
117        if width >= u32::BITS { u32::MAX } else { ((1u32 << width) - 1) << c2min }
118    }
119}
120
121#[derive(Clone, Copy, Default)]
122struct ColorBox {
123    c0min: i32,
124    c0max: i32,
125    c1min: i32,
126    c1max: i32,
127    c2min: i32,
128    c2max: i32,
129    volume: i64,
130    colorcount: i64,
131}
132
133pub struct Quantizer {
134    histogram: Box<[u16; HIST_ELEMS]>,
135    fserrors: Vec<i16>,
136    error_limiter: Vec<i32>,
137    on_odd_row: bool,
138    palette: Palette,
139}
140
141impl Quantizer {
142    pub fn new(reference: &RgbaImage) -> Self {
143        let mut q = Quantizer {
144            histogram: Box::new([0; HIST_ELEMS]),
145            fserrors: Vec::new(),
146            error_limiter: Vec::new(),
147            on_odd_row: false,
148            palette: Palette::new(),
149        };
150
151        q.init_error_limit();
152
153        let width = reference.width() as usize;
154        q.fserrors = vec![0i16; (width + 2) * 3];
155
156        q.prescan_quantize(reference);
157        q.select_colors(MAXNUMCOLORS);
158        q.zero_histogram();
159
160        q
161    }
162
163    fn init_error_limit(&mut self) {
164        self.error_limiter = vec![0i32; (MAXJSAMPLE * 2 + 1) as usize];
165        let table_offset = MAXJSAMPLE as usize;
166
167        const STEPSIZE: i32 = (MAXJSAMPLE + 1) / 16;
168
169        let mut out: i32 = 0;
170
171        // 1:1 up to +- MAXJSAMPLE/16
172        for inp in 0..STEPSIZE {
173            self.error_limiter[table_offset.wrapping_add(inp as usize)] = out;
174            self.error_limiter[table_offset.wrapping_sub(inp as usize)] = -out;
175            out += 1;
176        }
177
178        // 1:2 up to +- 3*MAXJSAMPLE/16
179        for inp in STEPSIZE..(STEPSIZE * 3) {
180            self.error_limiter[table_offset.wrapping_add(inp as usize)] = out;
181            self.error_limiter[table_offset.wrapping_sub(inp as usize)] = -out;
182            if inp & 1 == 0 {
183                out += 1;
184            }
185        }
186
187        // Clamp the rest
188        for inp in (STEPSIZE * 3)..=MAXJSAMPLE {
189            self.error_limiter[table_offset.wrapping_add(inp as usize)] = out;
190            self.error_limiter[table_offset.wrapping_sub(inp as usize)] = -out;
191        }
192    }
193
194    fn zero_histogram(&mut self) {
195        self.histogram.fill(0);
196    }
197
198    #[inline(always)]
199    const fn histogram_index(c0: usize, c1: usize, c2: usize) -> usize {
200        (c0 * HIST_C1_ELEMS + c1) * HIST_C2_ELEMS + c2
201    }
202
203    #[inline(always)]
204    fn histogram_index_of(pixel: &[u8; 4]) -> usize {
205        Self::histogram_index(
206            (pixel[0] as usize) >> C0_SHIFT,
207            (pixel[1] as usize) >> C1_SHIFT,
208            (pixel[2] as usize) >> C2_SHIFT,
209        )
210    }
211
212    fn prescan_quantize(&mut self, img: &RgbaImage) {
213        // The arithmetic per pixel is trivial; what costs is the scatter into
214        // the histogram. Neighbouring pixels of a photo usually land in the
215        // same cell, so a single table turns the scan into one long chain of
216        // store-to-load forwarded increments. Scattering even and odd pixels
217        // into two tables halves that chain; the tables are folded back
218        // together afterwards in one linear pass.
219        //
220        // Two lanes is the sweet spot: four makes the working set larger than
221        // the level of cache that keeps up, and loses more than the shorter
222        // chain wins.
223        let (pairs, tail) = img.as_raw().as_chunks::<8>();
224        let mut odd_counts = vec![0u16; HIST_ELEMS];
225
226        for pair in pairs {
227            let (pixels, _) = pair.as_chunks::<4>();
228            let even = Self::histogram_index_of(&pixels[0]);
229            let odd = Self::histogram_index_of(&pixels[1]);
230
231            let cell = &mut self.histogram[even];
232            if *cell < u16::MAX {
233                *cell += 1;
234            }
235
236            let cell = &mut odd_counts[odd];
237            if *cell < u16::MAX {
238                *cell += 1;
239            }
240        }
241
242        for pixel in tail.as_chunks::<4>().0 {
243            let cell = &mut self.histogram[Self::histogram_index_of(pixel)];
244            if *cell < u16::MAX {
245                *cell += 1;
246            }
247        }
248
249        // Folding the two halves back together is a linear, vectorizable pass.
250        for (cell, odd) in self.histogram.iter_mut().zip(odd_counts.iter()) {
251            *cell = cell.saturating_add(*odd);
252        }
253    }
254
255    fn find_biggest_color_pop(boxlist: &[ColorBox], numboxes: usize) -> Option<usize> {
256        let mut maxc: i64 = 0;
257        let mut which = None;
258
259        for (i, bx) in boxlist[..numboxes].iter().enumerate() {
260            if bx.colorcount > maxc && bx.volume > 0 {
261                which = Some(i);
262                maxc = bx.colorcount;
263            }
264        }
265
266        which
267    }
268
269    fn find_biggest_volume(boxlist: &[ColorBox], numboxes: usize) -> Option<usize> {
270        let mut maxv: i64 = 0;
271        let mut which = None;
272
273        for (i, bx) in boxlist[..numboxes].iter().enumerate() {
274            if bx.volume > maxv {
275                which = Some(i);
276                maxv = bx.volume;
277            }
278        }
279
280        which
281    }
282
283    fn update_box(&self, occupancy: &Occupancy, boxp: &mut ColorBox) {
284        let original = *boxp;
285        let mask = Occupancy::c2_mask(original.c2min, original.c2max);
286
287        // The scan is entirely branch-free: every row of the box is one masked
288        // load, a popcount and two ORs. Rather than widening six bounds as it
289        // goes, it collects which c0/c1/c2 coordinates are occupied as bitmaps
290        // and reads the bounds off them once at the end.
291        let mut colorcount: i64 = 0;
292        let mut c0_used: u32 = 0;
293        let mut c1_used: u64 = 0;
294        let mut c2_used: u32 = 0;
295
296        for c0 in original.c0min..=original.c0max {
297            let mut plane: u32 = 0;
298
299            for c1 in original.c1min..=original.c1max {
300                let row = occupancy.row(c0, c1) & mask;
301                colorcount += row.count_ones() as i64;
302                plane |= row;
303                c1_used |= ((row != 0) as u64) << c1;
304            }
305
306            c2_used |= plane;
307            c0_used |= ((plane != 0) as u32) << c0;
308        }
309
310        if colorcount != 0 {
311            boxp.c0min = c0_used.trailing_zeros() as i32;
312            boxp.c0max = (u32::BITS - 1 - c0_used.leading_zeros()) as i32;
313            boxp.c1min = c1_used.trailing_zeros() as i32;
314            boxp.c1max = (u64::BITS - 1 - c1_used.leading_zeros()) as i32;
315            boxp.c2min = c2_used.trailing_zeros() as i32;
316            boxp.c2max = (u32::BITS - 1 - c2_used.leading_zeros()) as i32;
317        }
318
319        let dist0 = ((boxp.c0max - boxp.c0min) << C0_SHIFT) as i64 * C0_SCALE as i64;
320        let dist1 = ((boxp.c1max - boxp.c1min) << C1_SHIFT) as i64 * C1_SCALE as i64;
321        let dist2 = ((boxp.c2max - boxp.c2min) << C2_SHIFT) as i64 * C2_SCALE as i64;
322        boxp.volume = dist0 * dist0 + dist1 * dist1 + dist2 * dist2;
323        boxp.colorcount = colorcount;
324    }
325
326    fn median_cut(
327        &self,
328        occupancy: &Occupancy,
329        boxlist: &mut [ColorBox],
330        mut numboxes: usize,
331        desired_colors: usize,
332    ) -> usize {
333        while numboxes < desired_colors {
334            // Select box to split
335            let b1_idx = if numboxes * 2 <= desired_colors {
336                Self::find_biggest_color_pop(boxlist, numboxes)
337            } else {
338                Self::find_biggest_volume(boxlist, numboxes)
339            };
340
341            let b1_idx = match b1_idx {
342                Some(idx) => idx,
343                None => break,
344            };
345
346            let b1 = boxlist[b1_idx];
347            let b2_idx = numboxes;
348            boxlist[b2_idx] = b1;
349
350            let c0 = ((b1.c0max - b1.c0min) << C0_SHIFT) * C0_SCALE;
351            let c1 = ((b1.c1max - b1.c1min) << C1_SHIFT) * C1_SCALE;
352            let c2 = ((b1.c2max - b1.c2min) << C2_SHIFT) * C2_SCALE;
353
354            let mut cmax = c1;
355            let mut n = 1;
356            if c2 > cmax {
357                cmax = c2;
358                n = 2;
359            }
360            if c0 > cmax {
361                n = 0;
362            }
363
364            match n {
365                0 => {
366                    let lb = (b1.c0max + b1.c0min) / 2;
367                    boxlist[b1_idx].c0max = lb;
368                    boxlist[b2_idx].c0min = lb + 1;
369                }
370                1 => {
371                    let lb = (b1.c1max + b1.c1min) / 2;
372                    boxlist[b1_idx].c1max = lb;
373                    boxlist[b2_idx].c1min = lb + 1;
374                }
375                2 => {
376                    let lb = (b1.c2max + b1.c2min) / 2;
377                    boxlist[b1_idx].c2max = lb;
378                    boxlist[b2_idx].c2min = lb + 1;
379                }
380                _ => unreachable!(),
381            }
382
383            self.update_box(occupancy, &mut boxlist[b1_idx]);
384            self.update_box(occupancy, &mut boxlist[b2_idx]);
385            numboxes += 1;
386        }
387
388        numboxes
389    }
390
391    fn compute_color(&self, occupancy: &Occupancy, boxp: &ColorBox) -> (u8, u8, u8) {
392        let mut total: i64 = 0;
393        let mut c0total: i64 = 0;
394        let mut c1total: i64 = 0;
395        let mut c2total: i64 = 0;
396
397        let mask = Occupancy::c2_mask(boxp.c2min, boxp.c2max);
398
399        for c0 in boxp.c0min..=boxp.c0max {
400            for c1 in boxp.c1min..=boxp.c1max {
401                // Walking the set bits visits only the occupied cells, so the
402                // empty ones never reach the histogram at all.
403                let mut row = occupancy.row(c0, c1) & mask;
404                while row != 0 {
405                    let c2 = row.trailing_zeros() as i32;
406                    row &= row - 1;
407
408                    let count = self.histogram
409                        [Self::histogram_index(c0 as usize, c1 as usize, c2 as usize)]
410                        as i64;
411                    total += count;
412                    c0total += ((c0 << C0_SHIFT) + (1 << (C0_SHIFT - 1))) as i64 * count;
413                    c1total += ((c1 << C1_SHIFT) + (1 << (C1_SHIFT - 1))) as i64 * count;
414                    c2total += ((c2 << C2_SHIFT) + (1 << (C2_SHIFT - 1))) as i64 * count;
415                }
416            }
417        }
418
419        if total > 0 {
420            (
421                ((c0total + (total >> 1)) / total) as u8,
422                ((c1total + (total >> 1)) / total) as u8,
423                ((c2total + (total >> 1)) / total) as u8,
424            )
425        } else {
426            (255, 255, 255)
427        }
428    }
429
430    fn select_colors(&mut self, desired_colors: usize) {
431        let mut boxlist = vec![ColorBox::default(); desired_colors];
432
433        // Initialize one box containing whole space
434        boxlist[0] = ColorBox {
435            c0min: 0,
436            c0max: (MAXJSAMPLE >> C0_SHIFT) as i32,
437            c1min: 0,
438            c1max: (MAXJSAMPLE >> C1_SHIFT) as i32,
439            c2min: 0,
440            c2max: (MAXJSAMPLE >> C2_SHIFT) as i32,
441            volume: 0,
442            colorcount: 0,
443        };
444
445        let occupancy = Occupancy::from_histogram(&self.histogram);
446
447        self.update_box(&occupancy, &mut boxlist[0]);
448        let numboxes = self.median_cut(&occupancy, &mut boxlist, 1, desired_colors);
449
450        for i in 0..numboxes {
451            let (r, g, b) = self.compute_color(&occupancy, &boxlist[i]);
452            self.palette.set(i, r, g, b);
453        }
454        self.palette.colors_total = numboxes;
455    }
456
457    pub fn palette(&self) -> &Palette {
458        &self.palette
459    }
460
461    pub fn palette_mut(&mut self) -> &mut Palette {
462        &mut self.palette
463    }
464
465    fn find_nearby_colors(
466        &self,
467        minc0: i32,
468        minc1: i32,
469        minc2: i32,
470        colorlist: &mut [u8; MAXNUMCOLORS],
471    ) -> usize {
472        let numcolors = self.palette.colors_total;
473
474        let maxc0 = minc0 + ((1 << BOX_C0_SHIFT) - (1 << C0_SHIFT));
475        let centerc0 = (minc0 + maxc0) >> 1;
476        let maxc1 = minc1 + ((1 << BOX_C1_SHIFT) - (1 << C1_SHIFT));
477        let centerc1 = (minc1 + maxc1) >> 1;
478        let maxc2 = minc2 + ((1 << BOX_C2_SHIFT) - (1 << C2_SHIFT));
479        let centerc2 = (minc2 + maxc2) >> 1;
480
481        // A weighted squared distance never exceeds (255 * 3)^2 * 3, so the
482        // whole computation fits in an i32 - half the traffic of the i64 it
483        // used to run in, and a min-reduce the compiler can vectorize.
484        let mut mindist = [0i32; MAXNUMCOLORS];
485        let mut minmaxdist: i32 = i32::MAX;
486
487        for i in 0..numcolors {
488            let x0 = self.palette.red[i] as i32;
489            let (min_dist0, max_dist0) =
490                Self::compute_dist_component(x0, minc0, maxc0, centerc0, C0_SCALE);
491
492            let x1 = self.palette.green[i] as i32;
493            let (min_dist1, max_dist1) =
494                Self::compute_dist_component(x1, minc1, maxc1, centerc1, C1_SCALE);
495
496            let x2 = self.palette.blue[i] as i32;
497            let (min_dist2, max_dist2) =
498                Self::compute_dist_component(x2, minc2, maxc2, centerc2, C2_SCALE);
499
500            mindist[i] = min_dist0 + min_dist1 + min_dist2;
501            let max_dist = max_dist0 + max_dist1 + max_dist2;
502            if max_dist < minmaxdist {
503                minmaxdist = max_dist;
504            }
505        }
506
507        let mut ncolors = 0;
508        for i in 0..numcolors {
509            if mindist[i] <= minmaxdist {
510                colorlist[ncolors] = i as u8;
511                ncolors += 1;
512            }
513        }
514
515        ncolors
516    }
517
518    fn compute_dist_component(
519        x: i32,
520        minc: i32,
521        maxc: i32,
522        centerc: i32,
523        scale: i32,
524    ) -> (i32, i32) {
525        if x < minc {
526            let tdist = (x - minc) * scale;
527            let min_dist = tdist * tdist;
528            let tdist = (x - maxc) * scale;
529            let max_dist = tdist * tdist;
530            (min_dist, max_dist)
531        } else if x > maxc {
532            let tdist = (x - maxc) * scale;
533            let min_dist = tdist * tdist;
534            let tdist = (x - minc) * scale;
535            let max_dist = tdist * tdist;
536            (min_dist, max_dist)
537        } else {
538            let tdist = if x <= centerc { (x - maxc) * scale } else { (x - minc) * scale };
539            (0, tdist * tdist)
540        }
541    }
542
543    fn find_best_colors(
544        &self,
545        minc0: i32,
546        minc1: i32,
547        minc2: i32,
548        numcolors: usize,
549        colorlist: &[u8; MAXNUMCOLORS],
550        bestcolor: &mut [u8; BOX_ELEMS],
551    ) {
552        // The distances here stay in i64 on purpose, even though they would fit
553        // in an i32. Every cell is a compare against the running best that
554        // almost never wins after the first candidate color, so the branchy
555        // scalar loop below is the shape we want. With i32 distances - and in
556        // particular with the winning color staged in a second i32 array next
557        // to it - LLVM instead turns the whole 128-cell update into an
558        // unconditional vector compare-and-select. That trades a
559        // near-perfectly-predicted branch for a load/select/store on every
560        // cell, which is a large loss on the aarch64 macro runners the
561        // benchmarks are measured on (`find_best_colors` 842 us -> 1.3 ms),
562        // however it may look on an x86 dev box. Keep this loop scalar.
563        let mut bestdist = [i64::MAX; BOX_ELEMS];
564
565        const STEP_C0: i64 = ((1 << C0_SHIFT) * C0_SCALE) as i64;
566        const STEP_C1: i64 = ((1 << C1_SHIFT) * C1_SCALE) as i64;
567        const STEP_C2: i64 = ((1 << C2_SHIFT) * C2_SCALE) as i64;
568
569        for i in 0..numcolors {
570            let icolor = colorlist[i];
571            let r = self.palette.red[icolor as usize] as i32;
572            let g = self.palette.green[icolor as usize] as i32;
573            let b = self.palette.blue[icolor as usize] as i32;
574
575            let mut inc0 = (minc0 - r) as i64 * C0_SCALE as i64;
576            let mut dist0 = inc0 * inc0;
577            let mut inc1 = (minc1 - g) as i64 * C1_SCALE as i64;
578            dist0 += inc1 * inc1;
579            let mut inc2 = (minc2 - b) as i64 * C2_SCALE as i64;
580            dist0 += inc2 * inc2;
581
582            inc0 = inc0 * (2 * STEP_C0) + STEP_C0 * STEP_C0;
583            inc1 = inc1 * (2 * STEP_C1) + STEP_C1 * STEP_C1;
584            inc2 = inc2 * (2 * STEP_C2) + STEP_C2 * STEP_C2;
585
586            let mut bptr_idx = 0;
587            let mut xx0 = inc0;
588
589            for _ic0 in 0..BOX_C0_ELEMS {
590                let mut dist1 = dist0;
591                let mut xx1 = inc1;
592
593                for _ic1 in 0..BOX_C1_ELEMS {
594                    let mut dist2 = dist1;
595                    let mut xx2 = inc2;
596
597                    for _ic2 in 0..BOX_C2_ELEMS {
598                        if dist2 < bestdist[bptr_idx] {
599                            bestdist[bptr_idx] = dist2;
600                            bestcolor[bptr_idx] = icolor;
601                        }
602                        dist2 += xx2;
603                        xx2 += 2 * STEP_C2 * STEP_C2;
604                        bptr_idx += 1;
605                    }
606                    dist1 += xx1;
607                    xx1 += 2 * STEP_C1 * STEP_C1;
608                }
609                dist0 += xx0;
610                xx0 += 2 * STEP_C0 * STEP_C0;
611            }
612        }
613    }
614
615    fn fill_inverse_cmap(&mut self, c0: i32, c1: i32, c2: i32) {
616        // These are small, fixed-size and dead by the end of the call, so they
617        // live on the stack. Heap-allocating and zeroing them on every call
618        // was pure overhead on the cold-cache path.
619        let mut colorlist = [0u8; MAXNUMCOLORS];
620        let mut bestcolor = [0u8; BOX_ELEMS];
621
622        let bc0 = c0 >> BOX_C0_LOG as i32;
623        let bc1 = c1 >> BOX_C1_LOG as i32;
624        let bc2 = c2 >> BOX_C2_LOG as i32;
625
626        let minc0 = (bc0 << BOX_C0_SHIFT) + (1 << (C0_SHIFT - 1));
627        let minc1 = (bc1 << BOX_C1_SHIFT) + (1 << (C1_SHIFT - 1));
628        let minc2 = (bc2 << BOX_C2_SHIFT) + (1 << (C2_SHIFT - 1));
629
630        let numcolors = self.find_nearby_colors(minc0, minc1, minc2, &mut colorlist);
631        self.find_best_colors(minc0, minc1, minc2, numcolors, &colorlist, &mut bestcolor);
632
633        let base_c0 = (bc0 << BOX_C0_LOG as i32) as usize;
634        let base_c1 = (bc1 << BOX_C1_LOG as i32) as usize;
635        let base_c2 = (bc2 << BOX_C2_LOG as i32) as usize;
636
637        let mut cptr_idx = 0;
638        for ic0 in 0..BOX_C0_ELEMS {
639            for ic1 in 0..BOX_C1_ELEMS {
640                for ic2 in 0..BOX_C2_ELEMS {
641                    let histogram_index =
642                        Self::histogram_index(base_c0 + ic0, base_c1 + ic1, base_c2 + ic2);
643                    self.histogram[histogram_index] = bestcolor[cptr_idx] as u16 + 1;
644                    cptr_idx += 1;
645                }
646            }
647        }
648    }
649
650    pub fn quantize_no_dither(&mut self, img: &RgbaImage) -> Vec<u8> {
651        let width = img.width() as usize;
652        let height = img.height() as usize;
653        let mut output = vec![0u8; width * height];
654
655        let (pixels, _) = img.as_raw().as_chunks::<4>();
656
657        for (out_row, src_row) in output.chunks_exact_mut(width).zip(pixels.chunks_exact(width)) {
658            for (out, pixel) in out_row.iter_mut().zip(src_row) {
659                let c0 = (pixel[0] as usize) >> C0_SHIFT;
660                let c1 = (pixel[1] as usize) >> C1_SHIFT;
661                let c2 = (pixel[2] as usize) >> C2_SHIFT;
662
663                let histogram_index = Self::histogram_index(c0, c1, c2);
664                let mut cached = self.histogram[histogram_index];
665                if cached == 0 {
666                    self.fill_inverse_cmap(c0 as i32, c1 as i32, c2 as i32);
667                    cached = self.histogram[histogram_index];
668                }
669
670                *out = (cached - 1) as u8;
671            }
672        }
673
674        output
675    }
676
677    pub fn quantize_fs_dither(&mut self, img: &RgbaImage) -> Vec<u8> {
678        let width = img.width() as usize;
679        let height = img.height() as usize;
680        let mut output = vec![0u8; width * height];
681
682        // Taking the error buffer out of `self` for the length of the scan
683        // lets the column loop work on a plain `&mut [i16]` while still being
684        // able to call `fill_inverse_cmap` on `self` when the colormap misses.
685        // It also reuses the previous frame's allocation.
686        let mut fserrors = std::mem::take(&mut self.fserrors);
687        fserrors.clear();
688        fserrors.resize((width + 2) * 3, 0);
689        self.on_odd_row = false;
690
691        let table_offset = MAXJSAMPLE as usize;
692        let (pixels, _) = img.as_raw().as_chunks::<4>();
693
694        for row in 0..height {
695            // One bounds check per row instead of a bounds-checked
696            // `get_pixel` and a bounds-checked `output[row * width + x]` per
697            // pixel.
698            let src_row = &pixels[row * width..row * width + width];
699            let out_row = &mut output[row * width..row * width + width];
700
701            let (dir, start_col, end_col, errorptr_start) = if self.on_odd_row {
702                (-1i32, width as i32 - 1, -1i32, (width + 1) * 3)
703            } else {
704                (1i32, 0i32, width as i32, 0usize)
705            };
706
707            let mut cur0: i32 = 0;
708            let mut cur1: i32 = 0;
709            let mut cur2: i32 = 0;
710            let mut belowerr0: i32 = 0;
711            let mut belowerr1: i32 = 0;
712            let mut belowerr2: i32 = 0;
713            let mut bpreverr0: i32 = 0;
714            let mut bpreverr1: i32 = 0;
715            let mut bpreverr2: i32 = 0;
716
717            let mut col = start_col;
718            let mut errorptr = errorptr_start as i32;
719            let dir3 = dir * 3;
720
721            while col != end_col {
722                let x = col as usize;
723                let pixel = &src_row[x];
724
725                // Add error from previous and below
726                let ep_idx = (errorptr + dir3) as usize;
727                let below = &fserrors[ep_idx..ep_idx + 3];
728                cur0 = (cur0 + below[0] as i32 + 8) >> 4;
729                cur1 = (cur1 + below[1] as i32 + 8) >> 4;
730                cur2 = (cur2 + below[2] as i32 + 8) >> 4;
731
732                cur0 = self.error_limiter[table_offset.wrapping_add(cur0 as usize)];
733                cur1 = self.error_limiter[table_offset.wrapping_add(cur1 as usize)];
734                cur2 = self.error_limiter[table_offset.wrapping_add(cur2 as usize)];
735
736                cur0 += pixel[0] as i32;
737                cur1 += pixel[1] as i32;
738                cur2 += pixel[2] as i32;
739                cur0 = cur0.clamp(0, 255);
740                cur1 = cur1.clamp(0, 255);
741                cur2 = cur2.clamp(0, 255);
742
743                let c0 = (cur0 as usize) >> C0_SHIFT;
744                let c1 = (cur1 as usize) >> C1_SHIFT;
745                let c2 = (cur2 as usize) >> C2_SHIFT;
746
747                let histogram_index = Self::histogram_index(c0, c1, c2);
748                let mut cached = self.histogram[histogram_index];
749                if cached == 0 {
750                    self.fill_inverse_cmap(c0 as i32, c1 as i32, c2 as i32);
751                    cached = self.histogram[histogram_index];
752                }
753
754                let pixcode = (cached - 1) as usize;
755                out_row[x] = pixcode as u8;
756
757                cur0 -= self.palette.red[pixcode] as i32;
758                cur1 -= self.palette.green[pixcode] as i32;
759                cur2 -= self.palette.blue[pixcode] as i32;
760
761                let here = &mut fserrors[errorptr as usize..errorptr as usize + 3];
762
763                let mut bnexterr = cur0;
764                let mut delta = cur0 * 2;
765                cur0 += delta; // 3x
766                here[0] = (bpreverr0 + cur0) as i16;
767                cur0 += delta; // 5x
768                bpreverr0 = belowerr0 + cur0;
769                belowerr0 = bnexterr;
770                cur0 += delta; // 7x
771
772                bnexterr = cur1;
773                delta = cur1 * 2;
774                cur1 += delta;
775                here[1] = (bpreverr1 + cur1) as i16;
776                cur1 += delta;
777                bpreverr1 = belowerr1 + cur1;
778                belowerr1 = bnexterr;
779                cur1 += delta;
780
781                bnexterr = cur2;
782                delta = cur2 * 2;
783                cur2 += delta;
784                here[2] = (bpreverr2 + cur2) as i16;
785                cur2 += delta;
786                bpreverr2 = belowerr2 + cur2;
787                belowerr2 = bnexterr;
788                cur2 += delta;
789
790                col += dir;
791                errorptr += dir3;
792            }
793
794            let tail = &mut fserrors[errorptr as usize..errorptr as usize + 3];
795            tail[1] = bpreverr1 as i16;
796            tail[2] = bpreverr2 as i16;
797
798            self.on_odd_row = !self.on_odd_row;
799        }
800
801        self.fserrors = fserrors;
802
803        output
804    }
805
806    pub fn quantize(&mut self, img: &RgbaImage, dither: bool) -> Vec<u8> {
807        if dither { self.quantize_fs_dither(img) } else { self.quantize_no_dither(img) }
808    }
809
810    pub fn sync_palette_from(&mut self, other: &Palette) {
811        self.palette = other.clone();
812    }
813}
814
815pub fn sync_palette(from: &Palette, to: &mut Palette) {
816    *to = from.clone();
817}