From 3763c29b4dfb98dcda491ecc78914a00b3ea06dd Mon Sep 17 00:00:00 2001
From: Matt Tracy <matt.r.tracy@gmail.com>
Date: Tue, 6 Feb 2024 16:40:20 -0800
Subject: [PATCH] Add Support for IMAP node.

The 'IMAP' node is used to map voxel indexes (the 'i' value on each
Voxel structure) to the color palette in a way that is *not* simply
value-to-value.

This structure has the following behavior if it appears in a file:
https://github.com/ephtracy/voxel-model/issues/19#issuecomment-739324018

When importing a file which contains this, imported models will not
match the expected values in the palette unless the IMAP is respected.
---
 src/dot_vox_data.rs |  2 ++
 src/lib.rs          |  1 +
 src/palette.rs      |  5 +++++
 src/parser.rs       | 22 ++++++++++++++++++++--
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/dot_vox_data.rs b/src/dot_vox_data.rs
index 7a64c08..22e0e63 100644
--- a/src/dot_vox_data.rs
+++ b/src/dot_vox_data.rs
@@ -10,6 +10,8 @@ pub struct DotVoxData {
     pub models: Vec<Model>,
     /// A `Vec` containing the colour palette as 32-bit integers
     pub palette: Vec<Color>,
+    /// Map from voxel values to their *true* palette value.
+    pub palette_map: Vec<u8>,
     /// A `Vec` containing all the [`Material`]s set.
     pub materials: Vec<Material>,
     /// Scene. The first node in this list is always the root node.
diff --git a/src/lib.rs b/src/lib.rs
index 2db72a0..0c3747f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -324,6 +324,7 @@ mod tests {
                 ],
             }],
             palette,
+            palette_map: vec![],
             materials,
             scenes,
             layers,
diff --git a/src/palette.rs b/src/palette.rs
index 882ea82..04f0432 100644
--- a/src/palette.rs
+++ b/src/palette.rs
@@ -20,6 +20,7 @@ fn parse_color(input: &[u8]) -> IResult<&[u8], Color> {
     Ok((input, Color { r, g, b, a }))
 }
 
+
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct Color {
     pub r: u8,
@@ -38,3 +39,7 @@ impl From<&Color> for [u8; 4] {
         [color.r, color.g, color.b, color.a]
     }
 }
+
+pub fn parse_imap(input: &[u8]) -> IResult<&[u8], Vec<u8>> {
+    all_consuming(many0(le_u8))(input)
+}
\ No newline at end of file
diff --git a/src/parser.rs b/src/parser.rs
index 1188fdb..5d167fe 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -11,7 +11,7 @@ use nom::{
     sequence::pair,
     IResult,
 };
-use std::{mem::size_of, str, str::Utf8Error};
+use std::{mem::size_of, str::{self, Utf8Error}};
 
 #[cfg(feature = "ahash")]
 use ahash::AHashMap as HashMap;
@@ -32,6 +32,7 @@ pub enum Chunk {
     GroupNode(SceneGroup),
     ShapeNode(SceneShape),
     Layer(RawLayer),
+    Imap(Vec<usize>),
     Unknown(String),
     Invalid(Vec<u8>),
 }
@@ -205,6 +206,7 @@ fn map_chunk_to_data(version: u32, main: Chunk) -> DotVoxData {
             let mut materials: Vec<Material> = vec![];
             let mut scene: Vec<SceneNode> = vec![];
             let mut layers: Vec<Layer> = Vec::new();
+            let mut palette_map: Option<Vec<u8>> = None;
 
             for chunk in children {
                 match chunk {
@@ -241,11 +243,17 @@ fn map_chunk_to_data(version: u32, main: Chunk) -> DotVoxData {
                                 layer.id
                             );
                         }
-
                         layers.push(Layer {
                             attributes: layer.attributes,
                         });
                     }
+                    Chunk::Imap(imap) => {
+                        let mut remap: Vec<u8> = vec![0;256];
+                        imap.iter().enumerate().for_each(|(idx, val)|{
+                            remap[*val] = idx as u8;
+                        });
+                        palette_map = Some(remap);
+                    }
                     _ => debug!("Unmapped chunk {:?}", chunk),
                 }
             }
@@ -254,6 +262,7 @@ fn map_chunk_to_data(version: u32, main: Chunk) -> DotVoxData {
                 version,
                 models,
                 palette: palette_holder,
+                palette_map: palette_map.unwrap_or((0u8..=255u8).collect()),
                 materials,
                 scenes: scene,
                 layers,
@@ -263,6 +272,7 @@ fn map_chunk_to_data(version: u32, main: Chunk) -> DotVoxData {
             version,
             models: vec![],
             palette: vec![],
+            palette_map: vec![],
             materials: vec![],
             scenes: vec![],
             layers: vec![],
@@ -290,6 +300,7 @@ fn build_chunk(id: &str, chunk_content: &[u8], children_size: u32, child_content
             "nGRP" => build_scene_group_chunk(chunk_content),
             "nSHP" => build_scene_shape_chunk(chunk_content),
             "LAYR" => build_layer_chunk(chunk_content),
+            "IMAP" => build_imap_chunk(chunk_content),
             _ => {
                 debug!("Unknown childless chunk {:?}", id);
                 Chunk::Unknown(id.to_owned())
@@ -370,6 +381,13 @@ fn build_layer_chunk(chunk_content: &[u8]) -> Chunk {
     }
 }
 
+fn build_imap_chunk(chunk_content: &[u8]) -> Chunk {
+    match palette::parse_imap(chunk_content) {
+        Ok((_, result)) => Chunk::Imap(result.into_iter().map(Into::<usize>::into).collect()),
+        _ => Chunk::Invalid(chunk_content.to_vec())
+    }
+}
+
 pub fn parse_material(i: &[u8]) -> IResult<&[u8], Material> {
     let (i, (id, properties)) = pair(le_u32, parse_dict)(i)?;
     Ok((i, Material { id, properties }))