From fbf18cc16d96c350eb73b115018f19272f96e8d4 Mon Sep 17 00:00:00 2001
From: kralverde <80051564+kralverde@users.noreply.github.com>
Date: Fri, 18 Apr 2025 06:26:34 -1000
Subject: [PATCH] Harden palette (de)serialization (#733)

* initial_commit

* double check palette data
---
 pumpkin-util/src/lib.rs            |  6 ++++-
 pumpkin-world/src/chunk/palette.rs | 35 ++++++++++++++++--------------
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/pumpkin-util/src/lib.rs b/pumpkin-util/src/lib.rs
index 935c5635c..5ea90502b 100644
--- a/pumpkin-util/src/lib.rs
+++ b/pumpkin-util/src/lib.rs
@@ -43,7 +43,11 @@ macro_rules! read_data_from_file {
 /// The minimum number of bits required to represent this number
 #[inline]
 pub fn encompassing_bits(count: usize) -> u8 {
-    count.ilog2() as u8 + if count.is_power_of_two() { 0 } else { 1 }
+    if count == 1 {
+        1
+    } else {
+        count.ilog2() as u8 + if count.is_power_of_two() { 0 } else { 1 }
+    }
 }
 
 #[derive(PartialEq, Serialize, Deserialize, Clone)]
diff --git a/pumpkin-world/src/chunk/palette.rs b/pumpkin-world/src/chunk/palette.rs
index e45bcc9b5..23b2fd0d0 100644
--- a/pumpkin-world/src/chunk/palette.rs
+++ b/pumpkin-world/src/chunk/palette.rs
@@ -26,19 +26,6 @@ pub struct HeterogeneousPaletteData<V: Hash + Eq + Copy, const DIM: usize> {
 }
 
 impl<V: Hash + Eq + Copy, const DIM: usize> HeterogeneousPaletteData<V, DIM> {
-    fn from_cube(cube: Box<AbstractCube<V, DIM>>) -> Self {
-        let counts =
-            cube.as_flattened()
-                .as_flattened()
-                .iter()
-                .fold(HashMap::new(), |mut acc, key| {
-                    acc.entry(*key).and_modify(|count| *count += 1).or_insert(1);
-                    acc
-                });
-
-        Self { cube, counts }
-    }
-
     fn get(&self, x: usize, y: usize, z: usize) -> V {
         debug_assert!(x < DIM);
         debug_assert!(y < DIM);
@@ -81,6 +68,23 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
     pub const SIZE: usize = DIM;
     pub const VOLUME: usize = DIM * DIM * DIM;
 
+    fn from_cube(cube: Box<AbstractCube<V, DIM>>) -> Self {
+        let counts =
+            cube.as_flattened()
+                .as_flattened()
+                .iter()
+                .fold(HashMap::new(), |mut acc, key| {
+                    acc.entry(*key).and_modify(|count| *count += 1).or_insert(1);
+                    acc
+                });
+
+        if counts.len() == 1 {
+            Self::Homogeneous(*counts.keys().next().unwrap())
+        } else {
+            Self::Heterogeneous(Box::new(HeterogeneousPaletteData { cube, counts }))
+        }
+    }
+
     fn bits_per_entry(&self) -> u8 {
         match self {
             Self::Homogeneous(_) => 0,
@@ -181,7 +185,7 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
                     });
                 });
 
-            Self::Heterogeneous(Box::new(HeterogeneousPaletteData::from_cube(cube)))
+            Self::from_cube(cube)
         }
     }
 
@@ -202,8 +206,7 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
                 if value != *original {
                     let mut cube = Box::new([[[*original; DIM]; DIM]; DIM]);
                     cube[y][z][x] = value;
-                    let data = HeterogeneousPaletteData::from_cube(cube);
-                    *self = Self::Heterogeneous(Box::new(data));
+                    *self = Self::from_cube(cube);
                 }
             }
             Self::Heterogeneous(data) => {