diff --git a/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/PalettedContainerMixin.java b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/PalettedContainerMixin.java new file mode 100644 index 000000000..2be1a452f --- /dev/null +++ b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/PalettedContainerMixin.java @@ -0,0 +1,147 @@ +package net.caffeinemc.mods.lithium.mixin.chunk.serialization; + +import net.caffeinemc.mods.lithium.common.world.chunk.CompactingPackedIntegerArray; +import net.caffeinemc.mods.lithium.common.world.chunk.LithiumHashPalette; +import net.caffeinemc.mods.lithium.mixin.util.accessors.StrategyAccessor; +import net.minecraft.util.BitStorage; +import net.minecraft.util.SimpleBitStorage; +import net.minecraft.util.ZeroBitStorage; +import net.minecraft.world.level.chunk.*; +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.callback.CallbackInfo; + +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.stream.LongStream; + +/** + * Makes a number of patches to {@link PalettedContainer} to speed up integer array compaction. While I/O operations + * in Minecraft 1.15+ are handled off-thread, NBT serialization is not and happens on the main server thread. + */ +@Mixin(PalettedContainer.class) +public abstract class PalettedContainerMixin { + private static final ThreadLocal CACHED_ARRAY_4096 = ThreadLocal.withInitial(() -> new short[4096]); + private static final ThreadLocal CACHED_ARRAY_64 = ThreadLocal.withInitial(() -> new short[64]); + + @Shadow + public abstract void acquire(); + + @Shadow + protected abstract T get(int index); + + @Shadow + private volatile PalettedContainer.Data data; + + @Shadow + public abstract void release(); + + /** + * This patch incorporates a number of changes to significantly reduce the time needed to serialize. + * - If a palette only contains one entry, do not attempt to repack it + * - The packed integer array is iterated over using a specialized consumer instead of a naive for-loop. + * - A temporary fixed array is used to cache palette lookups and remaps while compacting a data array. + * - If the palette didn't change after compaction, avoid the step of re-packing the integer array and instead do + * a simple memory copy. + * + * @reason Optimize serialization + * @author JellySquid + */ + @Overwrite + public PalettedContainerRO.PackedData pack(Strategy strategy) { + this.acquire(); + + // The palette that will be serialized + LithiumHashPalette hashPalette = null; + Optional data = Optional.empty(); + List elements = null; + + final Palette palette = this.data.palette(); + final BitStorage storage = this.data.storage(); + if (storage instanceof ZeroBitStorage || palette.getSize() == 1) { + // If the palette only contains one entry, don't attempt to repack it. + elements = List.of(palette.valueFor(0)); + } else if (palette instanceof LithiumHashPalette lithiumHashPalette) { + hashPalette = lithiumHashPalette; + } + + if (elements == null) { + LithiumHashPalette compactedPalette = new LithiumHashPalette<>(storage.getBits()); + short[] array = this.getOrCreate(strategy.entryCount()); + + ((CompactingPackedIntegerArray) storage).lithium$compact(this.data.palette(), compactedPalette, array); + + // If the palette didn't change during compaction, do a simple copy of the data array + Configuration origConfig; + if (hashPalette != null && hashPalette.getSize() == compactedPalette.getSize() && + !(origConfig = ((StrategyAccessor) strategy).lithium$getConfigurationForPaletteSize(hashPalette.getSize())).alwaysRepack() && storage.getBits() == origConfig.bitsInStorage()) { // paletteSize can de-sync from palette - see https://github.com/CaffeineMC/lithium-fabric/issues/279 + data = this.asOptional(storage.getRaw().clone()); + elements = hashPalette.getElements(); + } else { + int bits = ((StrategyAccessor) strategy).lithium$getConfigurationForPaletteSize(compactedPalette.getSize()).bitsInStorage(); + if (bits != 0) { + // Re-pack the integer array as the palette has changed size + SimpleBitStorage copy = new SimpleBitStorage(bits, array.length); + for (int i = 0; i < array.length; ++i) { + copy.set(i, array[i]); + } + + // We don't need to clone the data array as we are the sole owner of it + data = this.asOptional(copy.getRaw()); + } + + elements = compactedPalette.getElements(); + } + } + + this.release(); + return new PalettedContainerRO.PackedData<>(elements, data); + } + + private Optional asOptional(long[] data) { + return Optional.of(Arrays.stream(data)); + } + + private short[] getOrCreate(int size) { + return switch (size) { + case 64 -> CACHED_ARRAY_64.get(); + case 4096 -> CACHED_ARRAY_4096.get(); + default -> new short[size]; + }; + } + + /** + * If we know the palette will contain a fixed number of elements, we can make a significant optimization by counting + * blocks with a simple array instead of a integer map. Since palettes make no guarantee that they are bounded, + * we have to try and determine for each implementation type how many elements there are. + * + * @author JellySquid + */ + @Inject(method = "count(Lnet/minecraft/world/level/chunk/PalettedContainer$CountConsumer;)V", at = @At("HEAD"), cancellable = true) + public void count(PalettedContainer.CountConsumer consumer, CallbackInfo ci) { + int len = this.data.palette().getSize(); + + // Do not allocate huge arrays if we're using a large palette + if (len > 4096) { + return; + } + + short[] counts = new short[len]; + + this.data.storage().getAll(i -> counts[i]++); + + for (int i = 0; i < counts.length; i++) { + T obj = this.data.palette().valueFor(i); + + if (obj != null) { + consumer.accept(obj, counts[i]); + } + } + + ci.cancel(); + } +} diff --git a/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/SimpleBitStorageMixin.java b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/SimpleBitStorageMixin.java new file mode 100644 index 000000000..9b1809de1 --- /dev/null +++ b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/SimpleBitStorageMixin.java @@ -0,0 +1,75 @@ +package net.caffeinemc.mods.lithium.mixin.chunk.serialization; + +import net.caffeinemc.mods.lithium.common.world.chunk.CompactingPackedIntegerArray; +import net.minecraft.util.SimpleBitStorage; +import net.minecraft.world.level.chunk.Palette; +import net.minecraft.world.level.chunk.PaletteResize; +import org.spongepowered.asm.mixin.Final; +import org.spongepowered.asm.mixin.Mixin; +import org.spongepowered.asm.mixin.Shadow; + +/** + * Extends {@link SimpleBitStorage} with a special compaction method defined in {@link CompactingPackedIntegerArray}. + */ +@Mixin(SimpleBitStorage.class) +public abstract class SimpleBitStorageMixin implements CompactingPackedIntegerArray { + @Shadow + @Final + private int size; + + @Shadow + @Final + private int bits; + + @Shadow + @Final + private long mask; + + @Shadow + @Final + private int valuesPerLong; + + @Shadow + @Final + private long[] data; + + @Override + public void lithium$compact(Palette srcPalette, Palette dstPalette, short[] out) { + if (this.size >= Short.MAX_VALUE) { + throw new IllegalStateException("Array too large"); + } + + if (this.size != out.length) { + throw new IllegalStateException("Array size mismatch"); + } + + PaletteResize noResizeExpected = PaletteResize.noResizeExpected(); + + short[] mappings = new short[(int) (this.mask + 1)]; + + int idx = 0; + + for (long word : this.data) { + long bits = word; + + for (int elementIdx = 0; elementIdx < this.valuesPerLong; ++elementIdx) { + int value = (int) (bits & this.mask); + int remappedId = mappings[value]; + + if (remappedId == 0) { + remappedId = dstPalette.idFor(srcPalette.valueFor(value), noResizeExpected) + 1; + mappings[value] = (short) remappedId; + } + + out[idx] = (short) (remappedId - 1); + bits >>= this.bits; + + ++idx; + + if (idx >= this.size) { + return; + } + } + } + } +} diff --git a/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/package-info.java b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/package-info.java new file mode 100644 index 000000000..c059dce18 --- /dev/null +++ b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/chunk/serialization/package-info.java @@ -0,0 +1,4 @@ +@MixinConfigOption(description = "Optimizes chunk palette compaction when serializing chunks") +package net.caffeinemc.mods.lithium.mixin.chunk.serialization; + +import net.caffeinemc.gradle.MixinConfigOption; \ No newline at end of file diff --git a/common/src/main/java/net/caffeinemc/mods/lithium/mixin/util/accessors/StrategyAccessor.java b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/util/accessors/StrategyAccessor.java new file mode 100644 index 000000000..01e1c627c --- /dev/null +++ b/common/src/main/java/net/caffeinemc/mods/lithium/mixin/util/accessors/StrategyAccessor.java @@ -0,0 +1,12 @@ +package net.caffeinemc.mods.lithium.mixin.util.accessors; + +import net.minecraft.world.level.chunk.Configuration; +import net.minecraft.world.level.chunk.Strategy; +import org.spongepowered.asm.mixin.Mixin; +import org.spongepowered.asm.mixin.gen.Invoker; + +@Mixin(Strategy.class) +public interface StrategyAccessor { + @Invoker("getConfigurationForPaletteSize") + Configuration lithium$getConfigurationForPaletteSize(int i); +} diff --git a/common/src/main/resources/lithium.mixins.json b/common/src/main/resources/lithium.mixins.json index 922a80991..29fa47e12 100644 --- a/common/src/main/resources/lithium.mixins.json +++ b/common/src/main/resources/lithium.mixins.json @@ -98,6 +98,8 @@ "chunk.palette.StrategyMixin", "chunk.palette.StrategyMixin$Strategy1Mixin", "chunk.palette.StrategyMixin$Strategy2Mixin", + "chunk.serialization.PalettedContainerMixin", + "chunk.serialization.SimpleBitStorageMixin", "collections.attributes.AttributeMapMixin", "collections.block_entity_tickers.LevelChunkMixin", "collections.brain.BrainMixin", @@ -178,6 +180,7 @@ "util.accessors.LevelAccessor", "util.accessors.PersistentEntitySectionManagerAccessor", "util.accessors.ServerLevelAccessor", + "util.accessors.StrategyAccessor", "util.block_entity_retrieval.LevelMixin", "util.block_tracking.BlockStateBaseMixin", "util.block_tracking.LevelChunkSectionMixin", diff --git a/lithium-fabric-mixin-config.md b/lithium-fabric-mixin-config.md index 1adcd98d1..5d9dcda38 100644 --- a/lithium-fabric-mixin-config.md +++ b/lithium-fabric-mixin-config.md @@ -214,6 +214,10 @@ Skip bounds validation when accessing blocks (default: `true`) Replaces the vanilla hash palette with an optimized variant +### `mixin.chunk.serialization` +(default: `true`) +Optimizes chunk palette compaction when serializing chunks + ### `mixin.collections` (default: `true`) Various collection optimizations