Optimize biome translation to take advantage of palettes

connector/src/main/java/org/geysermc/connector/network/translators/world/BiomeTranslator.java

@@ -25,20 +25,25 @@
 
 package org.geysermc.connector.network.translators.world;
 
+import com.github.steveice10.mc.protocol.data.game.chunk.BitStorage;
 import com.github.steveice10.mc.protocol.data.game.chunk.DataPalette;
+import com.github.steveice10.mc.protocol.data.game.chunk.palette.GlobalPalette;
+import com.github.steveice10.mc.protocol.data.game.chunk.palette.Palette;
 import com.github.steveice10.mc.protocol.data.game.chunk.palette.SingletonPalette;
 import com.github.steveice10.opennbt.tag.builtin.*;
 import it.unimi.dsi.fastutil.ints.Int2IntMap;
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.ints.IntList;
 import it.unimi.dsi.fastutil.ints.IntLists;
 import org.geysermc.connector.network.session.GeyserSession;
 import org.geysermc.connector.network.translators.world.chunk.BlockStorage;
 import org.geysermc.connector.network.translators.world.chunk.GeyserChunkSection;
+import org.geysermc.connector.network.translators.world.chunk.bitarray.BitArray;
+import org.geysermc.connector.network.translators.world.chunk.bitarray.BitArrayVersion;
 import org.geysermc.connector.network.translators.world.chunk.bitarray.SingletonBitArray;
 import org.geysermc.connector.registry.Registries;
 import org.geysermc.connector.utils.MathUtils;
 
-// Based off of ProtocolSupport's LegacyBiomeData.java:
-// https://github.com/ProtocolSupport/ProtocolSupport/blob/b2cad35977f3fcb65bee57b9e14fc9c975f71d32/src/protocolsupport/protocol/typeremapper/legacy/LegacyBiomeData.java
 // Array index formula by https://wiki.vg/Chunk_Format
 public class BiomeTranslator {
 
@@ -87,11 +92,41 @@ public class BiomeTranslator {
         // As of 1.17.10: the client expects the same format as a chunk but filled with biomes
         // As of 1.18 this is the same as Java Edition
 
-        if (biomeData.getPalette() instanceof SingletonPalette palette) {
+        Palette palette = biomeData.getPalette();
+        if (palette instanceof SingletonPalette) {
             int biomeId = biomeTranslations.get(palette.idToState(0));
             return new BlockStorage(SingletonBitArray.INSTANCE, IntLists.singleton(biomeId));
         } else {
-            BlockStorage storage = new BlockStorage(0);
+            BlockStorage storage;
+            if (!(palette instanceof GlobalPalette)) {
+                // Prevent resizing by allocating what we can ahead of time
+                BitStorage bitStorage = biomeData.getStorage();
+                int size = palette.size();
+                BitArray bitArray = BitArrayVersion.forBitsCeil(bitStorage.getBitsPerEntry())
+                        .createArray(BlockStorage.SIZE);
+
+                IntList bedrockPalette = new IntArrayList(size);
+
+                for (int i = 0; i < size; i++) {
+                    int javaId = palette.idToState(i);
+                    bedrockPalette.add(biomeTranslations.get(javaId));
+                }
+
+                // Each section of biome corresponding to a chunk section contains 4 * 4 * 4 entries
+                for (int i = 0; i < 64; i++) {
+                    int idx = bitStorage.get(i);
+                    int x = i & 3;
+                    int y = (i >> 4) & 3;
+                    int z = (i >> 2) & 3;
+                    // Convert biome coordinates into block coordinates
+                    // Bedrock expects a full 4096 blocks
+                    multiplyIdToStorage(bitArray, idx, x, y, z);
+                }
+
+                storage = new BlockStorage(bitArray, bedrockPalette);
+            } else {
+                storage = new BlockStorage(0);
+                BitArray bitArray = storage.getBitArray();
 
                 // Each section of biome corresponding to a chunk section contains 4 * 4 * 4 entries
                 for (int i = 0; i < 64; i++) {
@@ -104,15 +139,20 @@ public class BiomeTranslator {
                     int idx = storage.idFor(biomeId);
                     // Convert biome coordinates into block coordinates
                     // Bedrock expects a full 4096 blocks
-                for (int blockX = x << 2; blockX < (x << 2) + 4; blockX++) {
+                    multiplyIdToStorage(bitArray, idx, x, y, z);
-                    for (int blockZ = z << 2; blockZ < (z << 2) + 4; blockZ++) {
-                        for (int blockY = y << 2; blockY < (y << 2) + 4; blockY++) {
-                            storage.getBitArray().set(GeyserChunkSection.blockPosition(blockX, blockY, blockZ), idx);
-                        }
-                    }
                 }
             }
             return storage;
         }
     }
+
+    private static void multiplyIdToStorage(BitArray bitArray, int idx, int x, int y, int z) {
+        for (int blockX = x << 2; blockX < (x << 2) + 4; blockX++) {
+            for (int blockZ = z << 2; blockZ < (z << 2) + 4; blockZ++) {
+                for (int blockY = y << 2; blockY < (y << 2) + 4; blockY++) {
+                    bitArray.set(GeyserChunkSection.blockPosition(blockX, blockY, blockZ), idx);
+                }
+            }
+        }
+    }
 }