From ea7897934f93243764ac8e4566e0475557bdf68e Mon Sep 17 00:00:00 2001
From: dordsor21 <dordsor21@gmail.com>
Date: Fri, 3 Jul 2020 13:49:17 +0100
Subject: [PATCH] Implement ability to use -w flag on //drain fixes #464

---
 .../java/com/sk89q/worldedit/EditSession.java     | 15 ++++++++-------
 .../function/pattern/WaterloggedRemover.java      |  2 +-
 2 files changed, 9 insertions(+), 8 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 b3679aacb..ce1b78f05 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java
@@ -63,6 +63,7 @@ import com.sk89q.worldedit.function.block.BlockReplace;
 import com.sk89q.worldedit.function.block.Naturalizer;
 import com.sk89q.worldedit.function.generator.ForestGenerator;
 import com.sk89q.worldedit.function.generator.GardenPatchGenerator;
+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.ExistingBlockMask;
@@ -135,6 +136,7 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
@@ -1624,13 +1626,12 @@ public class EditSession extends PassthroughExtent implements AutoCloseable {
         checkNotNull(origin);
         checkArgument(radius >= 0, "radius >= 0 required");
 
-        Mask liquidMask;
-        // Not thread safe, use hardcoded liquidmask
-//        if (getWorld() != null) {
-//            liquidMask = getWorld().createLiquidMask();
-//        } else {
-//        }
-        liquidMask = new BlockTypeMask(this, BlockTypes.LAVA, BlockTypes.WATER);
+        Mask liquidMask = new BlockTypeMask(this, BlockTypes.LAVA, BlockTypes.WATER);
+        if (waterlogged) {
+            Map<String, String> stateMap = new HashMap<>();
+            stateMap.put("waterlogged", "true");
+            liquidMask = new MaskUnion(liquidMask, new BlockStateMask(this, stateMap, true));
+        }
         Mask mask = new MaskIntersection(
                 new BoundedHeightMask(0, getWorld().getMaxY()),
                 new RegionMask(new EllipsoidRegion(null, origin, Vector3.at(radius, radius, radius))),
diff --git a/worldedit-core/src/main/java/com/sk89q/worldedit/function/pattern/WaterloggedRemover.java b/worldedit-core/src/main/java/com/sk89q/worldedit/function/pattern/WaterloggedRemover.java
index 886706f08..a8610b80d 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/function/pattern/WaterloggedRemover.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/function/pattern/WaterloggedRemover.java
@@ -45,7 +45,7 @@ public class WaterloggedRemover extends AbstractExtentPattern {
 
         // init
         for (int i = 0; i < remap.length; i++) {
-            BlockState state = remap[i];
+            BlockState state = BlockTypesCache.states[i];
             BlockType type = state.getBlockType();
             if (!type.hasProperty(PropertyKey.WATERLOGGED)) {
                 continue;