fix: save block entities before their chunk is dropped (#2856)

crates/pumpkin/src/block/entities/mod.rs

@@ -433,3 +433,40 @@ pub fn create_block_entity(
         _ => None,
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::{BlockEntity, block_entity_from_nbt, furnace::FurnaceBlockEntity};
+    use pumpkin_data::{item::Item, item_stack::ItemStack};
+    use pumpkin_nbt::compound::NbtCompound;
+    use pumpkin_util::math::position::BlockPos;
+    use pumpkin_world::inventory::Inventory;
+    use std::sync::Arc;
+
+    /// A loaded block entity is serialized back into its chunk with
+    /// `write_internal`, so whatever it holds has to survive that round trip or
+    /// it is gone the next time the chunk is read.
+    #[tokio::test]
+    async fn furnace_contents_survive_a_chunk_round_trip() {
+        let position = BlockPos::new(0, 100, 0);
+        let furnace = Arc::new(FurnaceBlockEntity::new(position));
+        furnace
+            .set_stack(0, ItemStack::new(5, &Item::DIAMOND))
+            .await;
+
+        let mut nbt = NbtCompound::new();
+        furnace.write_internal(&mut nbt).await;
+
+        let inventory = block_entity_from_nbt(&nbt).and_then(BlockEntity::get_inventory);
+        assert!(
+            inventory.is_some(),
+            "furnace should be readable back from its own NBT"
+        );
+
+        if let Some(inventory) = inventory {
+            let stack = inventory.get_stack(0).await;
+            assert_eq!(stack.get_item().id, Item::DIAMOND.id);
+            assert_eq!(stack.item_count, 5);
+        }
+    }
+}

crates/pumpkin/src/world/mod.rs

@@ -385,6 +385,15 @@ impl World {
             self.save_entity(entity).await;
         }
 
+        let chunks: Vec<Vector2<i32>> = self
+            .block_entities
+            .iter()
+            .map(|chunk_block_entities| *chunk_block_entities.key())
+            .collect();
+        for chunk_pos in chunks {
+            self.save_block_entities(&chunk_pos).await;
+        }
+
         // Save portal POI to disk
         let save_result = self.portal_poi.lock().await.save_all();
         if let Err(e) = save_result {
@@ -412,6 +421,27 @@ impl World {
         chunk.mark_dirty(true);
     }
 
+    /// Serializes the live block entities of a chunk back into that chunk's block
+    /// entity data. The live map is the source of truth while a chunk is loaded -
+    /// `get_block_entity` takes the saved NBT out of the chunk when it wakes an
+    /// entity up - so this has to run before the chunk is dropped, or everything
+    /// the entity did since it was loaded is lost.
+    async fn save_block_entities(&self, chunk_pos: &Vector2<i32>) {
+        let Some(block_entities) = self
+            .block_entities
+            .get(chunk_pos)
+            .map(|chunk_block_entities| chunk_block_entities.values().cloned().collect::<Vec<_>>())
+        else {
+            return;
+        };
+
+        for block_entity in block_entities {
+            let mut nbt = NbtCompound::new();
+            block_entity.write_internal(&mut nbt).await;
+            self.add_block_entity_nbt(block_entity.get_position(), &nbt);
+        }
+    }
+
     /// Sends an entity status update to all players tracking the specified entity.
     pub fn send_entity_status(
         &self,
@@ -4431,6 +4461,7 @@ impl World {
         }
 
         for chunk_pos in &chunks_set {
+            self.save_block_entities(chunk_pos).await;
             self.block_entities.remove(chunk_pos);
         }
     }