Improve API design and reduce duplicate code

worldedit-bukkit/adapters/adapter-1_17_1/src/main/java/com/sk89q/worldedit/bukkit/adapter/impl/fawe/v1_17_R1_2/PaperweightFaweAdapter.java

@@ -88,6 +88,7 @@ import org.bukkit.craftbukkit.v1_17_R1.util.CraftNamespacedKey;
 import org.bukkit.entity.Player;
 
 import javax.annotation.Nullable;
+import java.lang.invoke.MethodHandle;
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
 import java.util.ArrayList;
@@ -555,6 +556,11 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
         return ((CraftWorld) world).getHandle();
     }
 
+    @Override
+    protected MethodHandle getCurrentWorldData() {
+        return null;
+    }
+
     @Override
     public List<org.bukkit.entity.Entity> getEntities(org.bukkit.World world) {
         // Quickly add each entity to a list copy.

worldedit-bukkit/adapters/adapter-1_18_2/src/main/java/com/sk89q/worldedit/bukkit/adapter/impl/fawe/v1_18_R2/PaperweightFaweAdapter.java

@@ -89,6 +89,7 @@ import org.bukkit.craftbukkit.v1_18_R2.util.CraftNamespacedKey;
 import org.bukkit.entity.Player;
 
 import javax.annotation.Nullable;
+import java.lang.invoke.MethodHandle;
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
 import java.util.ArrayList;
@@ -547,6 +548,11 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
         return ((CraftWorld) world).getHandle();
     }
 
+    @Override
+    protected MethodHandle getCurrentWorldData() {
+        return null;
+    }
+
     @Override
     public List<org.bukkit.entity.Entity> getEntities(org.bukkit.World world) {
         // Quickly add each entity to a list copy.

worldedit-bukkit/adapters/adapter-1_19_4/src/main/java/com/sk89q/worldedit/bukkit/adapter/impl/fawe/v1_19_R3/PaperweightFaweAdapter.java

@@ -78,6 +78,7 @@ import org.bukkit.World;
 import org.bukkit.block.data.BlockData;
 import org.bukkit.craftbukkit.v1_19_R3.CraftServer;
 import org.bukkit.craftbukkit.v1_19_R3.CraftWorld;
+import org.bukkit.craftbukkit.v1_19_R3.block.CraftBlockState;
 import org.bukkit.craftbukkit.v1_19_R3.block.data.CraftBlockData;
 import org.bukkit.craftbukkit.v1_19_R3.entity.CraftEntity;
 import org.bukkit.craftbukkit.v1_19_R3.entity.CraftPlayer;
@@ -129,45 +130,8 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
     private boolean initialised = false;
     private Map<String, List<Property<?>>> allBlockProperties = null;
 
-    // Folia - START
-    private MethodHandle currentWorldData;
-
-    private Class<?> regionizedWorldData;
-
-    private Field captureTreeGeneration;
-    private Field captureBlockStates;
-    private Field capturedBlockStates;
-
     public PaperweightFaweAdapter() throws NoSuchFieldException, NoSuchMethodException {
         this.parent = new PaperweightAdapter();
-        if (this.parent.isFolia()) {
-            Method getCurrentWorldData = ServerLevel.class.getSuperclass().getDeclaredMethod(
-                    "getCurrentWorldData"
-            );
-            getCurrentWorldData.setAccessible(true);
-            try {
-                currentWorldData = MethodHandles.lookup().unreflect(getCurrentWorldData);
-            } catch (IllegalAccessException e) {
-            }
-
-            try {
-                regionizedWorldData = Class.forName("io.papermc.paper.threadedregions.RegionizedWorldData");
-            } catch (ClassNotFoundException e) {
-            }
-            if (regionizedWorldData != null) {
-                final Field captureTreeGeneration = regionizedWorldData.getDeclaredField("captureTreeGeneration");
-                captureTreeGeneration.setAccessible(true);
-                this.captureTreeGeneration = captureTreeGeneration;
-
-                final Field captureBlockStates = regionizedWorldData.getDeclaredField("captureBlockStates");
-                captureBlockStates.setAccessible(true);
-                this.captureBlockStates = captureBlockStates;
-
-                final Field capturedBlockStates = regionizedWorldData.getDeclaredField("capturedBlockStates");
-                capturedBlockStates.setAccessible(true);
-                this.capturedBlockStates = capturedBlockStates;
-            }
-        }
     }
 
     @Nullable
@@ -630,8 +594,19 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
 
     @Override
     protected void preCaptureStates(final ServerLevel serverLevel) {
-        serverLevel.captureTreeGeneration = true;
-        serverLevel.captureBlockStates = true;
+        if (this.isFolia()) {
+            try {
+                var data = currentWorldData.invoke(serverLevel);
+                captureBlockStates.setBoolean(data, true);
+                captureTreeGeneration.setBoolean(data, true);
+            } catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            serverLevel.captureTreeGeneration = true;
+            serverLevel.captureBlockStates = true;
+        }
+
     }
 
     @Override
@@ -639,16 +614,54 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
         return new ArrayList<>(serverLevel.capturedBlockStates.values());
     }
 
