fix block entities not being saved after chunk unload

(cherry picked from commit db8d44e)
Avanatiker · Feb 5, 2024 · 958fc4c · 958fc4c
1 parent b583736
commit 958fc4c
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 11 deletions.
diff --git a/common/src/main/kotlin/org/waste/of/time/Events.kt b/common/src/main/kotlin/org/waste/of/time/Events.kt
@@ -13,23 +13,23 @@ import net.minecraft.entity.player.PlayerEntity
 import net.minecraft.util.hit.BlockHitResult
 import net.minecraft.world.World
 import net.minecraft.world.chunk.WorldChunk
+import org.waste.of.time.WorldTools.CAPTURE_KEY
+import org.waste.of.time.WorldTools.CONFIG_KEY
+import org.waste.of.time.WorldTools.config
+import org.waste.of.time.WorldTools.mc
+import org.waste.of.time.gui.ManagerScreen
+import org.waste.of.time.manager.BarManager.updateCapture
 import org.waste.of.time.manager.CaptureManager
 import org.waste.of.time.manager.CaptureManager.capturing
 import org.waste.of.time.manager.CaptureManager.currentLevelName
 import org.waste.of.time.manager.MessageManager
 import org.waste.of.time.manager.MessageManager.translateHighlight
 import org.waste.of.time.manager.StatisticManager
-import org.waste.of.time.WorldTools.CAPTURE_KEY
-import org.waste.of.time.WorldTools.CONFIG_KEY
-import org.waste.of.time.WorldTools.config
-import org.waste.of.time.WorldTools.mc
+import org.waste.of.time.storage.StorageFlow
 import org.waste.of.time.storage.cache.EntityCacheable
+import org.waste.of.time.storage.cache.HotCache
 import org.waste.of.time.storage.serializable.PlayerStoreable
 import org.waste.of.time.storage.serializable.RegionBasedChunk
-import org.waste.of.time.gui.ManagerScreen
-import org.waste.of.time.manager.BarManager.updateCapture
-import org.waste.of.time.storage.cache.HotCache
-import org.waste.of.time.storage.StorageFlow
 import java.awt.Color
 
 object Events {
@@ -38,7 +38,9 @@ object Events {
     }
 
     fun onChunkUnload(chunk: WorldChunk) {
-        RegionBasedChunk(chunk).apply {
+        val regionBasedChunk = HotCache.chunks[chunk.pos] ?: RegionBasedChunk(chunk)
+        regionBasedChunk.apply {
+            cacheBlockEntities()
             emit()
             flush()
         }

diff --git a/common/src/main/kotlin/org/waste/of/time/storage/serializable/RegionBasedChunk.kt b/common/src/main/kotlin/org/waste/of/time/storage/serializable/RegionBasedChunk.kt
@@ -4,6 +4,7 @@ import net.minecraft.SharedConstants
 import net.minecraft.block.Block
 import net.minecraft.block.BlockState
 import net.minecraft.block.Blocks
+import net.minecraft.block.entity.BlockEntity
 import net.minecraft.fluid.Fluid
 import net.minecraft.nbt.NbtCompound
 import net.minecraft.nbt.NbtList
@@ -12,6 +13,7 @@ import net.minecraft.nbt.NbtOps
 import net.minecraft.registry.Registries
 import net.minecraft.registry.RegistryKeys
 import net.minecraft.text.MutableText
+import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.ChunkPos
 import net.minecraft.util.math.ChunkSectionPos
 import net.minecraft.world.ChunkSerializer
@@ -32,6 +34,18 @@ import org.waste.of.time.storage.RegionBased
 import org.waste.of.time.storage.cache.HotCache
 
 data class RegionBasedChunk(val chunk: WorldChunk) : RegionBased, Cacheable {
+    // storing a reference to the block entities in the chunk
+    // so it doesn't get removed out from under us while the chunk is unloaded
+    private var cachedBlockEntities: Map<BlockPos, BlockEntity>? = null
+
+    fun cacheBlockEntities() {
+        cachedBlockEntities = HashMap(chunk.blockEntities)
+    }
+
+    private fun getBlockEntities() : Map<BlockPos, BlockEntity> {
+        return cachedBlockEntities ?: chunk.blockEntities
+    }
+
     override fun shouldStore() = config.capture.chunks
 
     override val verboseInfo: MutableText
@@ -104,8 +118,8 @@ data class RegionBasedChunk(val chunk: WorldChunk) : RegionBased, Cacheable {
         }
 
         put("block_entities", NbtList().apply {
-            chunk.blockEntityPositions.mapNotNull {
-                chunk.getPackedBlockEntityNbt(it)
+            getBlockEntities().entries.mapNotNull {
+                getPackedBlockEntityNbt(it)
             }.forEach { add(it) }
         })
 
@@ -115,6 +129,13 @@ data class RegionBasedChunk(val chunk: WorldChunk) : RegionBased, Cacheable {
         // skip structures
     }
 
+    private fun getPackedBlockEntityNbt(entry: Map.Entry<BlockPos, BlockEntity>): NbtCompound? {
+        val blockEntity: BlockEntity = entry.value
+        var nbtCompound: NbtCompound = blockEntity.createNbtWithIdentifyingData()
+        nbtCompound.putBoolean("keepPacked", false)
+        return nbtCompound
+    }
+
     private fun generateSections(chunk: WorldChunk) = NbtList().apply {
         val biomeRegistry = chunk.world.registryManager.get(RegistryKeys.BIOME)
         val biomeCodec = PalettedContainer.createReadableContainerCodec(