Use ChunkStatusTracker for invalidation of ChunkSectionChangeCallbacks

CaffeineMC · Jun 17, 2024 · 382b871 · 382b871
1 parent bc98a03
commit 382b871
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 18 deletions.
diff --git a/lithium-mixin-config.md b/lithium-mixin-config.md
@@ -473,7 +473,8 @@ Allows access to existing BlockEntities without creating new ones
 (default: `true`)  
 Chunk sections count certain blocks inside them and provide a method to quickly check whether a chunk contains any of these blocks. Furthermore, chunk sections can notify registered listeners about certain blocks being placed or broken.  
 Requirements:
-- `mixin.util.data_storage=true`  
+- `mixin.util.data_storage=true`
+- `mixin.util.chunk_status_tracking=true`  
 
 ### `mixin.util.chunk_access`
 (default: `true`)  

diff --git a/src/main/java/me/jellysquid/mods/lithium/common/block/BlockListeningSection.java b/src/main/java/me/jellysquid/mods/lithium/common/block/BlockListeningSection.java
@@ -2,10 +2,11 @@
 
 import me.jellysquid.mods.lithium.common.entity.block_tracking.SectionedBlockChangeTracker;
 import net.minecraft.util.math.ChunkSectionPos;
+import net.minecraft.world.World;
 
 public interface BlockListeningSection {
 
-    void lithium$addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker);
+    void lithium$addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker, long sectionPos, World world);
 
     void lithium$removeFromCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker);
 

diff --git a/...a/me/jellysquid/mods/lithium/common/entity/block_tracking/ChunkSectionChangeCallback.java b/...a/me/jellysquid/mods/lithium/common/entity/block_tracking/ChunkSectionChangeCallback.java
@@ -1,22 +1,55 @@
 package me.jellysquid.mods.lithium.common.entity.block_tracking;
 
+import it.unimi.dsi.fastutil.longs.Long2ReferenceOpenHashMap;
 import me.jellysquid.mods.lithium.common.block.BlockListeningSection;
 import me.jellysquid.mods.lithium.common.block.BlockStateFlags;
 import me.jellysquid.mods.lithium.common.block.ListeningBlockStatePredicate;
+import me.jellysquid.mods.lithium.common.util.Pos;
+import me.jellysquid.mods.lithium.common.world.LithiumData;
+import me.jellysquid.mods.lithium.common.world.chunk.ChunkStatusTracker;
 import net.minecraft.util.math.ChunkSectionPos;
