Add optional draining of water-supplying blocks

worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java

@@ -61,6 +61,7 @@
 import com.sk89q.worldedit.function.mask.BlockStateMask;
 import com.sk89q.worldedit.function.mask.BlockTypeMask;
 import com.sk89q.worldedit.function.mask.BoundedHeightMask;
+import com.sk89q.worldedit.function.mask.WaterSupplyingBlockMask;
 import com.sk89q.worldedit.function.mask.ExistingBlockMask;
 import com.sk89q.worldedit.function.mask.Mask;
 import com.sk89q.worldedit.function.mask.MaskIntersection;
@@ -1654,20 +1655,38 @@ public int drainArea(BlockVector3 origin, double radius) throws MaxChangedBlocks
      * @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 with optional removal of waterlogged and water-supplying 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 waterSupplyingBlocks true to also remove water-supplying blocks such as kelp and bubble columns
+     * @return number of blocks affected
+     * @throws MaxChangedBlocksException thrown if too many blocks are changed
+     */
+    public int drainArea(BlockVector3 origin, double radius, boolean waterlogged, boolean waterSupplyingBlocks)
+            throws MaxChangedBlocksException {
         checkNotNull(origin);
         checkArgument(radius >= 0, "radius >= 0 required");
 
-        Mask waterloggedMask = null;
+        Mask drainMask = getWorld().createLiquidMask();
+        if (waterSupplyingBlocks) {
+            drainMask = new MaskUnion(drainMask, new WaterSupplyingBlockMask(this));
+        }
         if (waterlogged) {
             Map<String, String> stateMap = new HashMap<>();
             stateMap.put("waterlogged", "true");
-            waterloggedMask = new BlockStateMask(this, stateMap, true);
+            drainMask = new MaskUnion(drainMask, new BlockStateMask(this, stateMap, true));
         }
+
         MaskIntersection mask = new MaskIntersection(
                 new BoundedHeightMask(getWorld().getMinY(), getWorld().getMaxY()),
                 new RegionMask(new EllipsoidRegion(null, origin, Vector3.at(radius, radius, radius))),
-                waterlogged ? new MaskUnion(getWorld().createLiquidMask(), waterloggedMask)
-                            : getWorld().createLiquidMask());
+                drainMask);
 
         BlockReplace replace;
         if (waterlogged) {

worldedit-core/src/main/java/com/sk89q/worldedit/command/UtilityCommands.java

@@ -142,10 +142,13 @@ public int drain(Actor actor, LocalSession session, EditSession editSession,
                      @Arg(desc = "The radius to drain")
                          double radius,
                      @Switch(name = 'w', desc = "Also un-waterlog blocks")
-                         boolean waterlogged) throws WorldEditException {
+                         boolean waterlogged,
+                     @Switch(name = 'b', desc = "Also remove blocks that supply water (e.g. kelp and bubble columns)")
+                         boolean waterSupplyingBlocks) throws WorldEditException {
         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, waterSupplyingBlocks);
         actor.printInfo(TranslatableComponent.of("worldedit.drain.drained", TextComponent.of(affected)));
         return affected;
     }

worldedit-core/src/main/java/com/sk89q/worldedit/function/mask/WaterSupplyingBlockMask.java

@@ -0,0 +1,47 @@
+/*
+ * WorldEdit, a Minecraft world manipulation toolkit
+ * Copyright (C) sk89q <http://www.sk89q.com>
+ * Copyright (C) WorldEdit team and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.sk89q.worldedit.function.mask;
+
+import com.sk89q.worldedit.extent.Extent;
+import com.sk89q.worldedit.math.BlockVector3;
+import com.sk89q.worldedit.world.fluid.FluidType;
+import com.sk89q.worldedit.world.fluid.FluidTypes;
+
+/**
+ * Matches blocks which supply water, including non-liquid blocks such as
+ * bubble columns and kelp.
+ */
+public class WaterSupplyingBlockMask extends AbstractExtentMask {
+
+    /**
+     * Create a mask for blocks whose fluid state is water.
+     *
+     * @param extent the extent containing the blocks
+     */
+    public WaterSupplyingBlockMask(Extent extent) {
+        super(extent);
+    }
+
+    @Override
+    public boolean test(BlockVector3 vector) {
+        FluidType fluid = getExtent().getBlock(vector).getFluidState().getType();
+        return FluidTypes.WATER.equals(fluid) || FluidTypes.FLOWING_WATER.equals(fluid);
+    }
+}

worldedit-core/src/test/java/com/sk89q/worldedit/function/mask/WaterSupplyingBlockMaskTest.java

@@ -0,0 +1,70 @@
+/*
+ * WorldEdit, a Minecraft world manipulation toolkit
+ * Copyright (C) sk89q <http://www.sk89q.com>
+ * Copyright (C) WorldEdit team and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.sk89q.worldedit.function.mask;
+
+import com.sk89q.worldedit.BaseWorldEditTest;
+import com.sk89q.worldedit.extent.Extent;
+import com.sk89q.worldedit.math.BlockVector3;
+import com.sk89q.worldedit.world.block.BlockState;
+import com.sk89q.worldedit.world.fluid.FluidState;
+import com.sk89q.worldedit.world.fluid.FluidType;
+import com.sk89q.worldedit.world.fluid.FluidTypes;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+class WaterSupplyingBlockMaskTest extends BaseWorldEditTest {
+
+    private static final BlockVector3 POSITION = BlockVector3.at(0, 64, 0);
+
+    @Test
+    void matchesStillWater() {
+        assertTrue(maskFor(FluidTypes.WATER).test(POSITION));
+    }
+
+    @Test
+    void matchesFlowingWater() {
+        assertTrue(maskFor(FluidTypes.FLOWING_WATER).test(POSITION));
+    }
+
+    @Test
+    void doesNotMatchLava() {
+        assertFalse(maskFor(FluidTypes.LAVA).test(POSITION));
+        assertFalse(maskFor(FluidTypes.FLOWING_LAVA).test(POSITION));
+    }
+
+    @Test
+    void doesNotMatchEmptyFluid() {
+        assertFalse(maskFor(FluidTypes.EMPTY).test(POSITION));
+    }
+
+    private static WaterSupplyingBlockMask maskFor(FluidType fluidType) {
+        Extent extent = mock(Extent.class);
+        BlockState block = mock(BlockState.class);
+        FluidState fluid = mock(FluidState.class);
+        when(extent.getBlock(POSITION)).thenReturn(block);
+        when(block.getFluidState()).thenReturn(fluid);
+        when(fluid.getType()).thenReturn(fluidType);
+        return new WaterSupplyingBlockMask(extent);
+    }
+}