Slightly refactor ChunkSectionMixin, partially implement block listen…

…ing invalidation on chunk unload
CaffeineMC · Apr 8, 2024 · 11eb628 · 11eb628
1 parent be8ee4d
commit 11eb628
Show file tree

Hide file tree

Showing 8 changed files with 71 additions and 34 deletions.
diff --git a/src/main/java/me/jellysquid/mods/lithium/common/ai/pathing/PathNodeCache.java b/src/main/java/me/jellysquid/mods/lithium/common/ai/pathing/PathNodeCache.java
@@ -43,7 +43,7 @@ public static boolean isSectionSafeAsNeighbor(ChunkSection section) {
         }
 
         if (BlockStateFlags.ENABLED) {
-            return !((BlockCountingSection) section).mayContainAny(BlockStateFlags.PATH_NOT_OPEN);
+            return !((BlockCountingSection) section).lithium$mayContainAny(BlockStateFlags.PATH_NOT_OPEN);
         }
         return !isChunkSectionDangerousNeighbor(section);
     }

diff --git a/src/main/java/me/jellysquid/mods/lithium/common/block/BlockCountingSection.java b/src/main/java/me/jellysquid/mods/lithium/common/block/BlockCountingSection.java
@@ -1,5 +1,5 @@
 package me.jellysquid.mods.lithium.common.block;
 
 public interface BlockCountingSection {
-    boolean mayContainAny(TrackedBlockStatePredicate trackedBlockStatePredicate);
+    boolean lithium$mayContainAny(TrackedBlockStatePredicate trackedBlockStatePredicate);
 }
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
@@ -1,9 +1,13 @@
 package me.jellysquid.mods.lithium.common.block;
 
 import me.jellysquid.mods.lithium.common.entity.block_tracking.SectionedBlockChangeTracker;
+import net.minecraft.util.math.ChunkSectionPos;
 
 public interface BlockListeningSection {
 
-    void addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker);
-    void removeFromCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker);
+    void lithium$addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker);
+
+    void lithium$removeFromCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker);
+
+    void lithium$invalidateListeningSection(ChunkSectionPos sectionPos);
 }
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
@@ -3,6 +3,7 @@
 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 net.minecraft.util.math.ChunkSectionPos;
 
 import java.util.ArrayList;
 
@@ -23,7 +24,7 @@ public short onBlockChange(int flagIndex, BlockListeningSection section) {
         for (int i = 0; i < sectionedBlockChangeTrackers.size(); i++) {
             sectionedBlockChangeTrackers.get(i).setChanged(section);
         }
-        this.listeningMask &= ~(1 << flagIndex);
+        this.listeningMask &= (short) ~(1 << flagIndex);
 
         return this.listeningMask;
     }
@@ -36,7 +37,7 @@ public short addTracker(SectionedBlockChangeTracker tracker, ListeningBlockState
         }
         sectionedBlockChangeTrackers.add(tracker);
 
-        this.listeningMask |= (1 << blockGroupIndex);
+        this.listeningMask |= (short) (1 << blockGroupIndex);
         return this.listeningMask;
     }
 
@@ -46,9 +47,23 @@ public short removeTracker(SectionedBlockChangeTracker tracker, ListeningBlockSt
         if (sectionedBlockChangeTrackers != null) {
             sectionedBlockChangeTrackers.remove(tracker);
             if (sectionedBlockChangeTrackers.isEmpty()) {
-                this.listeningMask &= ~(1 << blockGroup.getIndex());
+                this.listeningMask &= (short) ~(1 << blockGroup.getIndex());
             }
         }
         return this.listeningMask;
     }
+
+    public void onChunkSectionInvalidated(ChunkSectionPos sectionPos) {
+        for (int flagIndex = 0; flagIndex < this.trackers.length; flagIndex++) {
+            ArrayList<SectionedBlockChangeTracker> sectionedBlockChangeTrackers = this.trackers[flagIndex];
+            this.trackers[flagIndex] = null;
+            if (sectionedBlockChangeTrackers != null) {
+                //noinspection ForLoopReplaceableByForEach
+                for (int i = 0; i < sectionedBlockChangeTrackers.size(); i++) {
+                    sectionedBlockChangeTrackers.get(i).onChunkSectionInvalidated(sectionPos);
+                }
+            }
+        }
+        this.listeningMask = 0;
+    }
 }
diff --git a/.../me/jellysquid/mods/lithium/common/entity/block_tracking/SectionedBlockChangeTracker.java b/.../me/jellysquid/mods/lithium/common/entity/block_tracking/SectionedBlockChangeTracker.java
@@ -73,10 +73,12 @@ public void register() {
                         ChunkSection section = sectionArray[Pos.SectionYIndex.fromSectionCoord(trackedSections.world(), y)];
 
                         BlockListeningSection blockListeningSection = (BlockListeningSection) section;
-                        blockListeningSection.addToCallback(this.blockGroup, this);
+                        blockListeningSection.lithium$addToCallback(this.blockGroup, this);
                     }
                 }
             }