+import net.minecraft.world.World;
+import net.minecraft.world.chunk.ChunkSection;
 
 import java.util.ArrayList;
 
 public final class ChunkSectionChangeCallback {
     private final ArrayList<SectionedBlockChangeTracker>[] trackers;
     private short listeningMask;
 
+    static {
+        if (BlockListeningSection.class.isAssignableFrom(ChunkSection.class)) {
+            ChunkStatusTracker.registerUnloadCallback((serverWorld, chunkPos) -> {
+                Long2ReferenceOpenHashMap<ChunkSectionChangeCallback> changeCallbacks = ((LithiumData) serverWorld).lithium$getData().chunkSectionChangeCallbacks();
+                int x = chunkPos.x;
+                int z = chunkPos.z;
+                for (int y = Pos.SectionYCoord.getMinYSection(serverWorld); y <= Pos.SectionYCoord.getMaxYSectionInclusive(serverWorld); y++) {
+                    ChunkSectionPos chunkSectionPos = ChunkSectionPos.from(x, y, z);
+                    ChunkSectionChangeCallback chunkSectionChangeCallback = changeCallbacks.remove(chunkSectionPos.asLong());
+                    if (chunkSectionChangeCallback != null) {
+                        chunkSectionChangeCallback.onChunkSectionInvalidated(chunkSectionPos);
+                    }
+                }
+            });
+        }
+    }
+
     public ChunkSectionChangeCallback() {
         //noinspection unchecked
         this.trackers = new ArrayList[BlockStateFlags.NUM_LISTENING_FLAGS];
         this.listeningMask = 0;
     }
 
+    public static ChunkSectionChangeCallback create(long sectionPos, World world) {
+        ChunkSectionChangeCallback chunkSectionChangeCallback = new ChunkSectionChangeCallback();
+        Long2ReferenceOpenHashMap<ChunkSectionChangeCallback> changeCallbacks = ((LithiumData) world).lithium$getData().chunkSectionChangeCallbacks();
+        ChunkSectionChangeCallback previous = changeCallbacks.put(sectionPos, chunkSectionChangeCallback);
+        if (previous != null) {
+            previous.onChunkSectionInvalidated(ChunkSectionPos.from(sectionPos));
+        }
+        return chunkSectionChangeCallback;
+    }
+
     public short onBlockChange(int flagIndex, BlockListeningSection section) {
         ArrayList<SectionedBlockChangeTracker> sectionedBlockChangeTrackers = this.trackers[flagIndex];
         this.trackers[flagIndex] = null;

diff --git a/.../me/jellysquid/mods/lithium/common/entity/block_tracking/SectionedBlockChangeTracker.java b/.../me/jellysquid/mods/lithium/common/entity/block_tracking/SectionedBlockChangeTracker.java
@@ -59,23 +59,25 @@ public void register() {
             WorldSectionBox trackedSections = this.trackedWorldSections;
             for (int x = trackedSections.chunkX1(); x < trackedSections.chunkX2(); x++) {
                 for (int z = trackedSections.chunkZ1(); z < trackedSections.chunkZ2(); z++) {
-                    Chunk chunk = trackedSections.world().getChunk(x, z, ChunkStatus.FULL, false);
+                    World world = trackedSections.world();
+                    Chunk chunk = world.getChunk(x, z, ChunkStatus.FULL, false);
                     ChunkSection[] sectionArray = chunk == null ? null : chunk.getSectionArray();
                     for (int y = trackedSections.chunkY1(); y < trackedSections.chunkY2(); y++) {
-                        if (Pos.SectionYCoord.getMinYSection(trackedSections.world()) > y || Pos.SectionYCoord.getMaxYSectionExclusive(trackedSections.world()) <= y) {
+                        if (Pos.SectionYCoord.getMinYSection(world) > y || Pos.SectionYCoord.getMaxYSectionExclusive(world) <= y) {
                             continue;
                         }
+                        ChunkSectionPos sectionPos = ChunkSectionPos.from(x, y, z);
                         if (sectionArray == null) {
                             if (this.sectionsNotListeningTo == null) {
                                 this.sectionsNotListeningTo = new ArrayList<>();
                             }
-                            this.sectionsNotListeningTo.add(ChunkSectionPos.from(x, y, z));
+                            this.sectionsNotListeningTo.add(sectionPos);
                             continue;
                         }
-                        ChunkSection section = sectionArray[Pos.SectionYIndex.fromSectionCoord(trackedSections.world(), y)];
+                        ChunkSection section = sectionArray[Pos.SectionYIndex.fromSectionCoord(world, y)];
 
                         BlockListeningSection blockListeningSection = (BlockListeningSection) section;
-                        blockListeningSection.lithium$addToCallback(this.blockGroup, this);
+                        blockListeningSection.lithium$addToCallback(this.blockGroup, this, ChunkSectionPos.asLong(x, y, z), world);
                     }
                 }
             }
@@ -123,24 +125,25 @@ public void listenToAllSections() {
             for (int i = notListeningTo.size() - 1; i >= 0; i--) {
                 changed = true;
                 ChunkSectionPos chunkSectionPos = notListeningTo.get(i);
-                Chunk chunk = this.trackedWorldSections.world().getChunk(chunkSectionPos.getX(), chunkSectionPos.getZ(), ChunkStatus.FULL, false);
+                World world = this.trackedWorldSections.world();
+                Chunk chunk = world.getChunk(chunkSectionPos.getX(), chunkSectionPos.getZ(), ChunkStatus.FULL, false);
                 if (chunk != null) {
                     notListeningTo.remove(i);
                 } else {
                     //Chunk not loaded, cannot listen to all sections.
                     return;
                 }
-                ChunkSection section = chunk.getSectionArray()[Pos.SectionYIndex.fromSectionCoord(this.trackedWorldSections.world(), chunkSectionPos.getY())];
+                ChunkSection section = chunk.getSectionArray()[Pos.SectionYIndex.fromSectionCoord(world, chunkSectionPos.getY())];
                 BlockListeningSection blockListeningSection = (BlockListeningSection) section;
-                blockListeningSection.lithium$addToCallback(this.blockGroup, this);
+                blockListeningSection.lithium$addToCallback(this.blockGroup, this, chunkSectionPos.asLong(), world);
             }
         }
         if (this.sectionsUnsubscribed != null) {
             ArrayList<BlockListeningSection> unsubscribed = this.sectionsUnsubscribed;
             for (int i = unsubscribed.size() - 1; i >= 0; i--) {
                 changed = true;
                 BlockListeningSection blockListeningSection = unsubscribed.remove(i);
-                blockListeningSection.lithium$addToCallback(this.blockGroup, this);
+                blockListeningSection.lithium$addToCallback(this.blockGroup, this, Long.MIN_VALUE, null);
             }
         }
         this.isListeningToAll = true;

diff --git a/src/main/java/me/jellysquid/mods/lithium/common/world/LithiumData.java b/src/main/java/me/jellysquid/mods/lithium/common/world/LithiumData.java
@@ -3,6 +3,7 @@
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.longs.Long2ReferenceOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet;
+import me.jellysquid.mods.lithium.common.entity.block_tracking.ChunkSectionChangeCallback;
 import me.jellysquid.mods.lithium.common.entity.block_tracking.SectionedBlockChangeTracker;
 import me.jellysquid.mods.lithium.common.entity.movement_tracker.SectionedEntityMovementTracker;
 import me.jellysquid.mods.lithium.common.util.deduplication.LithiumInterner;
@@ -32,15 +33,19 @@ record Data(
             LithiumInterner<SectionedBlockChangeTracker> blockChangeTrackers,
 
             // Entity movement tracker deduplication
-            LithiumInterner<SectionedEntityMovementTracker<?, ?>> entityMovementTrackers
+            LithiumInterner<SectionedEntityMovementTracker<?, ?>> entityMovementTrackers,
+
+            // Block ChunkSection listeners
+            Long2ReferenceOpenHashMap<ChunkSectionChangeCallback> chunkSectionChangeCallbacks
     ) {
         public Data(World world) {
             this(
                     new Long2ReferenceOpenHashMap<>(),
                     world.getRegistryManager().getOptionalWrapper(RegistryKeys.BANNER_PATTERN).map(Raid::getOminousBanner).orElse(null),
                     new ReferenceOpenHashSet<>(),
                     new LithiumInterner<>(),
-                    new LithiumInterner<>()
+                    new LithiumInterner<>(),
+                    new Long2ReferenceOpenHashMap<>()
             );
         }
     }

diff --git a/src/main/java/me/jellysquid/mods/lithium/mixin/util/block_tracking/ChunkSectionMixin.java b/src/main/java/me/jellysquid/mods/lithium/mixin/util/block_tracking/ChunkSectionMixin.java
@@ -6,6 +6,7 @@
 import net.minecraft.block.BlockState;
 import net.minecraft.network.PacketByteBuf;
 import net.minecraft.util.math.ChunkSectionPos;
+import net.minecraft.world.World;
 import net.minecraft.world.chunk.ChunkSection;
 import net.minecraft.world.chunk.PalettedContainer;
 import org.spongepowered.asm.mixin.Final;
@@ -143,9 +144,12 @@ private void updateFlagCounters(int x, int y, int z, BlockState newState, boolea
     }
 
     @Override
-    public void lithium$addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker) {
+    public void lithium$addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker, long sectionPos, World world) {
         if (this.changeListener == null) {
-            this.changeListener = new ChunkSectionChangeCallback();
+            if (sectionPos == Long.MIN_VALUE || world == null) {
+                throw new IllegalArgumentException("Expected world and section pos during intialization!");
+            }
+            this.changeListener = ChunkSectionChangeCallback.create(sectionPos, world);
         }
 
         this.listeningMask = this.changeListener.addTracker(tracker, blockGroup);
@@ -161,8 +165,6 @@ private void updateFlagCounters(int x, int y, int z, BlockState newState, boolea
     @Override
     @Unique
     public void lithium$invalidateListeningSection(ChunkSectionPos sectionPos) {
-        //TODO call this on chunk unload. Entities should already be unloaded, but just to be safe, try to unregister too
-
         if ((this.listeningMask) != 0) {
             this.changeListener.onChunkSectionInvalidated(sectionPos);
         }

diff --git a/src/main/java/me/jellysquid/mods/lithium/mixin/util/block_tracking/package-info.java b/src/main/java/me/jellysquid/mods/lithium/mixin/util/block_tracking/package-info.java
@@ -2,7 +2,10 @@
         description = "Chunk sections count certain blocks inside them and provide a method to quickly check whether a" +
                 " chunk contains any of these blocks. Furthermore, chunk sections can notify registered listeners about" +
                 " certain blocks being placed or broken.",
-        depends = @MixinConfigDependency(dependencyPath = "mixin.util.data_storage")
+        depends = {
+                @MixinConfigDependency(dependencyPath = "mixin.util.data_storage"),
+                @MixinConfigDependency(dependencyPath = "mixin.util.chunk_status_tracking")
+        }
 )
 package me.jellysquid.mods.lithium.mixin.util.block_tracking;