Add fireblanket optimization

README.md

@@ -3,4 +3,6 @@
 A mod made to assist in the Steam 'n' Rails modpack by fixing bugs, lag from unneeded things and small tweaks or game additions.
 
 ## License
-Railways Tweaks is licensed under the MIT license. See [LICENSE](LICENSE) for more information.
+Railways Tweaks is licensed under the AGPL license. See [LICENSE](LICENSE) for more information.
+
+Code for LevelChunkSectionMixin is from FireBlanket, which is licensed under AGPL. See [FireBlanket's License](https://github.com/ModFest/fireblanket/blob/b1180ea2cee7153aeb3f408bda946e1a52162ef7/LICENSE.AGPL.md)

src/main/java/dev/ithundxr/railwaystweaks/RailwayTweaksMixinPlugin.java

@@ -0,0 +1,36 @@
+package dev.ithundxr.railwaystweaks;
+
+import org.objectweb.asm.tree.ClassNode;
+import org.spongepowered.asm.mixin.extensibility.IMixinConfigPlugin;
+import org.spongepowered.asm.mixin.extensibility.IMixinInfo;
+
+import java.util.List;
+import java.util.Set;
+
+public class RailwayTweaksMixinPlugin implements IMixinConfigPlugin {
+    private static final boolean FLATTEN_CHUNK_PALETTES = Boolean.getBoolean("railwayTweaks.fireblanket.flattenChunkPalettes");
+
+    @Override
+    public void onLoad(String mixinPackage) { } // NO-OP
+
+    @Override
+    public String getRefMapperConfig() { return null; } // DEFAULT
+
+    @Override
+    public boolean shouldApplyMixin(String targetClassName, String mixinClassName) {
+        if (mixinClassName.contains("LevelChunkSectionMixin")) return FLATTEN_CHUNK_PALETTES;
+        return true;
+    }
+
+    @Override
+    public void acceptTargets(Set<String> myTargets, Set<String> otherTargets) { } // NO-OP
+
+    @Override
+    public List<String> getMixins() { return null; } // DEFAULT
+
+    @Override
+    public void preApply(String targetClassName, ClassNode targetClass, String mixinClassName, IMixinInfo mixinInfo) { } // NO-OP
+
+    @Override
+    public void postApply(String targetClassName, ClassNode targetClass, String mixinClassName, IMixinInfo mixinInfo) { } // NO-OP
+}

src/main/java/dev/ithundxr/railwaystweaks/mixin/LevelChunkSectionMixin.java

@@ -0,0 +1,240 @@
+// Licensed under AGPL, https://github.com/ModFest/fireblanket/blob/b1180ea2cee7153aeb3f408bda946e1a52162ef7/LICENSE.AGPL.md
+
+package dev.ithundxr.railwaystweaks.mixin;
+
+import dev.ithundxr.railwaystweaks.mixinsupport.FlatBlockstateArray;
+import net.minecraft.CrashReport;
+import net.minecraft.CrashReportCategory;
+import net.minecraft.ReportedException;
+import net.minecraft.network.FriendlyByteBuf;
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.state.BlockState;
+import net.minecraft.world.level.chunk.LevelChunkSection;
+import net.minecraft.world.level.chunk.PalettedContainer;
+import net.minecraft.world.level.material.FluidState;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Overwrite;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.Redirect;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import java.util.Arrays;
+import java.util.BitSet;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+/**
+ * Augments the chunk palette to primarily use a flat array to store blockstate ids, speeding up blockstate access and
+ * setting. As many minecraft structures rely on the palette, when the palette is queried, for saving or networking,
+ * it flushes the data from the flat array back into the palette.
+ */
+@Mixin(LevelChunkSection.class)
+public abstract class LevelChunkSectionMixin {
+    // Will be real due to mixin plugin
+    private static final int MASK_BITS = 1048575;
+
+    @Shadow @Final private PalettedContainer<BlockState> states;
+    @Shadow private short nonEmptyBlockCount;
+    @Shadow private short tickingBlockCount;
+    @Shadow private short tickingFluidCount;
+
+    // 20 bits per block, so 3 blocks per long. ceil(4096/3) --> 1366
+    private final long[] railwayTweaks$denseBlockStorage = new long[1366];
+    private final BitSet railwayTweaks$dirty = new BitSet(4096);
+
+    private final AtomicLong railwayTweaks$stamp = new AtomicLong();
+
+    @Redirect(method = "<init>(Lnet/minecraft/world/level/chunk/PalettedContainer;Lnet/minecraft/world/level/chunk/PalettedContainerRO;)V", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/level/chunk/LevelChunkSection;recalcBlockCounts()V"))
+    private void railwayTweaks$setupState(LevelChunkSection instance) {
+        PalettedContainer<BlockState> container = this.states;
+
+        railwayTweaks$applyFromPalette(container);
+    }
+
+    private void railwayTweaks$applyFromPalette(PalettedContainer<BlockState> container) {
+        for (int y = 0; y < 16; y++) {
+            for (int x = 0; x < 16; x++) {
+                for (int z = 0; z < 16; z++) {
+                    setBlockState(x, y, z, container.get(x, y, z), false);
+                }
+            }
+        }
+
+        railwayTweaks$dirty.clear();
+    }
+
+    /**
+     * @author Jasmine
+     * @reason Optimized with flat array
+     */
+    @Overwrite
+    public BlockState setBlockState(int x, int y, int z, BlockState state, boolean lock) {
+        // Calculate math needed for both stages before the barrier
+        int rawIdx = arrayIndex(x, y, z);
+        int arrIdx = rawIdx / 3;
+        int shlIdx = rawIdx % 3;
+        int shift = 20 * shlIdx;
+
+        // Grab a stamp to compare against after we're done writing, just in case anything weird happens
+        long stamp = this.railwayTweaks$stamp.incrementAndGet();
+
+        BlockState oldState;
+        try {
+            // Get the old state that we already see
+            long oldBits = this.railwayTweaks$denseBlockStorage[arrIdx];
+            oldState = FlatBlockstateArray.FROM_ID[((int) (oldBits >>> shift) & MASK_BITS)];
+
+            // Make data for the new state
+            long newId = Block.BLOCK_STATE_REGISTRY.getId(state);
+            long newBitsIn = newId << shift;
+            long mask = (long) MASK_BITS << shift;
+
+            // Replace old location with zeros, then apply the new bits
+            long newBits = (oldBits & ~mask) | (newBitsIn & mask);
+
+            // Commit to memory, mark dirty
+            this.railwayTweaks$denseBlockStorage[arrIdx] = newBits;
+            railwayTweaks$dirty.set(rawIdx);
+        } finally {
+            // Nightmare scenario. The stamp is not the same as the one we obtained at the start, so
+            // we were probably written to concurrently. This is bad, we need to stop immediately.
+            if (this.railwayTweaks$stamp.get() != stamp) {
+                Map<Thread, StackTraceElement[]> traces = Thread.getAllStackTraces();
+                String dumps = traces.entrySet().stream()
+                        .map(LevelChunkSectionMixin::formatThreadDump)
+                        .collect(Collectors.joining("\n"));
+
+                String error = "Accessing ChunkSection from multiple threads!";
+                CrashReport crashReport = new CrashReport(error, new IllegalStateException(error));
+                CrashReportCategory crashReportCategory = crashReport.addCategory("Thread dumps");
+                crashReportCategory.setDetail("Thread dumps", dumps);
+                throw new ReportedException(crashReport);
+            }
+        }
+
+        // Vanilla counting logic
+
+        FluidState fluidState = oldState.getFluidState();
+        FluidState fluidState2 = state.getFluidState();
+        if (!oldState.isAir()) {
+            --this.nonEmptyBlockCount;
+            if (oldState.isRandomlyTicking()) {
+                --this.tickingBlockCount;
+            }
+        }
+
+        if (!fluidState.isEmpty()) {
+            --this.tickingFluidCount;
+        }
+
+        if (!state.isAir()) {
+            ++this.nonEmptyBlockCount;
+            if (state.isRandomlyTicking()) {
+                ++this.tickingBlockCount;
+            }
+        }
+
+        if (!fluidState2.isEmpty()) {
+            ++this.tickingFluidCount;
+        }
+
+        return oldState;
+    }
+
+    private static String formatThreadDump(Map.Entry<Thread, StackTraceElement[]> e) {
+        return e.getKey().getName() + ": \n\tat " + Arrays.stream(e.getValue()).map(Object::toString).collect(Collectors.joining("\n\tat "));
+    }
+
+    /**
+     * @author Jasmine
+     * @reason Optimized with flat array
+     */
+    @Overwrite
+    public BlockState getBlockState(int x, int y, int z) {
+        int rawIdx = arrayIndex(x, y, z);
+
+        // Optimized divmod routine
+        int divRes = (rawIdx * 0xAAAB) >>> 17;             // rawIdx / 3;
+        int modRes = -((divRes + (divRes << 1)) - rawIdx);  // rawIdx % 3;
+
+        long rawVal = (this.railwayTweaks$denseBlockStorage[divRes] >>> (20L * modRes)) & MASK_BITS;
+
+        return FlatBlockstateArray.FROM_ID[((int) rawVal)];
+    }
+
+    private static int arrayIndex(int x, int y, int z) {
+        // Y before XZ is typically the access pattern found in extended periods of blockstate lookup
+        return y * 256 + x * 16 + z;
+    }
+
+    /**
+     * @author Jasmine
+     * @reason Optimized with flat array. Vanilla duplicates the getBlockState logic, whereas here just calls the
+     * method for simplicity
+     */
+    @Overwrite
+    public FluidState getFluidState(int x, int y, int z) {
+        return getBlockState(x, y, z).getFluidState();
+    }
+
+    /**
+     * @author Jasmine
+     * @reason Flush all updates to the container
+     */
+    @Overwrite
+    public PalettedContainer<BlockState> getStates() {
+        if (railwayTweaks$dirty.cardinality() > 0) {
+            for (int i = 0; i < 4096; i++) {
+                if (railwayTweaks$dirty.get(i)) {
+                    int y = (i >> 8) & 15;
+                    int x = (i >> 4) & 15;
+                    int z = (i >> 0) & 15;
+
+                    this.states.getAndSetUnchecked(x, y, z, getBlockState(x, y, z));
+                }
+            }
+        }
+
+        railwayTweaks$dirty.clear();
+
+        return this.states;
+    }
+
+    // Originally this was one injector, but broke horribly, so we're doing it the spacious way instead
+    @Inject(method = "write", at = @At("HEAD"))
+    private void railwayTweaks$resetUnderlyingStateToPacket(FriendlyByteBuf buf, CallbackInfo ci) {
+        getStates();
+    }
+
+    @Inject(method = "getSerializedSize", at = @At("HEAD"))
+    private void railwayTweaks$resetUnderlyingStateGetPacketSize(CallbackInfoReturnable<Integer> cir) {
+        getStates();
+    }
+
+    @Inject(method = "maybeHas", at = @At("HEAD"))
+    private void railwayTweaks$resetUnderlyingStateHasAny(Predicate<BlockState> predicate, CallbackInfoReturnable<Boolean> cir) {
+        getStates();
+    }
+
+    // Dear god please no one use this on their client. Support is provided for completeness.
+    @Inject(method = "read", at = @At("TAIL"))
+    private void railwayTweaks$resetForPacketBadTerrible(FriendlyByteBuf buf, CallbackInfo ci) {
+        railwayTweaks$applyFromPalette(this.states);
+    }
+
+    /**
+     * @author Jasmine
+     * @reason We don't really need this
+     */
+    @Overwrite
+    public void recalcBlockCounts() {
+        // Only ever used in one place, which is redirected, so I figure it's better to implement this on a need basis.
+        throw new UnsupportedOperationException("Not implemented");
+    }
+}

src/main/java/dev/ithundxr/railwaystweaks/mixinsupport/FlatBlockstateArray.java

@@ -0,0 +1,18 @@
+package dev.ithundxr.railwaystweaks.mixinsupport;
+
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.state.BlockState;
+
+public class FlatBlockstateArray {
+	public static BlockState[] FROM_ID;
+
+	public static void apply() {
+		int size = Block.BLOCK_STATE_REGISTRY.size();
+
+		FROM_ID = new BlockState[size];
+		int i = 0;
+		for (BlockState b : Block.BLOCK_STATE_REGISTRY) {
+			FROM_ID[i++] = b;
+		}
+	}
+}

src/main/resources/railwaystweaks.mixins.json

@@ -2,6 +2,7 @@
   "required": true,
   "minVersion": "0.8",
   "package": "dev.ithundxr.railwaystweaks.mixin",
+  "plugin": "dev.ithundxr.railwaystweaks.RailwayTweaksMixinPlugin",
   "compatibilityLevel": "JAVA_17",
   "mixins": [
     "DedicatedServerMixin",