+            this.isListeningToAll = (this.sectionsNotListeningTo == null || this.sectionsNotListeningTo.isEmpty())
+                    && (this.sectionsUnsubscribed == null || this.sectionsUnsubscribed.isEmpty());
             this.setChanged(this.getWorldTime());
         }
         this.timesRegistered++;
@@ -103,7 +105,7 @@ public void unregister() {
                     ChunkSection section = sectionArray[Pos.SectionYIndex.fromSectionCoord(world, y)];
 
                     BlockListeningSection blockListeningSection = (BlockListeningSection) section;
-                    blockListeningSection.removeFromCallback(this.blockGroup, this);
+                    blockListeningSection.lithium$removeFromCallback(this.blockGroup, this);
                 }
             }
         }
@@ -128,15 +130,15 @@ public void listenToAllSections() {
                 }
                 ChunkSection section = chunk.getSectionArray()[Pos.SectionYIndex.fromSectionCoord(this.trackedWorldSections.world(), chunkSectionPos.getY())];
                 BlockListeningSection blockListeningSection = (BlockListeningSection) section;
-                blockListeningSection.addToCallback(this.blockGroup, this);
+                blockListeningSection.lithium$addToCallback(this.blockGroup, this);
             }
         }
         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.addToCallback(this.blockGroup, this);
+                blockListeningSection.lithium$addToCallback(this.blockGroup, this);
             }
         }
         this.isListeningToAll = true;
@@ -192,4 +194,13 @@ public boolean equals(Object obj) {
     public int hashCode() {
         return this.getClass().hashCode() ^ this.trackedWorldSections.hashCode() ^ this.blockGroup.hashCode();
     }
+
+    public void onChunkSectionInvalidated(ChunkSectionPos sectionPos) {
+        if (this.sectionsNotListeningTo == null) {
+            this.sectionsNotListeningTo = new ArrayList<>();
+        }
+        this.sectionsNotListeningTo.add(sectionPos);
+        this.setChanged(this.getWorldTime());
+        this.isListeningToAll = false;
+    }
 }
