Add waterlogging support

connector/src/main/java/org/geysermc/connector/network/translators/block/BlockTranslator.java

@@ -34,10 +34,7 @@ import com.nukkitx.nbt.tag.CompoundTag;
 import com.nukkitx.nbt.tag.ListTag;
 import gnu.trove.map.TObjectIntMap;
 import gnu.trove.map.hash.TObjectIntHashMap;
-import it.unimi.dsi.fastutil.ints.Int2IntMap;
+import it.unimi.dsi.fastutil.ints.*;
-import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
-import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
-import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import org.geysermc.connector.console.GeyserLogger;
 import org.geysermc.connector.utils.Toolbox;
 
@@ -47,11 +44,11 @@ import java.util.*;
 public class BlockTranslator {
     public static final ListTag<CompoundTag> BLOCKS;
     public static final BlockState AIR = new BlockState(0);
+    public static final int BEDROCK_WATER_ID;
 
     private static final Int2IntMap JAVA_TO_BEDROCK_BLOCK_MAP = new Int2IntOpenHashMap();
     private static final Int2ObjectMap<BlockState> BEDROCK_TO_JAVA_BLOCK_MAP = new Int2ObjectOpenHashMap<>();
-    private static final Int2IntMap JAVA_TO_BEDROCK_LIQUID_MAP = new Int2IntOpenHashMap();
+    private static final IntSet WATERLOGGED = new IntOpenHashSet();
-    private static final Int2ObjectMap<BlockState> BEDROCK_TO_JAVA_LIQUID_MAP = new Int2ObjectOpenHashMap<>();
 
     private static final int BLOCK_STATE_VERSION = 17760256;
 
@@ -84,6 +81,7 @@ public class BlockTranslator {
         TObjectIntMap<CompoundTag> addedStatesMap = new TObjectIntHashMap<>(512, 0.5f, -1);
         List<CompoundTag> paletteList = new ArrayList<>();
 
+        int waterRuntimeId = -1;
         int javaRuntimeId = -1;
         int bedrockRuntimeId = 0;
         Iterator<Map.Entry<String, JsonNode>> blocksIterator = blocks.fields();
@@ -93,6 +91,18 @@ public class BlockTranslator {
             String javaId = entry.getKey();
             CompoundTag blockTag = buildBedrockState(entry.getValue());
 
+            if ("minecraft:water[level=0]".equals(javaId)) {
+                waterRuntimeId = bedrockRuntimeId;
+            }
+            boolean waterlogged = entry.getValue().has("waterlogged") && entry.getValue().get("waterlogged").booleanValue();
+
+            if (waterlogged) {
+                BEDROCK_TO_JAVA_BLOCK_MAP.putIfAbsent(bedrockRuntimeId | 1 << 31, new BlockState(javaRuntimeId));
+                WATERLOGGED.add(javaRuntimeId);
+            } else {
+                BEDROCK_TO_JAVA_BLOCK_MAP.putIfAbsent(bedrockRuntimeId, new BlockState(javaRuntimeId));
+            }
+
             CompoundTag runtimeTag = blockStateMap.remove(blockTag);
             if (runtimeTag != null) {
                 addedStatesMap.put(blockTag, bedrockRuntimeId);
@@ -107,11 +117,15 @@ public class BlockTranslator {
                 continue;
             }
             JAVA_TO_BEDROCK_BLOCK_MAP.put(javaRuntimeId, bedrockRuntimeId);
-            BEDROCK_TO_JAVA_BLOCK_MAP.put(bedrockRuntimeId, new BlockState(javaRuntimeId));
 
             bedrockRuntimeId++;
         }
 
+        if (waterRuntimeId == -1) {
+            throw new AssertionError("Unable to find water in palette");
+        }
+        BEDROCK_WATER_ID = waterRuntimeId;
+
         paletteList.addAll(blockStateMap.values()); // Add any missing mappings that could crash the client
 
         BLOCKS = new ListTag<>("", CompoundTag.class, paletteList);
@@ -161,11 +175,11 @@ public class BlockTranslator {
         return BEDROCK_TO_JAVA_BLOCK_MAP.get(bedrockId);
     }
 
-    public static int getBedrockWaterLoggedId(BlockState state) {
+    public static boolean isWaterlogged(BlockState state) {
-        return JAVA_TO_BEDROCK_LIQUID_MAP.getOrDefault(state.getId(), -1);
+        return WATERLOGGED.contains(state.getId());
     }
 
-    public static BlockState getJavaLiquidState(int bedrockId) {
+    public static BlockState getJavaWaterloggedState(int bedrockId) {
-        return BEDROCK_TO_JAVA_LIQUID_MAP.get(bedrockId);
+        return BEDROCK_TO_JAVA_BLOCK_MAP.get(1 << 31 | bedrockId);
     }
 }

connector/src/main/java/org/geysermc/connector/utils/ChunkUtils.java

@@ -38,6 +38,8 @@ import org.geysermc.connector.network.translators.TranslatorsInit;
 import org.geysermc.connector.network.translators.block.BlockTranslator;
 import org.geysermc.connector.world.chunk.ChunkSection;
 
+import static org.geysermc.connector.network.translators.block.BlockTranslator.BEDROCK_WATER_ID;
+
 public class ChunkUtils {
     public static ChunkData translateToBedrock(Column column) {
         ChunkData chunkData = new ChunkData();
@@ -63,9 +65,8 @@ public class ChunkUtils {
 
                         section.getBlockStorageArray()[0].setFullBlock(ChunkSection.blockPosition(x, y, z), id);
 
-                        int waterloggedId = BlockTranslator.getBedrockWaterLoggedId(blockState);
+                        if (BlockTranslator.isWaterlogged(blockState)) {
-                        if (waterloggedId != -1) {
+                            section.getBlockStorageArray()[1].setFullBlock(ChunkSection.blockPosition(x, y, z), BEDROCK_WATER_ID);
-                            section.getBlockStorageArray()[1].setFullBlock(ChunkSection.blockPosition(x, y, z), waterloggedId);
                         }
                     }
                 }
@@ -88,9 +89,8 @@ public class ChunkUtils {
         UpdateBlockPacket waterPacket = new UpdateBlockPacket();
         waterPacket.setDataLayer(1);
         waterPacket.setBlockPosition(pos);
-        int waterloggedId = BlockTranslator.getBedrockWaterLoggedId(blockState);
+        if (BlockTranslator.isWaterlogged(blockState)) {
-        if (waterloggedId != -1) {
+            waterPacket.setRuntimeId(BEDROCK_WATER_ID);
-            waterPacket.setRuntimeId(waterloggedId);
         } else {
             waterPacket.setRuntimeId(0);
         }