From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Jake Potrebic <jake.m.potrebic@gmail.com>
Date: Wed, 6 Oct 2021 20:50:48 -0700
Subject: [PATCH] Fix upstreams block state factories

Sometimes, blocks are changed and then logic is called before the associated
tile entity is removed. When this happens, the factories were relying on the
block at the position, not the tile entity. This change prioritizes using the
tile entity type to determine the block state factory and falls back on
the material type of the block at that location.

diff --git a/src/main/java/net/minecraft/world/level/block/entity/BlockEntity.java b/src/main/java/net/minecraft/world/level/block/entity/BlockEntity.java
index b0174aedb7358af1e80278e2f5f13e00c35ab3c6..d62181bd8bccfcfdd7da8f635bdf7ebc36294705 100644
--- a/src/main/java/net/minecraft/world/level/block/entity/BlockEntity.java
+++ b/src/main/java/net/minecraft/world/level/block/entity/BlockEntity.java
@@ -250,7 +250,7 @@ public abstract class BlockEntity {
         // Paper end
         if (this.level == null) return null;
         org.bukkit.block.Block block = this.level.getWorld().getBlockAt(this.worldPosition.getX(), this.worldPosition.getY(), this.worldPosition.getZ());
-        if (block.getType() == org.bukkit.Material.AIR) return null;
+        // if (block.getType() == org.bukkit.Material.AIR) return null; // Paper - actually get the tile entity if it still exists
         org.bukkit.block.BlockState state = block.getState(useSnapshot); // Paper
         if (state instanceof InventoryHolder) return (InventoryHolder) state;
         return null;
diff --git a/src/main/java/org/bukkit/craftbukkit/block/CraftBlockEntityState.java b/src/main/java/org/bukkit/craftbukkit/block/CraftBlockEntityState.java
index 87c25170fbe8b0591d452612496ee1a627138de7..a2894f02ceb7c58f6eafe055e1ff47b197b100f2 100644
--- a/src/main/java/org/bukkit/craftbukkit/block/CraftBlockEntityState.java
+++ b/src/main/java/org/bukkit/craftbukkit/block/CraftBlockEntityState.java
@@ -6,7 +6,7 @@ import org.bukkit.World;
 import org.bukkit.block.TileState;
 import org.bukkit.persistence.PersistentDataContainer;
 
-public class CraftBlockEntityState<T extends BlockEntity> extends CraftBlockState implements TileState {
+public abstract class CraftBlockEntityState<T extends BlockEntity> extends CraftBlockState implements TileState { // Paper - revert upstream's revert of the block state changes
 
     private final T tileEntity;
     private final T snapshot;
diff --git a/src/main/java/org/bukkit/craftbukkit/block/CraftBlockStates.java b/src/main/java/org/bukkit/craftbukkit/block/CraftBlockStates.java
index e609cfeaf5aa6807f57360dde9b0dccf40a23eb1..594b071cbb25d6e8f184c8558c3d1f490fe4712f 100644
--- a/src/main/java/org/bukkit/craftbukkit/block/CraftBlockStates.java
+++ b/src/main/java/org/bukkit/craftbukkit/block/CraftBlockStates.java
@@ -19,6 +19,7 @@ import net.minecraft.world.level.block.entity.BeehiveBlockEntity;
 import net.minecraft.world.level.block.entity.BellBlockEntity;
 import net.minecraft.world.level.block.entity.BlastFurnaceBlockEntity;
 import net.minecraft.world.level.block.entity.BlockEntity;
+import net.minecraft.world.level.block.entity.BlockEntityType; // Paper
 import net.minecraft.world.level.block.entity.BrewingStandBlockEntity;
 import net.minecraft.world.level.block.entity.CampfireBlockEntity;
 import net.minecraft.world.level.block.entity.ChestBlockEntity;
@@ -108,187 +109,57 @@ public final class CraftBlockStates {
     private static final BlockStateFactory<?> DEFAULT_FACTORY = new BlockStateFactory<CraftBlockState>(CraftBlockState.class) {
         @Override
         public CraftBlockState createBlockState(World world, BlockPos blockPosition, net.minecraft.world.level.block.state.BlockState blockData, BlockEntity tileEntity) {
-            // SPIGOT-6754, SPIGOT-6817: Restore previous behaviour for tile entities with removed blocks (loot generation post-destroy)
-            if (tileEntity != null) {
-                // block with unhandled TileEntity:
-                return new CraftBlockEntityState<>(world, tileEntity);
-            }
+            // Paper - revert upstream's revert of the block state changes. Block entities that have already had the block type set to AIR are still valid, upstream decided to ignore them
             Preconditions.checkState(tileEntity == null, "Unexpected BlockState for %s", CraftMagicNumbers.getMaterial(blockData.getBlock()));
             return new CraftBlockState(world, blockPosition, blockData);
         }
     };
+    // Paper start
+    private static final Map<BlockEntityType<?>, BlockStateFactory<?>> FACTORIES_BY_BLOCK_ENTITY_TYPE = new HashMap<>();
+    private static void register(BlockEntityType<?> type, BlockStateFactory<?> factory) {
+        FACTORIES_BY_BLOCK_ENTITY_TYPE.put(type, factory);
+    }
+    // Paper end
 
     static {
-        register(
-                Arrays.asList(
-                        Material.ACACIA_SIGN,
-                        Material.ACACIA_WALL_SIGN,
-                        Material.BIRCH_SIGN,
-                        Material.BIRCH_WALL_SIGN,
-                        Material.CRIMSON_SIGN,
-                        Material.CRIMSON_WALL_SIGN,
-                        Material.DARK_OAK_SIGN,
-                        Material.DARK_OAK_WALL_SIGN,
-                        Material.JUNGLE_SIGN,
-                        Material.JUNGLE_WALL_SIGN,
-                        Material.MANGROVE_SIGN,
-                        Material.MANGROVE_WALL_SIGN,
-                        Material.OAK_SIGN,
-                        Material.OAK_WALL_SIGN,
-                        Material.SPRUCE_SIGN,
-                        Material.SPRUCE_WALL_SIGN,
-                        Material.WARPED_SIGN,
-                        Material.WARPED_WALL_SIGN
-                ), CraftSign.class, CraftSign::new, SignBlockEntity::new
-        );
-
-        register(
-                Arrays.asList(
-                        Material.CREEPER_HEAD,
-                        Material.CREEPER_WALL_HEAD,
-                        Material.DRAGON_HEAD,
-                        Material.DRAGON_WALL_HEAD,
-                        Material.PLAYER_HEAD,
-                        Material.PLAYER_WALL_HEAD,
-                        Material.SKELETON_SKULL,
-                        Material.SKELETON_WALL_SKULL,
-                        Material.WITHER_SKELETON_SKULL,
-                        Material.WITHER_SKELETON_WALL_SKULL,
-                        Material.ZOMBIE_HEAD,
-                        Material.ZOMBIE_WALL_HEAD
-                ), CraftSkull.class, CraftSkull::new, SkullBlockEntity::new
-        );
-
-        register(
-                Arrays.asList(
-                        Material.COMMAND_BLOCK,
-                        Material.REPEATING_COMMAND_BLOCK,
-                        Material.CHAIN_COMMAND_BLOCK
-                ), CraftCommandBlock.class, CraftCommandBlock::new, CommandBlockEntity::new
-        );
-
-        register(
-                Arrays.asList(
-                        Material.BLACK_BANNER,
-                        Material.BLACK_WALL_BANNER,
-                        Material.BLUE_BANNER,
-                        Material.BLUE_WALL_BANNER,
-                        Material.BROWN_BANNER,
-                        Material.BROWN_WALL_BANNER,
-                        Material.CYAN_BANNER,
-                        Material.CYAN_WALL_BANNER,
-                        Material.GRAY_BANNER,
-                        Material.GRAY_WALL_BANNER,
-                        Material.GREEN_BANNER,
-                        Material.GREEN_WALL_BANNER,
-                        Material.LIGHT_BLUE_BANNER,
-                        Material.LIGHT_BLUE_WALL_BANNER,
-                        Material.LIGHT_GRAY_BANNER,
-                        Material.LIGHT_GRAY_WALL_BANNER,
-                        Material.LIME_BANNER,
-                        Material.LIME_WALL_BANNER,
-                        Material.MAGENTA_BANNER,
-                        Material.MAGENTA_WALL_BANNER,
-                        Material.ORANGE_BANNER,
-                        Material.ORANGE_WALL_BANNER,
-                        Material.PINK_BANNER,
-                        Material.PINK_WALL_BANNER,
-                        Material.PURPLE_BANNER,
-                        Material.PURPLE_WALL_BANNER,
-                        Material.RED_BANNER,
-                        Material.RED_WALL_BANNER,
-                        Material.WHITE_BANNER,
-                        Material.WHITE_WALL_BANNER,
-                        Material.YELLOW_BANNER,
-                        Material.YELLOW_WALL_BANNER
-                ), CraftBanner.class, CraftBanner::new, BannerBlockEntity::new
-        );
-
-        register(
-                Arrays.asList(
-                        Material.SHULKER_BOX,
-                        Material.WHITE_SHULKER_BOX,
-                        Material.ORANGE_SHULKER_BOX,
-                        Material.MAGENTA_SHULKER_BOX,
-                        Material.LIGHT_BLUE_SHULKER_BOX,
-                        Material.YELLOW_SHULKER_BOX,
-                        Material.LIME_SHULKER_BOX,
-                        Material.PINK_SHULKER_BOX,
-                        Material.GRAY_SHULKER_BOX,
-                        Material.LIGHT_GRAY_SHULKER_BOX,
-                        Material.CYAN_SHULKER_BOX,
-                        Material.PURPLE_SHULKER_BOX,
-                        Material.BLUE_SHULKER_BOX,
-                        Material.BROWN_SHULKER_BOX,
-                        Material.GREEN_SHULKER_BOX,
-                        Material.RED_SHULKER_BOX,
-                        Material.BLACK_SHULKER_BOX
-                ), CraftShulkerBox.class, CraftShulkerBox::new, ShulkerBoxBlockEntity::new
-        );
-
-        register(
-                Arrays.asList(
-                        Material.BLACK_BED,
-                        Material.BLUE_BED,
-                        Material.BROWN_BED,
-                        Material.CYAN_BED,
-                        Material.GRAY_BED,
-                        Material.GREEN_BED,
-                        Material.LIGHT_BLUE_BED,
-                        Material.LIGHT_GRAY_BED,
-                        Material.LIME_BED,
-                        Material.MAGENTA_BED,
-                        Material.ORANGE_BED,
-                        Material.PINK_BED,
-                        Material.PURPLE_BED,
-                        Material.RED_BED,
-                        Material.WHITE_BED,
-                        Material.YELLOW_BED
-                ), CraftBed.class, CraftBed::new, BedBlockEntity::new
-        );
-
-        register(
-                Arrays.asList(
-                        Material.BEEHIVE,
-                        Material.BEE_NEST
-                ), CraftBeehive.class, CraftBeehive::new, BeehiveBlockEntity::new
-        );
-
-        register(
-                Arrays.asList(
-                        Material.CAMPFIRE,
-                        Material.SOUL_CAMPFIRE
-                ), CraftCampfire.class, CraftCampfire::new, CampfireBlockEntity::new
-        );
-
-        register(Material.BARREL, CraftBarrel.class, CraftBarrel::new, BarrelBlockEntity::new);
-        register(Material.BEACON, CraftBeacon.class, CraftBeacon::new, BeaconBlockEntity::new);
-        register(Material.BELL, CraftBell.class, CraftBell::new, BellBlockEntity::new);
-        register(Material.BLAST_FURNACE, CraftBlastFurnace.class, CraftBlastFurnace::new, BlastFurnaceBlockEntity::new);
-        register(Material.BREWING_STAND, CraftBrewingStand.class, CraftBrewingStand::new, BrewingStandBlockEntity::new);
-        register(Material.CHEST, CraftChest.class, CraftChest::new, ChestBlockEntity::new);
-        register(Material.COMPARATOR, CraftComparator.class, CraftComparator::new, ComparatorBlockEntity::new);
-        register(Material.CONDUIT, CraftConduit.class, CraftConduit::new, ConduitBlockEntity::new);
-        register(Material.DAYLIGHT_DETECTOR, CraftDaylightDetector.class, CraftDaylightDetector::new, DaylightDetectorBlockEntity::new);
-        register(Material.DISPENSER, CraftDispenser.class, CraftDispenser::new, DispenserBlockEntity::new);
-        register(Material.DROPPER, CraftDropper.class, CraftDropper::new, DropperBlockEntity::new);
-        register(Material.ENCHANTING_TABLE, CraftEnchantingTable.class, CraftEnchantingTable::new, EnchantmentTableBlockEntity::new);
-        register(Material.ENDER_CHEST, CraftEnderChest.class, CraftEnderChest::new, EnderChestBlockEntity::new);
-        register(Material.END_GATEWAY, CraftEndGateway.class, CraftEndGateway::new, TheEndGatewayBlockEntity::new);
-        register(Material.END_PORTAL, CraftEndPortal.class, CraftEndPortal::new, TheEndPortalBlockEntity::new);
-        register(Material.FURNACE, CraftFurnaceFurnace.class, CraftFurnaceFurnace::new, FurnaceBlockEntity::new);
-        register(Material.HOPPER, CraftHopper.class, CraftHopper::new, HopperBlockEntity::new);
-        register(Material.JIGSAW, CraftJigsaw.class, CraftJigsaw::new, JigsawBlockEntity::new);
-        register(Material.JUKEBOX, CraftJukebox.class, CraftJukebox::new, JukeboxBlockEntity::new);
-        register(Material.LECTERN, CraftLectern.class, CraftLectern::new, LecternBlockEntity::new);
-        register(Material.MOVING_PISTON, CraftMovingPiston.class, CraftMovingPiston::new, PistonMovingBlockEntity::new);
-        register(Material.SCULK_CATALYST, CraftSculkCatalyst.class, CraftSculkCatalyst::new, SculkCatalystBlockEntity::new);
-        register(Material.SCULK_SENSOR, CraftSculkSensor.class, CraftSculkSensor::new, SculkSensorBlockEntity::new);
-        register(Material.SCULK_SHRIEKER, CraftSculkShrieker.class, CraftSculkShrieker::new, SculkShriekerBlockEntity::new);
-        register(Material.SMOKER, CraftSmoker.class, CraftSmoker::new, SmokerBlockEntity::new);
-        register(Material.SPAWNER, CraftCreatureSpawner.class, CraftCreatureSpawner::new, SpawnerBlockEntity::new);
-        register(Material.STRUCTURE_BLOCK, CraftStructureBlock.class, CraftStructureBlock::new, StructureBlockEntity::new);
-        register(Material.TRAPPED_CHEST, CraftChest.class, CraftChest::new, TrappedChestBlockEntity::new);
+        // Paper start - simplify
+        register(BlockEntityType.SIGN, CraftSign.class, CraftSign::new);
+        register(BlockEntityType.SKULL, CraftSkull.class, CraftSkull::new);
+        register(BlockEntityType.COMMAND_BLOCK, CraftCommandBlock.class, CraftCommandBlock::new);
+        register(BlockEntityType.BANNER, CraftBanner.class, CraftBanner::new);
+        register(BlockEntityType.SHULKER_BOX, CraftShulkerBox.class, CraftShulkerBox::new);
+        register(BlockEntityType.BED, CraftBed.class, CraftBed::new);
+        register(BlockEntityType.BEEHIVE, CraftBeehive.class, CraftBeehive::new);
+        register(BlockEntityType.CAMPFIRE, CraftCampfire.class, CraftCampfire::new);
+        register(BlockEntityType.BARREL, CraftBarrel.class, CraftBarrel::new);
+        register(BlockEntityType.BEACON, CraftBeacon.class, CraftBeacon::new);
+        register(BlockEntityType.BELL, CraftBell.class, CraftBell::new);
+        register(BlockEntityType.BLAST_FURNACE, CraftBlastFurnace.class, CraftBlastFurnace::new);
+        register(BlockEntityType.BREWING_STAND, CraftBrewingStand.class, CraftBrewingStand::new);
+        register(BlockEntityType.CHEST, CraftChest.class, CraftChest::new);
+        register(BlockEntityType.COMPARATOR, CraftComparator.class, CraftComparator::new);
+        register(BlockEntityType.CONDUIT, CraftConduit.class, CraftConduit::new);
+        register(BlockEntityType.DAYLIGHT_DETECTOR, CraftDaylightDetector.class, CraftDaylightDetector::new);
+        register(BlockEntityType.DISPENSER, CraftDispenser.class, CraftDispenser::new);
+        register(BlockEntityType.DROPPER, CraftDropper.class, CraftDropper::new);
+        register(BlockEntityType.ENCHANTING_TABLE, CraftEnchantingTable.class, CraftEnchantingTable::new);
+        register(BlockEntityType.ENDER_CHEST, CraftEnderChest.class, CraftEnderChest::new);
+        register(BlockEntityType.END_GATEWAY, CraftEndGateway.class, CraftEndGateway::new);
+        register(BlockEntityType.END_PORTAL, CraftEndPortal.class, CraftEndPortal::new);
+        register(BlockEntityType.FURNACE, CraftFurnaceFurnace.class, CraftFurnaceFurnace::new);
+        register(BlockEntityType.HOPPER, CraftHopper.class, CraftHopper::new);
+        register(BlockEntityType.JIGSAW, CraftJigsaw.class, CraftJigsaw::new);
+        register(BlockEntityType.JUKEBOX, CraftJukebox.class, CraftJukebox::new);
+        register(BlockEntityType.LECTERN, CraftLectern.class, CraftLectern::new);
+        register(BlockEntityType.PISTON, CraftMovingPiston.class, CraftMovingPiston::new);
+        register(BlockEntityType.SCULK_CATALYST, CraftSculkCatalyst.class, CraftSculkCatalyst::new);
+        register(BlockEntityType.SCULK_SENSOR, CraftSculkSensor.class, CraftSculkSensor::new);
+        register(BlockEntityType.SCULK_SHRIEKER, CraftSculkShrieker.class, CraftSculkShrieker::new);
+        register(BlockEntityType.SMOKER, CraftSmoker.class, CraftSmoker::new);
+        register(BlockEntityType.MOB_SPAWNER, CraftCreatureSpawner.class, CraftCreatureSpawner::new);
+        register(BlockEntityType.STRUCTURE_BLOCK, CraftStructureBlock.class, CraftStructureBlock::new);
+        register(BlockEntityType.TRAPPED_CHEST, CraftChest.class, CraftChest::new);
+        // Paper end
     }
 
     private static void register(Material blockType, BlockStateFactory<?> factory) {
@@ -296,30 +167,33 @@ public final class CraftBlockStates {
     }
 
     private static <T extends BlockEntity, B extends CraftBlockEntityState<T>> void register(
-            Material blockType,
-            Class<B> blockStateType,
-            BiFunction<World, T, B> blockStateConstructor,
-            BiFunction<BlockPos, net.minecraft.world.level.block.state.BlockState, T> tileEntityConstructor
-    ) {
-        CraftBlockStates.register(Collections.singletonList(blockType), blockStateType, blockStateConstructor, tileEntityConstructor);
-    }
-
-    private static <T extends BlockEntity, B extends CraftBlockEntityState<T>> void register(
-            List<Material> blockTypes,
+            net.minecraft.world.level.block.entity.BlockEntityType<? extends T> blockEntityType, // Paper
             Class<B> blockStateType,
-            BiFunction<World, T, B> blockStateConstructor,
-            BiFunction<BlockPos, net.minecraft.world.level.block.state.BlockState, T> tileEntityConstructor
+            BiFunction<World, T, B> blockStateConstructor // Paper
     ) {
-        BlockStateFactory<B> factory = new BlockEntityStateFactory<>(blockStateType, blockStateConstructor, tileEntityConstructor);
-        for (Material blockType : blockTypes) {
-            CraftBlockStates.register(blockType, factory);
+        // Paper start
+        BlockStateFactory<B> factory = new BlockEntityStateFactory<>(blockStateType, blockStateConstructor, blockEntityType::create);
+        for (net.minecraft.world.level.block.Block block : blockEntityType.validBlocks) {
+            CraftBlockStates.register(CraftMagicNumbers.getMaterial(block), factory);
         }
+        CraftBlockStates.register(blockEntityType, factory);
+        // Paper end
     }
 
     private static BlockStateFactory<?> getFactory(Material material) {
         return CraftBlockStates.FACTORIES.getOrDefault(material, DEFAULT_FACTORY);
     }
 
+    // Paper start
+    private static BlockStateFactory<?> getFactory(Material material, BlockEntityType<?> type) {
+        if (type != null) {
+            return CraftBlockStates.FACTORIES_BY_BLOCK_ENTITY_TYPE.getOrDefault(type, getFactory(material));
+        } else {
+            return getFactory(material);
+        }
+    }
+    // Paper end
+
     public static Class<? extends CraftBlockState> getBlockStateType(Material material) {
         Preconditions.checkNotNull(material, "material is null");
         return CraftBlockStates.getFactory(material).blockStateType;
@@ -335,6 +209,13 @@ public final class CraftBlockStates {
         return null;
     }
 
+    // Paper start
+    public static Class<? extends CraftBlockState> getBlockStateType(BlockEntityType<?> blockEntityType) {
+        Preconditions.checkNotNull(blockEntityType, "blockEntityType is null");
+        return CraftBlockStates.getFactory(null, blockEntityType).blockStateType;
+    }
+    // Paper end
+
     public static BlockState getBlockState(Block block) {
         // Paper start
         return CraftBlockStates.getBlockState(block, true);
@@ -392,7 +273,7 @@ public final class CraftBlockStates {
         if (world != null && tileEntity == null && CraftBlockStates.isTileEntityOptional(material)) {
             factory = CraftBlockStates.DEFAULT_FACTORY;
         } else {
-            factory = CraftBlockStates.getFactory(material);
+            factory = CraftBlockStates.getFactory(material, tileEntity != null ? tileEntity.getType() : null); // Paper
         }
         return factory.createBlockState(world, blockPosition, blockData, tileEntity);
     }
diff --git a/src/test/java/org/bukkit/craftbukkit/block/BlockStateTest.java b/src/test/java/org/bukkit/craftbukkit/block/BlockStateTest.java
index 15f760c4c71a0493c7c18c3908d229f6b8778a43..95e4a4080a4a168b53f69d9cd7d67c27263c1b7f 100644
--- a/src/test/java/org/bukkit/craftbukkit/block/BlockStateTest.java
+++ b/src/test/java/org/bukkit/craftbukkit/block/BlockStateTest.java
@@ -45,4 +45,11 @@ public class BlockStateTest extends AbstractTestingBase {
             }
         }
     }
+
+    @Test
+    public void testBlockEntityTypes() {
+        for (var blockEntityType : Registry.BLOCK_ENTITY_TYPE) {
+            org.junit.Assert.assertNotNull(CraftBlockStates.getBlockStateType(blockEntityType));
+        }
+    }
 }
