From 2c4d96ea1722ea284b1245802cf22e67f082a273 Mon Sep 17 00:00:00 2001
From: NichtStudioCode <51272202+NichtStudioCode@users.noreply.github.com>
Date: Thu, 12 Sep 2024 10:27:07 +0200
Subject: [PATCH] Fix #452

---
 .../chunksection/LevelChunkSectionWrapper.java | 18 ++++++++++++++++++
 .../nova/util/reflection/ReflectionUtils.kt    |  9 +++++++++
 2 files changed, 27 insertions(+)

diff --git a/nova/src/main/kotlin/xyz/xenondevs/nova/patch/impl/worldgen/chunksection/LevelChunkSectionWrapper.java b/nova/src/main/kotlin/xyz/xenondevs/nova/patch/impl/worldgen/chunksection/LevelChunkSectionWrapper.java
index ba69cab44a0..4319ee53875 100644
--- a/nova/src/main/kotlin/xyz/xenondevs/nova/patch/impl/worldgen/chunksection/LevelChunkSectionWrapper.java
+++ b/nova/src/main/kotlin/xyz/xenondevs/nova/patch/impl/worldgen/chunksection/LevelChunkSectionWrapper.java
@@ -32,6 +32,7 @@
 @ApiStatus.Internal
 public class LevelChunkSectionWrapper extends LevelChunkSection {
     
+    // Paper
     private static final MethodHandle GET_STATES;
     private static final MethodHandle GET_BIOMES;
     private static final MethodHandle GET_NON_EMPTY_BLOCK_COUNT;
@@ -44,9 +45,14 @@ public class LevelChunkSectionWrapper extends LevelChunkSection {
     private static final MethodHandle SET_SPECIAL_COLLIDING_BLOCKS;
     private static final Field TICKING_BLOCKS;
     
+    // Pufferfish
+    private static final MethodHandle GET_FLUID_STATE_COUNT;
+    private static final MethodHandle SET_FLUID_STATE_COUNT;
+    
     static {
         try {
             var lookup = MethodHandles.privateLookupIn(LevelChunkSection.class, MethodHandles.lookup());
+            
             GET_STATES = lookup.findGetter(LevelChunkSection.class, "states", PalettedContainer.class);
             GET_BIOMES = lookup.findGetter(LevelChunkSection.class, "biomes", PalettedContainer.class);
             GET_NON_EMPTY_BLOCK_COUNT = lookup.findGetter(LevelChunkSection.class, "nonEmptyBlockCount", short.class);
@@ -58,6 +64,15 @@ public class LevelChunkSectionWrapper extends LevelChunkSection {
             GET_SPECIAL_COLLIDING_BLOCKS = lookup.findGetter(LevelChunkSection.class, "specialCollidingBlocks", int.class);
             SET_SPECIAL_COLLIDING_BLOCKS = lookup.findSetter(LevelChunkSection.class, "specialCollidingBlocks", int.class);
             TICKING_BLOCKS = ReflectionUtils.getField(LevelChunkSection.class, "tickingBlocks");
+            
+            Field fluidStateCount = ReflectionUtils.getFieldOrNull(LevelChunkSection.class, "fluidStateCount");
+            if (fluidStateCount != null) {
+                GET_FLUID_STATE_COUNT = lookup.unreflectGetter(fluidStateCount);
+                SET_FLUID_STATE_COUNT = lookup.unreflectSetter(fluidStateCount);
+            } else {
+                GET_FLUID_STATE_COUNT = null;
+                SET_FLUID_STATE_COUNT = null;
+            }
         } catch (IllegalAccessException | NoSuchFieldException e) {
             throw new RuntimeException(e);
         }
@@ -152,6 +167,9 @@ private void copyBlockCounts() {
             SET_TICKING_BLOCK_COUNT.invoke(this, GET_TICKING_BLOCK_COUNT.invoke(delegate));
             SET_TICKING_FLUID_COUNT.invoke(this, GET_TICKING_FLUID_COUNT.invoke(delegate));
             SET_SPECIAL_COLLIDING_BLOCKS.invoke(this, GET_SPECIAL_COLLIDING_BLOCKS.invoke(delegate));
+            if (GET_FLUID_STATE_COUNT != null && SET_FLUID_STATE_COUNT != null) {
+                SET_FLUID_STATE_COUNT.invoke(this, GET_FLUID_STATE_COUNT.invoke(delegate));
+            }
         } catch (Throwable e) {
             throw new RuntimeException(e);
         }
diff --git a/nova/src/main/kotlin/xyz/xenondevs/nova/util/reflection/ReflectionUtils.kt b/nova/src/main/kotlin/xyz/xenondevs/nova/util/reflection/ReflectionUtils.kt
index aa1060790c7..3cd5ed2d23c 100644
--- a/nova/src/main/kotlin/xyz/xenondevs/nova/util/reflection/ReflectionUtils.kt
+++ b/nova/src/main/kotlin/xyz/xenondevs/nova/util/reflection/ReflectionUtils.kt
@@ -135,6 +135,15 @@ internal object ReflectionUtils {
         return field
     }
     
+    @JvmStatic
+    fun getFieldOrNull(clazz: Class<*>, name: String): Field? {
+         try {
+            return getField(clazz, name)
+        } catch (_: NoSuchFieldException) {
+            return null
+        }
+    }
+    
     @JvmStatic
     fun getServerSoftwareField(clazz: KClass<*>, name: String, serverSoftware: ServerSoftware): Field? {
         if (serverSoftware !in ServerUtils.SERVER_SOFTWARE.superSoftwares)