Improve palette optimiser (#768)

Make the whole thing more consistent and better at creating optimal
palettes.

- [x] Changelog updated / no changelog update needed

CHANGELOG.md

@@ -11,6 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added support for s3m and mod format files to `agb-tracker`.
 
+### Changed
+
+- Changed how 16 colour palettes are optimised to give better average case results. You should find that
+  either your palettes will always import, or never import correctly.
+
 ## [0.21.0] - 2024/09/24
 
 ### Added

agb-image-converter/Cargo.toml

@@ -20,3 +20,7 @@ proc-macro2 = "1"
 quote = "1"
 asefile = "0.3.8"
 fontdue = "0.9"
+pagination-packing = "2.1.0"
+
+[dev-dependencies]
+quickcheck = "1"

agb-image-converter/src/colour.rs

@@ -1,13 +1,25 @@
-use std::str::FromStr;
+use std::{fmt, str::FromStr};
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct Colour {
     pub r: u8,
     pub g: u8,
     pub b: u8,
     pub a: u8,
 }
 
+impl fmt::Debug for Colour {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "#{:02x}{:02x}{:02x}", self.r, self.g, self.b)?;
+
+        if self.a != 0xff {
+            write!(f, "{:02x}", self.a)?;
+        }
+
+        Ok(())
+    }
+}
+
 impl Colour {
     pub fn from_rgb(r: u8, g: u8, b: u8, a: u8) -> Self {
         Colour { r, g, b, a }
@@ -38,3 +50,26 @@ impl FromStr for Colour {
         Ok(Colour::from_rgb(r, g, b, 255))
     }
 }
+
+#[cfg(test)]
+impl quickcheck::Arbitrary for Colour {
+    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+        Self::from_rgb(
+            quickcheck::Arbitrary::arbitrary(g),
+            quickcheck::Arbitrary::arbitrary(g),
+            quickcheck::Arbitrary::arbitrary(g),
+            quickcheck::Arbitrary::arbitrary(g),
+        )
+    }
+
+    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
+        Box::new(
+            vec![
+                Colour::from_rgb(0, 0, 0, 0),
+                Colour::from_rgb(self.r, self.g, self.b, 0),
+                *self,
+            ]
+            .into_iter(),
+        )
+    }
+}

agb-image-converter/src/deduplicator.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, hash::BuildHasher};
+use std::collections::BTreeMap;
 
 use crate::{colour::Colour, image_loader::Image};
 
@@ -42,7 +42,7 @@ pub struct DeduplicatedData {
     pub transformation: Transformation,
 }
 
-#[derive(Clone, PartialEq, Eq, Hash)]
+#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 struct Tile {
     data: [Colour; 64],
 }