+    // Folia - Start
     @Override
     protected void postCaptureBlockStates(final ServerLevel serverLevel) {
-        serverLevel.captureBlockStates = false;
-        serverLevel.captureTreeGeneration = false;
-        serverLevel.capturedBlockStates.clear();
+        if (this.isFolia()) {
+            try {
+                var data = currentWorldData.invoke(serverLevel);
+                captureBlockStates.setBoolean(data, false);
+                captureTreeGeneration.setBoolean(data, false);
+                ((Map<BlockPos, CraftBlockState>) capturedBlockStates.get(data)).clear();
+            } catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            serverLevel.captureBlockStates = false;
+            serverLevel.captureTreeGeneration = false;
+            serverLevel.capturedBlockStates.clear();
+        }
     }
 
+    // Folia - End
+
     @Override
     protected ServerLevel getServerLevel(final World world) {
         return ((CraftWorld) world).getHandle();
     }
 
+    // Folia - Start
+    @Override
+    protected MethodHandle getCurrentWorldData() {
+        if (this.isFolia()) {
+            Method getCurrentWorldData;
+            try {
+                getCurrentWorldData = ServerLevel.class.getSuperclass().getDeclaredMethod(
+                        "getCurrentWorldData"
+                );
+            } catch (NoSuchMethodException e) {
+                return null;
+            }
+            getCurrentWorldData.setAccessible(true);
+
+            try {
+                currentWorldData = MethodHandles.lookup().unreflect(getCurrentWorldData);
+            } catch (IllegalAccessException e) {
+            }
+        }
+        return null;
+    }
+
+    // Folia - End
+
 }

worldedit-bukkit/adapters/adapter-1_20/src/main/java/com/sk89q/worldedit/bukkit/adapter/impl/fawe/v1_20_R1/PaperweightFaweAdapter.java

@@ -77,6 +77,7 @@ import org.bukkit.World;
 import org.bukkit.block.data.BlockData;
 import org.bukkit.craftbukkit.v1_20_R1.CraftServer;
 import org.bukkit.craftbukkit.v1_20_R1.CraftWorld;
+import org.bukkit.craftbukkit.v1_20_R1.block.CraftBlockState;
 import org.bukkit.craftbukkit.v1_20_R1.block.data.CraftBlockData;
 import org.bukkit.craftbukkit.v1_20_R1.entity.CraftEntity;
 import org.bukkit.craftbukkit.v1_20_R1.entity.CraftPlayer;
@@ -128,47 +129,9 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
     private boolean initialised = false;
     private Map<String, List<Property<?>>> allBlockProperties = null;
 
-    // Folia - START
-    private MethodHandle currentWorldData;
-
-    private Class<?> regionizedWorldData;
-
-    private Field captureTreeGeneration;
-    private Field captureBlockStates;
-    private Field capturedBlockStates;
-
 
     public PaperweightFaweAdapter() throws NoSuchFieldException, NoSuchMethodException {
         this.parent = new PaperweightAdapter();
-
-        if (this.parent.isFolia()) {
-            Method getCurrentWorldData = ServerLevel.class.getSuperclass().getDeclaredMethod(
-                    "getCurrentWorldData"
-            );
-            getCurrentWorldData.setAccessible(true);
-            try {
-                currentWorldData = MethodHandles.lookup().unreflect(getCurrentWorldData);
-            } catch (IllegalAccessException e) {
-            }
-
-            try {
-                regionizedWorldData = Class.forName("io.papermc.paper.threadedregions.RegionizedWorldData");
-            } catch (ClassNotFoundException e) {
-            }
-            if (regionizedWorldData != null) {
-                final Field captureTreeGeneration = regionizedWorldData.getDeclaredField("captureTreeGeneration");
-                captureTreeGeneration.setAccessible(true);
-                this.captureTreeGeneration = captureTreeGeneration;
-
-                final Field captureBlockStates = regionizedWorldData.getDeclaredField("captureBlockStates");
-                captureBlockStates.setAccessible(true);
-                this.captureBlockStates = captureBlockStates;
-
-                final Field capturedBlockStates = regionizedWorldData.getDeclaredField("capturedBlockStates");
-                capturedBlockStates.setAccessible(true);
-                this.capturedBlockStates = capturedBlockStates;
-            }
-        }
     }
 
     @Nullable
@@ -528,29 +491,79 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
         return CraftItemStack.asCraftMirror(stack);
     }
 
