From 28128414815959079496eb9833383d335cebe37d Mon Sep 17 00:00:00 2001
From: dordsor21 <dordsor21@gmail.com>
Date: Fri, 3 Jul 2020 14:01:16 +0100
Subject: [PATCH] Add -p flag to //drain to allow kelp and watergrass to be
 removed. properly fixed #464

---
 .../java/com/sk89q/worldedit/EditSession.java | 22 ++++++++++++++++++-
 .../worldedit/command/UtilityCommands.java    |  6 +++--
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java b/worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java
index ce1b78f05..f309f27a3 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java
@@ -1623,10 +1623,30 @@ public class EditSession extends PassthroughExtent implements AutoCloseable {
      * @throws MaxChangedBlocksException thrown if too many blocks are changed
      */
     public int drainArea(BlockVector3 origin, double radius, boolean waterlogged) throws MaxChangedBlocksException {
+        return drainArea(origin, radius, waterlogged, false);
+    }
+
+    /**
+     * Drain nearby pools of water or lava, optionally removed waterlogged states from blocks.
+     *
+     * @param origin the origin to drain from, which will search a 3x3 area
+     * @param radius the radius of the removal, where a value should be 0 or greater
+     * @param waterlogged true to make waterlogged blocks non-waterlogged as well
+     * @param plants true to remove underwater plants
+     * @return number of blocks affected
+     * @throws MaxChangedBlocksException thrown if too many blocks are changed
+     */
+    public int drainArea(BlockVector3 origin, double radius, boolean waterlogged, boolean plants) throws MaxChangedBlocksException {
         checkNotNull(origin);
         checkArgument(radius >= 0, "radius >= 0 required");
 
-        Mask liquidMask = new BlockTypeMask(this, BlockTypes.LAVA, BlockTypes.WATER);
+        Mask liquidMask;
+        if (plants) {
+             liquidMask = new BlockTypeMask(this, BlockTypes.LAVA, BlockTypes.WATER,
+                 BlockTypes.KELP_PLANT, BlockTypes.KELP, BlockTypes.SEAGRASS, BlockTypes.TALL_SEAGRASS);
+        } else {
+            liquidMask = new BlockTypeMask(this, BlockTypes.LAVA, BlockTypes.WATER);
+        }
         if (waterlogged) {
             Map<String, String> stateMap = new HashMap<>();
             stateMap.put("waterlogged", "true");
diff --git a/worldedit-core/src/main/java/com/sk89q/worldedit/command/UtilityCommands.java b/worldedit-core/src/main/java/com/sk89q/worldedit/command/UtilityCommands.java
index 58e20bc6c..b320dc300 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/command/UtilityCommands.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/command/UtilityCommands.java
@@ -326,11 +326,13 @@ public class UtilityCommands {
                      @Arg(desc = "The radius to drain")
                          Expression radiusExp,
                      @Switch(name = 'w', desc = "Also un-waterlog blocks")
-                         boolean waterlogged) throws WorldEditException {
+                         boolean waterlogged,
+                     @Switch(name = 'p', desc = "Also remove water plants")
+                         boolean plants) throws WorldEditException {
         double radius = radiusExp.evaluate();
         radius = Math.max(0, radius);
         we.checkMaxRadius(radius);
-        int affected = editSession.drainArea(session.getPlacementPosition(actor), radius, waterlogged);
+        int affected = editSession.drainArea(session.getPlacementPosition(actor), radius, waterlogged, plants);
         actor.printInfo(TranslatableComponent.of("worldedit.drain.drained", TextComponent.of(affected)));
         return affected;
     }