From eb5fee86923b2f8fd1b3f1cd72b5bd071cc50fd8 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Thu, 6 Jul 2023 17:27:03 +0100
Subject: [PATCH] Fixed #134

---
 .../unicopia/entity/Living.java               | 30 +++++++++++++++++++
 .../unicopia/item/AlicornAmuletItem.java      | 22 ++++----------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/main/java/com/minelittlepony/unicopia/entity/Living.java b/src/main/java/com/minelittlepony/unicopia/entity/Living.java
index 2bf7cbf2..7977898e 100644
--- a/src/main/java/com/minelittlepony/unicopia/entity/Living.java
+++ b/src/main/java/com/minelittlepony/unicopia/entity/Living.java
@@ -32,7 +32,11 @@ import com.minelittlepony.unicopia.server.world.DragonBreathStore;
 import com.minelittlepony.unicopia.trinkets.TrinketsDelegate;
 import com.minelittlepony.unicopia.util.*;
 
+import it.unimi.dsi.fastutil.floats.Float2ObjectFunction;
 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.data.*;
 import net.minecraft.entity.player.PlayerEntity;
@@ -46,6 +50,7 @@ import net.minecraft.server.world.ServerWorld;
 import net.minecraft.sound.SoundEvents;
 import net.minecraft.util.Hand;
 import net.minecraft.util.math.BlockPos;
+import net.minecraft.util.math.MathHelper;
 import net.minecraft.util.math.Vec3d;
 
 public abstract class Living<T extends LivingEntity> implements Equine<T>, Caster<T>, Transmittable {
@@ -214,6 +219,31 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
         updateDragonBreath();
     }
 
+    public void updateAttributeModifier(UUID id, EntityAttribute attribute, float desiredValue, Float2ObjectFunction<EntityAttributeModifier> modifierSupplier, boolean permanent) {
+        @Nullable
+        EntityAttributeInstance instance = 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) {
         return !isInvisible()
             && getSpellSlot()
diff --git a/src/main/java/com/minelittlepony/unicopia/item/AlicornAmuletItem.java b/src/main/java/com/minelittlepony/unicopia/item/AlicornAmuletItem.java
index 1afa15e9..87dd68f1 100644
--- a/src/main/java/com/minelittlepony/unicopia/item/AlicornAmuletItem.java
+++ b/src/main/java/com/minelittlepony/unicopia/item/AlicornAmuletItem.java
@@ -17,6 +17,7 @@ import com.minelittlepony.unicopia.particle.UParticles;
 import com.minelittlepony.unicopia.trinkets.TrinketsDelegate;
 import com.minelittlepony.unicopia.util.VecHelper;
 
+import it.unimi.dsi.fastutil.floats.Float2ObjectFunction;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
 import net.fabricmc.fabric.api.item.v1.FabricItemSettings;
@@ -55,6 +56,9 @@ public class AlicornAmuletItem extends AmuletItem implements ItemTracker.Trackab
             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);
+    };
 
     public AlicornAmuletItem(FabricItemSettings settings) {
         super(settings, 0, ImmutableMultimap.of());
@@ -159,12 +163,7 @@ public class AlicornAmuletItem extends AmuletItem implements ItemTracker.Trackab
         }
 
         EFFECT_SCALES.keySet().forEach(attribute -> {
-            EntityAttributeInstance instance = entity.getAttributeInstance(attribute);
-            @Nullable
-            EntityAttributeModifier modifier = instance.getModifier(EFFECT_UUID);
-            if (modifier != null) {
-                instance.removeModifier(modifier);
-            }
+            wearer.updateAttributeModifier(EFFECT_UUID, attribute, 0F, EFFECT_FACTORY, false);
         });
     }
 
@@ -268,16 +267,7 @@ public class AlicornAmuletItem extends AmuletItem implements ItemTracker.Trackab
         if (fullSecond) {
             EFFECT_SCALES.entrySet().forEach(attribute -> {
                 float seconds = (float)attachedTicks / ItemTracker.SECONDS;
-                EntityAttributeInstance instance = living.asEntity().getAttributeInstance(attribute.getKey());
-                @Nullable
-                EntityAttributeModifier modifier = instance.getModifier(EFFECT_UUID);
-                float desiredValue = attribute.getValue() * seconds;
-                if (!MathHelper.approximatelyEquals(desiredValue, modifier == null ? 0 : modifier.getValue())) {
-                    if (modifier != null) {
-                        instance.removeModifier(modifier);
-                    }
-                    instance.addTemporaryModifier(new EntityAttributeModifier(EFFECT_UUID, "Alicorn Amulet Modifier", attribute.getValue() * seconds, EntityAttributeModifier.Operation.ADDITION));
-                }
+                living.updateAttributeModifier(EFFECT_UUID, attribute.getKey(), attribute.getValue() * seconds, EFFECT_FACTORY, false);
             });
         }
     }