From 11f50f8b533a931459a09cc3f070f7231fa1f801 Mon Sep 17 00:00:00 2001
From: Aurora <21148213+aurorasmiles@users.noreply.github.com>
Date: Mon, 23 Nov 2020 19:04:34 +0100
Subject: [PATCH] Add 1.16 tree types (#757)

---
 .../function/generator/ForestGenerator.java   | 34 ++++++++-----------
 .../sk89q/worldedit/util/TreeGenerator.java   | 11 +++++-
 2 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/worldedit-core/src/main/java/com/sk89q/worldedit/function/generator/ForestGenerator.java b/worldedit-core/src/main/java/com/sk89q/worldedit/function/generator/ForestGenerator.java
index 50756d9b4..b6427453f 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/function/generator/ForestGenerator.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/function/generator/ForestGenerator.java
@@ -54,28 +54,22 @@ public class ForestGenerator implements RegionFunction {
         BlockState block = editSession.getBlock(position);
         BlockType t = block.getBlockType();
 
-        switch (t.getInternalId()) {
-            case BlockID.GRASS_BLOCK:
-            case BlockID.DIRT:
-            case BlockID.PODZOL:
-            case BlockID.COARSE_DIRT:
+        if (t.getMaterial().isSolid()) {
             return treeType.generate(editSession, position.add(0, 1, 0));
-            default:
-            if (t.getMaterial().isReplacedDuringPlacement()) {
-                // since the implementation's tree generators generally don't generate in non-air spots,
-                // we trick editsession history here in the first call
-                editSession.setBlock(position, BlockTypes.AIR.getDefaultState());
-                // and then trick the generator here by directly setting into the world
-                editSession.getWorld().setBlock(position, BlockTypes.AIR.getDefaultState());
-                // so that now the generator can generate the tree
-                boolean success = treeType.generate(editSession, position);
-                if (!success) {
-                    editSession.setBlock(position, block); // restore on failure
-                }
-                return success;
-            } else { // Trees won't grow on this!
-                return false;
+        } else if (t.getMaterial().isReplacedDuringPlacement()) {
+            // since the implementation's tree generators generally don't generate in non-air spots,
+            // we trick editsession history here in the first call
+            editSession.setBlock(position, BlockTypes.AIR.getDefaultState());
+            // and then trick the generator here by directly setting into the world
+            editSession.getWorld().setBlock(position, BlockTypes.AIR.getDefaultState());
+            // so that now the generator can generate the tree
+            boolean success = treeType.generate(editSession, position);
+            if (!success) {
+                editSession.setBlock(position, block); // restore on failure
             }
+            return success;
+        } else { // Trees won't grow on this!
+            return false;
         }
     }
 }
diff --git a/worldedit-core/src/main/java/com/sk89q/worldedit/util/TreeGenerator.java b/worldedit-core/src/main/java/com/sk89q/worldedit/util/TreeGenerator.java
index c67731f56..b1a5799ec 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/util/TreeGenerator.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/util/TreeGenerator.java
@@ -82,10 +82,12 @@ public class TreeGenerator {
         JUNGLE_BUSH("Jungle bush", "junglebush", "jungleshrub"),
         RED_MUSHROOM("Red mushroom", "redmushroom", "redgiantmushroom"),
         BROWN_MUSHROOM("Brown mushroom", "brownmushroom", "browngiantmushroom"),
+        CRIMSON_FUNGUS("Crimson fungus", "crimsonfungus", "rednethermushroom"),
+        WARPED_FUNGUS("Warped fungus", "warpedfungus", "greennethermushroom"),
         RANDOM_MUSHROOM("Random mushroom", "randmushroom", "randommushroom") {
             @Override
             public boolean generate(EditSession editSession, BlockVector3 pos) throws MaxChangedBlocksException {
-                TreeType[] choices = { RED_MUSHROOM, BROWN_MUSHROOM };
+                TreeType[] choices = { RED_MUSHROOM, BROWN_MUSHROOM, CRIMSON_FUNGUS, WARPED_FUNGUS };
                 return choices[TreeGenerator.RANDOM.nextInt(choices.length)].generate(editSession, pos);
             }
         },
@@ -99,6 +101,13 @@ public class TreeGenerator {
                 return true;
             }
         },
+        CHORUS_PLANT("Chorus plant", "chorusplant") {
+            @Override
+            public boolean generate(EditSession editSession, BlockVector3 pos) throws MaxChangedBlocksException {
+                // chorus plants have to generate starting in the end stone itself, not the air above the ground
+                return editSession.getWorld().generateTree(this, editSession, pos.subtract(0, 1, 0));
+            }
+        },
         RANDOM("Random tree", "rand", "random") {
             @Override
             public boolean generate(EditSession editSession, BlockVector3 pos) throws MaxChangedBlocksException {