@@ -99,32 +99,21 @@ pub(crate) fn deduplicate_image(input: &Image, can_flip: bool) -> (Image, Vec<De
     let mut deduplication_data = vec![];
 
     let all_tiles = Tile::split_image(input);
-    let mut existing_tiles = HashMap::new();
-
-    let hasher = std::collections::hash_map::RandomState::new();
+    let mut existing_tiles = BTreeMap::new();
 
     for tile in all_tiles {
         let (tile, transformation) = if can_flip {
             let vflipped = tile.vflipped();
             let hflipped = tile.hflipped();
             let vhflipped = vflipped.hflipped();
 
-            // find the one with the smallest hash
-            let tile_hash = hasher.hash_one(&tile);
-            let vflipped_hash = hasher.hash_one(&vflipped);
-            let hflipped_hash = hasher.hash_one(&hflipped);
-            let vhflipped_hash = hasher.hash_one(&vhflipped);
-
-            let minimum = tile_hash
-                .min(vflipped_hash)
-                .min(hflipped_hash)
-                .min(vhflipped_hash);
+            let minimum = (&tile).min(&vflipped).min(&hflipped).min(&vhflipped);
 
-            if minimum == tile_hash {
+            if minimum == &tile {
                 (tile, Transformation::none())
-            } else if minimum == vflipped_hash {
+            } else if minimum == &vflipped {
                 (vflipped, Transformation::vflipped())
-            } else if minimum == hflipped_hash {
+            } else if minimum == &hflipped {
                 (hflipped, Transformation::hflipped())
             } else {
                 (vhflipped, Transformation::vhflipped())

agb-image-converter/src/lib.rs

@@ -244,7 +244,9 @@ fn include_gfx_from_config(
         }
     }
 
-    let optimisation_results = optimiser.optimise_palettes();
+    let optimisation_results = optimiser
+        .optimise_palettes()
+        .expect("Failed to optimised palettes");
     let optimisation_results = palette256.extend_results(&optimisation_results);
 
     let mut image_code = vec![];
@@ -377,7 +379,9 @@ pub fn include_aseprite_inner(input: TokenStream) -> TokenStream {
         }
     }
 
-    let optimised_results = optimiser.optimise_palettes();
+    let optimised_results = optimiser
+        .optimise_palettes()
+        .expect("Failed to optimise palettes");
 
     let (palette_data, tile_data, assignments) = palette_tile_data(&optimised_results, &images);
 

agb-image-converter/src/palette16.rs

@@ -1,10 +1,13 @@
 use crate::colour::Colour;
-use std::collections::HashSet;
+use std::{
+    collections::{BTreeSet, HashSet},
+    fmt,
+};
 
 const MAX_COLOURS: usize = 256;
 const MAX_COLOURS_PER_PALETTE: usize = 16;
 
-#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Eq, PartialEq, Hash, PartialOrd, Ord)]
 pub(crate) struct Palette16 {
     colours: Vec<Colour>,
 }
@@ -62,12 +65,15 @@ impl Palette16 {
         self.colours.iter()
     }
 
-    fn union_length(&self, other: &Palette16) -> usize {
-        self.colours
+    fn with_transparent(&self, transparent_colour: Colour) -> Self {
+        let mut new_colours = self.colours.clone();
+        let transparent_colour_index = new_colours
             .iter()
-            .chain(&other.colours)
-            .collect::<HashSet<_>>()
-            .len()
+            .position(|&c| c == transparent_colour)
+            .expect("Could not find tranparent colour in palette");
+        new_colours.swap(0, transparent_colour_index);
+
+        Self::from(&new_colours)
     }
 
     fn is_satisfied_by(&self, other: &Palette16) -> bool {
@@ -87,6 +93,20 @@ impl IntoIterator for Palette16 {
     }
 }
 
+impl<'a, T> From<T> for Palette16
+where
+    T: IntoIterator<Item = &'a Colour>,
+{
+    fn from(value: T) -> Self {
+        let mut palette = Palette16::new();
+        for colour in value.into_iter() {
+            palette.add_colour(*colour);
+        }
+
+        palette
+    }
+}
+
 pub(crate) struct Palette16Optimiser {
     palettes: Vec<Palette16>,
     colours: Vec<Colour>,
@@ -125,88 +145,129 @@ impl Palette16Optimiser {
         }
     }
 
-    pub fn optimise_palettes(&self) -> Palette16OptimisationResults {
-        let mut assignments = vec![0; self.palettes.len()];
-        let mut optimised_palettes = vec![];
+    pub fn optimise_palettes(&self) -> Result<Palette16OptimisationResults, DoesNotFitError> {
+        let transparent_colour = self
+            .transparent_colour
+            .unwrap_or_else(|| Colour::from_rgb(255, 0, 255, 0));
 
-        let mut unsatisfied_palettes = self
+        let palettes_to_optimise = self
             .palettes
             .iter()
             .cloned()
-            .collect::<HashSet<Palette16>>();
-
-        while !unsatisfied_palettes.is_empty() {
-            let palette = self.find_maximal_palette_for(&unsatisfied_palettes);
-
-            unsatisfied_palettes.retain(|test_palette| !test_palette.is_satisfied_by(&palette));
-
-            optimised_palettes.push(palette);
-
-            if optimised_palettes.len() == MAX_COLOURS / MAX_COLOURS_PER_PALETTE {
-                panic!("Failed to find covering palettes");
-            }
+            .map(|mut palette| {
+                // ensure each palette we're creating the covering for has the transparent colour in it
+                palette.add_colour(transparent_colour);
+                palette
+            })
+            .collect::<BTreeSet<Palette16>>()
+            .into_iter()
+            .map(|palette| palette.colours)
+            .collect::<Vec<_>>();
+
+        let packed_palettes =
+            pagination_packing::overload_and_remove::<_, _, Vec<_>>(&palettes_to_optimise, 16);
+
+        let optimised_palettes = packed_palettes
+            .iter()
+            .map(|packed_palette| {
+                let colours = packed_palette.unique_symbols(&palettes_to_optimise);
+                Palette16::from(colours).with_transparent(transparent_colour)
+            })
+            .collect::<Vec<_>>();
+
+        if optimised_palettes.len() > 16 {
+            return Err(DoesNotFitError(packed_palettes.len()));
         }
 
+        let mut assignments = vec![0; self.palettes.len()];
+
         for (i, overall_palette) in self.palettes.iter().enumerate() {
             assignments[i] = optimised_palettes
                 .iter()
                 .position(|palette| overall_palette.is_satisfied_by(palette))
                 .unwrap();
         }
 
-        Palette16OptimisationResults {
+        Ok(Palette16OptimisationResults {
             optimised_palettes,
             assignments,
             transparent_colour: self.transparent_colour,
-        }
+        })
     }
+}
 
-    fn find_maximal_palette_for(&self, unsatisfied_palettes: &HashSet<Palette16>) -> Palette16 {
-        let mut palette = Palette16::new();
+pub struct DoesNotFitError(pub usize);
 
-        palette.add_colour(
-            self.transparent_colour
-                .unwrap_or_else(|| Colour::from_rgb(255, 0, 255, 0)),
-        );
-
-        loop {
-            let mut colour_usage = vec![0; MAX_COLOURS];
-            let mut a_colour_is_used = false;
-
-            for current_palette in unsatisfied_palettes {
-                if palette.union_length(current_palette) > MAX_COLOURS_PER_PALETTE {
-                    continue;
-                }
-
-                for colour in &current_palette.colours {
-                    if palette.colours.contains(colour) {
-                        continue;
-                    }
-
-                    if let Some(colour_index) = self.colours.iter().position(|c| c == colour) {
-                        colour_usage[colour_index] += 1;
-                        a_colour_is_used = true;
-                    }
-                }
+impl fmt::Display for DoesNotFitError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Could not fit colours into palette, needed {} bins but can have at most 16",
+            self.0
+        )
+    }
+}
+
+impl fmt::Debug for DoesNotFitError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{self}")
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use quickcheck::{quickcheck, Arbitrary};
+
+    use super::*;
+
+    quickcheck! {
+        fn less_than_256_colours_always_fits(palettes: Vec<Palette16>, transparent_colour: Colour) -> bool {
+            let mut optimiser = Palette16Optimiser::new(Some(transparent_colour));
+            for palette in palettes.clone().into_iter().take(16) {
+                optimiser.add_palette(palette);
+            }
+
+            let Ok(optimisation_results) = optimiser.optimise_palettes() else {
+                return false
+            };
+
+            check_palette_invariants(palettes.iter().take(16), optimisation_results, transparent_colour)
+        }
+    }
+
+    fn check_palette_invariants<'a>(
+        palettes: impl Iterator<Item = &'a Palette16>,
+        optimisation_results: Palette16OptimisationResults,
+        transparent_colour: Colour,
+    ) -> bool {
+        for (i, palette) in palettes.enumerate() {
+            let optimised_palette =
+                &optimisation_results.optimised_palettes[optimisation_results.assignments[i]];
+            if !palette.is_satisfied_by(optimised_palette) {
+                return false;
             }
 
-            if !a_colour_is_used {
-                return palette;
+            if optimised_palette.colour_index(transparent_colour) != 0 {
+                return false;
             }
+        }
 
-            let best_index = colour_usage
-                .iter()
-                .enumerate()
-                .max_by(|(_, usage1), (_, usage2)| usage1.cmp(usage2))
-                .unwrap()
-                .0;
+        true
+    }
 
-            let best_colour = self.colours[best_index];
+    impl Arbitrary for Palette16 {
+        fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+            let mut palette = Palette16::new();
 
-            palette.add_colour(best_colour);
-            if palette.colours.len() == MAX_COLOURS_PER_PALETTE {
-                return palette;
+            let size: usize = Arbitrary::arbitrary(g);
+            // never entirely fill the palette, will give at most 15 colours
+            let size = size.rem_euclid(16);
+
+            for _ in 0..size {
+                palette.add_colour(Arbitrary::arbitrary(g));
             }
+
+            palette
         }
     }
 }