+    // Folia - Start
     @Override
     protected void preCaptureStates(final ServerLevel serverLevel) {
-        serverLevel.captureTreeGeneration = true;
-        serverLevel.captureBlockStates = true;
+        if (this.isFolia()) {
+            try {
+                var data = currentWorldData.invoke(serverLevel);
+                captureBlockStates.setBoolean(data, true);
+                captureTreeGeneration.setBoolean(data, true);
+                ((Map<BlockPos, CraftBlockState>) capturedBlockStates.get(data)).clear();
+            } catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            serverLevel.captureTreeGeneration = true;
+            serverLevel.captureBlockStates = true;
+        }
     }
+    // Folia - End
 
     @Override
     protected List<org.bukkit.block.BlockState> getCapturedBlockStatesCopy(final ServerLevel serverLevel) {
         return new ArrayList<>(serverLevel.capturedBlockStates.values());
     }
 
+    // Folia - Start
     @Override
     protected void postCaptureBlockStates(final ServerLevel serverLevel) {
-        serverLevel.captureBlockStates = false;
-        serverLevel.captureTreeGeneration = false;
-        serverLevel.capturedBlockStates.clear();
+        if (this.isFolia()) {
+            try {
+                var data = currentWorldData.invoke(serverLevel);
+                captureBlockStates.setBoolean(data, false);
+                captureTreeGeneration.setBoolean(data, false);
+                ((Map<BlockPos, CraftBlockState>) capturedBlockStates.get(data)).clear();
+            } catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            serverLevel.captureBlockStates = false;
+            serverLevel.captureTreeGeneration = false;
+            serverLevel.capturedBlockStates.clear();
+        }
     }
+    // Folia - End
 
     @Override
     protected ServerLevel getServerLevel(final World world) {
         return ((CraftWorld) world).getHandle();
     }
 
+    // Folia - Start
+    @Override
+    protected MethodHandle getCurrentWorldData() {
+        if (this.isFolia()) {
+            Method getCurrentWorldData;
+            try {
+                getCurrentWorldData = ServerLevel.class.getSuperclass().getDeclaredMethod(
+                        "getCurrentWorldData"
+                );
+            } catch (NoSuchMethodException e) {
+                return null;
+            }
+            getCurrentWorldData.setAccessible(true);
+
+            try {
+                currentWorldData = MethodHandles.lookup().unreflect(getCurrentWorldData);
+            } catch (IllegalAccessException e) {
+            }
+        }
+        return null;
+    }
+
+    // Folia - End
+
     @Override
     public BaseItemStack adapt(org.bukkit.inventory.ItemStack itemStack) {
         final ItemStack nmsStack = CraftItemStack.asNMSCopy(itemStack);

worldedit-bukkit/adapters/adapter-1_20_2/src/main/java/com/sk89q/worldedit/bukkit/adapter/impl/fawe/v1_20_R2/PaperweightFaweAdapter.java

@@ -76,6 +76,7 @@ import org.bukkit.World;
 import org.bukkit.block.data.BlockData;
 import org.bukkit.craftbukkit.v1_20_R2.CraftServer;
 import org.bukkit.craftbukkit.v1_20_R2.CraftWorld;
+import org.bukkit.craftbukkit.v1_20_R2.block.CraftBlockState;
 import org.bukkit.craftbukkit.v1_20_R2.block.data.CraftBlockData;
 import org.bukkit.craftbukkit.v1_20_R2.entity.CraftEntity;
 import org.bukkit.craftbukkit.v1_20_R2.entity.CraftPlayer;
@@ -84,6 +85,8 @@ import org.bukkit.craftbukkit.v1_20_R2.util.CraftNamespacedKey;
 import org.bukkit.entity.Player;
 
 import javax.annotation.Nullable;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
@@ -492,8 +495,20 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
 
     @Override
     protected void preCaptureStates(final ServerLevel serverLevel) {
-        serverLevel.captureTreeGeneration = true;
-        serverLevel.captureBlockStates = true;
+        // Folia - Start
+        if (this.isFolia()) {
+            try {
+                var data = currentWorldData.invoke(serverLevel);
+                captureBlockStates.setBoolean(data, true);
+                captureTreeGeneration.setBoolean(data, true);
+            } catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            serverLevel.captureTreeGeneration = true;
+            serverLevel.captureBlockStates = true;
+        }
+        // Folia - End
     }
 
     @Override
@@ -503,9 +518,22 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
 
     @Override
     protected void postCaptureBlockStates(final ServerLevel serverLevel) {
-        serverLevel.captureBlockStates = false;
-        serverLevel.captureTreeGeneration = false;
-        serverLevel.capturedBlockStates.clear();
+        // Folia - Start
+        if (this.isFolia()) {
+            try {
+                var data = currentWorldData.invoke(serverLevel);
+                captureBlockStates.setBoolean(data, false);
+                captureTreeGeneration.setBoolean(data, false);
+                ((Map<BlockPos, CraftBlockState>) capturedBlockStates.get(data)).clear();
+            } catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            serverLevel.captureBlockStates = false;
+            serverLevel.captureTreeGeneration = false;
+            serverLevel.capturedBlockStates.clear();
+        }
+        // Folia - End
     }
 
     @Override
@@ -513,6 +541,29 @@ public final class PaperweightFaweAdapter extends FaweAdapter<net.minecraft.nbt.
         return ((CraftWorld) world).getHandle();
     }
 
