From 3850944a819768e9aa63e73c686a4ad31c18c578 Mon Sep 17 00:00:00 2001
From: Jesse Boyd <jessepaleg@gmail.com>
Date: Sat, 29 Jun 2019 04:49:48 +1000
Subject: [PATCH] Optimize waterlog remover

---
 .../function/mask/BlockMaskBuilder.java       |  4 --
 .../function/pattern/WaterloggedRemover.java  | 40 +++++++++++++++++--
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/worldedit-core/src/main/java/com/sk89q/worldedit/function/mask/BlockMaskBuilder.java b/worldedit-core/src/main/java/com/sk89q/worldedit/function/mask/BlockMaskBuilder.java
index b73b1d5d2..db4ea88e9 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/function/mask/BlockMaskBuilder.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/function/mask/BlockMaskBuilder.java
@@ -257,10 +257,6 @@ public class BlockMaskBuilder {
         this.bitSets = bitSets;
     }
 
-    public BlockMaskBuilder parse(String input) {
-        return this;
-    }
-
     public BlockMaskBuilder addAll() {
         for (int i = 0; i < bitSets.length; i++) {
             bitSets[i] = BlockMask.ALL;
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 51f8c0f67..e05e2d229 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
@@ -20,27 +20,59 @@
 package com.sk89q.worldedit.function.pattern;
 
 import com.sk89q.worldedit.extent.Extent;
+import com.sk89q.worldedit.function.mask.BlockMaskBuilder;
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.registry.state.Property;
+import com.sk89q.worldedit.registry.state.PropertyKey;
 import com.sk89q.worldedit.world.block.BaseBlock;
+import com.sk89q.worldedit.world.block.BlockState;
+import com.sk89q.worldedit.world.block.BlockType;
 import com.sk89q.worldedit.world.block.BlockTypes;
 
+import java.lang.ref.SoftReference;
+import java.util.Map;
+import java.util.function.BiPredicate;
+import java.util.function.Predicate;
+
 /**
  * Removes the waterlogged state from blocks if possible. If not possible, returns air.
  */
 public class WaterloggedRemover extends AbstractExtentPattern {
 
+    private static SoftReference<BlockState[]> cache = new SoftReference<>(null);
+
+    private synchronized BlockState[] getRemap() {
+         BlockState[] remap = this.cache.get();
+         if (remap != null) return remap;
+         this.cache = new SoftReference<>(remap = new BlockState[BlockTypes.states.length]);
+
+         // init
+        for (int i = 0; i < remap.length; i++) {
+            BlockState state = remap[i];
+            BlockType type = state.getBlockType();
+            if (!type.hasProperty(PropertyKey.WATERLOGGED)) {
+                continue;
+            }
+            if (state.getState(PropertyKey.WATERLOGGED) == Boolean.TRUE) {
+                remap[i] = state.with(PropertyKey.WATERLOGGED, false);
+            }
+        }
+        return remap;
+    }
+
+    private final BlockState[] remap;
+
     public WaterloggedRemover(Extent extent) {
         super(extent);
+        this.remap = getRemap();
     }
 
     @Override
     public BaseBlock apply(BlockVector3 position) {
         BaseBlock block = getExtent().getFullBlock(position);
-        @SuppressWarnings("unchecked")
-        Property<Object> prop = (Property<Object>) block.getBlockType().getPropertyMap().getOrDefault("waterlogged", null);
-        if (prop != null) {
-            return block.with(prop, false);
+        BlockState newState = remap[block.getOrdinal()];
+        if (newState != null) {
+            return newState.toBaseBlock(block.getNbtData());
         }
         return BlockTypes.AIR.getDefaultState().toBaseBlock();
     }