CatServer/patches/net/minecraft/world/gen/MapGenCaves.java.patch

--- ../src-base/minecraft/net/minecraft/world/gen/MapGenCaves.java
+++ ../src-work/minecraft/net/minecraft/world/gen/MapGenCaves.java
@@ -140,9 +140,7 @@
                             {
                                 if (l1 >= 0 && l1 < 256)
                                 {
-                                    IBlockState iblockstate = p_180702_5_.getBlockState(j1, l1, k1);
-
-                                    if (iblockstate.getBlock() == Blocks.FLOWING_WATER || iblockstate.getBlock() == Blocks.WATER)
+                                    if (isOceanBlock(p_180702_5_, j1, l1, k1, p_180702_3_, p_180702_4_))
                                     {
                                         flag3 = true;
                                     }
@@ -180,28 +178,12 @@
                                             IBlockState iblockstate1 = p_180702_5_.getBlockState(j3, j2, i2);
                                             IBlockState iblockstate2 = (IBlockState)MoreObjects.firstNonNull(p_180702_5_.getBlockState(j3, j2 + 1, i2), BLK_AIR);
 
-                                            if (iblockstate1.getBlock() == Blocks.GRASS || iblockstate1.getBlock() == Blocks.MYCELIUM)
+                                            if (isTopBlock(p_180702_5_, j3, j2, i2, p_180702_3_, p_180702_4_))
                                             {
                                                 flag1 = true;
                                             }
 
-                                            if (this.canReplaceBlock(iblockstate1, iblockstate2))
-                                            {
-                                                if (j2 - 1 < 10)
-                                                {
-                                                    p_180702_5_.setBlockState(j3, j2, i2, BLK_LAVA);
-                                                }
-                                                else
-                                                {
-                                                    p_180702_5_.setBlockState(j3, j2, i2, BLK_AIR);
-
-                                                    if (flag1 && p_180702_5_.getBlockState(j3, j2 - 1, i2).getBlock() == Blocks.DIRT)
-                                                    {
-                                                        blockpos$mutableblockpos.setPos(j3 + p_180702_3_ * 16, 0, i2 + p_180702_4_ * 16);
-                                                        p_180702_5_.setBlockState(j3, j2 - 1, i2, this.world.getBiome(blockpos$mutableblockpos).topBlock.getBlock().getDefaultState());
-                                                    }
-                                                }
-                                            }
+                                            digBlock(p_180702_5_, j3, j2, i2, p_180702_3_, p_180702_4_, flag1, iblockstate1, iblockstate2);
                                         }
                                     }
                                 }
@@ -299,4 +281,66 @@
             }
         }
     }
+
+    protected boolean isOceanBlock(ChunkPrimer data, int x, int y, int z, int chunkX, int chunkZ)
+    {
+        net.minecraft.block.Block block = data.getBlockState(x, y, z).getBlock();
+        return block== Blocks.FLOWING_WATER || block == Blocks.WATER;
+    }
+
+    //Exception biomes to make sure we generate like vanilla
+    private boolean isExceptionBiome(net.minecraft.world.biome.Biome biome)
+    {
+        if (biome == net.minecraft.init.Biomes.BEACH) return true;
+        if (biome == net.minecraft.init.Biomes.DESERT) return true;
+        return false;
+    }
+
+    //Determine if the block at the specified location is the top block for the biome, we take into account
+    //Vanilla bugs to make sure that we generate the map the same way vanilla does.
+    private boolean isTopBlock(ChunkPrimer data, int x, int y, int z, int chunkX, int chunkZ)
+    {
+        net.minecraft.world.biome.Biome biome = world.getBiome(new BlockPos(x + chunkX * 16, 0, z + chunkZ * 16));
+        IBlockState state = data.getBlockState(x, y, z);
+        return (isExceptionBiome(biome) ? state.getBlock() == Blocks.GRASS : state.getBlock() == biome.topBlock);
+    }
+
+    /**
+     * Digs out the current block, default implementation removes stone, filler, and top block
+     * Sets the block to lava if y is less then 10, and air other wise.
+     * If setting to air, it also checks to see if we've broken the surface and if so
+     * tries to make the floor the biome's top block
+     *
+     * @param data Block data array
+     * @param index Pre-calculated index into block data
+     * @param x local X position
+     * @param y local Y position
+     * @param z local Z position
+     * @param chunkX Chunk X position
+     * @param chunkZ Chunk Y position
+     * @param foundTop True if we've encountered the biome's top block. Ideally if we've broken the surface.
+     */
+    protected void digBlock(ChunkPrimer data, int x, int y, int z, int chunkX, int chunkZ, boolean foundTop, IBlockState state, IBlockState up)
+    {
+        net.minecraft.world.biome.Biome biome = world.getBiome(new BlockPos(x + chunkX * 16, 0, z + chunkZ * 16));
+        IBlockState top = biome.topBlock;
+        IBlockState filler = biome.fillerBlock;
+
+        if (this.canReplaceBlock(state, up) || state.getBlock() == top.getBlock() || state.getBlock() == filler.getBlock())
+        {
+            if (y - 1 < 10)
+            {
+                data.setBlockState(x, y, z, BLK_LAVA);
+            }
+            else
+            {
+                data.setBlockState(x, y, z, BLK_AIR);
+
+                if (foundTop && data.getBlockState(x, y - 1, z).getBlock() == filler.getBlock())
+                {
+                    data.setBlockState(x, y - 1, z, top.getBlock().getDefaultState());
+                }
+            }
+        }
+    }
 }