From a433c341fc6be70c11ca50b80ab3bf808ae9dd8a Mon Sep 17 00:00:00 2001
From: Eclipse <eclipse@eclipseisoffline.xyz>
Date: Thu, 20 Jun 2024 08:12:58 +0000
Subject: [PATCH] Make sure block state is never null, remove now unnecessary
 null checks

---
 .../geyser/entity/type/player/SkullPlayerEntity.java   |  4 ++--
 .../java/org/geysermc/geyser/level/WorldManager.java   |  7 +------
 .../geysermc/geyser/level/block/type/BlockState.java   | 10 +++++++++-
 .../org/geysermc/geyser/session/cache/SkullCache.java  |  4 ++--
 .../level/block/entity/SkullBlockEntityTranslator.java |  6 +-----
 .../bedrock/entity/player/BedrockActionTranslator.java |  4 ++--
 .../java/level/JavaBlockDestructionTranslator.java     |  4 +---
 .../java/level/JavaBlockEntityDataTranslator.java      |  6 +++---
 .../java/level/JavaLevelChunkWithLightTranslator.java  |  2 +-
 9 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/entity/type/player/SkullPlayerEntity.java b/core/src/main/java/org/geysermc/geyser/entity/type/player/SkullPlayerEntity.java
index 0cf8399f7..1c5060de0 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/type/player/SkullPlayerEntity.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/type/player/SkullPlayerEntity.java
@@ -141,7 +141,7 @@ public class SkullPlayerEntity extends PlayerEntity {
         float rotation;
 
         BlockState blockState = skull.getBlockState();
-        if (blockState != null && blockState.block() instanceof WallSkullBlock) {
+        if (blockState.block() instanceof WallSkullBlock) {
             y += 0.25f;
             Direction direction = blockState.getValue(Properties.HORIZONTAL_FACING);
             rotation = WallSkullBlock.getDegrees(direction);
@@ -152,7 +152,7 @@ public class SkullPlayerEntity extends PlayerEntity {
                 case EAST -> x -= 0.24f;
             }
         } else {
-            rotation = (180f + (blockState != null ? blockState.getValue(Properties.ROTATION_16) * 22.5f : 0.0f)) % 360;
+            rotation = (180f + blockState.getValue(Properties.ROTATION_16, 0) * 22.5f) % 360;
         }
 
         moveAbsolute(Vector3f.from(x, y, z), rotation, 0, rotation, true, true);
diff --git a/core/src/main/java/org/geysermc/geyser/level/WorldManager.java b/core/src/main/java/org/geysermc/geyser/level/WorldManager.java
index a18db03fe..3670b6b73 100644
--- a/core/src/main/java/org/geysermc/geyser/level/WorldManager.java
+++ b/core/src/main/java/org/geysermc/geyser/level/WorldManager.java
@@ -33,7 +33,6 @@ import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.cloudburstmc.math.vector.Vector3i;
 import org.geysermc.erosion.util.BlockPositionIterator;
-import org.geysermc.geyser.level.block.type.Block;
 import org.geysermc.geyser.level.block.type.BlockState;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.player.GameMode;
@@ -65,11 +64,7 @@ public abstract class WorldManager {
 
     @NonNull
     public BlockState blockAt(GeyserSession session, int x, int y, int z) {
-        BlockState block = BlockState.of(this.getBlockAt(session, x, y, z));
-        if (block == null) {
-            return BlockState.of(Block.JAVA_AIR_ID);
-        }
-        return block;
+        return BlockState.of(this.getBlockAt(session, x, y, z));
     }
 
     /**
diff --git a/core/src/main/java/org/geysermc/geyser/level/block/type/BlockState.java b/core/src/main/java/org/geysermc/geyser/level/block/type/BlockState.java
index 2a4b1774d..636101b15 100644
--- a/core/src/main/java/org/geysermc/geyser/level/block/type/BlockState.java
+++ b/core/src/main/java/org/geysermc/geyser/level/block/type/BlockState.java
@@ -25,6 +25,7 @@
 
 package org.geysermc.geyser.level.block.type;
 
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.geysermc.geyser.level.block.property.Property;
 import org.geysermc.geyser.registry.BlockRegistries;
@@ -192,7 +193,14 @@ public final class BlockState {
         return builder.toString();
     }
 
+    /**
+     * Null-safe method that looks up a Java block state ID in the BLOCK_STATES registry, and defaults to air if not found.
+     *
+     * @param javaId the Java block state ID to look up.
+     * @return the corresponding block state, or air if the given ID wasn't registered and returned null.
+     */
+    @NonNull
     public static BlockState of(int javaId) {
-        return BlockRegistries.BLOCK_STATES.get(javaId);
+        return BlockRegistries.BLOCK_STATES.getOrDefault(javaId, BlockRegistries.BLOCK_STATES.get(Block.JAVA_AIR_ID));
     }
 }
diff --git a/core/src/main/java/org/geysermc/geyser/session/cache/SkullCache.java b/core/src/main/java/org/geysermc/geyser/session/cache/SkullCache.java
index 43ee38735..a40a1156d 100644
--- a/core/src/main/java/org/geysermc/geyser/session/cache/SkullCache.java
+++ b/core/src/main/java/org/geysermc/geyser/session/cache/SkullCache.java
@@ -254,10 +254,10 @@ public class SkullCache {
         CustomSkull customSkull = BlockRegistries.CUSTOM_SKULLS.get(skinHash);
         if (customSkull != null) {
             CustomBlockState customBlockState;
-            if (blockState != null && blockState.block() instanceof WallSkullBlock) {
+            if (blockState.block() instanceof WallSkullBlock) {
                 customBlockState = customSkull.getWallBlockState(WallSkullBlock.getDegrees(blockState));
             } else {
-                customBlockState = customSkull.getFloorBlockState(blockState != null ? blockState.getValue(Properties.ROTATION_16) : 0);
+                customBlockState = customSkull.getFloorBlockState(blockState.getValue(Properties.ROTATION_16));
             }
 
             return session.getBlockMappings().getCustomBlockStateDefinitions().get(customBlockState);
diff --git a/core/src/main/java/org/geysermc/geyser/translator/level/block/entity/SkullBlockEntityTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/level/block/entity/SkullBlockEntityTranslator.java
index 86b88c5c0..727d8fd4f 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/level/block/entity/SkullBlockEntityTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/level/block/entity/SkullBlockEntityTranslator.java
@@ -53,17 +53,13 @@ public class SkullBlockEntityTranslator extends BlockEntityTranslator implements
 
     @Override
     public void translateTag(GeyserSession session, NbtMapBuilder bedrockNbt, NbtMap javaNbt, BlockState blockState) {
-        if (blockState == null) {
-            bedrockNbt.putByte("SkullType", (byte) 0);
-            return;
-        }
         Integer rotation = blockState.getValue(Properties.ROTATION_16);
         if (rotation != null) {
             // Could be a wall skull block otherwise, which has rotation in its Bedrock state
             bedrockNbt.putFloat("Rotation", rotation * 22.5f);
         }
         bedrockNbt.putByte("SkullType", (byte) (blockState.block() instanceof SkullBlock skull ? skull.skullType().bedrockId() : 0));
-        if (blockState.getValue(Properties.POWERED)) {
+        if (blockState.getValue(Properties.POWERED, false)) {
             bedrockNbt.putBoolean("MouthMoving", true);
         }
     }
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/entity/player/BedrockActionTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/entity/player/BedrockActionTranslator.java
index 959797d41..6834d3190 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/entity/player/BedrockActionTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/bedrock/entity/player/BedrockActionTranslator.java
@@ -160,7 +160,7 @@ public class BedrockActionTranslator extends PacketTranslator<PlayerActionPacket
                 LevelEventPacket startBreak = new LevelEventPacket();
                 startBreak.setType(LevelEvent.BLOCK_START_BREAK);
                 startBreak.setPosition(vector.toFloat());
-                double breakTime = BlockUtils.getSessionBreakTime(session, BlockRegistries.BLOCK_STATES.getOrDefault(blockState, BlockState.of(Block.JAVA_AIR_ID)).block()) * 20;
+                double breakTime = BlockUtils.getSessionBreakTime(session, BlockState.of(blockState).block()) * 20;
 
                 // If the block is custom or the breaking item is custom, we must keep track of break time ourselves
                 GeyserItemStack item = session.getPlayerInventory().getItemInHand();
@@ -211,7 +211,7 @@ public class BedrockActionTranslator extends PacketTranslator<PlayerActionPacket
                 LevelEventPacket updateBreak = new LevelEventPacket();
                 updateBreak.setType(LevelEvent.BLOCK_UPDATE_BREAK);
                 updateBreak.setPosition(vectorFloat);
-                double breakTime = BlockUtils.getSessionBreakTime(session, BlockRegistries.BLOCK_STATES.getOrDefault(breakingBlock, BlockState.of(Block.JAVA_AIR_ID)).block()) * 20;
+                double breakTime = BlockUtils.getSessionBreakTime(session, BlockState.of(breakingBlock).block()) * 20;
 
 
                 // If the block is custom, we must keep track of when it should break ourselves
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockDestructionTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockDestructionTranslator.java
index 700f7e701..674a1c583 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockDestructionTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockDestructionTranslator.java
@@ -27,9 +27,7 @@ package org.geysermc.geyser.translator.protocol.java.level;
 
 import org.cloudburstmc.protocol.bedrock.data.LevelEvent;
 import org.cloudburstmc.protocol.bedrock.packet.LevelEventPacket;
-import org.geysermc.geyser.level.block.type.Block;
 import org.geysermc.geyser.level.block.type.BlockState;
-import org.geysermc.geyser.registry.BlockRegistries;
 import org.geysermc.geyser.registry.type.ItemMapping;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.translator.protocol.PacketTranslator;
@@ -43,7 +41,7 @@ public class JavaBlockDestructionTranslator extends PacketTranslator<Clientbound
     @Override
     public void translate(GeyserSession session, ClientboundBlockDestructionPacket packet) {
         int state = session.getGeyser().getWorldManager().getBlockAt(session, packet.getPosition().getX(), packet.getPosition().getY(), packet.getPosition().getZ());
-        int breakTime = (int) (65535 / Math.ceil(BlockUtils.getBreakTime(session, BlockRegistries.BLOCK_STATES.getOrDefault(state, BlockState.of(Block.JAVA_AIR_ID)).block(), ItemMapping.AIR, null, false) * 20));
+        int breakTime = (int) (65535 / Math.ceil(BlockUtils.getBreakTime(session, BlockState.of(state).block(), ItemMapping.AIR, null, false) * 20));
         LevelEventPacket levelEventPacket = new LevelEventPacket();
         levelEventPacket.setPosition(packet.getPosition().toFloat());
         levelEventPacket.setType(LevelEvent.BLOCK_START_BREAK);
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java
index 4210a8fe0..2b6f1d70e 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaBlockEntityDataTranslator.java
@@ -33,8 +33,8 @@ import org.cloudburstmc.protocol.bedrock.data.structure.StructureMirror;
 import org.cloudburstmc.protocol.bedrock.data.structure.StructureRotation;
 import org.cloudburstmc.protocol.bedrock.packet.ContainerOpenPacket;
 import org.cloudburstmc.protocol.bedrock.packet.UpdateBlockPacket;
+import org.geysermc.geyser.level.block.type.Block;
 import org.geysermc.geyser.level.block.type.BlockState;
-import org.geysermc.geyser.registry.BlockRegistries;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.translator.level.block.entity.BlockEntityTranslator;
 import org.geysermc.geyser.translator.level.block.entity.RequiresBlockState;
@@ -61,9 +61,9 @@ public class JavaBlockEntityDataTranslator extends PacketTranslator<ClientboundB
         // between Java block states and Bedrock block entity data
         BlockState blockState;
         if (translator instanceof RequiresBlockState) {
-            blockState = BlockRegistries.BLOCK_STATES.get(session.getGeyser().getWorldManager().getBlockAt(session, packet.getPosition()));
+            blockState = session.getGeyser().getWorldManager().blockAt(session, packet.getPosition());
         } else {
-            blockState = BlockRegistries.BLOCK_STATES.get(0); //TODO
+            blockState = BlockState.of(Block.JAVA_AIR_ID); //TODO
         }
 
         Vector3i position = packet.getPosition();
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java
index e37edef66..b60de0fbe 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/level/JavaLevelChunkWithLightTranslator.java
@@ -398,7 +398,7 @@ public class JavaLevelChunkWithLightTranslator extends PacketTranslator<Clientbo
 
                 // Get the Java block state ID from block entity position
                 DataPalette section = javaChunks[(y >> 4) - yOffset];
-                BlockState blockState = BlockRegistries.BLOCK_STATES.get(section.get(x, y & 0xF, z));
+                BlockState blockState = BlockState.of(section.get(x, y & 0xF, z));
 
                 // Note that, since 1.20.5, tags can be null, but Bedrock still needs a default tag to render the item
                 // Also, some properties - like banner base colors - are part of the tag and is processed here.