+    // Folia - Start
+    @Override
+    protected MethodHandle getCurrentWorldData() {
+        if (this.isFolia()) {
+            Method getCurrentWorldData;
+            try {
+                getCurrentWorldData = ServerLevel.class.getSuperclass().getDeclaredMethod(
+                        "getCurrentWorldData"
+                );
+            } catch (NoSuchMethodException e) {
+                return null;
+            }
+            getCurrentWorldData.setAccessible(true);
+
+            try {
+                currentWorldData = MethodHandles.lookup().unreflect(getCurrentWorldData);
+            } catch (IllegalAccessException e) {
+            }
+        }
+        return null;
+    }
+    // Folia - End
+
     @Override
     public BaseItemStack adapt(org.bukkit.inventory.ItemStack itemStack) {
         final ItemStack nmsStack = CraftItemStack.asNMSCopy(itemStack);

worldedit-bukkit/src/main/java/com/fastasyncworldedit/bukkit/adapter/FaweAdapter.java

@@ -11,6 +11,8 @@ import org.bukkit.TreeType;
 import org.bukkit.World;
 import org.bukkit.block.BlockState;
 
+import java.lang.invoke.MethodHandle;
+import java.lang.reflect.Field;
 import java.util.List;
 
 /**
@@ -21,6 +23,32 @@ import java.util.List;
  */
 public abstract class FaweAdapter<TAG, SERVER_LEVEL> extends CachedBukkitAdapter implements IDelegateBukkitImplAdapter<TAG> {
 
+
+    // Folia - Start
+    protected MethodHandle currentWorldData = getCurrentWorldData();
+
+    protected Class<?> regionizedWorldData = getRegionitedWorldData();
+
+    protected Field captureTreeGeneration = getCaptureTreeGeneration();
+
+    protected Field captureBlockStates = getCaptureBlockStates();
+
+    protected Field capturedBlockStates = getCapturedBlockStates();
+
+    private final boolean folia;
+
+    protected FaweAdapter() {
+        boolean isFolia = false;
+        try {
+            // Assume API is present
+            Class.forName("io.papermc.paper.threadedregions.scheduler.EntityScheduler");
+            isFolia = true;
+        } catch (Exception unused) {
+
+        }
+        this.folia = isFolia;
+    }
+
     @Override
     public boolean generateTree(
             final TreeGenerator.TreeType treeType,
@@ -69,4 +97,55 @@ public abstract class FaweAdapter<TAG, SERVER_LEVEL> extends CachedBukkitAdapter
 
     protected abstract SERVER_LEVEL getServerLevel(World world);
 
+    // Folia Support
+    protected abstract MethodHandle getCurrentWorldData();
+
+    private Class<?> getRegionitedWorldData() {
+        try {
+            regionizedWorldData = Class.forName("io.papermc.paper.threadedregions.RegionizedWorldData");
+        } catch (ClassNotFoundException e) {
+        }
+        return null;
+    }
+
+    private Field getCaptureTreeGeneration() {
+        Field captureTreeGeneration = null;
+        if (regionizedWorldData != null) {
+            try {
+                captureTreeGeneration = regionizedWorldData.getDeclaredField("captureTreeGeneration");
+                captureTreeGeneration.setAccessible(true);
+            } catch (NoSuchFieldException e) {
+            }
+        }
+        return captureTreeGeneration;
+    }
+
+    private Field getCaptureBlockStates() {
+        Field captureBlockStates = null;
+        if (regionizedWorldData != null) {
+            try {
+                captureBlockStates = regionizedWorldData.getDeclaredField("captureBlockStates");
+                captureBlockStates.setAccessible(true);
+            } catch (NoSuchFieldException e) {
+            }
+        }
+        return captureBlockStates;
+    }
+
+    private Field getCapturedBlockStates() {
+        Field capturedBlockStates = null;
+        if (regionizedWorldData != null) {
+            try {
+                capturedBlockStates = regionizedWorldData.getDeclaredField("capturedBlockStates");
+                capturedBlockStates.setAccessible(true);
+            } catch (NoSuchFieldException e) {
+            }
+        }
+        return capturedBlockStates;
+    }
+
+    protected boolean isFolia() {
+        return this.folia;
+    }
+
 }