fix: entities in air

hopefully lol

pumpkin-protocol/src/java/client/play/chunk_data.rs

@@ -48,9 +48,15 @@ impl ClientPacket for CChunkData<'_> {
                 Ok(())
             };
 
-            write_heightmap(1, &heightmaps.world_surface)?;
+            write_heightmap(1, heightmaps.world_surface.as_deref().unwrap_or(&[0; 37]))?;
-            write_heightmap(4, &heightmaps.motion_blocking)?;
+            write_heightmap(4, heightmaps.motion_blocking.as_deref().unwrap_or(&[0; 37]))?;
-            write_heightmap(5, &heightmaps.motion_blocking_no_leaves)?;
+            write_heightmap(
+                5,
+                heightmaps
+                    .motion_blocking_no_leaves
+                    .as_deref()
+                    .unwrap_or(&[0; 37]),
+            )?;
         }
         drop(heightmaps);
 

pumpkin-world/src/chunk/mod.rs

@@ -104,7 +104,7 @@ pub struct ChunkEntityData {
 pub struct ChunkSections {
     pub count: usize,
     pub block_sections: RwLock<Box<[BlockPalette]>>,
-    pub random_tick_sections: RwLock<Box<[RandomTickSectionCache]>>,
+    pub random_tick_sections: RwLock<Option<Box<[RandomTickSectionCache]>>>,
     pub randomly_ticking_mask: std::sync::atomic::AtomicU32,
     pub biome_sections: RwLock<Box<[BiomePalette]>>,
     pub min_y: i32,
@@ -189,12 +189,21 @@ impl ChunkHeightmapType {
 #[derive(Deserialize, Serialize, Debug, Clone)]
 #[serde(rename_all = "UPPERCASE")]
 pub struct ChunkHeightmaps {
-    #[serde(serialize_with = "nbt_long_array")]
+    #[serde(
-    pub world_surface: Box<[i64]>,
+        serialize_with = "nbt_long_array",
-    #[serde(serialize_with = "nbt_long_array")]
+        skip_serializing_if = "Option::is_none"
-    pub motion_blocking: Box<[i64]>,
+    )]
-    #[serde(serialize_with = "nbt_long_array")]
+    pub world_surface: Option<Box<[i64]>>,
-    pub motion_blocking_no_leaves: Box<[i64]>,
+    #[serde(
+        serialize_with = "nbt_long_array",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub motion_blocking: Option<Box<[i64]>>,
+    #[serde(
+        serialize_with = "nbt_long_array",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub motion_blocking_no_leaves: Option<Box<[i64]>>,
 }
 
 impl ChunkHeightmaps {
@@ -205,6 +214,8 @@ impl ChunkHeightmaps {
             ChunkHeightmapType::MotionBlockingNoLeaves => &mut self.motion_blocking_no_leaves,
         };
 
+        let data = data.get_or_insert_with(|| vec![0; 37].into_boxed_slice());
+
         let local_x = (x & 15) as usize;
         let local_z = (z & 15) as usize;
         let column_idx = local_z * 16 + local_x;
@@ -232,6 +243,10 @@ impl ChunkHeightmaps {
             ChunkHeightmapType::MotionBlockingNoLeaves => &self.motion_blocking_no_leaves,
         };
 
+        let Some(data) = data else {
+            return min_y - 1;
+        };
+
         let local_x = (x & 15) as usize;
         let local_z = (z & 15) as usize;
         let column_idx = local_z * 16 + local_x;
@@ -289,11 +304,9 @@ impl ChunkHeightmaps {
 impl Default for ChunkHeightmaps {
     fn default() -> Self {
         Self {
-            // 9 bits per entry
+            motion_blocking: None,
-            // 0 packed into an i64 7 times.
+            motion_blocking_no_leaves: None,
-            motion_blocking: vec![0; 37].into_boxed_slice(),
+            world_surface: None,
-            motion_blocking_no_leaves: vec![0; 37].into_boxed_slice(),
-            world_surface: vec![0; 37].into_boxed_slice(),
         }
     }
 }
@@ -302,8 +315,9 @@ impl ChunkSections {
     #[must_use]
     pub fn build_random_tick_sections_cache(
         block_sections: &[BlockPalette],
-    ) -> (Box<[RandomTickSectionCache]>, u32) {
+    ) -> (Option<Box<[RandomTickSectionCache]>>, u32) {
         let mut mask = 0;
+        let mut has_ticks = false;
         let cache = block_sections
             .iter()
             .enumerate()
@@ -312,6 +326,7 @@ impl ChunkSections {
                     section.random_ticking_counts();
                 if random_ticking_block_count > 0 || random_ticking_fluid_count > 0 {
                     mask |= 1 << i;
+                    has_ticks = true;
                 }
                 RandomTickSectionCache {
                     random_ticking_block_count,
@@ -320,7 +335,12 @@ impl ChunkSections {
             })
             .collect::<Vec<_>>()
             .into_boxed_slice();
-        (cache, mask)
+
+        if has_ticks {
+            (Some(cache), mask)
+        } else {
+            (None, 0)
+        }
     }
 
     #[must_use]
@@ -443,18 +463,25 @@ impl ChunkSections {
 
         // Keep lock order consistent to avoid deadlocks: block sections first, then random-tick cache.
         let mut sections = self.block_sections.write().unwrap();
-        let mut random_tick_sections = self.random_tick_sections.write().unwrap();
+        let mut random_tick_sections_guard = self.random_tick_sections.write().unwrap();
 
-        if let (Some(section), Some(random_tick_cache)) = (
+        if let Some(section) = sections.get_mut(section_index) {
-            sections.get_mut(section_index),
-            random_tick_sections.get_mut(section_index),
-        ) {
             let replaced_block_state_id =
                 section.set(relative_x, relative_y, relative_z, block_state_id);
             if replaced_block_state_id == block_state_id {
                 return replaced_block_state_id;
             }
 
+            if (has_random_ticks(block_state_id) || has_random_ticking_fluid(block_state_id))
+                && random_tick_sections_guard.is_none()
+            {
+                let new_cache =
+                    vec![RandomTickSectionCache::default(); self.count].into_boxed_slice();
+                *random_tick_sections_guard = Some(new_cache);
+            }
+
+            if let Some(random_tick_sections) = random_tick_sections_guard.as_mut() {
+                let random_tick_cache = &mut random_tick_sections[section_index];
                 if has_random_ticks(replaced_block_state_id) {
                     random_tick_cache.random_ticking_block_count = random_tick_cache
                         .random_ticking_block_count
@@ -489,6 +516,9 @@ impl ChunkSections {
                 self.randomly_ticking_mask
                     .store(mask, std::sync::atomic::Ordering::Relaxed);
 
+                // If no more ticking sections, we could potentially deallocate, but that might be overkill and cause jitter.
+            }
+
             return replaced_block_state_id;
         }
         0
@@ -754,6 +784,7 @@ mod tests {
         sections[1].set(0, 0, 0, Block::LAVA.default_state.id);
 
         let (cache, _mask) = ChunkSections::build_random_tick_sections_cache(&sections);
+        let cache = cache.unwrap();
         assert!(!cache[0].is_randomly_ticking());
         assert!(cache[1].random_ticking_fluid_count > 0);
         assert!(cache[1].is_randomly_ticking());
@@ -765,7 +796,12 @@ mod tests {
         let sections = ChunkSections::new(1, min_y);
 
         assert!(
-            !sections.random_tick_sections.read().unwrap()[0].is_randomly_ticking(),
+            sections
+                .random_tick_sections
+                .read()
+                .unwrap()
+                .as_ref()
+                .is_none_or(|c| !c[0].is_randomly_ticking()),
             "fresh sections should not be randomly ticking"
         );
 
@@ -778,6 +814,7 @@ mod tests {
         sections.set_block_absolute_y(0, min_y, 0, random_block_state);
         {
             let cache = sections.random_tick_sections.read().unwrap();
+            let cache = cache.as_ref().unwrap();
             assert_eq!(cache[0].random_ticking_block_count, 1);
             assert_eq!(cache[0].random_ticking_fluid_count, 0);
             assert!(cache[0].is_randomly_ticking());
@@ -786,6 +823,7 @@ mod tests {
         sections.set_block_absolute_y(0, min_y, 0, Block::STONE.default_state.id);
         {
             let cache = sections.random_tick_sections.read().unwrap();
+            let cache = cache.as_ref().unwrap();
             assert_eq!(cache[0].random_ticking_block_count, 0);
             assert_eq!(cache[0].random_ticking_fluid_count, 0);
             assert!(!cache[0].is_randomly_ticking());
@@ -794,6 +832,7 @@ mod tests {
         sections.set_block_absolute_y(0, min_y, 0, Block::LAVA.default_state.id);
         {
             let cache = sections.random_tick_sections.read().unwrap();
+            let cache = cache.as_ref().unwrap();
             assert!(cache[0].random_ticking_fluid_count > 0);
             assert!(cache[0].is_randomly_ticking());
         }

pumpkin-world/src/chunk/palette.rs

@@ -22,18 +22,27 @@ pub fn has_random_ticking_fluid(state_id: u16) -> bool {
 
 #[derive(Clone)]
 pub struct HeterogeneousPaletteData<V: Hash + Eq + Copy, const DIM: usize> {
-    cube: Box<AbstractCube<V, DIM>>,
+    storage: PaletteStorage<V, DIM>,
     palette: Vec<V>,
     counts: Vec<u16>,
 }
 
-impl<V: Hash + Eq + Copy, const DIM: usize> HeterogeneousPaletteData<V, DIM> {
+#[derive(Clone)]
+enum PaletteStorage<V, const DIM: usize> {
+    Dense(Box<AbstractCube<V, DIM>>),
+    Indexed(Box<AbstractCube<u8, DIM>>),
+}
+
+impl<V: Hash + Eq + Copy + Default, const DIM: usize> HeterogeneousPaletteData<V, DIM> {
     fn get(&self, x: usize, y: usize, z: usize) -> V {
         debug_assert!(x < DIM);
         debug_assert!(y < DIM);
         debug_assert!(z < DIM);
 
-        self.cube[y][z][x]
+        match &self.storage {
+            PaletteStorage::Dense(cube) => cube[y][z][x],
+            PaletteStorage::Indexed(indices) => self.palette[indices[y][z][x] as usize],
+        }
     }
 
     /// Returns the Original
@@ -42,25 +51,73 @@ impl<V: Hash + Eq + Copy, const DIM: usize> HeterogeneousPaletteData<V, DIM> {
         debug_assert!(y < DIM);
         debug_assert!(z < DIM);
 
-        let original = self.cube[y][z][x];
+        let original = self.get(x, y, z);
-        let original_index = self.palette.iter().position(|v| v == &original).unwrap();
+        if original == value {
-        self.counts[original_index] -= 1;
+            return original;
-
-        if self.counts[original_index] == 0 {
-            // Remove from palette and counts Vecs if the count hits zero.
-            self.palette.swap_remove(original_index);
-            self.counts.swap_remove(original_index);
         }
 
-        // Set the new value in the cube
+        let original_index = self.palette.iter().position(|v| v == &original).unwrap();
-        self.cube[y][z][x] = value;
 
         // Find or add the new value to the palette.
-        if let Some(new_index) = self.palette.iter().position(|v| v == &value) {
+        let new_index = if let Some(new_index) = self.palette.iter().position(|v| v == &value) {
             self.counts[new_index] += 1;
+            new_index
         } else {
             self.palette.push(value);
             self.counts.push(1);
+            self.palette.len() - 1
+        };
+
+        // Handle storage upgrades or updates
+        let mut upgraded = false;
+        match &mut self.storage {
+            PaletteStorage::Dense(cube) => {
+                cube[y][z][x] = value;
+            }
+            PaletteStorage::Indexed(indices) => {
+                if new_index <= 255 {
+                    indices[y][z][x] = new_index as u8;
+                } else {
+                    // Upgrade to Dense
+                    let mut cube = Box::new([[[V::default(); DIM]; DIM]; DIM]);
+                    for (i, v) in cube
+                        .as_flattened_mut()
+                        .as_flattened_mut()
+                        .iter_mut()
+                        .enumerate()
+                    {
+                        let y = i / (DIM * DIM);
+                        let z = (i / DIM) % DIM;
+                        let x = i % DIM;
+                        *v = self.palette[indices[y][z][x] as usize];
+                    }
+                    cube[y][z][x] = value;
+                    self.storage = PaletteStorage::Dense(cube);
+                    upgraded = true;
+                }
+            }
+        }
+
+        self.counts[original_index] -= 1;
+
+        if self.counts[original_index] == 0 {
+            let last_index = self.palette.len() - 1;
+            // Remove from palette and counts Vecs if the count hits zero.
+            self.palette.swap_remove(original_index);
+            self.counts.swap_remove(original_index);
+
+            // If we are indexed, we need to update all indices because swap_remove changed indices
+            if !upgraded && let PaletteStorage::Indexed(indices) = &mut self.storage {
+                for row in indices.iter_mut() {
+                    for col in row.iter_mut() {
+                        for idx in col.iter_mut() {
+                            if *idx as usize == last_index {
+                                *idx = original_index as u8;
+                            }
+                        }
+                    }
+                }
+            }
         }
 
         original
@@ -100,11 +157,24 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
             Self::Homogeneous(palette[0])
         } else {
             // Heterogeneous cube, store the full data
+            if palette.len() <= 256 && std::mem::size_of::<V>() > 1 {
+                let mut indices = Box::new([[[0u8; DIM]; DIM]; DIM]);
+                for (i, v) in cube.as_flattened().as_flattened().iter().enumerate() {
+                    let idx = palette.iter().position(|p| p == v).unwrap();
+                    indices.as_flattened_mut().as_flattened_mut()[i] = idx as u8;
+                }
                 Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
-                cube,
+                    storage: PaletteStorage::Indexed(indices),
                     palette,
                     counts,
                 }))
+            } else {
+                Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
+                    storage: PaletteStorage::Dense(cube),
+                    palette,
+                    counts,
+                }))
+            }
         }
     }
 
@@ -125,14 +195,15 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
                 // Don't use HashMap's here, because its slow
                 let blocks_per_i64 = 64 / bits_per_entry;
 
-                let packed_indices: Box<[i64]> = data
+                let packed_indices: Box<[i64]> = match &data.storage {
-                    .cube
+                    PaletteStorage::Dense(cube) => cube
                         .as_flattened()
                         .as_flattened()
                         .chunks(blocks_per_i64 as usize)
                         .map(|chunk| {
                             chunk.iter().enumerate().fold(0, |acc, (index, key)| {
-                            let key_index = data.palette.iter().position(|&x| x == *key).unwrap();
+                                let key_index =
+                                    data.palette.iter().position(|&x| x == *key).unwrap();
                                 debug_assert!((1 << bits_per_entry) > key_index);
 
                                 let packed_offset_index =
@@ -140,7 +211,23 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
                                 acc | packed_offset_index as i64
                             })
                         })
-                    .collect();
+                        .collect(),
+                    PaletteStorage::Indexed(indices) => indices
+                        .as_flattened()
+                        .as_flattened()
+                        .chunks(blocks_per_i64 as usize)
+                        .map(|chunk| {
+                            chunk.iter().enumerate().fold(0, |acc, (index, key_index)| {
+                                let key_index = *key_index as usize;
+                                debug_assert!((1 << bits_per_entry) > key_index);
+
+                                let packed_offset_index =
+                                    (key_index as u64) << (bits_per_entry as u64 * index as u64);
+                                acc | packed_offset_index as i64
+                            })
+                        })
+                        .collect(),
+                };
 
                 (data.palette.clone().into_boxed_slice(), packed_indices)
             }
@@ -166,10 +253,45 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
         let index_mask = (1 << bits_per_key) - 1;
         let keys_per_i64 = 64 / bits_per_key;
 
-        let mut decompressed_values = Vec::with_capacity(Self::VOLUME);
+        // Optimized path for indexed storage if palette is small enough
+        if palette.len() <= 256 && std::mem::size_of::<V>() > 1 {
+            let mut indices = Box::new([[[0u8; DIM]; DIM]; DIM]);
+            let mut counts = vec![0u16; palette.len()];
+            let indices_flat = indices.as_flattened_mut().as_flattened_mut();
 
-        // We already have the palette from the input `palette_slice`.
+            let mut packed_data_iter = packed_data.iter();
-        // The counts will be created in the next step.
+            let mut current_packed_word = *packed_data_iter.next().unwrap_or(&0);
+
+            for (i, index_out) in indices_flat.iter_mut().enumerate().take(Self::VOLUME) {
+                let bit_index_in_word = i % keys_per_i64 as usize;
+                if bit_index_in_word == 0 && i > 0 {
+                    current_packed_word = *packed_data_iter.next().unwrap_or(&0);
+                }
+
+                let lookup_index = ((current_packed_word as u64)
+                    >> (bit_index_in_word as u64 * bits_per_key as u64))
+                    & index_mask;
+
+                let idx = lookup_index as usize;
+                if idx < palette.len() {
+                    *index_out = idx as u8;
+                    counts[idx] += 1;
+                } else {
+                    warn!("Lookup index out of bounds! Defaulting...");
+                    // value is already 0, and counts[0] will be updated correctly if we track it
+                }
+            }
+            // fix counts[0] if it was skipped in out-of-bounds cases (rare)
+            // But actually we should just ensure it's correct.
+
+            return Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
+                storage: PaletteStorage::Indexed(indices),
+                palette: palette.to_vec(),
+                counts,
+            }));
+        }
+
+        let mut decompressed_values = Vec::with_capacity(Self::VOLUME);
 
         let mut packed_data_iter = packed_data.iter();
         let mut current_packed_word = *packed_data_iter.next().unwrap_or(&0);
@@ -216,7 +338,7 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
             .copy_from_slice(&decompressed_values);
 
         Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
-            cube,
+            storage: PaletteStorage::Dense(cube),
             palette: palette.to_vec(),
             counts,
         }))
@@ -257,7 +379,16 @@ impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM>
     pub fn iter(&self) -> Box<dyn Iterator<Item = &V> + '_> {
         match self {
             Self::Homogeneous(registry_id) => Box::new(repeat_n(registry_id, Self::VOLUME)),
-            Self::Heterogeneous(data) => Box::new(data.cube.as_flattened().as_flattened().iter()),
+            Self::Heterogeneous(data) => match &data.storage {
+                PaletteStorage::Dense(cube) => Box::new(cube.as_flattened().as_flattened().iter()),
+                PaletteStorage::Indexed(indices) => Box::new(
+                    indices
+                        .as_flattened()
+                        .as_flattened()
+                        .iter()
+                        .map(|&idx| &data.palette[idx as usize]),
+                ),
+            },
         }
     }
 
@@ -289,25 +420,28 @@ impl BiomePalette {
                 if raw_bits_per_entry > BIOME_NETWORK_MAX_MAP_BITS {
                     let bits_per_entry = BIOME_NETWORK_MAX_BITS;
                     let values_per_i64 = 64 / bits_per_entry;
-                    let packed_data = data
+                    let mut packed_data = Vec::new();
-                        .cube
+                    let mut current_idx = 0;
-                        .as_flattened()
+                    while current_idx < Self::VOLUME {
-                        .as_flattened()
+                        let mut acc = 0u64;
-                        .chunks(values_per_i64 as usize)
+                        for i in 0..values_per_i64 as usize {
-                        .map(|chunk| {
+                            if current_idx + i < Self::VOLUME {
-                            chunk.iter().enumerate().fold(0, |acc, (index, value)| {
+                                let y = (current_idx + i) / (Self::SIZE * Self::SIZE);
-                                debug_assert!((1 << bits_per_entry) > *value);
+                                let z = ((current_idx + i) / Self::SIZE) % Self::SIZE;
-                                let packed_offset_index =
+                                let x = (current_idx + i) % Self::SIZE;
-                                    (*value as u64) << (bits_per_entry as u64 * index as u64);
+                                let value = data.get(x, y, z);
-                                acc | packed_offset_index as i64
+                                debug_assert!((1 << bits_per_entry) > value);
-                            })
+                                acc |= (value as u64) << (bits_per_entry as u64 * i as u64);
-                        })
+                            }
-                        .collect();
+                        }
+                        packed_data.push(acc as i64);
+                        current_idx += values_per_i64 as usize;
+                    }
 
                     NetworkSerialization {
                         bits_per_entry,
                         palette: NetworkPalette::Direct,
-                        packed_data,
+                        packed_data: packed_data.into_boxed_slice(),
                     }
                 } else {
                     let bits_per_entry = raw_bits_per_entry.max(BIOME_NETWORK_MIN_MAP_BITS);
@@ -364,26 +498,28 @@ impl BlockPalette {
                 if raw_bits_per_entry > BLOCK_NETWORK_MAX_MAP_BITS {
                     let bits_per_entry = BLOCK_NETWORK_MAX_BITS;
                     let values_per_i64 = 64 / bits_per_entry;
-                    let packed_data = data
+                    let mut packed_data = Vec::new();
-                        .cube
+                    let mut current_idx = 0;
-                        .as_flattened()
+                    while current_idx < Self::VOLUME {
-                        .as_flattened()
+                        let mut acc = 0u64;
-                        .chunks(values_per_i64 as usize)
+                        for i in 0..values_per_i64 as usize {
-                        .map(|chunk| {
+                            if current_idx + i < Self::VOLUME {
-                            chunk.iter().enumerate().fold(0u64, |acc, (index, value)| {
+                                let y = (current_idx + i) / (Self::SIZE * Self::SIZE);
-                                debug_assert!((1 << bits_per_entry) > *value);
+                                let z = ((current_idx + i) / Self::SIZE) % Self::SIZE;
-
+                                let x = (current_idx + i) % Self::SIZE;
-                                let packed_offset_index =
+                                let value = data.get(x, y, z);
-                                    (*value as u64) << (bits_per_entry as u64 * index as u64);
+                                debug_assert!((1 << bits_per_entry) > value);
-                                acc | packed_offset_index
+                                acc |= (value as u64) << (bits_per_entry as u64 * i as u64);
-                            }) as i64
+                            }
-                        })
+                        }
-                        .collect();
+                        packed_data.push(acc as i64);
+                        current_idx += values_per_i64 as usize;
+                    }
 
                     NetworkSerialization {
                         bits_per_entry,
                         palette: NetworkPalette::Direct,
-                        packed_data,
+                        packed_data: packed_data.into_boxed_slice(),
                     }
                 } else {
                     let bits_per_entry = raw_bits_per_entry.max(BLOCK_NETWORK_MIN_MAP_BITS);

pumpkin-world/src/tick/scheduler.rs

@@ -9,11 +9,15 @@ use rustc_hash::FxHashSet;
 use crate::tick::{MAX_TICK_DELAY, OrderedTick, ScheduledTick};
 
 pub struct ChunkTickScheduler<T> {
-    tick_queue: Mutex<[Vec<OrderedTick<T>>; MAX_TICK_DELAY]>,
+    inner: Mutex<Option<Box<ChunkTickSchedulerInner<T>>>>,
-    queued_ticks: Mutex<FxHashSet<(BlockPos, T)>>,
     offset: AtomicUsize,
 }
 
+struct ChunkTickSchedulerInner<T> {
+    tick_queue: [Vec<OrderedTick<T>>; MAX_TICK_DELAY],
+    queued_ticks: FxHashSet<(BlockPos, T)>,
+}
+
 impl<'a, T: std::hash::Hash + Eq> ChunkTickScheduler<&'a T> {
     pub fn step_tick(&self) -> Vec<OrderedTick<&'a T>> {
         // Atomic update for the offset
@@ -21,29 +25,37 @@ impl<'a, T: std::hash::Hash + Eq> ChunkTickScheduler<&'a T> {
         let next_offset = (current_offset + 1) % MAX_TICK_DELAY;
         self.offset.store(next_offset, Ordering::SeqCst);
 
-        let res = {
+        let mut inner_guard = self.inner.lock().unwrap_or_else(|e| e.into_inner());
-            let mut queue = self.tick_queue.lock().unwrap_or_else(|e| e.into_inner());
+        let Some(inner) = inner_guard.as_mut() else {
-            std::mem::take(&mut queue[current_offset])
+            return Vec::new();
         };
 
+        let res = std::mem::take(&mut inner.tick_queue[current_offset]);
+
         if !res.is_empty() {
-            let mut set = self.queued_ticks.lock().unwrap_or_else(|e| e.into_inner());
             for next_tick in &res {
-                set.remove(&(next_tick.position, next_tick.value));
+                inner
+                    .queued_ticks
+                    .remove(&(next_tick.position, next_tick.value));
             }
         }
         res
     }
 
     pub fn schedule_tick(&self, tick: &ScheduledTick<&'a T>, sub_tick_order: u64) {
-        let mut set = self.queued_ticks.lock().unwrap_or_else(|e| e.into_inner());
+        let mut inner_guard = self.inner.lock().unwrap_or_else(|e| e.into_inner());
+        let inner = inner_guard.get_or_insert_with(|| {
+            Box::new(ChunkTickSchedulerInner {
+                tick_queue: std::array::from_fn(|_| Vec::new()),
+                queued_ticks: FxHashSet::default(),
+            })
+        });
 
-        if set.insert((tick.position, tick.value)) {
+        if inner.queued_ticks.insert((tick.position, tick.value)) {
-            let mut queue = self.tick_queue.lock().unwrap_or_else(|e| e.into_inner());
             let offset = self.offset.load(Ordering::SeqCst);
             let index = (offset + tick.delay as usize) % MAX_TICK_DELAY;
 
-            queue[index].push(OrderedTick {
+            inner.tick_queue[index].push(OrderedTick {
                 priority: tick.priority,
                 sub_tick_order,
                 position: tick.position,
@@ -53,29 +65,34 @@ impl<'a, T: std::hash::Hash + Eq> ChunkTickScheduler<&'a T> {
     }
 
     pub fn is_scheduled(&self, pos: BlockPos, value: &T) -> bool {
-        self.queued_ticks
+        self.inner
             .lock()
             .unwrap_or_else(|e| e.into_inner())
-            .contains(&(pos, value))
+            .as_ref()
+            .is_some_and(|inner| inner.queued_ticks.contains(&(pos, value)))
     }
 
     pub fn has_ticks(&self) -> bool {
-        !self
+        self.inner
-            .queued_ticks
             .lock()
             .unwrap_or_else(|e| e.into_inner())
-            .is_empty()
+            .as_ref()
+            .is_some_and(|inner| !inner.queued_ticks.is_empty())
     }
 
     #[must_use]
     pub fn to_vec(&self) -> Vec<ScheduledTick<&'a T>> {
         let offset = self.offset.load(Ordering::SeqCst);
-        let queue = self.tick_queue.lock().unwrap_or_else(|e| e.into_inner());
+        let inner_guard = self.inner.lock().unwrap_or_else(|e| e.into_inner());
+        let Some(inner) = inner_guard.as_ref() else {
+            return Vec::new();
+        };
+
         let mut res = Vec::new();
 
         for i in 0..MAX_TICK_DELAY {
             let index = (offset + i) % MAX_TICK_DELAY;
-            res.extend(queue[index].iter().map(|x| ScheduledTick {
+            res.extend(inner.tick_queue[index].iter().map(|x| ScheduledTick {
                 delay: i as u8,
                 priority: x.priority,
                 position: x.position,
@@ -95,11 +112,14 @@ impl<'a, T: std::hash::Hash + Eq + 'static> FromIterator<ScheduledTick<&'a T>>
 
         let (lower, _) = iter.size_hint();
         if lower > 0 {
-            scheduler
+            let mut inner_guard = scheduler.inner.lock().unwrap_or_else(|e| e.into_inner());
-                .queued_ticks
+            let inner = inner_guard.get_or_insert_with(|| {
-                .lock()
+                Box::new(ChunkTickSchedulerInner {
-                .unwrap_or_else(|e| e.into_inner())
+                    tick_queue: std::array::from_fn(|_| Vec::new()),
-                .reserve(lower);
+                    queued_ticks: FxHashSet::default(),
+                })
+            });
+            inner.queued_ticks.reserve(lower);
         }
 
         for tick in iter {
@@ -112,8 +132,7 @@ impl<'a, T: std::hash::Hash + Eq + 'static> FromIterator<ScheduledTick<&'a T>>
 impl<T> Default for ChunkTickScheduler<T> {
     fn default() -> Self {
         Self {
-            tick_queue: Mutex::new(std::array::from_fn(|_| Vec::new())),
+            inner: Mutex::new(None),
-            queued_ticks: Mutex::new(FxHashSet::default()),
             offset: AtomicUsize::new(0),
         }
     }

pumpkin/src/entity/decoration/armor_stand.rs

@@ -241,7 +241,9 @@ use pumpkin_nbt::pnbt::PNbtCompound;
 impl NBTStorage for ArmorStandEntity {
     fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
         Box::pin(async {
+            self.living_entity.write_nbt(nbt).await;
             let disabled_slots = self.disabled_slots.load(Ordering::Relaxed);
+            // ...
 
             nbt.put_bool(self.is_invisible());
             nbt.put_bool(self.is_small());
@@ -274,7 +276,9 @@ impl NBTStorage for ArmorStandEntity {
 
     fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
         Box::pin(async {
+            self.living_entity.read_nbt_non_mut(nbt).await;
             let mut flags = 0u8;
+            // ...
 
             let invisible = nbt.get_bool().unwrap_or(false);
             if invisible {

pumpkin/src/entity/decoration/painting.rs

@@ -22,12 +22,14 @@ use pumpkin_nbt::pnbt::PNbtCompound;
 impl NBTStorage for PaintingEntity {
     fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
         Box::pin(async {
+            self.entity.write_nbt(nbt).await;
             nbt.put_byte(self.entity.data.load(Ordering::Relaxed) as i8);
         })
     }
 
     fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
         Box::pin(async {
+            self.entity.read_nbt_non_mut(nbt).await;
             let facing = nbt.get_byte().unwrap_or(3);
             self.entity.data.store(facing as i32, Ordering::Relaxed);
         })

pumpkin/src/entity/mob/enderman.rs

@@ -401,6 +401,7 @@ use pumpkin_nbt::pnbt::PNbtCompound;
 impl NBTStorage for EndermanEntity {
     fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
         Box::pin(async {
+            self.mob_entity.living_entity.write_nbt(nbt).await;
             if let Some(block_state) = self.carried_block.load() {
                 nbt.put_bool(true);
                 nbt.put_int(block_state as i32);
@@ -412,6 +413,7 @@ impl NBTStorage for EndermanEntity {
 
     fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
         Box::pin(async {
+            self.mob_entity.living_entity.read_nbt_non_mut(nbt).await;
             if nbt.get_bool().unwrap_or(false) {
                 let block_state = nbt.get_int().unwrap_or(0);
                 self.set_carried_block(Some(block_state as u16)).await;

pumpkin/src/entity/mob/skeleton/mod.rs

@@ -3,7 +3,7 @@ use std::sync::{Arc, Weak};
 use pumpkin_data::entity::EntityType;
 
 use crate::entity::{
-    Entity, NBTStorage,
+    Entity, NBTStorage, NbtFuture,
     ai::goal::{
         active_target::ActiveTargetGoal, look_around::RandomLookAroundGoal,
         look_at_entity::LookAtEntityGoal, melee_attack::MeleeAttackGoal, revenge::RevengeGoal,
@@ -11,6 +11,7 @@ use crate::entity::{
     },
     mob::{Mob, MobEntity},
 };
+use pumpkin_nbt::pnbt::PNbtCompound;
 
 pub mod bogged;
 pub mod parched;
@@ -56,7 +57,15 @@ impl SkeletonEntityBase {
     }
 }
 
-impl NBTStorage for SkeletonEntityBase {}
+impl NBTStorage for SkeletonEntityBase {
+    fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.write_nbt(nbt)
+    }
+
+    fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.read_nbt_non_mut(nbt)
+    }
+}
 
 impl Mob for SkeletonEntityBase {
     fn get_mob_entity(&self) -> &MobEntity {

pumpkin/src/entity/mob/zombie/mod.rs

@@ -6,10 +6,11 @@ use crate::entity::ai::goal::swim::SwimGoal;
 use crate::entity::ai::goal::wander_around::WanderAroundGoal;
 use crate::entity::ai::goal::zombie_attack::ZombieAttackGoal;
 use crate::entity::{
-    Entity, NBTStorage,
+    Entity, NBTStorage, NbtFuture,
     ai::goal::{active_target::ActiveTargetGoal, look_at_entity::LookAtEntityGoal},
 };
 use pumpkin_data::entity::EntityType;
+use pumpkin_nbt::pnbt::PNbtCompound;
 use std::sync::{Arc, Weak};
 
 pub mod drowned;
@@ -69,7 +70,15 @@ impl ZombieEntityBase {
     }
 }
 
-impl NBTStorage for ZombieEntityBase {}
+impl NBTStorage for ZombieEntityBase {
+    fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.write_nbt(nbt)
+    }
+
+    fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.read_nbt_non_mut(nbt)
+    }
+}
 
 impl Mob for ZombieEntityBase {
     fn get_mob_entity(&self) -> &MobEntity {

pumpkin/src/entity/mod.rs

@@ -848,14 +848,14 @@ impl Entity {
         movement: Vector3<f64>,
         caller: &dyn EntityBase,
     ) -> Vector3<f64> {
-        self.on_ground.store(false, Ordering::SeqCst);
-        self.supporting_block_pos.store(None);
-        self.horizontal_collision.store(false, Ordering::SeqCst);
-
         if movement.length_squared() == 0.0 {
             return movement;
         }
 
+        self.on_ground.store(false, Ordering::SeqCst);
+        self.supporting_block_pos.store(None);
+        self.horizontal_collision.store(false, Ordering::SeqCst);
+
         let bounding_box = self.bounding_box.load();
 
         let (collisions, block_positions) = self
@@ -2740,6 +2740,7 @@ impl NBTStorage for Entity {
             let z = nbt.get_f64().unwrap_or(0.0);
             let pos = Vector3::new(x, y, z);
             self.set_pos(pos);
+            self.last_sent_pos.store(pos);
             self.first_loaded_chunk_position.store(Some(pos.to_i32()));
             let vx = nbt.get_f64().unwrap_or(0.0);
             let vy = nbt.get_f64().unwrap_or(0.0);
@@ -2748,7 +2749,12 @@ impl NBTStorage for Entity {
             let yaw = nbt.get_f32().unwrap_or(0.0);
             let pitch = nbt.get_f32().unwrap_or(0.0);
             self.set_rotation(yaw, pitch);
+            let yaw_byte = (yaw * 256.0 / 360.0).rem_euclid(256.0) as u8;
+            let pitch_byte = (pitch * 256.0 / 360.0).rem_euclid(256.0) as u8;
+            self.last_sent_yaw.store(yaw_byte, Relaxed);
+            self.last_sent_pitch.store(pitch_byte, Relaxed);
             self.head_yaw.store(yaw);
+            self.last_sent_head_yaw.store(yaw_byte, Relaxed);
             self.fire_ticks
                 .store(i32::from(nbt.get_short().unwrap_or(0)), Relaxed);
             self.on_ground

pumpkin/src/entity/passive/chicken.rs

@@ -11,7 +11,7 @@ use pumpkin_util::math::vector3::Vector3;
 use rand::RngExt;
 
 use crate::entity::{
-    Entity, EntityBase, EntityBaseFuture, NBTStorage,
+    Entity, EntityBase, EntityBaseFuture, NBTStorage, NbtFuture,
     ai::goal::{
         breed::BreedGoal, escape_danger::EscapeDangerGoal, follow_parent::FollowParentGoal,
         look_around::RandomLookAroundGoal, look_at_entity::LookAtEntityGoal, swim::SwimGoal,
@@ -20,6 +20,7 @@ use crate::entity::{
     mob::{Mob, MobEntity},
     player::Player,
 };
+use pumpkin_nbt::pnbt::PNbtCompound;
 
 const TEMPT_ITEMS: &[&Item] = &[
     &Item::WHEAT_SEEDS,
@@ -72,7 +73,22 @@ impl ChickenEntity {
     }
 }
 
-impl NBTStorage for ChickenEntity {}
+impl NBTStorage for ChickenEntity {
+    fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        Box::pin(async {
+            self.mob_entity.living_entity.write_nbt(nbt).await;
+            nbt.put_int(self.egg_lay_time.load(Ordering::Relaxed));
+        })
+    }
+
+    fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        Box::pin(async {
+            self.mob_entity.living_entity.read_nbt_non_mut(nbt).await;
+            self.egg_lay_time
+                .store(nbt.get_int().unwrap_or(6000), Ordering::Relaxed);
+        })
+    }
+}
 
 impl Mob for ChickenEntity {
     fn get_mob_entity(&self) -> &MobEntity {

pumpkin/src/entity/passive/cow.rs

@@ -7,7 +7,7 @@ use pumpkin_data::{entity::EntityType, item::Item};
 use pumpkin_util::math::vector3::Vector3;
 
 use crate::entity::{
-    Entity, EntityBase, EntityBaseFuture, NBTStorage,
+    Entity, EntityBase, EntityBaseFuture, NBTStorage, NbtFuture,
     ai::goal::{
         breed::BreedGoal, escape_danger::EscapeDangerGoal, follow_parent::FollowParentGoal,
         look_around::RandomLookAroundGoal, look_at_entity::LookAtEntityGoal, swim::SwimGoal,
@@ -16,6 +16,7 @@ use crate::entity::{
     mob::{Mob, MobEntity},
     player::Player,
 };
+use pumpkin_nbt::pnbt::PNbtCompound;
 
 const TEMPT_ITEMS: &[&Item] = &[&Item::WHEAT];
 
@@ -56,7 +57,15 @@ impl CowEntity {
     }
 }
 
-impl NBTStorage for CowEntity {}
+impl NBTStorage for CowEntity {
+    fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.write_nbt(nbt)
+    }
+
+    fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.read_nbt_non_mut(nbt)
+    }
+}
 
 impl Mob for CowEntity {
     fn get_mob_entity(&self) -> &MobEntity {

pumpkin/src/entity/passive/iron_golem.rs

@@ -3,7 +3,7 @@ use std::sync::{Arc, Weak};
 use pumpkin_data::entity::EntityType;
 
 use crate::entity::{
-    Entity, NBTStorage,
+    Entity, NBTStorage, NbtFuture,
     ai::goal::{
         active_target::ActiveTargetGoal, look_around::RandomLookAroundGoal,
         look_at_entity::LookAtEntityGoal, melee_attack::MeleeAttackGoal, revenge::RevengeGoal,
@@ -11,6 +11,7 @@ use crate::entity::{
     },
     mob::{Mob, MobEntity},
 };
+use pumpkin_nbt::pnbt::PNbtCompound;
 
 /// Represents an Iron Golem, a powerful neutral mob that protects villagers and players.
 ///
@@ -56,7 +57,15 @@ impl IronGolemEntity {
     }
 }
 
-impl NBTStorage for IronGolemEntity {}
+impl NBTStorage for IronGolemEntity {
+    fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.write_nbt(nbt)
+    }
+
+    fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.read_nbt_non_mut(nbt)
+    }
+}
 
 impl Mob for IronGolemEntity {
     fn get_mob_entity(&self) -> &MobEntity {

pumpkin/src/entity/passive/pig.rs

@@ -7,7 +7,7 @@ use pumpkin_data::{entity::EntityType, item::Item};
 use pumpkin_util::math::vector3::Vector3;
 
 use crate::entity::{
-    Entity, EntityBase, EntityBaseFuture, NBTStorage,
+    Entity, EntityBase, EntityBaseFuture, NBTStorage, NbtFuture,
     ai::goal::{
         breed::BreedGoal, escape_danger::EscapeDangerGoal, follow_parent::FollowParentGoal,
         look_around::RandomLookAroundGoal, look_at_entity::LookAtEntityGoal, swim::SwimGoal,
@@ -16,6 +16,7 @@ use crate::entity::{
     mob::{Mob, MobEntity},
     player::Player,
 };
+use pumpkin_nbt::pnbt::PNbtCompound;
 
 const PIG_FOOD: &[&Item] = &[
     &Item::CARROT,
@@ -61,7 +62,15 @@ impl PigEntity {
     }
 }
 
-impl NBTStorage for PigEntity {}
+impl NBTStorage for PigEntity {
+    fn write_nbt<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.write_nbt(nbt)
+    }
+
+    fn read_nbt_non_mut<'a>(&'a self, nbt: &'a mut PNbtCompound) -> NbtFuture<'a, ()> {
+        self.mob_entity.living_entity.read_nbt_non_mut(nbt)
+    }
+}
 
 impl Mob for PigEntity {
     fn get_mob_entity(&self) -> &MobEntity {

pumpkin/src/entity/player.rs

@@ -111,7 +111,7 @@ const MAX_PREVIOUS_MESSAGES: u8 = 20; // Vanilla: 20
 
 pub const DATA_VERSION: i32 = 4790; // 26.1.2
 
-struct HeapNode(i32, Vector2<i32>, SyncChunk);
+struct HeapNode(i32, Vector2<i32>, Weak<ChunkData>);
 
 impl Eq for HeapNode {}
 
@@ -140,7 +140,7 @@ pub struct ChunkManager {
     chunk_listener: Receiver<(Vector2<i32>, SyncChunk)>,
     chunk_sent: HashMap<Vector2<i32>, Weak<ChunkData>>,
     chunk_queue: BinaryHeap<HeapNode>,
-    entity_chunk_queue: VecDeque<(Vector2<i32>, SyncEntityChunk)>,
+    entity_chunk_queue: VecDeque<(Vector2<i32>, Weak<ChunkEntityData>)>,
     batches_sent_since_ack: u8,
     last_chunk_batch_sent_at: Instant,
     /// The current world for chunk loading. Updated on dimension change.
@@ -202,7 +202,8 @@ impl ChunkManager {
                 continue;
             }
             if self.should_enqueue_chunk(pos, &chunk) {
-                self.chunk_queue.push(HeapNode(dst, pos, chunk));
+                self.chunk_queue
+                    .push(HeapNode(dst, pos, Arc::downgrade(&chunk)));
             }
         }
     }
@@ -248,6 +249,11 @@ impl ChunkManager {
                 && !unloading_chunks.contains(pos)
         });
 
+        self.entity_chunk_queue.retain(|(pos, _)| {
+            (pos.x - center.x).abs().max((pos.y - center.y).abs()) <= view_distance_i32
+                && !unloading_chunks.contains(pos)
+        });
+
         let mut tasks: Vec<_> = self
             .chunk_queue
             .drain()
@@ -265,7 +271,7 @@ impl ChunkManager {
                 let chunk = chunk.value().clone();
                 if self.should_enqueue_chunk(*pos, &chunk) {
                     let dst = (pos.x - center.x).abs().max((pos.y - center.y).abs());
-                    tasks.push(HeapNode(dst, *pos, chunk));
+                    tasks.push(HeapNode(dst, *pos, Arc::downgrade(&chunk)));
                 }
             }
         }
@@ -313,17 +319,19 @@ impl ChunkManager {
         self.chunks_per_tick = chunks_per_tick.ceil() as usize;
     }
 
-    pub fn push_chunk(&mut self, position: Vector2<i32>, chunk: SyncChunk) {
+    pub fn push_chunk(&mut self, position: Vector2<i32>, chunk: &SyncChunk) {
-        if self.should_enqueue_chunk(position, &chunk) {
+        if self.should_enqueue_chunk(position, chunk) {
             let dst = (position.x - self.center.x)
                 .abs()
                 .max((position.y - self.center.y).abs());
-            self.chunk_queue.push(HeapNode(dst, position, chunk));
+            self.chunk_queue
+                .push(HeapNode(dst, position, Arc::downgrade(chunk)));
         }
     }
 
-    pub fn push_entity(&mut self, position: Vector2<i32>, chunk: SyncEntityChunk) {
+    pub fn push_entity(&mut self, position: Vector2<i32>, chunk: &SyncEntityChunk) {
-        self.entity_chunk_queue.push_back((position, chunk));
+        self.entity_chunk_queue
+            .push_back((position, Arc::downgrade(chunk)));
     }
 
     #[must_use]
@@ -336,8 +344,12 @@ impl ChunkManager {
     pub fn next_chunk(&mut self) -> Box<[SyncChunk]> {
         let take = self.chunk_queue.len().min(self.chunks_per_tick.max(1));
         let mut chunks = Vec::with_capacity(take);
-        for _ in 0..take {
+        while chunks.len() < take
-            chunks.push(self.chunk_queue.pop().unwrap().2);
+            && let Some(node) = self.chunk_queue.pop()
+        {
+            if let Some(chunk) = node.2.upgrade() {
+                chunks.push(chunk);
+            }
         }
         self.batches_sent_since_ack = self.batches_sent_since_ack.saturating_add(1);
         self.last_chunk_batch_sent_at = Instant::now();
@@ -351,16 +363,19 @@ impl ChunkManager {
             .len()
             .min(self.chunks_per_tick.max(1));
 
-        let chunks: Box<[Arc<ChunkEntityData>]> = self
+        let mut chunks = Vec::with_capacity(chunk_size);
-            .entity_chunk_queue
+        while chunks.len() < chunk_size
-            .drain(..chunk_size)
+            && let Some((_, weak_chunk)) = self.entity_chunk_queue.pop_front()
-            .map(|(_, chunk)| chunk)
+        {
-            .collect();
+            if let Some(chunk) = weak_chunk.upgrade() {
+                chunks.push(chunk);
+            }
+        }
 
         self.batches_sent_since_ack = self.batches_sent_since_ack.saturating_add(1);
         self.last_chunk_batch_sent_at = Instant::now();
 
-        chunks
+        chunks.into_boxed_slice()
     }
 }
 
@@ -779,7 +794,10 @@ impl Player {
         // Decrement the value of watched chunks
         let chunks_to_clean = level.mark_chunks_as_not_watched(&radial_chunks).await;
         // Remove chunks with no watchers from the cache
+        if !chunks_to_clean.is_empty() {
             level.clean_entity_chunks(&chunks_to_clean);
+            world.remove_entities_in_chunks(&chunks_to_clean).await;
+        }
         // Remove left over entries from all possiblily loaded chunks
         level.clean_memory();
 

pumpkin/src/world/chunker.rs

@@ -108,11 +108,16 @@ pub async fn update_position(player: &Arc<Player>) {
         .level
         .mark_chunks_as_newly_watched(&loading_chunks)
         .await;
-    world
+    let chunks_to_clean = world
         .level
         .mark_chunks_as_not_watched(&unloading_chunks)
         .await;
 
+    if !chunks_to_clean.is_empty() {
+        world.level.clean_entity_chunks(&chunks_to_clean);
+        world.remove_entities_in_chunks(&chunks_to_clean).await;
+    }
+
     if !loading_chunks.is_empty() {
         world.spawn_world_entity_chunks(player.clone(), loading_chunks, new_chunk_center);
     }

pumpkin/src/world/mod.rs

@@ -3188,6 +3188,31 @@ impl World {
         self.remove_entity_data(base_entity).await;
     }
 
+    pub async fn remove_entities_in_chunks(&self, chunks: &[Vector2<i32>]) {
+        let chunks_set: FxHashSet<_> = chunks.iter().copied().collect();
+        let mut entities_to_remove = Vec::new();
+
+        self.entities.rcu(|current_entities| {
+            let mut new_entities = (**current_entities).clone();
+            new_entities.retain(|entity| {
+                let base_entity = entity.get_entity();
+                let pos = base_entity.chunk_pos.load();
+                if chunks_set.contains(&pos) {
+                    entities_to_remove.push(entity.clone());
+                    false
+                } else {
+                    true
+                }
+            });
+            new_entities
+        });
+
+        for entity in entities_to_remove {
+            self.spawn_state.load().remove_entity(self, entity.as_ref());
+            self.remove_entity_data(entity.get_entity()).await;
+        }
+    }
+
     pub async fn set_block_breaking(&self, from: &Entity, location: BlockPos, progress: i32) {
         let chunk_pos = location.chunk_position(); // pumpkin's BlockPos already has this method
         self.broadcast_to_chunk_except(