new: collections.fluid_submersion: Use faster collection and avoid Fl…

…uidTag equality tests

lithium-mixin-config.md

@@ -220,6 +220,10 @@ The expensive check to see if a TypeFilterableList can be filtered by a specific
 (default: `true`)  
 Copy entity hashmap instead of duplicating the list using iteration  
 
+### `mixin.collections.fluid_submersion`
+(default: `true`)  
+Use ReferenceArraySet instead of HashSet to store the fluids the entity is currently submerged in.  
+
 ### `mixin.collections.gamerules`
 (default: `true`)  
 Uses fastutil hashmaps for gamerules  

src/main/java/me/jellysquid/mods/lithium/mixin/collections/fluid_submersion/EntityMixin.java

@@ -0,0 +1,44 @@
+package me.jellysquid.mods.lithium.mixin.collections.fluid_submersion;
+
+import it.unimi.dsi.fastutil.objects.ReferenceArraySet;
+import net.minecraft.entity.Entity;
+import net.minecraft.fluid.Fluid;
+import net.minecraft.registry.tag.TagKey;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Mutable;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.Redirect;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import java.util.Set;
+
+@Mixin(Entity.class)
+public class EntityMixin {
+
+    @Mutable
+    @Shadow
+    @Final
+    private Set<TagKey<Fluid>> submergedFluidTag;
+
+    @Inject(
+            method = "<init>",
+            at = @At("RETURN")
+    )
+    private void useReferenceArraySet(CallbackInfo ci) {
+        this.submergedFluidTag = new ReferenceArraySet<>(this.submergedFluidTag);
+    }
+
+    @Redirect(
+            method = "updateSubmergedInWaterState()V",
+            at = @At(value = "INVOKE", target = "Ljava/util/Set;clear()V"),
+            require = 0
+    )
+    private void clearIfNotEmpty(Set<?> set) {
+        if (!set.isEmpty()) {
+            set.clear();
+        }
+    }
+}

src/main/java/me/jellysquid/mods/lithium/mixin/collections/fluid_submersion/package-info.java

@@ -0,0 +1,4 @@
+@MixinConfigOption(description = "Use ReferenceArraySet instead of HashSet to store the fluids the entity is currently submerged in.")
+package me.jellysquid.mods.lithium.mixin.collections.fluid_submersion;
+
+import net.caffeinemc.gradle.MixinConfigOption;

src/main/resources/lithium.mixins.json

@@ -97,6 +97,7 @@
     "collections.entity_by_type.TypeFilterableListMixin",
     "collections.entity_filtering.TypeFilterableListMixin",
     "collections.entity_ticking.EntityListMixin",
+    "collections.fluid_submersion.EntityMixin",
     "collections.gamerules.GameRulesMixin",
     "collections.goals.GoalSelectorMixin",
     "collections.mob_spawning.PoolMixin",