From d94b5ae45bd7f4b2af28f111477cb835c4aceee0 Mon Sep 17 00:00:00 2001
From: Octavia Togami <octavia.togami@gmail.com>
Date: Tue, 7 Jun 2022 23:14:44 -0700
Subject: [PATCH] Sort of fix #2008

This fix only applies to Forge/Fabric, and it still produces floating
trees... but it's better than nothing
---
 .../sk89q/worldedit/fabric/FabricWorld.java   | 19 +++++++++++++++----
 .../com/sk89q/worldedit/forge/ForgeWorld.java | 18 +++++++++++++++---
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/worldedit-fabric/src/main/java/com/sk89q/worldedit/fabric/FabricWorld.java b/worldedit-fabric/src/main/java/com/sk89q/worldedit/fabric/FabricWorld.java
index 5e9ce89e2f..25734039ee 100644
--- a/worldedit-fabric/src/main/java/com/sk89q/worldedit/fabric/FabricWorld.java
+++ b/worldedit-fabric/src/main/java/com/sk89q/worldedit/fabric/FabricWorld.java
@@ -95,6 +95,8 @@
 import net.minecraft.world.level.dimension.LevelStem;
 import net.minecraft.world.level.levelgen.WorldGenSettings;
 import net.minecraft.world.level.levelgen.feature.ConfiguredFeature;
+import net.minecraft.world.level.levelgen.feature.Feature;
+import net.minecraft.world.level.levelgen.feature.TreeFeature;
 import net.minecraft.world.level.storage.LevelData;
 import net.minecraft.world.level.storage.LevelStorageSource;
 import net.minecraft.world.level.storage.ServerLevelData;
@@ -459,10 +461,19 @@ public boolean generateTree(TreeType type, EditSession editSession, BlockVector3
         if (type == TreeType.CHORUS_PLANT) {
             position = position.add(0, 1, 0);
         }
-        return generator != null && generator.place(
-            world, chunkManager.getGenerator(), random,
-            FabricAdapter.toBlockPos(position)
-        );
+        if (generator == null) {
+            return false;
+        }
+        BlockPos blockPos = FabricAdapter.toBlockPos(position);
+        // Manual override of the chunk generator to prevent the tree from generating in non-dirt
+        if (generator.feature() instanceof TreeFeature && !isTreePosition(world, blockPos)) {
+            return false;
+        }
+        return generator.place(world, chunkManager.getGenerator(), random, blockPos);
+    }
+
+    private static boolean isTreePosition(ServerLevel world, BlockPos blockPos) {
+        return world.isStateAtPosition(blockPos.below(), Feature::isDirt);
     }
 
     @Override
diff --git a/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java b/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java
index 6eaaefbc01..deaa5a0ddd 100644
--- a/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java
+++ b/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java
@@ -91,6 +91,8 @@
 import net.minecraft.world.level.dimension.LevelStem;
 import net.minecraft.world.level.levelgen.WorldGenSettings;
 import net.minecraft.world.level.levelgen.feature.ConfiguredFeature;
+import net.minecraft.world.level.levelgen.feature.Feature;
+import net.minecraft.world.level.levelgen.feature.TreeFeature;
 import net.minecraft.world.level.storage.LevelData;
 import net.minecraft.world.level.storage.LevelStorageSource;
 import net.minecraft.world.level.storage.PrimaryLevelData;
@@ -447,9 +449,19 @@ public boolean generateTree(TreeType type, EditSession editSession, BlockVector3
         if (type == TreeType.CHORUS_PLANT) {
             position = position.add(0, 1, 0);
         }
-        return generator != null && generator.place(
-            world, chunkManager.getGenerator(), random, ForgeAdapter.toBlockPos(position)
-        );
+        if (generator == null) {
+            return false;
+        }
+        BlockPos blockPos = ForgeAdapter.toBlockPos(position);
+        // Manual override of the chunk generator to prevent the tree from generating in non-dirt
+        if (generator.feature() instanceof TreeFeature && !isTreePosition(world, blockPos)) {
+            return false;
+        }
+        return generator.place(world, chunkManager.getGenerator(), random, blockPos);
+    }
+
+    private static boolean isTreePosition(ServerLevel world, BlockPos blockPos) {
+        return world.isStateAtPosition(blockPos.below(), Feature::isDirt);
     }
 
     @Override