Move entity attribute stuff to a separate class

src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/BubbleSpell.java

@@ -12,6 +12,7 @@ import com.minelittlepony.unicopia.ability.magic.spell.attribute.TooltipFactory;
 import com.minelittlepony.unicopia.ability.magic.spell.trait.SpellTraits;
 import com.minelittlepony.unicopia.ability.magic.spell.trait.Trait;
 import com.minelittlepony.unicopia.entity.*;
+import com.minelittlepony.unicopia.entity.AttributeContainer;
 import com.minelittlepony.unicopia.entity.mob.UEntityAttributes;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.network.track.DataTracker;
@@ -22,7 +23,6 @@ import com.minelittlepony.unicopia.projectile.ProjectileDelegate;
 import com.minelittlepony.unicopia.util.shape.Sphere;
 
 import net.minecraft.entity.Entity;
-import net.minecraft.entity.LivingEntity;
 import net.minecraft.entity.attribute.*;
 import net.minecraft.entity.attribute.EntityAttributeModifier.Operation;
 import net.minecraft.nbt.NbtCompound;
@@ -86,12 +86,8 @@ public class BubbleSpell extends AbstractSpell implements TimedSpell,
 
         Entity entity = source.asEntity();
 
-        if (entity instanceof LivingEntity l) {
+        if (source instanceof AttributeContainer l) {
-            MODIFIERS.forEach((attribute, modifier) -> {
+            l.applyAttributeModifiers(MODIFIERS, false, true);
-                if (!l.getAttributes().hasAttribute(attribute)) {
-                    l.getAttributeInstance(attribute).addPersistentModifier(modifier);
-                }
-            });
         }
         radius.set(Math.max(entity.getHeight(), entity.getWidth()) * 1.2F);
         source.playSound(USounds.ENTITY_PLAYER_UNICORN_TELEPORT, 1);
@@ -148,12 +144,8 @@ public class BubbleSpell extends AbstractSpell implements TimedSpell,
     @Override
     protected void onDestroyed(Caster<?> source) {
         super.onDestroyed(source);
-        if (source.asEntity() instanceof LivingEntity l) {
+        if (source instanceof AttributeContainer l) {
-            MODIFIERS.forEach((attribute, modifier) -> {
+            l.applyAttributeModifiers(MODIFIERS, false, false);
-                if (l.getAttributes().hasAttribute(attribute)) {
-                    l.getAttributeInstance(attribute).removeModifier(modifier);
-                }
-            });
         }
         source.playSound(USounds.ENTITY_PLAYER_UNICORN_TELEPORT, 1);
     }

src/main/java/com/minelittlepony/unicopia/entity/AttributeContainer.java

@@ -0,0 +1,71 @@
+package com.minelittlepony.unicopia.entity;
+
+import java.util.Map;
+import java.util.UUID;
+
+import org.jetbrains.annotations.Nullable;
+
+import it.unimi.dsi.fastutil.floats.Float2ObjectFunction;
+import net.minecraft.entity.attribute.EntityAttribute;
+import net.minecraft.entity.attribute.EntityAttributeInstance;
+import net.minecraft.entity.attribute.EntityAttributeModifier;
+import net.minecraft.util.math.MathHelper;
+
+public interface AttributeContainer {
+    @Nullable
+    EntityAttributeInstance getAttributeInstance(EntityAttribute attribute);
+
+    default void updateAttributeModifier(UUID id, EntityAttribute attribute, float desiredValue, Float2ObjectFunction<EntityAttributeModifier> modifierSupplier, boolean permanent) {
+        @Nullable
+        EntityAttributeInstance instance = getAttributeInstance(attribute);
+        if (instance == null) {
+            return;
+        }
+
+        @Nullable
+        EntityAttributeModifier modifier = instance.getModifier(id);
+
+        if (!MathHelper.approximatelyEquals(desiredValue, modifier == null ? 0 : modifier.getValue())) {
+            if (modifier != null) {
+                instance.removeModifier(modifier);
+            }
+
+            if (desiredValue != 0) {
+                if (permanent) {
+                    instance.addPersistentModifier(modifierSupplier.get(desiredValue));
+                } else {
+                    instance.addTemporaryModifier(modifierSupplier.get(desiredValue));
+                }
+            }
+        }
+    }
+
+    default void applyAttributeModifiers(Map<EntityAttribute, EntityAttributeModifier> modifiers, boolean permanent, boolean apply) {
+        modifiers.forEach((attribute, modifier) -> {
+            applyAttributeModifier(attribute, modifier, permanent, apply);
+        });
+    }
+
+    default void applyAttributeModifier(EntityAttribute attribute, EntityAttributeModifier modifier, boolean permanent, boolean apply) {
+        @Nullable
+        EntityAttributeInstance instance = getAttributeInstance(attribute);
+        if (instance == null) {
+            return;
+        }
+
+        @Nullable
+        boolean present = instance.hasModifier(modifier);
+
+        if (present != apply) {
+            if (apply) {
+                if (permanent) {
+                    instance.addPersistentModifier(modifier);
+                } else {
+                    instance.addTemporaryModifier(modifier);
+                }
+            } else {
+                instance.removeModifier(modifier);
+            }
+        }
+    }
+}

src/main/java/com/minelittlepony/unicopia/entity/Living.java

@@ -40,7 +40,6 @@ import com.minelittlepony.unicopia.projectile.ProjectileImpactListener;
 import com.minelittlepony.unicopia.server.world.DragonBreathStore;
 import com.minelittlepony.unicopia.util.*;
 
-import it.unimi.dsi.fastutil.floats.Float2ObjectFunction;
 import net.fabricmc.fabric.api.util.TriState;
 import net.minecraft.block.BlockState;
 import net.minecraft.enchantment.Enchantment;
@@ -48,7 +47,6 @@ import net.minecraft.enchantment.EnchantmentHelper;
 import net.minecraft.entity.*;
 import net.minecraft.entity.attribute.EntityAttribute;
 import net.minecraft.entity.attribute.EntityAttributeInstance;
-import net.minecraft.entity.attribute.EntityAttributeModifier;
 import net.minecraft.entity.damage.DamageSource;
 import net.minecraft.entity.damage.DamageTypes;
 import net.minecraft.entity.mob.HostileEntity;
@@ -73,7 +71,7 @@ import net.minecraft.util.math.Direction;
 import net.minecraft.util.math.MathHelper;
 import net.minecraft.util.math.Vec3d;
 
-public abstract class Living<T extends LivingEntity> implements Equine<T>, Caster<T> {
+public abstract class Living<T extends LivingEntity> implements Equine<T>, Caster<T>, AttributeContainer {
     protected final T entity;
 
     private final SpellInventory spells;
@@ -253,29 +251,9 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
         transportation.tick();
     }
 
-    public void updateAttributeModifier(UUID id, EntityAttribute attribute, float desiredValue, Float2ObjectFunction<EntityAttributeModifier> modifierSupplier, boolean permanent) {
+    @Override
-        @Nullable
+    public final @Nullable EntityAttributeInstance getAttributeInstance(EntityAttribute attribute) {
-        EntityAttributeInstance instance = asEntity().getAttributeInstance(attribute);
+        return asEntity().getAttributeInstance(attribute);
-        if (instance == null) {
-            return;
-        }
-
-        @Nullable
-        EntityAttributeModifier modifier = instance.getModifier(id);
-
-        if (!MathHelper.approximatelyEquals(desiredValue, modifier == null ? 0 : modifier.getValue())) {
-            if (modifier != null) {
-                instance.removeModifier(modifier);
-            }
-
-            if (desiredValue != 0) {
-                if (permanent) {
-                    instance.addPersistentModifier(modifierSupplier.get(desiredValue));
-                } else {
-                    instance.addTemporaryModifier(modifierSupplier.get(desiredValue));
-                }
-            }
-        }
     }
 
     public boolean canBeSeenBy(Entity entity) {

src/main/java/com/minelittlepony/unicopia/item/AlicornAmuletItem.java

@@ -23,6 +23,9 @@ import com.minelittlepony.unicopia.server.world.UnicopiaWorldProperties;
 import com.minelittlepony.unicopia.util.VecHelper;
 
 import it.unimi.dsi.fastutil.floats.Float2ObjectFunction;
+import it.unimi.dsi.fastutil.objects.Object2FloatMap;
+import it.unimi.dsi.fastutil.objects.Object2FloatMaps;
+import it.unimi.dsi.fastutil.objects.Object2FloatOpenHashMap;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
 import net.fabricmc.fabric.api.item.v1.FabricItemSettings;
@@ -56,13 +59,13 @@ import net.minecraft.world.World.ExplosionSourceType;
 
 public class AlicornAmuletItem extends AmuletItem implements ItemTracker.Trackable, ItemImpl.ClingyItem, ItemImpl.GroundTickCallback {
     private static final UUID EFFECT_UUID = UUID.fromString("c0a870f5-99ef-4716-a23e-f320ee834b26");
-    private static final Map<EntityAttribute, Float> EFFECT_SCALES = Map.of(
+    private static final Object2FloatMap<EntityAttribute> EFFECT_SCALES = Object2FloatMaps.unmodifiable(new Object2FloatOpenHashMap<>(Map.of(
             EntityAttributes.GENERIC_ATTACK_DAMAGE, 0.2F,
             EntityAttributes.GENERIC_ATTACK_KNOCKBACK, 0.05F,
             EntityAttributes.GENERIC_ATTACK_SPEED, 0.2F,
             EntityAttributes.GENERIC_ARMOR_TOUGHNESS, 0.001F,
             EntityAttributes.GENERIC_ARMOR, 0.01F
-    );
+    )));
     private static final Float2ObjectFunction<EntityAttributeModifier> EFFECT_FACTORY = v -> {
         return new EntityAttributeModifier(EFFECT_UUID, "Alicorn Amulet Modifier", v, EntityAttributeModifier.Operation.ADDITION);
     };
@@ -173,8 +176,8 @@ public class AlicornAmuletItem extends AmuletItem implements ItemTracker.Trackab
     }
 
     public static void updateAttributes(Living<?> wearer, float effectScale) {
-        EFFECT_SCALES.entrySet().forEach(attribute -> {
+        EFFECT_SCALES.object2FloatEntrySet().forEach(entry -> {
-            wearer.updateAttributeModifier(EFFECT_UUID, attribute.getKey(), attribute.getValue() * effectScale, EFFECT_FACTORY, false);
+            wearer.updateAttributeModifier(EFFECT_UUID, entry.getKey(), entry.getFloatValue() * effectScale, EFFECT_FACTORY, false);
         });
     }