Officially support large palettes

Cargo.toml

@@ -17,6 +17,8 @@ rust-version = "1.60"
 [features]
 default = ["threads"]
 threads = ["dep:rayon", "dep:num_cpus", "dep:thread_local"]
+# supports up to 2048 colors for palettes, but NOT FOR REMAPPING
+large_palettes = []
 
 # this is private and unstable for imagequant-sys only, do not use
 _internal_c_ffi = []

src/image.rs

@@ -1,7 +1,7 @@
 use crate::attr::Attributes;
 use crate::blur::{liq_blur, liq_max3, liq_min3};
 use crate::error::*;
-use crate::pal::{f_pixel, PalF, PalIndex, MAX_COLORS, MIN_OPAQUE_A, RGBA};
+use crate::pal::{f_pixel, PalF, PalIndexRemap, MAX_COLORS, MIN_OPAQUE_A, RGBA};
 use crate::remap::DitherMapMode;
 use crate::rows::{DynamicRows, PixelsSource};
 use crate::seacow::Pointer;
@@ -123,7 +123,7 @@ impl<'pixels> Image<'pixels> {
         true
     }
 
-    pub(crate) fn update_dither_map(&mut self, remapped_image: &RowBitmap<'_, PalIndex>, palette: &PalF, uses_background: bool) -> Result<(), Error> {
+    pub(crate) fn update_dither_map(&mut self, remapped_image: &RowBitmap<'_, PalIndexRemap>, palette: &PalF, uses_background: bool) -> Result<(), Error> {
         if self.edges.is_none() {
             self.contrast_maps()?;
         }

src/nearest.rs

@@ -6,6 +6,9 @@ use crate::{Error, OrdFloat};
 impl<'pal> Nearest<'pal> {
     #[inline(never)]
     pub fn new(palette: &'pal PalF) -> Result<Self, Error> {
+        if palette.len() > PalIndex::MAX as usize + 1 {
+            return Err(Error::Unsupported);
+        }
         let mut indexes: Vec<_> = (0..palette.len())
             .map(|idx| MapIndex { idx: idx as _ })
             .collect();
@@ -20,7 +23,7 @@ impl<'pal> Nearest<'pal> {
         for (i, color) in palette.as_slice().iter().enumerate() {
             let mut best = Visitor {
                 idx: 0, distance: f32::MAX, distance_squared: f32::MAX,
-                exclude: i as i16,
+                exclude: Some(i as PalIndex),
             };
             vp_search_node(&handle.root, color, &mut best);
             handle.nearest_other_color_dist[i] = best.distance_squared / 4.;
@@ -42,14 +45,14 @@ impl Nearest<'_> {
                 distance: guess_diff.sqrt(),
                 distance_squared: guess_diff,
                 idx: likely_colormap_index,
-                exclude: -1,
+                exclude: None,
             }
         } else {
-            Visitor { distance: f32::INFINITY, distance_squared: f32::INFINITY, idx: 0, exclude: -1, }
+            Visitor { distance: f32::INFINITY, distance_squared: f32::INFINITY, idx: 0, exclude: None, }
         };
 
         vp_search_node(&self.root, px, &mut best_candidate);
-        (best_candidate.idx as PalIndex, best_candidate.distance_squared)
+        (best_candidate.idx, best_candidate.distance_squared)
     }
 }
 
@@ -67,13 +70,13 @@ pub struct Visitor {
     pub distance: f32,
     pub distance_squared: f32,
     pub idx: PalIndex,
-    pub exclude: i16,
+    pub exclude: Option<PalIndex>,
 }
 
 impl Visitor {
     #[inline]
     fn visit(&mut self, distance: f32, distance_squared: f32, idx: PalIndex) {
-        if distance_squared < self.distance_squared && self.exclude != i16::from(idx) {
+        if distance_squared < self.distance_squared && self.exclude != Some(idx) {
             self.distance = distance;
             self.distance_squared = distance_squared;
             self.idx = idx;

src/pal.rs

@@ -197,11 +197,18 @@ impl PalPop {
     }
 }
 
-/// This could be increased to support > 256 colors
+#[cfg(feature = "large_palettes")]
+pub type PalIndex = u16;
+
+#[cfg(not(feature = "large_palettes"))]
 pub type PalIndex = u8;
+
+/// This could be increased to support > 256 colors in remapping too
+pub type PalIndexRemap = u8;
 pub type PalLen = u16;
 
-pub(crate) const MAX_COLORS: usize = 256;
+/// Palettes are stored on the stack, and really large ones will cause stack overflows
+pub(crate) const MAX_COLORS: usize = if PalIndex::MAX == 255 { 256 } else { 2048 };
 
 /// A palette of premultiplied ARGB 4xf32 colors in internal gamma
 #[derive(Clone)]
@@ -249,7 +256,7 @@ impl PalF {
         }
 
         // if using low quality, there's a chance mediancut won't create enough colors in the palette
-        let max_fixed_colors = fixed_colors.len().min(max_colors.into());
+        let max_fixed_colors = fixed_colors.len().min(max_colors as usize);
         if self.len() < max_fixed_colors {
             let needs_extra = max_fixed_colors - self.len();
             self.colors.extend(fixed_colors.iter().copied().take(needs_extra));
@@ -421,3 +428,14 @@ fn pal_test() {
         assert_eq!(int_pal[i as usize], RGBA::new(i, i, i, 100 + i / 2));
     }
 }
+
+#[test]
+#[cfg(feature = "large_palettes")]
+fn largepal() {
+    let gamma = gamma_lut(0.5);
+    let mut p = PalF::new();
+    for i in 0..1000 {
+        let rgba = RGBA::new(i as u8, (i/2) as u8, (i/4) as u8, 255);
+        p.push(f_pixel::from_rgba(&gamma, rgba), PalPop::new(1.));
+    }
+}

src/quant.rs

@@ -4,7 +4,7 @@ use crate::hist::HistogramInternal;
 use crate::image::Image;
 use crate::kmeans::Kmeans;
 use crate::mediancut::mediancut;
-use crate::pal::{PalIndex, PalF, PalLen, PalPop, Palette, LIQ_WEIGHT_MSE, MAX_COLORS, MAX_TRANSP_A, RGBA};
+use crate::pal::{PalIndexRemap, PalF, PalLen, PalPop, Palette, LIQ_WEIGHT_MSE, MAX_COLORS, MAX_TRANSP_A, RGBA};
 use crate::remap::{mse_to_standard_mse, DitherMapMode, Remapped};
 use crate::remap::{remap_to_palette, remap_to_palette_floyd};
 use crate::seacow::RowBitmapMut;
@@ -74,14 +74,14 @@ impl QuantizationResult {
     /// This is 100% redundant and unnecessary. This work is done anyway when remap is called.
     /// However, this can be called before calling `image.set_background()`, so it may allow better parallelization while the background is generated on another thread.
     #[doc(hidden)]
-    pub fn optionally_prepare_for_dithering_with_background_set(&mut self, image: &mut Image<'_>, output_buf: &mut [MaybeUninit<PalIndex>]) -> Result<(), Error> {
+    pub fn optionally_prepare_for_dithering_with_background_set(&mut self, image: &mut Image<'_>, output_buf: &mut [MaybeUninit<PalIndexRemap>]) -> Result<(), Error> {
         let mut output_pixels = RowBitmapMut::new_contiguous(output_buf, image.width());
         Self::optionally_generate_dither_map(self.use_dither_map, image, true, &mut output_pixels, &mut self.palette)?;
         Ok(())
     }
 
     #[inline(never)]
-    pub(crate) fn write_remapped_image_rows_internal(&mut self, image: &mut Image, mut output_pixels: RowBitmapMut<'_, MaybeUninit<PalIndex>>) -> Result<(), Error> {
+    pub(crate) fn write_remapped_image_rows_internal(&mut self, image: &mut Image, mut output_pixels: RowBitmapMut<'_, MaybeUninit<PalIndexRemap>>) -> Result<(), Error> {
         let progress_stage1 = if self.use_dither_map != DitherMapMode::None { 20 } else { 0 };
         if self.remap_progress(progress_stage1 as f32 * 0.25) {
             return Err(Error::Aborted);
@@ -115,7 +115,7 @@ impl QuantizationResult {
         Ok(())
     }
 
-    fn optionally_generate_dither_map(use_dither_map: DitherMapMode, image: &mut Image<'_>, uses_background: bool, output_pixels: &mut RowBitmapMut<'_, MaybeUninit<PalIndex>>, palette: &mut PalF) -> Result<Option<f64>, Error> {
+    fn optionally_generate_dither_map(use_dither_map: DitherMapMode, image: &mut Image<'_>, uses_background: bool, output_pixels: &mut RowBitmapMut<'_, MaybeUninit<PalIndexRemap>>, palette: &mut PalF) -> Result<Option<f64>, Error> {
         let is_image_huge = (image.px.width * image.px.height) > 2000 * 2000;
         let allow_dither_map = use_dither_map == DitherMapMode::Always || (!is_image_huge && use_dither_map != DitherMapMode::None);
         let generate_dither_map = allow_dither_map && image.dither_map.is_none();
@@ -231,7 +231,7 @@ impl QuantizationResult {
     /// Remap image into a palette + indices.
     ///
     /// Returns the palette and a 1-byte-per-pixel uncompressed bitmap
-    pub fn remapped(&mut self, image: &mut Image<'_>) -> Result<(Vec<RGBA>, Vec<PalIndex>), Error> {
+    pub fn remapped(&mut self, image: &mut Image<'_>) -> Result<(Vec<RGBA>, Vec<PalIndexRemap>), Error> {
         let mut buf = Vec::new();
         let pal = self.remap_into_vec(image, &mut buf)?;
         Ok((pal, buf))
@@ -243,7 +243,7 @@ impl QuantizationResult {
     ///
     /// Returns the palette.
     #[inline]
-    pub fn remap_into_vec(&mut self, image: &mut Image<'_>, buf: &mut Vec<PalIndex>) -> Result<Vec<RGBA>, Error> {
+    pub fn remap_into_vec(&mut self, image: &mut Image<'_>, buf: &mut Vec<PalIndexRemap>) -> Result<Vec<RGBA>, Error> {
         let len = image.width() * image.height();
         // Capacity is essential here, as it creates uninitialized buffer
         unsafe {
@@ -264,7 +264,7 @@ impl QuantizationResult {
     /// You should call [`palette()`][Self::palette] _after_ this call, but not before it,
     /// because remapping refines the palette.
     #[inline]
-    pub fn remap_into(&mut self, image: &mut Image<'_>, output_buf: &mut [MaybeUninit<PalIndex>]) -> Result<(), Error> {
+    pub fn remap_into(&mut self, image: &mut Image<'_>, output_buf: &mut [MaybeUninit<PalIndexRemap>]) -> Result<(), Error> {
         let required_size = (image.width()) * (image.height());
         let output_buf = output_buf.get_mut(0..required_size).ok_or(BufferTooSmall)?;
 

src/remap.rs

@@ -2,7 +2,7 @@ use crate::error::Error;
 use crate::image::Image;
 use crate::kmeans::Kmeans;
 use crate::nearest::Nearest;
-use crate::pal::{f_pixel, PalF, PalIndex, Palette, ARGBF, LIQ_WEIGHT_MSE, MIN_OPAQUE_A};
+use crate::pal::{f_pixel, PalF, PalIndexRemap, Palette, ARGBF, LIQ_WEIGHT_MSE, MIN_OPAQUE_A};
 use crate::quant::QuantizationResult;
 use crate::rayoff::*;
 use crate::rows::{temp_buf, DynamicRows};
@@ -24,11 +24,11 @@ pub(crate) struct Remapped {
 }
 
 #[inline(never)]
-pub(crate) fn remap_to_palette<'x, 'b: 'x>(px: &mut DynamicRows, background: Option<&mut Image<'_>>, output_pixels: &'x mut RowBitmapMut<'b, MaybeUninit<PalIndex>>, palette: &mut PalF) -> Result<(f64, RowBitmap<'x, PalIndex>), Error> {
+pub(crate) fn remap_to_palette<'x, 'b: 'x>(px: &mut DynamicRows, background: Option<&mut Image<'_>>, output_pixels: &'x mut RowBitmapMut<'b, MaybeUninit<PalIndexRemap>>, palette: &mut PalF) -> Result<(f64, RowBitmap<'x, PalIndexRemap>), Error> {
     let n = Nearest::new(palette)?;
     let colors = palette.as_slice();
     let palette_len = colors.len();
-    if palette_len > u8::MAX as usize + 1 {
+    if palette_len > PalIndexRemap::MAX as usize + 1 {
         return Err(Error::Unsupported);
     }
 
@@ -41,14 +41,14 @@ pub(crate) fn remap_to_palette<'x, 'b: 'x>(px: &mut DynamicRows, background: Opt
 
     let input_rows = px.rows_iter(&mut tls_tmp.1)?;
     let (background, transparent_index) = background.map(|background| {
-        (Some(background), n.search(&f_pixel::default(), 0).0)
+        (Some(background), n.search(&f_pixel::default(), 0).0 as PalIndexRemap)
     })
     .filter(|&(_, transparent_index)| colors[usize::from(transparent_index)].a < MIN_OPAQUE_A)
     .unwrap_or((None, 0));
     let background = background.map(|bg| bg.px.rows_iter(&mut tls_tmp.1)).transpose()?;
 
     if background.is_some() {
-        tls_tmp.0.update_color(f_pixel::default(), 1., transparent_index);
+        tls_tmp.0.update_color(f_pixel::default(), 1., transparent_index as _);
     }
 
     drop(tls_tmp);
@@ -69,7 +69,8 @@ pub(crate) fn remap_to_palette<'x, 'b: 'x>(px: &mut DynamicRows, background: Opt
 
         let mut last_match = 0;
         for (col, (inp, out)) in row_pixels.iter().zip(output_pixels_row).enumerate() {
-            let (matched, diff) = n.search(inp, last_match);
+            let (matched, diff) = n.search(inp, last_match as _);
+            let matched = matched as PalIndexRemap;
             last_match = matched;
             if let Some(bg) = bg_pixels.get(col) {
                 let bg_diff = bg.diff(inp);
@@ -81,7 +82,7 @@ pub(crate) fn remap_to_palette<'x, 'b: 'x>(px: &mut DynamicRows, background: Opt
             }
             remapping_error += f64::from(diff);
             out.write(matched);
-            kmeans.update_color(*inp, 1., matched);
+            kmeans.update_color(*inp, 1., matched as _);
         }
         remapping_error
     })
@@ -142,7 +143,7 @@ fn get_dithered_pixel(dither_level: f32, max_dither_error: f32, thiserr: f_pixel
 ///
 ///  If `output_image_is_remapped` is true, only pixels noticeably changed by error diffusion will be written to output image.
 #[inline(never)]
-pub(crate) fn remap_to_palette_floyd(input_image: &mut Image, mut output_pixels: RowBitmapMut<'_, MaybeUninit<PalIndex>>, palette: &PalF, quant: &QuantizationResult, max_dither_error: f32, output_image_is_remapped: bool) -> Result<(), Error> {
+pub(crate) fn remap_to_palette_floyd(input_image: &mut Image, mut output_pixels: RowBitmapMut<'_, MaybeUninit<PalIndexRemap>>, palette: &PalF, quant: &QuantizationResult, max_dither_error: f32, output_image_is_remapped: bool) -> Result<(), Error> {
     let progress_stage1 = if quant.use_dither_map != DitherMapMode::None { 20 } else { 0 };
 
     let width = input_image.width();
@@ -158,16 +159,13 @@ pub(crate) fn remap_to_palette_floyd(input_image: &mut Image, mut output_pixels:
 
     let n = Nearest::new(palette)?;
     let palette = palette.as_slice();
-    if palette.len() > u8::MAX as usize + 1 {
-        return Err(Error::Unsupported);
-    }
 
     let mut background = input_image.background.as_mut().map(|bg| {
         bg.px.prepare_iter(&mut temp_row, true)?;
         Ok::<_, Error>(&bg.px)
     }).transpose()?;
 
-    let transparent_index = if background.is_some() { n.search(&f_pixel::default(), 0).0 } else { 0 };
+    let transparent_index = if background.is_some() { n.search(&f_pixel::default(), 0).0 as PalIndexRemap } else { 0 };
     if background.is_some() && palette[transparent_index as usize].a > MIN_OPAQUE_A {
         background = None;
     }
@@ -233,7 +231,7 @@ pub(crate) fn remap_to_palette_floyd(input_image: &mut Image, mut output_pixels:
 }
 
 #[inline(never)]
-fn dither_row(row_pixels: &[f_pixel], output_pixels_row: &mut [MaybeUninit<PalIndex>], width: u32, dither_map: &[u8], base_dithering_level: f32, max_dither_error: f32, n: &Nearest, palette: &[f_pixel], transparent_index: PalIndex, bg_pixels: &[f_pixel], guess_from_remapped_pixels: bool, diffusion: &mut [f_pixel], even_row: bool) {
+fn dither_row(row_pixels: &[f_pixel], output_pixels_row: &mut [MaybeUninit<PalIndexRemap>], width: u32, dither_map: &[u8], base_dithering_level: f32, max_dither_error: f32, n: &Nearest, palette: &[f_pixel], transparent_index: PalIndexRemap, bg_pixels: &[f_pixel], guess_from_remapped_pixels: bool, diffusion: &mut [f_pixel], even_row: bool) {
     let width = width as usize;
     assert_eq!(row_pixels.len(), width);
     assert_eq!(output_pixels_row.len(), width);
@@ -265,8 +263,9 @@ fn dither_row(row_pixels: &[f_pixel], output_pixels_row: &mut [MaybeUninit<PalIn
         } else {
             last_match
         };
-        let (mut matched, dither_diff) = n.search(&spx, guessed_match);
-        last_match = matched;
+        let (matched, dither_diff) = n.search(&spx, guessed_match as _);
+        let mut matched = matched  as PalIndexRemap;
+        last_match = matched as PalIndexRemap;
         let mut output_px = palette[last_match as usize];
         if let Some(bg_pixel) = bg_pixels.get(col) {
             // if the background makes better match *with* dithering, it's a definitive win