diff --git a/...va/me/jellysquid/mods/lithium/common/entity/movement/ChunkAwareBlockCollisionSweeper.java b/...va/me/jellysquid/mods/lithium/common/entity/movement/ChunkAwareBlockCollisionSweeper.java
@@ -290,7 +290,7 @@ private static int expandMax(int coord) {
     private static boolean hasChunkSectionOversizedBlocks(Chunk chunk, int chunkY) {
         if (BlockStateFlags.ENABLED) {
             ChunkSection section = chunk.getSectionArray()[chunkY];
-            return section != null && ((BlockCountingSection) section).mayContainAny(BlockStateFlags.OVERSIZED_SHAPE);
+            return section != null && ((BlockCountingSection) section).lithium$mayContainAny(BlockStateFlags.OVERSIZED_SHAPE);
         }
         return true; //like vanilla, assume that a chunk section has oversized blocks, when the section mixin isn't loaded
     }

diff --git a/src/main/java/me/jellysquid/mods/lithium/mixin/entity/collisions/fluid/EntityMixin.java b/src/main/java/me/jellysquid/mods/lithium/mixin/entity/collisions/fluid/EntityMixin.java
@@ -28,7 +28,7 @@ public abstract class EntityMixin implements FluidCachingEntity {
     public abstract Box getBoundingBox();
 
     @Shadow
-    public World world;
+    private World world;
 
     @Shadow
     protected Object2DoubleMap<TagKey<Fluid>> fluidHeight;
@@ -63,7 +63,7 @@ public void tryShortcutFluidPushing(TagKey<Fluid> tag, double speed, CallbackInf
                 Chunk chunk = this.world.getChunk(chunkX, chunkZ);
                 for (int chunkYIndex = chunkYIndex1; chunkYIndex <= chunkYIndex2; chunkYIndex++) {
                     ChunkSection section = chunk.getSectionArray()[chunkYIndex];
-                    if (((BlockCountingSection) section).mayContainAny(blockStateFlag)) {
+                    if (((BlockCountingSection) section).lithium$mayContainAny(blockStateFlag)) {
                         //fluid found, cannot skip code
                         return;
                     }

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
@@ -5,6 +5,7 @@
 import me.jellysquid.mods.lithium.common.entity.block_tracking.SectionedBlockChangeTracker;
 import net.minecraft.block.BlockState;
 import net.minecraft.network.PacketByteBuf;
+import net.minecraft.util.math.ChunkSectionPos;
 import net.minecraft.world.chunk.ChunkSection;
 import net.minecraft.world.chunk.PalettedContainer;
 import org.spongepowered.asm.mixin.Final;
@@ -32,17 +33,31 @@ public abstract class ChunkSectionMixin implements BlockCountingSection, BlockLi
     private PalettedContainer<BlockState> blockStateContainer;
     @Unique
     private short[] countsByFlag = null;
+    @Unique
     private ChunkSectionChangeCallback changeListener;
+    @Unique
     private short listeningMask;
 
+    @Unique
+    private static void addToFlagCount(short[] countsByFlag, BlockState state, short change) {
+        int flags = ((BlockStateFlagHolder) state).lithium$getAllFlags();
+        int i;
+        while ((i = Integer.numberOfTrailingZeros(flags)) < 32 && i < countsByFlag.length) {
+            //either count up by one (prevFlag not set) or down by one (prevFlag set)
+            countsByFlag[i] += change;
+            flags &= ~(1 << i);
+        }
+    }
+
     @Override
-    public boolean mayContainAny(TrackedBlockStatePredicate trackedBlockStatePredicate) {
+    public boolean lithium$mayContainAny(TrackedBlockStatePredicate trackedBlockStatePredicate) {
         if (this.countsByFlag == null) {
             fastInitClientCounts();
         }
         return this.countsByFlag[trackedBlockStatePredicate.getIndex()] != (short) 0;
     }
 
+    @Unique
     private void fastInitClientCounts() {
         this.countsByFlag = new short[BlockStateFlags.NUM_TRACKED_FLAGS];
         for (TrackedBlockStatePredicate trackedBlockStatePredicate : BlockStateFlags.TRACKED_FLAGS) {
@@ -64,20 +79,10 @@ private void fastInitClientCounts() {
     private void initFlagCounters(PalettedContainer<BlockState> palettedContainer, PalettedContainer.Counter<BlockState> consumer) {
         palettedContainer.count((state, count) -> {
             consumer.accept(state, count);
-            addToFlagCount(this.countsByFlag, state, count);
+            addToFlagCount(this.countsByFlag, state, (short) count);
         });
     }
 
-    private static void addToFlagCount(short[] countsByFlag, BlockState state, int change) {
-        int flags = ((BlockStateFlagHolder) state).lithium$getAllFlags();
-        int i;
-        while ((i = Integer.numberOfTrailingZeros(flags)) < 32 && i < countsByFlag.length) {
-            //either count up by one (prevFlag not set) or down by one (prevFlag set)
-            countsByFlag[i] += change;
-            flags &= ~(1 << i);
-        }
-    }
-
     @Inject(method = "calculateCounts()V", at = @At("HEAD"))
     private void createFlagCounters(CallbackInfo ci) {
         this.countsByFlag = new short[BlockStateFlags.NUM_TRACKED_FLAGS];
@@ -122,7 +127,7 @@ private void updateFlagCounters(int x, int y, int z, BlockState newState, boolea
             int flagBit = 1 << flagIndex;
             //either count up by one (prevFlag not set) or down by one (prevFlag set)
             if ((flagsXOR & flagBit) != 0) {
-                countsByFlag[flagIndex] += 1 - (((prevFlags >>> flagIndex) & 1) << 1);
+                countsByFlag[flagIndex] += (short) (1 - (((prevFlags >>> flagIndex) & 1) << 1));
             }
             if ((this.listeningMask & flagBit) != 0) {
                 this.listeningMask = this.changeListener.onBlockChange(flagIndex, this);
@@ -132,7 +137,7 @@ private void updateFlagCounters(int x, int y, int z, BlockState newState, boolea
     }
 
     @Override
-    public void addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker) {
+    public void lithium$addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker) {
         if (this.changeListener == null) {
             this.changeListener = new ChunkSectionChangeCallback();
         }
@@ -141,17 +146,19 @@ public void addToCallback(ListeningBlockStatePredicate blockGroup, SectionedBloc
     }
 
     @Override
-    public void removeFromCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker) {
+    public void lithium$removeFromCallback(ListeningBlockStatePredicate blockGroup, SectionedBlockChangeTracker tracker) {
         if (this.changeListener != null) {
             this.listeningMask = this.changeListener.removeTracker(tracker, blockGroup);
         }
     }
 
-    private boolean isListening(ListeningBlockStatePredicate blockGroup) {
-        return (this.listeningMask & (1 << blockGroup.getIndex())) != 0;
-    }
+    @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
 
-    public void invalidateSection() {
-        //TODO on section unload, unregister all kinds of stuff
+        if ((this.listeningMask) != 0) {
+            this.changeListener.onChunkSectionInvalidated(sectionPos);
+        }
     }
 }