From 6d56e72d662a97334ed4573fab736755d4715ad4 Mon Sep 17 00:00:00 2001
From: Gwilym Inzani <email@gwilym.dev>
Date: Tue, 29 Aug 2023 14:39:41 +0100
Subject: [PATCH] Fix palette assignment

---
 agb-image-converter/src/lib.rs            | 12 ++++++++--
 agb-image-converter/src/rust_generator.rs | 27 +++++++++++++++++------
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/agb-image-converter/src/lib.rs b/agb-image-converter/src/lib.rs
index 16f19adb..d5fa34e4 100644
--- a/agb-image-converter/src/lib.rs
+++ b/agb-image-converter/src/lib.rs
@@ -494,7 +494,14 @@ fn palette_tile_data(
     let mut tile_data = Vec::new();
 
     for (image_idx, image) in images.iter().enumerate() {
-        add_image_to_tile_data(&mut tile_data, image, optimiser, image_idx, true)
+        add_image_to_tile_data(
+            &mut tile_data,
+            image,
+            optimiser,
+            image_idx,
+            true,
+            &(0..images.len()).collect::<Vec<_>>(),
+        );
     }
 
     let tile_data = collapse_to_4bpp(&tile_data);
@@ -517,6 +524,7 @@ fn add_image_to_tile_data(
     optimiser: &Palette16OptimisationResults,
     assignment_offset: usize,
     is_sprite: bool,
+    remap_index: &[usize],
 ) {
     let tile_size = 8;
     let tiles_x = image.width / tile_size;
@@ -527,7 +535,7 @@ fn add_image_to_tile_data(
             let assignment = if is_sprite {
                 assignment_offset
             } else {
-                y * tiles_x + x + assignment_offset
+                remap_index[y * tiles_x + x] + assignment_offset
             };
 
             let palette_index = optimiser.assignments[assignment];
diff --git a/agb-image-converter/src/rust_generator.rs b/agb-image-converter/src/rust_generator.rs
index aad0b001..bd6f243c 100644
--- a/agb-image-converter/src/rust_generator.rs
+++ b/agb-image-converter/src/rust_generator.rs
@@ -6,6 +6,7 @@ use crate::{image_loader::Image, ByteString};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 
+use std::collections::BTreeMap;
 use std::iter;
 
 pub(crate) fn generate_palette_code(
@@ -63,21 +64,33 @@ pub(crate) fn generate_code(
         )
     };
 
+    let remap_index = dedup_data
+        .iter()
+        .enumerate()
+        .map(|(i, data)| (data.new_index, i))
+        .collect::<BTreeMap<_, _>>(); // BTreeMap so that values below is in order
+
+    let remap_index = remap_index.values().cloned().collect::<Vec<_>>();
+
     let (tile_data, assignments) = if let Some(assignment_offset) = assignment_offset {
         let mut tile_data = Vec::new();
 
-        add_image_to_tile_data(&mut tile_data, &image, results, assignment_offset, false);
+        add_image_to_tile_data(
+            &mut tile_data,
+            &image,
+            results,
+            assignment_offset,
+            false,
+            &remap_index,
+        );
 
         let tile_data = collapse_to_4bpp(&tile_data);
 
         let num_tiles = image.width * image.height / 8usize.pow(2);
 
-        let assignments = results
-            .assignments
-            .iter()
-            .skip(assignment_offset)
-            .take(num_tiles)
-            .map(|&x| x as u8)
+        let all_assignments = &results.assignments[assignment_offset..];
+        let assignments = (0..num_tiles)
+            .map(|tile_id| all_assignments[remap_index[tile_id]] as u8)
             .collect();
 
         (tile_data, assignments)